First model version

app.py

@@ -0,0 +1,218 @@
+import os
+import streamlit as st
+import torch
+import torchvision.transforms as transforms
+from PIL import Image, ImageOps, ImageEnhance, ImageFilter
+import numpy as np
+import time
+import io
+
+
+# Metrics imports
+from skimage.metrics import peak_signal_noise_ratio as psnr
+from skimage.metrics import structural_similarity as ssim
+
+# Model imports
+from models.srcnn import SRCNN
+from models.vdsr import VDSR
+from models.edsr import EDSR
+
+# Cache for loaded models
+model_cache = {}
+
+def load_model(model_name):
+    """
+    Load super-resolution model with optional scale factor
+    
+    Args:
+        model_name (str): Name of the model (SRCNN, VDSR, EDSR)
+        scale_factor (int): Upscaling factor (2, 3, or 4)
+    
+    Returns:
+        torch.nn.Module: Loaded model
+    """
+    try:
+        # Check if model is already in the cache
+        if model_name in model_cache:
+            return model_cache[model_name]
+        
+        if model_name == 'SRCNN':
+            model = SRCNN()
+        elif model_name == 'VDSR':
+            model = VDSR()
+        else:
+            model = EDSR()
+            
+        # Load pre-trained weights if available
+        weight_path = f'checkpoints/{model_name.lower()}_best.pth'
+        if os.path.exists(weight_path):
+            model.load_state_dict(torch.load(weight_path, map_location=torch.device('cpu'), weights_only=True))
+        else:
+            st.warning(f"No pre-trained weights found for the {model_name} model. Using randomly initialized weights.")
+        
+        model.eval()
+        
+        # Cache the loaded model
+        model_cache[model_name] = model
+        return model
+    except Exception as e:
+        st.error(f"Error loading {model_name} model: {e}")
+        return None
+
+def process_image(image, model):
+    # Convert to YCbCr and extract Y channel
+    ycbcr = image.convert('YCbCr')
+    y, cb, cr = ycbcr.split()
+    
+    # Transform Y channel
+    transform = transforms.Compose([
+        transforms.ToTensor()
+    ])
+    
+    input_tensor = transform(y).unsqueeze(0)
+    
+    # Process through model
+    with torch.no_grad():
+        output = model(input_tensor)
+    
+    # Post-process output
+    output = output.squeeze().clamp(0, 1).numpy()
+    output_y = Image.fromarray((output * 255).astype(np.uint8))
+    
+    # Merge channels back
+    output_ycbcr = Image.merge('YCbCr', [output_y, cb, cr])
+    output_rgb = output_ycbcr.convert('RGB')
+    
+    return output_rgb
+
+def calculate_image_metrics(original, enhanced):
+    """
+    Calculate image quality metrics
+    
+    Args:
+        original (np.ndarray): Original image
+        enhanced (np.ndarray): Enhanced image
+    
+    Returns:
+        dict: Quality metrics
+    """
+    try:
+        # Ensure images are the same size
+        min_height = min(original.shape[0], enhanced.shape[0])
+        min_width = min(original.shape[1], enhanced.shape[1])
+        
+        # Resize images to the smallest common size
+        original = original[:min_height, :min_width]
+        enhanced = enhanced[:min_height, :min_width]
+        
+        # Calculate SSIM with an explicit window size
+        win_size = min(7, min(min_height, min_width))
+        if win_size % 2 == 0:
+            win_size -= 1  # Ensure odd window size
+        
+        return {
+            'PSNR': psnr(original, enhanced),
+            'SSIM': ssim(original, enhanced, multichannel=True, win_size=win_size, channel_axis=-1)
+        }
+    except Exception as e:
+        st.error(f"Error calculating metrics: {e}")
+        return {'PSNR': 0, 'SSIM': 0}
+
+def main():
+    st.set_page_config(
+        page_title="Super Resolution Comparison",
+        page_icon="🖼️",
+        layout="wide"
+    )
+    
+    st.title("🚀 Super Resolution Model Comparison")
+    st.write("Upload a low-resolution image and compare different super-resolution models.")
+    
+
+
+    
+    # File Upload
+    uploaded_file = st.file_uploader(
+        "Choose an image",
+        type=['png', 'jpg', 'jpeg'],
+        help="Upload a low-resolution image for enhancement"
+    )
+    
+    if uploaded_file is not None:
+        # Load input image
+        input_image = Image.open(uploaded_file)
+        input_array = np.array(input_image)
+        
+        st.subheader("📸 Original Image")
+        st.image(input_image, caption="Low-Resolution Input", use_column_width=True)
+        
+        # Model Names
+        model_names = ['SRCNN', 'VDSR', 'EDSR']
+        
+        # Performance and Quality Storage
+        processing_times = {}
+        quality_metrics = {}
+        enhanced_images = {}
+        
+        # Process images
+        columns = st.columns(len(model_names))
+        for i, model_name in enumerate(model_names):
+            with columns[i]:
+                st.subheader(f"{model_name} Model")
+                
+                # Load model
+                model = load_model(model_name)
+                
+                if model:
+                    # Time the processing
+                    start_time = time.time()
+                    enhanced_image = process_image(input_image, model)
+                    processing_time = time.time() - start_time
+                    
+                    if enhanced_image:
+                        # Display enhanced image
+                        st.image(enhanced_image, caption=f"{model_name} Output", use_column_width=True)
+                        
+                        # Calculate metrics
+                        enhanced_array = np.array(enhanced_image)
+                        metrics = calculate_image_metrics(input_array, enhanced_array)
+                        
+                        # Store results
+                        processing_times[model_name] = processing_time
+                        quality_metrics[model_name] = metrics
+                        enhanced_images[model_name] = enhanced_image
+        
+        # Performance Metrics Section
+        st.subheader("📊 Performance Metrics")
+        metric_cols = st.columns(len(model_names))
+        
+        for i, (model, time_val) in enumerate(processing_times.items()):
+            with metric_cols[i]:
+                st.metric(f"{model} Processing Time", f"{time_val:.4f} seconds")
+
+        # Quality Metrics Section
+        st.subheader("🔍 Image Quality Assessment")
+        quality_cols = st.columns(len(model_names))
+
+        for i, (model, metrics) in enumerate(quality_metrics.items()):
+            with quality_cols[i]:
+                st.metric(f"{model} PSNR", f"{metrics['PSNR']:.2f} dB")
+                st.metric(f"{model} SSIM", f"{metrics['SSIM']:.4f}")
+
+        # Download Section
+        st.subheader("💾 Download Enhanced Images")
+        download_cols = st.columns(len(model_names))
+
+        for i, (model, image) in enumerate(enhanced_images.items()):
+            with download_cols[i]:
+                buffered = io.BytesIO()
+                image.save(buffered, format="PNG")
+                st.download_button(
+                    label=f"Download {model} Image",
+                    data=buffered.getvalue(),
+                    file_name=f"{model}_enhanced.png",
+                    mime="image/png"
+                )
+
+if __name__ == "__main__":
+    main()

