Update app.py

app.py

@@ -1,39 +1,40 @@
+import os
 import gradio as gr
 from ultralytics import YOLO
-import torch
 from PIL import Image
-import os
 
-# Load model (you can switch to 'fasterrcnn' based loading if needed)
-model = YOLO("model/best.pt")  # fine-tuned YOLOv8 model
+# Prevent config error on HF Spaces
+os.environ["YOLO_CONFIG_DIR"] = "/tmp/ultralytics"
+
+model = YOLO("model/best.pt")
 
-# Define defect labels if custom
 DEFECT_LABELS = ['crack', 'spalling', 'rust', 'deformation']
 
-# Inference function
 def detect_defects(image):
     results = model(image)
-    annotated_img = results[0].plot()  # Draw boxes
-    predictions = results[0].boxes.data.cpu().numpy()  # [x1, y1, x2, y2, conf, class]
+    rendered = results[0].plot()
+    predictions = results[0].boxes.data.cpu().numpy()
 
-    # Create readable output
-    output = []
+    output_lines = []
     for pred in predictions:
-        x1, y1, x2, y2, conf, cls_id = pred
-        label = DEFECT_LABELS[int(cls_id)]
-        output.append(f"{label}: {conf:.2f}")
+        conf = float(pred[4])
+        class_id = int(pred[5])
+        label = DEFECT_LABELS[class_id] if class_id < len(DEFECT_LABELS) else f"Class {class_id}"
+        output_lines.append(f"{label}: {conf:.2f}")
 
-    return Image.fromarray(annotated_img), "\n".join(output)
+    return Image.fromarray(rendered), "\n".join(output_lines)
 
-# Gradio UI
-interface = gr.Interface(
+demo = gr.Interface(
     fn=detect_defects,
-    inputs=gr.Image(type="pil"),
-    outputs=[gr.Image(type="pil", label="Detected Image"), gr.Textbox(label="Detected Faults")],
-    title="Structural Defect Detection",
-    description="Upload drone-captured image to detect cracks, rust, spalling, and deformations.",
+    inputs=gr.Image(type="pil", label="Upload Drone Image"),
+    outputs=[
+        gr.Image(type="pil", label="Detected Structural Defects"),
+        gr.Textbox(label="Detected Faults and Confidence")
+    ],
+    title="🧱 Structural Defect Detection",
+    description="Upload drone-captured image to detect cracks, spalling, rust, and deformation.",
     allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch()