Update app.py

app.py

@@ -67,10 +67,9 @@ 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, anchor_boxes
+    global model
     if approach not in [1, 2, 3]:
         raise ValueError("Approach must be 1, 2, or 3.")
 
@@ -84,17 +83,10 @@ 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():
@@ -125,35 +117,12 @@ async def upload_file(approach: int = Form(...), file: UploadFile = File(...)):
     db.refresh(upload_obj)
     pdf_path = generate_pdf(upload_obj)
     db.close()
-    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, 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, 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
-
-def generate_pdf(upload_obj: Upload):
-    buffer = BytesIO()
-    doc = SimpleDocTemplate(buffer, pagesize=A4)
-    elements = []
-    styles = getSampleStyleSheet()
-    elements.append(Paragraph("Industrial Safety Report", styles["Title"]))
-    elements.append(Paragraph(f"Filename: {upload_obj.filename}", styles["Normal"]))
-    elements.append(Paragraph(f"Timestamp: {upload_obj.timestamp.strftime('%Y-%m-%d %H:%M:%S')}", styles["Normal"]))
-    data = [["Total Workers", upload_obj.total_workers], ["Total Helmets", upload_obj.total_helmets], ["Total Vests", upload_obj.total_vests]]
-    table = Table(data)
-    table.setStyle(TableStyle([("BACKGROUND", (0, 0), (-1, 0), colors.grey), ("TEXTCOLOR", (0, 0), (-1, 0), colors.whitesmoke), ("GRID", (0, 0), (-1, -1), 1, colors.black)]))
-    elements.append(table)
-    doc.build(elements)
-    buffer.seek(0)
-    pdf_path = os.path.join(PROCESSED_FOLDER, f"report_{upload_obj.id}.pdf")
-    with open(pdf_path, "wb") as f:
-        f.write(buffer.getvalue())
-    return pdf_path
+    return {
+        "message": "File processed successfully.",
+        "upload_id": upload_id,
+        "total_workers": upload_obj.total_workers,
+        "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
+    }