def test_list(self):
with deepstar_path():
VideoModel().insert('test1', 'test2')
FrameSetModel().insert(1)
frame_model = FrameModel()
frame_model.insert(1, 0)
frame_model.insert(1, 1)
frame_model.insert(1, 0)
result = frame_model.list(1)
self.assertEqual(len(result), 3)
self.assertEqual(result[0], (1, 1, 0))
self.assertEqual(result[1], (2, 1, 1))
self.assertEqual(result[2], (3, 1, 0))
result = frame_model.list(1, length=2)
self.assertEqual(len(result), 2)
self.assertEqual(result[0], (1, 1, 0))
self.assertEqual(result[1], (2, 1, 1))
result = frame_model.list(1, offset=1)
self.assertEqual(len(result), 2)
self.assertEqual(result[0], (2, 1, 1))
self.assertEqual(result[1], (3, 1, 0))
result = frame_model.list(1, length=1, offset=1)
self.assertEqual(len(result), 1)
self.assertEqual(result[0], (2, 1, 1))
result = frame_model.list(1, rejected=False)
self.assertEqual(len(result), 2)
self.assertEqual(result[0], (1, 1, 0))
self.assertEqual(result[1], (3, 1, 0))
def test_fk_frame_sets_on_delete_cascade(self):
with deepstar_path():
frame_set_model = FrameSetModel()
frame_set_model.insert(1)
frame_model = FrameModel()
frame_model.insert(1, 0)
self.assertEqual(frame_model.list(1), [(1, 1, 0)])
frame_set_model.delete(1)
self.assertIsNone(frame_model.list(1))
def select_export_dir(self, frame_set_ids, target_dir, opts={}):
"""
This method exports frame sets to a directory.
:param list frame_set_ids: The list of frame set IDs to export.
:param str target_dir: The directory to which to export the frame sets.
:param dict opts: The dict of options.
:raises: CommandLineRouteHandlerError
:rtype: None
"""
frame_set_model = FrameSetModel()
for frame_set_id in frame_set_ids:
result = frame_set_model.select(frame_set_id)
if result is None:
raise CommandLineRouteHandlerError(
f'Frame set with ID {frame_set_id:08d} not found')
frame_model = FrameModel()
length = int(os.environ.get('MODEL_LIST_LENGTH', '100'))
count = 0
for fid in frame_set_ids:
offset = 0
frame_set_path = FrameSetSubDir.path(fid)
while True:
results = frame_model.list(fid,
length=length,
offset=offset,
rejected=False)
if not results:
break
for frame_id, _, _ in results:
file_path = FrameFile().path(frame_set_path, frame_id,
'jpg')
if 'format' in opts:
filename = opts['format'] % frame_id
else:
filename = os.path.basename(file_path)
target_path = os.path.join(target_dir, filename)
shutil.copy(file_path, target_path)
count += 1
debug(
f'Frame with ID {frame_id:08d} at {file_path} '
f'exported to {target_path}', 4)
offset += length
debug(f'{count} frames were successfully exported to {target_dir}', 3)
def list(self, frame_set_id):
"""
This method lists frames in the frame collection for a frame set.
:param int frame_set_id: The frame set ID.
:raises: CommandLineRouteHandlerError
:rtype: None
"""
result = FrameSetModel().select(frame_set_id)
if result is None:
raise CommandLineRouteHandlerError(
f'Frame set with ID {frame_set_id:08d} not found')
frame_model = FrameModel()
count = frame_model.count(frame_set_id)
debug(f'{count} results', 3)
debug('id | fk_frame_sets | rejected | (width | height)', 3)
debug('------------------------------------------------', 3)
if count == 0:
return
length = int(os.environ.get('MODEL_LIST_LENGTH', '100'))
offset = 0
p1 = FrameSetSubDir.path(frame_set_id)
while True:
frames = frame_model.list(frame_set_id,
length=length,
offset=offset)
if not frames:
break
for frame in frames:
p2 = FrameFile.path(p1, frame[0], 'jpg')
height, width, _ = cv2.imread(p2).shape
debug(
f'{frame[0]} | {frame[1]} | {frame[2]} | ({width} | '
f'{height})', 3)
offset += length
def frame_set_select_extract(self, frame_set_id, opts):
"""
This method extracts a frame set to a transform set.
:param int frame_set_id: The frame set ID.
:param dict opts: The dict of options.
:rtype: int
"""
transform_set_id = TransformSetModel().insert('transform_set',
frame_set_id)
p1 = TransformSetSubDir.path(transform_set_id)
os.mkdir(p1)
frame_model = FrameModel()
transform_model = TransformModel()
length = int(os.environ.get('MODEL_LIST_LENGTH', '100'))
offset = 0
p2 = FrameSetSubDir.path(frame_set_id)
while True:
frames = frame_model.list(frame_set_id, length=length,
offset=offset, rejected=False)
if not frames:
break
for frame in frames:
transform_id = transform_model.insert(transform_set_id,
frame[0], None, 0)
p3 = FrameFile.path(p2, frame[0], 'jpg')
p4 = TransformFile.path(p1, transform_id, 'jpg')
shutil.copy(p3, p4)
debug(f'Transform with ID {transform_id:08d} at {p4} '
f'extracted from frame with ID {frame[0]:08d} at '
f'{p3}', 4)
offset += length
return transform_set_id
def select_merge(self, frame_set_ids, video_id=None, rejected=False):
"""
This method merges multiple frame sets into one frame set.
:param list(int) frame_set_ids: The frame set IDs.
:param int video_id: The video ID to which this frame set corresponds
(if any). The default value is None.
:param bool rejected: True if should include rejected frames else False
if should not. The default value is False.
:raises: CommandLineRouteHandlerError
:rtype: None
"""
frame_set_model = FrameSetModel()
for frame_set_id in frame_set_ids:
result = frame_set_model.select(frame_set_id)
if result is None:
raise CommandLineRouteHandlerError(
f'Frame set with ID {frame_set_id:08d} not found')
frame_set_id = frame_set_model.insert(video_id)
p1 = FrameSetSubDir.path(frame_set_id)
os.makedirs(p1)
frame_model = FrameModel()
length = int(os.environ.get('MODEL_LIST_LENGTH', '100'))
for frame_set_id_ in frame_set_ids:
offset = 0
p2 = FrameSetSubDir.path(frame_set_id_)
while True:
frames = frame_model.list(frame_set_id_,
length=length,
offset=offset,
rejected=rejected)
if not frames:
break
for frame in frames:
frame_id = frame_model.insert(frame_set_id, frame[2])
p3 = FrameFile.path(p2, frame[0], 'jpg')
p4 = FrameFile.path(p2, frame[0], 'jpg', '192x192')
p5 = FrameFile.path(p1, frame_id, 'jpg')
p6 = FrameFile.path(p1, frame_id, 'jpg', '192x192')
shutil.copy(p3, p5)
shutil.copy(p4, p6)
debug(
f'Frame with ID {frame[0]:08d} and thumbnail at '
f'{p3} and {p4} merged as ID {frame_id:08d} at '
f'{p5} and {p6}', 4)
offset += length
debug(f'frame_set_id={frame_set_id}, fk_videos={video_id}', 3)
