# Label_v6_OBB.py # --------------------- # Upgraded for YOLO-OBB (Oriented Bounding Boxes) # Features intuitive 3-step angled box drawing, dynamic view scaling, # interactive zooming, and multi-class tagging. import glob import os import random import cv2 import numpy as np # --- CONFIGURATION --- OUTPUT_DIRECTORY = "dataset_v6_obb" 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 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_obb_corners(p1, p2, p_mouse): """Calculates the 4 corners of an oriented bounding box from a baseline and mouse position.""" v = np.array(p2, dtype=float) - np.array(p1, dtype=float) length = np.linalg.norm(v) if length < 1e-5: return [p1, p1, p1, p1] u = v / length perp = np.array([-u[1], u[0]]) # Perpendicular unit vector v_mouse = np.array(p_mouse, dtype=float) - np.array(p1, dtype=float) d = np.dot(v_mouse, perp) # Signed distance from mouse to baseline c1 = np.array(p1, dtype=float) c2 = np.array(p2, dtype=float) c3 = c2 + d * perp c4 = c1 + d * perp return [ tuple(np.round(c1).astype(int)), tuple(np.round(c2).astype(int)), tuple(np.round(c3).astype(int)), tuple(np.round(c4).astype(int)) ] def get_dataset_counts(): """Counts the number of existing OBB 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() # OBB format has 1 class_id + 8 coordinate values (9 total) if len(parts) >= 9: class_id = int(parts[0]) if 0 <= class_id < len(CLASSIFICATION_OPTIONS): option = CLASSIFICATION_OPTIONS[class_id] counts[option] += 1 except Exception: pass 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-OBB annotations (4 corner points) from disk.""" 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[:2] try: with open(label_path, "r") as f: for line in f: parts = line.strip().split() if len(parts) >= 9: class_id = int(parts[0]) pts = [] for j in range(1, 9, 2): nx = float(parts[j]) ny = float(parts[j+1]) pts.append((int(nx * img_w), int(ny * img_h))) label = ( CLASSIFICATION_OPTIONS[class_id] if 0 <= class_id < len(CLASSIFICATION_OPTIONS) else f"ID {class_id}" ) param["annotations"].append((class_id, label, pts)) 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_step"] = 0 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.""" 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) 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)] zoom_str = "FULL" if param["crop_box"] is None else f"CROP ({view_scale:.1f}x VIEW)" step_msg = { 0: "[Drag] Set Length | [SPACE] Class | [R] Rand Crop", 1: "[Release] Lock Length Axis...", 2: "[Move] Set Width -> [Click] Save OBB | [Right-Click] Cancel" }[param["drawing_step"]] label_str = f"[{zoom_str}] ACTIVE: [{param['class_id']}] {param['session_label'].upper()} | {step_msg}" cv2.rectangle(frame, (0, 0), (14, HUD_HEIGHT), color, -1) cv2.line(frame, (0, HUD_HEIGHT - 1), (frame.shape[1], HUD_HEIGHT - 1), (70, 70, 70), 1) 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 oriented polygon overlays.""" h_frame, w_frame = param["frame"].shape[:2] 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) param["view_scale"] = view_scale 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) fill_overlay = display_img.copy() border_overlay = display_img.copy() text_frame = display_img.copy() for ann in param["annotations"]: class_id, label, pts = ann color = CLASS_COLORS[class_id % len(CLASS_COLORS)] scaled_pts = np.array([[int(pt[0] * view_scale), int(pt[1] * view_scale)] for pt in pts], dtype=np.int32) scaled_pts = scaled_pts.reshape((-1, 1, 2)) cv2.fillPoly(fill_overlay, [scaled_pts], color) cv2.polylines(border_overlay, [scaled_pts], True, color, 2) # Place text label near the first vertex dx1, dy1 = scaled_pts[0][0][0], scaled_pts[0][0][1] text_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.45, 1) bg_width, bg_height = text_size[0] + 6, 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.20, display_img, 0.80, 0, display_img) cv2.addWeighted(border_overlay, 0.70, display_img, 0.30, 0, display_img) cv2.addWeighted(text_frame, 0.30, display_img, 0.70, 0, display_img) final_display = cv2.copyMakeBorder( display_img, HUD_HEIGHT, 0, 0, 0, cv2.BORDER_CONSTANT, value=(30, 30, 30) ) 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-OBB text file with 8 corner coordinates and saves the image patch.""" 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}") else: cv2.imwrite(image_save_path, param["frame"]) img_h, img_w = param["frame"].shape[:2] with open(label_save_path, "w") as f: for ann in param["annotations"]: class_id, _, pts = ann norm_coords = [] for pt in pts: norm_coords.append(f"{pt[0] / img_w:.6f}") norm_coords.append(f"{pt[1] / img_h:.6f}") f.write(f"{class_id} " + " ".join(norm_coords) + "\n") print(f" Updated OBB dataset for {filename_base} ({len(param['annotations'])} boxes)") def mouse_click_callback(event, x, y, flags, param): """Handles 3-step OBB drawing, zooming, and undo actions.""" display_frame = param["display_frame"] window_name = param["window_name"] # --- Right Click: Cancel active drawing OR Undo last box --- if event == cv2.EVENT_RBUTTONDOWN: if param["drawing_step"] > 0: param["drawing_step"] = 0 print(" Cancelled active box drawing.") redraw_frame(param) elif param["annotations"]: removed = param["annotations"].pop() print(f" Undid last OBB ([{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 --- if event == cv2.EVENT_LBUTTONDOWN and not param["zooming"]: if param["drawing_step"] == 0: # Step 1: Start baseline length axis param["drawing_step"] = 1 param["p1_win"] = (x, y) redraw_frame(param) elif param["drawing_step"] == 2: # Step 3: Second click confirms width and locks OBB p1_f = win_to_frame_coords(*param["p1_win"], param) p2_f = win_to_frame_coords(*param["p2_win"], param) mouse_f = win_to_frame_coords(x, y, param) pts = get_obb_corners(p1_f, p2_f, mouse_f) param["annotations"].append((param["class_id"], param["session_label"], pts)) param["drawing_step"] = 0 save_annotations_to_file(param) redraw_frame(param) # --- Mouse Move: Live Preview --- elif event == cv2.EVENT_MOUSEMOVE: if param["zooming"]: temp = display_frame.copy() cv2.rectangle(temp, (param["win_ix"], param["win_iy"]), (x, y), (255, 255, 255), 2) cv2.imshow(window_name, temp) elif param["drawing_step"] == 1: # Preview length baseline temp = display_frame.copy() color = CLASS_COLORS[param["class_id"] % len(CLASS_COLORS)] cv2.line(temp, param["p1_win"], (x, y), color, 2) cv2.imshow(window_name, temp) elif param["drawing_step"] == 2: # Preview full oriented bounding box expanding to mouse position temp = display_frame.copy() color = CLASS_COLORS[param["class_id"] % len(CLASS_COLORS)] pts_win = get_obb_corners(param["p1_win"], param["p2_win"], (x, y)) poly = np.array(pts_win, dtype=np.int32).reshape((-1, 1, 2)) cv2.polylines(temp, [poly], True, color, 2) cv2.imshow(window_name, temp) # --- Left Click Release --- elif event == cv2.EVENT_LBUTTONUP: if param["zooming"]: param["zooming"] = False 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 and (y2 - y1) >= 10: if param["crop_box"] is not None: off_x, off_y = param["crop_box"][:2] x1, x2 = x1 + off_x, x2 + off_x y1, y2 = y1 + off_y, y2 + off_y apply_crop(param, (x1, y1, x2, y2)) else: cv2.imshow(window_name, display_frame) elif param["drawing_step"] == 1: # Step 2: Release locks baseline length, switches mouse to width mode if abs(x - param["p1_win"][0]) < 5 and abs(y - param["p1_win"][1]) < 5: # Ignore accidental micro-clicks param["drawing_step"] = 0 redraw_frame(param) else: param["p2_win"] = (x, y) param["drawing_step"] = 2 redraw_frame(param) def label_video_session(): print("=========================================") print(" AI YOLO-OBB DATASET LABELER v6 ") print("=========================================\n") mp4_files = glob.glob("*.mp4") if not mp4_files: print("Error: No .mp4 files found in current directory.") return selected_video = random.choice(mp4_files) print(f"-> Selected Video File: '{selected_video}'") counts = get_dataset_counts() print("\n-----------------------------------------") print(" CURRENT OBB DATASET STATUS ") print("-----------------------------------------") for option, count in counts.items(): print(f" • {option.capitalize():<12} : {count} OBBs logged") print("-----------------------------------------") session_class_id = 0 session_label = CLASSIFICATION_OPTIONS[0] 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)) num_frames = min(FRAMES_TO_EXTRACT, total_frames) frame_indices = sorted(random.sample(range(total_frames), num_frames)) 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 DRAW OBB BOXES") print("=" * 60) print("► 3-STEP ORIENTED BOX DRAWING:") print(" 1. [Left-Click & Drag] along the length (nose to tail of ship).") print(" 2. [Release Mouse] and move away to expand the box width.") print(" 3. [Left-Click Again] to lock and save the angled rectangle!") print("-" * 60) print("► GENERAL CONTROLS:") print(" [Right-Click] - Cancel active drawing OR Undo last box.") print(" [SPACEBAR] - SWITCH TARGET CLASS.") print(" [Shift + Drag] - Zoom into a custom cluster of ships.") print(" [R Key] - Jump to random 50% sub-crop.") print(" [X Key] - Reset zoom back to full screenshot.") print(" [ENTER Key] - JUMP to NEXT frame / Save background.") print(" [Right / Left Arrow] - Step FORWARD/BACKWARD 10 frames.") 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] window_name = "YOLO-OBB Labeler v6" 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: 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) cv2.setWindowTitle(window_name, f"YOLO-OBB Labeler v6 | Spot {i+1}/{num_frames} (Frame {current_frame})") session_state = { "full_frame": frame_resized, "frame": frame_resized.copy(), "display_frame": frame_resized.copy(), "crop_box": None, "view_scale": 1.0, "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, "drawing_step": 0, # 0: Idle, 1: Dragging baseline, 2: Setting width "zooming": False, "win_ix": -1, "win_iy": -1, "p1_win": (0, 0), "p2_win": (0, 0), "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) 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"]] redraw_frame(session_state) elif key == ord("r") or key == ord("R"): h, w = session_state["full_frame"].shape[:2] crop_w, crop_h = int(w * RANDOM_CROP_SIZE_RATIO), int(h * RANDOM_CROP_SIZE_RATIO) rx1, ry1 = random.randint(0, w - crop_w), random.randint(0, h - crop_h) apply_crop(session_state, (rx1, ry1, rx1 + crop_w, ry1 + crop_h)) elif key == ord("x") or key == ord("X"): if session_state["crop_box"] is not None: apply_crop(session_state, None) elif key == ord("\t"): redraw_frame(session_state, display_tags=False) elif key in [13, 10]: i += 1 if i < len(frame_indices): current_frame = frame_indices[i] break elif key in RIGHT_KEYS: if current_frame < total_frames - 10: current_frame += 10 break elif key in LEFT_KEYS: if current_frame > 0: current_frame -= 10 break elif key == ord("q") or key == ord("Q") or (key & 0xFF) in [ord("q"), ord("Q")]: cap.release() cv2.destroyAllWindows() return cap.release() cv2.destroyAllWindows() print(f"\nSession finished! OBB Dataset saved to '{OUTPUT_DIRECTORY}'.") if __name__ == "__main__": label_video_session()