Update app.py

app.py

@@ -55,20 +55,35 @@ POTATO_MODEL_PATH = "pot_ato_model.tflite"
 POTATO_CLASS_NAMES = ['Potato___Early_blight', 'Potato___Late_blight', 'Potato___healthy']
 
 
-# --- TOMATO CONFIGURATION ---
-TOMATO_MODEL_PATH = "TOMAto_model.tflite" 
+# --- LEMON CONFIGURATION ---
+LEMON_MODEL_PATH = "lemon_disease_model.tflite"
+
+# Replace these with the exact 9 lines from your lemon_labels.txt file
+LEMON_CLASS_NAMES = [
+    "Anthracnose",
+    "Bacterial Blight",
+    "Citrus Canker",
+    "Curl Virus",
+    "Deficiency Leaf",
+    "Dry Leaf",
+    "Healthy Leaf",
+    "Sooty Mould",
+    "Spider Mites"
+]
+
+# --- Tomato CONFIGURATION ---
+TOMATO_MODEL_PATH = "lemon_disease_model.tflite" 
 
 TOMATO_CLASS_NAMES = [
-    'Tomato___Bacterial_spot', 
-    'Tomato___Early_blight', 
-    'Tomato___Late_blight', 
-    'Tomato___Leaf_Mold', 
-    'Tomato___Septoria_leaf_spot', 
-    'Tomato___Spider_mites Two-spotted_spider_mite', 
-    'Tomato___Target_Spot', 
-    'Tomato___Tomato_Yellow_Leaf_Curl_Virus', 
-    'Tomato___Tomato_mosaic_virus', 
-    'Tomato___healthy'
+    'Anthracnose', 
+    'Bacterial Blight', 
+    'Citrus Canker', 
+    'Curl Virus', 
+    'Deficiency Leaf', 
+    'Dry Leaf', 
+    'Healthy Leaf', 
+    'Sooty Mould', 
+    'Spider Mites'
 ]
 
 
@@ -107,6 +122,12 @@ tomato_interpreter.allocate_tensors()
 tomato_input_details = tomato_interpreter.get_input_details()
 tomato_output_details = tomato_interpreter.get_output_details()
 
+# --- LOAD LEMON MODEL ---
+lemon_interpreter = tf.lite.Interpreter(model_path=LEMON_MODEL_PATH)
+lemon_interpreter.allocate_tensors()
+lemon_input_details = lemon_interpreter.get_input_details()
+lemon_output_details = lemon_interpreter.get_output_details()
+
 # --- LOAD RICE TFLITE MODEL ---
 rice_interpreter = tf.lite.Interpreter(model_path=RICE_MODEL_PATH)
 rice_interpreter.allocate_tensors()
@@ -254,6 +275,46 @@ async def predict_potato(file: UploadFile = File(...)):
         "confidence": confidence
     }   
 
+@app.post("/predict_lemon")
+async def predict_lemon(file: UploadFile = File(...)): 
+    try:
+        # --- PREPROCESS IMAGE ---
+        image_data = await file.read()
+        image = Image.open(io.BytesIO(image_data)).convert("RGB")
+        
+        if not is_leaf(image):
+            return JSONResponse(status_code=400, content={
+            "error": "Not a leaf",
+            "message": "Please upload a clear image of a PLANT leaf."
+        })
+
+        image = image.resize((224, 224))
+        
+        # EfficientNetB0 Normalization
+        input_data = np.array(image).astype(np.float32)
+        input_data = np.expand_dims(input_data, axis=0)
+
+        # --- RUN INFERENCE ---
+        lemon_interpreter.set_tensor(lemon_input_details[0]['index'], input_data)
+        lemon_interpreter.invoke()
+        output_data = lemon_interpreter.get_tensor(lemon_output_details[0]['index'])
+
+        # --- PROCESS RESULTS ---
+        prediction = np.argmax(output_data[0])
+        confidence = float(np.max(output_data[0]))
+        predicted_class = LEMON_CLASS_NAMES[prediction]
+        
+        display_name = predicted_class
+
+    
+        return {
+                "class": display_name,
+                "confidence": confidence
+            }
+            
+    except Exception as e:
+        return {"error": str(e), "message": "The code crashed before finishing the prediction."}
+
 # --- Tomato's Code ---
 @app.post("/predict_tomato")
 async def predict_tomato(file: UploadFile = File(...)):