def do(self, frame_id: int, frame: np.ndarray, *_,
**__) -> typing.Optional[np.ndarray]:
super().do(frame_id, frame, *_, **__)
black = np.zeros([*frame.shape, 3], np.uint8)
white = black + 255
black_ssim = toolbox.compare_ssim(black, frame)
white_ssim = toolbox.compare_ssim(white, frame)
logger.debug(f'black: {black_ssim}; white: {white_ssim}')
self.result[frame_id] = {
'black': black_ssim,
'white': white_ssim,
}
return
def _classify_frame(self,
frame: np.ndarray,
video_cap: cv2.VideoCapture,
threshold: float = None,
*_, **__) -> str:
if not threshold:
threshold = 0.85
frame = toolbox.compress_frame(
frame,
self.compress_rate,
self.target_size,
)
result = list()
for each_stage_name, each_stage_pic_list in self.read(video_cap):
each_result = list()
for target_pic in each_stage_pic_list:
target_pic = toolbox.compress_frame(target_pic, self.compress_rate, self.target_size)
each_pic_ssim = toolbox.compare_ssim(frame, target_pic)
each_result.append(each_pic_ssim)
ssim = max(each_result)
result.append((each_stage_name, ssim))
logger.debug(f'stage [{each_stage_name}]: {ssim}')
result = max(result, key=lambda x: x[1])
if result[1] < threshold:
logger.debug('not a known stage, set it -1')
result = ('-1', result[1])
return result[0]
def compare_frame_list(
self, src: typing.List[np.ndarray],
target: typing.List[np.ndarray]) -> typing.List[float]:
"""
core method about how to compare two lists of ndarray and get their ssim/mse/psnr
you can overwrite this method to implement your own algo
see https://github.com/williamfzc/stagesepx/issues/136
:param src:
:param target:
:return:
"""
# find the min ssim and the max mse / psnr
ssim = 1.0
mse = 0.0
psnr = 0.0
for part_index, (each_start, each_end) in enumerate(zip(src, target)):
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}"
)
return [ssim, mse, psnr]
def _compare_frame_list(
src: typing.List[np.ndarray],
target: typing.List[np.ndarray]) -> typing.List[float]:
# find the min ssim and the max mse / psnr
ssim = 1.0
mse = 0.0
psnr = 0.0
for part_index, (each_start,
each_end) in enumerate(zip(src, target)):
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}"
)
return [ssim, mse, psnr]
def diff(
self, another: 'BaseClassifier'
) -> typing.Dict[str, typing.Dict[str, float]]:
""" would be better to use 'diff' in VideoCutResult """
assert another._data, 'must load data first'
result_dict = dict()
self_data = dict(self.read())
for each_self_stage, each_self_data in self_data.items():
another_data = dict(another.read())
each_self_data_pic = next(each_self_data)
each_stage_dict = dict()
# compare with all the stages
for each_another_stage, each_another_data in another_data.items():
# compare with all the pictures in same stage, and pick the max one
max_ssim = -1
for each_pic in each_another_data:
ssim = toolbox.compare_ssim(
each_pic,
each_self_data_pic,
)
if ssim > max_ssim:
max_ssim = ssim
each_stage_dict[each_another_stage] = max_ssim
result_dict[each_self_stage] = each_stage_dict
return result_dict
def is_loop(self, threshold: float = None, **_) -> bool:
if not threshold:
threshold = constants.DEFAULT_THRESHOLD
operator = self.video.get_operator()
start_frame = operator.get_frame_by_id(self.start)
end_frame = operator.get_frame_by_id(self.end)
return toolbox.compare_ssim(start_frame.data, end_frame.data) > threshold
def _classify_frame(self,
frame: VideoFrame,
threshold: float = None,
*_,
**__) -> str:
if not threshold:
threshold = 0.85
result = list()
for each_stage_name, each_stage_pic_list in self.read():
each_result = list()
for target_pic in each_stage_pic_list:
# apply hooks
target_pic = self._apply_hook(VideoFrame(-1, -1.0, target_pic))
target_pic = target_pic.data
each_pic_ssim = toolbox.compare_ssim(frame.data, target_pic)
each_result.append(each_pic_ssim)
ssim = max(each_result)
result.append((each_stage_name, ssim))
logger.debug(f"stage [{each_stage_name}]: {ssim}")
result = max(result, key=lambda x: x[1])
if result[1] < threshold:
logger.debug("not a known stage, set it -1")
result = ("-1", result[1])
return result[0]
def _classify_frame(self,
frame: np.ndarray,
threshold: float = None,
*_,
**__) -> str:
if not threshold:
threshold = 0.85
frame = toolbox.compress_frame(frame)
result = list()
for each_stage_name, each_stage_pic_list in self.data.items():
each_result = list()
for each in each_stage_pic_list:
target_pic = cv2.imread(each.as_posix())
target_pic = toolbox.compress_frame(target_pic)
each_pic_ssim = toolbox.compare_ssim(frame, target_pic)
each_result.append(each_pic_ssim)
ssim = max(each_result)
result.append((each_stage_name, ssim))
logger.debug(f'stage [{each_stage_name}]: {ssim}')
result = max(result, key=lambda x: x[1])
if result[1] < threshold:
logger.debug('not a known stage, set it -1')
result = ('-1', result[1])
return result[0]
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 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 _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_range_list(self,
video: VideoObject,
block: int = None,
*args,
**kwargs) -> typing.List[VideoCutRange]:
range_list: typing.List[VideoCutRange] = list()
logger.info(
f"total frame count: {video.frame_count}, size: {video.frame_size}"
)
# load the first two frames
video_operator = video.get_operator()
cur_frame = video_operator.get_frame_by_id(1)
next_frame = video_operator.get_frame_by_id(1 + self.step)
# hook
cur_frame.data = self._apply_hook(cur_frame.frame_id, cur_frame.data)
# check block
if not block:
block = 2
if not self.is_block_valid(cur_frame.data, block):
logger.warning(
"array split does not result in an equal division, set block to 1"
)
block = 1
while True:
# hook
next_frame.data = self._apply_hook(next_frame.frame_id,
next_frame.data, *args,
**kwargs)
logger.debug(
f"computing {cur_frame.frame_id}({cur_frame.timestamp}) & {next_frame.frame_id}({next_frame.timestamp}) ..."
)
start_part_list = self.pic_split(cur_frame.data, block)
end_part_list = self.pic_split(next_frame.data, block)
# find the min ssim and the max mse / psnr
ssim = 1.0
mse = 0.0
psnr = 0.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 {cur_frame.frame_id} & {next_frame.frame_id}: ssim={ssim}; mse={mse}; psnr={psnr}"
)
range_list.append(
VideoCutRange(
video,
start=cur_frame.frame_id,
end=next_frame.frame_id,
ssim=[ssim],
mse=[mse],
psnr=[psnr],
start_time=cur_frame.timestamp,
end_time=next_frame.timestamp,
))
# load the next one
cur_frame = next_frame
next_frame = video_operator.get_frame_by_id(next_frame.frame_id +
self.step)
if next_frame is None:
break
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
