Update app.py

app.py

@@ -1,21 +1,20 @@
+
 import gradio as gr
 import torch
 import numpy as np
 import cv2
-import matplotlib.pyplot as plt
 from arch import SegFormerUNet
 from albumentations import Compose, Resize, Normalize
 from albumentations.pytorch import ToTensorV2
 
+# Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = SegFormerUNet().to(device)
-checkpoint_path = "model/segformer_unet_focal_loss_97_63.pth"
-checkpoint = torch.load(checkpoint_path, map_location=device)
+checkpoint = torch.load("model/segformer_unet_focal_loss_97_63.pth", map_location=device)
 model.load_state_dict(checkpoint)
 model.eval()
-print("Model weights loaded successfully!")
 
-# Image Transformation
+# Transform
 transform = Compose([
     Resize(256, 256),
     Normalize(mean=[0.5], std=[0.5]),
@@ -23,26 +22,23 @@ transform = Compose([
 ])
 
 def process_image(image):
-    """Process uploaded image, perform segmentation, and compute energy output."""
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     transformed = transform(image=image)['image'].unsqueeze(0).to(device)
-    
     with torch.no_grad():
         output = model(transformed)
         pred_mask = torch.sigmoid(output).squeeze().cpu().numpy()
         pred_mask = (pred_mask > 0.5).astype(np.uint8)
-    
     area_m2 = np.sum(pred_mask) * (0.125 ** 2)
     energy_kwh = area_m2 * 0.19 * 1676.2 * 0.935 / 1000
-    
     return pred_mask * 255, f"Estimated Solar Panel Area: {area_m2:.2f} m²", f"Estimated Energy Output: {energy_kwh:.2f} MWh per year"
 
+# Gradio Interface
 demo = gr.Interface(
     fn=process_image,
-    inputs=gr.Image(type="numpy"),
+    inputs=gr.Image(type="numpy", value="PV08_315125_1194663.bmp"),  # <-- default image here
     outputs=[gr.Image(type="numpy"), gr.Text(), gr.Text()],
     title="Solar Panel Segmentation",
-    description="Upload an image to detect solar panels and estimate energy output.",
+    description="Upload an image to detect solar panels and estimate energy output."
 )
 
-demo.launch()
+demo.launch()