# Label_v5.py # --------------------- # Upgraded with Multi-Class Tagging, Spacebar Cycling, and Transparent Box Blending 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 video to 1/2 size for easier viewing # 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 load_existing_annotations(param): """Loads existing YOLO annotations from disk when landing on a frame.""" video_base = os.path.splitext(os.path.basename(param["video_name"]))[0] filename_base = f"{video_base}_frame_{param['frame_idx']}" 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 draw_hud(frame, param): """Draws an informative status banner at the top of the screen.""" color = CLASS_COLORS[param["class_id"] % len(CLASS_COLORS)] mode_str = "ALL-IN-ONE" if param["multi_mode"] else "SINGLE-CLASS" label_str = f"MODE: {mode_str} | ACTIVE TAG: [{param['class_id']}] {param['session_label'].upper()} | Press [SPACE] to cycle class" # Banner background cv2.rectangle(frame, (0, 0), (frame.shape[1], 30), (30, 30, 30), -1) # Active color indicator block cv2.rectangle(frame, (0, 0), (12, 30), color, -1) # HUD Text cv2.putText( frame, label_str, (22, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA, ) def redraw_frame(param, display_tags = True): """Regenerates the display frame using semi-transparent boxes and color coding.""" display_frame = param["frame"].copy() # 1. Create overlay layers for transparent fills and borders fill_overlay = display_frame.copy() border_overlay = display_frame.copy() text_frame = display_frame.copy() for ann in param["annotations"]: class_id, label, x1, y1, x2, y2 = ann color = CLASS_COLORS[class_id % len(CLASS_COLORS)] # Faint fill inside the box cv2.rectangle(fill_overlay, (x1, y1), (x2, y2), color, -1) # Thicker border line cv2.rectangle(border_overlay, (x1, y1), (x2, y2), 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 # Add text label cv2.rectangle( text_frame, (x1, max(0, y1 - bg_height)), (x1 + bg_width, max(0, y1)), color, -1, ) cv2.putText( text_frame, label, (x1 + 3, max(12, y1 - 4)), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 0), 1, cv2.LINE_AA, ) # 2. Blend overlays with the base frame (15% opacity fill, 55% opacity borders) if display_tags: cv2.addWeighted(fill_overlay, 0.15, display_frame, 0.85, 0, display_frame) cv2.addWeighted(border_overlay, 0.55, display_frame, 0.45, 0, display_frame) cv2.addWeighted(text_frame, 0.25, display_frame, 0.75, 0, display_frame) # Draw status HUD draw_hud(display_frame, param) param["display_frame"] = display_frame cv2.imshow(param["window_name"], display_frame) def save_annotations_to_file(param): """Overwrites the YOLO text file and updates the image file on disk.""" video_base = os.path.splitext(os.path.basename(param["video_name"]))[0] filename_base = f"{video_base}_frame_{param['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") # If all annotations were undone, delete the files 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 clean untouched frame image cv2.imwrite(image_save_path, param["frame"]) # Write/Overwrite the text file 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)") param["box_drawn"] = True def mouse_click_callback(event, x, y, flags, param): """Handles click-and-drag bounding box creation and right-click undo actions.""" 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 frame = param["frame"] 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 # --- Left Click Down: Start Dragging --- if event == cv2.EVENT_LBUTTONDOWN: param["drawing"] = True param["ix"], param["iy"] = x, y # --- Mouse Move: Live Semi-Transparent Drag Preview --- elif event == cv2.EVENT_MOUSEMOVE: if param["drawing"]: temp_frame = display_frame.copy() overlay = temp_frame.copy() color = CLASS_COLORS[param["class_id"] % len(CLASS_COLORS)] cv2.rectangle(overlay, (param["ix"], param["iy"]), (x, y), color, 3) cv2.addWeighted(overlay, 0.55, temp_frame, 0.45, 0, temp_frame) cv2.imshow(window_name, temp_frame) # --- Left Click Release: Lock Box & Save --- elif event == cv2.EVENT_LBUTTONUP: param["drawing"] = False ix, iy = param["ix"], param["iy"] x1, x2 = min(ix, x), max(ix, x) y1, y2 = min(iy, y), max(iy, y) # Ignore tiny accidental micro-clicks if (x2 - x1) < 5 or (y2 - y1) < 5: cv2.imshow(window_name, display_frame) return 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)) # Append using the CURRENTLY ACTIVE class_id and label 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" + "=" * 50) print(" HOW TO LABEL") print("=" * 50) 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("-" * 50) print("► CONTROLS:") print(" [Click & Drag] - Draw a custom box around the target (auto-saves).") print(" [Right-Click] - Undo / delete the last drawn box on the current frame.") print(" [SPACEBAR] - SWITCH TARGET CLASS (Cycles through options live).") print(" [ENTER Key] - JUMP to the NEXT random video search spot.") print(" [Right Arrow] - Move FORWARD 10 frames (sequential tracking, no save).") print(" [Left Arrow] - Move BACKWARD 10 frames (sequential tracking, no save).") print(" [Q Key] - Quit and save all progress up to this point.") print("=" * 50) 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] 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() window_name = f"YOLO Labeler v5 | 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, "multi_mode": multi_mode, "box_drawn": False, "drawing": False, "ix": -1, "iy": -1, "annotations": [], } # Load any existing boxes for this frame so you don't overwrite previous work load_existing_annotations(session_state) redraw_frame(session_state) cv2.setMouseCallback(window_name, mouse_click_callback, session_state) while True: cv2.moveWindow(window_name, 0, 0) 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 # --- TAB: Hide the labels temporarialy. --- 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": 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() 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 cv2.destroyWindow(window_name) cap.release() cv2.destroyAllWindows() print(f"\nSession finished! Dataset updated successfully inside '{OUTPUT_DIRECTORY}'.") if __name__ == "__main__": label_video_session()