Update app.py

app.py

@@ -483,8 +483,7 @@ from detectron2.data import MetadataCatalog
 
 
 
-
-import gradio as gr
+import streamlit as st
 import numpy as np
 import cv2
 import torch
@@ -502,15 +501,20 @@ import tensorflow as tf
 warnings.filterwarnings("ignore")
 tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
 
-# Load models
-model_name = load_model('name_model_inception.h5')
-model_quality = load_model('type_model_inception.h5')
+@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}_model.pth"
+    cfg.merge_from_file(f"{fruit_name.lower()}_config.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)
@@ -539,31 +543,35 @@ def predict_fruit(img):
     predicted_name = label_map_name[np.argmax(pred_name, axis=1)[0]]
     predicted_quality = label_map_quality[np.argmax(pred_quality, axis=1)[0]]
 
-    result = f"Fruit Type: {predicted_name}\nFruit Quality: {predicted_quality}"
+    return predicted_name, predicted_quality, img
 
-    # Damage detection for specific fruits
-    if predicted_name.lower() in ["kaki", "tomato", "strawberry", "pepper", "pear", "peach", "papaya", "watermelon", "grape", "banana", "cucumber"] and predicted_quality in ["Mild", "Rotten"]:
-        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"))
-        result_image = out.get_image()
-    else:
-        result_image = np.array(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.")
 
-    return result, result_image
+    uploaded_file = st.file_uploader("Choose a fruit image...", type=["jpg", "jpeg", "png"])
 
-iface = gr.Interface(
-    fn=predict_fruit,
-    inputs=gr.Image(),
-    outputs=[gr.Textbox(), gr.Image()],
-    title="Fruit Quality and Damage Detection",
-    description="Upload an image of a fruit to detect its type, quality, and potential damage."
-)
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
 
-iface.launch()
+        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()