Upload joint_transforms.py

joint_transforms.py

@@ -0,0 +1,72 @@
+# joint_transforms.py
+
+import math
+import random
+import numpy as np
+from PIL import Image, ImageFilter
+
+class Compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, img, gt):
+        for t in self.transforms:
+            img, gt = t(img, gt)
+        return img, gt
+
+class RandomScaleCrop(object):
+    """多尺度缩放裁剪（同时处理图像和标签）"""
+    def __init__(self, base_size=352, crop_size=352, scale_factor=[0.75, 1.0, 1.25]):
+        self.base_size = base_size
+        self.crop_size = crop_size
+        self.scale_factor = scale_factor
+
+    def __call__(self, img, gt):
+        # 随机选择缩放比例
+        sf = random.choice(self.scale_factor)
+        new_size = int(self.base_size * sf)
+
+        # 缩放
+        img = img.resize((new_size, new_size), Image.BILINEAR)
+        gt = gt.resize((new_size, new_size), Image.NEAREST)
+
+        # 随机裁剪
+        x = random.randint(0, new_size - self.crop_size)
+        y = random.randint(0, new_size - self.crop_size)
+        img = img.crop((x, y, x+self.crop_size, y+self.crop_size))
+        gt = gt.crop((x, y, x+self.crop_size, y+self.crop_size))
+
+        return img, gt
+
+class RandomRotate(object):
+    """随机旋转（保持图像和标签同步）"""
+    def __init__(self, degree=30):
+        self.degree = degree
+
+    def __call__(self, img, gt):
+        rotate_degree = random.uniform(-self.degree, self.degree)
+        img = img.rotate(rotate_degree, Image.BILINEAR)
+        gt = gt.rotate(rotate_degree, Image.NEAREST)
+        return img, gt
+
+class RandomGaussianBlur(object):
+    """随机高斯模糊（仅对图像处理）"""
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, img, gt):
+        if random.random() < self.p:
+            img = img.filter(ImageFilter.GaussianBlur(
+                radius=random.uniform(0.5, 2.0)))
+        return img, gt
+
+class RandomHorizontallyFlip(object):
+    """随机水平翻转"""
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, img, gt):
+        if random.random() < self.p:
+            img = img.transpose(Image.FLIP_LEFT_RIGHT)
+            gt = gt.transpose(Image.FLIP_LEFT_RIGHT)
+        return img, gt