Update app.py

app.py

@@ -3,6 +3,8 @@ from PIL import Image
 import base64
 import io
 import json
+import os
+import argparse
 
 # Load model
 model = pipeline("image-classification", model="linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification", top_k=1)
@@ -17,17 +19,15 @@ treatment_map = load_map("treatments.txt")
 fertilizer_map = load_map("fertilizers.txt")
 
 with open("critical_diseases.txt", "r") as f:
-    critical_diseases = set(line.strip() for line in f)
+    critical_diseases = set(line.strip() for line in f if line.strip())
 
 def predict_disease(base64_img):
-    # Decode base64
     try:
         img_bytes = base64.b64decode(base64_img)
         image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
     except Exception as e:
         return {"error": f"Invalid image: {str(e)}"}
 
-    # Predict
     try:
         result = model(image)[0]
         label = result["label"]
@@ -43,7 +43,6 @@ def predict_disease(base64_img):
             "fertilizer_recommendation": fertilizer
         }
 
-        # Alerts (for integration)
         if confidence < 0.6 or disease in critical_diseases:
             output["alert"] = True
 
@@ -51,9 +50,26 @@ def predict_disease(base64_img):
     except Exception as e:
         return {"error": f"Prediction failed: {str(e)}"}
 
-# Test hook
+# CLI Entry Point
 if __name__ == "__main__":
-    with open("example_image.b64", "r") as f:
-        base64_img = f.read()
-    response = predict_disease(base64_img)
-    print(json.dumps(response, indent=2))
+    parser = argparse.ArgumentParser(description="AI Tree Disease Predictor")
+    parser.add_argument("--image", type=str, help="Path to base64 file or image to encode")
+    parser.add_argument("--raw", action="store_true", help="Indicates the --image is a raw image, not base64 text")
+    args = parser.parse_args()
+
+    if not args.image:
+        print("❗ Please provide an image file using --image")
+        exit(1)
+
+    try:
+        if args.raw:
+            with open(args.image, "rb") as img_file:
+                base64_img = base64.b64encode(img_file.read()).decode('utf-8')
+        else:
+            with open(args.image, "r") as f:
+                base64_img = f.read().strip()
+
+        result = predict_disease(base64_img)
+        print(json.dumps(result, indent=2))
+    except FileNotFoundError:
+        print(f"❌ File not found: {args.image}")