【数据集标注制作】视频剪切标注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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207