From cc53baf7c539c14122a00968a2f416f2dd8c50a8 Mon Sep 17 00:00:00 2001 From: bionickatana Date: Wed, 15 Jul 2026 12:15:40 -0600 Subject: [PATCH] Updated labeling script --- label_v4.py | 349 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 349 insertions(+) create mode 100644 label_v4.py diff --git a/label_v4.py b/label_v4.py new file mode 100644 index 0000000..6154de7 --- /dev/null +++ b/label_v4.py @@ -0,0 +1,349 @@ +# Label_v4.py +# --------------------- + + +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", +] +# --------------------- + +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 click-and-drag bounding box creation. + + Calculates dynamic YOLO coordinates, saves the frame, and appends the + new dynamic coordinates to the YOLO annotation file. + """ + if param["session_label"] == "background": + if event == cv2.EVENT_LBUTTONDOWN: + 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"] + + # --- 1. Start drawing on Left Mouse Click Down --- + if event == cv2.EVENT_LBUTTONDOWN: + param["drawing"] = True + param["ix"], param["iy"] = x, y + + # --- 2. Dynamic Box Preview during Mouse Drag --- + elif event == cv2.EVENT_MOUSEMOVE: + if param["drawing"]: + # Copy display_frame to wipe previous frame's dynamic preview lines + temp_frame = display_frame.copy() + cv2.rectangle(temp_frame, (param["ix"], param["iy"]), (x, y), (0, 255, 0), 2) + cv2.imshow(window_name, temp_frame) + + # --- 3. Save Coordinates on Mouse Release --- + elif event == cv2.EVENT_LBUTTONUP: + param["drawing"] = False + ix, iy = param["ix"], param["iy"] + + # Handle negative dragging (drawing from bottom-right to top-left) + x1, x2 = min(ix, x), max(ix, x) + y1, y2 = min(iy, y), max(iy, y) + + # Skip accidental tiny clicks (smaller than 5x5 pixels) + width = x2 - x1 + height = y2 - y1 + if width < 5 or height < 5: + cv2.imshow(window_name, display_frame) + return + + # Ensure box corners do not exceed image boundaries + img_h, img_w = frame.shape[0], frame.shape[1] + x1 = max(0, min(x1, img_w - 1)) + y1 = max(0, min(y1, img_h - 1)) + x2 = max(0, min(x2, img_w - 1)) + y2 = max(0, min(y2, img_h - 1)) + + # Recalculate dimensions after boundary clipping + width = x2 - x1 + height = y2 - y1 + + # --- YOLO Normalized Coordinate Conversion --- + 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 = width / img_w + height_norm = 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 + 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 permanent visual feedback box on the master display frame + cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText( + display_frame, + param["session_label"], + (x1, max(15, y1 - 5)), # Prevents text from drawing off-screen at the top + 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("-" * 50) + print("► CONTROLS:") + print(" [Click & Drag] - Draw a custom box around the target (auto-saves frame on release).") + print(" [Right Arrow] - Move FORWARD 10 frames (sequential tracking, no save).") + print(" [Left Arrow] - Move BACKWARD 10 frames (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 SCALE_FACTOR 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) + + # Expanded state to hold drawing tracks + 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, + "drawing": False, # True when dragging mouse + "ix": -1, # Starting X + "iy": -1, # Starting Y + } + + 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 + 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 + 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()