HadesStarClassifier/label_v2.py

import glob
import os
import random
import cv2

# --- CONFIGURATION ---
OUTPUT_DIRECTORY = "dataset"
FRAMES_TO_EXTRACT = 20  # Number of random frames to pull
SCALE_FACTOR = 0.5  # Resize video to 1/2 size for easier viewing

# Define the EXACT size of the cropped images you want to save
BOX_WIDTH = 64
BOX_HEIGHT = 64

# Defined classification options
CLASSIFICATION_OPTIONS = ["guardian", "interceptor", "sentinel", "rocket", "bomber", "background", "asteroid"]
# ---------------------


def get_dataset_counts():
    """Counts the number of existing images in each classification folder."""
    counts = {}
    for option in CLASSIFICATION_OPTIONS:
        folder_path = os.path.join(OUTPUT_DIRECTORY, option)
        if os.path.exists(folder_path):
            # Count only common image files
            files = [
                f
                for f in os.listdir(folder_path)
                if f.lower().endswith((".jpg", ".jpeg", ".png"))
            ]
            counts[option] = len(files)
        else:
            counts[option] = 0
    return counts


def mouse_click_callback(event, x, y, flags, param):
    """Handles mouse clicks.

    Centers a fixed-size box on the click, crops it, and saves it.
    """
    if event == cv2.EVENT_LBUTTONDOWN:
        frame = param["frame"]
        display_frame = param["display_frame"]
        label_dir = param["label_dir"]
        frame_idx = param["frame_idx"]
        window_name = param["window_name"]

        # Calculate top-left corner based on centering the box on the click
        x1 = x - BOX_WIDTH // 2
        y1 = y - BOX_HEIGHT // 2
        x2 = x1 + BOX_WIDTH
        y2 = y1 + BOX_HEIGHT

        # --- Edge Case Handling ---
        if x1 < 0:
            x1 = 0
            x2 = BOX_WIDTH
        if y1 < 0:
            y1 = 0
            y2 = BOX_HEIGHT
        if x2 > frame.shape[1]:
            x2 = frame.shape[1]
            x1 = x2 - BOX_WIDTH
        if y2 > frame.shape[0]:
            y2 = frame.shape[0]
            y1 = y2 - BOX_HEIGHT

        # Extract the crop from the scaled frame
        crop = frame[y1:y2, x1:x2]

        # Save the crop
        count = param["counter"]
        filename = f"frame_{frame_idx}_crop_{count}.jpg"
        save_path = os.path.join(label_dir, filename)
        cv2.imwrite(save_path, crop)

        print(f" Saved crop {count} to {save_path}")
        param["counter"] += 1

        # Draw visual feedback (a green box) on the screen where you clicked
        cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.imshow(window_name, display_frame)
        cv2.moveWindow(window_name, 0, 0)


def label_video_session():
    print("=========================================")
    print("       AI DATASET SESSION LABELER        ")
    print("=========================================\n")

    # 1. Automatically find and select a random MP4 video
    mp4_files = glob.glob("*.mp4")
    if not mp4_files:
        print(
            "Error: No .mp4 files found in the current directory.\n"
            "Please place this script in the same folder as your video files."
        )
        return

    selected_video = random.choice(mp4_files)
    print(f"-> Selected Video File: '{selected_video}'")

    # 2. Fetch and display current dataset balance dashboard
    counts = get_dataset_counts()
    print("\n-----------------------------------------")
    print("      CURRENT DATASET BALANCE STATUS     ")
    print("-----------------------------------------")
    for option, count in counts.items():
        print(f"  • {option.capitalize():<12} : {count} items")
    print("-----------------------------------------")
    print(
        "💡 Tip: Try to choose options with lower counts to keep data balanced!"
    )

    # 3. Force selection of a valid label
    print("\nAvailable Classification Targets:")
    for idx, option in enumerate(CLASSIFICATION_OPTIONS, start=1):
        print(f"  [{idx}] {option} (Current: {counts[option]})")

    while True:
        try:
            choice = input(
                f"\nEnter the number (1-{len(CLASSIFICATION_OPTIONS)}) of what you are labeling: "
            ).strip()
            choice_idx = int(choice) - 1
            if 0 <= choice_idx < len(CLASSIFICATION_OPTIONS):
                session_label = CLASSIFICATION_OPTIONS[choice_idx]
                break
            else:
                print("Invalid selection. Please choose a number from the list.")
        except ValueError:
            print("Invalid input. Please enter a number.")

    # Open video and check validity
    cap = cv2.VideoCapture(selected_video)
    if not cap.isOpened():
        print(f"Error: Could not open video file {selected_video}")
        return

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if total_frames == 0:
        print("Error: Video has 0 frames.")
        return

    # Select random unique frames
    num_frames = min(FRAMES_TO_EXTRACT, total_frames)
    frame_indices = random.sample(range(total_frames), num_frames)
    frame_indices.sort()

    # Set up directories
    label_dir = os.path.join(OUTPUT_DIRECTORY, session_label)
    os.makedirs(label_dir, exist_ok=True)

    # Comprehensive User Instructions
    print("\n" + "=" * 50)
    print("              HOW TO LABEL")
    print("=" * 50)
    print(f"• TARGET OBJECT: {session_label.upper()}")
    print(f"• SESSION SCOPE: {num_frames} random frames from '{selected_video}'")
    print("• CROPPING SIZE: 64x64 pixels (centered automatically around your click)")
    print("-" * 50)
    print("► CONTROLS:")
    print("  [Left-Click]     - Click directly on the target object to save a snapshot.")
    print("                     (You can click multiple objects if they appear on screen)")
    print("  [SPACE or ENTER] - Advance to the next random video frame.")
    print("  [Q Key]          - Quit and save all progress up to this point.")
    print("=" * 50)
    input("\nPress ENTER when you are ready to begin...")

    for i, frame_idx in enumerate(frame_indices):
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
        ret, frame = cap.read()
        if not ret:
            continue

        # Resize the frame by 1/2 for laptop viewability
        frame_resized = cv2.resize(
            frame, (0, 0), fx=SCALE_FACTOR, fy=SCALE_FACTOR
        )

        # Create a copy for drawing the green boxes dynamically
        frame_display = frame_resized.copy()

        window_name = (
            f"Labeling: {session_label} | Frame {i+1}/{num_frames}"
        )
        cv2.namedWindow(window_name)

        # Dictionary to pass state/variables into the mouse callback function
        session_state = {
            "frame": frame_resized,
            "display_frame": frame_display,
            "label_dir": label_dir,
            "frame_idx": frame_idx,
            "window_name": window_name,
            "counter": 0,
        }

        # Bind the mouse click event to the window
        cv2.setMouseCallback(window_name, mouse_click_callback, session_state)

        # Keep window open until user hits Space/Enter or 'q'
        cv2.imshow(window_name, frame_display)
        while True:
            cv2.moveWindow(window_name, 0, 0)
            key = cv2.waitKey(1) & 0xFF
            if key == 13 or key == 32:
                break
            elif key == ord("q") or key == ord("Q"):
                print("\nQuitting session early... Saving progress.")
                cap.release()
                cv2.destroyAllWindows()
                return

        cv2.destroyWindow(window_name)

    cap.release()
    cv2.destroyAllWindows()
    print(
        f"\nSession finished! All crops saved successfully in '{label_dir}'."
    )


if __name__ == "__main__":
    label_video_session()