diff --git a/label_v4.py b/label_v4.py index 6154de7..b4d7904 100644 --- a/label_v4.py +++ b/label_v4.py @@ -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() + + # 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) - Calculates dynamic YOLO coordinates, saves the frame, and appends the - new dynamic coordinates to the YOLO annotation file. - """ +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)