Update app.py

app.py

@@ -16,58 +16,58 @@
 
 
 
-import streamlit as st
-from tensorflow.keras.models import load_model
-from tensorflow.keras.preprocessing import image
-import numpy as np
-from PIL import Image
+# import streamlit as st
+# from tensorflow.keras.models import load_model
+# from tensorflow.keras.preprocessing import image
+# import numpy as np
+# from PIL import Image
 
-# Load the pre-trained models
-@st.cache_resource
-def load_models():
-    model1 = load_model('name_model_inception.h5')  # Update with your Hugging Face model path
-    model2 = load_model('type_model_inception.h5')  # Update with your Hugging Face model path
-    return model1, model2
+# # Load the pre-trained models
+# @st.cache_resource
+# def load_models():
+#     model1 = load_model('name_model_inception.h5')  # Update with your Hugging Face model path
+#     model2 = load_model('type_model_inception.h5')  # Update with your Hugging Face model path
+#     return model1, model2
 
-model1, model2 = load_models()
+# model1, model2 = load_models()
 
-# Label mappings
-label_map1 = {
-    0: "Banana", 1: "Cucumber", 2: "Grape", 3: "Kaki", 4: "Papaya",
-    5: "Peach", 6: "Pear", 7: "Pepper", 8: "Strawberry", 9: "Watermelon", 10: "Tomato"
-}
+# # Label mappings
+# label_map1 = {
+#     0: "Banana", 1: "Cucumber", 2: "Grape", 3: "Kaki", 4: "Papaya",
+#     5: "Peach", 6: "Pear", 7: "Pepper", 8: "Strawberry", 9: "Watermelon", 10: "Tomato"
+# }
 
-label_map2 = {
-    0: "Good", 1: "Mild", 2: "Rotten"
-}
+# label_map2 = {
+#     0: "Good", 1: "Mild", 2: "Rotten"
+# }
 
-# Streamlit app layout
-st.title("Fruit Classifier")
+# # Streamlit app layout
+# st.title("Fruit Classifier")
 
-# Upload image
-uploaded_file = st.file_uploader("Choose an image of a fruit", type=["jpg", "jpeg", "png"])
+# # Upload image
+# uploaded_file = st.file_uploader("Choose an image of a fruit", type=["jpg", "jpeg", "png"])
 
-if uploaded_file is not None:
-    # Display the uploaded image
-    img = Image.open(uploaded_file)
-    st.image(img, caption="Uploaded Image", use_column_width=True)
+# if uploaded_file is not None:
+#     # Display the uploaded image
+#     img = Image.open(uploaded_file)
+#     st.image(img, caption="Uploaded Image", use_column_width=True)
     
-    # Preprocess the image
-    img = img.resize((224, 224))  # Resize image to match the model input
-    img_array = image.img_to_array(img)
-    img_array = np.expand_dims(img_array, axis=0)
-    img_array = img_array / 255.0  # Normalize the image
+#     # Preprocess the image
+#     img = img.resize((224, 224))  # Resize image to match the model input
+#     img_array = image.img_to_array(img)
+#     img_array = np.expand_dims(img_array, axis=0)
+#     img_array = img_array / 255.0  # Normalize the image
 
-    # Make predictions
-    pred1 = model1.predict(img_array)
-    pred2 = model2.predict(img_array)
+#     # Make predictions
+#     pred1 = model1.predict(img_array)
+#     pred2 = model2.predict(img_array)
 
-    predicted_class1 = np.argmax(pred1, axis=1)
-    predicted_class2 = np.argmax(pred2, axis=1)
+#     predicted_class1 = np.argmax(pred1, axis=1)
+#     predicted_class2 = np.argmax(pred2, axis=1)
 
-    # Display results
-    st.write(f"**Type Detection**: {label_map1[predicted_class1[0]]}")
-    st.write(f"**Condition Detection**: {label_map2[predicted_class2[0]]}")
+#     # Display results
+#     st.write(f"**Type Detection**: {label_map1[predicted_class1[0]]}")
+#     st.write(f"**Condition Detection**: {label_map2[predicted_class2[0]]}")
 
 
 
@@ -79,6 +79,18 @@ if uploaded_file is not None:
 
 
 
+
+
+
+
+
+
+
+
+
+
+
+
 # import streamlit as st
 # import numpy as np
 # import cv2
@@ -184,3 +196,138 @@ if uploaded_file is not None:
 
 #     except Exception as e:
 #         st.error(f"An error occurred during processing: {str(e)}")
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+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()
+
+# 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)
+
+            # 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
+
+            predictor = load_detectron_model()
+
+            # 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)}")