数据集修改

整体改动目标是：从“RGB + t + s，然后网络内部生成偏振图”改成“直接读取 PolaRGB 数据集里的 RGB + 4 张真实偏振图，再送进网络”。

1. 数据集加载
改的是：

basicsr/data/paired_image_dataset.py

新增/改造了 Dataset_PolaRGBImage。

原来代码按这种结构读：

dataroot_lq dataroot_gt dataroot_t dataroot_s

也就是 lq + gt + t + s。

现在改为支持 HuggingFace PolaRGB 官方结构：

PolaRGB/train/easy/input/00/0000_rgb.png PolaRGB/train/easy/input/00/0000_000.png PolaRGB/train/easy/input/00/0000_045.png PolaRGB/train/easy/input/00/0000_090.png PolaRGB/train/easy/input/00/0000_135.png PolaRGB/train/easy/gt/00/0000_rgb.png

loader 会自动做：

• 找所有 *_rgb.png 作为 RGB 输入
• 用同一个编号找 _000 / _045 / _090 / _135
• 找对应 scene 下的 GT RGB
• train 同时读 easy + hard
• val 读 test/input 和 test/gt
• 对 6 张图同步 padding、random crop、augmentation
• 返回：

{ 'lq': rgb, 'gt': gt, 'p0': img_000, 'p45': img_045, 'p90': img_090, 'p135': img_135 }

2. 网络输入
改的是：

basicsr/models/archs/PolarFMamba_arch.py

原来 PolarFMamba.forward 是：

forward(x, t, s)

内部调用：

generate_polarized_images(x, t, s)

也就是用 RGB、透射图、语义图合成偏振图。

现在改成：

forward(x, p0, p45, p90, p135)

删除了训练流程里的偏振生成逻辑。四张真实偏振图直接进入 PolarFrequencyDecomposition。

同时偏振通道数从原来的 3 张 RGB 偏振图：

3 * 3 = 9 channels

改成 4 张 RGB 偏振图：

4 * 3 = 12 channels

所以加了：

polar_in_channels: 12

网络里对应：

nn.Conv2d(polar_in_channels, 3, 1)

3. 训练数据流
改的是：

basicsr/models/image_restoration_model.py

新增了：

feed_train_polargbdata(...) optimize_parameters_polargb(...)

训练时把：

lq, gt, p0, p45, p90, p135

放到 GPU，然后调用：

self.net_g(self.lq, self.p0, self.p45, self.p90, self.p135)

验证时也支持这五输入，不再只支持 lq + t + s。

另外 mixup 原来只混：

gt, lq

现在 PolaRGB 分支会同步混：

gt, lq, p0, p45, p90, p135

否则 RGB 和偏振图会错位。

4. PolarFMamba_train.py
改的是：

basicsr/PolarFMamba_train.py

训练循环里原来取：

lq = train_data['lq'] gt = train_data['gt'] t = train_data['t'] s = train_data['s']

现在改成：

lq = train_data['lq'] gt = train_data['gt'] p0 = train_data['p0'] p45 = train_data['p45'] p90 = train_data['p90'] p135 = train_data['p135']

progressive crop / mini batch 抽样也同步处理四张偏振图。

最后调用：

model.feed_train_polargbdata(...) model.optimize_parameters_polargb(...)

5. 配置 yml
Options/PolarFMamba_cityscapes_v5.yml 改成：

type: Dataset_PolaRGBImage pola_rgb_layout: official dataroot_gt: /home/student_account/student27/PolaRGB/train dataroot_lq: /home/student_account/student27/PolaRGB/train subsets: [easy, hard]

验证集：

dataroot_gt: /home/student_account/student27/PolaRGB/test dataroot_lq: /home/student_account/student27/PolaRGB/test

网络：

network_g: type: PolarFMamba d_model: 32 polar_in_channels: 12

四卡：

num_gpu: 4

总结
原始流程：

RGB + t + s -> 网络内部随机/物理模型生成偏振图 -> PolarFMamba

现在流程：

PolaRGB 数据集 RGB + 000 + 045 + 090 + 135 -> 直接输入 PolarFMamba