Update app.py

app.py

@@ -9,9 +9,7 @@ from ultralytics import YOLO
 import time
 from transformers import DonutProcessor, VisionEncoderDecoderModel
 from PIL import Image
-import os
-# Set cache directory to a local path with write permissions
-os.environ['TRANSFORMERS_CACHE'] = '/.cache'
+
 
 # Load the Donut model and processor (adjust the model name if needed)
 processor = DonutProcessor.from_pretrained('naver-clova-ix/donut-base')
@@ -30,7 +28,7 @@ vegetables = [
     "corn", "cucumber", "drumstick", "garlic", "ginger", "ladies finger",
     "lemons", "Onion red", "potato", "sweet potato", "tomato", "Zuchini"
 ]
-model = keras.models.load_model('fresh_model.keras')
+model = YOLO('fresh_model.pt')
 # Load the YOLO model
 yolo_model = YOLO('another_model.pt')  # Adjust the path as needed
 
@@ -51,25 +49,19 @@ def predict():
         # Convert the NumPy array into an OpenCV image (BGR format)
         image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
 
-        # Convert to RGB format
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
-        # Resize the image to the expected input size of the model (e.g., 225x225)
-        image_resized = cv2.resize(image, (225, 225))
-
-        # Normalize the image for vegetable classification
-        image_resized = image_resized.astype('float32') / 255.0  # Normalize to [0, 1]
-        image_resized = np.expand_dims(image_resized, axis=0)  # Add batch dimension
+        # Perform inference using the YOLO model
+        results = model(image)  # Pass the image to the YOLO model
 
-        # Make a prediction using the vegetable classification model
-        predictions = model.predict(image_resized)
-        predicted_class = np.argmax(predictions, axis=1)[0]
-        time.sleep(3)
+        # 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 the class index to the vegetable label
-        prediction_label = vegetables[predicted_class]
+        # Map class indices to labels if needed
+        # Assuming you have a mapping for YOLO classes
+        detected_classes = [str(cls) for cls in unique_classes]
+        time.sleep(3)# Convert indices to strings
 
-        return jsonify({'prediction': prediction_label})
+        return jsonify({'predictions': detected_classes})
     except Exception as e:
         return jsonify({'error': str(e)}), 500