Spaces:
Running
Running
| 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() |