Upload 13 files

checkpoints/discriminator_epoch_10.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a93bbd6522431f63abe8c6821a17efba7e8b7751314a69db2caf2a14e3bda5e
+size 1106902

checkpoints/discriminator_epoch_15.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f5af7468f085e6e6d8b8058a0d629151243199d0a994be72f1f9e270d241e77f
+size 1106902

checkpoints/discriminator_epoch_2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:58d206e029ec8e09a3bda2034cbd1a1170848b5cdcc0861f280890186aa3043c
+size 1106807

checkpoints/generator_epoch_10.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a8feff031f9337d18f659075b7a0db41f19d782f4367c026a4cc374f8de2232
+size 6096658

checkpoints/generator_epoch_15.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce25de423b699eef2ffcc51585b049a95efae370520c640834e892a391070654
+size 6096658

checkpoints/generator_epoch_2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:95f21a754252b0d804fe63660e802f2c7fe435d599d8ba431f58e420c704947d
+size 6096580

dataset.py

@@ -0,0 +1,273 @@
+import os
+import glob
+from PIL import Image
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+import random # Needed for random cropping
+
+# --- Updated SRDataset Class ---
+class SRDataset(Dataset):
+    """
+    Custom Dataset for Super-Resolution.
+    Loads HR/LR pairs and returns fixed-size patches.
+    """
+    def __init__(self, hr_dir, lr_dir, scale_factor, patch_size_lr=48, transform=None):
+        """
+        Args:
+            hr_dir (str): Directory with all HR images.
+            lr_dir (str): Directory with all LR images (corresponding to hr_dir).
+            scale_factor (int): The upscaling factor.
+            patch_size_lr (int): The size (height and width) of the LR patch to crop.
+                                 HR patch size will be patch_size_lr * scale_factor.
+            transform (callable, optional): Optional transform (e.g., data augmentation like flips).
+        """
+        super(SRDataset, self).__init__() # Call parent constructor
+        self.hr_dir = hr_dir
+        self.lr_dir = lr_dir
+        self.scale_factor = scale_factor
+        self.patch_size_lr = patch_size_lr
+        self.patch_size_hr = patch_size_lr * scale_factor
+        self.transform = transform
+
+        # Find all image files (png, jpg, jpeg) in the LR directory
+        self.lr_image_files = sorted(
+            glob.glob(os.path.join(lr_dir, '*.png')) +
+            glob.glob(os.path.join(lr_dir, '*.jpg')) +
+            glob.glob(os.path.join(lr_dir, '*.jpeg'))
+        )
+
+        if not self.lr_image_files:
+            raise FileNotFoundError(f"No images found in LR directory: {lr_dir}. Check path and image extensions.")
+
+        # --- (Optional Verification Step - can be kept or removed) ---
+        if self.lr_image_files:
+            # ... (verification code from previous version can go here if desired) ...
+            pass
+
+        print(f"Found {len(self.lr_image_files)} image pairs in HR='{hr_dir}', LR='{lr_dir}'")
+        print(f"Using LR patch size: {self.patch_size_lr}x{self.patch_size_lr}, HR patch size: {self.patch_size_hr}x{self.patch_size_hr}")
+
+    def __len__(self):
+        return len(self.lr_image_files)
+
+    @staticmethod
+    def get_patch(lr_img, hr_img, patch_size_lr, scale_factor):
+        """
+        Randomly crops corresponding patches from LR and HR images.
+
+        Args:
+            lr_img (PIL.Image): Low-resolution image.
+            hr_img (PIL.Image): High-resolution image.
+            patch_size_lr (int): The desired height/width of the LR patch.
+            scale_factor (int): The upscaling factor.
+
+        Returns:
+            tuple: (lr_patch, hr_patch) PIL.Image objects.
+        """
+        lr_w, lr_h = lr_img.size
+        hr_w, hr_h = hr_img.size
+        patch_size_hr = patch_size_lr * scale_factor
+
+        # Ensure HR image dimensions are consistent with LR and scale factor
+        if hr_w != lr_w * scale_factor or hr_h != lr_h * scale_factor:
+            # Simple fallback: resize HR image to expected size if mismatch occurs
+            # This might happen with imperfect downscaling or odd original dimensions
+            # print(f"Warning: HR/LR size mismatch ({hr_img.size} vs {lr_img.size} * {scale_factor}). Resizing HR image.")
+            hr_img = hr_img.resize((lr_w * scale_factor, lr_h * scale_factor), resample=Image.BICUBIC)
+
+        # Choose random top-left corner for LR patch
+        # Ensure the patch fits within the image boundaries
+        if lr_w < patch_size_lr or lr_h < patch_size_lr:
+             # If LR image is smaller than patch size, resize LR and corresponding HR region
+             # This ensures __getitem__ always returns tensors of the target patch size
+             lr_img = lr_img.resize((max(lr_w, patch_size_lr), max(lr_h, patch_size_lr)), resample=Image.BICUBIC)
+             hr_img = hr_img.resize((lr_img.width * scale_factor, lr_img.height * scale_factor), resample=Image.BICUBIC)
+             lr_w, lr_h = lr_img.size # Update dimensions
+
+
+        lr_x = random.randrange(0, lr_w - patch_size_lr + 1)
+        lr_y = random.randrange(0, lr_h - patch_size_lr + 1)
+
+        # Calculate corresponding top-left corner for HR patch
+        hr_x = lr_x * scale_factor
+        hr_y = lr_y * scale_factor
+
+        # Crop patches
+        # PIL crop format is (left, upper, right, lower)
+        lr_patch = lr_img.crop((lr_x, lr_y, lr_x + patch_size_lr, lr_y + patch_size_lr))
+        hr_patch = hr_img.crop((hr_x, hr_y, hr_x + patch_size_hr, hr_y + patch_size_hr))
+
+        return lr_patch, hr_patch
+
+    @staticmethod
+    def augment_patch(lr_patch, hr_patch):
+         """Applies simple random augmentations (flip, rotation)."""
+         # Random horizontal flip
+         if random.random() < 0.5:
+             lr_patch = lr_patch.transpose(Image.FLIP_LEFT_RIGHT)
+             hr_patch = hr_patch.transpose(Image.FLIP_LEFT_RIGHT)
+
+         # Random vertical flip (less common, can sometimes be excluded)
+         # if random.random() < 0.5:
+         #     lr_patch = lr_patch.transpose(Image.FLIP_TOP_BOTTOM)
+         #     hr_patch = hr_patch.transpose(Image.FLIP_TOP_BOTTOM)
+
+         # Random 90-degree rotation
+         # rot_choice = random.choice([0, 1, 2, 3]) # 0: 0 deg, 1: 90 deg, 2: 180 deg, 3: 270 deg
+         # if rot_choice != 0:
+         #      lr_patch = lr_patch.rotate(90 * rot_choice, expand=True) # expand=True might change size if not square
+         #      hr_patch = hr_patch.rotate(90 * rot_choice, expand=True)
+
+         return lr_patch, hr_patch
+
+
+    def __getitem__(self, idx):
+        # Get the full LR image path
+        lr_path = self.lr_image_files[idx]
+        try:
+            lr_img = Image.open(lr_path).convert('RGB')
+        except Exception as e:
+            print(f"Error opening LR image {lr_path}: {e}")
+            # Decide how to handle: return None, raise error, or return dummy
+            # Returning None requires careful handling in the DataLoader collate_fn or training loop
+            return None # Let collate_fn handle this potentially
+
+        # Construct the corresponding full HR image path
+        base_name = os.path.basename(lr_path)
+        hr_path = os.path.join(self.hr_dir, base_name)
+
+        # Handle potential alternative HR filenames
+        if not os.path.exists(hr_path):
+            base, ext = os.path.splitext(base_name)
+            if f'x{self.scale_factor}' in base:
+                hr_name = base.replace(f'x{self.scale_factor}', '') + ext
+                hr_path_alt = os.path.join(self.hr_dir, hr_name)
+                if os.path.exists(hr_path_alt):
+                    hr_path = hr_path_alt
+                else:
+                    print(f"ERROR in __getitem__: Cannot find corresponding HR for LR: {lr_path}")
+                    return None # Indicate error
+            else:
+                 print(f"ERROR in __getitem__: Cannot find corresponding HR for LR: {lr_path}")
+                 return None # Indicate error
+
+        try:
+            hr_img = Image.open(hr_path).convert('RGB')
+        except Exception as e:
+            print(f"Error opening HR image {hr_path}: {e}")
+            return None # Indicate error
+
+
+        # --- Get Corresponding Patches ---
+        try:
+            lr_patch, hr_patch = self.get_patch(lr_img, hr_img, self.patch_size_lr, self.scale_factor)
+        except ValueError as e: # Catch randrange error if patch size > image size after potential resize
+             print(f"Error getting patch for {lr_path} (maybe image is smaller than patch size?): {e}")
+             return None
+
+
+        # --- Apply Augmentations (Optional) ---
+        lr_patch, hr_patch = self.augment_patch(lr_patch, hr_patch)
+
+
+        # --- Apply Custom Transform if provided ---
+        # (Currently we pass None, but this is where you'd integrate albumentations etc.)
+        if self.transform:
+            # A typical transform might operate on numpy arrays
+            # lr_np = np.array(lr_patch)
+            # hr_np = np.array(hr_patch)
+            # transformed = self.transform(image=lr_np, mask=hr_np) # Example syntax
+            # lr_patch = Image.fromarray(transformed['image'])
+            # hr_patch = Image.fromarray(transformed['mask'])
+            pass # Placeholder
+
+
+        # --- Convert Patches to Tensors ---
+        to_tensor = transforms.ToTensor() # Converts PIL image (HWC) [0, 255] to Tensor (CHW) [0.0, 1.0]
+        lr_tensor = to_tensor(lr_patch)
+        hr_tensor = to_tensor(hr_patch)
+
+
+        return {'lr': lr_tensor, 'hr': hr_tensor}
+
+# --- Example Usage (for testing the definition) ---
+if __name__ == '__main__':
+    print("--- Testing SRDataset with Patching ---")
+    hr_data_dir = './datasets/DIV2K/HR_extracted/DIV2K_train_HR' # Modify if needed
+    lr_data_dir = './datasets/DIV2K/DIV2K_train_LR_bicubic/X4' # Modify if needed
+    scale = 4
+    lr_patch_size = 48 # Common LR patch size for SR tasks
+
+    if not os.path.isdir(hr_data_dir): print(f"ERROR: HR dir not found: '{hr_data_dir}'")
+    if not os.path.isdir(lr_data_dir): print(f"ERROR: LR dir not found: '{lr_data_dir}'")
+
+    try:
+        dataset = SRDataset(hr_dir=hr_data_dir, lr_dir=lr_data_dir,
+                            scale_factor=scale, patch_size_lr=lr_patch_size)
+
+        if len(dataset) > 0:
+            print(f"\nSuccessfully loaded dataset with {len(dataset)} image pairs.")
+
+            # Test getting a single item (patch pair)
+            print("\n--- Testing __getitem__ ---")
+            num_test_items = 5
+            for i in range(min(num_test_items, len(dataset))):
+                 item = dataset[i]
+                 if item is None:
+                      print(f"Item {i}: Returned None (Error occurred)")
+                      continue
+
+                 lr_p = item['lr']
+                 hr_p = item['hr']
+                 print(f"Item {i}: LR Patch Shape={lr_p.shape}, HR Patch Shape={hr_p.shape}")
+
+                 # Verify shapes
+                 expected_hr_shape = (3, lr_patch_size * scale, lr_patch_size * scale)
+                 if lr_p.shape != (3, lr_patch_size, lr_patch_size) or hr_p.shape != expected_hr_shape:
+                      print(f"  WARNING: Shape mismatch! LR={lr_p.shape}, HR={hr_p.shape}, Expected HR={expected_hr_shape}")
+
+            # Test DataLoader with a simple collate function that filters Nones
+            print("\n--- Testing DataLoader with Patches ---")
+            from torch.utils.data import DataLoader
+
+            # Define a collate_fn that filters out None values returned by __getitem__
+            def collate_fn_filter_none(batch):
+                batch = list(filter(lambda x: x is not None, batch))
+                if not batch: # If all items in the batch failed
+                    return None
+                # Use default collate on the filtered batch
+                return torch.utils.data.dataloader.default_collate(batch)
+
+            # Use batch_size=4 for testing
+            dataloader = DataLoader(dataset, batch_size=4, shuffle=True,
+                                    num_workers=0, collate_fn=collate_fn_filter_none)
+
+            num_test_batches = 3
+            batch_count = 0
+            for batch in dataloader:
+                 if batch_count >= num_test_batches:
+                      break
+                 if batch is None:
+                      print(f"Skipping an entirely problematic batch.")
+                      continue
+
+                 lr_batch = batch['lr']
+                 hr_batch = batch['hr']
+                 print(f"Batch {batch_count}: LR Batch Shape={lr_batch.shape}, HR Batch Shape={hr_batch.shape}")
+                 batch_count += 1
+
+            if batch_count > 0:
+                 print("DataLoader test with patches successful.")
+            else:
+                 print("DataLoader test: Could not retrieve any valid batches.")
+
+        else:
+            print("\nDataset loaded but is empty.")
+
+    except FileNotFoundError as e:
+         print(f"\nERROR initializing dataset: {e}")
+    except Exception as e:
+         print(f"\nAn unexpected error occurred during dataset testing: {e}")
+
+    print("\n--- SRDataset Test Finished ---")

