616 lines
24 KiB
Python
616 lines
24 KiB
Python
# Label_v5.py
|
|
# ---------------------
|
|
# 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
|
|
import random
|
|
import cv2
|
|
|
|
# --- CONFIGURATION ---
|
|
OUTPUT_DIRECTORY = "dataset_v5"
|
|
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 (Cyan, Magenta, Yellow, Red, Green, Orange, Purple)
|
|
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_dataset_counts():
|
|
"""Counts the number of existing bounding 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()
|
|
if parts:
|
|
class_id = int(parts[0])
|
|
if 0 <= class_id < len(CLASSIFICATION_OPTIONS):
|
|
option = CLASSIFICATION_OPTIONS[class_id]
|
|
counts[option] += 1
|
|
except Exception:
|
|
pass # Skip corrupted or unreadable text files safely
|
|
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 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"] = []
|
|
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}"
|
|
)
|
|
param["annotations"].append((class_id, label, x1, y1, x2, y2))
|
|
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"] = 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"
|
|
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"
|
|
|
|
# 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,
|
|
(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 transparent box overlays."""
|
|
h_frame, w_frame = param["frame"].shape[:2]
|
|
|
|
# --- 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)]
|
|
|
|
# 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
|
|
|
|
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.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 in the top padded area
|
|
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 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 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}")
|
|
param["box_drawn"] = False
|
|
else:
|
|
# Save the ORIGINAL un-scaled working frame (keeps dataset images small & standard)
|
|
cv2.imwrite(image_save_path, param["frame"])
|
|
|
|
# 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"]:
|
|
class_id, _, x1, y1, x2, y2 = ann
|
|
|
|
true_center_x = (x1 + x2) / 2.0
|
|
true_center_y = (y1 + y2) / 2.0
|
|
width = x2 - x1
|
|
height = y2 - y1
|
|
|
|
x_center_norm = true_center_x / img_w
|
|
y_center_norm = true_center_y / img_h
|
|
width_norm = width / img_w
|
|
height_norm = 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" 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 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
|
|
|
|
display_frame = param["display_frame"]
|
|
window_name = param["window_name"]
|
|
|
|
# --- Right Click: Undo Last Box ---
|
|
if event == cv2.EVENT_RBUTTONDOWN:
|
|
if param["annotations"]:
|
|
removed = param["annotations"].pop()
|
|
print(f" Undid last box ([{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: Start Label Box Dragging ---
|
|
if event == cv2.EVENT_LBUTTONDOWN:
|
|
param["drawing"] = True
|
|
param["win_ix"], param["win_iy"] = x, y
|
|
|
|
# --- Mouse Move: Live Preview in UI Space ---
|
|
elif event == cv2.EVENT_MOUSEMOVE:
|
|
if param["drawing"] or param["zooming"]:
|
|
temp_frame = display_frame.copy()
|
|
overlay = temp_frame.copy()
|
|
|
|
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)
|
|
|
|
# --- 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)
|
|
|
|
if (x2 - x1) < 10 or (y2 - y1) < 10:
|
|
cv2.imshow(window_name, display_frame)
|
|
return
|
|
|
|
# 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)
|
|
|
|
param["annotations"].append(
|
|
(param["class_id"], param["session_label"], x1, y1, x2, y2)
|
|
)
|
|
|
|
save_annotations_to_file(param)
|
|
redraw_frame(param)
|
|
|
|
|
|
def label_video_session():
|
|
print("=========================================")
|
|
print(" AI YOLO DATASET LABELER v5.0 ")
|
|
print("=========================================\n")
|
|
|
|
mp4_files = glob.glob("*.mp4")
|
|
if not mp4_files:
|
|
print(
|
|
"Error: No .mp4 files found in the current directory.\n"
|
|
"Please place this script in the same folder as your video files."
|
|
)
|
|
return
|
|
|
|
selected_video = random.choice(mp4_files)
|
|
print(f"-> Selected Video File: '{selected_video}'")
|
|
|
|
counts = get_dataset_counts()
|
|
print("\n-----------------------------------------")
|
|
print(" CURRENT DATASET BALANCE STATUS ")
|
|
print("-----------------------------------------")
|
|
for option, count in counts.items():
|
|
print(f" • {option.capitalize():<12} : {count} boxes logged")
|
|
print("-----------------------------------------")
|
|
|
|
print("\nAvailable Classification Modes & Targets:")
|
|
print(" [0] ALL-IN-ONE MODE (Tag multiple classes per frame; cycle with SPACE)")
|
|
for idx, option in enumerate(CLASSIFICATION_OPTIONS, start=1):
|
|
print(f" [{idx}] {option} only (Current Boxes: {counts[option]})")
|
|
|
|
while True:
|
|
try:
|
|
choice = input(
|
|
f"\nEnter mode selection (0-{len(CLASSIFICATION_OPTIONS)}): "
|
|
).strip()
|
|
choice_idx = int(choice)
|
|
if choice_idx == 0:
|
|
multi_mode = True
|
|
session_class_id = 0
|
|
session_label = CLASSIFICATION_OPTIONS[0]
|
|
break
|
|
elif 1 <= choice_idx <= len(CLASSIFICATION_OPTIONS):
|
|
multi_mode = False
|
|
session_class_id = choice_idx - 1
|
|
session_label = CLASSIFICATION_OPTIONS[session_class_id]
|
|
break
|
|
else:
|
|
print("Invalid selection. Please choose a number from the list.")
|
|
except ValueError:
|
|
print("Invalid input. Please enter a valid number.")
|
|
|
|
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))
|
|
if total_frames == 0:
|
|
print("Error: Video has 0 frames.")
|
|
return
|
|
|
|
num_frames = min(FRAMES_TO_EXTRACT, total_frames)
|
|
frame_indices = random.sample(range(total_frames), num_frames)
|
|
frame_indices.sort()
|
|
|
|
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 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("-" * 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]
|
|
RIGHT_KEYS = [83, 3, 2490368, 65363, 63235]
|
|
|
|
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()
|
|
if not ret:
|
|
print(f" Warning: Could not read frame {current_frame}. Moving to next random anchor.")
|
|
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)
|
|
frame_display = frame_resized.copy()
|
|
|
|
cv2.setWindowTitle(window_name, f"YOLO Labeler v5 | Spot {i+1}/{num_frames} (Frame {current_frame})")
|
|
|
|
session_state = {
|
|
"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,
|
|
"window_name": window_name,
|
|
"video_name": selected_video,
|
|
"class_id": session_class_id,
|
|
"session_label": session_label,
|
|
"multi_mode": multi_mode,
|
|
"box_drawn": False,
|
|
"drawing": False,
|
|
"zooming": False,
|
|
"win_ix": -1,
|
|
"win_iy": -1,
|
|
"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)
|
|
|
|
# --- SPACEBAR (32): Cycle through classes ---
|
|
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"]]
|
|
print(f" Switched target class to: [{session_state['class_id']}] {session_state['session_label'].upper()}")
|
|
redraw_frame(session_state)
|
|
continue
|
|
|
|
# --- 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)
|
|
|
|
# --- ENTER (13 or 10): Next Frame / Save ---
|
|
elif key in [13, 10]:
|
|
if session_label == "background":
|
|
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, session_state["frame"])
|
|
open(label_save_path, "a").close()
|
|
print(f" Registered negative background frame: {filename_base}")
|
|
|
|
i += 1
|
|
if i < len(frame_indices):
|
|
current_frame = frame_indices[i]
|
|
break
|
|
|
|
# --- RIGHT ARROW: Step Forward ---
|
|
elif key in RIGHT_KEYS:
|
|
if current_frame < total_frames - 10:
|
|
current_frame += 10
|
|
else:
|
|
print(" Already at the very last frame of the video!")
|
|
continue
|
|
break
|
|
|
|
# --- LEFT ARROW: Step Backward ---
|
|
elif key in LEFT_KEYS:
|
|
if current_frame > 0:
|
|
current_frame -= 10
|
|
else:
|
|
print(" Already at the first frame of the video!")
|
|
continue
|
|
break
|
|
|
|
# --- Q KEY: Quit ---
|
|
elif key == ord("q") or key == ord("Q") or (key & 0xFF) in [ord("q"), ord("Q")]:
|
|
print("\nQuitting session early... Saving progress.")
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
return
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
print(f"\nSession finished! Dataset updated successfully inside '{OUTPUT_DIRECTORY}'.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
label_video_session()
|