Update app.py

app.py

@@ -67,9 +67,10 @@ class_names = ['H', 'V', 'W']
 num_classes = len(class_names)
 num_anchors = 9
 model = None
+anchor_boxes = None
 
 def prepare_model(approach: int):
-    global model
+    global model, anchor_boxes
     if approach not in [1, 2, 3]:
         raise ValueError("Approach must be 1, 2, or 3.")
 
@@ -83,10 +84,17 @@ def prepare_model(approach: int):
     if not os.path.exists(weight_path):
         raise FileNotFoundError(f"Weight file not found: {weight_path}")
 
+    anchor_boxes = np.array([
+        np.array([[10, 13], [16, 30], [33, 23]]) / 32,
+        np.array([[30, 61], [62, 45], [59, 119]]) / 16,
+        np.array([[116, 90], [156, 198], [373, 326]]) / 8
+    ], dtype="float64")
+
     input_tensor = Input(shape=(input_shape[0], input_shape[1], 3))
     num_out_filters = (num_anchors // 3) * (5 + num_classes)
     model = yolo_body(input_tensor, num_out_filters)
     model.load_weights(weight_path)
+    print(f"YOLO model (Approach={approach}) loaded successfully.")
 
 @app.on_event("startup")
 def on_startup():
@@ -120,13 +128,13 @@ async def upload_file(approach: int = Form(...), file: UploadFile = File(...)):
     return {"message": "File processed successfully.", "upload_id": upload_id, "pdf_path": pdf_path}
 
 def run_detection_on_frame(frame: np.ndarray, upload_id: int, db: Session) -> np.ndarray:
-    global model
+    global model, anchor_boxes
     ih, iw = frame.shape[:2]
     resized = letterbox_image(frame, input_shape)
     resized_expanded = np.expand_dims(resized, 0)
     image_data = np.array(resized_expanded) / 255.0
     prediction = model.predict(image_data)
-    boxes = detection(prediction, None, len(class_names), (ih, iw), input_shape, 50, 0.3, 0.45, False)[0].numpy()
+    boxes = detection(prediction, anchor_boxes, len(class_names), (ih, iw), input_shape, 50, 0.3, 0.45, False)[0].numpy()
     for box in boxes:
         x1, y1, x2, y2, _, cls_id = map(int, box)
     return frame