Update app.py

app.py

@@ -1,111 +1,39 @@
 import gradio as gr
-import torch
 from ultralytics import YOLO
+import torch
 from PIL import Image
-import numpy as np
-import cv2
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+import os
 
-# Load the fine-tuned YOLOv8 model from Hugging Face
-try:
-    model = YOLO("path/to/your/fine-tuned-yolov8-model.pt")  # Replace with your model path
-    logger.info("Model loaded successfully")
-except Exception as e:
-    logger.error(f"Failed to load model: {e}")
-    raise
+# Load model (you can switch to 'fasterrcnn' based loading if needed)
+model = YOLO("model/best.pt")  # fine-tuned YOLOv8 model
 
-# Define fault types and severity mapping
-FAULT_CLASSES = {0: "Crack", 1: "Spalling", 2: "Corrosion"}
-SE prints
-VERITY_THRESHOLDS = {
-    "Crack": {"Minor": 0.3, "Moderate": 0.6, "Severe": 0.9},
-    "Spalling": {"Minor": 0.4, "Moderate": 0.7, "Severe": 0.95},
-    "Corrosion": {"Minor": 0.35, "Moderate": 0.65, "Severe": 0.9}
-}
-
-def preprocess_image(image):
-    """Preprocess the input image: resize and normalize."""
-    try:
-        img = np.array(image)
-        img = cv2.resize(img, (640, 640))  # YOLOv8 default input size
-        img = img / 255.0  # Normalize to [0, 1]
-        logger.info("Image preprocessed successfully")
-        return img
-    except Exception as e:
-        logger.error(f"Preprocessing failed: {e}")
-        raise
-
-def get_severity(fault_type, confidence):
-    """Determine severity based on confidence score."""
-    thresholds = SEVERITY_THRESHOLDS.get(fault_type, {"Minor": 0.3, "Moderate": 0.6, "Severe": 0.9})
-    if confidence >= thresholds["Severe"]:
-        return "Severe"
-    elif confidence >= thresholds["Moderate"]:
-        return "Moderate"
-    elif confidence >= thresholds["Minor"]:
-        return "Minor"
-    return "None"
+# Define defect labels if custom
+DEFECT_LABELS = ['crack', 'spalling', 'rust', 'deformation']
 
+# Inference function
 def detect_defects(image):
-    """Detect structural defects and return annotated image with results."""
-    try:
-        # Preprocess image
-        img_processed = preprocess_image(image)
-        
-        # Run inference
-        results = model.predict(img_processed, conf=0.5)  # Confidence threshold
-        logger.info("Inference completed")
-        
-        # Process results
-        annotated_img = results[0].plot()  # Get annotated image with bounding boxes
-        detections = []
-        
-        for result in results[0].boxes:
-            class_id = int(result.cls)
-            fault_type = FAULT_CLASSES.get(class_id, "Unknown")
-            confidence = float(result.conf)
-            severity = get_severity(fault_type, confidence)
-            detections.append(f"{fault_type} (Severity: {severity}, Confidence: {confidence:.2f})")
-        
-        # Convert annotated image to PIL for Gradio
-        annotated_img = Image.fromarray(annotated_img)
-        
-        # Return annotated image and detection results
-        return annotated_img, "\n".join(detections) if detections else "No defects detected."
-    except Exception as e:
-        logger.error(f"Detection failed: {e}")
-        return None, f"Error: {str(e)}"
-
-# Gradio interface using Blocks
-with gr.Blocks(title="Structural Defect Detection with YOLOv8") as demo:
-    gr.Markdown("""
-    # Structural Defect Detection
-    Upload a high-resolution drone-captured image to detect structural defects (e.g., cracks, spalling, corrosion) with severity levels.
-    """)
-    
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="pil", label="Upload Drone-Captured Image")
-            submit_button = gr.Button("Detect Defects")
-        with gr.Column():
-            output_image = gr.Image(type="pil", label="Annotated Image")
-            output_text = gr.Textbox(label="Detected Faults and Severity")
-    
-    submit_button.click(
-        fn=detect_defects,
-        inputs=input_image,
-        outputs=[output_image, output_text]
-    )
+    results = model(image)
+    annotated_img = results[0].plot()  # Draw boxes
+    predictions = results[0].boxes.data.cpu().numpy()  # [x1, y1, x2, y2, conf, class]
+
+    # Create readable output
+    output = []
+    for pred in predictions:
+        x1, y1, x2, y2, conf, cls_id = pred
+        label = DEFECT_LABELS[int(cls_id)]
+        output.append(f"{label}: {conf:.2f}")
+
+    return Image.fromarray(annotated_img), "\n".join(output)
+
+# Gradio UI
+interface = 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.",
+    allow_flagging="never"
+)
 
-# Launch the app
 if __name__ == "__main__":
-    try:
-        demo.launch()
-        logger.info("Gradio app launched successfully")
-    except Exception as e:
-        logger.error(f"Failed to launch Gradio app: {e}")
-        raise
+    interface.launch()