Spaces:
Running
Running
File size: 3,600 Bytes
2940390 d049335 2940390 d049335 2940390 d049335 2940390 |
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 |
import torch
import gradio as gr
from torchvision import transforms
from PIL import Image
import numpy as np
from utils.utils import load_restore_ckpt, load_embedder_ckpt
import os
from gradio_imageslider import ImageSlider
# Enforce CPU usage
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
embedder_model_path = "ckpts/embedder_model.tar" # Update with actual path to embedder checkpoint
restorer_model_path = "ckpts/onerestore_cdd-11.tar" # Update with actual path to restorer checkpoint
# Load models on CPU only
embedder = load_embedder_ckpt(device, freeze_model=True, ckpt_name=embedder_model_path)
restorer = load_restore_ckpt(device, freeze_model=True, ckpt_name=restorer_model_path)
# Define image preprocessing and postprocessing
transform_resize = transforms.Compose([
transforms.Resize([224,224]),
transforms.ToTensor()
])
def postprocess_image(tensor):
image = tensor.squeeze(0).cpu().detach().numpy()
image = (image) * 255 # Assuming output in [-1, 1], rescale to [0, 255]
image = np.clip(image, 0, 255).astype("uint8") # Clip values to [0, 255]
return Image.fromarray(image.transpose(1, 2, 0)) # Reorder to (H, W, C)
# Define the enhancement function
def enhance_image(image, degradation_type=None):
# Preprocess the image
input_tensor = torch.Tensor((np.array(image)/255).transpose(2, 0, 1)).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
lq_em = transform_resize(image).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
lq_em = transform_resize(image).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
# Generate embedding
if degradation_type == "auto" or degradation_type is None:
text_embedding, _, [text] = embedder(lq_em, 'image_encoder')
else:
text_embedding, _, [text] = embedder([degradation_type], 'text_encoder')
# Model inference
with torch.no_grad():
enhanced_tensor = restorer(input_tensor, text_embedding)
# Postprocess the output
return (image, postprocess_image(enhanced_tensor)), text
# Define the Gradio interface
def inference(image, degradation_type=None):
return enhance_image(image, degradation_type)
#### Image,Prompts examples
examples = [
['image/low_haze_rain_00469_01_lq.png'],
['image/low_haze_snow_00337_01_lq.png'],
]
# Create the Gradio app interface using updated API
interface = gr.Interface(
fn=inference,
inputs=[
gr.Image(type="pil", value="image/low_haze_rain_00469_01_lq.png"), # Image input
gr.Dropdown(['auto', 'low', 'haze', 'rain', 'snow',\
'low_haze', 'low_rain', 'low_snow', 'haze_rain',\
'haze_snow', 'low_haze_rain', 'low_haze_snow'], label="Degradation Type", value="auto") # Manual or auto degradation
],
outputs=[
ImageSlider(label="Restored Image",
type="pil",
show_download_button=True,
), # Enhanced image outputImageSlider(type="pil", show_download_button=True, ),
gr.Textbox(label="Degradation Type") # Display the estimated degradation type
],
title="Image Restoration with OneRestore",
description="Upload an image and enhance it using OneRestore model. You can choose to let the model automatically estimate the degradation type or set it manually.",
examples=examples,
)
# Launch the app
if __name__ == "__main__":
interface.launch()
|