419 lines
17 KiB
Python
419 lines
17 KiB
Python
# Validate_v1.py
|
|
# ---------------------
|
|
# Full-featured YOLO Dataset Validator, Editor, and Sample Cleaner
|
|
|
|
import glob
|
|
import math
|
|
import os
|
|
import cv2
|
|
|
|
# --- CONFIGURATION ---
|
|
DATASET_DIRECTORY = "dataset_v5" # Folder containing /images and /labels
|
|
|
|
CLASSIFICATION_OPTIONS = [
|
|
"guardian",
|
|
"interceptor",
|
|
"sentinel",
|
|
"rocket",
|
|
"bomber",
|
|
]
|
|
|
|
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 load_annotations_for_image(param):
|
|
"""Loads existing YOLO annotations from disk for the current image."""
|
|
image_path = param["image_files"][param["current_idx"]]
|
|
base_name = os.path.splitext(os.path.basename(image_path))[0]
|
|
label_path = os.path.join(param["labels_dir"], f"{base_name}.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}"
|
|
)
|
|
# Stored as a LIST so we can mutate coordinates in-place during resizing/moving
|
|
param["annotations"].append([class_id, label, x1, y1, x2, y2])
|
|
except Exception as e:
|
|
print(f" Error loading annotations for {base_name}: {e}")
|
|
|
|
|
|
def save_annotations_to_file(param):
|
|
"""Overwrites the YOLO text file on disk with the current box list."""
|
|
image_path = param["image_files"][param["current_idx"]]
|
|
base_name = os.path.splitext(os.path.basename(image_path))[0]
|
|
label_save_path = os.path.join(param["labels_dir"], f"{base_name}.txt")
|
|
|
|
if len(param["annotations"]) == 0:
|
|
if os.path.exists(label_save_path):
|
|
os.remove(label_save_path)
|
|
print(f" Removed empty annotation file for {base_name}")
|
|
else:
|
|
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 = abs(x2 - x1)
|
|
height = abs(y2 - y1)
|
|
|
|
x_center_norm = max(0.0, min(1.0, true_center_x / img_w))
|
|
y_center_norm = max(0.0, min(1.0, true_center_y / img_h))
|
|
width_norm = max(0.0, min(1.0, width / img_w))
|
|
height_norm = max(0.0, min(1.0, 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" Saved {len(param['annotations'])} boxes for {base_name}")
|
|
|
|
|
|
def delete_current_sample(param):
|
|
"""Deletes the current image and its associated .txt label file completely from disk."""
|
|
image_path = param["image_files"][param["current_idx"]]
|
|
base_name = os.path.splitext(os.path.basename(image_path))[0]
|
|
label_path = os.path.join(param["labels_dir"], f"{base_name}.txt")
|
|
|
|
if os.path.exists(image_path):
|
|
os.remove(image_path)
|
|
if os.path.exists(label_path):
|
|
os.remove(label_path)
|
|
|
|
print(f"\n [NUKE] Eradicating sample: {base_name}.jpg and associated labels.")
|
|
param["image_files"].pop(param["current_idx"])
|
|
|
|
if param["current_idx"] >= len(param["image_files"]):
|
|
param["current_idx"] = max(0, len(param["image_files"]) - 1)
|
|
|
|
|
|
def draw_hud(frame, param):
|
|
"""Draws status banner and control instructions at the top."""
|
|
color = CLASS_COLORS[param["class_id"] % len(CLASS_COLORS)]
|
|
total = len(param["image_files"])
|
|
idx = param["current_idx"] + 1
|
|
label_str = f"VALIDATING {idx}/{total} | ACTIVE DRAW TAG: [{param['class_id']+1}] {param['session_label'].upper()} | [SPACE] Cycle | [D] Delete Sample"
|
|
|
|
cv2.rectangle(frame, (0, 0), (frame.shape[1], 30), (30, 30, 30), -1)
|
|
cv2.rectangle(frame, (0, 0), (12, 30), color, -1)
|
|
cv2.putText(
|
|
frame, label_str, (22, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.48, (255, 255, 255), 1, cv2.LINE_AA
|
|
)
|
|
|
|
|
|
def redraw_frame(param):
|
|
"""Renders boxes, transparent fills, labels, and interactive corner grab handles."""
|
|
display_frame = param["frame"].copy()
|
|
fill_overlay = display_frame.copy()
|
|
border_overlay = display_frame.copy()
|
|
|
|
for ann in param["annotations"]:
|
|
class_id, label, x1, y1, x2, y2 = ann
|
|
color = CLASS_COLORS[class_id % len(CLASS_COLORS)]
|
|
cv2.rectangle(fill_overlay, (x1, y1), (x2, y2), color, -1)
|
|
cv2.rectangle(border_overlay, (x1, y1), (x2, y2), color, 2)
|
|
|
|
cv2.addWeighted(fill_overlay, 0.20, display_frame, 0.80, 0, display_frame)
|
|
cv2.addWeighted(border_overlay, 0.70, display_frame, 0.30, 0, display_frame)
|
|
|
|
# Draw corner handles and text banners
|
|
for ann in param["annotations"]:
|
|
class_id, label, x1, y1, x2, y2 = ann
|
|
color = CLASS_COLORS[class_id % len(CLASS_COLORS)]
|
|
|
|
# Interactive corner grab dots
|
|
corners = [(x1, y1), (x2, y1), (x2, y2), (x1, y2)]
|
|
for cx, cy in corners:
|
|
cv2.circle(display_frame, (cx, cy), 5, (255, 255, 255), -1)
|
|
cv2.circle(display_frame, (cx, cy), 5, color, 1)
|
|
|
|
# Label banner
|
|
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(display_frame, (x1, max(0, y1 - bg_height)), (x1 + bg_width, max(0, y1)), color, -1)
|
|
cv2.putText(
|
|
display_frame, label, (x1 + 3, max(12, y1 - 4)), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 0), 1, cv2.LINE_AA
|
|
)
|
|
|
|
draw_hud(display_frame, param)
|
|
param["display_frame"] = display_frame
|
|
cv2.imshow(param["window_name"], display_frame)
|
|
|
|
|
|
def get_hovered_box_element(x, y, annotations):
|
|
"""Determines if the mouse is over a corner (resize), inside a box (move), or in empty space (draw)."""
|
|
# Check in reverse order so top-most rendered boxes are grabbed first
|
|
for idx in range(len(annotations) - 1, -1, -1):
|
|
_, _, x1, y1, x2, y2 = annotations[idx]
|
|
corners = [(x1, y1), (x2, y1), (x2, y2), (x1, y2)]
|
|
|
|
# Check distance to corners (12-pixel grab radius)
|
|
for c_idx, (cx, cy) in enumerate(corners):
|
|
if math.hypot(x - cx, y - cy) <= 12:
|
|
return ("resize", idx, c_idx)
|
|
|
|
# Check if inside box body
|
|
min_x, max_x = min(x1, x2), max(x1, x2)
|
|
min_y, max_y = min(y1, y2), max(y1, y2)
|
|
if min_x <= x <= max_x and min_y <= y <= max_y:
|
|
return ("move", idx, (x - x1, y - y1))
|
|
|
|
return ("draw", None, None)
|
|
|
|
|
|
def mouse_click_callback(event, x, y, flags, param):
|
|
"""Handles dragging to create, resize, or move boxes, and right-click deletion."""
|
|
param["mouse_x"], param["mouse_y"] = x, y
|
|
frame = param["frame"]
|
|
display_frame = param["display_frame"]
|
|
window_name = param["window_name"]
|
|
|
|
# --- Right Click: Delete specific hovered box or undo last ---
|
|
if event == cv2.EVENT_RBUTTONDOWN:
|
|
action, idx, _ = get_hovered_box_element(x, y, param["annotations"])
|
|
if idx is not None:
|
|
removed = param["annotations"].pop(idx)
|
|
print(f" Deleted hovered box: [{removed[0]}] {removed[1]}")
|
|
elif param["annotations"]:
|
|
removed = param["annotations"].pop()
|
|
print(f" Undid last box: [{removed[0]}] {removed[1]}")
|
|
save_annotations_to_file(param)
|
|
redraw_frame(param)
|
|
return
|
|
|
|
# --- Left Click Down: Initiate Action ---
|
|
if event == cv2.EVENT_LBUTTONDOWN:
|
|
action, idx, meta = get_hovered_box_element(x, y, param["annotations"])
|
|
param["action"] = action
|
|
param["active_idx"] = idx
|
|
param["action_meta"] = meta
|
|
param["ix"], param["iy"] = x, y
|
|
|
|
# --- Mouse Move: Live Preview for Drawing, Resizing, or Moving ---
|
|
elif event == cv2.EVENT_MOUSEMOVE:
|
|
if param["action"] == "draw":
|
|
temp_frame = display_frame.copy()
|
|
color = CLASS_COLORS[param["class_id"] % len(CLASS_COLORS)]
|
|
cv2.rectangle(temp_frame, (param["ix"], param["iy"]), (x, y), color, 2)
|
|
cv2.imshow(window_name, temp_frame)
|
|
|
|
elif param["action"] == "resize" and param["active_idx"] is not None:
|
|
c_idx = param["action_meta"]
|
|
ann = param["annotations"][param["active_idx"]]
|
|
# Update specific corner coordinates live
|
|
if c_idx == 0: ann[2], ann[3] = x, y
|
|
elif c_idx == 1: ann[4], ann[3] = x, y
|
|
elif c_idx == 2: ann[4], ann[5] = x, y
|
|
elif c_idx == 3: ann[2], ann[5] = x, y
|
|
redraw_frame(param)
|
|
|
|
elif param["action"] == "move" and param["active_idx"] is not None:
|
|
offset_x, offset_y = param["action_meta"]
|
|
ann = param["annotations"][param["active_idx"]]
|
|
width = ann[4] - ann[2]
|
|
height = ann[5] - ann[3]
|
|
|
|
new_x1 = max(0, min(x - offset_x, frame.shape[1] - abs(width)))
|
|
new_y1 = max(0, min(y - offset_y, frame.shape[0] - abs(height)))
|
|
ann[2] = new_x1
|
|
ann[3] = new_y1
|
|
ann[4] = new_x1 + width
|
|
ann[5] = new_y1 + height
|
|
redraw_frame(param)
|
|
|
|
# --- Left Click Release: Lock Coordinates & Save ---
|
|
elif event == cv2.EVENT_LBUTTONUP:
|
|
img_h, img_w = frame.shape[:2]
|
|
|
|
if param["action"] == "draw":
|
|
x1, x2 = min(param["ix"], x), max(param["ix"], x)
|
|
y1, y2 = min(param["iy"], y), max(param["iy"], y)
|
|
if (x2 - x1) >= 5 and (y2 - y1) >= 5:
|
|
param["annotations"].append([
|
|
param["class_id"], param["session_label"],
|
|
max(0, x1), max(0, y1), min(img_w - 1, x2), min(img_h - 1, y2)
|
|
])
|
|
|
|
elif param["action"] in ["resize", "move"] and param["active_idx"] is not None:
|
|
# Normalize inverted coordinates if dragged across axes
|
|
ann = param["annotations"][param["active_idx"]]
|
|
x1, x2 = min(ann[2], ann[4]), max(ann[2], ann[4])
|
|
y1, y2 = min(ann[3], ann[5]), max(ann[3], ann[5])
|
|
ann[2], ann[3] = max(0, x1), max(0, y1)
|
|
ann[4], ann[5] = min(img_w - 1, x2), min(img_h - 1, y2)
|
|
|
|
param["action"] = None
|
|
param["active_idx"] = None
|
|
save_annotations_to_file(param)
|
|
redraw_frame(param)
|
|
|
|
|
|
def validate_dataset_session():
|
|
print("=========================================")
|
|
print(" YOLO DATASET VALIDATOR v1.0 ")
|
|
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):
|
|
print(f"Error: Could not find '{images_dir}'. Please check DATASET_DIRECTORY.")
|
|
return
|
|
|
|
image_files = sorted(glob.glob(os.path.join(images_dir, "*.jpg")) + glob.glob(os.path.join(images_dir, "*.png")))
|
|
if not image_files:
|
|
print(f"Error: No images found inside '{images_dir}'.")
|
|
return
|
|
|
|
print(f"Found {len(image_files)} images in '{DATASET_DIRECTORY}'.")
|
|
print("\n" + "=" * 55)
|
|
print(" VALIDATOR CONTROLS")
|
|
print("=" * 55)
|
|
print(" [Left-Click Corner] - Grab white dot to RESIZE box.")
|
|
print(" [Left-Click Inside] - Click inside box body to MOVE it.")
|
|
print(" [Left-Click Empty] - Drag in empty space to DRAW new box.")
|
|
print(" [Right-Click Box] - Delete hovered box (or undo last).")
|
|
print(" [Number Keys 1-5] - Instantly change class of hovered box.")
|
|
print(" [SPACEBAR] - Cycle active target tag for new boxes.")
|
|
print(" [ENTER / Right] - Save and move to NEXT image.")
|
|
print(" [Left Arrow] - Save and move to PREVIOUS image.")
|
|
print(" [D Key / DEL] - NUKE/DELETE current image & label from disk.")
|
|
print(" [Q Key] - Quit validator.")
|
|
print("=" * 55)
|
|
input("\nPress ENTER when you are ready to start validating...")
|
|
|
|
LEFT_KEYS = [81, 2, 2424832, 65361, 63234]
|
|
RIGHT_KEYS = [83, 3, 2490368, 65363, 63235]
|
|
|
|
session_state = {
|
|
"image_files": image_files,
|
|
"labels_dir": labels_dir,
|
|
"current_idx": 0,
|
|
"class_id": 0,
|
|
"session_label": CLASSIFICATION_OPTIONS[0],
|
|
"window_name": "YOLO Validator & Editor",
|
|
"action": None,
|
|
"active_idx": None,
|
|
"action_meta": None,
|
|
"ix": -1, "iy": -1,
|
|
"mouse_x": -1, "mouse_y": -1,
|
|
"annotations": [],
|
|
}
|
|
|
|
cv2.namedWindow(session_state["window_name"])
|
|
cv2.setMouseCallback(session_state["window_name"], mouse_click_callback, session_state)
|
|
|
|
while session_state["current_idx"] < len(session_state["image_files"]):
|
|
current_img_path = session_state["image_files"][session_state["current_idx"]]
|
|
frame = cv2.imread(current_img_path)
|
|
|
|
if frame is None:
|
|
print(f" Warning: Could not read {current_img_path}. Skipping.")
|
|
session_state["image_files"].pop(session_state["current_idx"])
|
|
continue
|
|
|
|
session_state["frame"] = frame
|
|
load_annotations_for_image(session_state)
|
|
redraw_frame(session_state)
|
|
|
|
while True:
|
|
cv2.moveWindow(session_state["window_name"], 0, 0)
|
|
key = cv2.waitKeyEx(1)
|
|
|
|
# --- SPACEBAR: Cycle active drawing 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)
|
|
continue
|
|
|
|
# --- NUMBER KEYS (1-9): Reclassify hovered box or switch active class ---
|
|
elif ord("1") <= key <= ord(str(len(CLASSIFICATION_OPTIONS))):
|
|
new_class_id = key - ord("1")
|
|
mx, my = session_state["mouse_x"], session_state["mouse_y"]
|
|
action, idx, _ = get_hovered_box_element(mx, my, session_state["annotations"])
|
|
|
|
if idx is not None:
|
|
# Modify class of existing box under cursor
|
|
session_state["annotations"][idx][0] = new_class_id
|
|
session_state["annotations"][idx][1] = CLASSIFICATION_OPTIONS[new_class_id]
|
|
print(f" Reclassified box to: [{new_class_id+1}] {CLASSIFICATION_OPTIONS[new_class_id].upper()}")
|
|
save_annotations_to_file(session_state)
|
|
else:
|
|
# Switch active drawing tool
|
|
session_state["class_id"] = new_class_id
|
|
session_state["session_label"] = CLASSIFICATION_OPTIONS[new_class_id]
|
|
redraw_frame(session_state)
|
|
continue
|
|
|
|
# --- D KEY or DEL: Nuke current sample completely from disk ---
|
|
elif key in [ord("d"), ord("D"), 255, 3014656, 65535]:
|
|
delete_current_sample(session_state)
|
|
if not session_state["image_files"]:
|
|
print(" All samples have been deleted!")
|
|
cv2.destroyAllWindows()
|
|
return
|
|
break
|
|
|
|
# --- ENTER or RIGHT ARROW: Next Image ---
|
|
elif key in [13, 10] or key in RIGHT_KEYS:
|
|
if session_state["current_idx"] < len(session_state["image_files"]) - 1:
|
|
session_state["current_idx"] += 1
|
|
else:
|
|
print(" Reached the end of the dataset!")
|
|
break
|
|
|
|
# --- LEFT ARROW: Previous Image ---
|
|
elif key in LEFT_KEYS:
|
|
if session_state["current_idx"] > 0:
|
|
session_state["current_idx"] -= 1
|
|
else:
|
|
print(" Already at the first image!")
|
|
break
|
|
|
|
# --- Q KEY: Quit ---
|
|
elif key in [ord("q"), ord("Q")] or (key & 0xFF) in [ord("q"), ord("Q")]:
|
|
print("\nExiting validation session. All changes saved.")
|
|
cv2.destroyAllWindows()
|
|
return
|
|
|
|
cv2.destroyAllWindows()
|
|
print("\nDataset validation complete!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
validate_dataset_session()
|