def test_get_frame(): with toolbox.video_capture(VIDEO_PATH) as cap: first = 5 second = 8 toolbox.video_jump(cap, first) actual = toolbox.get_frame_time(cap, first) should = toolbox.get_current_frame_time(cap) # should be frame 5 assert actual == should assert actual - 0.16 < 0.01 # 5 -> 8 -> 5 frame = toolbox.get_frame(cap, second, True) assert frame is not None # grab, and recover # the next frame will be 5 # the current frame is 4 assert toolbox.get_current_frame_id(cap) == first - 1 # 5 -> 8 frame = toolbox.get_frame(cap, second) assert frame is not None assert toolbox.get_current_frame_id(cap) == second # cur_time = toolbox.get_current_frame_time(cap) toolbox.get_frame_time(cap, second, True) assert toolbox.get_current_frame_time(cap) == cur_time
def get_frame_by_id(self, frame_id: int) -> typing.Optional[VideoFrame]: if frame_id > self.get_length(): return None with toolbox.video_capture(self.video.path) as cap: toolbox.video_jump(cap, frame_id) success, frame = cap.read() video_frame = VideoFrame.init(cap, frame) if success else None return video_frame
def read_from_list(self, data: typing.List[int], video_cap: cv2.VideoCapture = None, *_, **__): cur_frame_id = toolbox.get_current_frame_id(video_cap) data = (toolbox.get_frame(video_cap, each) for each in data) toolbox.video_jump(video_cap, cur_frame_id) return data
def thumbnail( self, target_range: VideoCutRange, to_dir: str = None, compress_rate: float = None, is_vertical: bool = None, *_, **__, ) -> np.ndarray: """ build a thumbnail, for easier debug or something else :param target_range: VideoCutRange :param to_dir: your thumbnail will be saved to this path :param compress_rate: float, 0 - 1, about thumbnail's size, default to 0.1 (1/10) :param is_vertical: direction :return: """ if not compress_rate: compress_rate = 0.1 # direction if is_vertical: stack_func = np.vstack def get_split_line(f): return np.zeros((5, f.shape[1])) else: stack_func = np.hstack def get_split_line(f): return np.zeros((f.shape[0], 5)) frame_list = list() with toolbox.video_capture(self.video.path) as cap: toolbox.video_jump(cap, target_range.start) ret, frame = cap.read() count = 1 length = target_range.get_length() while ret and count <= length: frame = toolbox.compress_frame(frame, compress_rate) frame_list.append(frame) frame_list.append(get_split_line(frame)) ret, frame = cap.read() count += 1 merged = stack_func(frame_list) # create parent dir if to_dir: target_path = os.path.join( to_dir, f"thumbnail_{target_range.start}-{target_range.end}.png" ) cv2.imwrite(target_path, merged) logger.debug(f"save thumbnail to {target_path}") return merged
def classify(self, video_path: str, limit_range: typing.List[VideoCutRange] = None, step: int = None, *args, **kwargs) -> typing.List[ClassifierResult]: """ start classification :param video_path: path to target video :param limit_range: frames in these range will be ignored :param step: step between frames, default to 1 :param args: :param kwargs: :return: """ logger.debug(f'classify with {self.__class__.__name__}') if not step: step = 1 final_result: typing.List[ClassifierResult] = list() with toolbox.video_capture(video_path) as cap: ret, frame = cap.read() while ret: frame_id = toolbox.get_current_frame_id(cap) frame_timestamp = toolbox.get_current_frame_time(cap) if limit_range: if not any( [each.contain(frame_id) for each in limit_range]): logger.debug( f'frame {frame_id} ({frame_timestamp}) not in target range, skip' ) final_result.append( ClassifierResult(video_path, frame_id, frame_timestamp, '-1')) ret, frame = cap.read() continue # hook frame = self._apply_hook(frame_id, frame, *args, **kwargs) result = self._classify_frame(frame_id, frame, cap, *args, **kwargs) logger.debug( f'frame {frame_id} ({frame_timestamp}) belongs to {result}' ) final_result.append( ClassifierResult(video_path, frame_id, frame_timestamp, result)) toolbox.video_jump(cap, frame_id + step) ret, frame = cap.read() return final_result
def convert_video_into_ssim_list( self, video_path: str) -> typing.List[VideoCutRange]: ssim_list = list() with toolbox.video_capture(video_path) as cap: # get video info frame_count = toolbox.get_frame_count(cap) frame_size = toolbox.get_frame_size(cap) logger.debug( f'total frame count: {frame_count}, size: {frame_size}') # load the first two frames _, start = cap.read() start_frame_id = toolbox.get_current_frame_id(cap) toolbox.video_jump(cap, self.step) ret, end = cap.read() end_frame_id = toolbox.get_current_frame_id(cap) # compress start = toolbox.compress_frame(start, compress_rate=self.compress_rate) while ret: end = toolbox.compress_frame(end, compress_rate=self.compress_rate) ssim = toolbox.compare_ssim(start, end) logger.debug( f'ssim between {start_frame_id} & {end_frame_id}: {ssim}') ssim_list.append( VideoCutRange( video_path, start=start_frame_id, end=end_frame_id, ssim=ssim, )) # load the next one start = end start_frame_id, end_frame_id = end_frame_id, end_frame_id + self.step toolbox.video_jump(cap, end_frame_id) ret, end = cap.read() return ssim_list
def classify(self, video_path: str, limit_range: typing.List[VideoCutRange] = None, step: int = None, *args, **kwargs) -> typing.List[ClassifierResult]: logger.debug(f'classify with {self.__class__.__name__}') assert self.data, 'should load data first' if not step: step = 1 final_result: typing.List[ClassifierResult] = list() with toolbox.video_capture(video_path) as cap: ret, frame = cap.read() while ret: frame_id = toolbox.get_current_frame_id(cap) frame_timestamp = toolbox.get_current_frame_time(cap) if limit_range: if not any( [each.contain(frame_id) for each in limit_range]): logger.debug( f'frame {frame_id} ({frame_timestamp}) not in target range, skip' ) final_result.append( ClassifierResult(video_path, frame_id, frame_timestamp, '-1')) ret, frame = cap.read() continue result = self._classify_frame(frame, *args, **kwargs) logger.debug( f'frame {frame_id} ({frame_timestamp}) belongs to {result}' ) final_result.append( ClassifierResult(video_path, frame_id, frame_timestamp, result)) toolbox.video_jump(cap, frame_id + step - 1) ret, frame = cap.read() return final_result
def _convert_video_into_range_list(self, video: VideoObject, block: int = None, *args, **kwargs) -> typing.List[VideoCutRange]: if not block: block = 2 range_list: typing.List[VideoCutRange] = list() with toolbox.video_capture(video.path) as cap: logger.debug( f'total frame count: {video.frame_count}, size: {video.frame_size}' ) # load the first two frames _, start = cap.read() start_frame_id = toolbox.get_current_frame_id(cap) start_frame_time = toolbox.get_current_frame_time(cap) toolbox.video_jump(cap, self.step + 1) ret, end = cap.read() end_frame_id = toolbox.get_current_frame_id(cap) end_frame_time = toolbox.get_current_frame_time(cap) # hook start = self._apply_hook(start_frame_id, start) # check block if not self.is_block_valid(start, block): logger.warning( 'array split does not result in an equal division, set block to 1' ) block = 1 while ret: # hook end = self._apply_hook(end_frame_id, end, *args, **kwargs) logger.debug( f'computing {start_frame_id}({start_frame_time}) & {end_frame_id}({end_frame_time}) ...' ) start_part_list = self.pic_split(start, block) end_part_list = self.pic_split(end, block) # find the min ssim and the max mse / psnr ssim = 1. mse = 0. psnr = 0. for part_index, (each_start, each_end) in enumerate( zip(start_part_list, end_part_list)): part_ssim = toolbox.compare_ssim(each_start, each_end) if part_ssim < ssim: ssim = part_ssim # mse is very sensitive part_mse = toolbox.calc_mse(each_start, each_end) if part_mse > mse: mse = part_mse part_psnr = toolbox.calc_psnr(each_start, each_end) if part_psnr > psnr: psnr = part_psnr logger.debug( f'part {part_index}: ssim={part_ssim}; mse={part_mse}; psnr={part_psnr}' ) logger.debug( f'between {start_frame_id} & {end_frame_id}: ssim={ssim}; mse={mse}; psnr={psnr}' ) range_list.append( VideoCutRange( video, start=start_frame_id, end=end_frame_id, ssim=[ssim], mse=[mse], psnr=[psnr], start_time=start_frame_time, end_time=end_frame_time, )) # load the next one start = end start_frame_id, end_frame_id = end_frame_id, end_frame_id + self.step start_frame_time = end_frame_time toolbox.video_jump(cap, end_frame_id) ret, end = cap.read() end_frame_time = toolbox.get_current_frame_time(cap) return range_list
def convert_video_into_ssim_list(self, video_path: str, block: int = None, **kwargs) -> typing.List[VideoCutRange]: if not block: block = 1 ssim_list = list() with toolbox.video_capture(video_path) as cap: # get video info frame_count = toolbox.get_frame_count(cap) frame_size = toolbox.get_frame_size(cap) logger.debug(f'total frame count: {frame_count}, size: {frame_size}') # load the first two frames _, start = cap.read() start_frame_id = toolbox.get_current_frame_id(cap) start_frame_time = toolbox.get_current_frame_time(cap) toolbox.video_jump(cap, self.step + 1) ret, end = cap.read() end_frame_id = toolbox.get_current_frame_id(cap) end_frame_time = toolbox.get_current_frame_time(cap) # compress start = toolbox.compress_frame(start, **kwargs) # split func # width > height if frame_size[0] > frame_size[1]: split_func = np.hsplit else: split_func = np.vsplit logger.debug(f'split function: {split_func.__name__}') while ret: end = toolbox.compress_frame(end, **kwargs) ssim = 0 start_part_list = split_func(start, block) end_part_list = split_func(end, block) for part_index, (each_start, each_end) in enumerate(zip(start_part_list, end_part_list)): part_ssim = toolbox.compare_ssim(each_start, each_end) ssim += part_ssim logger.debug(f'part {part_index}: {part_ssim}') ssim = ssim / block logger.debug(f'ssim between {start_frame_id} & {end_frame_id}: {ssim}') ssim_list.append( VideoCutRange( video_path, start=start_frame_id, end=end_frame_id, ssim=[ssim], start_time=start_frame_time, end_time=end_frame_time, ) ) # load the next one start = end start_frame_id, end_frame_id = end_frame_id, end_frame_id + self.step start_frame_time = end_frame_time toolbox.video_jump(cap, end_frame_id) ret, end = cap.read() end_frame_time = toolbox.get_current_frame_time(cap) return ssim_list