loss.py

@@ -0,0 +1,138 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import vgg19, VGG19_Weights
+from torchvision import transforms
+
+class PerceptualLoss(nn.Module):
+    """
+    Calculates the VGG perceptual loss.
+
+    Uses features from the VGG19 network pretrained on ImageNet.
+    Compares features from specific layers for the generated and target images.
+    """
+    def __init__(self, feature_layers=None, use_l1=True, device='cpu'):
+        """
+        Args:
+            feature_layers (list of int, optional): Indices of VGG19 feature layers to use.
+                                                    Defaults correspond to layers before pool1, pool2, pool3, pool4.
+                                                    Specifically: relu1_1, relu2_1, relu3_1, relu4_1 in many implementations.
+                                                    VGG19 structure: layer indices relate to `features` module.
+            use_l1 (bool): If True, use L1 loss between features. If False, use L2 (MSE) loss.
+            device (str): 'cuda' or 'cpu'.
+        """
+        super(PerceptualLoss, self).__init__()
+        
+        # Load pre-trained VGG19 model
+        # Ensure you have torchvision installed: pip install torchvision
+        try:
+            # Recommended way with modern torchvision
+            weights = VGG19_Weights.IMAGENET1K_V1
+            self.vgg = vgg19(weights=weights).features
+            self.preprocess = weights.transforms() # Get the preprocessing expected by the model
+        except AttributeError:
+             # Fallback for older torchvision versions (might require manual weight download if not cached)
+            print("Warning: Using older torchvision VGG19 loading method. Consider upgrading torchvision.")
+            self.vgg = vgg19(pretrained=True).features
+            # Define standard ImageNet normalization manually if transform isn't available
+            self.preprocess = transforms.Compose([
+                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            ])
+
+        self.vgg.eval() # Set VGG to evaluation mode
+        for param in self.vgg.parameters():
+            param.requires_grad = False # Freeze VGG parameters
+
+        self.vgg = self.vgg.to(device)
+        self.device = device
+
+        # Define the layers to extract features from
+        # Common choices are layers before max pooling
+        # VGG19 features structure indices:
+        # ReLU1_1: 1, ReLU2_1: 6, ReLU3_1: 11, ReLU4_1: 20, ReLU5_1: 29 (Sometimes ReLU5 used too)
+        if feature_layers is None:
+            # These indices correspond to the output of Conv layers before MaxPool
+            # Specifically: conv1_1(0), conv2_1(5), conv3_1(10), conv4_1(19), conv5_1(28)
+            # Often the ReLU output right after is used: 1, 6, 11, 20, 29
+            self.feature_layers = {1, 6, 11, 20} # Using ReLU outputs before pooling layers 1-4
+            # Alternative common set often cited as relu5_4 (index 35 or 36 depending on source):
+            # self.feature_layers = {35} # Or use a specific high-level layer
+        else:
+            self.feature_layers = set(feature_layers)
+
+        self.loss_fn = nn.L1Loss() if use_l1 else nn.MSELoss()
+
+        print(f"PerceptualLoss: Using VGG19 features from layers: {sorted(list(self.feature_layers))}")
+        print(f"PerceptualLoss: Using {'L1' if use_l1 else 'L2'} distance.")
+
+
+    def forward(self, generated, target):
+        """
+        Compute the perceptual loss.
+
+        Args:
+            generated (torch.Tensor): The generated image tensor (B, C, H, W). Values [0, 1].
+            target (torch.Tensor): The target (ground truth) image tensor (B, C, H, W). Values [0, 1].
+
+        Returns:
+            torch.Tensor: The calculated perceptual loss.
+        """
+        # Ensure inputs are on the correct device
+        generated = generated.to(self.device)
+        target = target.to(self.device)
+
+        # Preprocess images for VGG
+        # VGG expects inputs normalized based on ImageNet stats
+        # The transform might handle dtype and range, but let's be explicit
+        generated_norm = self.preprocess(generated)
+        target_norm = self.preprocess(target)
+
+        # Extract features
+        loss = 0.0
+        current_layer_idx = 0
+        max_needed_layer = max(self.feature_layers) if self.feature_layers else 0
+
+        # Iterate through VGG layers, extracting features only from specified layers
+        for layer in self.vgg:
+            # Compute features for both images up to the current layer
+            generated_norm = layer(generated_norm)
+            target_norm = layer(target_norm)
+
+            # If the current layer index is one we want to use for loss calculation
+            if current_layer_idx in self.feature_layers:
+                loss += self.loss_fn(generated_norm, target_norm)
+
+            # Stop iterating if we've passed the last needed layer
+            if current_layer_idx >= max_needed_layer:
+                break
+
+            current_layer_idx += 1
+
+        return loss
+
+
+# --- Example Usage (for testing the definition) ---
+if __name__ == '__main__':
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    # Create dummy images (Batch Size, Channels, Height, Width)
+    # Note: Images should be in the range [0, 1] for standard transforms
+    dummy_generated = torch.rand(2, 3, 96, 96).to(device) # Example size (must match target)
+    dummy_target = torch.rand(2, 3, 96, 96).to(device)
+
+    # Instantiate the loss function
+    # Default layers: {1, 6, 11, 20} (Relu1_1, Relu2_1, Relu3_1, Relu4_1 outputs)
+    perceptual_loss_l1 = PerceptualLoss(device=device, use_l1=True)
+    # Example with different layers and L2 loss
+    # perceptual_loss_l2 = PerceptualLoss(feature_layers={35}, device=device, use_l1=False)
+
+    # Calculate loss
+    loss_val_l1 = perceptual_loss_l1(dummy_generated, dummy_target)
+    # loss_val_l2 = perceptual_loss_l2(dummy_generated, dummy_target)
+
+    print(f"\nCalculated Perceptual Loss (L1, default layers): {loss_val_l1.item()}")
+    # print(f"Calculated Perceptual Loss (L2, layer 35): {loss_val_l2.item()}")
+
+    assert loss_val_l1.item() >= 0, "Loss should be non-negative"
+    print("\nPerceptualLoss definition test successful!")

