Update app.py

app.py

@@ -3,60 +3,76 @@ import cv2
 import gradio as gr
 from ultralytics import YOLO
 
-# --- Fix 1: Unzip model file (if needed) ---
-# Hugging Face Spaces will automatically unzip uploaded .zip files,
-# so you can skip the !unzip command in Python code.
-
-# --- Fix 2: Use correct model path ---
-# Place your model file in the same directory as app.py and use relative path
-MODEL_PATH = "yolov8_model.pt"  # Make sure this matches your uploaded filename
-
-# --- Fix 3: Add error handling ---
-if not os.path.exists(MODEL_PATH):
-    raise FileNotFoundError(f"Model file not found at {MODEL_PATH}. "
-                          "Please upload your model file to the Space.")
-
 # Load model
+MODEL_PATH = "yolov8_model.pt"
+if not os.path.exists(MODEL_PATH):
+    raise FileNotFoundError(f"Model file {MODEL_PATH} not found. Please upload it to your Space.")
 model = YOLO(MODEL_PATH)
 
 def detect_defects(image):
-    """Run object detection on input image and return annotated result."""
+    """Run object detection and return annotated image with results."""
     try:
         # Run inference
         results = model(image)
         
-        # Plot results (with bounding boxes)
-        annotated_img = results[0].plot(line_width=2)
+        # Get detection results
+        boxes = results[0].boxes
+        class_names = model.names
+        detections = []
+        
+        for box in boxes:
+            class_id = int(box.cls)
+            confidence = float(box.conf)
+            detections.append({
+                "class": class_names[class_id],
+                "confidence": confidence,
+                "box": box.xyxy[0].tolist()  # [x1,y1,x2,y2]
+            })
         
-        # Convert BGR to RGB (OpenCV uses BGR by default)
+        # Plot results
+        annotated_img = results[0].plot(line_width=2, font_size=10)
         annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
-        return annotated_img
+        
+        return annotated_img, detections
     
     except Exception as e:
         print(f"Error during detection: {e}")
-        return image  # Return original image if error occurs
-
-# Gradio UI
-title = "🔍 Steel Surface Defect Detector (YOLOv8)"
-description = "Upload a steel surface image to detect defects (crazing, scratches, etc.)."
+        return image, []
 
-# --- Fix 4: Use relative paths for example images ---
-examples = [
-    "example1.jpg",  # Make sure these exist in your Space
-    "example2.jpg"
-]
+# Custom CSS for better UI
+css = """
+.gradio-container {max-width: 800px !important}
+.output-image {border-radius: 8px}
+"""
 
-# Create interface
-interface = gr.Interface(
-    fn=detect_defects,
-    inputs=gr.Image(type="numpy"),
-    outputs=gr.Image(type="numpy"),
-    title=title,
-    description=description,
-    examples=examples,
-    allow_flagging="never"
-)
+# Interface components
+with gr.Blocks(css=css, title="Steel Defect Detector") as app:
+    gr.Markdown("## 🔍 Steel Surface Defect Detector (YOLOv8)")
+    gr.Markdown("Upload or drag-and-drop an image of steel surface to detect defects")
+    
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="Input Image", type="numpy")
+            submit_btn = gr.Button("Detect Defects", variant="primary")
+        
+        with gr.Column():
+            image_output = gr.Image(label="Detected Defects", elem_classes=["output-image"])
+            json_output = gr.JSON(label="Detection Results")
+    
+    # Example images
+    gr.Examples(
+        examples=[["example1.jpg"], ["example2.jpg"]],
+        inputs=image_input,
+        outputs=[image_output, json_output],
+        fn=detect_defects,
+        cache_examples=False
+    )
+    
+    submit_btn.click(
+        fn=detect_defects,
+        inputs=image_input,
+        outputs=[image_output, json_output]
+    )
 
-# Launch app
 if __name__ == "__main__":
-    interface.launch(debug=True)
+    app.launch(debug=True)