Update main.py

main.py

@@ -1,27 +1,43 @@
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import FileResponse
-from fastapi.responses import JSONResponse
+from fastapi.responses import FileResponse, JSONResponse
 import os
 import uvicorn
 
 app = FastAPI()
 
-# Directory where model files are stored
 MODEL_DIRECTORY = "models"
 
+# Code-to-plant mapping starting from "000001"
+CODE_TO_PLANT = {
+    "000001": "Apple",
+    "000002": "Cherry",
+    "000003": "Corn",
+    "000004": "Grape",
+    "000005": "Peach",
+    "000006": "Pepperbell",
+    "000007": "Potato",
+    "000008": "Rice",
+    "000009": "Strawberry",
+    "000010": "Tomato"
+}
+
 @app.get("/health")
 async def api_health_check():
     return JSONResponse(content={"status": "Service is running"})
 
-@app.get("/download/{plant_name}")
-async def download_model(plant_name: str):
+@app.get("/download_by_code/{code}")
+async def download_model_by_code(code: str):
+    plant_name = CODE_TO_PLANT.get(code)
+    if not plant_name:
+        raise HTTPException(status_code=404, detail="Invalid 6-digit code provided")
+
     filename = f"model_{plant_name}.tflite"
     file_path = os.path.join(MODEL_DIRECTORY, filename)
 
-    # Check if file exists
     if not os.path.isfile(file_path):
         raise HTTPException(status_code=404, detail=f"Model file '{filename}' not found")
 
     return FileResponse(file_path, filename=filename, media_type="application/octet-stream")
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)