File size: 3,696 Bytes
3fa54b5
9980bbc
18b1364
9980bbc
 
 
 
3fa54b5
9980bbc
 
 
 
 
 
 
 
 
 
 
 
 
3fa54b5
9980bbc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3074852
9980bbc
 
 
 
 
 
 
3074852
9980bbc
 
 
 
 
 
18b1364
9980bbc
18b1364
bbf7bd0
18b1364
 
 
 
 
 
 
 
 
 
9980bbc
 
 
 
18b1364
 
 
 
9980bbc
 
 
 
 
3074852
18b1364
 
 
 
9980bbc
 
18b1364
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import gradio as gr
from PIL import Image
import torch
import torchvision.transforms as transforms
from torchvision.models import resnet50
import torch.nn.functional as F
import numpy as np

# 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

# Custom remastering function with multiple options
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 sharpness
    image = transforms.functional.adjust_sharpness(image, sharpness_factor=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)

    return graded_image

# 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)

    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
with gr.Blocks() as demo:
    with gr.Row():
        image_input = gr.Image(label="Upload Image", type="pil")
        image_output = gr.Image(label="Output Image")
    
    with gr.Row():
        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")
        
        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")
    
    process_button = gr.Button("Process Image")
    
    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],
        outputs=image_output
    )

demo.launch()