Update app.py

app.py

@@ -517,7 +517,7 @@ def load_models():
 
 model_name, model_quality = load_models()
 
-# Setup Detectron2 configuration for watermelon and banana
+# Setup Detectron2 configuration for watermelon and tomato
 @st.cache_resource
 def load_detectron_model(model_type):
     cfg = get_cfg()
@@ -526,7 +526,7 @@ def load_detectron_model(model_type):
         cfg.MODEL.WEIGHTS = "Watermelon_model.pth"
     elif model_type == "tomato":
         cfg.merge_from_file("tomato.yaml")
-        cfg.MODEL.WEIGHTS = "tomato_model.pth"
+        cfg.MODEL.WEIGHTS = "Tomato_model.pth"
     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
     cfg.MODEL.DEVICE = 'cpu'  # Use CPU for inference
     predictor = DefaultPredictor(cfg)
@@ -551,7 +551,7 @@ if uploaded_file is not None:
 
         # Predict fruit name
         pred_name = model_name.predict(img_array)
-        predicted_name = 'Banana'  # Example: Modify based on the actual model's output
+        predicted_name = 'Tomato'  # Example: Modify based on the actual model's output
 
         # Predict fruit quality
         pred_quality = model_quality.predict(img_array)
@@ -577,8 +577,8 @@ if uploaded_file is not None:
             # Load the image again for the mask detection (Detectron2 requires the original image)
             im = cv2.imdecode(np.fromstring(uploaded_file.read(), np.uint8), 1)
 
-            # Check if the predicted fruit is Banana or Watermelon and load the correct model
-            if predicted_name == "tomato":
+            # Check if the predicted fruit is Tomato or Watermelon and load the correct model
+            if predicted_name == "Tomato":
                 predictor, cfg = load_detectron_model("tomato")
             elif predicted_name == "Watermelon":
                 predictor, cfg = load_detectron_model("watermelon")
@@ -596,3 +596,4 @@ if uploaded_file is not None:
     except Exception as e:
         st.error(f"An error occurred during processing: {str(e)}")
 
+