File size: 4,091 Bytes
3fa54b5
5d26a73
18b1364
772383d
 
5d26a73
772383d
 
 
3fa54b5
772383d
 
 
5d26a73
772383d
 
 
 
 
 
 
 
 
 
 
 
 
 
5d26a73
 
3fa54b5
5d26a73
9980bbc
5d26a73
 
 
9980bbc
 
5d26a73
 
9980bbc
772383d
5d26a73
 
772383d
5d26a73
 
772383d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9980bbc
5d26a73
772383d
5d26a73
772383d
5d26a73
9980bbc
 
18b1364
 
bbf7bd0
18b1364
 
 
 
 
 
 
 
 
 
5d26a73
18b1364
 
 
 
9980bbc
772383d
9980bbc
 
 
3074852
18b1364
 
 
 
772383d
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
92
93
94
95
96
97
98
99
100
101
102
103
104
import gradio as gr
from PIL import Image, ImageEnhance
import torch
import torch.nn.functional as F
from torchvision import transforms
from torchvision.models import resnet34
from torchvision.models.segmentation import deeplabv3_resnet50
import numpy as np
import cv2

# Load a pre-trained ResNet model for remastering
resnet_model = resnet34(pretrained=True)
resnet_model.eval()

# Load a pre-trained DeepLab model for segmentation (optional for advanced remastering)
deeplab_model = deeplabv3_resnet50(pretrained=True)
deeplab_model.eval()

# Define the upscaling function using super-resolution techniques
def upscale_image(image, upscale_factor=2):
    # Convert the image to a tensor and upscale it using a neural network
    preprocess = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.unsqueeze(0))
    ])
    img_tensor = preprocess(image)
    upscaled_tensor = F.interpolate(img_tensor, scale_factor=upscale_factor, mode='bicubic', align_corners=False)
    upscaled_image = transforms.ToPILImage()(upscaled_tensor.squeeze())
    
    return upscaled_image

# 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):
    # Adjust color range
    enhancer = ImageEnhance.Color(image)
    image = enhancer.enhance(color_range)

    # Adjust sharpness
    enhancer = ImageEnhance.Sharpness(image)
    image = enhancer.enhance(sharpness)

    # Apply a simulated HDR effect using tone mapping
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(hdr_intensity)
    
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(color_grading)
    
    # Optional: Use segmentation to remaster specific regions
    input_tensor = transforms.ToTensor()(image).unsqueeze(0)
    with torch.no_grad():
        output = deeplab_model(input_tensor)['out'][0]
    output_predictions = output.argmax(0)
    
    # Process each segmented region (e.g., sky, water) differently (optional)
    # Example: Apply a slight blur to the sky region to create a dreamy effect
    mask = output_predictions.byte().cpu().numpy()
    segmented_image = np.array(image)
    segmented_image[mask == 15] = cv2.GaussianBlur(segmented_image[mask == 15], (5, 5), 0)
    
    final_image = Image.fromarray(segmented_image)
    
    return final_image

# Process function for Gradio
def process_image(image, upscale=False, upscale_factor=2, remaster=False, color_range=1.0, sharpness=1.0, hdr_intensity=1.0, tone_mapping=1.0, color_grading=1.0):
    if upscale:
        image = upscale_image(image, upscale_factor)
    
    if remaster:
        image = remaster_image(image, color_range, sharpness, hdr_intensity, tone_mapping, color_grading)
    
    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(1, 8, value=2, label="Upscale Factor")
        
        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="Sharpness")
            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, remaster_checkbox, color_range, sharpness, hdr_intensity, tone_mapping, color_grading],
        outputs=image_output
    )

demo.launch()