Update app.py

app.py

@@ -6,11 +6,11 @@ from PIL import Image
 import numpy as np
 import os
 
-# Check if CUDA is available
+# ✅ Check if CUDA is available
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 print(f"Using device: {device}")
 
-# Model definition
+# ✅ Model definition
 class DeblurNet(nn.Module):
     def __init__(self):
         super(DeblurNet, self).__init__()
@@ -53,13 +53,19 @@ class DeblurNet(nn.Module):
         x = self.final_conv(x)
         return torch.tanh(x)
 
-# Load model
+# ✅ Load model
 model = DeblurNet().to(device)
 model_path = os.path.join('model', 'best_deblur_model.pth')
-model.load_state_dict(torch.load(model_path, map_location=device))
-model.eval()
 
-# Image processing functions
+# Ensure model path exists before loading
+if os.path.exists(model_path):
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.eval()
+    print("Model loaded successfully.")
+else:
+    print(f"Model file not found at {model_path}. Please check the path.")
+
+# ✅ Image processing functions
 transform = transforms.Compose([
     transforms.Resize((256, 256)),
     transforms.ToTensor(),
@@ -67,6 +73,7 @@ transform = transforms.Compose([
 ])
 
 def postprocess_image(tensor):
+    """Post-process the output tensor into a displayable image."""
     tensor = tensor * 0.5 + 0.5
     tensor = torch.clamp(tensor, 0, 1)
     image = tensor.cpu().detach().numpy()
@@ -74,11 +81,12 @@ def postprocess_image(tensor):
     return (image * 255).astype(np.uint8)
 
 def deblur_image(filepath):
-    if filepath is None:
+    """Deblurs the uploaded image."""
+    if not filepath:
         return None
-    
+
     try:
-        # Load the image from the filepath
+        # Load image from filepath
         input_image = Image.open(filepath).convert("RGB")
 
         # Save original size
@@ -91,24 +99,30 @@ def deblur_image(filepath):
         with torch.no_grad():
             output_tensor = model(input_tensor)
 
-        # Postprocess
+        # Post-process
         output_image = postprocess_image(output_tensor[0])
 
         # Resize back to original size
         output_image = Image.fromarray(output_image).resize(original_size)
+
         return np.array(output_image)
 
     except Exception as e:
         print(f"Error processing image: {e}")
         return None
 
-# ✅ Gradio interface using gr.File with correct type
+# ✅ Custom CSS for styling and hiding share button
 custom_css = """
 /* Hide Gradio's footer and header */
 footer, header, .gradio-footer, .gradio-header {
     display: none !important;
 }
 
+/* Hide share button */
+.share-wrap {
+    display: none !important;
+}
+
 /* Non-draggable images */
 img {
     pointer-events: none !important;
@@ -140,11 +154,6 @@ body, .gradio-container {
     color: white !important;
     border: 1px solid #333333 !important;
 }
-
-/* Hide the share button icon in the output box */
-.gradio-container .gr-image-output .wrap.svelte-1ipelgc {
-    display: none !important;
-}
 """
 
 # ✅ Gradio interface
@@ -157,5 +166,6 @@ demo = gr.Interface(
     css=custom_css
 )
 
+# ✅ Launch Gradio app
 if __name__ == "__main__":
     demo.launch()