Detectron2水下图像检测：在线图像增强集成问题求助

原代码核心问题 #

• 错误使用DatasetMapper参数：augmentations仅接受Detectron2官方的Augmentation子类实例，无法直接传入普通函数
• 变量未定义：map_enhance中annos变量未声明，需从输入的dataset_dict["annotations"]提取标注
• 缺失概率触发逻辑：未实现“以指定概率应用增强”的核心需求
• 冗余IO操作：重复读取图像+无意义深拷贝导致训练耗时激增

修正后的完整实现 #

import copy
import cv2
from detectron2.data import DatasetMapper, build_detection_train_loader
from detectron2.data.transforms import Augmentation, Transform
from detectron2.engine import DefaultTrainer
from detectron2.utils import utils

# 自定义带概率控制的水下图像增强类，适配Detectron2增强体系
class ProbabilisticUnderwaterEnhance(Augmentation):
    def __init__(self, prob=0.5):
        self.prob = prob  # 增强触发概率，范围0-1

    def get_transform(self, image):
        # 按设定概率随机触发增强
        if utils.random.random() < self.prob:
            enhanced_img = self._apply_underwater_enhance(image)
            return Transform(lambda img: enhanced_img)
        # 不触发时返回恒等变换，不修改图像
        return Transform(lambda img: img)

    def _apply_underwater_enhance(self, image):
        # 替换为你实际的水下图像增强逻辑（比如颜色校正、CLAHE、去雾等）
        # 示例：RGB空间CLAHE增强，适配水下图像低对比度问题
        img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8,8))
        # 对RGB三通道分别应用CLAHE
        img_rgb[:, :, 0] = clahe.apply(img_rgb[:, :, 0])
        img_rgb[:, :, 1] = clahe.apply(img_rgb[:, :, 1])
        img_rgb[:, :, 2] = clahe.apply(img_rgb[:, :, 2])
        # 转回Detectron2默认的BGR格式
        return cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)

# 自定义Trainer，正确集成实时增强
class TrainerEnhance(DefaultTrainer):
    @classmethod
    def build_train_loader(cls, cfg):
        # 复用Detectron2默认的训练增强流水线（随机翻转、缩放等）
        base_augmentations = cls.build_augmentation(cfg, is_train=True)
        # 开启增强时，加入自定义概率增强模块
        if cfg.INPUT.ENHANCE_IMAGE:
            base_augmentations.append(ProbabilisticUnderwaterEnhance(prob=cfg.INPUT.ENHANCE_PROB))
        
        # 初始化标准DatasetMapper，传入完整增强列表
        mapper = DatasetMapper(
            cfg,
            is_train=True,
            augmentations=base_augmentations,
            image_format=cfg.INPUT.FORMAT
        )
        return build_detection_train_loader(cfg, mapper=mapper)

关键优化说明 #

• 符合Detectron2规范：通过继承Augmentation类实现增强，完美融入官方数据加载流水线
• 概率控制：通过prob参数灵活设置增强触发比例，满足“一定概率实时增强”需求
• 性能优化：移除冗余的图像重复读取和深拷贝操作，降低训练耗时
• 可扩展性：_apply_underwater_enhance方法可直接替换为任意水下增强算法（如UDN、WaterNet等）

配置补充 #

需在训练配置中添加以下参数（可通过代码或配置文件设置）：

# 开启实时增强
cfg.INPUT.ENHANCE_IMAGE = True
# 设置增强触发概率（示例：70%的图像应用增强）
cfg.INPUT.ENHANCE_PROB = 0.7

内容的提问来源于stack exchange，提问作者muhammad ali