Update app.py

app.py

@@ -1,150 +1,78 @@
-import os
 import base64
 import torch
-import numpy as np
-import faiss
-import json
 
 from fastapi import FastAPI
 from pydantic import BaseModel
-from contextlib import asynccontextmanager
-from huggingface_hub import snapshot_download
-from sentence_transformers import SentenceTransformer
 from PIL import Image
 from io import BytesIO
 
-from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 
 # ─────────────────────────────
-# CONFIG
+# MODEL (REPLACE WITH YOUR CLASSIFIER)
 # ─────────────────────────────
-MODEL_REPO = "Rady10/Plant-Disease-Qwen3VL-2B"
-RAG_REPO = "Rady10/Agriculture-Rag-Data-Index"
+MODEL_REPO = "Rady10/Plant-Disease-Qwen3VL-2B"  # or your plant model
 
-DEVICE = "cpu"
+app = FastAPI(title="🌿 Plant Disease Classifier")
 
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-# ─────────────────────────────
-# GLOBALS
 # ─────────────────────────────
-model = None
-processor = None
-
-faiss_index = None
-rag_chunks = None
-embedder = None
-
-# ─────────────────────────────
-# FASTAPI APP
-# ─────────────────────────────
-app = FastAPI(title="🌿 Plant Disease Vision API")
-
+# LOAD MODEL
 # ─────────────────────────────
-# LOAD MODELS ONCE
-# ─────────────────────────────
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-
-    global model, processor, faiss_index, rag_chunks, embedder
-
-    print("Loading vision model...")
-
-    processor = AutoProcessor.from_pretrained(
-        MODEL_REPO,
-        trust_remote_code=True
-    )
-
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_REPO,
-        torch_dtype=torch.float32,
-        device_map="cpu",
-        trust_remote_code=True
-    )
-
-    model.eval()
+print("Loading classifier...")
 
-    # ───── RAG (optional but included) ─────
-    print("Loading RAG...")
+processor = AutoImageProcessor.from_pretrained(MODEL_REPO)
+model = AutoModelForImageClassification.from_pretrained(MODEL_REPO)
 
-    rag_dir = snapshot_download(
-        repo_id=RAG_REPO,
-        repo_type="dataset",
-        local_dir="./rag"
-    )
+model.eval()
 
-    faiss_index = faiss.read_index(
-        os.path.join(rag_dir, "agro.index")
-    )
-
-    with open(os.path.join(rag_dir, "chunks.json"), "r", encoding="utf-8") as f:
-        rag_chunks = json.load(f)
-
-    embedder = SentenceTransformer(
-        "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
-    )
-
-    print("ALL LOADED")
-
-    yield
-
-app = FastAPI(lifespan=lifespan)
+print("Model loaded")
 
 # ─────────────────────────────
-# REQUEST MODEL
+# REQUEST FORMAT
 # ─────────────────────────────
-class VisionRequest(BaseModel):
-    image: str   # base64
-    text: str = ""
+class Request(BaseModel):
+    image: str  # base64 string
 
 # ─────────────────────────────
-# IMAGE DECODER
+# DECODE IMAGE
 # ─────────────────────────────
-def decode_image(base64_str):
-    img_data = base64.b64decode(base64_str)
-    return Image.open(BytesIO(img_data)).convert("RGB")
+def decode_image(b64):
+    return Image.open(
+        BytesIO(base64.b64decode(b64))
+    ).convert("RGB")
 
 # ─────────────────────────────
-# GENERATION
+# PREDICTION FUNCTION
 # ─────────────────────────────
-def generate(image, text):
+def predict(image):
 
-    if text.strip() == "":
-        text = "What disease is shown in this plant image?"
+    inputs = processor(images=image, return_tensors="pt")
 
-    inputs = processor(
-        text=text,
-        images=image,
-        return_tensors="pt"
-    )
+    with torch.no_grad():
+        outputs = model(**inputs)
 
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    logits = outputs.logits
 
-    with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            temperature=0.7,
-            top_p=0.9
-        )
-
-    return processor.batch_decode(
-        output,
-        skip_special_tokens=True
-    )[0]
+    pred_id = logits.argmax(-1).item()
+
+    label = model.config.id2label[pred_id]
+
+    return label
 
 # ─────────────────────────────
-# API ROUTES
+# API ENDPOINT
 # ────────���────────────────────
-@app.get("/")
-def root():
-    return {"status": "vision api running"}
-
-@app.post("/analyze")
-def analyze(req: VisionRequest):
+@app.post("/predict")
+def predict_api(req: Request):
 
     image = decode_image(req.image)
 
-    result = generate(image, req.text)
+    label = predict(image)
+
+    return {
+        "prediction": label
+    }
 
-    return {"response": result}
+@app.get("/")
+def home():
+    return {"status": "classifier running"}