# 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