checkpoints/edsr_best.pth

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

checkpoints/srcnn_best.pth

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

checkpoints/vdsr_best.pth

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

inference.py

@@ -0,0 +1,86 @@
+# inference.py
+import os
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
+import streamlit as st
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import numpy as np
+from models.srcnn import SRCNN
+from models.vdsr import VDSR
+from models.edsr import EDSR
+
+def load_model(model_name):
+    if model_name == 'SRCNN':
+        model = SRCNN()
+    elif model_name == 'VDSR':
+        model = VDSR()
+    else:
+        model = EDSR()
+    
+    model.load_state_dict(torch.load(f'checkpoints/{model_name.lower()}_best.pth', map_location=torch.device('cpu')))
+    model.eval()
+    return model
+
+def process_image(image, model):
+    # Convert to YCbCr and extract Y channel
+    ycbcr = image.convert('YCbCr')
+    y, cb, cr = ycbcr.split()
+    
+    # Transform Y channel
+    transform = transforms.Compose([
+        transforms.ToTensor()
+    ])
+    
+    input_tensor = transform(y).unsqueeze(0)
+    
+    # Process through model
+    with torch.no_grad():
+        output = model(input_tensor)
+    
+    # Post-process output
+    output = output.squeeze().clamp(0, 1).numpy()
+    output_y = Image.fromarray((output * 255).astype(np.uint8))
+    
+    # Merge channels back
+    output_ycbcr = Image.merge('YCbCr', [output_y, cb, cr])
+    output_rgb = output_ycbcr.convert('RGB')
+    
+    return output_rgb
+
+def main():
+    st.title("Super Resolution Model Comparison")
+    st.write("Upload a low-resolution image to compare SRCNN, VDSR, and EDSR models")
+    
+    # File uploader
+    uploaded_file = st.file_uploader("Choose an image", type=['png', 'jpg', 'jpeg'])
+    
+    if uploaded_file is not None:
+        # Load and display input image
+        input_image = Image.open(uploaded_file)
+        st.subheader("Input Image")
+        st.image(input_image, caption="Original Image")
+        
+        # Process with each model
+        col1, col2, col3 = st.columns(3)
+        
+        with col1:
+            st.subheader("SRCNN")
+            model = load_model('SRCNN')
+            srcnn_output = process_image(input_image, model)
+            st.image(srcnn_output, caption="SRCNN Output")
+        
+        with col2:
+            st.subheader("VDSR")
+            model = load_model('VDSR')
+            vdsr_output = process_image(input_image, model)
+            st.image(vdsr_output, caption="VDSR Output")
+        
+        with col3:
+            st.subheader("EDSR")
+            model = load_model('EDSR')
+            edsr_output = process_image(input_image, model)
+            st.image(edsr_output, caption="EDSR Output")
+
+if __name__ == "__main__":
+    main()

