Update app.py

app.py

@@ -1,12 +1,13 @@
+import gradio as gr
 from PIL import Image, ImageDraw
 import torch
 from torchvision import models, transforms
 
-# Load pretrained Faster R-CNN model
+# ✅ Directly load a pretrained model (you would need to fine-tune this model on structural defects)
 model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
 model.eval()
 
-# COCO Class Labels
+# Define labels (this could be adjusted if you have specific defect labels)
 COCO_INSTANCE_CATEGORY_NAMES = [
     '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
     'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter',
@@ -23,12 +24,11 @@ COCO_INSTANCE_CATEGORY_NAMES = [
 
 # Image transformations
 transform = transforms.Compose([
-    transforms.Resize((800, 800)),  # Resize to 800x800 for faster processing
     transforms.ToTensor(),
 ])
 
+# Default severity labels (can be customized based on detection confidence)
 def get_severity(score):
-    """Assign severity based on confidence score."""
     if score >= 0.9:
         return "High"
     elif score >= 0.7:
@@ -37,7 +37,6 @@ def get_severity(score):
         return "Low"
 
 def detect_defects(image):
-    # Resize and transform the image
     image_tensor = transform(image).unsqueeze(0)
 
     with torch.no_grad():
@@ -49,7 +48,7 @@ def detect_defects(image):
     output = []
     for i in range(len(predictions[0]['boxes'])):
         score = predictions[0]['scores'][i].item()
-        if score < 0.7:  # Filter low-confidence predictions
+        if score < 0.7:  # filter low-confidence predictions
             continue
 
         box = predictions[0]['boxes'][i].tolist()
@@ -64,8 +63,22 @@ def detect_defects(image):
             "severity": severity
         })
 
-        # Draw bounding box and label
         draw.rectangle(box, outline="red", width=3)
         draw.text((box[0], box[1]), f"{label}: {severity}", fill="red")
 
     return result_image, output
+
+# Gradio Interface
+demo = gr.Interface(
+    fn=detect_defects,
+    inputs=gr.Image(type="pil", label="Upload Drone Image"),
+    outputs=[
+        gr.Image(label="Detection Result"),
+        gr.JSON(label="Detected Faults with Severity")
+    ],
+    title="Structural Defect Detection (Demo)",
+    description="Detects general objects using Faster R-CNN pretrained on COCO. This model needs fine-tuning for structural defects like cracks, rust, and spalling."
+)
+
+if __name__ == "__main__":
+    demo.launch()