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()