# Validator_OBB.py # --------------------- # Dedicated Dataset Validator & Editor for YOLO-OBB # Features: # - Step through existing dataset images and label files # - Point-and-click deletion: Right-click inside any box to erase it # - Add missing boxes using the 3-step OBB drawing workflow # - Interactive zoom and scaling for inspecting dense clusters import glob import os import cv2 import numpy as np # --- CONFIGURATION --- DATASET_DIRECTORY = "dataset_v6_obb" # Target dataset folder to validate # UI Display Constraints MIN_WINDOW_WIDTH = 800 MIN_WINDOW_HEIGHT = 600 MAX_VIEW_SCALE = 10.0 HUD_HEIGHT = 40 # Defined classification options (Must match your training classes exactly) 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]]) v_mouse = np.array(p_mouse, dtype=float) - np.array(p1, dtype=float) d = np.dot(v_mouse, perp) 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 load_annotations(label_path, img_w, img_h): """Loads existing 8-point YOLO-OBB annotations from disk.""" annotations = [] if os.path.exists(label_path): 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}" ) annotations.append((class_id, label, pts)) except Exception as e: print(f" Error loading {label_path}: {e}") return annotations def save_annotations(label_path, annotations, img_w, img_h): """Overwrites the YOLO-OBB text file with the modified annotation list.""" try: with open(label_path, "w") as f: for ann in annotations: class_id, _, pts = ann norm_coords = [] for pt in pts: # Clamp coordinates between 0.0 and 1.0 to prevent out-of-bounds errors nx = max(0.0, min(1.0, pt[0] / float(img_w))) ny = max(0.0, min(1.0, pt[1] / float(img_h))) norm_coords.append(f"{nx:.6f}") norm_coords.append(f"{ny:.6f}") f.write(f"{class_id} " + " ".join(norm_coords) + "\n") print(f" Saved changes -> {os.path.basename(label_path)} ({len(annotations)} boxes)") except Exception as e: print(f" Failed to save 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: 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 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)" file_name = os.path.basename(param["image_list"][param["img_idx"]]) step_msg = { 0: "[Right-Click Box] Delete | [Drag] Add New | [SPACE] Class | [A/D] Prev/Next", 1: "[Release] Lock Length Axis...", 2: "[Move] Set Width -> [Click] Save OBB | [Right-Click] Cancel" }[param["drawing_step"]] label_str = f"[{zoom_str}] {file_name} ({len(param['annotations'])} boxes) | 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.42, (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)] # If currently cropped, translate coordinates relative to the crop box display_pts = pts if param["crop_box"] is not None: off_x, off_y = param["crop_box"][:2] display_pts = [(p[0] - off_x, p[1] - off_y) for p in pts] scaled_pts = np.array([[int(pt[0] * view_scale), int(pt[1] * view_scale)] for pt in display_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) 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 validator_mouse_callback(event, x, y, flags, param): """Handles point-and-click box deletion, 3-step OBB drawing, and zooming.""" display_frame = param["display_frame"] window_name = param["window_name"] img_h, img_w = param["full_frame"].shape[:2] current_label_path = param["label_list"][param["img_idx"]] # --- Right Click: Cancel drawing OR Delete Clicked Box --- if event == cv2.EVENT_RBUTTONDOWN: if param["drawing_step"] > 0: param["drawing_step"] = 0 print(" Cancelled active box drawing.") redraw_frame(param) return if not param["annotations"]: print(" No boxes on this image to delete!") return # Get exact mouse click position in native image coordinates click_f = win_to_frame_coords(x, y, param) if param["crop_box"] is not None: click_f = (click_f[0] + param["crop_box"][0], click_f[1] + param["crop_box"][1]) deleted = False # Iterate backwards to test the top-most drawn box first for idx in range(len(param["annotations"]) - 1, -1, -1): pts = param["annotations"][idx][2] poly = np.array(pts, dtype=np.int32) # Point-in-Polygon test: returns positive if inside, 0 if on edge, negative if outside if cv2.pointPolygonTest(poly, click_f, False) >= 0: removed = param["annotations"].pop(idx) print(f" Deleted box under cursor: [{removed[0]}] {removed[1]}") save_annotations(current_label_path, param["annotations"], img_w, img_h) redraw_frame(param) deleted = True break if not deleted: # If they didn't click inside a specific box, fallback to undoing the last drawn box removed = param["annotations"].pop() print(f" Undid last box: [{removed[0]}] {removed[1]}") save_annotations(current_label_path, param["annotations"], img_w, img_h) redraw_frame(param) return # --- Middle Click OR Shift + Left Click: 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: Add New Box --- if event == cv2.EVENT_LBUTTONDOWN and not param["zooming"]: if param["drawing_step"] == 0: param["drawing_step"] = 1 param["p1_win"] = (x, y) redraw_frame(param) elif param["drawing_step"] == 2: 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) # If zoomed in, translate coordinates back to full image space if param["crop_box"] is not None: off_x, off_y = param["crop_box"][:2] p1_f = (p1_f[0] + off_x, p1_f[1] + off_y) p2_f = (p2_f[0] + off_x, p2_f[1] + off_y) mouse_f = (mouse_f[0] + off_x, mouse_f[1] + off_y) 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(current_label_path, param["annotations"], img_w, img_h) 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: 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: 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: if abs(x - param["p1_win"][0]) < 5 and abs(y - param["p1_win"][1]) < 5: param["drawing_step"] = 0 redraw_frame(param) else: param["p2_win"] = (x, y) param["drawing_step"] = 2 redraw_frame(param) def validate_dataset(): print("=========================================") print(" YOLO-OBB DATASET VALIDATOR ") print("=========================================\n") images_dir = os.path.join(DATASET_DIRECTORY, "images") labels_dir = os.path.join(DATASET_DIRECTORY, "labels") if not os.path.exists(images_dir) or not os.path.exists(labels_dir): print(f"Error: Could not find '{images_dir}' or '{labels_dir}'.") print("Please ensure DATASET_DIRECTORY is set to the correct folder path.") return # Find all supported image extensions image_files = sorted( glob.glob(os.path.join(images_dir, "*.jpg")) + glob.glob(os.path.join(images_dir, "*.png")) + glob.glob(os.path.join(images_dir, "*.jpeg")) ) if not image_files: print(f"Error: No images found inside '{images_dir}'.") return # Map image files to corresponding .txt label paths label_files = [] for img_path in image_files: base_name = os.path.splitext(os.path.basename(img_path))[0] label_files.append(os.path.join(labels_dir, f"{base_name}.txt")) print(f"Found {len(image_files)} images in '{DATASET_DIRECTORY}'.") print("\n" + "=" * 60) print(" VALIDATOR CONTROLS") print("=" * 60) print("► NAVIGATION:") print(" [D Key] / [Right Arrow] / [ENTER] - Go to NEXT image.") print(" [A Key] / [Left Arrow] - Go to PREVIOUS image.") print(" [Q Key] - Quit validation session.") print("-" * 60) print("► EDITING & DELETING:") print(" [Right-Click inside a Box] - DELETE that specific bounding box!") print(" [Left-Click & Drag] - Add a missing box (3-step drawing).") print(" [SPACEBAR] - Cycle active target class.") print("-" * 60) print("► VIEW / ZOOM:") print(" [Shift + Drag] - Zoom into a dense cluster.") print(" [X Key] - Reset zoom back to full screen.") print(" [TAB Key] - Toggle text labels on/off.") print("=" * 60) input("\nPress ENTER to start validating...") LEFT_KEYS = [ord("a"), ord("A"), 81, 2, 2424832, 65361, 63234] RIGHT_KEYS = [ord("d"), ord("D"), 13, 10, 83, 3, 2490368, 65363, 63235] img_idx = 0 window_name = "YOLO-OBB Dataset Validator" cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE) cv2.moveWindow(window_name, 20, 20) while 0 <= img_idx < len(image_files): img_path = image_files[img_idx] label_path = label_files[img_idx] frame = cv2.imread(img_path) if frame is None: print(f"Warning: Failed to load image {img_path}. Skipping.") img_idx += 1 continue img_h, img_w = frame.shape[:2] annotations = load_annotations(label_path, img_w, img_h) cv2.setWindowTitle(window_name, f"Validator | Image {img_idx + 1}/{len(image_files)} | {os.path.basename(img_path)}") session_state = { "full_frame": frame, "frame": frame.copy(), "display_frame": frame.copy(), "crop_box": None, "view_scale": 1.0, "image_list": image_files, "label_list": label_files, "img_idx": img_idx, "window_name": window_name, "class_id": 0, "session_label": CLASSIFICATION_OPTIONS[0], "drawing_step": 0, "zooming": False, "win_ix": -1, "win_iy": -1, "p1_win": (0, 0), "p2_win": (0, 0), "annotations": annotations, } redraw_frame(session_state) cv2.setMouseCallback(window_name, validator_mouse_callback, session_state) while True: key = cv2.waitKeyEx(1) # SPACE: Switch Class 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) # X: Reset Zoom elif key == ord("x") or key == ord("X"): if session_state["crop_box"] is not None: apply_crop(session_state, None) # TAB: Toggle Labels elif key == ord("\t"): redraw_frame(session_state, display_tags=False) # NEXT IMAGE elif key in RIGHT_KEYS: if img_idx < len(image_files) - 1: img_idx += 1 else: print(" Reached the end of the dataset!") break # PREVIOUS IMAGE elif key in LEFT_KEYS: if img_idx > 0: img_idx -= 1 else: print(" Already at the first image!") break # QUIT elif key == ord("q") or key == ord("Q") or (key & 0xFF) in [ord("q"), ord("Q")]: cv2.destroyAllWindows() print("\nValidator session closed.") return cv2.destroyAllWindows() print("\nDataset validation complete!") if __name__ == "__main__": validate_dataset()