Update app.py

app.py

@@ -45,14 +45,16 @@ def predict():
 
         # Perform inference using the YOLO model
         results = model(image)  # Pass the image to the YOLO model
-
-        # Extract the detected classes
-        class_indices = results[0].boxes.cls.cpu().numpy().astype(int)  # Get the class indices
-        unique_classes = np.unique(class_indices)  # Get unique classes detected
-
-        # Map class indices to labels if needed
-        # Assuming you have a mapping for YOLO classes
-        detected_classes = [str(cls) for cls in unique_classes]
+        detections = results[0].boxes
+        detected_vegetables = []
+        for box in detections:
+            if box.conf > CONFIDENCE_THRESHOLD:
+                class_index = int(box.cls[0].item())  # Get class index as an integer
+                
+                # Use model's class names dynamically
+                vegetable_name = model.names[class_index]  # Get the name directly from the model
+                detected_vegetables.append(vegetable_name)
+        detected_classes=detected_vegetables[0]
 
         return jsonify({'predictions': detected_classes})
     except Exception as e:
@@ -105,7 +107,7 @@ def ocr():
         # Perform text detection using YOLO
         results = model(image)
         detections = results[0].boxes
-        high_conf_detections = [box for box in detections if box.conf > 0.85]
+        high_conf_detections = [box for box in detections if box.conf > 0.55]
 
         # Check if there are any high-confidence detections
         text_present = len(high_conf_detections) > 0