def load_frames(self, *args, **kwargs): logger.info(f"start loading {self.path} to memory ...") data: typing.List[VideoFrame] = [] with toolbox.video_capture(self.path) as cap: # the first success, frame = cap.read() while success: frame_object = VideoFrame.init(cap, frame) # apply hooks for each_hook in self._hook_list: frame_object = each_hook.do(frame_object, *args, **kwargs) data.append(frame_object) # read the next one success, frame = cap.read() # calculate memory cost each_cost = data[0].data.nbytes logger.debug(f"single frame cost: {each_cost} bytes") total_cost = each_cost * self.frame_count logger.debug(f"total frame cost: {total_cost} bytes") # lock the order self.data = tuple(data) # fix the length ( the last frame may be broken sometimes ) self.frame_count = len(data) # and size (reversed, see: https://github.com/williamfzc/stagesepx/issues/132) self.frame_size = data[0].data.shape[::-1] logger.info( f"frames loaded. frame count: {self.frame_count}, size: {self.frame_size}, memory cost: {total_cost} bytes" )
def contain_image(self, image_path: str = None, image_object: np.ndarray = None, *args, **kwargs) -> typing.Dict: assert image_path or image_object, 'should fill image_path or image_object' if image_path: logger.debug(f'found image path, use it first: {image_path}') assert os.path.isfile( image_path), f'image {image_path} not existed' image_object = toolbox.imread(image_path) image_object = toolbox.turn_grey(image_object) # TODO use client or itself..? fi = FindIt(engine=['template']) fi_template_name = 'default' fi.load_template(fi_template_name, pic_object=image_object) with toolbox.video_capture(self.video.path) as cap: target_id = self.pick(*args, **kwargs)[0] frame = toolbox.get_frame(cap, target_id) frame = toolbox.turn_grey(frame) result = fi.find(str(target_id), target_pic_object=frame) return result['data'][fi_template_name]['TemplateEngine']
def load_frames(self): # TODO full frames list can be very huge, for some devices logger.info(f"start loading {self.path} to memory ...") data: typing.List[VideoFrame] = [] with toolbox.video_capture(self.path) as cap: success, frame = cap.read() while success: frame_object = VideoFrame.init(cap, frame) data.append(frame_object) success, frame = cap.read() # calculate memory cost each_cost = data[0].data.nbytes logger.debug(f"single frame cost: {each_cost} bytes") total_cost = each_cost * self.frame_count logger.debug(f"total frame cost: {total_cost} bytes") logger.info( f"frames loaded. frame count: {self.frame_count}. memory cost: {total_cost} bytes" ) # lock the order self.data = tuple(data) # fix the length ( the last frame may be broken sometimes ) self.frame_count = len(data)
def pick_and_save(self, range_list: typing.List[VideoCutRange], frame_count: int, to_dir: str = None, compress_rate: float = None, *args, **kwargs) -> str: stage_list = list() for index, each_range in enumerate(range_list): picked = each_range.pick(frame_count, *args, **kwargs) logger.info(f'pick {picked} in range {each_range}') stage_list.append((index, picked)) # create parent dir if not to_dir: to_dir = toolbox.get_timestamp_str() os.makedirs(to_dir, exist_ok=True) for each_stage_id, each_frame_list in stage_list: # create sub dir each_stage_dir = os.path.join(to_dir, str(each_stage_id)) os.makedirs(each_stage_dir, exist_ok=True) with toolbox.video_capture(self.video_path) as cap: for each_frame_id in each_frame_list: each_frame_path = os.path.join(each_stage_dir, f'{uuid.uuid4()}.png') each_frame = toolbox.get_frame(cap, each_frame_id - 1) if compress_rate: each_frame = toolbox.compress_frame(each_frame, compress_rate) cv2.imwrite(each_frame_path, each_frame) logger.debug(f'frame [{each_frame_id}] saved to {each_frame_path}') return to_dir
def __init__( self, path: typing.Union[str, os.PathLike], pre_load: bool = None, fps: int = None, *_, **__, ): assert os.path.isfile(path), f"video [{path}] not existed" self.path: str = str(path) self.data: typing.Optional[typing.Tuple[VideoFrame]] = tuple() self._hook_list: typing.List["BaseHook"] = [] self.fps: int = fps if fps: video_path = os.path.join(tempfile.mkdtemp(), f"tmp_{fps}.mp4") logger.debug(f"convert video, and bind path to {video_path}") toolbox.fps_convert(fps, self.path, video_path, constants.FFMPEG) self.path = video_path with toolbox.video_capture(self.path) as cap: self.frame_count = toolbox.get_frame_count(cap) self.frame_size = toolbox.get_frame_size(cap) if pre_load: self.load_frames() logger.info( f"video object generated, length: {self.frame_count}, size: {self.frame_size}" )
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_stable_stage_sample( data_list: typing.List[ClassifierResult], *args, **kwargs ) -> np.ndarray: last = data_list[0] picked: typing.List[ClassifierResult] = [last] for each in data_list: # ignore unstable stage if each.stage == UNSTABLE_FLAG: continue if last.stage != each.stage: last = each picked.append(each) def get_split_line(f): return np.zeros((f.shape[0], 5)) with toolbox.video_capture(last.video_path) as cap: frame_list: typing.List[np.ndarray] = [] for each in picked: frame = toolbox.get_frame(cap, each.frame_id) frame = toolbox.compress_frame(frame, *args, **kwargs) split_line = get_split_line(frame) frame_list.append(frame) frame_list.append(split_line) return np.hstack(frame_list)
def contain_image(self, image_path: str = None, image_object: np.ndarray = None, threshold: float = None, *args, **kwargs): assert image_path or image_object, 'should fill image_path or image_object' if not threshold: threshold = 0.99 if image_path: logger.debug(f'found image path, use it first: {image_path}') assert os.path.isfile( image_path), f'image {image_path} not existed' image_object = cv2.imread(image_path) image_object = toolbox.turn_grey(image_object) # TODO use client or itself..? fi = FindIt(engine=['template']) fi_template_name = 'default' fi.load_template(fi_template_name, pic_object=image_object) with toolbox.video_capture(self.video.path) as cap: target_id = self.pick(*args, **kwargs)[0] frame = toolbox.get_frame(cap, target_id) frame = toolbox.turn_grey(frame) result = fi.find(str(target_id), target_pic_object=frame) find_result = result['data'][fi_template_name]['TemplateEngine'] position = find_result['target_point'] sim = find_result['target_sim'] logger.debug(f'position: {position}, sim: {sim}') return sim > threshold
def __init__(self, path: typing.Union[bytes, str, os.PathLike], *_, **__): assert os.path.isfile(path), f"video [{path}] not existed" self.path: str = str(path) self.data: typing.Optional[typing.Tuple[VideoFrame]] = tuple() with toolbox.video_capture(path) as cap: self.frame_count = toolbox.get_frame_count(cap) self.frame_size = toolbox.get_frame_size(cap)
def __init__(self, path: str = None, *_, **__): assert os.path.isfile(path), f"video [{path}] not existed" self.path = path self.data: typing.Tuple[VideoFrame] = tuple() with toolbox.video_capture(path) as cap: self.frame_count = toolbox.get_frame_count(cap) self.frame_size = toolbox.get_frame_size(cap)
def to_video_frame(self, *args, **kwargs) -> VideoFrame: # VideoFrame has `data` # SingleClassifierResult has `stage` with toolbox.video_capture(self.video_path) as cap: frame = toolbox.get_frame(cap, self.frame_id) compressed = toolbox.compress_frame(frame, *args, **kwargs) return VideoFrame(self.frame_id, self.timestamp, compressed)
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 is_loop(self, threshold: float = None, **_) -> bool: if not threshold: threshold = 0.95 with toolbox.video_capture(video_path=self.video_path) as cap: start_frame = toolbox.get_frame(cap, self.start) end_frame = toolbox.get_frame(cap, self.end) start_frame, end_frame = map(toolbox.compress_frame, (start_frame, end_frame)) return toolbox.compare_ssim(start_frame, end_frame) > threshold
def __init__(self, path: str = None, *_, **__): assert os.path.isfile(path), f'video [{path}] not existed' self.path = path with toolbox.video_capture(path) as cap: self.frame_count = toolbox.get_frame_count(cap) self.frame_size = toolbox.get_frame_size(cap)
def get_frames(self, frame_id_list: typing.List[int], *_, **__) -> typing.List[toolbox.VideoFrame]: """ return a list of VideoFrame, usually works with pick """ out = list() with toolbox.video_capture(self.video.path) as cap: for each_id in frame_id_list: timestamp = toolbox.get_frame_time(cap, each_id) frame = toolbox.get_frame(cap, each_id) out.append(toolbox.VideoFrame(each_id, timestamp, frame)) return out
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 load_frames(self): # TODO full frames list can be very huge data: typing.List[VideoFrame] = [] with toolbox.video_capture(self.path) as cap: success, frame = cap.read() while success: frame_object = VideoFrame.init(cap, frame) data.append(frame_object) success, frame = cap.read() # lock the order self.data = tuple(data) logger.info(f"load frames to memory. total frame count: {len(data)}")
def to_video_frame(self, *args, **kwargs) -> VideoFrame: # VideoFrame has `data` # SingleClassifierResult has `stage` (data is optional) # already have data if self.data is not None: return VideoFrame(self.frame_id, self.timestamp, self.data) # no data with toolbox.video_capture(self.video_path) as cap: frame = toolbox.get_frame(cap, self.frame_id) compressed = toolbox.compress_frame(frame, *args, **kwargs) return VideoFrame(self.frame_id, self.timestamp, compressed)
def pick_and_save(self, range_list: typing.List[VideoCutRange], frame_count: int, to_dir: str = None, # in kwargs # compress_rate: float = None, # target_size: typing.Tuple[int, int] = None, # to_grey: bool = None, *args, **kwargs) -> str: """ pick some frames from range, and save them as files :param range_list: VideoCutRange list :param frame_count: default to 3, and finally you will get 3 frames for each range :param to_dir: will saved to this path :param args: :param kwargs: :return: """ stage_list = list() for index, each_range in enumerate(range_list): picked = each_range.pick(frame_count, *args, **kwargs) logger.info(f'pick {picked} in range {each_range}') stage_list.append((index, picked)) # create parent dir if not to_dir: to_dir = toolbox.get_timestamp_str() os.makedirs(to_dir, exist_ok=True) for each_stage_id, each_frame_list in stage_list: # create sub dir each_stage_dir = os.path.join(to_dir, str(each_stage_id)) os.makedirs(each_stage_dir, exist_ok=True) with toolbox.video_capture(self.video_path) as cap: for each_frame_id in each_frame_list: each_frame_path = os.path.join(each_stage_dir, f'{uuid.uuid4()}.png') each_frame = toolbox.get_frame(cap, each_frame_id - 1) each_frame = toolbox.compress_frame(each_frame, **kwargs) cv2.imwrite(each_frame_path, each_frame) logger.debug(f'frame [{each_frame_id}] saved to {each_frame_path}') return to_dir
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 test_get_frame(): with toolbox.video_capture(VIDEO_PATH) as cap: first = 5 second = 8 actual = toolbox.get_frame_time(cap, first) should = toolbox.get_current_frame_time(cap) # should be frame 5 assert actual == should # 5 -> 8 -> 5 frame = toolbox.get_frame(cap, second, True) assert frame is not None assert toolbox.get_current_frame_id(cap) == first # 5 -> 8 frame = toolbox.get_frame(cap, second) assert frame is not None assert toolbox.get_current_frame_id(cap) == second
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_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
def _read_from_file(self) -> typing.Generator[VideoFrame, None, None]: with toolbox.video_capture(self.path) as cap: success, frame = cap.read() while success: yield VideoFrame.init(cap, frame) success, frame = cap.read()
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