Create main.py

main.py

@@ -0,0 +1,63 @@
+# main.py
+
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import JSONResponse
+import uvicorn
+import shutil
+import os
+import uuid
+import cv2
+import numpy as np
+import base64
+
+from detectron_infer import predict as detectron_predict
+from yolo_infer import predict_yolo
+
+app = FastAPI()
+
+UPLOAD_DIR = "/tmp/uploads"
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+@app.get("/")
+async def root():
+    return {"message": "API is up and running!"}
+
+@app.post("/predict/")
+async def predict_endpoint(
+    file: UploadFile = File(...),
+    model: str = Form(...),  # either "detectron" or "yolo"
+):
+    # Save uploaded image
+    file_ext = file.filename.split('.')[-1]
+    filename = f"{uuid.uuid4()}.{file_ext}"
+    file_path = os.path.join(UPLOAD_DIR, filename)
+
+    with open(file_path, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
+
+    try:
+        if model == "Anomaly":
+            result_img, predictions = detectron_predict(file_path)
+        elif model == "Numbering":
+            result_img, predictions = predict_yolo(file_path)
+        else:
+            return JSONResponse({"error": "Invalid model choice"}, status_code=400)
+
+        # Encode image to bytes (optional)
+        _, img_encoded = cv2.imencode(".jpg", result_img)
+        img_bytes = img_encoded.tobytes()
+
+        img_b64 = base64.b64encode(img_bytes).decode('utf-8')
+
+        return {
+          "predictions": predictions,
+          "image_base64": img_b64
+              }
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
+    finally:
+        os.remove(file_path)  # Clean up uploaded file
+
+
+# Only for testing locally
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)