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
