506 lines
19 KiB
Python
506 lines
19 KiB
Python
# 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()
|