Update app.py

app.py

@@ -29,19 +29,18 @@ logger = logging.getLogger(__name__)
 # Global lock for thread safety
 processing_lock = threading.Lock()
 
-# Load custom best.pt model with fallback
-model = None
+# Load custom best.pt model with validation
 try:
     model = YOLO("best.pt")
-    logger.info("Custom best.pt model loaded successfully")
+    logger.info("Attempting to load best.pt model")
+    if not os.path.exists("best.pt"):
+        logger.error("best.pt file not found in the root directory")
+        raise FileNotFoundError("best.pt not found")
+    model_info = model.model.names
+    logger.info(f"best.pt loaded successfully. Detected classes: {model_info}")
 except Exception as e:
     logger.error(f"Failed to load best.pt model: {str(e)}")
-    try:
-        model = YOLO("yolov8n.pt")
-        logger.info("Fallback to yolov8n.pt model loaded successfully")
-    except Exception as e:
-        logger.error(f"Failed to load fallback yolov8n.pt model: {str(e)}")
-        raise Exception(f"Model loading failed: {str(e)}")
+    raise Exception(f"Model loading failed: {str(e)}")
 
 # Define categories and relevant objects
 CATEGORIES = {
@@ -53,7 +52,7 @@ CATEGORIES = {
 
 # Color mapping for bounding boxes
 COLORS = [
-    (0, 255, 0),      # Green
+    (0, 0, 255),      # Red for Toll Booth emphasis
     (255, 0, 0),      # Blue
     (255, 182, 193),  # Light Red
     (255, 165, 0),    # Orange
@@ -62,7 +61,7 @@ COLORS = [
     (255, 215, 0)     # Gold
 ]
 
-# Enhanced detection for "Construction Not Started" (e.g., open areas or signage)
+# Enhanced detection for "Construction Not Started"
 def detect_construction_not_started(frame):
     logger.info("Detecting Construction Not Started areas")
     try:
@@ -93,7 +92,7 @@ def detect_construction_indicators(frame, x1, y1, x2, y2):
     contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     for cnt in contours:
         area = cv2.contourArea(cnt)
-        if 500 < area < 5000:  # Small objects like cones or signs
+        if 500 < area < 5000:
             return True
     return False
 
@@ -101,7 +100,7 @@ def detect_construction_indicators(frame, x1, y1, x2, y2):
 def process_frame(frame, category, class_colors, object_counts):
     annotated_frame = frame.copy()
     with processing_lock:
-        if category != "Operations Maintenance" and model is not None:
+        if category != "Operations Maintenance":
             try:
                 results = model(frame, verbose=False, device="cpu", imgsz=640)
                 logger.info(f"Processed frame with {len(results)} detections")
@@ -118,10 +117,9 @@ def process_frame(frame, category, class_colors, object_counts):
                         w, h = x2 - x1, y2 - y1
                         conf = float(box.conf)
                         if label == "Toll Booth" and conf > 0.9 and w * h > 5000:
-                            if detect_construction_indicators(frame, x1, y1, x2, y2):
-                                label = "Under Construction (Toll Booth)"
-                                conf = 0.95
-                            logger.info(f"Detected {label} with confidence {conf:.2f} at ({x1}, {y1}, {x2}, {y2})")
+                            logger.info(f"Detected Toll Booth at ({x1}, {y1}, {x2}, {y2}) with conf {conf}")
+                            label = "Under Construction (Toll Booth)"  # Force label for clarity
+                            conf = 0.95
                         object_counts[label] = object_counts.get(label, 0) + 1
                         if label not in class_colors:
                             class_colors[label] = COLORS[len(class_colors) % len(COLORS)]
@@ -130,8 +128,7 @@ def process_frame(frame, category, class_colors, object_counts):
                         cv2.putText(annotated_frame, f"{label} {conf:.2f}", (x1, y1-10),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
 
-        # Custom detection for Construction Not Started
-        if category == "Under Construction" and model is not None:
+        if category == "Under Construction":
             construction_dets = detect_construction_not_started(frame)
             for i, (x1, y1, x2, y2, label, conf) in enumerate(construction_dets):
                 object_counts[label] = object_counts.get(label, 0) + 1
@@ -144,7 +141,7 @@ def process_frame(frame, category, class_colors, object_counts):
 
     return annotated_frame
 
-# Optimized video processing with slower speed
+# Optimized video processing
 def process_video(video_path, category, output_dir="outputs"):
     logger.info(f"Processing video: {video_path} for category: {category}")
     try:
@@ -157,19 +154,19 @@ def process_video(video_path, category, output_dir="outputs"):
         frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         input_fps = int(cap.get(cv2.CAP_PROP_FPS)) or 30
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        max_frames = min(total_frames, input_fps * 40)  # Limit to 40 seconds
+        max_frames = min(total_frames, input_fps * 40)
 
         os.makedirs(output_dir, exist_ok=True)
         output_video_path = os.path.join(output_dir, f"{category}_output.mp4")
-        output_fps = 15  # Slower frame rate for clearer AI reading
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # H.264 compatible
+        output_fps = 15
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
         out = cv2.VideoWriter(output_video_path, fourcc, output_fps, (frame_width, frame_height))
 
         object_counts = {}
         screenshots = []
         frame_count = 0
         class_colors = {}
-        screenshot_interval = max(1, total_frames // 20)  # 20 screenshots
+        screenshot_interval = max(1, total_frames // 20)
 
         with ThreadPoolExecutor(max_workers=2) as executor:
             while cap.isOpened() and frame_count < max_frames:
@@ -218,7 +215,7 @@ def create_pie_chart(object_counts, output_dir, category):
         logger.error(f"Pie chart creation failed: {str(e)}")
         return None, f"Error creating pie chart: {str(e)}"
 
-# Generate PDF report with 3 screenshots per page
+# Generate PDF report
 def create_pdf_report(pie_path, screenshots, output_dir, category):
     logger.info(f"Creating PDF report for category: {category}")
     try:
@@ -226,16 +223,13 @@ def create_pdf_report(pie_path, screenshots, output_dir, category):
         c = canvas.Canvas(pdf_path, pagesize=letter)
         width, height = letter
 
-        # Title and timestamp
         c.setFont("Helvetica-Bold", 18)
         c.drawCentredString(width/2, height - 0.75*inch, f"AI Video Analysis Report - {category}")
         c.setFont("Helvetica", 12)
         c.drawString(1*inch, height - 1.25*inch, f"Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S IST')}")
 
-        # Process Description
-        c.drawString(1*inch, height - 1.75*inch, "Process: Video analyzed frame-by-frame with custom best.pt model, focusing on Toll Booths. Slower output at 15 FPS, 20 screenshots captured.")
+        c.drawString(1*inch, height - 1.75*inch, "Process: Video analyzed with best.pt, focusing on Toll Booths with 'Under Construction' indicators. 15 FPS, 20 screenshots.")
 
-        # Logs from app.log
         c.setFont("Helvetica-Bold", 12)
         c.drawString(1*inch, height - 2.25*inch, "Processing Logs")
         c.setFont("Helvetica", 10)
@@ -248,14 +242,12 @@ def create_pdf_report(pie_path, screenshots, output_dir, category):
                     c.showPage()
                     log_y = height - 1*inch
 
-        # Pie chart
         c.showPage()
         if pie_path and os.path.exists(pie_path):
             c.drawImage(pie_path, 1*inch, height - 4.5*inch, width=4*inch, height=4*inch, preserveAspectRatio=True)
         else:
             c.drawString(1*inch, height - 4*inch, "No pie chart available")
 
-        # Object counts table
         c.setFont("Helvetica-Bold", 12)
         c.drawString(5.5*inch, height - 4.5*inch, "Detected Objects")
         c.setFont("Helvetica", 10)
@@ -270,7 +262,6 @@ def create_pdf_report(pie_path, screenshots, output_dir, category):
                     c.showPage()
                     table_y = height - 1*inch
 
-        # Step-by-step screenshots (3 per page)
         c.showPage()
         c.setFont("Helvetica-Bold", 12)
         c.drawString(1*inch, height - 0.75*inch, "Step-by-Step Screenshots")
@@ -338,9 +329,9 @@ def compare_videos(video1_counts, video2_counts):
         logger.error(f"Video comparison failed: {str(e)}")
         return None, f"Error comparing videos: {str(e)}"
 
-# Main Gradio function with object_counts passed to PDF
+# Main Gradio function
 def analyze_videos(video1, video2, category):
-    logger.info(f"Starting analysis for category: {category}, Video1: {video1}, Video2: {video2}")
+    logger.info(f"Starting analysis at {datetime.now().strftime('%Y-%m-%d %H:%M:%S IST')} for category: {category}, Video1: {video1}, Video2: {video2}")
     try:
         if not video1:
             logger.error("No primary video provided")
@@ -350,13 +341,11 @@ def analyze_videos(video1, video2, category):
         os.makedirs(output_dir, exist_ok=True)
         logger.info(f"Output directory created: {output_dir}")
 
-        # Process videos
         output_videos = []
         all_screenshots = []
         all_object_counts = {}
         errors = []
 
-        # Process video1 for selected category or all
         if category == "All":
             categories = CATEGORIES.keys()
         else:
@@ -375,7 +364,6 @@ def analyze_videos(video1, video2, category):
                     errors.append(error)
                     logger.error(error)
 
-        # Process video2 for comparison (same category)
         comparison_df = None
         if video2 and category != "All":
             _, video2_counts, _, error = process_video(video2, category, output_dir)
@@ -388,7 +376,6 @@ def analyze_videos(video1, video2, category):
                     errors.append(comp_error)
                     logger.error(comp_error)
 
-        # Generate pie chart and PDF for each category
         pdf_files = []
         output_files = []
         global object_counts
@@ -413,7 +400,6 @@ def analyze_videos(video1, video2, category):
 
         output_files.extend(output_videos + all_screenshots)
 
-        # Create ZIP file
         zip_path, zip_error = create_zip_file(output_files, output_dir)
         if zip_error:
             errors.append(zip_error)
@@ -426,7 +412,7 @@ def analyze_videos(video1, video2, category):
         logger.error(f"Analysis failed: {str(e)}")
         return [], [], None, [], None, f"Error in analysis: {str(e)}"
 
-# Custom CSS for enhanced UI
+# Custom CSS
 custom_css = """
 #video-gallery, .gallery { 
     display: grid;