# 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()