Update app.py

app.py

@@ -24,7 +24,8 @@ CONFIG = {
     "VIOLATION_LABELS": {
         0: "no_helmet",
         1: "no_harness",
-        2: "unsafe_posture"
+        2: "unsafe_posture",
+        3: "unsafe_zone"  # Ignored in scoring
     },
     "SF_CREDENTIALS": {
         "username": "prashanth1ai@safety.com",
@@ -33,7 +34,8 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/",
-    "MAX_PROCESSING_TIME": 60  # Set to 60 seconds for 1-minute video
+    "FRAME_SKIP": 5,
+    "MAX_PROCESSING_TIME": 30
 }
 
 # Setup logging
@@ -178,7 +180,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         logger.info(f"Salesforce record created: {record_id}")
 
         if pdf_file:
-            uploaded_url = upload_pdf_to_salesforce(sf, pdf_file, record_id)
+            uploaded_url = upload_pdf_to_salesforce(sf, pdf_file, report_id)
             if uploaded_url:
                 try:
                     sf.Safety_Video_Report__c.update(record_id, {"PDF_Report_URL__c": uploaded_url})
@@ -196,17 +198,17 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
 # ==========================
 # Safety Score Calculation
 # ==========================
-def calculate_safety_score(violations, total_frames):
+def calculate_safety_score(violations):
     penalties = {
-        "no_helmet": 25 / max(total_frames, 1),  # Normalize penalty by video length
-        "no_harness": 30 / max(total_frames, 1),
-        "unsafe_posture": 20 / max(total_frames, 1)
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20
     }
     score = 100
     for v in violations:
-        penalty = penalties.get(v["violation"], 0)
-        score -= penalty
-    return max(round(score, 2), 0)
+        if v["violation"] in penalties:
+            score -= penalties[v["violation"]]
+    return max(score, 0)
 
 # ==========================
 # Video Processing
@@ -224,30 +226,28 @@ def process_video(video_data):
 
         violations, snapshots = [], []
         frame_count = 0
-        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
-        fps = video.get(cv2.CAP_PROP_FPS)
-        video_duration = total_frames / max(fps, 1)
-        logger.info(f"Video duration: {video_duration:.2f} seconds, Total frames: {total_frames}, FPS: {fps:.2f}")
-
         start_time = time.time()
-        max_frames = int(fps * CONFIG["MAX_PROCESSING_TIME"])  # Max frames for 60 seconds
+        fps = video.get(cv2.CAP_PROP_FPS)
+        max_frames = int(CONFIG["MAX_PROCESSING_TIME"] * fps)
 
         while True:
             ret, frame = video.read()
-            if not ret:
-                logger.info(f"End of video reached at frame {frame_count}")
-                break
-            if frame_count >= max_frames:
-                logger.info(f"Reached maximum processing limit at frame {frame_count}")
+            if not ret or frame_count >= max_frames:
                 break
 
+            if frame_count % CONFIG["FRAME_SKIP"] != 0:
+                frame_count += 1
+                continue
+
             results = model(frame, device=device)
             seen_violations = set()
             for result in results:
                 for box in result.boxes:
                     cls, conf = int(box.cls), float(box.conf)
                     label = CONFIG["VIOLATION_LABELS"].get(cls, f"class_{cls}")
-                    if label in seen_violations or label not in ["no_helmet", "no_harness", "unsafe_posture"]:
+                    if label not in ["no_helmet", "no_harness", "unsafe_posture"]:
+                        continue
+                    if label in seen_violations:
                         continue
                     seen_violations.add(label)
 
@@ -256,10 +256,9 @@ def process_video(video_data):
                         "violation": label,
                         "confidence": round(conf, 2),
                         "bounding_box": [round(x, 2) for x in box.xywh.cpu().numpy()[0]],
-                        "timestamp": frame_count / max(fps, 1)
+                        "timestamp": frame_count / fps
                     }
                     violations.append(violation)
-                    logger.debug(f"Violation detected: {label} at frame {frame_count}, timestamp {violation['timestamp']:.2f}s")
 
                     snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], f"snapshot_{frame_count}_{label}.jpg")
                     cv2.imwrite(snapshot_path, frame)
@@ -274,13 +273,20 @@ def process_video(video_data):
 
             frame_count += 1
 
-        processing_time = time.time() - start_time
-        logger.info(f"Processed {frame_count} frames in {processing_time:.2f} seconds")
-
         video.release()
         os.remove(video_path)
-        score = calculate_safety_score(violations, frame_count)
-        logger.info(f"Safety score calculated: {score}% with {len(violations)} violations")
+
+        if not violations:
+            logger.info("No violations detected, skipping score calculation and reporting")
+            return {
+                "violations": [],
+                "snapshots": [],
+                "score": 100,
+                "salesforce_record_id": None,
+                "violation_details_url": ""
+            }
+
+        score = calculate_safety_score(violations)
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
         report_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
