Upload 4 files

apply_mask.py

@@ -0,0 +1,72 @@
+import cv2
+import os
+import glob
+import numpy as np
+
+def apply_mask_and_crop(input_folder, mask_path, output_folder):
+    """
+    Apply binary mask to all images, crop to masked region, and save to output folder
+    
+    Args:
+        input_folder (str): Path to folder containing input images
+        mask_path (str): Path to binary mask image
+        output_folder (str): Path to save cropped masked images
+    """
+    # Load and prepare mask
+    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
+    if mask is None:
+        raise ValueError(f"Could not load mask from {mask_path}")
+    
+    _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
+    
+    # Create output directory if it doesn't exist
+    os.makedirs(output_folder, exist_ok=True)
+    
+    # Get list of image files
+    image_files = glob.glob(os.path.join(input_folder, "*.jpg")) + \
+                 glob.glob(os.path.join(input_folder, "*.png")) + \
+                 glob.glob(os.path.join(input_folder, "*.bmp"))
+    
+    if not image_files:
+        print(f"No images found in {input_folder}")
+        return
+    
+    print(f"Found {len(image_files)} images to process")
+    
+    for img_path in image_files:
+        # Load image
+        img = cv2.imread(img_path)
+        if img is None:
+            print(f"Warning: Could not read image {img_path}")
+            continue
+        
+        # Resize mask if dimensions don't match
+        if img.shape[:2] != binary_mask.shape[:2]:
+            resized_mask = cv2.resize(binary_mask, (img.shape[1], img.shape[0]))
+        else:
+            resized_mask = binary_mask
+        
+        # Apply mask
+        masked_img = cv2.bitwise_and(img, img, mask=resized_mask)
+        
+        # Find contours to get bounding box of mask
+        contours, _ = cv2.findContours(resized_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        if not contours:
+            print(f"No mask area found in {img_path} - skipping")
+            continue
+        
+        # Get bounding rectangle of largest contour
+        x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea))
+        
+        # Crop to masked region
+        cropped_img = masked_img[y:y+h, x:x+w]
+        
+        # Create output path (preserve original filename)
+        filename = os.path.basename(img_path)
+        output_path = os.path.join(output_folder, filename)
+        
+        # Save cropped image
+        cv2.imwrite(output_path, cropped_img)
+        # print(f"Processed and saved: {output_path}")
+    
+    print(f"\nProcessing complete! Saved {len(image_files)} cropped images to {output_folder}")

compose_video.py

@@ -0,0 +1,56 @@
+import os
+import cv2
+import numpy as np
+
+def compose_final_video(mask_path, frames_folder, extracted_frames_folder, output_path, fps=24):
+    """
+    Composes the final heatmap video by placing each heatmap frame
+    on top of its corresponding original frame using the table region defined by the mask.
+    """
+    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
+    if mask is None:
+        raise ValueError(f"Mask not found at {mask_path}")
+
+    _, binary_mask = cv2.threshold(mask, 200, 255, cv2.THRESH_BINARY)
+    height, width = binary_mask.shape
+
+    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        raise ValueError("No white area found in mask")
+
+    table_contour = max(contours, key=cv2.contourArea)
+    x, y, w, h = cv2.boundingRect(table_contour)
+
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    video_writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    heatmap_files = sorted([f for f in os.listdir(frames_folder) if f.endswith(('.bmp', '.jpg', '.png'))])
+    original_files = sorted([f for f in os.listdir(extracted_frames_folder) if f.endswith(('.bmp', '.jpg', '.png'))])
+
+    if len(heatmap_files) != len(original_files):
+        raise ValueError("Mismatch between number of heatmap frames and original frames")
+
+    for heatmap_file, original_file in zip(heatmap_files, original_files):
+        heatmap = cv2.imread(os.path.join(frames_folder, heatmap_file))
+        original = cv2.imread(os.path.join(extracted_frames_folder, original_file))
+
+        if heatmap is None or original is None:
+            continue
+
+        ph, pw = heatmap.shape[:2]
+        if pw > w or ph > h:
+            scale = min(w / pw, h / ph)
+            heatmap = cv2.resize(heatmap, (int(pw * scale), int(ph * scale)))
+            pw, ph = heatmap.shape[1], heatmap.shape[0]
+
+        x_offset = x + (w - pw) // 2
+        y_offset = y + (h - ph) // 2
+
+        result = original.copy()
+        result[y_offset:y_offset+ph, x_offset:x_offset+pw] = heatmap
+        cv2.drawContours(result, [table_contour], -1, (0, 255, 0), 2)
+
+        video_writer.write(result)
+
+    video_writer.release()
+    print("[Video] Final video saved to:", output_path)

extract_frames.py

@@ -0,0 +1,59 @@
+import cv2
+import os
+
+def video_to_keyframes(video_path, output_folder, frame_interval=1, prefix='frame'):
+    """
+    Convert video to keyframes (individual frames) and save as images
+    
+    Args:
+        video_path (str): Path to input video file
+        output_folder (str): Directory to save extracted frames
+        frame_interval (int): Extract every nth frame (1=every frame)
+        prefix (str): Prefix for saved frame files
+    """
+    # Create output directory if it doesn't exist
+    os.makedirs(output_folder, exist_ok=True)
+    
+    # Open the video file
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise ValueError(f"Could not open video {video_path}")
+    
+    # Get video properties
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    duration = total_frames / fps
+    
+    print(f"Video Info:")
+    print(f"- Path: {video_path}")
+    print(f"- FPS: {fps}")
+    print(f"- Total Frames: {total_frames}")
+    print(f"- Duration: {duration:.2f} seconds")
+    print(f"- Extracting every {frame_interval} frame(s)")
+    
+    frame_count = 0
+    saved_count = 0
+    
+    while True:
+        ret, frame = cap.read()
+        
+        # Stop if we've reached the end of the video
+        if not ret:
+            break
+            
+        # Only process frames at the specified interval
+        if frame_count % frame_interval == 0:
+            # Save frame as image file
+            frame_filename = f"{prefix}_{saved_count:06d}.jpg"
+            output_path = os.path.join(output_folder, frame_filename)
+            cv2.imwrite(output_path, frame)
+            saved_count += 1
+            
+            # Print progress every 100 frames
+            if saved_count % 100 == 0:
+                print(f"Saved frame {saved_count} (original frame {frame_count})")
+        
+        frame_count += 1
+    
+    # Release resources
+    cap.release()

run_gmm.py

@@ -0,0 +1,31 @@
+import os
+import cv2
+import glob
+from GMM import GMM
+
+def run_gmm_inference(input_dir, output_dir):
+    """
+    Applies GMM model inference on all frames in input_dir,
+    saves results with heatmap overlays to output_dir.
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    test_files = sorted(glob.glob(os.path.join(input_dir, '*.jpg')))
+
+    if not test_files:
+        raise ValueError("No input images found for GMM inference.")
+
+    heatmap = None
+    gmm = GMM.load_model("gmm_model.joblib")
+
+    for index, file in enumerate(test_files):
+        # print(f"[GMM] Processing: {file}")
+        img = cv2.imread(file)
+        if img is None:
+            print(f"Warning: Skipped unreadable file {file}")
+            continue
+
+        result_img, heatmap = gmm.infer(img, heatmap)
+        output_path = os.path.join(output_dir, f"{index:05d}.bmp")
+        cv2.imwrite(output_path, result_img)
+
+    print("[GMM] Inference complete. Output saved to:", output_dir)