Spaces:
Running
Running
import gradio as gr | |
from PIL import Image, ImageEnhance | |
import numpy as np | |
import torch | |
import torchvision.transforms as transforms | |
from torchvision.models import resnet34 | |
# 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 | |
# 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) | |
# For HDR simulation and tone mapping, we're using simple brightness adjustments | |
enhancer = ImageEnhance.Brightness(image) | |
image = enhancer.enhance(hdr_intensity) | |
# Simulate color grading by adjusting contrast | |
enhancer = ImageEnhance.Contrast(image) | |
image = enhancer.enhance(color_grading) | |
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) | |
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") | |
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") | |
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, sharpness, contrast, brightness, | |
remaster_checkbox, color_range, hdr_intensity, tone_mapping, color_grading], | |
outputs=image_output | |
) | |
demo.launch() |