Create app.py

app.py

@@ -0,0 +1,101 @@
+import torch
+from PIL import Image
+from RealESRGAN import RealESRGAN
+import gradio as gr
+import numpy as np
+import tempfile
+import time
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+def load_model(scale):
+    model = RealESRGAN(device, scale=scale)
+    weights_path = f'weights/RealESRGAN_x{scale}.pth'
+    try:
+        model.load_weights(weights_path, download=True)
+        print(f"Weights for scale {scale} loaded successfully.")
+    except Exception as e:
+        print(f"Error loading weights for scale {scale}: {e}")
+        model.load_weights(weights_path, download=False)
+    return model
+
+model2 = load_model(2)
+model4 = load_model(4)
+model8 = load_model(8)
+
+def enhance_image(image, scale):
+    try:
+        print(f"Enhancing image with scale {scale}...")
+        start_time = time.time()
+        image_np = np.array(image.convert('RGB'))
+        print(f"Image converted to numpy array: shape {image_np.shape}, dtype {image_np.dtype}")
+        
+        if scale == '2x':
+            result = model2.predict(image_np)
+        elif scale == '4x':
+            result = model4.predict(image_np)
+        else:
+            result = model8.predict(image_np)
+            
+        enhanced_image = Image.fromarray(np.uint8(result))
+        print(f"Image enhanced in {time.time() - start_time:.2f} seconds")
+        return enhanced_image
+    except Exception as e:
+        print(f"Error enhancing image: {e}")
+        return image
+
+def muda_dpi(input_image, dpi):
+    dpi_tuple = (dpi, dpi)
+    image = Image.fromarray(input_image.astype('uint8'), 'RGB')
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+    image.save(temp_file, format='PNG', dpi=dpi_tuple)
+    temp_file.close()
+    return Image.open(temp_file.name)
+
+def resize_image(input_image, width, height):
+    image = Image.fromarray(input_image.astype('uint8'), 'RGB')
+    resized_image = image.resize((width, height))
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+    resized_image.save(temp_file, format='PNG')
+    temp_file.close()
+    return Image.open(temp_file.name)
+
+def process_image(input_image, enhance, scale, adjust_dpi, dpi, resize, width, height):
+    original_image = Image.fromarray(input_image.astype('uint8'), 'RGB')
+    
+    if enhance:
+        original_image = enhance_image(original_image, scale)
+    
+    if adjust_dpi:
+        original_image = muda_dpi(np.array(original_image), dpi)
+        
+    if resize:
+        original_image = resize_image(np.array(original_image), width, height)
+        
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+    original_image.save(temp_file.name)
+    return original_image, temp_file.name
+
+iface = gr.Interface(
+    fn=process_image,
+    inputs=[
+        gr.Image(label="Upload"),
+        gr.Checkbox(label="Enhance Image"),
+        gr.Radio(['2x', '4x', '8x'], type="value", value='2x', label='Select Resolution model'),
+        gr.Checkbox(label="Apply DPI"),
+        gr.Number(label="DPI", value=300),
+        gr.Checkbox(label="Apply Resize"),
+        gr.Number(label="Width", value=512),
+        gr.Number(label="Height", value=512)
+    ],
+    outputs=[
+        gr.Image(label="Final Image"),
+        gr.File(label="Download Final Image")
+    ],
+    title="Image Enhancer",
+    description="Sorry for the inconvenience. The model is currently running on the CPU, which might affect performance. We appreciate your understanding.",
+    theme="NoCrypt/miku"
+    ]
+)
+
+iface.launch(debug=True)