def test_delete(self):
with deepstar_path():
frame_set_model = FrameSetModel()
frame_set_model.insert(1)
result = frame_set_model.select(1)
self.assertEqual(result, (1, 1))
result = frame_set_model.delete(1)
self.assertTrue(result)
result = frame_set_model.select(1)
self.assertIsNone(result)
def test_delete_one(self):
with deepstar_path():
with mock.patch.dict(os.environ, {'DEBUG_LEVEL': '0'}):
route_handler = VideoCommandLineRouteHandler()
video_0001 = os.path.dirname(os.path.realpath(__file__)) + '/../../support/video_0001.mp4' # noqa
route_handler.insert_file(video_0001)
route_handler.select_extract([1])
# files
self.assertTrue(os.path.exists(FrameSetSubDir.path(1)))
# db
frame_set_model = FrameSetModel()
self.assertIsNotNone(frame_set_model.select(1))
frame_model = FrameModel()
self.assertIsNotNone(frame_model.select(1))
self.assertIsNotNone(frame_model.select(2))
self.assertIsNotNone(frame_model.select(3))
self.assertIsNotNone(frame_model.select(4))
self.assertIsNotNone(frame_model.select(5))
args = ['main.py', 'delete', 'frame_sets', '1']
opts = {}
route_handler = FrameSetCommandLineRouteHandler()
try:
sys.stdout = StringIO()
route_handler.handle(args, opts)
actual = sys.stdout.getvalue().strip()
finally:
sys.stdout = sys.__stdout__
# stdout
self.assertEqual(actual, 'Frame set 1 was successfully deleted')
# db
self.assertIsNone(frame_set_model.select(1))
self.assertIsNone(frame_model.select(1))
self.assertIsNone(frame_model.select(2))
self.assertIsNone(frame_model.select(3))
self.assertIsNone(frame_model.select(4))
self.assertIsNone(frame_model.select(5))
# files
self.assertFalse(os.path.exists(FrameSetSubDir.path(1)))
def test_select(self):
with deepstar_path():
VideoModel().insert('test1', 'test2')
frame_set_model = FrameSetModel()
frame_set_model.insert(1)
result = frame_set_model.select(1)
self.assertEqual(result, (1, 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 select_clone(self, frame_set_ids):
"""
This method clones frame sets to new frame sets.
:param list(int) frame_set_ids: The frame set IDs.
: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')
for frame_set_id in frame_set_ids:
result = frame_set_model.select(frame_set_id)
self.select_merge([frame_set_id], result[1], rejected=True)
def delete(self, frame_set_ids):
"""
This method deletes a frame set from the frame set collection.
:param list(int) frame_set_ids: The frame set IDs.
: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')
for frame_set_id in frame_set_ids:
frame_set_model.delete(frame_set_id)
shutil.rmtree(FrameSetSubDir.path(frame_set_id))
debug(f'Frame set {frame_set_id} was successfully deleted', 3)
def select_extract(self, frame_set_ids, name, opts):
"""
This method extracts portions of frames to transform sets.
:param list(int) frame_set_ids: The frame set IDs.
:param str name: The name of the frame set extraction plugin to use.
: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')
plugin = Plugin.get('frame_set_select_extract', name)
if plugin is None:
raise CommandLineRouteHandlerError(
f"'{name}' is not a valid frame set extraction plugin name")
transform_set_model = TransformSetModel()
for frame_set_id in frame_set_ids:
transform_set_id = plugin().frame_set_select_extract(
frame_set_id, opts)
result = transform_set_model.select(transform_set_id)
debug(
f'transform_set_id={result[0]}, name={result[1]}, '
f'fk_frame_sets={result[2]}, '
f'fk_prev_transform_sets={result[3]}', 3)
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)