Update app.py

app.py

@@ -5,7 +5,7 @@ import tensorflow as tf
 from fastapi import FastAPI, File, UploadFile, Form, HTTPException
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
-from typing import Dict, Any
+from typing import Dict, Any, List
 from datetime import datetime, timezone
 from io import BytesIO
 from sqlalchemy import create_engine, Column, Integer, String, DateTime, ForeignKey, Text, func
@@ -93,38 +93,31 @@ 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
+def generate_pdf(upload_id: int, total_workers: int, total_helmets: int, total_vests: int) -> str:
+    pdf_path = os.path.join(PROCESSED_FOLDER, f"report_{upload_id}.pdf")
+    doc = SimpleDocTemplate(pdf_path, pagesize=A4)
+    styles = getSampleStyleSheet()
+    elements = []
+    elements.append(Paragraph("Industrial Safety Detection Report", styles["Title"]))
+    elements.append(Spacer(1, 12))
+    data = [["Total Workers", total_workers],
+            ["Total Helmets Detected", total_helmets],
+            ["Total Vests Detected", total_vests]]
+    table = Table(data, colWidths=[200, 100])
+    table.setStyle(TableStyle([
+        ('BACKGROUND', (0, 0), (-1, 0), colors.grey),
+        ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
+        ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
+        ('FONTNAME', (0, 0), (-1, -1), 'Helvetica-Bold'),
+        ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
+        ('BACKGROUND', (0, 1), (-1, -1), colors.beige),
+    ]))
+    elements.append(table)
+    doc.build(elements)
+    return pdf_path
 
 @app.post("/upload")
 async def upload_file(approach: int = Form(...), file: UploadFile = File(...)):
-    global model
     db = SessionLocal()
     filename = file.filename
     filepath = os.path.join(UPLOAD_FOLDER, filename)
@@ -143,6 +136,7 @@ async def upload_file(approach: int = Form(...), file: UploadFile = File(...)):
     processed_filename = f"processed_{filename}"
     processed_path = os.path.join(PROCESSED_FOLDER, processed_filename)
     cv2.imwrite(processed_path, processed_img)
+    pdf_path = generate_pdf(upload_id, upload_obj.total_workers, upload_obj.total_helmets, upload_obj.total_vests)
     db.refresh(upload_obj)
     db.close()
     return {
@@ -151,6 +145,6 @@ async def upload_file(approach: int = Form(...), file: UploadFile = File(...)):
         "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 [],
-        "processed_image": processed_path
-    }
+        "processed_image": processed_path,
+        "pdf_report": pdf_path
+    }