CBT_project/utils/helpers.py

import errno
import os
import time

import numpy as np
import matplotlib.pyplot as plt


def retry_robust(fn, *args, retries=20, delay=0.5, **kwargs):
    """對 ENOENT/EEXIST 重試：NFS 上輸出樹被外部刪除（或多 worker 併發建同一
    output 目錄）會有短暫的 ENOENT 窗口，重試可恢復；其他錯誤直接丟出。"""
    for i in range(retries):
        try:
            return fn(*args, **kwargs)
        except OSError as e:
            if e.errno not in (errno.ENOENT, errno.EEXIST) or i == retries - 1:
                raise
            time.sleep(delay)


def get_unique_filepath(path: str) -> str:
    """
    如果檔案已存在，自動加 _1, _2... 避免覆蓋
    """
    current_path = path
    count = 1

    base_name, file_ext = os.path.splitext(path)

    while os.path.exists(current_path):
        current_path = f"{base_name}_{count}{file_ext}"
        count += 1

    return current_path

def save_with_unique_name(folder, label_str, way, diameter_l, length_l, diameter_r, length_r, swarm_size, max_iter):
    
    base_name = f'{label_str}_{way}_L{diameter_l}_{length_l}_R{diameter_r}_{length_r}_{swarm_size}_{max_iter}.png'
    file_path = os.path.join(folder, base_name)
    
    count = 1
    while os.path.exists(file_path):
        file_name, file_ext = os.path.splitext(base_name)
        file_path = os.path.join(folder, f"{file_name}_{count}{file_ext}")
        count += 1

    return file_path

def pad_rgb_to_shape(rgb, target_hw, pad_value=0.0):
    """Pad RGB image (H,W,3) to target (H,W), centered."""
    h, w, c = rgb.shape
    H, W = target_hw
    assert c == 3

    if h > H or w > W:
        raise ValueError(f"target {target_hw} smaller than rgb {(h,w)}")

    pad_h1 = (H - h) // 2
    pad_h2 = H - h - pad_h1
    pad_w1 = (W - w) // 2
    pad_w2 = W - w - pad_w1

    return np.pad(
        rgb,
        ((pad_h1, pad_h2), (pad_w1, pad_w2), (0, 0)),
        mode="constant",
        constant_values=pad_value
    )