Updated labeling script
This commit is contained in:
+349
@@ -0,0 +1,349 @@
|
||||
# Label_v4.py
|
||||
# ---------------------
|
||||
|
||||
|
||||
import glob
|
||||
import os
|
||||
import random
|
||||
import cv2
|
||||
|
||||
# --- CONFIGURATION ---
|
||||
OUTPUT_DIRECTORY = "dataset"
|
||||
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
|
||||
|
||||
# Define the EXACT size of the bounding box around your click point
|
||||
BOX_WIDTH = 64
|
||||
BOX_HEIGHT = 64
|
||||
|
||||
# Defined classification options (YOLO maps these to IDs: 0, 1, 2, 3...)
|
||||
CLASSIFICATION_OPTIONS = [
|
||||
"guardian",
|
||||
"interceptor",
|
||||
"sentinel",
|
||||
"rocket",
|
||||
"bomber",
|
||||
]
|
||||
# ---------------------
|
||||
|
||||
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):
|
||||
# Scan all text annotation files in the YOLO labels directory
|
||||
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 mouse_click_callback(event, x, y, flags, param):
|
||||
"""Handles click-and-drag bounding box creation.
|
||||
|
||||
Calculates dynamic YOLO coordinates, saves the frame, and appends the
|
||||
new dynamic coordinates to the YOLO annotation file.
|
||||
"""
|
||||
if param["session_label"] == "background":
|
||||
if event == cv2.EVENT_LBUTTONDOWN:
|
||||
print(
|
||||
" Info: In background mode, do not click targets. Just press SPACE to save frame."
|
||||
)
|
||||
return
|
||||
|
||||
frame = param["frame"]
|
||||
display_frame = param["display_frame"]
|
||||
frame_idx = param["frame_idx"]
|
||||
window_name = param["window_name"]
|
||||
|
||||
# --- 1. Start drawing on Left Mouse Click Down ---
|
||||
if event == cv2.EVENT_LBUTTONDOWN:
|
||||
param["drawing"] = True
|
||||
param["ix"], param["iy"] = x, y
|
||||
|
||||
# --- 2. Dynamic Box Preview during Mouse Drag ---
|
||||
elif event == cv2.EVENT_MOUSEMOVE:
|
||||
if param["drawing"]:
|
||||
# Copy display_frame to wipe previous frame's dynamic preview lines
|
||||
temp_frame = display_frame.copy()
|
||||
cv2.rectangle(temp_frame, (param["ix"], param["iy"]), (x, y), (0, 255, 0), 2)
|
||||
cv2.imshow(window_name, temp_frame)
|
||||
|
||||
# --- 3. Save Coordinates on Mouse Release ---
|
||||
elif event == cv2.EVENT_LBUTTONUP:
|
||||
param["drawing"] = False
|
||||
ix, iy = param["ix"], param["iy"]
|
||||
|
||||
# Handle negative dragging (drawing from bottom-right to top-left)
|
||||
x1, x2 = min(ix, x), max(ix, x)
|
||||
y1, y2 = min(iy, y), max(iy, y)
|
||||
|
||||
# Skip accidental tiny clicks (smaller than 5x5 pixels)
|
||||
width = x2 - x1
|
||||
height = y2 - y1
|
||||
if width < 5 or height < 5:
|
||||
cv2.imshow(window_name, display_frame)
|
||||
return
|
||||
|
||||
# Ensure box corners do not exceed image boundaries
|
||||
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))
|
||||
|
||||
# Recalculate dimensions after boundary clipping
|
||||
width = x2 - x1
|
||||
height = y2 - y1
|
||||
|
||||
# --- YOLO Normalized Coordinate Conversion ---
|
||||
true_center_x = (x1 + x2) / 2.0
|
||||
true_center_y = (y1 + y2) / 2.0
|
||||
|
||||
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
|
||||
|
||||
# Clean video name to ensure a unique, safe filename
|
||||
video_base = os.path.splitext(os.path.basename(param["video_name"]))[0]
|
||||
filename_base = f"{video_base}_frame_{frame_idx}"
|
||||
|
||||
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"
|
||||
)
|
||||
|
||||
# 1. Save the full frame image
|
||||
cv2.imwrite(image_save_path, frame)
|
||||
|
||||
# 2. Append the target annotation to the YOLO text file
|
||||
class_id = param["class_id"]
|
||||
with open(label_save_path, "a") as f:
|
||||
f.write(
|
||||
f"{class_id} {x_center_norm:.6f} {y_center_norm:.6f} {width_norm:.6f} {height_norm:.6f}\n"
|
||||
)
|
||||
|
||||
print(
|
||||
f" Logged {param['session_label']} box coordinate to {label_save_path}"
|
||||
)
|
||||
param["box_drawn"] = True
|
||||
|
||||
# Draw permanent visual feedback box on the master display frame
|
||||
cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||||
cv2.putText(
|
||||
display_frame,
|
||||
param["session_label"],
|
||||
(x1, max(15, y1 - 5)), # Prevents text from drawing off-screen at the top
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.4,
|
||||
(0, 255, 0),
|
||||
1,
|
||||
)
|
||||
cv2.imshow(window_name, display_frame)
|
||||
|
||||
def label_video_session():
|
||||
print("=========================================")
|
||||
print(" AI YOLO DATASET LABELER ")
|
||||
print("=========================================\n")
|
||||
|
||||
# 1. Automatically find and select a random MP4 video
|
||||
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}'")
|
||||
|
||||
# 2. Fetch and display current dataset balance dashboard
|
||||
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(
|
||||
"💡 Tip: Try to choose options with lower counts to keep data balanced!"
|
||||
)
|
||||
|
||||
# 3. Force selection of a valid label
|
||||
print("\nAvailable Classification Targets:")
|
||||
for idx, option in enumerate(CLASSIFICATION_OPTIONS, start=1):
|
||||
print(f" [{idx}] {option} (Current Boxes: {counts[option]})")
|
||||
|
||||
while True:
|
||||
try:
|
||||
choice = input(
|
||||
f"\nEnter the number (1-{len(CLASSIFICATION_OPTIONS)}) of what you are labeling: "
|
||||
).strip()
|
||||
choice_idx = int(choice) - 1
|
||||
if 0 <= choice_idx < len(CLASSIFICATION_OPTIONS):
|
||||
session_label = CLASSIFICATION_OPTIONS[choice_idx]
|
||||
session_class_id = choice_idx
|
||||
break
|
||||
else:
|
||||
print("Invalid selection. Please choose a number from the list.")
|
||||
except ValueError:
|
||||
print("Invalid input. Please enter a number.")
|
||||
|
||||
# Open video and check validity
|
||||
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
|
||||
|
||||
# Select random unique anchor frames to jump between
|
||||
num_frames = min(FRAMES_TO_EXTRACT, total_frames)
|
||||
frame_indices = random.sample(range(total_frames), num_frames)
|
||||
frame_indices.sort()
|
||||
|
||||
# Set up standard YOLO dataset directory branches
|
||||
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)
|
||||
|
||||
# Comprehensive User Instructions
|
||||
print("\n" + "=" * 50)
|
||||
print(" HOW TO LABEL")
|
||||
print("=" * 50)
|
||||
print(f"• TARGET OBJECT: {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 frame on release).")
|
||||
print(" [Right Arrow] - Move FORWARD 10 frames (sequential tracking, no save).")
|
||||
print(" [Left Arrow] - Move BACKWARD 10 frames (sequential tracking, no save).")
|
||||
print(" [SPACE / ENTER] - Jump completely to the NEXT random video search spot.")
|
||||
print(" [Q Key] - Quit and save all progress up to this point.")
|
||||
print("=" * 50)
|
||||
input("\nPress ENTER when you are ready to begin...")
|
||||
|
||||
# Cross-platform arrow key mappings for waitKeyEx()
|
||||
LEFT_KEYS = [81, 2, 2424832, 65361, 63234]
|
||||
RIGHT_KEYS = [83, 3, 2490368, 65363, 63235]
|
||||
|
||||
i = 0
|
||||
current_frame = frame_indices[i]
|
||||
|
||||
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
|
||||
|
||||
# Resize the frame by SCALE_FACTOR for laptop viewability
|
||||
frame_resized = cv2.resize(
|
||||
frame, (0, 0), fx=SCALE_FACTOR, fy=SCALE_FACTOR
|
||||
)
|
||||
frame_display = frame_resized.copy()
|
||||
|
||||
window_name = f"Labeling: {session_label} | Spot {i+1}/{num_frames} (Frame {current_frame})"
|
||||
cv2.namedWindow(window_name)
|
||||
|
||||
# Expanded state to hold drawing tracks
|
||||
session_state = {
|
||||
"frame": frame_resized,
|
||||
"display_frame": frame_display,
|
||||
"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,
|
||||
"box_drawn": False,
|
||||
"drawing": False, # True when dragging mouse
|
||||
"ix": -1, # Starting X
|
||||
"iy": -1, # Starting Y
|
||||
}
|
||||
|
||||
cv2.setMouseCallback(window_name, mouse_click_callback, session_state)
|
||||
|
||||
# Keep window open until user acts
|
||||
cv2.imshow(window_name, frame_display)
|
||||
while True:
|
||||
cv2.moveWindow(window_name, 0, 0)
|
||||
key = cv2.waitKeyEx(1)
|
||||
|
||||
# Enter (13) or Space (32) -> Jump to the next random anchor spot
|
||||
if key == 13 or key == 32:
|
||||
if session_label == "background":
|
||||
video_base = os.path.splitext(
|
||||
os.path.basename(selected_video)
|
||||
)[0]
|
||||
filename_base = f"{video_base}_frame_{current_frame}"
|
||||
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)
|
||||
open(label_save_path, "a").close() # Creates empty file
|
||||
print(f" Registered negative background frame: {filename_base}")
|
||||
|
||||
i += 1
|
||||
if i < len(frame_indices):
|
||||
current_frame = frame_indices[i]
|
||||
break
|
||||
|
||||
# Right Arrow -> Move 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 -> Move 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
|
||||
|
||||
# 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
|
||||
|
||||
cv2.destroyWindow(window_name)
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
print(f"\nSession finished! Dataset updated successfully inside '{OUTPUT_DIRECTORY}'.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
label_video_session()
|
||||
Reference in New Issue
Block a user