Update app.py

app.py

@@ -3,24 +3,37 @@ from PIL import Image
 import base64
 import io
 import json
-import os
 import argparse
+import os
 
 # Load model
-model = pipeline("image-classification", model="linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification", top_k=1)
+model = pipeline(
+    "image-classification",
+    model="linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification",
+    top_k=1
+)
 
 # Load mapping data
 def load_map(file_path):
-    with open(file_path, 'r', encoding='utf-8') as f:
-        return dict(line.strip().split(":", 1) for line in f if ":" in line)
+    try:
+        with open(file_path, 'r', encoding='utf-8') as f:
+            return dict(line.strip().split(":", 1) for line in f if ":" in line)
+    except Exception as e:
+        print(f"⚠️ Failed to load {file_path}: {e}")
+        return {}
 
 disease_map = load_map("diseases.txt")
 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 if line.strip())
+# Load critical diseases
+try:
+    with open("critical_diseases.txt", "r", encoding='utf-8') as f:
+        critical_diseases = set(line.strip() for line in f if line.strip())
+except:
+    critical_diseases = set()
 
+# Main prediction function
 def predict_disease(base64_img):
     try:
         img_bytes = base64.b64decode(base64_img)
@@ -50,15 +63,15 @@ def predict_disease(base64_img):
     except Exception as e:
         return {"error": f"Prediction failed: {str(e)}"}
 
-# CLI Entry Point
+# CLI Interface
 if __name__ == "__main__":
-    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")
+    parser = argparse.ArgumentParser(description="🌿 AI-Powered Tree Disease Detector")
+    parser.add_argument("--image", type=str, help="Path to image or base64 file", required=True)
+    parser.add_argument("--raw", action="store_true", help="Indicates that the file is a raw image (e.g., JPG, PNG)")
     args = parser.parse_args()
 
-    if not args.image:
-        print("❗ Please provide an image file using --image")
+    if not os.path.isfile(args.image):
+        print(f"❌ File not found: {args.image}")
         exit(1)
 
     try:
@@ -66,10 +79,10 @@ if __name__ == "__main__":
             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:
+            with open(args.image, "r", encoding="utf-8") 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}")
+    except Exception as e:
+        print(json.dumps({"error": f"Unexpected error: {str(e)}"}, indent=2))