Update app.py

app.py

@@ -280,6 +280,14 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
 def load_model():
+    try:
+        # Check for timm dependency
+        import timm
+        logger.info("timm library is available.")
+    except ImportError as e:
+        logger.error("timm library is not installed. Install it with: pip install timm")
+        raise ImportError("timm is required for DetrConvEncoder. Run `pip install timm` and restart your runtime.") from e
+
     try:
         processor = DetrImageProcessor.from_pretrained(CONFIG["MODEL_NAME"])
         model = DetrForObjectDetection.from_pretrained(CONFIG["MODEL_NAME"]).to(device)
@@ -289,7 +297,7 @@ def load_model():
         logger.info(f"Model classes: {model.config.id2label}")
         return processor, model
     except Exception as e:
-        logger.error(f"Failed to load model: {e}")
+        logger.error(f"Failed to load model: {str(e)}")
         raise
 
 processor, model = load_model()
@@ -309,21 +317,21 @@ def is_unsafe_posture(box, frame_shape):
     height = y2 - y1
     width = x2 - x1
     aspect_ratio = height / max(width, 1)
-    return aspect_ratio > 2.0  # Tall, narrow box suggests bending/unsafe posture
+    return aspect_ratio > 2.0
 
 def is_improper_tool_use(person_box, tool_box):
     """Placeholder for improper tool use. Fine-tune DETR for specific tools."""
     person_center = ((person_box[0] + person_box[2]) / 2, (person_box[1] + person_box[3]) / 2)
     tool_center = ((tool_box[0] + tool_box[2]) / 2, (tool_box[1] + tool_box[3]) / 2)
     dist = distance.euclidean(person_center, tool_center)
-    return dist > 100  # Tool too far from person
+    return dist > 100
 
 def is_unsafe_zone(person_box, frame_shape):
     """Check if person is in restricted area (e.g., top-left quadrant)."""
     px, py, pw, ph = person_box
     frame_h, frame_w = frame_shape
     person_center = (px + pw / 2, py + ph / 2)
-    unsafe_zone = (0, 0, 0.5, 0.5)  # Top-left quadrant
+    unsafe_zone = (0, 0, 0.5, 0.5)
     return (unsafe_zone[0] * frame_w < person_center[0] < unsafe_zone[2] * frame_w and
             unsafe_zone[1] * frame_h < person_center[1] < unsafe_zone[3] * frame_h)
 
@@ -487,7 +495,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         logger.info(f"Creating Salesforce record with data: {record_data}")
         try:
             record = sf.Safety_Video_Report__c.create(record_data)
-            logger.info(f"Created Safety_Video_Report__c record: {record['id']}")
+            logger.info(f"Created record: {record['id']}")
         except Exception as e:
             logger.error(f"Failed to create Safety_Video_Report__c: {e}")
             record = sf.Account.create({"Name": f"Safety_Report_{int(time.time())}"})
@@ -497,12 +505,12 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             uploaded_url = upload_pdf_to_salesforce(sf, pdf_file, record_id)
             if uploaded_url:
                 try:
-                    sf.Safety_Video_Report__c.update(record_id, {"PDF_Report_URL__c": uploaded_url})
+                    sf.Safety_Video_Report__c.update(record_id, {"PDF_Report_URL": uploaded_url})
                     logger.info(f"Updated record {record_id} with PDF URL: {uploaded_url}")
                 except Exception as e:
                     logger.error(f"Failed to update Safety_Video_Report__c: {e}")
                     sf.Account.update(record_id, {"Description": uploaded_url})
-                    logger.info(f"Updated Account record {record_id} with PDF URL")
+                    logger.info(f"Updated account record {record_id} with PDF URL")
                 pdf_url = uploaded_url
         return record_id, pdf_url
     except Exception as e:
@@ -626,7 +634,7 @@ def process_video(video_data, temp_dir):
                     break
                 ret, frame = cap.read()
                 if not ret:
-                    logger.warning(f"Failed to read frame {frame_idx}. Skipping.")
+                    logger.warning(f"Failed to read frame {frame_idx}. Skipping...")
                     break
                 original_frame = frame.copy()
                 frame = preprocess_frame(frame)
@@ -636,7 +644,7 @@ def process_video(video_data, temp_dir):
                 batch_frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
                 batch_indices.append(frame_idx)
                 batch_originals.append(original_frame)
-                processed_frames += 1
+                processed_frames += frame_skip + 1
 
             if not batch_frames:
                 logger.info("No more frames to process.")
@@ -663,7 +671,7 @@ def process_video(video_data, temp_dir):
                 progress = (processed_frames / total_frames) * 100
                 elapsed_time = current_time - start_time
                 fps_processed = processed_frames / elapsed_time if elapsed_time > 0 else 0
-                yield f"Processing video... {progress:.1f}% complete (Frame {processed_frames}/{total_frames}, {fps_processed:.1f} FPS)", "", "", "", ""
+                yield f"Processing video... {progress:.1f}% complete (Frame {processed_frames}/{total_frames}, {fps_processed:.1f} FPS)", "", "", "", "", ""
                 last_yield_time = current_time
 
             for i, (result, frame_idx, original_frame) in enumerate(zip(results, batch_indices, batch_originals)):