@@ -289,20 +295,16 @@ def process_video(video_data):
             "snapshots": snapshots,
             "score": score,
             "salesforce_record_id": report_id,
-            "violation_details_url": final_pdf_url,
-            "processed_frames": frame_count,
-            "video_duration": video_duration
+            "violation_details_url": final_pdf_url
         }
     except Exception as e:
-        logger.error(f"Error processing video: {e}", exc_info=True)
+        logger.error(f"Error processing video: {e}")
         return {
             "violations": [],
             "snapshots": [],
-            "score": 0,
+            "score": 100,
             "salesforce_record_id": None,
-            "violation_details_url": "",
-            "processed_frames": 0,
-            "video_duration": 0
+            "violation_details_url": ""
         }
 
 # ==========================
@@ -310,7 +312,7 @@ def process_video(video_data):
 # ==========================
 def gradio_interface(video_file):
     if not video_file:
-        return "No file uploaded.", "", "No file uploaded.", "", "", ""
+        return "No file uploaded.", "", "No file uploaded.", "", ""
     try:
         with open(video_file, "rb") as f:
             video_data = f.read()
@@ -318,25 +320,19 @@ def gradio_interface(video_file):
 
         violation_table = "No violations detected."
         if result["violations"]:
-            header = "| Violation        | Timestamp | Confidence | Bounding Box             | Violation Details       |\n"
-            separator = "|------------------|-----------|------------|--------------------------|-------------------------|\n"
+            header = "| Violation     | Timestamp | Confidence | Bounding Box             | Violation Details       |\n"
+            separator = "|---------------|-----------|------------|--------------------------|-------------------------|\n"
             rows = []
             for v in result["violations"]:
-                violation_name = v["violation"]
-                if violation_name == "no_helmet":
-                    violation_name = "Missing Helmet"
-                elif violation_name == "no_harness":
-                    violation_name = "Missing Harness"
-                else:
-                    violation_name = "Unsafe Posture"
-                row = f"| {violation_name:<16} | {v['timestamp']:.2f}s  | {v['confidence']:.2f}      | {v['bounding_box']} | {result['violation_details_url']} |"
+                violation_name = v["violation"].replace("no_", "").replace("unsafe_", "")
+                row = f"| {violation_name:<13} | {v['timestamp']:.2f}s  | {v['confidence']:.2f}      | {v['bounding_box']} | {result['violation_details_url']} |"
                 rows.append(row)
             violation_table = header + separator + "\n".join(rows)
 
         snapshots_text = "No snapshots captured."
         if result["snapshots"]:
             snapshots_text = "\n".join(
-                f"- Snapshot for {s['violation'].replace('no_helmet', 'Missing Helmet').replace('no_harness', 'Missing Harness').replace('unsafe_posture', 'Unsafe Posture')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
+                f"- Snapshot for {s['violation'].replace('no_', '').replace('unsafe_', '')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
                 for s in result["snapshots"]
             )
 
@@ -345,26 +341,24 @@ def gradio_interface(video_file):
             f"Safety Score: {result['score']}%",
             snapshots_text,
             f"Salesforce Record ID: {result['salesforce_record_id'] or 'N/A'}",
-            result["violation_details_url"] or "N/A",
-            f"Processed {result['processed_frames']} frames (Video duration: {result['video_duration']:.2f} seconds)"
+            result["violation_details_url"] or "N/A"
         )
     except Exception as e:
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
-        return f"Error: {str(e)}", "", "Error in processing.", "", "", ""
+        return f"Error: {str(e)}", "", "Error in processing.", "", ""
 
 interface = gr.Interface(
     fn=gradio_interface,
-    inputs=gr.Video(label="Upload Site Video (Up to 1 Minute)"),
+    inputs=gr.Video(label="Upload Site Video"),
     outputs=[
         gr.Markdown(label="Detected Safety Violations"),
         gr.Textbox(label="Compliance Score"),
         gr.Markdown(label="Snapshots"),
         gr.Textbox(label="Salesforce Record ID"),
-        gr.Textbox(label="Violation Details URL"),
-        gr.Textbox(label="Processing Summary")
+        gr.Textbox(label="Violation Details URL")
     ],
     title="Worksite Safety Violation Analyzer",
-    description="Upload a video (up to 1 minute) to detect safety violations (missing helmet, missing harness, unsafe posture). Violations are only recorded when safety gear is missing or posture is unsafe."
+    description="Upload site videos to detect safety violations (e.g., no helmet, no harness, unsafe posture)."
 )
 
 if __name__ == "__main__":