dataset.py

import os
import random
from PIL import Image, ImageFilter, ImageEnhance
import torch
from torch.utils.data import Dataset
import torchvision.transforms.functional as TF
import numpy as np


class HandwriteDataset(Dataset):
    def __init__(self, img_dir, label_dir=None, size=(576, 576), is_train=True, augment_level='medium'):
        """
        Args:
            img_dir (str): Path to the directory containing input images.
            label_dir (str, optional): Path to the directory containing label images.
            size (tuple): Target size (height, width) for resizing. 576 is divisible by 16.
            is_train (bool): Whether to apply data augmentation.
            augment_level (str): Augmentation intensity - 'light', 'medium', 'heavy'
        """
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.size = size
        self.is_train = is_train
        self.augment_level = augment_level
        self.images = sorted([f for f in os.listdir(img_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
        
        # Augmentation probability settings based on level
        self.aug_probs = {
            'light': {'flip': 0.3, 'rotate': 0.3, 'color': 0.2, 'noise': 0.1, 'blur': 0.1, 'crop': 0.2, 'affine': 0.1, 'perspective': 0.1},
            'medium': {'flip': 0.5, 'rotate': 0.5, 'color': 0.4, 'noise': 0.3, 'blur': 0.3, 'crop': 0.4, 'affine': 0.3, 'perspective': 0.2},
            'heavy': {'flip': 0.5, 'rotate': 0.7, 'color': 0.6, 'noise': 0.5, 'blur': 0.5, 'crop': 0.5, 'affine': 0.5, 'perspective': 0.4}
        }[augment_level]

    def __len__(self):
        return len(self.images)
    
    def _sync_random_crop(self, image, label, crop_ratio_range=(0.7, 0.95)):
        """Random crop with same region for both image and label"""
        w, h = image.size
        crop_ratio = random.uniform(*crop_ratio_range)
        new_w, new_h = int(w * crop_ratio), int(h * crop_ratio)
        
        left = random.randint(0, w - new_w)
        top = random.randint(0, h - new_h)
        
        image = image.crop((left, top, left + new_w, top + new_h))
        label = label.crop((left, top, left + new_w, top + new_h))
        return image, label
    
    def _sync_affine(self, image, label, degrees=15, translate=(0.1, 0.1), scale=(0.9, 1.1), shear=10):
        """Synchronized affine transformation"""
        angle = random.uniform(-degrees, degrees)
        translations = (random.uniform(-translate[0], translate[0]) * image.size[0],
                       random.uniform(-translate[1], translate[1]) * image.size[1])
        scale_factor = random.uniform(*scale)
        shear_angle = random.uniform(-shear, shear)
        
        image = TF.affine(image, angle=angle, translate=translations, scale=scale_factor, shear=shear_angle)
        label = TF.affine(label, angle=angle, translate=translations, scale=scale_factor, shear=shear_angle)
        return image, label
    
    def _sync_perspective(self, image, label, distortion_scale=0.2):
        """Synchronized perspective transformation"""
        width, height = image.size
        half_height = height // 2
        half_width = width // 2
        
        topleft = [
            int(random.uniform(0, distortion_scale * half_width)),
            int(random.uniform(0, distortion_scale * half_height))
        ]
        topright = [
            int(random.uniform(width - distortion_scale * half_width, width)),
            int(random.uniform(0, distortion_scale * half_height))
        ]
        botright = [
            int(random.uniform(width - distortion_scale * half_width, width)),
            int(random.uniform(height - distortion_scale * half_height, height))
        ]
        botleft = [
            int(random.uniform(0, distortion_scale * half_width)),
            int(random.uniform(height - distortion_scale * half_height, height))
        ]
        
        startpoints = [[0, 0], [width, 0], [width, height], [0, height]]
        endpoints = [topleft, topright, botright, botleft]
        
        image = TF.perspective(image, startpoints, endpoints)
        label = TF.perspective(label, startpoints, endpoints)
        return image, label
    
    def _add_gaussian_noise(self, image, mean=0, std_range=(0.01, 0.05)):
        """Add Gaussian noise to image only (not label)"""
        img_array = np.array(image).astype(np.float32) / 255.0
        std = random.uniform(*std_range)
        noise = np.random.normal(mean, std, img_array.shape)
        noisy = np.clip(img_array + noise, 0, 1)
        return Image.fromarray((noisy * 255).astype(np.uint8))
    
    def _adjust_color(self, image):
        """Adjust brightness, contrast, saturation for input image only"""
        # Brightness
        brightness_factor = random.uniform(0.8, 1.2)
        image = ImageEnhance.Brightness(image).enhance(brightness_factor)
        
        # Contrast
        contrast_factor = random.uniform(0.8, 1.2)
        image = ImageEnhance.Contrast(image).enhance(contrast_factor)
        
        # Saturation
        if random.random() > 0.5:
            saturation_factor = random.uniform(0.7, 1.3)
            image = ImageEnhance.Color(image).enhance(saturation_factor)
        
        return image
    
    def _add_blur(self, image, radius_range=(0.5, 1.5)):
        """Add Gaussian blur to input image only"""
        radius = random.uniform(*radius_range)
        return image.filter(ImageFilter.GaussianBlur(radius=radius))
    
    def _add_jpeg_compression(self, image, quality_range=(50, 85)):
        """Simulate JPEG compression artifacts"""
        import io
        quality = random.randint(*quality_range)
        buffer = io.BytesIO()
        image.save(buffer, format='JPEG', quality=quality)
        buffer.seek(0)
        return Image.open(buffer).convert('RGB')

    def __getitem__(self, idx):
        img_name = self.images[idx]
        img_path = os.path.join(self.img_dir, img_name)
        image = Image.open(img_path).convert("RGB")

        label = None
        if self.label_dir:
            # Construct label path. Input is .jpg, label is .png
            base_name = os.path.splitext(img_name)[0]
            label_name = base_name + ".png"
            label_path = os.path.join(self.label_dir, label_name)
            
            if not os.path.exists(label_path):
                label_path = os.path.join(self.label_dir, base_name + ".jpg")
            
            if os.path.exists(label_path):
                label = Image.open(label_path).convert("RGB")
            else:
                raise FileNotFoundError(f"Label not found for {img_name}")

        # Data Augmentation for training (applied BEFORE resize for better quality)
        if self.is_train and label:
            probs = self.aug_probs
            
            # ========== Synchronized geometric transforms ==========
            # Random crop (same region for both)
            if random.random() < probs['crop']:
                image, label = self._sync_random_crop(image, label)
            
            # Random horizontal flip
            if random.random() < probs['flip']:
                image = TF.hflip(image)
                label = TF.hflip(label)
            
            # Random vertical flip
            if random.random() < probs['flip']:
                image = TF.vflip(image)
                label = TF.vflip(label)
            
            # Random rotation
            if random.random() < probs['rotate']:
                angle = random.uniform(-15, 15)
                image = TF.rotate(image, angle, fill=255)  # White fill for document
                label = TF.rotate(label, angle, fill=255)
            
            # Random affine transformation
            if random.random() < probs['affine']:
                image, label = self._sync_affine(image, label)
            
            # Random perspective transformation
            if random.random() < probs['perspective']:
                image, label = self._sync_perspective(image, label)
            
            # ========== Input-only transforms (simulate real-world degradation) ==========
            # Color/brightness adjustments
            if random.random() < probs['color']:
                image = self._adjust_color(image)
            
            # Gaussian noise
            if random.random() < probs['noise']:
                image = self._add_gaussian_noise(image)
            
            # Gaussian blur
            if random.random() < probs['blur']:
                image = self._add_blur(image)
            
            # JPEG compression artifacts (occasional)
            if random.random() < probs['noise'] * 0.5:
                image = self._add_jpeg_compression(image)

        # Resize to target size
        image = TF.resize(image, self.size)
        if label:
            label = TF.resize(label, self.size)

        # To Tensor
        image = TF.to_tensor(image)
        if label:
            label = TF.to_tensor(label)

        if label is not None:
            return image, label
        return image, img_name