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