# Label_v5.py # --------------------- # Upgraded with Dedicated Header Bar, Dynamic View Scaling for Small Crops, # Enforced Minimum Window Sizes, Interactive Zoom, and Multi-Class Tagging. import glob import os import random import cv2 # --- CONFIGURATION --- OUTPUT_DIRECTORY = "dataset_v5" FRAMES_TO_EXTRACT = 20 # Number of random anchor frames to find targets SCALE_FACTOR = 0.5 # Resize initial video to 1/2 size for easier viewing RANDOM_CROP_SIZE_RATIO = 0.5 # Random crop takes 50% width x 50% height of the frame # UI Display Constraints (Prevents tiny windows & makes small crops easy to see) MIN_WINDOW_WIDTH = 800 MIN_WINDOW_HEIGHT = 600 MAX_VIEW_SCALE = 10.0 # Maximum zoom multiplier for visual display HUD_HEIGHT = 40 # Height of the dedicated header bar above the image # Defined classification options (YOLO maps these to IDs: 0, 1, 2, 3...) CLASSIFICATION_OPTIONS = [ "guardian", "interceptor", "sentinel", "rocket", "bomber", ] # Distinct BGR color palette for each class (Cyan, Magenta, Yellow, Red, Green, Orange, Purple) CLASS_COLORS = [ (255, 255, 0), # 0: Guardian (Cyan) (255, 0, 255), # 1: Interceptor (Magenta) (0, 255, 255), # 2: Sentinel (Yellow) (0, 0, 255), # 3: Rocket (Bright Red) (0, 255, 0), # 4: Bomber (Bright Green) (0, 165, 255), # 5: Orange (255, 0, 128), # 6: Purple ] # --------------------- 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): 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 get_current_filename_base(param): """Generates a unique filename base, appending crop coordinates if zoomed in.""" video_base = os.path.splitext(os.path.basename(param["video_name"]))[0] base = f"{video_base}_frame_{param['frame_idx']}" if param["crop_box"] is not None: cx1, cy1, cx2, cy2 = param["crop_box"] base += f"_crop_{cx1}_{cy1}_{cx2}_{cy2}" return base def load_existing_annotations(param): """Loads existing YOLO annotations from disk for the current frame/crop view.""" filename_base = get_current_filename_base(param) label_path = os.path.join(param["labels_dir"], f"{filename_base}.txt") param["annotations"] = [] if os.path.exists(label_path): img_h, img_w = param["frame"].shape[0], param["frame"].shape[1] try: with open(label_path, "r") as f: for line in f: parts = line.strip().split() if len(parts) >= 5: class_id = int(parts[0]) x_center_norm = float(parts[1]) y_center_norm = float(parts[2]) width_norm = float(parts[3]) height_norm = float(parts[4]) width = width_norm * img_w height = height_norm * img_h true_center_x = x_center_norm * img_w true_center_y = y_center_norm * img_h x1 = int(true_center_x - width / 2.0) y1 = int(true_center_y - height / 2.0) x2 = int(true_center_x + width / 2.0) y2 = int(true_center_y + height / 2.0) label = ( CLASSIFICATION_OPTIONS[class_id] if 0 <= class_id < len(CLASSIFICATION_OPTIONS) else f"ID {class_id}" ) param["annotations"].append((class_id, label, x1, y1, x2, y2)) except Exception as e: print(f" Could not load existing annotations: {e}") def apply_crop(param, crop_coords=None): """Updates the working frame to either a cropped sub-region or the full screenshot.""" if crop_coords is None: param["crop_box"] = None param["frame"] = param["full_frame"].copy() else: cx1, cy1, cx2, cy2 = crop_coords h, w = param["full_frame"].shape[:2] cx1, cx2 = max(0, min(cx1, w - 1)), max(0, min(cx2, w)) cy1, cy2 = max(0, min(cy1, h - 1)), max(0, min(cy2, h)) if (cx2 - cx1) < 15 or (cy2 - cy1) < 15: print(" Crop area too small! Aborting zoom.") return param["crop_box"] = (cx1, cy1, cx2, cy2) param["frame"] = param["full_frame"][cy1:cy2, cx1:cx2].copy() param["drawing"] = False param["zooming"] = False load_existing_annotations(param) redraw_frame(param) def win_to_frame_coords(win_x, win_y, param): """Converts UI window mouse coordinates back to native working frame coordinates.""" # Subtract the dedicated header bar height from y adj_y = max(0, win_y - HUD_HEIGHT) scale = param.get("view_scale", 1.0) frame_x = int(win_x / scale) frame_y = int(adj_y / scale) # Clamp to actual working frame dimensions h, w = param["frame"].shape[:2] frame_x = max(0, min(frame_x, w - 1)) frame_y = max(0, min(frame_y, h - 1)) return frame_x, frame_y def draw_hud(frame, param, view_scale): """Draws an informative status banner in the dedicated top header area.""" color = CLASS_COLORS[param["class_id"] % len(CLASS_COLORS)] mode_str = "ALL-IN-ONE" if param["multi_mode"] else "SINGLE-CLASS" zoom_str = "FULL" if param["crop_box"] is None else f"CROP ({view_scale:.1f}x VIEW)" label_str = f"[{zoom_str}] | ACTIVE: [{param['class_id']}] {param['session_label'].upper()} | [SPACE] Class | [R] Rand Crop | [X] Reset Zoom" # Active color indicator block on the left edge of the banner cv2.rectangle(frame, (0, 0), (14, HUD_HEIGHT), color, -1) # Draw a subtle separator line between HUD and image cv2.line(frame, (0, HUD_HEIGHT - 1), (frame.shape[1], HUD_HEIGHT - 1), (70, 70, 70), 1) # HUD Text centered vertically in the 40px banner cv2.putText( frame, label_str, (24, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1, cv2.LINE_AA, ) def redraw_frame(param, display_tags=True): """Regenerates the display frame using dynamic view scaling and transparent box overlays.""" h_frame, w_frame = param["frame"].shape[:2] # --- DYNAMIC VIEW SCALING --- # Calculate scale needed to meet minimum comfortable window viewing size scale_w = MIN_WINDOW_WIDTH / float(w_frame) scale_h = MIN_WINDOW_HEIGHT / float(h_frame) view_scale = max(1.0, min(scale_w, scale_h)) view_scale = min(MAX_VIEW_SCALE, view_scale) # Cap extreme zooms param["view_scale"] = view_scale # Scale up the working frame ONLY for display purposes disp_w = int(w_frame * view_scale) disp_h = int(h_frame * view_scale) display_img = cv2.resize(param["frame"], (disp_w, disp_h), interpolation=cv2.INTER_LINEAR) # --- OVERLAY LAYERS --- fill_overlay = display_img.copy() border_overlay = display_img.copy() text_frame = display_img.copy() for ann in param["annotations"]: class_id, label, x1, y1, x2, y2 = ann color = CLASS_COLORS[class_id % len(CLASS_COLORS)] # Map working frame box coordinates to the scaled display dimensions dx1, dy1 = int(x1 * view_scale), int(y1 * view_scale) dx2, dy2 = int(x2 * view_scale), int(y2 * view_scale) cv2.rectangle(fill_overlay, (dx1, dy1), (dx2, dy2), color, -1) cv2.rectangle(border_overlay, (dx1, dy1), (dx2, dy2), color, 3) text_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.45, 1) bg_width = text_size[0] + 6 bg_height = text_size[1] + 6 cv2.rectangle( text_frame, (dx1, max(0, dy1 - bg_height)), (dx1 + bg_width, max(0, dy1)), color, -1, ) cv2.putText( text_frame, label, (dx1 + 3, max(12, dy1 - 4)), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 0), 1, cv2.LINE_AA, ) if display_tags: cv2.addWeighted(fill_overlay, 0.15, display_img, 0.85, 0, display_img) cv2.addWeighted(border_overlay, 0.55, display_img, 0.45, 0, display_img) cv2.addWeighted(text_frame, 0.25, display_img, 0.75, 0, display_img) # --- DEDICATED HEADER BAR --- # Add top border padding for the HUD banner so it NEVER covers the image final_display = cv2.copyMakeBorder( display_img, HUD_HEIGHT, 0, 0, 0, cv2.BORDER_CONSTANT, value=(30, 30, 30), ) # Draw status HUD in the top padded area draw_hud(final_display, param, view_scale) param["display_frame"] = final_display cv2.imshow(param["window_name"], final_display) def save_annotations_to_file(param): """Overwrites the YOLO text file and saves the un-scaled image patch to disk.""" filename_base = get_current_filename_base(param) 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") if len(param["annotations"]) == 0: if os.path.exists(label_save_path): os.remove(label_save_path) if os.path.exists(image_save_path): os.remove(image_save_path) print(f" Removed empty annotation files for {filename_base}") param["box_drawn"] = False else: # Save the ORIGINAL un-scaled working frame (keeps dataset images small & standard) cv2.imwrite(image_save_path, param["frame"]) # Write YOLO coordinates normalized against the working frame dimensions img_h, img_w = param["frame"].shape[0], param["frame"].shape[1] with open(label_save_path, "w") as f: for ann in param["annotations"]: class_id, _, x1, y1, x2, y2 = ann true_center_x = (x1 + x2) / 2.0 true_center_y = (y1 + y2) / 2.0 width = x2 - x1 height = y2 - y1 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 f.write( f"{class_id} {x_center_norm:.6f} {y_center_norm:.6f} {width_norm:.6f} {height_norm:.6f}\n" ) print(f" Updated dataset files for {filename_base} ({len(param['annotations'])} boxes | Saved at {img_w}x{img_h}px)") param["box_drawn"] = True def mouse_click_callback(event, x, y, flags, param): """Handles labeling drags, zoom dragging, and undo actions in native UI window coordinates.""" if param["session_label"] == "background": if event == cv2.EVENT_LBUTTONDOWN: print(" Info: In background mode, do not click targets. Just press ENTER to save frame.") return display_frame = param["display_frame"] window_name = param["window_name"] # --- Right Click: Undo Last Box --- if event == cv2.EVENT_RBUTTONDOWN: if param["annotations"]: removed = param["annotations"].pop() print(f" Undid last box ([{removed[0]}] {removed[1]})") save_annotations_to_file(param) redraw_frame(param) else: print(" No boxes left to undo on this frame!") return # --- Middle Click OR Shift + Left Click: Start Zoom/Crop Box --- is_shift_click = (event == cv2.EVENT_LBUTTONDOWN) and (flags & cv2.EVENT_FLAG_SHIFTKEY) if event == cv2.EVENT_MBUTTONDOWN or is_shift_click: param["zooming"] = True param["win_ix"], param["win_iy"] = x, y return # --- Standard Left Click Down: Start Label Box Dragging --- if event == cv2.EVENT_LBUTTONDOWN: param["drawing"] = True param["win_ix"], param["win_iy"] = x, y # --- Mouse Move: Live Preview in UI Space --- elif event == cv2.EVENT_MOUSEMOVE: if param["drawing"] or param["zooming"]: temp_frame = display_frame.copy() overlay = temp_frame.copy() color = (255, 255, 255) if param["zooming"] else CLASS_COLORS[param["class_id"] % len(CLASS_COLORS)] cv2.rectangle(overlay, (param["win_ix"], param["win_iy"]), (x, y), color, 2 if param["zooming"] else 3) cv2.addWeighted(overlay, 0.55, temp_frame, 0.45, 0, temp_frame) cv2.imshow(window_name, temp_frame) # --- Middle Click OR Shift + Left Click Release: Apply Zoom Crop --- elif event == cv2.EVENT_MBUTTONUP or (event == cv2.EVENT_LBUTTONUP and param["zooming"]): param["zooming"] = False # Convert UI window coordinates back to native working frame coordinates x1, y1 = win_to_frame_coords(param["win_ix"], param["win_iy"], param) x2, y2 = win_to_frame_coords(x, y, param) x1, x2 = min(x1, x2), max(x1, x2) y1, y2 = min(y1, y2), max(y1, y2) if (x2 - x1) < 10 or (y2 - y1) < 10: cv2.imshow(window_name, display_frame) return # If already cropped, translate local coordinates back to full_frame coordinates if param["crop_box"] is not None: offset_x, offset_y = param["crop_box"][0], param["crop_box"][1] x1, x2 = x1 + offset_x, x2 + offset_x y1, y2 = y1 + offset_y, y2 + offset_y print(f" Zooming into region: X({x1}->{x2}), Y({y1}->{y2})") apply_crop(param, (x1, y1, x2, y2)) # --- Left Click Release: Lock Label Box & Save --- elif event == cv2.EVENT_LBUTTONUP and param["drawing"]: param["drawing"] = False # Ignore tiny accidental micro-clicks in UI window space if abs(x - param["win_ix"]) < 5 or abs(y - param["win_iy"]) < 5: cv2.imshow(window_name, display_frame) return # Convert UI window coordinates back to native working frame coordinates x1, y1 = win_to_frame_coords(param["win_ix"], param["win_iy"], param) x2, y2 = win_to_frame_coords(x, y, param) x1, x2 = min(x1, x2), max(x1, x2) y1, y2 = min(y1, y2), max(y1, y2) param["annotations"].append( (param["class_id"], param["session_label"], x1, y1, x2, y2) ) save_annotations_to_file(param) redraw_frame(param) def label_video_session(): print("=========================================") print(" AI YOLO DATASET LABELER v5.0 ") print("=========================================\n") 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}'") 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("\nAvailable Classification Modes & Targets:") print(" [0] ALL-IN-ONE MODE (Tag multiple classes per frame; cycle with SPACE)") for idx, option in enumerate(CLASSIFICATION_OPTIONS, start=1): print(f" [{idx}] {option} only (Current Boxes: {counts[option]})") while True: try: choice = input( f"\nEnter mode selection (0-{len(CLASSIFICATION_OPTIONS)}): " ).strip() choice_idx = int(choice) if choice_idx == 0: multi_mode = True session_class_id = 0 session_label = CLASSIFICATION_OPTIONS[0] break elif 1 <= choice_idx <= len(CLASSIFICATION_OPTIONS): multi_mode = False session_class_id = choice_idx - 1 session_label = CLASSIFICATION_OPTIONS[session_class_id] break else: print("Invalid selection. Please choose a number from the list.") except ValueError: print("Invalid input. Please enter a valid number.") 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 num_frames = min(FRAMES_TO_EXTRACT, total_frames) frame_indices = random.sample(range(total_frames), num_frames) frame_indices.sort() 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) print("\n" + "=" * 60) print(" HOW TO LABEL & ZOOM") print("=" * 60) print(f"• MODE: {'ALL-IN-ONE (Multi-Class)' if multi_mode else session_label.upper()}") print(f"• SESSION SCOPE: {num_frames} random search zones in '{selected_video}'") print("-" * 60) print("► LABELING CONTROLS:") print(" [Left-Click & Drag] - Draw box around target (auto-saves to active view).") print(" [Right-Click] - Undo / delete last box on the current view.") print(" [SPACEBAR] - SWITCH TARGET CLASS (Cycles through options live).") print("-" * 60) print("► ZOOM & CROP CONTROLS (Isolate moving targets from default clutter!):") print(" [Shift + Click & Drag] - Draw a custom zoom box to isolate an area.") print(" [Middle-Click & Drag] - Alternative shortcut to draw a custom zoom box.") print(" [R Key] - Jump to a Random 50x50% sub-crop of the screen.") print(" [X Key] - Reset zoom back to the full screenshot view.") print("-" * 60) print("► NAVIGATION CONTROLS:") print(" [ENTER Key] - JUMP to NEXT random anchor (resets zoom).") print(" [Right Arrow] / [Left] - Step FORWARD/BACKWARD 10 frames (resets zoom).") print(" [Q Key] - Quit and save all progress.") print("=" * 60) input("\nPress ENTER when you are ready to begin...") LEFT_KEYS = [81, 2, 2424832, 65361, 63234] RIGHT_KEYS = [83, 3, 2490368, 65363, 63235] i = 0 current_frame = frame_indices[i] # Create window once outside loop so user can drag & resize it freely window_name = "YOLO Labeler v5 | Active Session" cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE) cv2.moveWindow(window_name, 20, 20) 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 frame_resized = cv2.resize(frame, (0, 0), fx=SCALE_FACTOR, fy=SCALE_FACTOR) frame_display = frame_resized.copy() cv2.setWindowTitle(window_name, f"YOLO Labeler v5 | Spot {i+1}/{num_frames} (Frame {current_frame})") session_state = { "full_frame": frame_resized, # Stores original uncropped frame "frame": frame_resized.copy(), # Current working frame (full or cropped) "display_frame": frame_display, "crop_box": None, # Tuple (x1, y1, x2, y2) if currently zoomed "view_scale": 1.0, # Multiplier for display upscaling "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, "multi_mode": multi_mode, "box_drawn": False, "drawing": False, "zooming": False, "win_ix": -1, "win_iy": -1, "annotations": [], } load_existing_annotations(session_state) redraw_frame(session_state) cv2.setMouseCallback(window_name, mouse_click_callback, session_state) while True: key = cv2.waitKeyEx(1) # --- SPACEBAR (32): Cycle through classes --- if key == 32 or key == ord(" "): session_state["class_id"] = (session_state["class_id"] + 1) % len(CLASSIFICATION_OPTIONS) session_state["session_label"] = CLASSIFICATION_OPTIONS[session_state["class_id"]] print(f" Switched target class to: [{session_state['class_id']}] {session_state['session_label'].upper()}") redraw_frame(session_state) continue # --- R KEY: Generate a Random Sub-Crop --- elif key == ord("r") or key == ord("R"): h, w = session_state["full_frame"].shape[:2] crop_w = int(w * RANDOM_CROP_SIZE_RATIO) crop_h = int(h * RANDOM_CROP_SIZE_RATIO) rx1 = random.randint(0, w - crop_w) ry1 = random.randint(0, h - crop_h) print(f" Applying random sub-crop at X({rx1}), Y({ry1})") apply_crop(session_state, (rx1, ry1, rx1 + crop_w, ry1 + crop_h)) continue # --- X KEY: Reset Zoom / Return to Full Frame --- elif key == ord("x") or key == ord("X"): if session_state["crop_box"] is not None: print(" Resetting zoom back to full screenshot.") apply_crop(session_state, None) continue # --- TAB: Hide the labels temporarily --- elif key == ord("\t"): redraw_frame(session_state, display_tags=False) # --- ENTER (13 or 10): Next Frame / Save --- elif key in [13, 10]: if session_label == "background": filename_base = get_current_filename_base(session_state) 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, session_state["frame"]) open(label_save_path, "a").close() print(f" Registered negative background frame: {filename_base}") i += 1 if i < len(frame_indices): current_frame = frame_indices[i] break # --- RIGHT ARROW: Step 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: Step 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 # --- Q KEY: 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 cap.release() cv2.destroyAllWindows() print(f"\nSession finished! Dataset updated successfully inside '{OUTPUT_DIRECTORY}'.") if __name__ == "__main__": label_video_session()