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