models.py

@@ -0,0 +1,181 @@
+import torch
+import torch.nn as nn
+import math
+import os
+
+# --- ResidualBlock, Upsampler, and Generator classes remain the same ---
+class ResidualBlock(nn.Module):
+    def __init__(self, num_features, kernel_size=3, bn=False, act=nn.ReLU(True), res_scale=1.0):
+        super(ResidualBlock, self).__init__()
+        padding = kernel_size // 2
+        m = []
+        m.append(nn.Conv2d(num_features, num_features, kernel_size, padding=padding))
+        if bn: m.append(nn.BatchNorm2d(num_features))
+        m.append(act)
+        m.append(nn.Conv2d(num_features, num_features, kernel_size, padding=padding))
+        if bn: m.append(nn.BatchNorm2d(num_features))
+        self.body = nn.Sequential(*m)
+        self.res_scale = res_scale
+    def forward(self, x):
+        res = self.body(x).mul(self.res_scale)
+        res += x
+        return res
+
+class Upsampler(nn.Module):
+    def __init__(self, scale_factor, num_features, act=nn.ReLU(True)):
+        super(Upsampler, self).__init__()
+        m = []
+        m.append(nn.Conv2d(num_features, num_features * (scale_factor ** 2), kernel_size=3, padding=1))
+        m.append(nn.PixelShuffle(scale_factor))
+        if act: m.append(act)
+        self.body = nn.Sequential(*m)
+    def forward(self, x):
+        return self.body(x)
+
+class Generator(nn.Module):
+    def __init__(self, scale_factor=4, in_channels=3, out_channels=3, num_features=64, num_res_blocks=16, res_scale=1.0):
+        super(Generator, self).__init__()
+        self.scale_factor = scale_factor
+        act = nn.ReLU(True)
+        self.head = nn.Conv2d(in_channels, num_features, kernel_size=3, padding=1)
+        res_blocks = [ResidualBlock(num_features, kernel_size=3, act=act, res_scale=res_scale) for _ in range(num_res_blocks)]
+        res_blocks.append(nn.Conv2d(num_features, num_features, kernel_size=3, padding=1))
+        self.body = nn.Sequential(*res_blocks)
+        m_tail = []
+        if (scale_factor & (scale_factor - 1)) == 0:
+            for _ in range(int(math.log2(scale_factor))):
+                m_tail.append(Upsampler(scale_factor=2, num_features=num_features, act=None))
+        elif scale_factor == 3:
+             m_tail.append(Upsampler(scale_factor=3, num_features=num_features, act=None))
+        else:
+            raise NotImplementedError(f"Scale factor {scale_factor} not directly supported by this simple upsampler.")
+        self.tail = nn.Sequential(*m_tail)
+        self.final_conv = nn.Conv2d(num_features, out_channels, kernel_size=3, padding=1)
+
+    def forward(self, lr_img):
+        x = self.head(lr_img)
+        res = self.body(x)
+        res += x
+        x = self.tail(res)
+        x = self.final_conv(x)
+        return x
+
+# +++ NEW Discriminator Class +++
+class Discriminator(nn.Module):
+    """
+    Simple CNN Discriminator Network (PatchGAN style is common but this is simpler).
+    Takes an image (real HR or generated SR) and outputs a single logit.
+    """
+    def __init__(self, in_channels=3, num_features_start=64, num_blocks=4):
+        super(Discriminator, self).__init__()
+
+        # Initial block
+        layers = [
+            nn.Conv2d(in_channels, num_features_start, kernel_size=3, stride=1, padding=1),
+            nn.LeakyReLU(0.2, inplace=True)
+        ]
+
+        current_features = num_features_start
+        for i in range(num_blocks):
+            stride = 1 if i % 2 == 0 else 2 # Downsample every other block
+            next_features = current_features * 2 if stride == 2 else current_features
+            layers.extend([
+                nn.Conv2d(current_features, next_features, kernel_size=3, stride=stride, padding=1),
+                nn.BatchNorm2d(next_features), # BatchNorm is common in discriminators
+                nn.LeakyReLU(0.2, inplace=True)
+            ])
+            current_features = next_features
+
+        self.features = nn.Sequential(*layers)
+
+        # Classifier part - adjust input features based on final conv output size
+        # We need to know the output size of the feature extractor to define the Linear layer.
+        # Using AdaptiveAvgPool2d makes it independent of the input image size.
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.classifier = nn.Sequential(
+            nn.Linear(current_features, 100), # Example intermediate size
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(100, 1) # Output a single logit (no sigmoid here)
+        )
+
+    def forward(self, img):
+        """
+        Args:
+            img (torch.Tensor): Input image tensor (B, C, H, W), either real HR or fake SR.
+        Returns:
+            torch.Tensor: Output logits (B, 1). Higher values -> more likely "real".
+        """
+        batch_size = img.size(0)
+        features = self.features(img)
+        pooled = self.avgpool(features)
+        # Flatten the output of avgpool for the linear layer
+        pooled = pooled.view(batch_size, -1)
+        output = self.classifier(pooled)
+        return output
+
+# --- Main block for testing and saving ---
+if __name__ == '__main__':
+    # --- Generator Test (as before) ---
+    SCALE = 4
+    GEN_FEATURES = 64
+    GEN_RES_BLOCKS = 8
+    save_dir = "saved_models"
+    os.makedirs(save_dir, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    # Dummy LR input for Generator
+    gen_batch_size = 1
+    lr_height = 32
+    lr_width = 32
+    in_channels = 3
+    dummy_lr = torch.randn(gen_batch_size, in_channels, lr_height, lr_width).to(device)
+    print(f"Dummy LR input shape (Generator): {dummy_lr.shape}")
+
+    generator = Generator(scale_factor=SCALE, num_features=GEN_FEATURES, num_res_blocks=GEN_RES_BLOCKS).to(device)
+    generator.eval()
+    with torch.no_grad():
+        output_sr = generator(dummy_lr)
+    print(f"Output SR shape (Generator): {output_sr.shape}")
+    # ... (rest of generator verification and saving code remains here) ...
+    print("\nGenerator definition test successful!")
+    num_params_gen = sum(p.numel() for p in generator.parameters() if p.requires_grad)
+    print(f"Generator - Number of trainable parameters: {num_params_gen:,}")
+    # ... (Saving code as before) ...
+
+    print("\n--- Testing Discriminator ---")
+    # --- Discriminator Test ---
+    DISC_FEATURES = 64 # Starting features for discriminator
+    DISC_BLOCKS = 3   # Number of conv blocks in discriminator
+
+    # Dummy HR/SR input for Discriminator (must match Generator's output size)
+    disc_batch_size = 4 # Can be different from generator test batch size
+    hr_height = output_sr.shape[2] # Use the calculated HR height
+    hr_width = output_sr.shape[3]  # Use the calculated HR width
+    dummy_hr = torch.randn(disc_batch_size, in_channels, hr_height, hr_width).to(device)
+    print(f"Dummy HR/SR input shape (Discriminator): {dummy_hr.shape}")
+
+    # Instantiate the Discriminator
+    discriminator = Discriminator(in_channels=in_channels,
+                                num_features_start=DISC_FEATURES,
+                                num_blocks=DISC_BLOCKS).to(device)
+    discriminator.eval() # Set to evaluation mode for testing
+
+    # print(discriminator) # Optional: Print structure
+
+    # Perform a forward pass
+    with torch.no_grad():
+        output_logits = discriminator(dummy_hr)
+
+    print(f"Output Logits shape (Discriminator): {output_logits.shape}")
+
+    # Verify output shape
+    expected_disc_shape = (disc_batch_size, 1)
+    assert output_logits.shape == expected_disc_shape, \
+        f"Discriminator output shape mismatch! Expected {expected_disc_shape}, got {output_logits.shape}"
+
+    print("Discriminator definition test successful!")
+
+    # Optional: Count parameters
+    num_params_disc = sum(p.numel() for p in discriminator.parameters() if p.requires_grad)
+    print(f"Discriminator - Number of trainable parameters: {num_params_disc:,}")

prep.py

@@ -0,0 +1,252 @@
+import os
+import glob
+import zipfile
+import requests
+import argparse
+from PIL import Image
+from tqdm import tqdm
+
+# --- Helper Functions ---
+
+def download_file(url, dest_path, chunk_size=8192):
+    """Downloads a file from a URL to a destination path with progress bar."""
+    try:
+        response = requests.get(url, stream=True, timeout=30) # Added timeout
+        response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)
+
+        total_size = int(response.headers.get('content-length', 0))
+        
+        print(f"Downloading {os.path.basename(dest_path)} ({total_size / (1024*1024):.2f} MB)...")
+        with open(dest_path, 'wb') as f, tqdm(
+            desc=os.path.basename(dest_path),
+            total=total_size,
+            unit='iB',
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as bar:
+            for chunk in response.iter_content(chunk_size=chunk_size):
+                size = f.write(chunk)
+                bar.update(size)
+        print(f"Download complete: {dest_path}")
+        return True
+
+    except requests.exceptions.RequestException as e:
+        print(f"Error downloading {url}: {e}")
+        # Clean up partially downloaded file if it exists
+        if os.path.exists(dest_path):
+            os.remove(dest_path)
+        return False
+    except Exception as e:
+        print(f"An unexpected error occurred during download: {e}")
+        if os.path.exists(dest_path):
+            os.remove(dest_path)
+        return False
+
+
+def unzip_file(zip_path, extract_to):
+    """Unzips a file to a specified directory."""
+    print(f"Extracting {os.path.basename(zip_path)} to {extract_to}...")
+    try:
+        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+            # You could add a progress bar here for large zips if needed
+            # using zip_ref.infolist() and iterating extraction, but
+            # extractall is usually efficient enough.
+            zip_ref.extractall(extract_to)
+        print("Extraction complete.")
+        return True
+    except zipfile.BadZipFile:
+        print(f"Error: Invalid or corrupted zip file: {zip_path}")
+        return False
+    except Exception as e:
+        print(f"An error occurred during extraction: {e}")
+        return False
+
+def find_image_dir(base_path, expected_subdir_suffix='_HR'):
+    """
+    Tries to find the actual directory containing images after extraction.
+    Handles cases where unzip creates an extra top-level folder.
+    """
+    # Check if images are directly in base_path
+    if glob.glob(os.path.join(base_path, '*.png')) or \
+       glob.glob(os.path.join(base_path, '*.jpg')) or \
+       glob.glob(os.path.join(base_path, '*.jpeg')):
+        return base_path
+
+    # Check common pattern: base_path/DatasetName_HR/
+    potential_dirs = [d for d in glob.glob(os.path.join(base_path, '*')) if os.path.isdir(d)]
+    if len(potential_dirs) == 1:
+         subdir = potential_dirs[0]
+         # Check if this subdir contains images or ends with the expected suffix
+         if subdir.endswith(expected_subdir_suffix) or \
+            glob.glob(os.path.join(subdir, '*.png')) or \
+            glob.glob(os.path.join(subdir, '*.jpg')) or \
+            glob.glob(os.path.join(subdir, '*.jpeg')):
+             print(f"Found image directory: {subdir}")
+             return subdir
+
+    # Fallback if specific pattern not found, maybe it's still just base_path
+    print(f"Warning: Could not definitively locate image subdirectory in {base_path}. Assuming images are directly within or in a single nested folder.")
+    # If we found exactly one directory, return that, otherwise return the original path
+    return potential_dirs[0] if len(potential_dirs) == 1 else base_path
+
+
+def downsample_images(hr_dir, lr_dir, scale_factor):
+    """Downsamples HR images using bicubic interpolation."""
+    if not os.path.exists(lr_dir):
+        os.makedirs(lr_dir)
+        print(f"Created LR directory: {lr_dir}")
+
+    hr_images = glob.glob(os.path.join(hr_dir, '*.png')) + \
+                glob.glob(os.path.join(hr_dir, '*.jpg')) + \
+                glob.glob(os.path.join(hr_dir, '*.jpeg'))
+
+    if not hr_images:
+        print(f"Error: No images found in the determined HR directory: {hr_dir}")
+        return False
+
+    print(f"Found {len(hr_images)} HR images in {hr_dir}. Starting downsampling (x{scale_factor})...")
+
+    processed_count = 0
+    for hr_path in tqdm(hr_images, desc=f"Downsampling x{scale_factor}"):
+        try:
+            hr_img = Image.open(hr_path).convert('RGB') # Ensure RGB
+            hr_width, hr_height = hr_img.size
+
+            lr_width = hr_width // scale_factor
+            lr_height = hr_height // scale_factor
+
+            if lr_width == 0 or lr_height == 0:
+                print(f"\nWarning: Image {os.path.basename(hr_path)} is too small ({hr_width}x{hr_height}) for scale factor {scale_factor}. Skipping.")
+                continue
+
+            lr_img = hr_img.resize((lr_width, lr_height), resample=Image.BICUBIC)
+
+            base_name = os.path.basename(hr_path)
+            lr_save_path = os.path.join(lr_dir, base_name)
+            lr_img.save(lr_save_path)
+            processed_count += 1
+
+        except Exception as e:
+            print(f"\nError processing {hr_path}: {e}")
+
+    print(f"Downsampling complete. Processed {processed_count}/{len(hr_images)} images.")
+    return processed_count > 0
+
+
+# --- Main Execution ---
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Download and prepare dataset for Super-Resolution.")
+    parser.add_argument('--url', type=str, default='https://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip', help='URL of the dataset zip file (default: DIV2K Train HR).')
+    parser.add_argument('--base_dir', type=str, default='./datasets', help='Base directory to store datasets.')
+    parser.add_argument('--dataset_name', type=str, default='DIV2K', help='Name for the dataset folder.')
+    parser.add_argument('--scale', type=int, default=4, help='Downsampling scale factor (e.g., 4 for x4).')
+    parser.add_argument('--force', action='store_true', help='Force redownload and reprocessing even if data exists.')
+
+    args = parser.parse_args()
+
+    # --- Define Paths ---
+    dataset_base_path = os.path.join(args.base_dir, args.dataset_name)
+    zip_filename = os.path.basename(args.url)
+    zip_save_path = os.path.join(dataset_base_path, zip_filename)
+    hr_extract_base = os.path.join(dataset_base_path, 'HR_extracted') # Temp extraction location
+    # We will determine the *actual* HR image dir after extraction
+    lr_save_dir = os.path.join(dataset_base_path, f'DIV2K_train_LR_bicubic/X{args.scale}') # Following previous convention
+
+    print(f"--- Configuration ---")
+    print(f"Dataset URL: {args.url}")
+    print(f"Base Directory: {args.base_dir}")
+    print(f"Dataset Name: {args.dataset_name}")
+    print(f"Target Scale: x{args.scale}")
+    print(f"Zip Save Path: {zip_save_path}")
+    print(f"Initial Extract Path: {hr_extract_base}")
+    print(f"LR Save Path: {lr_save_dir}")
+    print(f"Force Re-run: {args.force}")
+    print(f"--------------------")
+
+    # --- Create Base Directory ---
+    os.makedirs(dataset_base_path, exist_ok=True)
+
+    # --- Step 1: Download ---
+    hr_dir_exists = os.path.isdir(hr_extract_base) # Check if base extraction dir exists
+    download_needed = not os.path.exists(zip_save_path) or args.force
+
+    if download_needed:
+        if args.force and os.path.exists(zip_save_path):
+            print("Force enabled: Removing existing zip file...")
+            os.remove(zip_save_path)
+        if not download_file(args.url, zip_save_path):
+            print("Exiting due to download failure.")
+            exit(1)
+    elif hr_dir_exists: # If zip exists and hr dir exists, assume download & unzip ok unless forced
+         print("Zip file already exists. Skipping download (use --force to override).")
+    else: # Zip exists but HR dir doesn't - need to unzip
+        print("Zip file found, but extraction directory missing. Will proceed to unzip.")
+
+
+    # --- Step 2: Unzip ---
+    # Check if the *potential* content directory already exists. Be a bit lenient here.
+    # A more robust check would be to look inside the zip first or check for specific files.
+    unzip_needed = not hr_dir_exists or args.force
+
+    actual_hr_dir = None # Will store the path to the actual images
+
+    if unzip_needed:
+        if args.force and hr_dir_exists:
+             print("Force enabled: Removing existing extraction directory...")
+             import shutil
+             shutil.rmtree(hr_extract_base) # Careful! Removes directory and contents
+
+        if not os.path.exists(zip_save_path):
+            print("Error: Zip file not found, cannot unzip. Please check download step or path.")
+            exit(1)
+
+        os.makedirs(hr_extract_base, exist_ok=True) # Ensure extraction target exists
+        if not unzip_file(zip_save_path, hr_extract_base):
+             print("Exiting due to extraction failure.")
+             exit(1)
+        # Find the actual directory containing images post-extraction
+        actual_hr_dir = find_image_dir(hr_extract_base, expected_subdir_suffix=f'{args.dataset_name}_HR') # e.g., DIV2K_HR
+        if not actual_hr_dir or not (glob.glob(os.path.join(actual_hr_dir, '*.png')) or glob.glob(os.path.join(actual_hr_dir, '*.jpg'))):
+            print(f"Error: Could not locate the directory with HR images within {hr_extract_base} after extraction.")
+            exit(1)
+        print(f"Located HR images in: {actual_hr_dir}")
+
+    else:
+        print("HR extraction directory already exists. Skipping unzip (use --force to override).")
+        # Try to find the HR dir even if we skipped unzipping
+        actual_hr_dir = find_image_dir(hr_extract_base, expected_subdir_suffix=f'{args.dataset_name}_HR')
+        if not actual_hr_dir:
+             print(f"Error: Could not locate the directory with HR images within existing {hr_extract_base}.")
+             exit(1)
+        print(f"Using existing HR images from: {actual_hr_dir}")
+
+
+    # --- Step 3: Process (Downsample) ---
+    lr_dir_exists_and_populated = os.path.isdir(lr_save_dir) and len(os.listdir(lr_save_dir)) > 0
+    processing_needed = not lr_dir_exists_and_populated or args.force
+
+    if processing_needed:
+        if args.force and lr_dir_exists_and_populated:
+             print("Force enabled: Removing existing LR directory...")
+             import shutil
+             shutil.rmtree(lr_save_dir) # Careful!
+
+        if not actual_hr_dir:
+             print("Error: Cannot proceed with downsampling, HR image directory not determined.")
+             exit(1)
+
+        if not downsample_images(actual_hr_dir, lr_save_dir, args.scale):
+             print("Downsampling process failed or produced no images.")
+             # Optionally exit here depending on desired behavior
+             # exit(1)
+        else:
+             print("Downsampling finished successfully.")
+    else:
+        print("LR directory already exists and is populated. Skipping downsampling (use --force to override).")
+
+
+    print("\n--- Script Finished ---")
+    print(f"HR images should be available in/under: {actual_hr_dir}")
+    print(f"LR images (x{args.scale}) should be available in: {lr_save_dir}")
+    print("You can now use these directories with the SRDataset class.")

saved_models/generator_x4_f64_b8_untrained.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ea5ec7bb7ec436c504f98cf3380a7b2258bf3730cf4ae726b838dd7df52d0b1
+size 3717459

saved_models/generator_x4_f64_b8_untrained.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de3bc0ab6790bede102d5a40fd5122bbff83e05b22331f9cc983eb76aace56db
+size 3722508

train.py

@@ -0,0 +1,307 @@
+# RUN python train.py --epochs 2 --batch_size 2 --subset 10 --num_workers 0 --cpu --patch_size 48
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import os
+import argparse
+from tqdm import tqdm
+import time
+
+# Import custom modules
+from dataset import SRDataset # Make sure dataset.py is in the same directory
+from models import Generator, Discriminator # Make sure models.py is in the same directory
+from loss import PerceptualLoss # Make sure loss.py is in the same directory
+
+def train(args):
+    # --- 1. Setup ---
+    device = torch.device("cuda" if torch.cuda.is_available() and not args.cpu else "cpu")
+    print(f"Using device: {device}")
+
+    # Create directories for saving models and potentially logs/outputs
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    # --- 2. Data ---
+    print("Loading dataset...")
+    # Note: args.hr_dir and args.lr_dir are assumed to be valid paths by this point
+    # due to checks in the __main__ block
+    try:
+        train_dataset = SRDataset(hr_dir=args.hr_dir, lr_dir=args.lr_dir, scale_factor=args.scale, patch_size_lr=args.patch_size)
+    except FileNotFoundError as e:
+        print(f"Error creating dataset: {e}")
+        print("Please ensure the specified HR and LR directories contain correctly named image files.")
+        exit(1)
+    except Exception as e:
+        print(f"An unexpected error occurred while creating the dataset: {e}")
+        exit(1)
+
+
+    # Use a smaller subset for initial testing on CPU if needed
+    if args.subset > 0 and args.subset < len(train_dataset):
+         print(f"Using a subset of {args.subset} images for training.")
+         indices = torch.randperm(len(train_dataset))[:args.subset]
+         train_dataset = torch.utils.data.Subset(train_dataset, indices)
+    elif args.subset >= len(train_dataset) and len(train_dataset) > 0 :
+        print(f"Subset size ({args.subset}) is >= dataset size ({len(train_dataset)}). Using full dataset.")
+
+
+    if len(train_dataset) == 0:
+        print(f"Error: Dataset is empty after attempting to load. Please check HR dir '{args.hr_dir}' and LR dir '{args.lr_dir}'")
+        return
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers, # Set to 0 if you encounter issues on Windows/macOS
+        pin_memory=True if device == 'cuda' else False # pin_memory only useful for GPU
+    )
+    print(f"Dataset loaded: {len(train_dataset)} training images.")
+    print(f"Dataloader: {len(train_loader)} batches per epoch.")
+
+
+    # --- 3. Models ---
+    print("Initializing models...")
+    generator = Generator(scale_factor=args.scale,
+                          num_features=args.gen_features,
+                          num_res_blocks=args.gen_blocks).to(device)
+
+    discriminator = Discriminator(in_channels=3, # Assuming RGB input for discriminator
+                                  num_features_start=args.disc_features,
+                                  num_blocks=args.disc_blocks).to(device)
+
+    print(f"Generator params: {sum(p.numel() for p in generator.parameters()):,}")
+    print(f"Discriminator params: {sum(p.numel() for p in discriminator.parameters()):,}")
+
+    # --- 4. Loss Functions ---
+    print("Initializing loss functions...")
+    # Content Loss (Pixel-wise) - L1 is common for SR
+    content_loss_criterion = nn.L1Loss().to(device)
+
+    # Adversarial Loss - Measures how well G fools D and D identifies fakes
+    adversarial_loss_criterion = nn.BCEWithLogitsLoss().to(device) # More stable than BCELoss + Sigmoid
+
+    # Perceptual Loss (VGG-based)
+    try:
+        perceptual_loss_criterion = PerceptualLoss(device=device, use_l1=True) # Using L1 feature distance
+    except Exception as e:
+        print(f"Error initializing Perceptual Loss (check VGG weights download/torchvision install): {e}")
+        exit(1)
+
+
+    # --- 5. Optimizers ---
+    print("Initializing optimizers...")
+    optimizer_g = optim.Adam(generator.parameters(), lr=args.lr_gen, betas=(0.9, 0.999))
+    optimizer_d = optim.Adam(discriminator.parameters(), lr=args.lr_disc, betas=(0.9, 0.999))
+
+    # --- Optional: Learning Rate Scheduler ---
+    # Example: scheduler_g = optim.lr_scheduler.StepLR(optimizer_g, step_size=args.lr_decay_step, gamma=0.5)
+    # Example: scheduler_d = optim.lr_scheduler.StepLR(optimizer_d, step_size=args.lr_decay_step, gamma=0.5)
+
+    # --- 6. Training Loop ---
+    print("\n--- Starting Training ---")
+    start_time = time.time()
+
+    for epoch in range(1, args.epochs + 1):
+        generator.train()   # Set generator to training mode
+        discriminator.train() # Set discriminator to training mode
+        epoch_loss_g = 0.0
+        epoch_loss_d = 0.0
+        epoch_start_time = time.time()
+
+        progress_bar = tqdm(train_loader, desc=f"Epoch {epoch}/{args.epochs}", leave=True) # leave=True to keep bar after epoch
+
+        for batch_idx, batch in enumerate(progress_bar):
+            # Ensure batch is valid (dataset loader might return None on error in __getitem__)
+            if batch is None:
+                print(f"Warning: Skipping problematic batch at index {batch_idx}")
+                continue
+
+            try:
+                lr_images = batch['lr'].to(device) # Low-resolution images
+                hr_images = batch['hr'].to(device) # High-resolution (ground truth) images
+            except KeyError as e:
+                print(f"Error accessing batch data: {e}. Check SRDataset's __getitem__ return format.")
+                continue # Skip this batch
+
+            # Create labels for adversarial loss
+            # Real labels = 1, Fake labels = 0
+            # Add some noise or use soft labels (e.g., 0.9 instead of 1.0) can sometimes help stabilize GAN training
+            real_labels = torch.ones((hr_images.size(0), 1)).to(device)
+            fake_labels = torch.zeros((hr_images.size(0), 1)).to(device)
+
+            # ---------------------
+            #  Train Discriminator
+            # ---------------------
+            optimizer_d.zero_grad()
+
+            # Generate fake HR images
+            # Use torch.no_grad() for generator forward pass when only training discriminator
+            with torch.no_grad():
+                 fake_sr_images = generator(lr_images) # No need to detach() if already in no_grad context
+
+            # Loss for real images
+            real_logits = discriminator(hr_images)
+            loss_d_real = adversarial_loss_criterion(real_logits, real_labels)
+
+            # Loss for fake images
+            fake_logits = discriminator(fake_sr_images) # Use the generated fakes
+            loss_d_fake = adversarial_loss_criterion(fake_logits, fake_labels)
+
+            # Total discriminator loss
+            loss_d = (loss_d_real + loss_d_fake) / 2
+
+            # Backpropagate and update Discriminator
+            loss_d.backward()
+            # Optional: Gradient clipping for Discriminator (can help stability)
+            # torch.nn.utils.clip_grad_norm_(discriminator.parameters(), max_norm=1.0)
+            optimizer_d.step()
+
+
+            # -----------------
+            #  Train Generator
+            # (Typically done less frequently than discriminator, e.g., every k steps,
+            # but for simplicity here we do it every step)
+            # -----------------
+            optimizer_g.zero_grad()
+
+            # Generate fake HR images (this time track gradients for G)
+            generated_sr_images = generator(lr_images)
+
+            # --- Calculate Generator Losses ---
+            # 1. Content Loss (e.g., L1 distance between generated and real HR)
+            loss_content = content_loss_criterion(generated_sr_images, hr_images)
+
+            # 2. Perceptual Loss (VGG feature distance)
+            loss_perceptual = perceptual_loss_criterion(generated_sr_images, hr_images)
+
+            # 3. Adversarial Loss (how well G fools D)
+            # We want the discriminator to output 'real' (1) for the generated images
+            # Pass generated images through the discriminator (ensure D is not in no_grad context here)
+            generated_logits = discriminator(generated_sr_images)
+            loss_adversarial = adversarial_loss_criterion(generated_logits, real_labels) # Use real_labels!
+
+            # --- Combine Generator Losses ---
+            # Weights control the balance between pixel accuracy, perceptual quality, and realism
+            loss_g = (args.lambda_content * loss_content +
+                      args.lambda_percep * loss_perceptual +
+                      args.lambda_adv * loss_adversarial)
+
+            # Backpropagate and update Generator
+            loss_g.backward()
+             # Optional: Gradient clipping for Generator
+            # torch.nn.utils.clip_grad_norm_(generator.parameters(), max_norm=1.0)
+            optimizer_g.step()
+
+            # --- Update running losses and progress bar ---
+            epoch_loss_g += loss_g.item()
+            epoch_loss_d += loss_d.item()
+            progress_bar.set_postfix({
+                'Loss G': f"{loss_g.item():.4f}",
+                'Loss D': f"{loss_d.item():.4f}",
+                # Optional: Show individual components of G loss
+                # 'L_Cont': f"{loss_content.item():.4f}",
+                # 'L_Perc': f"{loss_perceptual.item():.4f}",
+                # 'L_Adv': f"{loss_adversarial.item():.4f}"
+            })
+
+        # --- End of Epoch ---
+        avg_loss_g = epoch_loss_g / len(train_loader) if len(train_loader) > 0 else 0
+        avg_loss_d = epoch_loss_d / len(train_loader) if len(train_loader) > 0 else 0
+        epoch_time = time.time() - epoch_start_time
+
+        # Optional: Update learning rate schedulers
+        # scheduler_g.step()
+        # scheduler_d.step()
+        # current_lr_g = optimizer_g.param_groups[0]['lr']
+
+        print(f"\nEpoch {epoch}/{args.epochs} | Time: {epoch_time:.2f}s | Avg Loss G: {avg_loss_g:.4f} | Avg Loss D: {avg_loss_d:.4f}")
+
+        # --- Save Checkpoint ---
+        if epoch % args.save_interval == 0 or epoch == args.epochs:
+            gen_path = os.path.join(args.save_dir, f"generator_epoch_{epoch}.pth")
+            disc_path = os.path.join(args.save_dir, f"discriminator_epoch_{epoch}.pth")
+            try:
+                torch.save(generator.state_dict(), gen_path)
+                torch.save(discriminator.state_dict(), disc_path)
+                print(f"Checkpoint saved for epoch {epoch} to '{args.save_dir}'")
+            except Exception as e:
+                print(f"Error saving checkpoint for epoch {epoch}: {e}")
+
+    # --- End of Training ---
+    total_time = time.time() - start_time
+    print(f"\n--- Training Finished ---")
+    print(f"Total time: {total_time // 3600:.0f}h {(total_time % 3600) // 60:.0f}m {total_time % 60:.2f}s")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Train SRGAN Model')
+
+    # --- Data Args ---
+    parser.add_argument('--hr_dir', type=str,
+                        default='./datasets/DIV2K/HR_extracted/DIV2K_train_HR',
+                        help='Path to high-resolution training images')
+    parser.add_argument('--lr_dir', type=str, default=None, # Default to None, will be auto-set
+                        help='Path to low-resolution training images (auto-set if None)')
+    parser.add_argument('--scale', type=int, default=4, help='Upscaling factor')
+    parser.add_argument('--batch_size', type=int, default=16, help='Training batch size (reduce for CPU/low VRAM)')
+    parser.add_argument('--subset', type=int, default=0, help='Use only N images for debugging (0 to use all)')
+    parser.add_argument('--num_workers', type=int, default=0, help='Number of workers for DataLoader (set to 0 for Mac/Windows usually)')
+    parser.add_argument('--patch_size', type=int, default=48, help='Size (height/width) of LR patches for training') # NEW ARGUMENT
+
+    # --- Model Args ---
+    parser.add_argument('--gen_features', type=int, default=64, help='Number of features in Generator')
+    parser.add_argument('--gen_blocks', type=int, default=16, help='Number of residual blocks in Generator (reduce for faster training/less memory)')
+    parser.add_argument('--disc_features', type=int, default=64, help='Number of starting features in Discriminator')
+    parser.add_argument('--disc_blocks', type=int, default=3, help='Number of conv blocks in Discriminator')
+
+    # --- Training Args ---
+    parser.add_argument('--epochs', type=int, default=100, help='Number of training epochs')
+    parser.add_argument('--lr_gen', type=float, default=1e-4, help='Learning rate for Generator')
+    parser.add_argument('--lr_disc', type=float, default=1e-4, help='Learning rate for Discriminator')
+    parser.add_argument('--lambda_content', type=float, default=0.01, help='Weight for Content Loss (L1)') # SRGAN paper uses 1e-2 for L1/MSE when combined with VGG
+    parser.add_argument('--lambda_percep', type=float, default=1.0, help='Weight for Perceptual Loss') # SRGAN paper uses 1.0
+    parser.add_argument('--lambda_adv', type=float, default=0.001, help='Weight for Adversarial Loss') # SRGAN paper uses 1e-3
+
+    # --- Other Args ---
+    parser.add_argument('--save_dir', type=str, default='checkpoints', help='Directory to save model checkpoints')
+    parser.add_argument('--save_interval', type=int, default=10, help='Save checkpoint every N epochs')
+    parser.add_argument('--cpu', action='store_true', help='Force training on CPU')
+    # parser.add_argument('--load_checkpoint', type=str, default=None, help='Path to checkpoint file to resume training') # Example for adding resume functionality
+
+    args = parser.parse_args()
+
+    # --- Set and Validate Directories ---
+    # Auto-set LR directory based on scale IF it wasn't provided via command line
+    if args.lr_dir is None:
+        args.lr_dir = f'./datasets/DIV2K/DIV2K_train_LR_bicubic/X{args.scale}'
+        print(f"LR directory not provided, automatically setting based on scale {args.scale} to: {args.lr_dir}")
+
+    # Validate HR directory
+    if not os.path.isdir(args.hr_dir):
+         print(f"\nERROR: High-Resolution directory not found at '{args.hr_dir}'")
+         print("Please ensure the directory exists or provide the correct path using --hr_dir.")
+         exit(1) # Exit if the directory is invalid
+    # Validate LR directory
+    if not os.path.isdir(args.lr_dir):
+         print(f"\nERROR: Low-Resolution directory not found at '{args.lr_dir}'")
+         print(f"Please ensure the directory exists (check scale factor {args.scale}?) or provide the correct path using --lr_dir.")
+         exit(1) # Exit if the directory is invalid
+
+    print("\n--- Training Configuration ---")
+    # Print configuration cleanly
+    config_dict = vars(args)
+    # Calculate terminal width for better formatting (optional)
+    try:
+        term_width = os.get_terminal_size().columns
+    except OSError:
+        term_width = 80 # Default if terminal size unavailable
+
+    print("-" * term_width)
+    for key, value in config_dict.items():
+        print(f"{key:<25}: {value}") # Format for alignment
+    print("-" * term_width)
+
+
+    # Start the training process
+    train(args)