Update app.py

app.py

@@ -93,6 +93,35 @@ def on_startup():
     fix_tf_gpu()
     prepare_model(approach=1)
 
+def run_detection_on_frame(frame: np.ndarray, upload_id: int, db: Session) -> np.ndarray:
+    global model, input_shape, class_names
+    ih, iw = frame.shape[:2]
+    resized = letterbox_image(frame, input_shape)
+    image_data = np.expand_dims(resized, 0) / 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, score, cls_id = box
+        x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
+        cls_id = int(cls_id)
+        label = class_names[cls_id]
+        color = (0, 255, 0) if label == 'W' else (255, 0, 0) if label == 'H' else (0, 0, 255)
+        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+        cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
+        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
+
 @app.post("/upload")
 async def upload_file(approach: int = Form(...), file: UploadFile = File(...)):
     global model
@@ -115,7 +144,6 @@ async def upload_file(approach: int = Form(...), file: UploadFile = File(...)):
     processed_path = os.path.join(PROCESSED_FOLDER, processed_filename)
     cv2.imwrite(processed_path, processed_img)
     db.refresh(upload_obj)
-    pdf_path = generate_pdf(upload_obj)
     db.close()
     return {
         "message": "File processed successfully.",
@@ -124,5 +152,5 @@ async def upload_file(approach: int = Form(...), file: UploadFile = File(...)):
         "total_helmets": upload_obj.total_helmets,
         "total_vests": upload_obj.total_vests,
         "detected_items": upload_obj.worker_images.split(',') if upload_obj.worker_images else [],
-        "pdf_download_link": pdf_path
+        "processed_image": processed_path
     }