feat(actpred): 训练时安全图像增强(ColorJitter),支持任意 frame_len
训练集启用 --augment 后对每帧施加亮度/对比度/饱和度/色调变换, 同一样本内所有帧共享参数(模拟连续帧光照不变),验证集不增强。 仅颜色空间变换,无空间变换,动作标签不受影响。 删除 t1.mt(@TEST 块无对应插件,录制时报错)。
This commit is contained in:
1 parent
85133868f3
commit
5af863a2c5
5 files changed
+136
-23
No files matched your search
+14
-1
@@ -1,4 +1,5 @@
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
@@ -7,6 +8,7 @@ from PIL import Image
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision import transforms
|
||||
from torchvision.transforms import functional as TF
|
||||
|
||||
|
||||
def _natural_key(p):
|
||||
@@ -19,12 +21,13 @@ class ActionDataset(Dataset):
|
||||
image_width: int = 640, image_height: int = 480,
|
||||
normalize_actions: bool = False,
|
||||
action_stats_path: str = None, episode_indices: list = None,
|
||||
frame_len: int = 3):
|
||||
frame_len: int = 3, augment: bool = False):
|
||||
self.history_len = history_len
|
||||
self.action_dim = action_dim
|
||||
self.image_width = image_width
|
||||
self.image_height = image_height
|
||||
self.frame_len = frame_len
|
||||
self.augment = augment
|
||||
|
||||
self.transform = transforms.Compose([
|
||||
transforms.Resize((image_height, image_width)),
|
||||
@@ -92,8 +95,18 @@ class ActionDataset(Dataset):
|
||||
def __getitem__(self, idx):
|
||||
image_paths, action_history, target = self.samples[idx]
|
||||
frames = []
|
||||
if self.augment:
|
||||
bf = random.uniform(0.8, 1.2)
|
||||
cf = random.uniform(0.8, 1.2)
|
||||
sf = random.uniform(0.8, 1.2)
|
||||
hf = random.uniform(-0.05, 0.05)
|
||||
for p in image_paths:
|
||||
image = Image.open(p).convert("RGB")
|
||||
if self.augment:
|
||||
image = TF.adjust_brightness(image, bf)
|
||||
image = TF.adjust_contrast(image, cf)
|
||||
image = TF.adjust_saturation(image, sf)
|
||||
image = TF.adjust_hue(image, hf)
|
||||
frames.append(self.transform(image))
|
||||
images = torch.stack(frames) # (F, 3, S, S)
|
||||
action_history = torch.tensor(action_history, dtype=torch.float32)
|
||||
|
||||
Reference in new issue
Block a user