【数据集标注制作】视频剪切标注1——类DarkLabel软件

视频标注

用于从视频中标注数据，用于YOLO网络的目标检测。旨在实现单次鼠标标注能生成多张被标注图像，实现数据集快速制作！

1. 从视频中，通过鼠标框选指定区域，形成掩码box
2. 鼠标选定区域后，根据设定的成像尺寸，在选定区域周围随机实施多个角度的剪切，制作数据集。（重在减少鼠标标注操作，丰富数据集图片数量）
3. 剪切出的图片，每张都带掩码box、原图和效果图。
4. 标注过程视频可暂停，生成图像文件自动命名。（鼠标单击过后，或者键盘按键，都可控制视频暂停和继续）
5. 不足之处：不可控制视频帧移动

程序文件：https://download.csdn.net/download/tjb132/88514408?spm=1001.2014.3001.5503

import threading
import queue

class Mask_video:
    def __init__(self, urls, label, images_save_path):
        self.urls = urls
        self.video_idx = 0

        self.lock = threading.Lock()
        self.drawing = False
        self.start_point = (-1, -1)
        self.end_point = (-1, -1)
        self.cropped_image = None
        self.paused = False  # 是否暂停视频播放

        self.expanded_cropped_image = None
        self.expanded_cropped_image_mask_pts = [0,0,0,0]

        self.images_save_root = images_save_path
        self.images_data_save_path = os.path.join(self.images_save_root, 'images')
        self.images_label_save_path = os.path.join(self.images_save_root, 'labels')
        self.images_label_mask_path = os.path.join(self.images_save_root, 'mask')
        os.makedirs(self.images_data_save_path, exist_ok=True)
        os.makedirs(self.images_label_save_path, exist_ok=True)
        os.makedirs(self.images_label_mask_path, exist_ok=True)
        self.images_idx = 0
        self.label = label


        # 创建一个窗口并设置鼠标事件回调函数
        cv.namedWindow('Video')
        cv.setMouseCallback('Video', self.draw_rectangle)

    pass

    def draw_rectangle(self, event, x, y, flags, param):
        if event == cv.EVENT_LBUTTONDOWN:
            with self.lock:
                self.drawing = True
                self.start_point = (x, y)
            self.pause_video()

        elif event == cv.EVENT_MOUSEMOVE:
            if self.drawing:
                with self.lock:
                    self.end_point = (x, y)

        if event == cv.EVENT_LBUTTONUP:
            with self.lock:
                self.drawing = False
                self.end_point = (x, y)
                # print('=== 1 ===',  self.start_point, self.end_point)
                # 归一化坐标，确保 start_point 包含左上角坐标，end_point 包含右下角坐标
                start_point = (
                min(self.start_point[0], self.end_point[0]), min(self.start_point[1], self.end_point[1]))
                end_point = (
                max(self.start_point[0], self.end_point[0]), max(self.start_point[1], self.end_point[1]))
                self.start_point, self.end_point = start_point, end_point
                # print('=== 2 ===', self.start_point, self.end_point)
                if (self.start_point[0]-self.end_point[0])==0 or (self.start_point[1]-self.end_point[1])==0:
                    self.cropped_image = None
                    return
                else:
                    # 裁剪图像并显示
                    self.cropped_image = self.frame[self.start_point[1]:self.end_point[1],
                                         self.start_point[0]:self.end_point[0]]
                    # cv.imshow('Cropped Image', self.cropped_image)

            for i in range(4): self.crop_and_random_expand()
            self.cropped_image = None

            self.resume_video()

    def pause_video(self):
        with self.lock:
            self.paused = True

    def resume_video(self):
        with self.lock:
            self.paused = False

    def crop_and_random_expand(self):
        """ 在指定区域的附近，实施随机剪裁，生成图像 """
        # with self.lock:
        if self.cropped_image is not None:
            # 定义扩展的像素范围
            expand_range1 = np.random.randint(0, self.start_point[0])  # 您可以根据需要调整这个值
            expanded_x1 = max(self.start_point[0] - expand_range1, 0)
            expand_range2 = np.random.randint(0, self.frame.shape[1]-self.end_point[0])  # 您可以根据需要调整这个值
            expanded_x2 = min(self.end_point[0] + expand_range2, self.frame.shape[1])
            expand_range3 = np.random.randint(0, self.start_point[1])  # 您可以根据需要调整这个值
            expanded_y1 = max(self.start_point[1] - expand_range3, 0)
            expand_range4 = np.random.randint(0, self.frame.shape[0] - self.end_point[1])  # 您可以根据需要调整这个值
            expanded_y2 = min(self.end_point[1] + expand_range4, self.frame.shape[0])
            expanded_cropped_image = self.frame[expanded_y1:expanded_y2, expanded_x1:expanded_x2]
            yh1, xw1 = expanded_cropped_image.shape[:2]
            expanded_cropped_image = cv.resize(expanded_cropped_image, (640, 640))
            yh2, xw2 = expanded_cropped_image.shape[:2]
            self.expanded_cropped_image = expanded_cropped_image.copy()

            new_pts = [expand_range1, expand_range3, self.end_point[0]-self.start_point[0], self.end_point[1]-self.start_point[1]]
            new_pts = [new_pts[0]*xw2/xw1, new_pts[1]*yh2/yh1, new_pts[2]*xw2/xw1, new_pts[3]*yh2/yh1]
            new_pts = np.array(new_pts, dtype=np.int32)

            self.expanded_cropped_image_mask_pts = new_pts
            cv.rectangle(expanded_cropped_image, (new_pts[0],new_pts[1]), (new_pts[0]+new_pts[2], new_pts[1]+new_pts[3]), (0, 255, 0), 2)
            cv.imshow('expanded_cropped_image', expanded_cropped_image)
            self.save_image(self.expanded_cropped_image, expanded_cropped_image)
            return expanded_cropped_image
        else:
            return None

    def run_video_crop(self):
        while True:
            if not self.paused:
                ret, self.frame = self.cap.read()

                if not ret:
                    print("无法读取视频帧")
                    break
            img = self.frame.copy()
            if self.start_point != (-1, -1) and self.end_point != (-1, -1):
                # 在帧上绘制方框
                with self.lock:
                    cv.rectangle(img, self.start_point, self.end_point, (0, 255, 0), 2)

            cv.imshow('Video', img)

            key = cv.waitKey(20)
            if key & 0xFF == ord('q'):  # 退出剪裁软件
                break
            elif key & 0xFF == ord('s'):   # 暂停视频
                self.pause_video()
            elif key & 0xFF == ord('d'):  # 继续播放视频
                self.resume_video()
            # cv.waitKey(20)

        self.cap.release()
        # cv.destroyAllWindows()

    def run(self):
        for i in range(len(self.urls)):
            url = self.urls[i]
            self.video_idx = i
            try:
                cap = cv.VideoCapture(url)
                ret, frame = cap.read()
                if ret:
                    self.cap = cap
                    # break
                    self.lock = threading.Lock()
                    self.drawing = False
                    self.start_point = (-1, -1)
                    self.end_point = (-1, -1)
                    self.cropped_image = None
                    self.paused = False  # 是否暂停视频播放
                    self.run_video_crop()
            except: pass
        cv.destroyAllWindows()

    def save_image(self, imgdata, maskdata):
        self.images_idx += 1
        imgfile = os.path.join(self.images_data_save_path, f"{str(self.video_idx)}_{self.images_idx}.jpg")
        labelfile = os.path.join(self.images_label_save_path, f"{str(self.video_idx)}_{self.images_idx}.txt")
        maskfile = os.path.join(self.images_label_mask_path, f"{str(self.video_idx)}_{self.images_idx}.jpg")

        print('\timages data jpg save in:', imgfile)
        cv.imwrite(imgfile, imgdata)
        print('\timages label txt save in:', labelfile)
        hy, wx = imgdata.shape[:2]
        x,y,w,h = self.expanded_cropped_image_mask_pts
        with open(labelfile, 'w') as f:
            data = f"{str(self.label)}\t{x / wx}\t{y / hy}\t{w / wx}\t{h / hy}" + "\n"
            f.write(data)
        print('\timages mask save in:', maskfile)
        cv.imwrite(maskfile, maskdata)






def run():
    print('==========   start system  ==============')
    import pandas as pd

    # 读取Excel文件
    excel_file = r'I:\python\02-job\h03090 data-output\video.xlsx'  # 将文件名替换为实际的Excel文件名
    df = pd.read_excel(io=excel_file)

    # 提取某一行的数据，例如第3行（索引为2）
    row_index = 3
    selected_row = df.iloc[row_index]

    # 打印提取的行数据
    print("提取的行数据:")
    print(selected_row)

    videos = str(selected_row['topVideo']).split(',')
    print(videos)
    data = {'id': selected_row['goods_id'], 'videos': videos}
    images_save_path = r'I:\python\02-job\h03090 data-output\new_images-labeled'
    mask = Mask_video(urls=videos, label=row_index, images_save_path=images_save_path)
    mask.run()
    pass