metrics.py

@@ -0,0 +1,123 @@
+# metrics.py
+import os
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
+import torch
+import torch.nn as nn
+import numpy as np
+from PIL import Image
+import torchvision.transforms as transforms
+from models.srcnn import SRCNN
+from models.vdsr import VDSR
+from models.edsr import EDSR
+import math
+from skimage.metrics import structural_similarity as ssim
+import matplotlib.pyplot as plt
+
+
+def calculate_psnr(img1, img2):
+    mse = torch.mean((img1 - img2) ** 2)
+    if mse == 0:
+        return float('inf')
+    return 20 * math.log10(1.0 / math.sqrt(mse.item()))
+
+def process_image(model, lr_img):
+    with torch.no_grad():
+        # Convert to YCbCr and extract Y channel
+        ycbcr = lr_img.convert('YCbCr')
+        y, cb, cr = ycbcr.split()
+        
+        # Transform Y channel
+        transform = transforms.Compose([transforms.ToTensor()])
+        input_tensor = transform(y).unsqueeze(0)
+        
+        # Process through model
+        output = model(input_tensor)
+        
+        # Post-process output
+        output = output.squeeze().clamp(0, 1).numpy()
+        output_y = Image.fromarray((output * 255).astype(np.uint8))
+        
+        # Merge channels back
+        output_ycbcr = Image.merge('YCbCr', [output_y, cb, cr])
+        return output_ycbcr.convert('RGB')
+
+def calculate_ssim(img1, img2):
+    # Move channel axis to the end for SSIM calculation
+    img1_np = img1.cpu().numpy().transpose(1, 2, 0)
+    img2_np = img2.cpu().numpy().transpose(1, 2, 0)
+    return ssim(img1_np, img2_np, data_range=1.0, channel_axis=2, win_size=7)
+
+def evaluate_models(test_image_path):
+    # Load models
+    models = {
+        'SRCNN': SRCNN(),
+        'VDSR': VDSR(),
+        'EDSR': EDSR()
+    }
+    
+    # Load weights
+    for name, model in models.items():
+        model.load_state_dict(torch.load(f'checkpoints/{name.lower()}_best.pth', weights_only=True))
+        model.eval()
+    
+    # Load test image
+    lr_img = Image.open(test_image_path)
+    hr_img = Image.open(test_image_path)  # Using same image as reference
+    
+    # Results dictionary
+    results = {model_name: {} for model_name in models.keys()}
+    
+    # Process image with each model and calculate metrics
+    for name, model in models.items():
+        # Generate SR image
+        sr_img = process_image(model, lr_img)
+        
+        # Convert images to tensors for metric calculation
+        transform = transforms.Compose([
+            transforms.Resize((256, 256)),  # Ensure minimum size for SSIM
+            transforms.ToTensor()
+        ])
+        
+        sr_tensor = transform(sr_img)
+        hr_tensor = transform(hr_img)
+        
+        # Calculate metrics
+        results[name]['PSNR'] = calculate_psnr(sr_tensor, hr_tensor)
+        results[name]['SSIM'] = calculate_ssim(sr_tensor, hr_tensor)
+        
+        # Save output images
+        sr_img.save(f'results/{name.lower()}_output.png')
+    
+    # Display results
+    print("\nModel Performance Metrics:")
+    print("-" * 50)
+    print(f"{'Model':<10} {'PSNR (dB)':<15} {'SSIM':<15}")
+    print("-" * 50)
+    
+    for model_name, metrics in results.items():
+        print(f"{model_name:<10} {metrics['PSNR']:<15.2f} {metrics['SSIM']:<15.4f}")
+    
+    # Plot results
+    plt.figure(figsize=(12, 6))
+    
+    # PSNR comparison
+    plt.subplot(1, 2, 1)
+    plt.bar(results.keys(), [m['PSNR'] for m in results.values()])
+    plt.title('PSNR Comparison')
+    plt.ylabel('PSNR (dB)')
+    
+    # SSIM comparison
+    plt.subplot(1, 2, 2)
+    plt.bar(results.keys(), [m['SSIM'] for m in results.values()])
+    plt.title('SSIM Comparison')
+    plt.ylabel('SSIM')
+    
+    plt.tight_layout()
+    plt.savefig('results/metrics_comparison.png')
+    plt.close()
+
+if __name__ == "__main__":
+    import os
+    os.makedirs('results', exist_ok=True)
+    test_image_path = r"data\DIV2K_train_LR_bicubic_X4\DIV2K_train_LR_bicubic\X4\0001x4.png"  # Replace with your test image path
+    evaluate_models(test_image_path)

