Update main.py

main.py

@@ -10,9 +10,8 @@ import cv2
 import numpy as np
 import base64
 
-from yolo_Numbering import predict as predict_yolo
-from yolo_infer import predict_yolo
-
+from yolo_numbering import predict_yolo as predict_yolo_numbering
+from yolo_anomaly import predict_yolo_anomaly
 app = FastAPI()
 
 UPLOAD_DIR = "/tmp/uploads"
@@ -36,9 +35,9 @@ async def predict_endpoint(
 
     try:
         if model == "Anomaly":
-            result_img, predictions = detectron_predict(file_path)
+            result_img, predictions = predict_yolo_anomaly(file_path)
         elif model == "Numbering":
-            result_img, predictions = predict_yolo(file_path)
+            result_img, predictions = predict_yolo_numbering(file_path)
         else:
             return JSONResponse({"error": "Invalid model choice"}, status_code=400)