Update app.py

app.py

@@ -11,48 +11,53 @@ import numpy as np
 random.seed(42)
 np.random.seed(42)
 
-# Load default YOLO model
+# Configuration
+MODEL_PATH = os.getenv("MODEL_PATH", "last.pt")
+ZERO_DEFECTS_PATH = "zero_defects.png"
+VALID_EXTENSIONS = [".jpg", ".jpeg", ".png"]
+
+# Load YOLO model
 try:
-    yolo_model = YOLO("last.pt")
+    if not os.path.exists(MODEL_PATH):
+        raise FileNotFoundError(f"Model file {MODEL_PATH} not found.")
+    yolo_model = YOLO(MODEL_PATH)
     print("YOLO model loaded successfully.")
-except FileNotFoundError:
-    print("Error: 'last.pt' not found.")
-    yolo_model = None
 except Exception as e:
-    print(f"An error occurred while loading the YOLO model: {e}")
+    print(f"Error loading YOLO model: {e}")
     yolo_model = None
 
 # Load SAHI model
 try:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     sahi_model = AutoDetectionModel.from_pretrained(
         model_type="ultralytics",
-        model_path="last.pt",  # same model used for consistency
+        model_path=MODEL_PATH,
         confidence_threshold=0.5,
-        device="cpu",
+        device=device,
     )
     print("SAHI model loaded successfully.")
 except Exception as e:
-    print(f"An error occurred while loading the SAHI model: {e}")
+    print(f"Error loading SAHI model: {e}")
     sahi_model = None
 
-# Prediction function
-def predict_and_show_bounding_boxes(image_path, model_choice):
+def predict_and_show_bounding_boxes(image_path, model_choice, conf_threshold=0.5):
+    # Validate image path
+    if not image_path or not any(image_path.lower().endswith(ext) for ext in VALID_EXTENSIONS):
+        return None, "Error: Invalid or unsupported image format."
+
     if model_choice == "YOLO":
         if yolo_model is None:
             return None, "Error: YOLO model not loaded."
-
         try:
             img = cv2.imread(image_path)
             if img is None:
-                return None, "Error: Could not load image"
-
-            results = yolo_model(image_path, conf=0.5)[0]
+                return None, "Error: Could not load image."
+            results = yolo_model(image_path, conf=conf_threshold)[0]
             boxes = results.boxes
-
             if len(boxes) == 0:
-                zero_defects_img = cv2.imread('zero_defects.png')
-                return zero_defects_img if zero_defects_img is not None else (None, "Error: Could not load zero defects image")
-
+                if os.path.exists(ZERO_DEFECTS_PATH):
+                    return cv2.imread(ZERO_DEFECTS_PATH), "No defects detected."
+                return img, "No defects detected (zero_defects.png not found)."
             for box in boxes:
                 xyxy = box.xyxy[0].tolist()
                 x_min, y_min, x_max, y_max = map(int, xyxy[:4])
@@ -61,16 +66,16 @@ def predict_and_show_bounding_boxes(image_path, model_choice):
                 cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
                 label = f"{results.names[cls]}: {conf:.2f}"
                 cv2.putText(img, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
-
-            return img
+            return cv2.cvtColor(img, cv2.COLOR_BGR2RGB), "Detection complete."
         except Exception as e:
-            return None, str(e)
-
+            return None, f"Error during YOLO prediction: {e}"
     elif model_choice == "SAHI":
         if sahi_model is None:
             return None, "Error: SAHI model not loaded."
-
         try:
+            img = cv2.imread(image_path)
+            if img is None:
+                return None, "Error: Could not load image."
             result = get_sliced_prediction(
                 image_path,
                 sahi_model,
@@ -79,29 +84,20 @@ def predict_and_show_bounding_boxes(image_path, model_choice):
                 overlap_height_ratio=0.2,
                 overlap_width_ratio=0.2,
             )
-
-            img = cv2.imread(image_path)
-            if img is None:
-                return None, "Error: Could not load image"
-
+            if len(result.object_prediction_list) == 0:
+                if os.path.exists(ZERO_DEFECTS_PATH):
+                    return cv2.imread(ZERO_DEFECTS_PATH), "No defects detected."
+                return img, "No defects detected (zero_defects.png not found)."
             for pred in result.object_prediction_list:
                 box = pred.bbox.to_xyxy()
                 x_min, y_min, x_max, y_max = map(int, box)
                 label = f"{pred.category.name}: {pred.score.value:.2f}"
                 cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (255, 0, 0), 2)
                 cv2.putText(img, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 2)
-
-            if len(result.object_prediction_list) == 0:
-                zero_defects_img = cv2.imread('zero_defects.png')
-                return zero_defects_img if zero_defects_img is not None else (None, "Error: Could not load zero defects image")
-
-            return img
+            return cv2.cvtColor(img, cv2.COLOR_BGR2RGB), "Detection complete."
         except Exception as e:
-            return None, str(e)
-
-    else:
-        return None, "Invalid model choice."
-
+            return None, f"Error during SAHI prediction: {e}"
+    return None, "Invalid model choice."
 
 # Gradio interface
 iface = gr.Interface(
@@ -109,10 +105,13 @@ iface = gr.Interface(
     inputs=[
         gr.Image(type="filepath", label="Upload Image"),
         gr.Radio(choices=["YOLO", "SAHI"], label="Choose Detection Mode", value="YOLO"),
+        gr.Slider(minimum=0.1, maximum=0.9, value=0.5, label="Confidence Threshold"),
     ],
     outputs=[gr.Image(label="Result"), gr.Textbox(label="Message")],
-    title="Defect Detection",
-    description="Upload an image and choose YOLO or SAHI for defect detection.",
+    title="PCB Defect Detection",
+    description="Upload a PCB image and choose YOLO (green boxes) or SAHI (red boxes) for defect detection. Adjust confidence threshold for sensitivity.",
 )
 
-iface.launch(share=True)
+if __name__ == "__main__":
+    share = os.getenv("HF_SHARE", "False").lower() == "true"
+    iface.launch(share=share)