我首先检查sources参数是目录还是文本文件 。 如果是一个目录 ， 我会读取目录中的所有内容（如果subdir_search参数为True ， 子目录也会包括在内） ， 否则我会读取文本文件中视频的路径 。 视频的路径存储在列表中 。 使用cur_pos以跟踪列表中的当前位置 。
该列表以batch_size为最大值进行迭代 ， 并检查以跳过错误视频或不存在的视频 。 它们被发送到letterbox函数 ， 以调整图像大小 。 这与原始版本相比没有任何变化 ， 除非所有视频都有故障/不可用 。
def letterbox(img, new_shape=(416, 416), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):# 将图像调整为32个像素倍数的矩形 /issues/232shape = img.shape[:2]# 当前形状 [height, width]if isinstance(new_shape, int):new_shape = (new_shape, new_shape)# 比例r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])if not scaleup:# 只按比例缩小 ， 不按比例放大（用于更好的测试图）r = min(r, 1.0)# 计算填充ratio = r, r# 宽高比new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]#填充if auto:# 最小矩形dw, dh = np.mod(dw, 64), np.mod(dh, 64)# 填充elif scaleFill:# 伸展dw, dh = 0.0, 0.0new_unpad = new_shaperatio = new_shape[0] / shape[1], new_shape[1] / shape[0]# 宽高比dw /= 2# 将填充分成两侧dh /= 2if shape[::-1] != new_unpad:# 改变大小img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))left, right = int(round(dw - 0.1)), int(round(dw + 0.1))img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=http://kandian.youth.cn/index/color)# 添加边界return img, ratio, (dw, dh)固定间隔检索帧函数update函数有一个小的变化 ， 我们另外存储了默认的图像大小 ， 以便在所有视频都被提取进行处理 ， 但由于长度不相等 ， 一个视频比另一个视频提前完成 。 当我解释代码的下一部分时 ， 它会更清楚 ， 那就是__next__ 函数 。
def update(self, index, cap, cur_pos):# 读取守护进程线程中的下一个帧n = 0while cap.isOpened():n += 1# _, self.imgs[index] = cap.read()cap.grab()if n == 4:# 每4帧读取一次_, self.data[index] = cap.retrieve()if self.def_img_size is None:self.def_img_size = self.data[index].shapen = 0time.sleep(0.01)# 等待迭代器如果帧存在 ， 它会像往常一样传递给letterbox函数 。 在frame为None的情况下 ， 这意味着视频已被完全处理 ， 我们检查列表中的所有视频是否都已被处理 。 如果有更多的视频要处理 ， cur_pos指针用于获取下一个可用视频的位置 。
如果不再从列表中提取视频 ， 但仍在处理某些视频 ， 则向其他处理组件发送一个空白帧 ， 即 ， 它根据其他批次中的剩余帧动态调整视频大小 。
def __next__(self):self.count += 1img0 = self.data.copy()img = []for i, x in enumerate(img0):if x is not None:img.append(letterbox(x, new_shape=self.img_size, auto=self.rect)[0])else:if self.cur_pos == self.n:if all( v is None for v in img0 ):cv2.destroyAllWindows()raise StopIterationelse:img0[i] = np.zeros(self.def_img_size)img.append(letterbox(img0[i], new_shape=self.img_size, auto=self.rect)[0])else:print('%g/%g: %s... ' % (self.cur_pos+1, self.n, self.sources[self.cur_pos]), end='')self.cap[i] = cv2.VideoCapture(self.sources[self.cur_pos])fldr_end_flg = 0while not self.cap[i].isOpened():print('Failed to open %s' % self.sources[self.cur_pos])self.cur_pos+=1if self.cur_pos == self.n:img0[i] = np.zeros(self.def_img_size)img.append(letterbox(img0[i], new_shape=self.img_size, auto=self.rect)[0])fldr_end_flg = 1breakself.cap[i] = cv2.VideoCapture(self.sources[self.cur_pos])if fldr_end_flg:continuew = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))fps = cap.get(cv2.CAP_PROP_FPS) % 100frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))_, self.data[i] = self.cap[i].read()# 保证第一帧img0[i] = self.data[i]img.append(letterbox(self.data[i], new_shape=self.img_size, auto=self.rect)[0])thread = Thread(target=self.update, args=([i, self.cap[i], self.cur_pos+1]), daemon=True)print(' success (%gx%g at %.2f FPS having %g frames).' % (w, h, fps, frames))self.cur_pos+=1thread.start()print('')# 新的一行# 堆叠img = np.stack(img, 0)# 转换img = img[:, :, :, ::-1].transpose(0, 3, 1, 2)# BGR 到 RGB, bsx3x416x416img = np.ascontiguousarray(img)return self.sources, img, img0, None