Added undo feature

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):
+def redraw_frame(param):
-    """Handles click-and-drag bounding box creation.
+    """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
+    # Draw all active boxes
-    new dynamic coordinates to the YOLO annotation file.
+    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
+        # Store the coordinates in our dynamic box list
-        width = x2 - x1
+        param["annotations"].append((
-        height = y2 - y1
+            param["class_id"],
-
-        # --- 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
+            x1, y1, x2, y2
-            cv2.FONT_HERSHEY_SIMPLEX,
+        ))
-            0.4,
+
-            (0, 255, 0),
+        # Save updates to disk and redraw
-            1,
+        save_annotations_to_file(param)
-        )
+        redraw_frame(param)
-        cv2.imshow(window_name, display_frame)
 
 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
+            "drawing": False,
-            "ix": -1,          # Starting X
+            "ix": -1,
-            "iy": -1,          # Starting Y
+            "iy": -1,
+            "annotations": [],  # Holds tuple info of drawn boxes for live Undo action
         }
 
         cv2.setMouseCallback(window_name, mouse_click_callback, session_state)