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 _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
