标注工具--抹除目标

恩，chatgpt帮我写的标注抹除的工具

import cv2
import numpy as np
import argparse
import os
"""
左键拖动	框选区域
左键单击	粘贴选区
z	撤销上一次操作
c	取消当前选区
s	手动保存
n	自动保存当前 → 加载下一张（若存在 _edited.jpg 优先）
p	自动保存当前 → 加载上一张（若存在 _edited.jpg 优先）
ESC	自动保存后退出
"""

class RegionMoveAnnotator:
    def __init__(self, img_dir):
        self.img_dir = img_dir
        self.img_paths = sorted([os.path.join(img_dir, f) for f in os.listdir(img_dir)
                                 if f.lower().endswith((".jpg", ".png"))])
        if not self.img_paths:
            raise ValueError("❌ 没有找到图片，请检查路径")
        
        # 创建同级目录保存标注结果
        base_dir = os.path.dirname(os.path.abspath(img_dir))
        folder_name = os.path.basename(img_dir.rstrip("/"))
        self.save_dir = os.path.join(base_dir, f"{folder_name}_edited")
        os.makedirs(self.save_dir, exist_ok=True)

        self.index = 0
        self.img = None
        self.clone = None
        self.start_point = None
        self.end_point = None
        self.selecting = False
        self.selected_roi = None
        self.show_preview = False
        self.undo_stack = []
        self.mouse_pos = (0, 0)
        self.load_image()

    def get_save_path(self, index):
        base_name = os.path.basename(self.img_paths[index])
        name, ext = os.path.splitext(base_name)
        return os.path.join(self.save_dir, f"{name}_edited{ext}")

    def load_image(self):
        save_path = self.get_save_path(self.index)
        if os.path.exists(save_path):
            self.img = cv2.imread(save_path)
            print(f"📷 Loaded edited: {save_path}")
        else:
            self.img = cv2.imread(self.img_paths[self.index])
            print(f"📷 Loaded original: {self.img_paths[self.index]}")
        self.clone = self.img.copy()
        self.undo_stack = []
        self.selected_roi = None
        self.show_preview = False

    def save_current(self):
        save_path = self.get_save_path(self.index)
        cv2.imwrite(save_path, self.img)
        print(f"💾 Saved: {save_path}")

    def draw_preview(self, frame):
        if self.selected_roi is None or not self.show_preview:
            return frame

        overlay = frame.copy()
        h, w = self.selected_roi.shape[:2]
        x, y = self.mouse_pos
        x1, x2 = x, x + w
        y1, y2 = y, y + h

        if x1 < 0 or y1 < 0:
            return frame
        if x2 > frame.shape[1] or y2 > frame.shape[0]:
            return frame

        roi_h = min(h, frame.shape[0] - y1)
        roi_w = min(w, frame.shape[1] - x1)

        alpha = 0.5
        overlay[y1:y1+roi_h, x1:x1+roi_w] = (
            overlay[y1:y1+roi_h, x1:x1+roi_w] * (1 - alpha)
            + self.selected_roi[:roi_h, :roi_w] * alpha
        ).astype(np.uint8)

        return overlay

    def mouse_callback(self, event, x, y, flags, param):
        self.mouse_pos = (x, y)
        if event == cv2.EVENT_LBUTTONDOWN:
            if self.selected_roi is not None and self.show_preview:
                self.undo_stack.append(self.img.copy())
                h, w = self.selected_roi.shape[:2]
                y1, y2 = y, y + h
                x1, x2 = x, x + w
                y2 = min(y2, self.img.shape[0])
                x2 = min(x2, self.img.shape[1])
                self.img[y1:y2, x1:x2] = self.selected_roi[:y2 - y1, :x2 - x1]
                self.selected_roi = None
                self.show_preview = False
                self.clone = self.img.copy()
                cv2.imshow("Annotator", self.img)
                print("✅ Paste done")
            else:
                self.start_point = (x, y)
                self.selecting = True

        elif event == cv2.EVENT_MOUSEMOVE:
            if self.selecting:
                temp = self.clone.copy()
                cv2.rectangle(temp, self.start_point, (x, y), (0, 255, 0), 2)
                cv2.imshow("Annotator", temp)
            elif self.selected_roi is not None and self.show_preview:
                cv2.imshow("Annotator", self.draw_preview(self.clone.copy()))

        elif event == cv2.EVENT_LBUTTONUP and self.selecting:
            self.end_point = (x, y)
            self.selecting = False
            x1, y1 = self.start_point
            x2, y2 = self.end_point
            x1, x2 = sorted([x1, x2])
            y1, y2 = sorted([y1, y2])
            if x2 - x1 > 0 and y2 - y1 > 0:
                self.selected_roi = self.clone[y1:y2, x1:x2].copy()
                self.show_preview = True
                print("✅ Region selected, move mouse and click to paste.")

    def run(self):
        cv2.namedWindow("Annotator", cv2.WINDOW_NORMAL)
        cv2.setMouseCallback("Annotator", self.mouse_callback)
        cv2.imshow("Annotator", self.img)

        while True:
            if self.selected_roi is not None and self.show_preview:
                cv2.imshow("Annotator", self.draw_preview(self.clone.copy()))
            key = cv2.waitKey(10) & 0xFF

            if key == 27:  # ESC
                self.save_current()
                break
            elif key == ord('z'):
                if self.undo_stack:
                    self.img = self.undo_stack.pop()
                    self.clone = self.img.copy()
                    cv2.imshow("Annotator", self.img)
                    print("↩ Undo")
            elif key == ord('c'):
                if self.selected_roi is not None:
                    self.selected_roi = None
                    self.show_preview = False
                    cv2.imshow("Annotator", self.clone)
                    print("❎ Selection canceled")
            elif key == ord('s'):
                self.save_current()
            elif key == ord('n'):
                self.save_current()
                self.index = (self.index + 1) % len(self.img_paths)
                self.load_image()
                cv2.imshow("Annotator", self.img)
            elif key == ord('p'):
                self.save_current()
                self.index = (self.index - 1) % len(self.img_paths)
                self.load_image()
                cv2.imshow("Annotator", self.img)

        cv2.destroyAllWindows()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Label & Move Tool")
    parser.add_argument('--images','-i', required=True, help="Image file or folder")
    args = parser.parse_args()
    img_dir = args.images  # 原图目录
    annotator = RegionMoveAnnotator(img_dir)
    annotator.run()