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