Update app.py

app.py

@@ -18,8 +18,7 @@ from retrying import retry
 # Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": os.getenv("SAFETY_MODEL_PATH", "models/yolov8_safety.pt"),
-    "FALLBACK_MODEL": "yolov8n.pt",
+    "MODEL_PATH": "yolov8n.pt",  # Force lightweight Nano model
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
@@ -34,12 +33,12 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/",
-    "FRAME_SKIP": 5,
-    "MAX_PROCESSING_TIME": 60  # Updated to 60 seconds
+    "FRAME_SKIP": 15,  # Increased to reduce frames processed
+    "MAX_PROCESSING_TIME": 25  # Cap video processing at 25s to leave time for reporting
 }
 
 # Setup logging
-logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
 # Ensure output directory exists
@@ -55,10 +54,9 @@ logger.info(f"Using device: {device}")
 # Model Loading
 # ==========================
 def load_model():
-    model_path = CONFIG["MODEL_PATH"] if os.path.isfile(CONFIG["MODEL_PATH"]) else CONFIG["FALLBACK_MODEL"]
     try:
-        model = YOLO(model_path).to(device)
-        logger.info(f"Model loaded: {model_path}")
+        model = YOLO(CONFIG["MODEL_PATH"]).to(device)
+        logger.info(f"Model loaded: {CONFIG['MODEL_PATH']}")
         return model
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
@@ -69,13 +67,12 @@ model = load_model()
 # ==========================
 # Salesforce Integration
 # ==========================
-@retry(stop_max_attempt_number=3, wait_fixed=2000)
+@retry(stop_max_attempt_number=2, wait_fixed=1000)
 def connect_to_salesforce():
     try:
         sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
         logger.info("Connected to Salesforce")
         sf.describe()
-        logger.debug("Salesforce object metadata fetched successfully")
         return sf
     except Exception as e:
         logger.error(f"Salesforce connection failed: {e}")
@@ -126,7 +123,7 @@ def generate_violation_pdf(violations, score):
         logger.error(f"Error generating PDF: {e}")
         return "", "", None
 
-@retry(stop_max_attempt_number=3, wait_fixed=2000)
+@retry(stop_max_attempt_number=2, wait_fixed=1000)
 def upload_pdf_to_salesforce(sf, pdf_file, report_id):
     try:
         if not pdf_file:
@@ -139,9 +136,8 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
             "VersionData": encoded_pdf,
             "FirstPublishLocationId": report_id
         }
-        logger.debug(f"Uploading PDF with data: {content_version_data}")
         content_version = sf.ContentVersion.create(content_version_data)
-        result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version['id']}'")
+        result = sf.query(f"SELECT Id FROM ContentVersion WHERE Id = '{content_version['id']}'")
         if not result['records']:
             logger.error("Failed to retrieve ContentVersion")
             return ""
@@ -152,7 +148,7 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
         logger.error(f"Error uploading PDF to Salesforce: {e}")
         return ""
 
-@retry(stop_max_attempt_number=3, wait_fixed=2000)
+@retry(stop_max_attempt_number=2, wait_fixed=1000)
 def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
     try:
         sf = connect_to_salesforce()
@@ -169,15 +165,14 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             "Status__c": "Pending",
             "PDF_Report_URL__c": pdf_url
         }
-        logger.debug(f"Attempting to create Salesforce record with data: {record_data}")
+        logger.info(f"Creating Salesforce record with data: {record_data}")
         try:
             record = sf.Safety_Video_Report__c.create(record_data)
-            logger.info(f"Successfully created Safety_Video_Report__c record: {record['id']}")
+            logger.info(f"Created Safety_Video_Report__c record: {record['id']}")
         except Exception as e:
             logger.error(f"Failed to create Safety_Video_Report__c: {e}")
-            # Fallback to Account object
             record = sf.Account.create({"Name": f"Safety_Report_{int(time.time())}"})
-            logger.warning(f"Fell back to creating Account record: {record['id']}")
+            logger.warning(f"Fell back to Account record: {record['id']}")
         record_id = record["id"]
 
         if pdf_file:
@@ -185,16 +180,16 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             if uploaded_url:
                 try:
                     sf.Safety_Video_Report__c.update(record_id, {"PDF_Report_URL__c": uploaded_url})
-                    logger.debug(f"Updated Safety_Video_Report__c record {record_id} with PDF 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.debug(f"Updated Account record {record_id} with PDF URL: {uploaded_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:
-        logger.error(f"Salesforce record creation failed: {e}", exc_info=True)
+        logger.error(f"Salesforce record creation failed: {e}")
         return None, ""
 
 # ==========================
@@ -228,8 +223,12 @@ def process_video(video_data):
 
         violations, snapshots = [], []
         frame_count = 0
+        start_time = time.time()
         fps = video.get(cv2.CAP_PROP_FPS)
-        max_frames = int(CONFIG["MAX_PROCESSING_TIME"] * fps)
+        max_frames = int(60 * fps)  # Process up to 1 minute
+
+        # Track one snapshot per violation type
+        snapshot_taken = {"no_helmet": False, "no_harness": False, "unsafe_posture": False}
 
         while True:
             ret, frame = video.read()
@@ -240,6 +239,11 @@ def process_video(video_data):
                 frame_count += 1
                 continue
 
+            # Stop if processing time exceeds 25 seconds
+            if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
+                logger.info("Processing time limit reached")
+                break
+
             results = model(frame, device=device)
             seen_violations = set()
             for result in results:
@@ -261,16 +265,19 @@ def process_video(video_data):
                     }
                     violations.append(violation)
 
-                    snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], f"snapshot_{frame_count}_{label}.jpg")
-                    cv2.imwrite(snapshot_path, frame)
-                    with open(snapshot_path, "rb") as img_file:
-                        img_base64 = base64.b64encode(img_file.read()).decode('utf-8')
-                    snapshots.append({
-                        "violation": label,
-                        "frame": frame_count,
-                        "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(snapshot_path)}",
-                        "snapshot_base64": f"data:image/jpeg;base64,{img_base64}"
-                    })
+                    # Save only one snapshot per violation type
+                    if not snapshot_taken[label]:
+                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], f"snapshot_{frame_count}_{label}.jpg")
+                        cv2.imwrite(snapshot_path, frame)
+                        with open(snapshot_path, "rb") as img_file:
+                            img_base64 = base64.b64encode(img_file.read()).decode('utf-8')
+                        snapshots.append({
+                            "violation": label,
+                            "frame": frame_count,
+                            "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(snapshot_path)}",
+                            "snapshot_base64": f"data:image/jpeg;base64,{img_base64}"
+                        })
+                        snapshot_taken[label] = True
 
             frame_count += 1
 
@@ -278,7 +285,7 @@ def process_video(video_data):
         os.remove(video_path)
 
         if not violations:
-            logger.info("No violations detected, skipping score calculation and reporting")
+            logger.info("No violations detected")
             return {
                 "violations": [],
                 "snapshots": [],
@@ -345,7 +352,7 @@ def gradio_interface(video_file):
             result["violation_details_url"] or "N/A"
         )
     except Exception as e:
-        logger.error(f"Error in Gradio interface: {e}", exc_info=True)
+        logger.error(f"Error in Gradio interface: {e}")
         return f"Error: {str(e)}", "", "Error in processing.", "", ""
 
 interface = gr.Interface(