@@ -681,8 +689,8 @@ def process_video(video_data, temp_dir):
                     bbox_xywh = [x + w/2, y + h/2, w, h]
 
                     if label_name in ["no_helmet", "no_harness"] and conf >= CONFIG["CONFIDENCE_THRESHOLDS"].get(label_name, 0.25):
-                        if label_name == "no_helmet" and not validate_helmet_detection(original_frame, bbox_xywh, conf):
-                            logger.info(f"Frame {frame_idx}: Helmet false positive filtered at {conf:.2f} confidence")
+                        if label_name == "no_helmet" and not validate_helmet(original_frame, bbox_xywh, conf):
+                            logger.info(f"Frame {frame_idx}: Height false positive violation filtered out at {conf:.2f} confidence")
                             continue
                         track_inputs.append({"bbox": bbox_xywh, "conf": conf, "cls": label_name})
                     elif label_name == "person":
@@ -690,7 +698,7 @@ def process_video(video_data, temp_dir):
                     elif label_name in ["hammer", "wrench"]:  # Example tools; update with your dataset
                         tool_boxes.append(bbox_xywh)
 
-                # Handle Unsafe Posture, Unsafe Zone, Improper Tool Use
+                # Handle Unsafe violations
                 for pbox in person_boxes:
                     if is_unsafe_posture(pbox, original_frame.shape[:2]):
                         track_inputs.append({"bbox": pbox, "conf": 0.9, "cls": "unsafe_posture"})
@@ -714,7 +722,6 @@ def process_video(video_data, temp_dir):
                     tracker_id = obj['id']
                     label = obj['cls']
                     conf = obj['score']
-                    bbox = obj['bbox']
                     
                     if label not in CONFIG["VIOLATION_LABELS"]:
                         continue
@@ -764,7 +771,7 @@ def process_video(video_data, temp_dir):
 
         if not violations:
             logger.info("No violations detected after processing")
-            yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"
+            yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A", ""
             return
 
         snapshots = []
@@ -865,12 +872,13 @@ def process_video(video_data, temp_dir):
             f"Safety Score: {score}%",
             snapshots_text,
             f"Salesforce Record ID: {record_id}",
-            final_pdf_url
+            final_pdf_url,
+            ""
         )
 
     except Exception as e:
         logger.error(f"Error processing video: {str(e)}", exc_info=True)
-        yield f"Error processing video: {str(e)}", "", "", "", ""
+        yield f"Error processing video: {str(e)}", "", "", "", "", ""
     finally:
         if video_path and os.path.exists(video_path):
             try:
@@ -886,7 +894,7 @@ def gradio_interface(video_file):
     local_video_path = None
     try:
         if not video_file:
-            return "No file uploaded.", "", "No file uploaded.", "", ""
+            return "No file uploaded.", "", "No file uploaded.", "", "", ""
         
         temp_dir = tempfile.mkdtemp(prefix="DETR_")
         logger.info(f"Created temporary directory for video processing: {temp_dir}")
@@ -896,7 +904,7 @@ def gradio_interface(video_file):
         logger.info(f"Read Gradio video file: {video_file}, size: {len(video_data)} bytes")
         
         if len(video_data) == 0:
-            return "Uploaded video file is empty.", "", "", "", ""
+            return "Uploaded video file is empty.", "", "", "", "", ""
 
         with tempfile.NamedTemporaryFile(suffix=".mp4", dir=temp_dir, delete=False) as temp_file:
             temp_file.write(video_data)
@@ -905,14 +913,14 @@ def gradio_interface(video_file):
         logger.info(f"Copied Gradio video to local temporary file: {local_video_path}")
 
         if not FFMPEG_AVAILABLE:
-            return "FFmpeg is not available in the environment. Please install FFmpeg to process videos.", "", "", "", ""
+            return "FFmpeg is not available in the environment. Please install FFmpeg to process videos.", "", "", "", "", ""
 
-        for status, score, snapshots_text, record_id, details_url in process_video(video_data, temp_dir):
-            yield status, score, snapshots_text, record_id, details_url
+        for status, score, snapshots_text, record_id, details_url, _ in process_video(video_data, temp_dir):
+            yield status, score, snapshots_text, record_id, details_url, ""
             
     except Exception as e:
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
-        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
+        yield f"Error: {str(e)}", "", "Error in processing.", "", "", ""
     finally:
         if local_video_path and os.path.exists(local_video_path):
             try:
@@ -936,7 +944,8 @@ interface = gr.Interface(
         gr.Textbox(label="Compliance Score"),
         gr.Markdown(label="Snapshots"),
         gr.Textbox(label="Salesforce Record ID"),
-        gr.Textbox(label="Violation Details URL")
+        gr.Textbox(label="Violation Details URL"),
+        gr.Textbox(label="Error Log", visible=False)
     ],
     title="Worksite Safety Violation Analyzer",
     description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use). Each unique violation is detected only once per worker.",