Update app.py

app.py

@@ -9,26 +9,38 @@ from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
 from tensorflow.keras.models import load_model
 from fpdf import FPDF
 import matplotlib.pyplot as plt
-# from ultralytics import YOLO  # Commented out for now
+from ultralytics import YOLO
 
 # -----------------------------
 # Flask Config
 # -----------------------------
 app = Flask(__name__)
 app.config["UPLOAD_FOLDER"] = "static/uploads"
+app.config["MAX_CONTENT_LENGTH"] = 16 * 1024 * 1024  # 16MB max file size
 os.makedirs(app.config["UPLOAD_FOLDER"], exist_ok=True)
 
 # -----------------------------
 # Load Keras classification model
 try:
     best_model = load_model("efficientnet_b0_best.keras")
-    print("✅ Model loaded successfully!")
+    print("✅ EfficientNet model loaded successfully!")
 except Exception as e:
-    print(f"❌ Error loading model: {e}")
+    print(f"❌ Error loading EfficientNet model: {e}")
     best_model = None
 
+# -----------------------------
+# Load YOLO model
+try:
+    yolo_model = YOLO("best.pt")
+    print("✅ YOLO model loaded successfully!")
+    print(f"YOLO model classes: {yolo_model.names}")
+except Exception as e:
+    print(f"❌ Error loading YOLO model: {e}")
+    yolo_model = None
+
 IMG_SIZE = 128
 
+# EfficientNet classes
 CLASS_LABELS = ['biological', 'brown-glass', 'cardboard', 'green-glass',
                 'metal', 'paper', 'plastic', 'shoes', 'trash', 'white-glass']
 
@@ -37,12 +49,83 @@ NON_RECYCLABLE = ["trash", "biological", "shoes"]
 
 stats = {}
 
+# YOLO detection confidence threshold
+YOLO_CONFIDENCE_THRESHOLD = 0.5
+
 # -----------------------------
-# Load YOLOv8 model (disabled for Hugging Face Spaces)
-yolo_model = YOLO("best.pt")  # Uncomment and add your YOLO model file if needed
+# YOLO object detection function
+def detect_objects_yolo(file_path):
+    """Detect objects in image using YOLO model"""
+    if yolo_model is None:
+        return None, "YOLO model not loaded"
+    
+    try:
+        # Read image
+        img = cv2.imread(file_path)
+        if img is None:
+            return None, "Could not read image"
+        
+        # Run YOLO detection
+        results = yolo_model(img)[0]
+        
+        detections = []
+        if results.boxes is not None and len(results.boxes) > 0:
+            boxes = results.boxes.xyxy.cpu().numpy()
+            confidences = results.boxes.conf.cpu().numpy()
+            class_ids = results.boxes.cls.cpu().numpy().astype(int)
+            
+            for i, box in enumerate(boxes):
+                if confidences[i] >= YOLO_CONFIDENCE_THRESHOLD:
+                    x1, y1, x2, y2 = map(int, box)
+                    class_name = yolo_model.names[class_ids[i]]
+                    confidence = confidences[i]
+                    
+                    detections.append({
+                        'class': class_name,
+                        'confidence': confidence,
+                        'bbox': [x1, y1, x2, y2]
+                    })
+        
+        return detections, None
+    except Exception as e:
+        return None, str(e)
 
 # -----------------------------
-# Preprocess image for classification
+# Draw bounding boxes on image
+def draw_detections(img_path, detections, output_path):
+    """Draw YOLO detections on image"""
+    try:
+        img = cv2.imread(img_path)
+        
+        for detection in detections:
+            x1, y1, x2, y2 = detection['bbox']
+            class_name = detection['class']
+            confidence = detection['confidence']
+            
+            # Choose color based on confidence
+            if confidence >= 0.80:
+                color = (0, 255, 0)   # Green - high confidence
+            elif confidence >= 0.60:
+                color = (0, 255, 255) # Yellow - medium confidence
+            else:
+                color = (0, 165, 255) # Orange - low confidence
+            
+            # Draw bounding box
+            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
+            
+            # Draw label background
+            label = f"{class_name} {confidence*100:.1f}%"
+            (text_width, text_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
+            cv2.rectangle(img, (x1, y1 - text_height - 10), (x1 + text_width, y1), color, -1)
+            
+            # Draw label text
+            cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2)
+        
+        cv2.imwrite(output_path, img)
+        return True
+    except Exception as e:
+        print(f"Error drawing detections: {e}")
+        return False
 def preprocess_image(file_path):
     img = cv2.imread(file_path)
     if img is None:
