commit 70fa8bfab7b0961a0a2d01e0340b200d4276e849 Author: bionickatana Date: Fri Jul 10 14:33:55 2026 -0600 Uploading classification script diff --git a/Screen_Recording_20260710_084027_Hades' Star_1.mp4 b/Screen_Recording_20260710_084027_Hades' Star_1.mp4 new file mode 100644 index 0000000..d3ecf7c Binary files /dev/null and b/Screen_Recording_20260710_084027_Hades' Star_1.mp4 differ diff --git a/Screen_Recording_20260710_101822_Hades' Star.mp4 b/Screen_Recording_20260710_101822_Hades' Star.mp4 new file mode 100644 index 0000000..88269a2 Binary files /dev/null and b/Screen_Recording_20260710_101822_Hades' Star.mp4 differ diff --git a/Screen_Recording_20260710_105358_Hades' Star.mp4 b/Screen_Recording_20260710_105358_Hades' Star.mp4 new file mode 100644 index 0000000..3c3e754 Binary files /dev/null and b/Screen_Recording_20260710_105358_Hades' Star.mp4 differ diff --git a/Screen_Recording_20260710_110910_Hades' Star.mp4 b/Screen_Recording_20260710_110910_Hades' Star.mp4 new file mode 100644 index 0000000..26ec27b Binary files /dev/null and b/Screen_Recording_20260710_110910_Hades' Star.mp4 differ diff --git a/Screen_Recording_20260710_141945_Hades' Star.mp4 b/Screen_Recording_20260710_141945_Hades' Star.mp4 new file mode 100644 index 0000000..2e83eee Binary files /dev/null and b/Screen_Recording_20260710_141945_Hades' Star.mp4 differ diff --git a/Screen_Recording_20260710_142109_Hades' Star.mp4 b/Screen_Recording_20260710_142109_Hades' Star.mp4 new file mode 100644 index 0000000..05b5595 Binary files /dev/null and b/Screen_Recording_20260710_142109_Hades' Star.mp4 differ diff --git a/Screen_Recording_20260710_142243_Hades' Star.mp4 b/Screen_Recording_20260710_142243_Hades' Star.mp4 new file mode 100644 index 0000000..c2a0a0d Binary files /dev/null and b/Screen_Recording_20260710_142243_Hades' Star.mp4 differ diff --git a/label_v2.py b/label_v2.py new file mode 100644 index 0000000..d88fe85 --- /dev/null +++ b/label_v2.py @@ -0,0 +1,227 @@ +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()