import os import glob import cv2 import numpy as np import torch import random from torch.utils.data import Dataset import albumentations as A from albumentations.pytorch import ToTensorV2 class YOLODataset(Dataset): def __init__(self, img_dir, label_dir, img_size=640, is_train=True): self.img_dir = img_dir self.label_dir = label_dir self.img_size = img_size self.is_train = is_train self.use_mosaic = is_train # 新增标志位，默认训练时开启 # 支持多种图片格式 self.img_files = sorted( glob.glob(os.path.join(img_dir, "*.jpg")) + glob.glob(os.path.join(img_dir, "*.png")) + glob.glob(os.path.join(img_dir, "*.jpeg")) ) # --- 1. 对齐 Ultralytics 的 Albumentations 配置 --- # default.yaml: hsv_h: 0.015, hsv_s: 0.7, hsv_v: 0.4 # OpenCV Hue range is [0, 179], Sat/Val is [0, 255] h_limit = int(0.015 * 179) # ~2 s_limit = int(0.7 * 255) # ~178 v_limit = int(0.4 * 255) # ~102 # default.yaml: translate: 0.1, scale: 0.5 (0.5~1.5), degrees: 0.0 if is_train: self.transform = A.Compose([ # 几何增强 (Mosaic 之后再做一次微调，或者处理非 Mosaic 的情况) # 注意：Mosaic 输出已经是大图，这里主要负责最后的随机扰动 A.Affine( scale=(0.5, 1.5), # scale: 0.5 translate_percent=(0.1, 0.1), # translate: 0.1 rotate=(-0, 0), # degrees: 0.0 (COCO default) shear=(-0, 0), # shear: 0.0 p=0.5 ), # 色彩增强 (严格对齐 default.yaml) A.HueSaturationValue( hue_shift_limit=h_limit, sat_shift_limit=s_limit, val_shift_limit=v_limit, p=0.5 ), A.Blur(p=0.01), A.MedianBlur(p=0.01), A.ToGray(p=0.01), A.CLAHE(p=0.01), # 翻转 A.HorizontalFlip(p=0.5), # fliplr: 0.5 # 最终处理 A.Resize(img_size, img_size), # 确保最后尺寸一致 A.Normalize(mean=(0, 0, 0), std=(1, 1, 1), max_pixel_value=255.0), ToTensorV2() ], bbox_params=A.BboxParams(format='yolo', min_visibility=0.0, label_fields=['class_labels'])) else: # 验证集：Letterbox (保持比例填充) self.transform = A.Compose([ A.LongestMaxSize(max_size=img_size), A.PadIfNeeded(min_height=img_size, min_width=img_size, border_mode=cv2.BORDER_CONSTANT), A.Normalize(mean=(0, 0, 0), std=(1, 1, 1), max_pixel_value=255.0), ToTensorV2() ], bbox_params=A.BboxParams(format='yolo', min_visibility=0.1, label_fields=['class_labels'])) def close_mosaic(self): """关闭 Mosaic 增强""" self.use_mosaic = False print("Mosaic augmentation disabled.") def __len__(self): return len(self.img_files) def load_image(self, index): """加载单张图片并调整长边到 img_size""" img_path = self.img_files[index] img = cv2.imread(img_path) if img is None: raise FileNotFoundError(f"Image not found: {img_path}") h, w = img.shape[:2] r = self.img_size / max(h, w) if r != 1: img = cv2.resize(img, (int(w * r), int(h * r)), interpolation=cv2.INTER_LINEAR) return img, (h, w), img.shape[:2] # img, original_hw, resized_hw def load_label(self, index, img_shape): """加载标签并归一化""" img_path = self.img_files[index] label_path = self._get_label_path(img_path) h, w = img_shape labels = [] if os.path.exists(label_path): with open(label_path, 'r') as f: for line in f: parts = line.strip().split() if len(parts) >= 5: cls = int(parts[0]) bx, by, bw, bh = map(float, parts[1:5]) labels.append([cls, bx, by, bw, bh]) return np.array(labels, dtype=np.float32) if labels else np.zeros((0, 5), dtype=np.float32) def load_mosaic(self, index): """ 实现 YOLO 的 Mosaic 增强 (4张图拼成一张) """ labels4 = [] s = self.img_size # 修复: 列表推导式需要遍历两次以生成 yc 和 xc yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in [-s // 2, -s // 2]] # 随机选3个额外索引 indices = [index] + [random.randint(0, len(self.img_files) - 1) for _ in range(3)] random.shuffle(indices) # 初始化大图 (2x size) img4 = np.full((s * 2, s * 2, 3), 114, dtype=np.uint8) for i, idx in enumerate(indices): # 加载图片 img, _, (h, w) = self.load_image(idx) # 放置位置: top-left, top-right, bottom-left, bottom-right if i == 0: # top left x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h elif i == 1: # top right x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h elif i == 2: # bottom left x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(yc + h, s * 2) x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h) elif i == 3: # bottom right x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(yc + h, s * 2) x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h) # 贴图 img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # type: ignore padw = x1a - x1b # type: ignore padh = y1a - y1b # type: ignore # 处理标签 labels = self.load_label(idx, (h, w)) if labels.size > 0: # Normalized xywh -> Pixel xywh labels[:, 1] = labels[:, 1] * w labels[:, 2] = labels[:, 2] * h labels[:, 3] = labels[:, 3] * w labels[:, 4] = labels[:, 4] * h # xywh -> xyxy (Pixel) xyxy = np.copy(labels) xyxy[:, 1] = labels[:, 1] - labels[:, 3] / 2 + padw xyxy[:, 2] = labels[:, 2] - labels[:, 4] / 2 + padh xyxy[:, 3] = labels[:, 1] + labels[:, 3] / 2 + padw xyxy[:, 4] = labels[:, 2] + labels[:, 4] / 2 + padh labels4.append(xyxy) # Concat labels if len(labels4): labels4 = np.concatenate(labels4, 0) # Clip to mosaic image border np.clip(labels4[:, 1:], 0, 2 * s, out=labels4[:, 1:]) # 转换回 Normalized xywh (相对于 2*s 的大图) # Albumentations 需要 normalized xywh new_labels = np.copy(labels4) w_mosaic, h_mosaic = s * 2, s * 2 # xyxy -> xywh new_labels[:, 1] = (labels4[:, 1] + labels4[:, 3]) / 2 / w_mosaic new_labels[:, 2] = (labels4[:, 2] + labels4[:, 4]) / 2 / h_mosaic new_labels[:, 3] = (labels4[:, 3] - labels4[:, 1]) / w_mosaic new_labels[:, 4] = (labels4[:, 4] - labels4[:, 2]) / h_mosaic return img4, new_labels else: return img4, np.zeros((0, 5)) def __getitem__(self, index): try: if self.is_train: # 修改判断逻辑：同时检查 is_train 和 use_mosaic if self.use_mosaic and random.random() < 1.0: img, labels = self.load_mosaic(index) else: img, _, _ = self.load_image(index) h, w = img.shape[:2] labels = self.load_label(index, (h, w)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # type: ignore else: img = cv2.imread(self.img_files[index]) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # type: ignore h, w = img.shape[:2] labels = self.load_label(index, (h, w)) # --- 修复开始: 清洗边界框 --- # labels 格式: [cls, x, y, w, h] (Normalized) valid_bboxes = [] valid_labels = [] h_img, w_img = img.shape[:2] for i in range(len(labels)): cls = labels[i, 0] x, y, w, h = labels[i, 1:] # 1. 限制在 [0, 1] 范围内 (处理 Mosaic 裁剪产生的越界) x1 = np.clip(x - w / 2, 0, 1) y1 = np.clip(y - h / 2, 0, 1) x2 = np.clip(x + w / 2, 0, 1) y2 = np.clip(y + h / 2, 0, 1) # 2. 重新计算宽高 w_new = x2 - x1 h_new = y2 - y1 # 3. 重新计算中心点 x_new = x1 + w_new / 2 y_new = y1 + h_new / 2 # 4. 过滤掉极小的框 (例如小于 2 个像素) if w_new * w_img > 2 and h_new * h_img > 2: valid_bboxes.append([x_new, y_new, w_new, h_new]) valid_labels.append(cls) if len(valid_bboxes) == 0: # 如果这张图的所有框都被过滤掉了，尝试下一张 return self.__getitem__((index + 1) % len(self)) bboxes = valid_bboxes class_labels = valid_labels # --- 修复结束 --- # 应用增强 transformed = self.transform(image=img, bboxes=bboxes, class_labels=class_labels) image = transformed['image'] bboxes = transformed['bboxes'] class_labels = transformed['class_labels'] # 构建 Target Tensor n = len(bboxes) targets = torch.zeros((n, 6)) if n > 0: targets[:, 1] = torch.tensor(class_labels) targets[:, 2:] = torch.tensor(bboxes) return image, targets, self.img_files[index] except Exception as e: # 打印更详细的错误信息以便调试，但不要中断训练 # print(f"Error loading data {index}: {e}") return self.__getitem__((index + 1) % len(self)) def _get_label_path(self, img_path): filename = os.path.basename(img_path).rsplit('.', 1)[0] + ".txt" return os.path.join(self.label_dir, filename) @staticmethod def collate_fn(batch): imgs, targets, paths = zip(*batch) imgs = torch.stack(imgs, 0) new_targets = [] for i, t in enumerate(targets): if t.shape[0] > 0: t[:, 0] = i new_targets.append(t) if new_targets: targets = torch.cat(new_targets, 0) else: targets = torch.zeros((0, 6)) return imgs, targets, paths