@@ -96,6 +179,8 @@ def generate_pdf_report():
 def index():
     if request.method == "POST":
         file = request.files.get("file")
+        detection_mode = request.form.get("detection_mode", "classification")
+        
         if not file or file.filename == "":
             return redirect(request.url)
 
@@ -109,31 +194,68 @@ def index():
             file_path = os.path.join(app.config["UPLOAD_FOLDER"], file.filename)
             file.save(file_path)
 
-            if best_model is None:
-                return render_template("index.html", error="Model not loaded. Please check if the model file exists.")
+            if detection_mode == "yolo" and yolo_model is not None:
+                # YOLO Object Detection Mode
+                detections, error = detect_objects_yolo(file_path)
+                
+                if error:
+                    return render_template("index.html", error=f"YOLO detection error: {error}")
+                
+                if detections:
+                    # Draw detections on image
+                    output_filename = f"detected_{file.filename}"
+                    output_path = os.path.join(app.config["UPLOAD_FOLDER"], output_filename)
+                    draw_detections(file_path, detections, output_path)
+                    
+                    # Log detections
+                    for detection in detections:
+                        log_prediction(detection['class'])
+                    
+                    return render_template(
+                        "yolo_result.html",
+                        original_image=file.filename,
+                        detected_image=output_filename,
+                        detections=detections,
+                        detection_count=len(detections)
+                    )
+                else:
+                    return render_template(
+                        "yolo_result.html",
+                        original_image=file.filename,
+                        detected_image=file.filename,
+                        detections=[],
+                        detection_count=0,
+                        message="No objects detected with sufficient confidence."
+                    )
+            
+            else:
+                # EfficientNet Classification Mode
+                if best_model is None:
+                    return render_template("index.html", error="Classification model not loaded. Please check if the model file exists.")
 
-            img_rgb, img_input = preprocess_image(file_path)
-            preds = best_model.predict(img_input)
-            class_idx = np.argmax(preds, axis=1)[0]
-            class_label = CLASS_LABELS[class_idx]
-            confidence = preds[0][class_idx]
+                img_rgb, img_input = preprocess_image(file_path)
+                preds = best_model.predict(img_input)
+                class_idx = np.argmax(preds, axis=1)[0]
+                class_label = CLASS_LABELS[class_idx]
+                confidence = preds[0][class_idx]
 
-            log_prediction(class_label)
+                log_prediction(class_label)
 
-            if class_label in RECYCLABLE:
-                bin_type = "Recyclable ♻️"
-            elif class_label in NON_RECYCLABLE:
-                bin_type = "Non-Recyclable 🗑️"
-            else:
-                bin_type = "Unknown ⚠️"
-
-            return render_template(
-                "result.html",
-                image=file.filename,
-                label=class_label,
-                confidence=f"{confidence*100:.2f}%",
-                bin_type=bin_type
-            )
+                if class_label in RECYCLABLE:
+                    bin_type = "Recyclable ♻️"
+                elif class_label in NON_RECYCLABLE:
+                    bin_type = "Non-Recyclable 🗑️"
+                else:
+                    bin_type = "Unknown ⚠️"
+
+                return render_template(
+                    "result.html",
+                    image=file.filename,
+                    label=class_label,
+                    confidence=f"{confidence*100:.2f}%",
+                    bin_type=bin_type
+                )
+                
         except Exception as e:
             return render_template("index.html", error=f"Error processing image: {str(e)}")