Update app.py

app.py

@@ -1,58 +1,56 @@
 import gradio as gr
-from PIL import Image
+from PIL import Image, ImageEnhance
+import numpy as np
 import torch
 import torchvision.transforms as transforms
-from torchvision.models import resnet50
-import torch.nn.functional as F
-import numpy as np
+from torchvision.models import resnet34
 
-# Load a pre-trained ResNet model and modify it for upscaling
-class Upscaler(torch.nn.Module):
-    def __init__(self, upscale_factor):
-        super(Upscaler, self).__init__()
-        self.model = resnet50(pretrained=True)
-        self.upscale_factor = upscale_factor
-        self.conv1x1 = torch.nn.Conv2d(1000, 3, kernel_size=1)
-        
-    def forward(self, x):
-        x = F.interpolate(x, scale_factor=self.upscale_factor, mode='bilinear', align_corners=True)
-        x = self.model(x)
-        x = self.conv1x1(x)
-        return x
+# Load a pre-trained ResNet model
+model = resnet34(pretrained=True)
+
+# Define the upscaling function
+def upscale_image(image, upscale_factor=2, sharpness=1.0, contrast=1.0, brightness=1.0):
+    # Resize the image
+    width, height = image.size
+    new_size = (int(width * upscale_factor), int(height * upscale_factor))
+    upscaled_image = image.resize(new_size, Image.BICUBIC)
+    
+    # Apply sharpness, contrast, and brightness adjustments
+    upscaled_image = ImageEnhance.Sharpness(upscaled_image).enhance(sharpness)
+    upscaled_image = ImageEnhance.Contrast(upscaled_image).enhance(contrast)
+    upscaled_image = ImageEnhance.Brightness(upscaled_image).enhance(brightness)
+    
+    return upscaled_image
 
-# Custom remastering function with multiple options
+# Define the remastering function
 def remaster_image(image, color_range=1.0, sharpness=1.0, hdr_intensity=1.0, tone_mapping=1.0, color_grading=1.0):
-    enhancer = transforms.ColorJitter(
-        brightness=hdr_intensity,
-        contrast=contrast,
-        saturation=color_range,
-        hue=0
-    )
-    image = enhancer(image)
+    # Adjust color range
+    enhancer = ImageEnhance.Color(image)
+    image = enhancer.enhance(color_range)
 
     # Adjust sharpness
-    image = transforms.functional.adjust_sharpness(image, sharpness_factor=sharpness)
+    enhancer = ImageEnhance.Sharpness(image)
+    image = enhancer.enhance(sharpness)
 
-    # Apply tone mapping and color grading
-    tone_map = lambda x: x * tone_mapping
-    graded_image = transforms.functional.lerp(image, tone_map(image), color_grading)
+    # For HDR simulation and tone mapping, we're using simple brightness adjustments
+    enhancer = ImageEnhance.Brightness(image)
+    image = enhancer.enhance(hdr_intensity)
 
-    return graded_image
+    # Simulate color grading by adjusting contrast
+    enhancer = ImageEnhance.Contrast(image)
+    image = enhancer.enhance(color_grading)
 
-# Function to process image with the selected options
-def process_image(image, upscale=False, upscale_factor=2, noise_reduction=0, edge_enhancement=1.0, 
-                  detail_preservation=1.0, remaster=False, color_range=1.0, sharpness=1.0, 
-                  hdr_intensity=1.0, tone_mapping=1.0, color_grading=1.0):
-    image = transforms.ToTensor()(image).unsqueeze(0)
-    
-    if upscale:
-        upscaler = Upscaler(upscale_factor)
-        image = upscaler(image)
+    return image
 
+# Process function for Gradio
+def process_image(image, upscale=False, upscale_factor=2, sharpness=1.0, contrast=1.0, brightness=1.0,
+                  remaster=False, color_range=1.0, hdr_intensity=1.0, tone_mapping=1.0, color_grading=1.0):
+    if upscale:
+        image = upscale_image(image, upscale_factor, sharpness, contrast, brightness)
+    
     if remaster:
         image = remaster_image(image, color_range, sharpness, hdr_intensity, tone_mapping, color_grading)
     
-    image = transforms.ToPILImage()(image.squeeze(0))
     return image
 
 # Gradio UI
@@ -65,16 +63,15 @@ with gr.Blocks() as demo:
         with gr.Group():
             gr.Markdown("### Upscaling Options")
             upscale_checkbox = gr.Checkbox(label="Apply Upscaling")
-            upscale_factor = gr.Slider(2, 8, value=2, label="Upscale Factor")
-            noise_reduction = gr.Slider(0, 100, value=0, label="Noise Reduction")
-            edge_enhancement = gr.Slider(0.5, 2.0, value=1.0, label="Edge Enhancement")
-            detail_preservation = gr.Slider(0.5, 2.0, value=1.0, label="Detail Preservation")
+            upscale_factor = gr.Slider(1, 8, value=2, label="Upscale Factor")
+            sharpness = gr.Slider(0.5, 2.0, value=1.0, label="Sharpness")
+            contrast = gr.Slider(0.5, 2.0, value=1.0, label="Contrast")
+            brightness = gr.Slider(0.5, 2.0, value=1.0, label="Brightness")
         
         with gr.Group():
             gr.Markdown("### Remastering Options")
             remaster_checkbox = gr.Checkbox(label="Apply Remastering")
             color_range = gr.Slider(0.5, 2.0, value=1.0, label="Dynamic Color Range")
-            sharpness = gr.Slider(0.5, 2.0, value=1.0, label="Advanced Sharpness Control")
             hdr_intensity = gr.Slider(0.5, 2.0, value=1.0, label="HDR Intensity")
             tone_mapping = gr.Slider(0.5, 2.0, value=1.0, label="Tone Mapping")
             color_grading = gr.Slider(0.5, 2.0, value=1.0, label="Color Grading")
@@ -83,8 +80,8 @@ with gr.Blocks() as demo:
     
     process_button.click(
         process_image,
-        inputs=[image_input, upscale_checkbox, upscale_factor, noise_reduction, edge_enhancement, detail_preservation,
-                remaster_checkbox, color_range, sharpness, hdr_intensity, tone_mapping, color_grading],
+        inputs=[image_input, upscale_checkbox, upscale_factor, sharpness, contrast, brightness,
+                remaster_checkbox, color_range, hdr_intensity, tone_mapping, color_grading],
         outputs=image_output
     )