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

updated labeler and added more training data

Browse Source
This commit is contained in:
bionickatana
2026-07-15 20:45:12 -06:00
parent 2f36829131
commit 55eca275fa
153 changed files with 472 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
label_v5.py
+229 -99
View File
@@ -1,6 +1,7 @@
# Label_v5.py
# ---------------------
# Upgraded with Multi-Class Tagging, Spacebar Cycling, and Transparent Box Blending
# Upgraded with Dedicated Header Bar, Dynamic View Scaling for Small Crops,
# Enforced Minimum Window Sizes, Interactive Zoom, and Multi-Class Tagging.
import glob
import os
@@ -10,7 +11,14 @@ import cv2
# --- CONFIGURATION ---
OUTPUT_DIRECTORY = "dataset_v5"
FRAMES_TO_EXTRACT = 20 # Number of random anchor frames to find targets
SCALE_FACTOR = 0.5 # Resize video to 1/2 size for easier viewing
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 = [
@@ -55,10 +63,19 @@ def get_dataset_counts():
return counts
def load_existing_annotations(param):
"""Loads existing YOLO annotations from disk when landing on a frame."""
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]
filename_base = f"{video_base}_frame_{param['frame_idx']}"
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 annotations from disk for the current frame/crop view."""
filename_base = get_current_filename_base(param)
label_path = os.path.join(param["labels_dir"], f"{filename_base}.txt")
param["annotations"] = []
@@ -95,95 +112,157 @@ def load_existing_annotations(param):
print(f" Could not load existing annotations: {e}")
def draw_hud(frame, param):
"""Draws an informative status banner at the top of the screen."""
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"] = False
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."""
# Subtract the dedicated header bar height from y
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)
# Clamp to actual working frame dimensions
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)]
mode_str = "ALL-IN-ONE" if param["multi_mode"] else "SINGLE-CLASS"
label_str = f"MODE: {mode_str} | ACTIVE TAG: [{param['class_id']}] {param['session_label'].upper()} | Press [SPACE] to cycle class"
zoom_str = "FULL" if param["crop_box"] is None else f"CROP ({view_scale:.1f}x VIEW)"
label_str = f"[{zoom_str}] | ACTIVE: [{param['class_id']}] {param['session_label'].upper()} | [SPACE] Class | [R] Rand Crop | [X] Reset Zoom"
# Banner background
cv2.rectangle(frame, (0, 0), (frame.shape[1], 30), (30, 30, 30), -1)
# Active color indicator block
cv2.rectangle(frame, (0, 0), (12, 30), color, -1)
# HUD Text
# Active color indicator block on the left edge of the banner
cv2.rectangle(frame, (0, 0), (14, HUD_HEIGHT), color, -1)
# Draw a subtle separator line between HUD and image
cv2.line(frame, (0, HUD_HEIGHT - 1), (frame.shape[1], HUD_HEIGHT - 1), (70, 70, 70), 1)
# HUD Text centered vertically in the 40px banner
cv2.putText(
frame,
label_str,
(22, 20),
(24, 25),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
0.45,
(255, 255, 255),
1,
cv2.LINE_AA,
)
def redraw_frame(param, display_tags = True):
"""Regenerates the display frame using semi-transparent boxes and color coding."""
display_frame = param["frame"].copy()
def redraw_frame(param, display_tags=True):
"""Regenerates the display frame using dynamic view scaling and transparent box overlays."""
h_frame, w_frame = param["frame"].shape[:2]
# 1. Create overlay layers for transparent fills and borders
fill_overlay = display_frame.copy()
border_overlay = display_frame.copy()
text_frame = display_frame.copy()
# --- DYNAMIC VIEW SCALING ---
# Calculate scale needed to meet minimum comfortable window viewing size
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) # Cap extreme zooms
param["view_scale"] = view_scale
# Scale up the working frame ONLY for display purposes
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)
# --- OVERLAY LAYERS ---
fill_overlay = display_img.copy()
border_overlay = display_img.copy()
text_frame = display_img.copy()
for ann in param["annotations"]:
class_id, label, x1, y1, x2, y2 = ann
color = CLASS_COLORS[class_id % len(CLASS_COLORS)]
# Faint fill inside the box
cv2.rectangle(fill_overlay, (x1, y1), (x2, y2), color, -1)
# Thicker border line
cv2.rectangle(border_overlay, (x1, y1), (x2, y2), color, 3)
# Map working frame box coordinates to the scaled display dimensions
dx1, dy1 = int(x1 * view_scale), int(y1 * view_scale)
dx2, dy2 = int(x2 * view_scale), int(y2 * view_scale)
cv2.rectangle(fill_overlay, (dx1, dy1), (dx2, dy2), color, -1)
cv2.rectangle(border_overlay, (dx1, dy1), (dx2, dy2), color, 3)
text_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.45, 1)
bg_width = text_size[0] + 6
bg_height = text_size[1] + 6
# Add text label
cv2.rectangle(
text_frame,
(x1, max(0, y1 - bg_height)),
(x1 + bg_width, max(0, y1)),
(dx1, max(0, dy1 - bg_height)),
(dx1 + bg_width, max(0, dy1)),
color,
-1,
)
cv2.putText(
text_frame,
label,
(x1 + 3, max(12, y1 - 4)),
(dx1 + 3, max(12, dy1 - 4)),
cv2.FONT_HERSHEY_SIMPLEX,
0.45,
(0, 0, 0),
1,
cv2.LINE_AA,
)
# 2. Blend overlays with the base frame (15% opacity fill, 55% opacity borders)
if display_tags:
cv2.addWeighted(fill_overlay, 0.15, display_frame, 0.85, 0, display_frame)
cv2.addWeighted(border_overlay, 0.55, display_frame, 0.45, 0, display_frame)
cv2.addWeighted(text_frame, 0.25, display_frame, 0.75, 0, display_frame)
cv2.addWeighted(fill_overlay, 0.15, display_img, 0.85, 0, display_img)
cv2.addWeighted(border_overlay, 0.55, display_img, 0.45, 0, display_img)
cv2.addWeighted(text_frame, 0.25, display_img, 0.75, 0, display_img)
# --- DEDICATED HEADER BAR ---
# Add top border padding for the HUD banner so it NEVER covers the image
final_display = cv2.copyMakeBorder(
display_img,
HUD_HEIGHT,
0,
0,
0,
cv2.BORDER_CONSTANT,
value=(30, 30, 30),
)
# Draw status HUD
draw_hud(display_frame, param)
# Draw status HUD in the top padded area
draw_hud(final_display, param, view_scale)
param["display_frame"] = display_frame
cv2.imshow(param["window_name"], display_frame)
param["display_frame"] = final_display
cv2.imshow(param["window_name"], final_display)
def save_annotations_to_file(param):
"""Overwrites the YOLO text file and updates the image file on disk."""
video_base = os.path.splitext(os.path.basename(param["video_name"]))[0]
filename_base = f"{video_base}_frame_{param['frame_idx']}"
"""Overwrites the YOLO text file and saves the un-scaled image patch to disk."""
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 all annotations were undone, delete the files
if len(param["annotations"]) == 0:
if os.path.exists(label_save_path):
os.remove(label_save_path)
@@ -192,10 +271,10 @@ def save_annotations_to_file(param):
print(f" Removed empty annotation files for {filename_base}")
param["box_drawn"] = False
else:
# Save the clean untouched frame image
# Save the ORIGINAL un-scaled working frame (keeps dataset images small & standard)
cv2.imwrite(image_save_path, param["frame"])
# Write/Overwrite the text file
# Write YOLO coordinates normalized against the working frame dimensions
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"]:
@@ -214,18 +293,17 @@ def save_annotations_to_file(param):
f.write(
f"{class_id} {x_center_norm:.6f} {y_center_norm:.6f} {width_norm:.6f} {height_norm:.6f}\n"
)
print(f" Updated dataset files for {filename_base} ({len(param['annotations'])} boxes)")
print(f" Updated dataset files for {filename_base} ({len(param['annotations'])} boxes | Saved at {img_w}x{img_h}px)")
param["box_drawn"] = True
def mouse_click_callback(event, x, y, flags, param):
"""Handles click-and-drag bounding box creation and right-click undo actions."""
"""Handles labeling drags, zoom dragging, and undo actions in native UI window coordinates."""
if param["session_label"] == "background":
if event == cv2.EVENT_LBUTTONDOWN:
print(" Info: In background mode, do not click targets. Just press ENTER to save frame.")
return
frame = param["frame"]
display_frame = param["display_frame"]
window_name = param["window_name"]
@@ -240,42 +318,67 @@ def mouse_click_callback(event, x, y, flags, param):
print(" No boxes left to undo on this frame!")
return
# --- Left Click Down: Start Dragging ---
# --- 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: Start Label Box Dragging ---
if event == cv2.EVENT_LBUTTONDOWN:
param["drawing"] = True
param["ix"], param["iy"] = x, y
param["win_ix"], param["win_iy"] = x, y
# --- Mouse Move: Live Semi-Transparent Drag Preview ---
# --- Mouse Move: Live Preview in UI Space ---
elif event == cv2.EVENT_MOUSEMOVE:
if param["drawing"]:
if param["drawing"] or param["zooming"]:
temp_frame = display_frame.copy()
overlay = temp_frame.copy()
color = CLASS_COLORS[param["class_id"] % len(CLASS_COLORS)]
cv2.rectangle(overlay, (param["ix"], param["iy"]), (x, y), color, 3)
color = (255, 255, 255) if param["zooming"] else CLASS_COLORS[param["class_id"] % len(CLASS_COLORS)]
cv2.rectangle(overlay, (param["win_ix"], param["win_iy"]), (x, y), color, 2 if param["zooming"] else 3)
cv2.addWeighted(overlay, 0.55, temp_frame, 0.45, 0, temp_frame)
cv2.imshow(window_name, temp_frame)
# --- Left Click Release: Lock Box & Save ---
elif event == cv2.EVENT_LBUTTONUP:
param["drawing"] = False
ix, iy = param["ix"], param["iy"]
# --- Middle Click OR Shift + Left Click Release: Apply Zoom Crop ---
elif event == cv2.EVENT_MBUTTONUP or (event == cv2.EVENT_LBUTTONUP and param["zooming"]):
param["zooming"] = False
# Convert UI window coordinates back to native working frame coordinates
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)
x1, x2 = min(ix, x), max(ix, x)
y1, y2 = min(iy, y), max(iy, y)
# Ignore tiny accidental micro-clicks
if (x2 - x1) < 5 or (y2 - y1) < 5:
if (x2 - x1) < 10 or (y2 - y1) < 10:
cv2.imshow(window_name, display_frame)
return
img_h, img_w = frame.shape[0], frame.shape[1]
x1 = max(0, min(x1, img_w - 1))
y1 = max(0, min(y1, img_h - 1))
x2 = max(0, min(x2, img_w - 1))
y2 = max(0, min(y2, img_h - 1))
# If already cropped, translate local coordinates back to full_frame coordinates
if param["crop_box"] is not None:
offset_x, offset_y = param["crop_box"][0], param["crop_box"][1]
x1, x2 = x1 + offset_x, x2 + offset_x
y1, y2 = y1 + offset_y, y2 + offset_y
print(f" Zooming into region: X({x1}->{x2}), Y({y1}->{y2})")
apply_crop(param, (x1, y1, x2, y2))
# --- Left Click Release: Lock Label Box & Save ---
elif event == cv2.EVENT_LBUTTONUP and param["drawing"]:
param["drawing"] = False
# Ignore tiny accidental micro-clicks in UI window space
if abs(x - param["win_ix"]) < 5 or abs(y - param["win_iy"]) < 5:
cv2.imshow(window_name, display_frame)
return
# Convert UI window coordinates back to native working frame coordinates
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)
# Append using the CURRENTLY ACTIVE class_id and label
param["annotations"].append(
(param["class_id"], param["session_label"], x1, y1, x2, y2)
)
@@ -353,21 +456,28 @@ def label_video_session():
os.makedirs(images_dir, exist_ok=True)
os.makedirs(labels_dir, exist_ok=True)
print("\n" + "=" * 50)
print(" HOW TO LABEL")
print("=" * 50)
print("\n" + "=" * 60)
print(" HOW TO LABEL & ZOOM")
print("=" * 60)
print(f"• MODE: {'ALL-IN-ONE (Multi-Class)' if multi_mode else session_label.upper()}")
print(f"• SESSION SCOPE: {num_frames} random search zones in '{selected_video}'")
print("-" * 50)
print("► CONTROLS:")
print(" [Click & Drag] - Draw a custom box around the target (auto-saves).")
print(" [Right-Click] - Undo / delete the last drawn box on the current frame.")
print(" [SPACEBAR] - SWITCH TARGET CLASS (Cycles through options live).")
print(" [ENTER Key] - JUMP to the NEXT random video search spot.")
print(" [Right Arrow] - Move FORWARD 10 frames (sequential tracking, no save).")
print(" [Left Arrow] - Move BACKWARD 10 frames (sequential tracking, no save).")
print(" [Q Key] - Quit and save all progress up to this point.")
print("=" * 50)
print("-" * 60)
print("► LABELING CONTROLS:")
print(" [Left-Click & Drag] - Draw box around target (auto-saves to active view).")
print(" [Right-Click] - Undo / delete last box on the current view.")
print(" [SPACEBAR] - SWITCH TARGET CLASS (Cycles through options live).")
print("-" * 60)
print("► ZOOM & CROP CONTROLS (Isolate moving targets from default clutter!):")
print(" [Shift + Click & Drag] - Draw a custom zoom box to isolate an area.")
print(" [Middle-Click & Drag] - Alternative shortcut to draw a custom zoom box.")
print(" [R Key] - Jump to a Random 50x50% sub-crop of the screen.")
print(" [X Key] - Reset zoom back to the full screenshot view.")
print("-" * 60)
print("► NAVIGATION CONTROLS:")
print(" [ENTER Key] - JUMP to NEXT random anchor (resets zoom).")
print(" [Right Arrow] / [Left] - Step FORWARD/BACKWARD 10 frames (resets zoom).")
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]
@@ -376,6 +486,11 @@ def label_video_session():
i = 0
current_frame = frame_indices[i]
# Create window once outside loop so user can drag & resize it freely
window_name = "YOLO Labeler v5 | Active Session"
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()
@@ -389,12 +504,14 @@ def label_video_session():
frame_resized = cv2.resize(frame, (0, 0), fx=SCALE_FACTOR, fy=SCALE_FACTOR)
frame_display = frame_resized.copy()
window_name = f"YOLO Labeler v5 | Spot {i+1}/{num_frames} (Frame {current_frame})"
cv2.namedWindow(window_name)
cv2.setWindowTitle(window_name, f"YOLO Labeler v5 | Spot {i+1}/{num_frames} (Frame {current_frame})")
session_state = {
"frame": frame_resized,
"full_frame": frame_resized, # Stores original uncropped frame
"frame": frame_resized.copy(), # Current working frame (full or cropped)
"display_frame": frame_display,
"crop_box": None, # Tuple (x1, y1, x2, y2) if currently zoomed
"view_scale": 1.0, # Multiplier for display upscaling
"images_dir": images_dir,
"labels_dir": labels_dir,
"frame_idx": current_frame,
@@ -405,19 +522,17 @@ def label_video_session():
"multi_mode": multi_mode,
"box_drawn": False,
"drawing": False,
"ix": -1,
"iy": -1,
"zooming": False,
"win_ix": -1,
"win_iy": -1,
"annotations": [],
}
# Load any existing boxes for this frame so you don't overwrite previous work
load_existing_annotations(session_state)
redraw_frame(session_state)
cv2.setMouseCallback(window_name, mouse_click_callback, session_state)
while True:
cv2.moveWindow(window_name, 0, 0)
key = cv2.waitKeyEx(1)
# --- SPACEBAR (32): Cycle through classes ---
@@ -428,19 +543,36 @@ def label_video_session():
redraw_frame(session_state)
continue
# --- TAB: Hide the labels temporarialy. ---
# --- R KEY: Generate a Random Sub-Crop ---
elif key == ord("r") or key == ord("R"):
h, w = session_state["full_frame"].shape[:2]
crop_w = int(w * RANDOM_CROP_SIZE_RATIO)
crop_h = int(h * RANDOM_CROP_SIZE_RATIO)
rx1 = random.randint(0, w - crop_w)
ry1 = random.randint(0, h - crop_h)
print(f" Applying random sub-crop at X({rx1}), Y({ry1})")
apply_crop(session_state, (rx1, ry1, rx1 + crop_w, ry1 + crop_h))
continue
# --- X KEY: Reset Zoom / Return to Full Frame ---
elif key == ord("x") or key == ord("X"):
if session_state["crop_box"] is not None:
print(" Resetting zoom back to full screenshot.")
apply_crop(session_state, None)
continue
# --- TAB: Hide the labels temporarily ---
elif key == ord("\t"):
redraw_frame(session_state, display_tags = False)
redraw_frame(session_state, display_tags=False)
# --- ENTER (13 or 10): Next Frame / Save ---
elif key in [13, 10]:
if session_label == "background":
video_base = os.path.splitext(os.path.basename(selected_video))[0]
filename_base = f"{video_base}_frame_{current_frame}"
filename_base = get_current_filename_base(session_state)
image_save_path = os.path.join(images_dir, f"{filename_base}.jpg")
label_save_path = os.path.join(labels_dir, f"{filename_base}.txt")
cv2.imwrite(image_save_path, frame_resized)
cv2.imwrite(image_save_path, session_state["frame"])
open(label_save_path, "a").close()
print(f" Registered negative background frame: {filename_base}")
@@ -474,8 +606,6 @@ def label_video_session():
cv2.destroyAllWindows()
return
cv2.destroyWindow(window_name)
cap.release()
cv2.destroyAllWindows()
print(f"\nSession finished! Dataset updated successfully inside '{OUTPUT_DIRECTORY}'.")