Update app.py

app.py

@@ -6,20 +6,27 @@ import json
 import argparse
 import os
 import warnings
+import sys
 
-# Suppress deprecated warning for feature extractor
-warnings.filterwarnings("ignore", category=FutureWarning)
+# Suppress warnings
+warnings.filterwarnings("ignore")
 
-# Load model pipeline
+# === Load model ===
 print("🔄 Loading model...")
-model = pipeline(
-    "image-classification",
-    model="linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification",
-    top_k=1
-)
+
+try:
+    model = pipeline(
+        "image-classification",
+        model="linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification",
+        top_k=1
+    )
+except Exception as e:
+    print(f"❌ Failed to load model: {e}")
+    sys.exit(1)
+
 print("✅ Model loaded.")
 
-# Safe map loader
+# === Mapping file loader ===
 def load_map(file_path):
     mapping = {}
     try:
@@ -32,12 +39,12 @@ def load_map(file_path):
         print(f"⚠️ Failed to load {file_path}: {e}")
     return mapping
 
-# Load external data
+# === Load mappings ===
 disease_map = load_map("diseases.txt")
 treatment_map = load_map("treatments.txt")
 fertilizer_map = load_map("fertilizers.txt")
 
-# Load critical diseases
+# === 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())
@@ -45,22 +52,21 @@ except Exception as e:
     print(f"⚠️ Failed to load critical_diseases.txt: {e}")
     critical_diseases = set()
 
-# Main prediction logic
+# === Prediction Function ===
 def predict_disease(base64_img):
     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)}"}
+        return {"error": f"Invalid image input: {str(e)}"}
 
     try:
         result = model(image)[0]
         label = result["label"]
         confidence = float(result["score"])
-
         disease = disease_map.get(label, label)
-        treatment = treatment_map.get(disease, "Consult an expert for treatment.")
-        fertilizer = fertilizer_map.get(disease, "Use a general NPK 10-10-10 fertilizer.")
+        treatment = treatment_map.get(disease, "Consult an expert.")
+        fertilizer = fertilizer_map.get(disease, "Use general NPK 10-10-10.")
 
         output = {
             "disease_prediction": disease,
@@ -73,28 +79,28 @@ def predict_disease(base64_img):
             output["alert"] = True
 
         return output
-
     except Exception as e:
         return {"error": f"Prediction failed: {str(e)}"}
 
-# CLI interface
+# === Main CLI ===
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
-        description=(
-            "🌿 AI-Powered Tree Disease Detector\n\n"
-            "Examples:\n"
-            "  python app.py --image leaf.jpg --raw\n"
-            "  python app.py --image encoded.txt"
-        ),
+        description="🌿 Plant Disease Detector\n\nUsage examples:\n  python app.py --image leaf.jpg --raw\n  python app.py --image encoded.txt",
         formatter_class=argparse.RawTextHelpFormatter
     )
-    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)")
+
+    parser.add_argument("--image", type=str, help="Path to image or base64 file", required=False)
+    parser.add_argument("--raw", action="store_true", help="Use this flag if the input is a raw image file (e.g., JPG, PNG)")
     args = parser.parse_args()
 
+    if not args.image:
+        print("❗ Error: You must provide an image file using --image\n")
+        parser.print_help()
+        sys.exit(1)
+
     if not os.path.isfile(args.image):
         print(f"❌ File not found: {args.image}")
-        exit(1)
+        sys.exit(1)
 
     try:
         if args.raw:
@@ -106,6 +112,5 @@ if __name__ == "__main__":
 
         result = predict_disease(base64_img)
         print(json.dumps(result, indent=2))
-
     except Exception as e:
         print(json.dumps({"error": f"Unexpected error: {str(e)}"}, indent=2))