Update app.py

app.py

@@ -68,20 +68,30 @@ num_classes = len(class_names)
 num_anchors = 9
 model = None
 
-def prepare_model():
+def prepare_model(approach: int):
     global model
+    if approach not in [1, 2, 3]:
+        raise ValueError("Approach must be 1, 2, or 3.")
+
+    weight_files = {
+        1: "pictor-ppe-v302-a1-yolo-v3-weights.h5",
+        2: "pictor-ppe-v302-a2-yolo-v3-weights.h5",
+        3: "pictor-ppe-v302-a3-yolo-v3-weights.h5",
+    }
+    
+    weight_path = os.path.join("model-data", "weights", weight_files[approach])
+    if not os.path.exists(weight_path):
+        raise FileNotFoundError(f"Weight file not found: {weight_path}")
+
     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)
-    weight_path = "model-data/weights/yolo_weights.h5"
-    if not os.path.exists(weight_path):
-        raise FileNotFoundError(f"Weight file not found: {weight_path}")
     model.load_weights(weight_path)
 
 @app.on_event("startup")
 def on_startup():
     fix_tf_gpu()
-    prepare_model()
+    prepare_model(approach=1)
 
 @app.post("/upload")
 async def upload_file(approach: int = Form(...), file: UploadFile = File(...)):
@@ -117,22 +127,8 @@ def run_detection_on_frame(frame: np.ndarray, upload_id: int, db: Session) -> np
     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()
-    workers, helmets, vests = [], [], []
     for box in boxes:
         x1, y1, x2, y2, _, cls_id = map(int, box)
-        label = class_names[cls_id]
-        if label == 'W':
-            workers.append((x1, y1, x2, y2))
-        elif label == 'H':
-            helmets.append((x1, y1, x2, y2))
-        elif label == 'V':
-            vests.append((x1, y1, x2, y2))
-    upload_obj = db.query(Upload).filter(Upload.id == upload_id).first()
-    if upload_obj:
-        upload_obj.total_workers += len(workers)
-        upload_obj.total_helmets += len(helmets)
-        upload_obj.total_vests += len(vests)
-        db.commit()
     return frame
 
 def generate_pdf(upload_obj: Upload):