commit 816c36d2871d208745f0919253de7150df3bfd45 Author: bionickatana Date: Sat Jul 11 10:12:36 2026 -0600 Setting up diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..ccd5125 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,2 @@ +dataset/** filter=lfs diff=lfs merge=lfs -text +*.mp4 filter=lfs diff=lfs merge=lfs -text diff --git a/label_v3.py b/label_v3.py new file mode 100644 index 0000000..14abc50 --- /dev/null +++ b/label_v3.py @@ -0,0 +1,327 @@ +import glob +import os +import random +import cv2 + +# --- CONFIGURATION --- +OUTPUT_DIRECTORY = "dataset" +FRAMES_TO_EXTRACT = 20 # Number of random anchor frames to find targets +SCALE_FACTOR = 0.5 # Resize video to 1/2 size for easier viewing + +# Define the EXACT size of the bounding box around your click point +BOX_WIDTH = 64 +BOX_HEIGHT = 64 + +# Defined classification options (YOLO maps these to IDs: 0, 1, 2, 3...) +CLASSIFICATION_OPTIONS = [ + "guardian", + "interceptor", + "sentinel", + "rocket", + "bomber", + "background", + "asteroid", +] +# --------------------- + +def get_dataset_counts(): + """Counts the number of existing bounding boxes by parsing YOLO text files.""" + counts = {option: 0 for option in CLASSIFICATION_OPTIONS} + labels_dir = os.path.join(OUTPUT_DIRECTORY, "labels") + + if os.path.exists(labels_dir): + # Scan all text annotation files in the YOLO labels directory + for txt_file in glob.glob(os.path.join(labels_dir, "*.txt")): + try: + with open(txt_file, "r") as f: + for line in f: + parts = line.strip().split() + if parts: + class_id = int(parts[0]) + if 0 <= class_id < len(CLASSIFICATION_OPTIONS): + option = CLASSIFICATION_OPTIONS[class_id] + counts[option] += 1 + except Exception: + pass # Skip corrupted or unreadable text files safely + return counts + +def mouse_click_callback(event, x, y, flags, param): + """Handles mouse clicks. + + Calculates YOLO coordinates, saves the full frame, and appends a line to + the text file. + """ + if event == cv2.EVENT_LBUTTONDOWN: + # Ignore clicks if the user is in "background" capture mode + if param["session_label"] == "background": + print( + " Info: In background mode, do not click targets. Just press SPACE to save frame." + ) + return + + frame = param["frame"] + display_frame = param["display_frame"] + frame_idx = param["frame_idx"] + window_name = param["window_name"] + + # Calculate top-left and bottom-right corners + 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 + + # --- YOLO Normalized Coordinate Conversion --- + img_h, img_w = frame.shape[0], frame.shape[1] + + true_center_x = (x1 + x2) / 2.0 + true_center_y = (y1 + y2) / 2.0 + + x_center_norm = true_center_x / img_w + y_center_norm = true_center_y / img_h + width_norm = BOX_WIDTH / img_w + height_norm = BOX_HEIGHT / img_h + + # Clean video name to ensure a unique, safe filename + video_base = os.path.splitext(os.path.basename(param["video_name"]))[0] + filename_base = f"{video_base}_frame_{frame_idx}" + + image_save_path = os.path.join( + param["images_dir"], f"{filename_base}.jpg" + ) + label_save_path = os.path.join( + param["labels_dir"], f"{filename_base}.txt" + ) + + # 1. Save the full frame image (overwriting is fine if clicking multiple items) + cv2.imwrite(image_save_path, frame) + + # 2. Append the target annotation to the YOLO text file + class_id = param["class_id"] + with open(label_save_path, "a") as f: + f.write( + f"{class_id} {x_center_norm:.6f} {y_center_norm:.6f} {width_norm:.6f} {height_norm:.6f}\n" + ) + + print( + f" Logged {param['session_label']} box coordinate to {label_save_path}" + ) + param["box_drawn"] = True + + # Draw visual feedback box and short class label on screen + cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText( + display_frame, + param["session_label"], + (x1, y1 - 5), + cv2.FONT_HERSHEY_SIMPLEX, + 0.4, + (0, 255, 0), + 1, + ) + cv2.imshow(window_name, display_frame) + +def label_video_session(): + print("=========================================") + print(" AI YOLO DATASET 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} boxes logged") + 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 Boxes: {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] + session_class_id = 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 anchor frames to jump between + num_frames = min(FRAMES_TO_EXTRACT, total_frames) + frame_indices = random.sample(range(total_frames), num_frames) + frame_indices.sort() + + # Set up standard YOLO dataset directory branches + images_dir = os.path.join(OUTPUT_DIRECTORY, "images") + labels_dir = os.path.join(OUTPUT_DIRECTORY, "labels") + os.makedirs(images_dir, exist_ok=True) + os.makedirs(labels_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 search zones in '{selected_video}'") + print(f"• BOUNDING BOX : {BOX_WIDTH}x{BOX_HEIGHT} pixels scaled to YOLO layout") + print("-" * 50) + print("► CONTROLS:") + print(" [Left-Click] - Click directly on target to tag position (auto-saves frame).") + print(" [Right Arrow] - Move FORWARD 1 individual frame (sequential tracking, no save).") + print(" [Left Arrow] - Move BACKWARD 1 individual frame (sequential tracking, no save).") + print(" [SPACE / ENTER] - Jump completely to the NEXT random video search spot.") + print(" [Q Key] - Quit and save all progress up to this point.") + print("=" * 50) + input("\nPress ENTER when you are ready to begin...") + + # Cross-platform arrow key mappings for waitKeyEx() + LEFT_KEYS = [81, 2, 2424832, 65361, 63234] + RIGHT_KEYS = [83, 3, 2490368, 65363, 63235] + + i = 0 + current_frame = frame_indices[i] + + while i < len(frame_indices): + cap.set(cv2.CAP_PROP_POS_FRAMES, current_frame) + ret, frame = cap.read() + if not ret: + print(f" Warning: Could not read frame {current_frame}. Moving to next random anchor.") + i += 1 + if i < len(frame_indices): + current_frame = frame_indices[i] + continue + + # Resize the frame by 1/2 for laptop viewability + frame_resized = cv2.resize( + frame, (0, 0), fx=SCALE_FACTOR, fy=SCALE_FACTOR + ) + frame_display = frame_resized.copy() + + window_name = f"Labeling: {session_label} | Spot {i+1}/{num_frames} (Frame {current_frame})" + cv2.namedWindow(window_name) + + session_state = { + "frame": frame_resized, + "display_frame": frame_display, + "images_dir": images_dir, + "labels_dir": labels_dir, + "frame_idx": current_frame, + "window_name": window_name, + "video_name": selected_video, + "class_id": session_class_id, + "session_label": session_label, + "box_drawn": False, + } + + cv2.setMouseCallback(window_name, mouse_click_callback, session_state) + + # Keep window open until user acts + cv2.imshow(window_name, frame_display) + while True: + cv2.moveWindow(window_name, 0, 0) + key = cv2.waitKeyEx(1) + + # Enter (13) or Space (32) -> Jump to the next random anchor spot + if key == 13 or key == 32: + if session_label == "background": + video_base = os.path.splitext( + os.path.basename(selected_video) + )[0] + filename_base = f"{video_base}_frame_{current_frame}" + image_save_path = os.path.join( + images_dir, f"{filename_base}.jpg" + ) + label_save_path = os.path.join( + labels_dir, f"{filename_base}.txt" + ) + + cv2.imwrite(image_save_path, frame_resized) + open(label_save_path, "a").close() # Creates empty file + print(f" Registered negative background frame: {filename_base}") + + i += 1 + if i < len(frame_indices): + current_frame = frame_indices[i] + break + + # Right Arrow -> Move forward exactly ONE individual sequential frame + elif key in RIGHT_KEYS: + if current_frame < total_frames - 10: + current_frame += 10 + else: + print(" Already at the very last frame of the video!") + continue + break + + # Left Arrow -> Move backward exactly ONE individual sequential frame + elif key in LEFT_KEYS: + if current_frame > 0: + current_frame -= 10 + else: + print(" Already at the first frame of the video!") + continue + break + + # Quit + elif key == ord("q") or key == ord("Q") or (key & 0xFF) in [ord("q"), 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! Dataset updated successfully inside '{OUTPUT_DIRECTORY}'.") + +if __name__ == "__main__": + label_video_session() diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..9bbe9e5 --- /dev/null +++ b/readme.md @@ -0,0 +1,16 @@ +# Install pip packages + +Probably in a virtual env... *only the opencv packages are needed.* Others are +optional. + +`pip install -r requirements.txt` + +# Run script: + +`python3 label_v2.py` + + +Try to click only on ones that are moving. Doing some while they are just +resting is okay but mostly I need the ones that are moving around on the +screen. Also try to only include one or if you include two, click on each +seperately.