Added undo feature
This commit is contained in:
+104
-68
@@ -26,6 +26,9 @@ CLASSIFICATION_OPTIONS = [
|
||||
]
|
||||
# ---------------------
|
||||
|
||||
# Label_v4_with_undo.py
|
||||
# ---------------------
|
||||
|
||||
def get_dataset_counts():
|
||||
"""Counts the number of existing bounding boxes by parsing YOLO text files."""
|
||||
counts = {option: 0 for option in CLASSIFICATION_OPTIONS}
|
||||
@@ -47,12 +50,74 @@ def get_dataset_counts():
|
||||
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.
|
||||
def redraw_frame(param):
|
||||
"""Regenerates the display frame using the current list of annotations."""
|
||||
# Start with a fresh, clean copy of the frame
|
||||
display_frame = param["frame"].copy()
|
||||
|
||||
Calculates dynamic YOLO coordinates, saves the frame, and appends the
|
||||
new dynamic coordinates to the YOLO annotation file.
|
||||
"""
|
||||
# Draw all active boxes
|
||||
for ann in param["annotations"]:
|
||||
class_id, label, x1, y1, x2, y2 = ann
|
||||
cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||||
cv2.putText(
|
||||
display_frame,
|
||||
label,
|
||||
(x1, max(15, y1 - 5)),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.4,
|
||||
(0, 255, 0),
|
||||
1,
|
||||
)
|
||||
|
||||
# Update state and window
|
||||
param["display_frame"] = display_frame
|
||||
cv2.imshow(param["window_name"], display_frame)
|
||||
|
||||
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']}"
|
||||
|
||||
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)
|
||||
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 frame image
|
||||
cv2.imwrite(image_save_path, param["frame"])
|
||||
|
||||
# Write/Overwrite the text file
|
||||
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
|
||||
|
||||
# YOLO Normalized Coordinate calculations
|
||||
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}")
|
||||
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."""
|
||||
if param["session_label"] == "background":
|
||||
if event == cv2.EVENT_LBUTTONDOWN:
|
||||
print(
|
||||
@@ -62,96 +127,65 @@ def mouse_click_callback(event, x, y, flags, param):
|
||||
|
||||
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 ---
|
||||
# --- Right Click: Undo Last Box ---
|
||||
if event == cv2.EVENT_RBUTTONDOWN:
|
||||
if param["annotations"]:
|
||||
removed = param["annotations"].pop()
|
||||
print(f" Undid last box ({removed[1]})")
|
||||
save_annotations_to_file(param)
|
||||
redraw_frame(param)
|
||||
else:
|
||||
print(" No boxes left to undo on this frame!")
|
||||
return
|
||||
|
||||
# --- Left Click Down: Start Dragging ---
|
||||
if event == cv2.EVENT_LBUTTONDOWN:
|
||||
param["drawing"] = True
|
||||
param["ix"], param["iy"] = x, y
|
||||
|
||||
# --- 2. Dynamic Box Preview during Mouse Drag ---
|
||||
# --- Mouse Move: Live Preview ---
|
||||
elif event == cv2.EVENT_MOUSEMOVE:
|
||||
if param["drawing"]:
|
||||
# Copy display_frame to wipe previous frame's dynamic preview lines
|
||||
# Wipe frame clean to draw fresh temporary drag 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 ---
|
||||
# --- Left Click Release: Lock Box & Save ---
|
||||
elif event == cv2.EVENT_LBUTTONUP:
|
||||
param["drawing"] = False
|
||||
ix, iy = param["ix"], param["iy"]
|
||||
|
||||
# Handle negative dragging (drawing from bottom-right to top-left)
|
||||
# Handle drawing in any physical direction
|
||||
x1, x2 = min(ix, x), max(ix, x)
|
||||
y1, y2 = min(iy, y), max(iy, y)
|
||||
|
||||
# Skip accidental tiny clicks (smaller than 5x5 pixels)
|
||||
# Ignore tiny accidental micro-clicks
|
||||
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
|
||||
# Keep box coordinates within 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,
|
||||
# Store the coordinates in our dynamic box list
|
||||
param["annotations"].append((
|
||||
param["class_id"],
|
||||
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)
|
||||
x1, y1, x2, y2
|
||||
))
|
||||
|
||||
# Save updates to disk and redraw
|
||||
save_annotations_to_file(param)
|
||||
redraw_frame(param)
|
||||
|
||||
def label_video_session():
|
||||
print("=========================================")
|
||||
@@ -232,7 +266,8 @@ def label_video_session():
|
||||
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(" [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(" [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.")
|
||||
@@ -266,7 +301,7 @@ def label_video_session():
|
||||
window_name = f"Labeling: {session_label} | Spot {i+1}/{num_frames} (Frame {current_frame})"
|
||||
cv2.namedWindow(window_name)
|
||||
|
||||
# Expanded state to hold drawing tracks
|
||||
# Expanded state tracking for multi-box lists
|
||||
session_state = {
|
||||
"frame": frame_resized,
|
||||
"display_frame": frame_display,
|
||||
@@ -278,9 +313,10 @@ def label_video_session():
|
||||
"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
|
||||
"drawing": False,
|
||||
"ix": -1,
|
||||
"iy": -1,
|
||||
"annotations": [], # Holds tuple info of drawn boxes for live Undo action
|
||||
}
|
||||
|
||||
cv2.setMouseCallback(window_name, mouse_click_callback, session_state)
|
||||
|
||||
Reference in New Issue
Block a user