Update app.py

app.py

@@ -2,36 +2,32 @@ import gradio as gr
 import cv2
 from ultralytics import YOLO
 
-# 1. Load the trained model
+# 1. Load your new masterpiece model
 model = YOLO("best.pt")
 
-# 2. CORRECT WAY TO OVERRIDE NAMES (Accessing the underlying model dictionary)
-# This bypasses the read-only property and fixes the cross-wiring safely.
-model.model.names = {0: "Smoke", 1: "Fire"}
-
 def predict_image(img):
     if img is None:
         return None, "No image uploaded."
     
-    # Convert Gradio's RGB image to BGR for YOLO processing
+    # Convert Gradio's RGB format to BGR for YOLO
     bgr_img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
     
-    # Run prediction (Confidence threshold at a balanced 0.25)
-    results = model.predict(source=bgr_img, conf=0.12)
+    # Run prediction (Balanced threshold at 0.25 confidence)
+    results = model.predict(source=bgr_img, conf=0.25)
     
-    # Get the correctly labeled visual image
+    # Get the visually annotated BGR image and map back to RGB for Gradio
     annotated_img_bgr = results[0].plot()
     annotated_img_rgb = cv2.cvtColor(annotated_img_bgr, cv2.COLOR_BGR2RGB)
     
-    # Extract the corrected classes safely
+    # Extract detected classes safely using native, pre-mapped indices
     detected_classes = []
     if results[0].boxes is not None:
         for box in results[0].boxes:
             cls_id = int(box.cls[0])
-            class_name = model.model.names[cls_id] 
+            class_name = model.names[cls_id] # Natively tracks 'Smoke' or 'Fire' perfectly
             detected_classes.append(class_name)
     
-    # Generate the accurate warning message
+    # Generate the warning message
     if len(detected_classes) == 0:
         status_warning = "✅ SYSTEM STATUS: Safe (No Fire or Smoke detected)"
     else:
@@ -41,10 +37,10 @@ def predict_image(img):
 
     return annotated_img_rgb, status_warning
 
-# Build the Gradio UI Layout
-with gr.Blocks(title="🔥 AI Fire & Smoke Detection System") as demo:
-    gr.Markdown("# 🔥 AI Fire & Smoke Detection System")
-    gr.Markdown("Upload an image or a frame from a security feed to test the custom-trained YOLOv8 model.")
+# Build the Masterpiece Gradio UI Layout
+with gr.Blocks(title="🔥 AI Fire & Smoke Detection System v2.0") as demo:
+    gr.Markdown("# 🔥 AI Fire & Smoke Detection System v2.0 (Masterpiece Edition)")
+    gr.Markdown("An advanced custom-trained YOLOv8 system optimized against glare, ambient lighting, and complex vapor patterns.")
     
     with gr.Row():
         with gr.Column():