extract_frames.py

import cv2
import os
import pathlib
import tkinter as tk
from tkinter import ttk
import threading
from multiprocessing import Manager, Pool
import subprocess
import sys


def get_video_properties(video_path):
    """Reads a video file and returns its properties or an error string."""
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        return (
            None,
            f"Could not open the file. It may be corrupt or an unsupported format.",
        )

    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()

    if fps <= 0:
        return None, "Could not determine a valid FPS (frames per second)."
    if frame_count <= 0:
        return None, "Video appears to be empty (zero frames)."

    return {"path": video_path, "fps": fps, "frame_count": frame_count}, None


def extract_frames(video_info, images_per_second, images_dir, progress_queue):
    """Extracts frames from a single video and reports progress via a queue."""
    # This function is now run in a separate process for each video.
    # We can add a message here to indicate which video is starting.
    progress_queue.put(f"Processing: {video_info['path'].name}...")

    cap = cv2.VideoCapture(str(video_info["path"]))
    if not cap.isOpened():
        print(f"Error: Could not open video {video_info['path'].name}")
        return

    frame_interval = video_info["fps"] / images_per_second

    frame_number = 0
    next_frame_to_capture = 0.0

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        if frame_number >= next_frame_to_capture:
            base_name = video_info["path"].stem
            output_filename = (
                images_dir / f"{base_name}_frame_{round(next_frame_to_capture)}.png"
            )
            cv2.imwrite(str(output_filename), frame)

            progress_queue.put(1)

            next_frame_to_capture += frame_interval

        frame_number += 1

    cap.release()


def main():
    """Main function to run the frame extraction process."""
    # Using pathlib for robust path handling
    base_dir = pathlib.Path(__file__).parent
    videos_dir = base_dir / "Videos"
    images_dir = base_dir / "Images"

    # Ensure output directory exists
    images_dir.mkdir(exist_ok=True)

    print("Analyzing videos in the 'Videos' folder...")
    video_files = [f for f in videos_dir.iterdir() if f.is_file()]

    if not video_files:
        print("No videos found in the 'Videos' folder. Exiting.")
        return

    # --- Pre-flight validation check ---
    print("Validating video files...")
    corrupted_files = {}
    video_properties = []
    for video_path in video_files:
        props, error = get_video_properties(video_path)
        if error:
            corrupted_files[video_path.name] = error
        else:
            video_properties.append(props)

    if corrupted_files:
        print("\n--- ERROR: Invalid Videos Detected ---")
        print("The script cannot continue. Please fix or remove the following files:")
        for filename, reason in corrupted_files.items():
            print(f"- {filename}: {reason}")
        return  # Exit before showing GUI

    if not video_properties:
        print("Could not find any valid videos to process. Exiting.")
        return

    # --- GUI for Interactive Input ---
    root = tk.Tk()
    root.title("Frame Extractor")
    root.geometry("400x250")

    chosen_rate = tk.DoubleVar(value=1.0)
    main_frame = ttk.Frame(root, padding="20")
    main_frame.pack(fill="both", expand=True)

    rate_label_var = tk.StringVar()
    ttk.Label(main_frame, textvariable=rate_label_var, font=("Helvetica", 10)).pack()

    total_images_var = tk.StringVar()
    ttk.Label(
        main_frame, textvariable=total_images_var, font=("Helvetica", 12, "bold")
    ).pack(pady=5)

    progress_queue = Manager().Queue()

    def calculate_total_images(rate):
        total_images = 0
        for props in video_properties:
            duration = props["frame_count"] / props["fps"]
            total_images += int(duration * rate)
        return total_images

    def update_display_labels(*args):
        rate = chosen_rate.get()
        if rate < 1:
            rate_text = f"1 frame every {1/rate:.1f} seconds"
        else:
            rate_text = f"{rate:.1f} frame{'s' if rate > 1 else ''} per second"
        rate_label_var.set(rate_text)
        total_images_var.set(
            f"Total Images to Extract: ~{calculate_total_images(rate)}"
        )

    slider = ttk.Scale(
        main_frame,
        from_=0.1,
        to=10.0,
        orient="horizontal",
        variable=chosen_rate,
        command=update_display_labels,
    )
    slider.pack(fill="x", pady=10)

    # --- Progress UI Elements (initially hidden) ---
    progress_var = tk.IntVar()
    progressbar = ttk.Progressbar(main_frame, variable=progress_var)
    status_var = tk.StringVar()
    status_label = ttk.Label(main_frame, textvariable=status_var)

    extracted_count = tk.IntVar()
    extracted_label_var = tk.StringVar()
    extracted_label = ttk.Label(main_frame, textvariable=extracted_label_var)

    def check_queue():
        try:
            msg = progress_queue.get_nowait()
            if isinstance(msg, str):
                if msg == "DONE":
                    status_var.set("Extraction Complete! Starting upload script...")

                    def proceed_to_upload():
                        root.destroy()

                        script_dir = pathlib.Path(__file__).parent
                        upload_script_path = script_dir / "upload_to_roboflow.py"
                        python_executable = sys.executable

                        if upload_script_path.exists():
                            print("\n--- Running Roboflow Upload Script ---")
                            subprocess.run(
                                [python_executable, str(upload_script_path)], check=True
                            )
                        else:
                            print(
                                "\n'upload_to_roboflow.py' not found. Skipping upload."
                            )

                    root.after(2000, proceed_to_upload)
                    return  # Stop checking
                else:
                    status_var.set(msg)
            else:
                extracted_count.set(extracted_count.get() + msg)
                progress_var.set(extracted_count.get())
                extracted_label_var.set(
                    f"Extracted: {extracted_count.get()} of ~{progressbar['maximum']}"
                )

        except Exception:  # Catches queue.Empty and other potential issues
            pass  # Continue checking
        root.after(100, check_queue)

    def manage_pool(tasks):
        """Manages the multiprocessing pool in a separate thread."""
        # Use all available CPU cores, or 1 if detection fails
        num_processes = os.cpu_count() or 1
        pool = Pool(processes=num_processes)

        # Dispatch all jobs to the pool
        pool.starmap(extract_frames, tasks)

        # Close the pool and wait for all processes to complete
        pool.close()
        pool.join()
        progress_queue.put("DONE")

    def start_extraction_process():
        rate = chosen_rate.get()

        # Reconfigure UI for processing
        slider.pack_forget()
        start_button.pack_forget()

        total_to_extract = calculate_total_images(rate)
        progressbar["maximum"] = total_to_extract

        progressbar.pack(fill="x", pady=10)
        status_label.pack()
        extracted_label.pack()
        extracted_label_var.set(f"Extracted: 0 of ~{total_to_extract}")

        # Prepare arguments for each worker process
        tasks = [
            (props, rate, images_dir, progress_queue) for props in video_properties
        ]

        # Start the pool manager in a separate thread so the GUI doesn't freeze
        threading.Thread(target=manage_pool, args=(tasks,), daemon=True).start()

        root.after(100, check_queue)

    start_button = ttk.Button(
        main_frame, text="Start Extraction", command=start_extraction_process
    )
    start_button.pack(pady=15)

    update_display_labels()
    root.mainloop()


if __name__ == "__main__":
    # On Windows, multiprocessing requires this protection
    main()