adding full MVP features for diagnosis and treatment

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import base64
+import regex as re
 from predictor import predict
 import torch
 import numpy as np
@@ -27,20 +28,93 @@ def add_bg_from_local(image_file):
 def header_white_bg(text, fontsize = 40, bold = True):
     st.markdown(
         f"""
-        <span style="background:rgba(255, 255, 255, 0.8); font-size:{fontsize}px; font-weight:{"bold" if bold else "normal"}">{text}</span>
+        <span style="background:rgba(255, 255, 255, 0.8); font-size:{fontsize}px; font-weight:{"bold" if bold else "normal"}; line-height: 2em">{text}</span>
         """,
         unsafe_allow_html=True
     )
 
+def header_red_bg(text, fontsize = 40, bold = True):
+    st.markdown(
+        f"""
+        <span style="background:rgba(184, 35, 35, 0.5); font-size:{fontsize}px; font-weight:{"bold" if bold else "normal"}; line-height: 1.25">{text}</span>
+        """,
+        unsafe_allow_html=True
+    )
+
+def header_green_bg(text, fontsize = 40, bold = True):
+    st.markdown(
+        f"""
+        <span style="background:rgba(26, 153, 58, 0.5); font-size:{fontsize}px; font-weight:{"bold" if bold else "normal"}; line-height: 1.25">{text}</span>
+        """,
+        unsafe_allow_html=True
+    )
+
+def plant_treatment_message(predicted_string):
+    if predicted_string == "Apple___Apple_scab":
+        return "Remove the infected leaves and fruit and apply a fungicide to prevent it from spreading."
+    elif predicted_string == "Apple___Black_rot":
+        return "Remove the infected branches and fruit and apply a fungicide to prevent it from spreading."
+    elif predicted_string == "Apple___Cedar_apple_rust":
+        return "Remove the infected branches and apply a fungicide to prevent it from spreading."
+    elif predicted_string == "Cherry_(including_sour)___Powdery_mildew":
+        return "Remove the infected leaves and apply a fungicide to prevent it from spreading."
+    elif predicted_string == "Corn_(maize)___Cercospora_leaf_spot Gray_leaf_spot":
+        return "Remove the infected leaves and apply a fungicide to prevent it from spreading."
+    elif predicted_string == "Corn_(maize)___Common_rust_":
+        return "Remove the infected leaves and apply a fungicide to prevent it from spreading."
+    elif predicted_string == "Corn_(maize)___Northern_Leaf_Blight":
+        return "Remove the infected leaves and apply a fungicide to prevent it from spreading."
+    elif predicted_string == "Grape___Black_rot":
+        return "Remove the infected branches and fruit and apply a fungicide to prevent it from spreading."
+    elif predicted_string == "Grape___Esca_(Black_Measles)":
+        return "Remove the infected branches and apply a fungicide to prevent it from spreading."
+    elif predicted_string == "Grape___Leaf_blight_(Isariopsis_Leaf_Spot)":
+        return "Remove the infected leaves and apply a fungicide to prevent it from spreading."
+    elif predicted_string == "Orange___Haunglongbing_(Citrus_greening)":
+        return "Remove the infected branches and apply a pesticide to prevent it from spreading."
+    elif predicted_string == "Peach___Bacterial_spot":
+        return "Remove the infected leaves and apply a copper fungicide to prevent it from spreading."
+    elif predicted_string == "Squash___Powdery_mildew":
+        return "This is a fungal disease that can cause white powdery spots on leaves and fruit. Consider removing infected plant parts and treating with a fungicide."
+    elif predicted_string == "Strawberry___Leaf_scorch":
+        return "This can be caused by drought, sunburn, or fungal diseases. Make sure Consider removing infected plant parts and treating with a fungicide."
+    elif predicted_string == "Tomato___Bacterial_spot":
+        return "This is a bacterial disease that can cause spots on leaves and fruit. Consider removing infected plant parts and treating with a copper-based fungicide."
+    elif predicted_string == "Tomato___Early_blight":
+        return "This is a fungal disease that can cause dark spots on leaves and stems. Consider removing infected plant parts and treating with a fungicide."
+    elif predicted_string == "Tomato___Late_blight":
+        return "This is a fungal disease that can cause rapid decay of foliage and fruit. Consider removing infected plant parts and treating with a fungicide."
+    elif predicted_string == "Tomato___Leaf_Mold":
+        return "This is a fungal disease that can cause brown spots on leaves. Consider removing infected plant parts and treating with a fungicide."
+    elif predicted_string == "Tomato___Septoria_leaf_spot":
+        return "This is a fungal disease that can cause brown spots with a yellow halo on leaves. Consider removing infected plant parts and treating with a fungicide."
+    elif predicted_string == "Tomato___Spider_mites Two-spotted_spider_mite":
+        return "These are tiny pests that can cause yellow spots on leaves and webbing. Consider removing infected plant parts and treating with an insecticide."
+    elif predicted_string == "Tomato___Target_Spot":
+        return "This is a fungal disease that can cause circular spots with a bullseye pattern on leaves. Consider removing infected plant parts and treating with a fungicide."
+    elif predicted_string == "Tomato___Tomato_Yellow_Leaf_Curl_Virus":
+        return "This is a viral disease that can cause yellowing and curling of leaves. Consider treating with a fungicide."
+
+def clean_prediction(prediction):
+    pattern = re.compile('(.*)___(.*)')
+    clean_predictions = []
+    for p in prediction:
+        r = pattern.search(p['predicted'])
+        plant = r.groups()[0].replace('_', ' ').lower()
+        diagnosis = r.groups()[1].replace('_', ' ').lower()
+        treatment = plant_treatment_message(p['predicted']) if diagnosis is not 'healthy' else None
+        clean_predictions.append([plant, diagnosis, "{0:.1f}%".format(float(p['probability']) * 100), treatment])
+
+    clean_predictions.sort(key=lambda x: x[2], reverse=True)
+    return clean_predictions
+
 def diagnose_health(file):
     prediction = predict(file)
