Update app.py

app.py

@@ -338,6 +338,133 @@
 
 
 
+# ///////////////////////////////////Working
+
+
+# import streamlit as st
+# import numpy as np
+# import cv2
+# import warnings
+# import os
+
+# # Suppress warnings
+# warnings.filterwarnings("ignore", category=FutureWarning)
+# warnings.filterwarnings("ignore", category=UserWarning)
+
+# # Try importing TensorFlow
+# try:
+#     from tensorflow.keras.models import load_model
+#     from tensorflow.keras.preprocessing import image
+# except ImportError:
+#     st.error("Failed to import TensorFlow. Please make sure it's installed correctly.")
+
+# # Try importing PyTorch and Detectron2
+# try:
+#     import torch
+#     import detectron2
+# except ImportError:
+#     with st.spinner("Installing PyTorch and Detectron2..."):
+#         os.system("pip install torch torchvision")
+#         os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'")
+
+#     import torch
+#     import detectron2
+
+# from detectron2.engine import DefaultPredictor
+# from detectron2.config import get_cfg
+# from detectron2.utils.visualizer import Visualizer
+# from detectron2.data import MetadataCatalog
+
+# # Load the trained models
+# @st.cache_resource
+# def load_models():
+#     try:
+#         model_path_name = 'name_model_inception.h5'
+#         model_path_quality = 'type_model_inception.h5'
+#         model_name = load_model(model_path_name)
+#         model_quality = load_model(model_path_quality)
+#         return model_name, model_quality
+#     except Exception as e:
+#         st.error(f"Failed to load models: {str(e)}")
+#         return None, None
+
+# model_name, model_quality = load_models()
+
+# # Setup Detectron2 configuration for watermelon
+# @st.cache_resource
+# def load_detectron_model():
+#     cfg = get_cfg()
+#     cfg.merge_from_file("watermelon.yaml")
+#     cfg.MODEL.WEIGHTS = "Watermelon_model.pth"
+#     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
+#     cfg.MODEL.DEVICE = 'cpu'  # Use CPU for inference
+#     predictor = DefaultPredictor(cfg)
+#     return predictor, cfg
+
+# predictor, cfg = load_detectron_model()
+
+# # Streamlit app title
+# st.title("Watermelon Quality and Damage Detection")
+
+# # Upload image
+# uploaded_file = st.file_uploader("Choose a watermelon image...", type=["jpg", "jpeg", "png"])
+
+# if uploaded_file is not None:
+#     try:
+#         # Load the image
+#         img = image.load_img(uploaded_file, target_size=(224, 224))
+#         img_array = image.img_to_array(img)
+#         img_array = np.expand_dims(img_array, axis=0)
+#         img_array /= 255.0
+
+#         # Display uploaded image
+#         st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+
+#         # Predict watermelon name
+#         pred_name = model_name.predict(img_array)
+#         predicted_name = 'Watermelon'
+
+#         # Predict watermelon quality
+#         pred_quality = model_quality.predict(img_array)
+#         predicted_class_quality = np.argmax(pred_quality, axis=1)
+
+#         # Define labels for watermelon quality
+#         label_map_quality = {
+#             0: "Good",
+#             1: "Mild",
+#             2: "Rotten"
+#         }
+
+#         predicted_quality = label_map_quality[predicted_class_quality[0]]
+
+#         # Display predictions
+#         st.write(f"Fruit Type Detection: {predicted_name}")
+#         st.write(f"Fruit Quality Classification: {predicted_quality}")
+
+#         # If the quality is 'Mild' or 'Rotten', pass the image to the mask detection model
+#         if predicted_quality in ["Mild", "Rotten"]:
+#             st.write("Passing the image to the mask detection model for damage detection...")
+
+#             # Load the image again for the mask detection (Detectron2 requires the original image)
+#             im = cv2.imdecode(np.fromstring(uploaded_file.read(), np.uint8), 1)
+
+#             # Run prediction on the image
+#             outputs = predictor(im)
+
+#             # Visualize the predictions
+#             v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=0.8)
+#             out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+
+#             # Display the output
+#             st.image(out.get_image()[:, :, ::-1], caption="Detected Damage", use_column_width=True)
+
+#     except Exception as e:
+#         st.error(f"An error occurred during processing: {str(e)}")
+
+
+
+
+
 
 
 
@@ -390,24 +517,26 @@ def load_models():
 
 model_name, model_quality = load_models()
 
-# Setup Detectron2 configuration for watermelon
+# Setup Detectron2 configuration for watermelon and banana
 @st.cache_resource
-def load_detectron_model():
+def load_detectron_model(model_type):
     cfg = get_cfg()
-    cfg.merge_from_file("watermelon.yaml")
-    cfg.MODEL.WEIGHTS = "Watermelon_model.pth"
+    if model_type == "watermelon":
+        cfg.merge_from_file("watermelon.yaml")
+        cfg.MODEL.WEIGHTS = "Watermelon_model.pth"
+    elif model_type == "tomato":
+        cfg.merge_from_file("tomato.yaml")
+        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)
     return predictor, cfg
 
-predictor, cfg = load_detectron_model()
-
 # Streamlit app title
-st.title("Watermelon Quality and Damage Detection")
+st.title("Fruit Quality and Damage Detection")
 
 # Upload image
-uploaded_file = st.file_uploader("Choose a watermelon image...", type=["jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("Choose a fruit image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
     try:
@@ -420,15 +549,15 @@ if uploaded_file is not None:
         # Display uploaded image
         st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
 
-        # Predict watermelon name
+        # Predict fruit name
         pred_name = model_name.predict(img_array)
-        predicted_name = 'Watermelon'
+        predicted_name = 'Banana'  # Example: Modify based on the actual model's output
 
-        # Predict watermelon quality
+        # Predict fruit quality
         pred_quality = model_quality.predict(img_array)
         predicted_class_quality = np.argmax(pred_quality, axis=1)
 
-        # Define labels for watermelon quality
+        # Define labels for fruit quality
         label_map_quality = {
             0: "Good",
             1: "Mild",
@@ -441,13 +570,19 @@ if uploaded_file is not None:
         st.write(f"Fruit Type Detection: {predicted_name}")
         st.write(f"Fruit Quality Classification: {predicted_quality}")
 
-        # If the quality is 'Mild' or 'Rotten', pass the image to the mask detection model
+        # If the quality is 'Mild' or 'Rotten', pass the image to the corresponding mask detection model
         if predicted_quality in ["Mild", "Rotten"]:
             st.write("Passing the image to the mask detection model for damage detection...")
 
             # 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":
+                predictor, cfg = load_detectron_model("tomato")
+            elif predicted_name == "Watermelon":
+                predictor, cfg = load_detectron_model("watermelon")
+
             # Run prediction on the image
             outputs = predictor(im)
 
@@ -460,3 +595,4 @@ if uploaded_file is not None:
 
     except Exception as e:
         st.error(f"An error occurred during processing: {str(e)}")
+