Update app.py

app.py

@@ -480,6 +480,108 @@ from detectron2.data import MetadataCatalog
 
 
 
+
+
+
+# import streamlit as st
+# import numpy as np
+# import cv2
+# import torch
+# from PIL import Image
+# from tensorflow.keras.models import load_model
+# from tensorflow.keras.preprocessing import image
+# from detectron2.engine import DefaultPredictor
+# from detectron2.config import get_cfg
+# from detectron2.utils.visualizer import Visualizer
+# from detectron2.data import MetadataCatalog
+
+# # Suppress warnings
+# import warnings
+# import tensorflow as tf
+# warnings.filterwarnings("ignore")
+# tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+
+# @st.cache_resource
+# def load_models():
+#     model_name = load_model('name_model_inception.h5')
+#     model_quality = load_model('type_model_inception.h5')
+#     return model_name, model_quality
+
+# model_name, model_quality = load_models()
+
+# # Detectron2 setup
+# @st.cache_resource
+# def load_detectron_model(fruit_name):
+#     cfg = get_cfg()
+#     cfg.merge_from_file(f"{fruit_name.lower()}.yaml")
+#     cfg.MODEL.WEIGHTS = f"{fruit_name.lower()}_model.pth"
+#     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
+#     cfg.MODEL.DEVICE = 'cpu'
+#     predictor = DefaultPredictor(cfg)
+#     return predictor, cfg
+
+# # Labels
+# label_map_name = {
+#     0: "Banana", 1: "Cucumber", 2: "Grape", 3: "Kaki", 4: "Papaya",
+#     5: "Peach", 6: "Pear", 7: "Pepper", 8: "Strawberry", 9: "Watermelon",
+#     10: "Tomato"
+# }
+# label_map_quality = {0: "Good", 1: "Mild", 2: "Rotten"}
+
+# def predict_fruit(img):
+#     # Preprocess image
+#     img = Image.fromarray(img.astype('uint8'), 'RGB')
+#     img = img.resize((224, 224))
+#     x = image.img_to_array(img)
+#     x = np.expand_dims(x, axis=0)
+#     x = x / 255.0
+
+#     # Predict
+#     pred_name = model_name.predict(x)
+#     pred_quality = model_quality.predict(x)
+
+#     predicted_name = label_map_name[np.argmax(pred_name, axis=1)[0]]
+#     predicted_quality = label_map_quality[np.argmax(pred_quality, axis=1)[0]]
+
+#     return predicted_name, predicted_quality, img
+
+# def main():
+#     st.title("Fruit Quality and Damage Detection")
+#     st.write("Upload an image of a fruit to detect its type, quality, and potential damage.")
+
+#     uploaded_file = st.file_uploader("Choose a fruit image...", type=["jpg", "jpeg", "png"])
+
+#     if uploaded_file is not None:
+#         image = Image.open(uploaded_file)
+#         st.image(image, caption="Uploaded Image", use_column_width=True)
+
+#         if st.button("Analyze"):
+#             predicted_name, predicted_quality, img = predict_fruit(np.array(image))
+
+#             st.write(f"Fruit Type: {predicted_name}")
+#             st.write(f"Fruit Quality: {predicted_quality}")
+
+#             if predicted_name.lower() in ["kaki", "tomato", "strawberry", "pepper", "pear", "peach", "papaya", "watermelon", "grape", "banana", "cucumber"] and predicted_quality in ["Mild", "Rotten"]:
+#                 st.write("Detecting damage...")
+#                 predictor, cfg = load_detectron_model(predicted_name)
+#                 outputs = predictor(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
+#                 v = Visualizer(np.array(img), MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=0.8)
+#                 out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+#                 st.image(out.get_image(), caption="Damage Detection Result", use_column_width=True)
+#             else:
+#                 st.write("No damage detection performed for this fruit or quality level.")
+
+# if __name__ == "__main__":
+#     main()
+
+
+
+
+
+
+
+
+
 
 
 
@@ -487,6 +589,7 @@ import streamlit as st
 import numpy as np
 import cv2
 import torch
+import os
 from PIL import Image
 from tensorflow.keras.models import load_model
 from tensorflow.keras.preprocessing import image
@@ -513,8 +616,10 @@ model_name, model_quality = load_models()
 @st.cache_resource
 def load_detectron_model(fruit_name):
     cfg = get_cfg()
-    cfg.merge_from_file(f"{fruit_name.lower()}.yaml")
-    cfg.MODEL.WEIGHTS = f"{fruit_name.lower()}_model.pth"
+    config_path = os.path.join('utils', f"{fruit_name.lower()}_config.yaml")
+    cfg.merge_from_file(config_path)
+    model_path = os.path.join('models', f"{fruit_name.lower()}_model.pth")
+    cfg.MODEL.WEIGHTS = model_path
     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
     cfg.MODEL.DEVICE = 'cpu'
     predictor = DefaultPredictor(cfg)
@@ -563,15 +668,16 @@ def main():
 
             if predicted_name.lower() in ["kaki", "tomato", "strawberry", "pepper", "pear", "peach", "papaya", "watermelon", "grape", "banana", "cucumber"] and predicted_quality in ["Mild", "Rotten"]:
                 st.write("Detecting damage...")
-                predictor, cfg = load_detectron_model(predicted_name)
-                outputs = predictor(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
-                v = Visualizer(np.array(img), MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=0.8)
-                out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
-                st.image(out.get_image(), caption="Damage Detection Result", use_column_width=True)
+                try:
+                    predictor, cfg = load_detectron_model(predicted_name)
+                    outputs = predictor(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
+                    v = Visualizer(np.array(img), MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=0.8)
+                    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+                    st.image(out.get_image(), caption="Damage Detection Result", use_column_width=True)
+                except Exception as e:
+                    st.error(f"Error in damage detection: {str(e)}")
             else:
                 st.write("No damage detection performed for this fruit or quality level.")
 
 if __name__ == "__main__":
-    main()
-
-
+    main()