-    predicted_strings = []
-    for p in prediction:
-        predicted_string = f"{p['predicted']}, Probability: {float(p['probability']):.2f}"
-        predicted_strings.append(predicted_string)
-    return predicted_strings
+    clean_predictions = clean_prediction(prediction)
+    return clean_predictions
 
 def app():
+
     add_bg_from_local('assets/background.png')
     header_white_bg(f'<span style="color:green">Plant</span><span style="color:orange">Dx</span><span style="color:green">: Diagnosis in a Snap!</span> ')
 
@@ -56,8 +130,26 @@ def app():
     if st.button("Get Diagnosis"):
         if uploaded_file is not None:
             # Diagnose plant health and display results
-            result = diagnose_health(uploaded_file)
-            st.success(f"Your plant is {result}")
+            results = diagnose_health(uploaded_file)
+
+            if results[0][1] == 'healthy':
+                header_green_bg(f"We believe this is a healthy {results[0][0]} plant with {results[0][2]} confidence. Keep up the good work with proper watering, sunlight, and nutrients.", fontsize=32, bold=False)
+            else:
+                header_red_bg(f"We believe this is an unhealthy {results[0][0]} plant with {results[0][1]}, with {results[0][2]} confidence. {results[0][3]}", fontsize=32, bold=False)
+
+            if len(results) > 1:
+                header_white_bg("Other potential diagnoses: ", fontsize=24)
+
+                for p in range(1, len(results)):
+                    if results[p][1] == 'healthy':
+                        header_white_bg(
+                            f"A healthy {results[p][0]} plant, {results[p][2]} confidence.",
+                            fontsize=20, bold=False)
+                    else:
+                        header_white_bg(
+                            f"An unhealthy {results[p][0]} plant with {results[p][1]}, {results[p][2]} confidence. {results[p][3] if results [p][3] else ''}",
+                            fontsize=20, bold=False)
+
         else:
             st.warning("Please upload an image of your plant first")
 

best_model.pth → model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ad9f39b9178b952df9bf40b334b72cee838fd7b0821d6da40f00cd1e3fe637c3
-size 21452911
+oid sha256:b141df83b46c70fcf350d4ee34e9d9943c6347b11d9cf9df1619eea5d5d97470
+size 21454657

predictor.py

@@ -7,7 +7,7 @@ def predict(image_file):
 
     #load model with params
     model = models.efficientnet_b0(weights=None)
-    model.load_state_dict(torch.load('best_model.pth', map_location=torch.device('cpu')), strict=False)
+    model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')), strict=False)
     device = torch.device('cpu')
 
     class_names = [
@@ -79,8 +79,8 @@ def predict(image_file):
             label_indx = indices[0][j]
             # print(label_indx)
             class_name = class_names[label_indx]
-            tmp_dct["predicted"] =  class_name
-            tmp_dct["probability"] =  probs[0][j]
+            tmp_dct["predicted"] = class_name
+            tmp_dct["probability"] = probs[0][j]
             tmp_lst.append(tmp_dct)
                 
                 # print(f"Prediction {j+1}: label index: {indices[i][j]}, probability: {probs[i][j]:.4f}")