module/image_utils.py

import numpy as np
import torch


def apply_gamma_tensor(img_tensor, low=0.15, high=0.85, target=0.5):
    """对图像张量进行Gamma校正，当亮度超出[low, high]范围时自动调整到target亮度"""
    arr = img_tensor.detach().cpu().numpy()
    arr = (arr * 0.5) + 0.5
    arr = np.clip(arr.squeeze(0).transpose(1, 2, 0), 0.0, 1.0)

    lum = float(np.dot(arr[..., :3], [0.299, 0.587, 0.114]).mean())

    if low <= lum <= high:
        return img_tensor, lum, False, 1.0

    eps = 1e-6
    gamma = np.log(target + eps) / np.log(lum + eps)
    gamma = float(np.clip(gamma, 0.5, 2.5))

    arr = np.clip(arr ** gamma, 0.0, 1.0)
    tensor = torch.from_numpy(((arr - 0.5) / 0.5).transpose(2, 0, 1)).unsqueeze(0).float()
    return tensor, lum, True, gamma