models/edsr.py

@@ -0,0 +1,51 @@
+# models/edsr.py
+import torch
+import torch.nn as nn
+
+class ResBlock(nn.Module):
+    def __init__(self, n_feats, kernel_size, bias=True, res_scale=1):
+        super(ResBlock, self).__init__()
+        m = []
+        for i in range(2):
+            m.append(nn.Conv2d(n_feats, n_feats, kernel_size, padding=(kernel_size//2), bias=bias))
+            if i == 0:
+                m.append(nn.ReLU(True))
+        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 EDSR(nn.Module):
+    def __init__(self, n_resblocks=16, n_feats=64, scale=4):
+        super(EDSR, self).__init__()
+        kernel_size = 3
+        self.scale = scale
+        
+        # Define head module
+        m_head = [nn.Conv2d(1, n_feats, kernel_size, padding=(kernel_size//2))]
+
+        # Define body module
+        m_body = [
+            ResBlock(n_feats, kernel_size) \
+            for _ in range(n_resblocks)
+        ]
+        m_body.append(nn.Conv2d(n_feats, n_feats, kernel_size, padding=(kernel_size//2)))
+
+        # Define tail module
+        m_tail = [
+            nn.Conv2d(n_feats, 1, kernel_size, padding=(kernel_size//2))
+        ]
+
+        self.head = nn.Sequential(*m_head)
+        self.body = nn.Sequential(*m_body)
+        self.tail = nn.Sequential(*m_tail)
+
+    def forward(self, x):
+        x = self.head(x)
+        res = self.body(x)
+        res += x
+        x = self.tail(res)
+        return x

models/srcnn.py

@@ -0,0 +1,16 @@
+# models/srcnn.py
+import torch.nn as nn
+
+class SRCNN(nn.Module):
+    def __init__(self):
+        super(SRCNN, self).__init__()
+        self.conv1 = nn.Conv2d(1, 64, kernel_size=9, padding=4)
+        self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0)
+        self.conv3 = nn.Conv2d(32, 1, kernel_size=5, padding=2)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.conv2(x))
+        x = self.conv3(x)
+        return x

models/vdsr.py

@@ -0,0 +1,39 @@
+# models/vdsr.py
+import torch.nn as nn
+from math import sqrt
+
+class Conv_ReLU_Block(nn.Module):
+    def __init__(self):
+        super(Conv_ReLU_Block, self).__init__()
+        self.conv = nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False)
+        self.relu = nn.ReLU(inplace=True)
+        
+    def forward(self, x):
+        return self.relu(self.conv(x))
+
+class VDSR(nn.Module):
+    def __init__(self):
+        super(VDSR, self).__init__()
+        self.residual_layer = self.make_layer(Conv_ReLU_Block, 18)
+        self.input = nn.Conv2d(1, 64, kernel_size=3, padding=1, bias=False)
+        self.output = nn.Conv2d(64, 1, kernel_size=3, padding=1, bias=False)
+        self.relu = nn.ReLU(inplace=True)
+        
+        # Initialize weights
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, sqrt(2. / n))
+                
+    def make_layer(self, block, num_of_layer):
+        layers = []
+        for _ in range(num_of_layer):
+            layers.append(block())
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        residual = x
+        out = self.relu(self.input(x))
+        out = self.residual_layer(out)
+        out = self.output(out)
+        return out + residual

requirements.txt

@@ -0,0 +1,6 @@
+streamlit 
+torch 
+numpy 
+torchvision
+
+scikit-image

train.py

@@ -0,0 +1,137 @@
+# train.py
+import os
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from utils.dataset import DIV2KDataset
+from models.srcnn import SRCNN
+from models.vdsr import VDSR
+from models.edsr import EDSR
+import math
+import numpy as np
+
+class EarlyStopping:
+    def __init__(self, patience=7, min_delta=0.01, min_psnr_improvement=0.1):
+        self.patience = patience
+        self.min_delta = min_delta
+        self.min_psnr_improvement = min_psnr_improvement
+        self.counter = 0
+        self.best_loss = None
+        self.best_psnr = None
+        self.early_stop = False
+        
+    def __call__(self, loss, psnr):
+        if self.best_loss is None:
+            self.best_loss = loss
+            self.best_psnr = psnr
+        elif (loss > self.best_loss - self.min_delta) and (psnr < self.best_psnr + self.min_psnr_improvement):
+            self.counter += 1
+            print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
+            if self.counter >= self.patience:
+                self.early_stop = True
+        else:
+            self.best_loss = min(loss, self.best_loss)
+            self.best_psnr = max(psnr, self.best_psnr)
+            self.counter = 0
+
+def calculate_psnr(img1, img2):
+    mse = torch.mean((img1 - img2) ** 2)
+    if mse == 0:
+        return float('inf')
+    return 20 * math.log10(1.0 / math.sqrt(mse.item()))
+
+def train_model(model_name, train_loader, val_loader, device, num_epochs=100):
+    # Initialize model
+    if model_name == 'srcnn':
+        model = SRCNN()
+    elif model_name == 'vdsr':
+        model = VDSR()
+    else:
+        model = EDSR()
+    
+    model = model.to(device)
+    criterion = nn.MSELoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.0001)
+    
+    # Initialize early stopping
+    early_stopping = EarlyStopping(patience=10, min_delta=0.00001, min_psnr_improvement=0.1)
+    best_psnr = 0
+    
+    for epoch in range(num_epochs):
+        # Training
+        model.train()
+        train_loss = 0
+        num_batches = 0
+        
+        for batch_idx, (lr_img, hr_img) in enumerate(train_loader):
+            lr_img, hr_img = lr_img.to(device), hr_img.to(device)
+            
+            optimizer.zero_grad()
+            output = model(lr_img)
+            loss = criterion(output, hr_img)
+            loss.backward()
+            optimizer.step()
+            
+            train_loss += loss.item()
+            num_batches += 1
+            
+            if batch_idx % 100 == 0:
+                print(f'Train Epoch: {epoch} [{batch_idx}/{len(train_loader)}]\tLoss: {loss.item():.6f}')
+        
+        avg_train_loss = train_loss / num_batches
+        
+        # Validation
+        model.eval()
+        val_psnr = 0
+        with torch.no_grad():
+            for lr_img, hr_img in val_loader:
+                lr_img, hr_img = lr_img.to(device), hr_img.to(device)
+                output = model(lr_img)
+                val_psnr += calculate_psnr(output, hr_img)
+        
+        val_psnr /= len(val_loader)
+        print(f'Epoch: {epoch}, Average Loss: {avg_train_loss:.6f}, Average PSNR: {val_psnr:.2f}dB')
+        
+        # Early stopping check
+        early_stopping(avg_train_loss, val_psnr)
+        if early_stopping.early_stop:
+            print(f"Early stopping triggered at epoch {epoch}")
+            break
+        
+        # Save best model
+        if val_psnr > best_psnr:
+            best_psnr = val_psnr
+            torch.save(model.state_dict(), f'checkpoints/{model_name}_best.pth')
+            print(f'Saved new best model with PSNR: {best_psnr:.2f}dB')
+
+def main():
+    # Setup
+    device = torch.device('cpu')
+    
+    # Data paths
+    train_hr_dir = 'data/DIV2K_train_HR/DIV2K_train_HR/'
+    train_lr_dir = 'data/DIV2K_train_LR_bicubic_X4/DIV2K_train_LR_bicubic/X4'
+    val_hr_dir = 'data/DIV2K_valid_HR/DIV2K_valid_HR'
+    val_lr_dir = 'data/DIV2K_valid_LR_bicubic_X4/DIV2K_valid_LR_bicubic/X4'
+    
+    # Create datasets
+    train_dataset = DIV2KDataset(train_hr_dir, train_lr_dir, patch_size=48)
+    val_dataset = DIV2KDataset(val_hr_dir, val_lr_dir, patch_size=48)
+    
+    # Create dataloaders
+    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=4)
+    val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)
+    
+    # Create checkpoints directory
+    os.makedirs('checkpoints', exist_ok=True)
+    
+    # Train models
+    models = ['edsr']
+    for model_name in models:
+        print(f'Training {model_name.upper()}...')
+        train_model(model_name, train_loader, val_loader, device)
+
+if __name__ == '__main__':
+    main()

utils/dataset.py

@@ -0,0 +1,53 @@
+# utils/dataset.py
+from torch.utils.data import Dataset
+from PIL import Image
+import torchvision.transforms as transforms
+import os
+
+class DIV2KDataset(Dataset):
+    def __init__(self, hr_dir, lr_dir, patch_size=96, upscale_factor=4):
+        self.hr_files = sorted(os.listdir(hr_dir))
+        self.lr_files = sorted(os.listdir(lr_dir))
+        self.hr_dir = hr_dir
+        self.lr_dir = lr_dir
+        self.patch_size = patch_size
+        self.upscale_factor = upscale_factor
+        
+        # LR transform
+        self.lr_transform = transforms.Compose([
+            transforms.Resize((patch_size//upscale_factor, patch_size//upscale_factor), 
+                            interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.ToTensor()
+        ])
+        
+        # HR transform
+        self.hr_transform = transforms.Compose([
+            transforms.Resize((patch_size, patch_size),
+                            interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.ToTensor()
+        ])
+        
+        # Upscale LR images to match HR size for SRCNN
+        self.lr_upscale = transforms.Compose([
+            transforms.Resize((patch_size, patch_size),
+                            interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.ToTensor()
+        ])
+        
+    def __getitem__(self, idx):
+        # Load images and convert to YCbCr
+        hr_img = Image.open(os.path.join(self.hr_dir, self.hr_files[idx])).convert('YCbCr')
+        lr_img = Image.open(os.path.join(self.lr_dir, self.lr_files[idx])).convert('YCbCr')
+        
+        # Extract Y channel
+        hr_y, _, _ = hr_img.split()
+        lr_y, _, _ = lr_img.split()
+        
+        # For SRCNN, we need to upscale LR images first
+        lr_y_upscaled = self.lr_upscale(lr_y)
+        hr_y_tensor = self.hr_transform(hr_y)
+        
+        return lr_y_upscaled, hr_y_tensor
+        
+    def __len__(self):
+        return len(self.hr_files)