bionickatana/HadesStarYOLO
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

new dataset

Browse Source
This commit is contained in:
bionickatana
2026-07-17 11:19:02 -06:00
parent 427205b264
commit 9b95523a7c
74 changed files with 1268 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
label_v6.py
+521
View File
@@ -0,0 +1,521 @@
# 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()