Update app.py

app.py

@@ -18,12 +18,6 @@ def load_models():
     model_name = 'Model/mango_model.h5'
     model = tf.keras.models.load_model(model_name) 
     return model
-
-def load_labels():
-    with open('dataset_labels.txt', 'r') as file:
-        data = file.read().splitlines()
-    mango_dict = dict(enumerate(data, 1))
-    return mango_dict
     
 def load_image():
      uploaded_file = st.file_uploader(label='Pick an image to test')
@@ -36,7 +30,7 @@ def load_image():
      else:
          return None
     
-def predict(model, categories, img):
+def predict(model, img):
     img_array = tf.keras.preprocessing.image.img_to_array(img)
     prediction = [img_array]
     prediction_test = [1]
@@ -44,31 +38,20 @@ def predict(model, categories, img):
     test_ds = test_ds.cache().batch(32).prefetch(buffer_size = tf.data.experimental.AUTOTUNE)
 
     prediction = model.predict(test_ds)
-    prediction_dict = dict(enumerate(prediction.flatten(), 1))
-    k = Counter(prediction_dict)
- 
-    # Finding 3 highest values
-    high = k.most_common(3)
 
-    percentages = []
-    flowers = []
-    for i in high:
-        key, value = i
-        flowers.append(categories[key])
-        percentages.append(np.round(value*100, 2))
-    return flowers, percentages
+    if prediction[0]>0.5:
+        return 'unripe'
+    else:
+        return 'ripe'
 
 def main():
     st.title('Mango Ripeness Classifier 🥭')
     model = load_models()
-    categories = load_labels()
     image = load_image()
     result = st.button('Run on image')
     if result:
         st.write('Calculating results...')
-        flowers, percentages = predict(model, categories, image)
-        st.text(flowers)
-        st.text(percentages)
+        st.write(predict(model, image))
 
 if __name__ == '__main__':
     main()