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 test_select(self):
with deepstar_path():
VideoModel().insert('test1', 'test2')
FrameSetModel().insert(1)
frame_model = FrameModel()
frame_model.insert(1, 0)
result = frame_model.select(1)
self.assertEqual(result, (1, 1, 0))
def put(self, frame_set_id, frame_id):
frame_model = FrameModel()
result = frame_model.update(frame_id,
request.get_json()['rejected'])
if result is False:
abort(404)
result = frame_model.select(frame_id)
if result is None:
abort(404)
return jsonify(result)
def insert_images(self, images_path, opts):
"""
This method inserts a frame set from a directory of images.
:param str images_path: The path to the directory containing the images
to insert.
:param dict opts: The dict of options.
:raises: CommandLineRouteHandlerError
:rtype: None
"""
if not os.path.isdir(images_path):
raise CommandLineRouteHandlerError(
f'The path at {images_path} is not a directory')
frame_set_id = FrameSetModel().insert(None)
p1 = FrameSetSubDir.path(frame_set_id)
os.makedirs(p1)
frame_model = FrameModel()
for image_path in glob.glob(os.path.join(images_path, '*')):
ext = os.path.splitext(image_path)[1].lower()
if ext not in ['.jpg', '.png']:
debug(
f"Skipped image at {image_path} because it does not "
f"have a '.jpg' or '.png' file extension", 4)
continue
image = cv2.imread(image_path)
frame_id = frame_model.insert(frame_set_id, 0)
p2 = FrameFile.path(p1, frame_id, 'jpg')
cv2.imwrite(p2, image, [cv2.IMWRITE_JPEG_QUALITY, 100])
thumbnail = imutils.resize(image, width=192, height=192)
p3 = FrameFile.path(p1, frame_id, 'jpg', '192x192')
cv2.imwrite(p3, thumbnail, [cv2.IMWRITE_JPEG_QUALITY, 100])
debug(
f'Image at {image_path} inserted with ID {frame_id:08d} at '
f'{p2} and {p3}', 4)
debug(f'frame_set_id={frame_set_id}, fk_videos=None', 3)
def test_insert_images_does_not_insert_non_images(self):
with deepstar_path():
tmpdir = os.environ['DEEPSTAR_PATH'] + '/test'
os.mkdir(tmpdir)
with open(os.path.join(tmpdir, 'test'), 'w') as file_:
file_.write('test')
with open(os.path.join(tmpdir, 'test.txt'), 'w') as file_:
file_.write('test')
args = ['main.py', 'insert', 'frame_sets', 'images', tmpdir]
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_id=1, fk_videos=None')
# db
result = FrameSetModel().select(1)
self.assertEqual(result, (1, None))
result = FrameModel().list(0)
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 test_video_select_extract(self):
with deepstar_path():
video_0001 = os.path.dirname(os.path.realpath(
__file__)) + '/../../support/video_0001.mp4' # noqa
shutil.copyfile(video_0001, VideoFile.path('video_0001.mp4'))
VideoModel().insert('test', 'video_0001.mp4')
frame_set_id = DefaultVideoSelectExtractPlugin(
).video_select_extract(1) # noqa
self.assertEqual(frame_set_id, 1)
# db
result = FrameSetModel().select(1)
self.assertEqual(result, (1, 1))
result = FrameModel().list(1)
self.assertEqual(len(result), 5)
self.assertEqual(result[0], (1, 1, 0))
self.assertEqual(result[1], (2, 1, 0))
self.assertEqual(result[2], (3, 1, 0))
self.assertEqual(result[3], (4, 1, 0))
self.assertEqual(result[4], (5, 1, 0))
# files
p1 = FrameSetSubDir.path(1)
# frames
self.assertIsInstance(cv2.imread(FrameFile.path(p1, 1, 'jpg')),
np.ndarray) # noqa
self.assertIsInstance(cv2.imread(FrameFile.path(p1, 2, 'jpg')),
np.ndarray) # noqa
self.assertIsInstance(cv2.imread(FrameFile.path(p1, 3, 'jpg')),
np.ndarray) # noqa
self.assertIsInstance(cv2.imread(FrameFile.path(p1, 4, 'jpg')),
np.ndarray) # noqa
self.assertIsInstance(cv2.imread(FrameFile.path(p1, 5, 'jpg')),
np.ndarray) # noqa
# thumbnails
self.assertEqual(
cv2.imread(FrameFile.path(p1, 1, 'jpg', '192x192')).shape[1],
192) # noqa
self.assertEqual(
cv2.imread(FrameFile.path(p1, 2, 'jpg', '192x192')).shape[1],
192) # noqa
self.assertEqual(
cv2.imread(FrameFile.path(p1, 3, 'jpg', '192x192')).shape[1],
192) # noqa
self.assertEqual(
cv2.imread(FrameFile.path(p1, 4, 'jpg', '192x192')).shape[1],
192) # noqa
self.assertEqual(
cv2.imread(FrameFile.path(p1, 5, 'jpg', '192x192')).shape[1],
192) # noqa
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_clone_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])
args = ['main.py', 'select', 'frame_sets', '1', 'clone']
opts = {}
route_handler = FrameSetCommandLineRouteHandler()
try:
sys.stdout = StringIO()
with mock.patch.dict(os.environ, {'MODEL_LIST_LENGTH': '4'}):
route_handler.handle(args, opts)
actual = sys.stdout.getvalue().strip()
finally:
sys.stdout = sys.__stdout__
# stdout
self.assertEqual(actual, 'frame_set_id=2, fk_videos=1')
# db
result = FrameSetModel().select(2)
self.assertEqual(result, (2, 1))
result = FrameModel().list(2)
self.assertEqual(len(result), 5)
self.assertEqual(result[0], (6, 2, 0))
self.assertEqual(result[1], (7, 2, 0))
self.assertEqual(result[2], (8, 2, 0))
self.assertEqual(result[3], (9, 2, 0))
self.assertEqual(result[4], (10, 2, 0))
# files
p1 = FrameSetSubDir.path(2)
# frames
self.assertTrue(os.path.isfile(FrameFile.path(p1, 6, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 7, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 8, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 9, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 10, 'jpg')))
# thumbnails
self.assertTrue(os.path.isfile(FrameFile.path(p1, 6, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 7, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 8, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 9, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 10, 'jpg', '192x192'))) # noqa
def test_frame_set_select_extract_face_rejected(self):
with deepstar_path():
video_0001 = os.path.dirname(os.path.realpath(__file__)) + '/../../support/video_0001.mp4' # noqa
shutil.copyfile(video_0001, VideoFile.path('video_0001.mp4'))
VideoModel().insert('test', 'video_0001.mp4')
DefaultVideoSelectExtractPlugin().video_select_extract(1) # noqa
frame_model = FrameModel()
frame_model.update(2, 1)
frame_model.update(4, 1)
with mock.patch.dict(os.environ, {'MODEL_LIST_LENGTH': '2'}):
transform_set_id = MTCNNFrameSetSelectExtractPlugin().frame_set_select_extract(1, {}) # noqa
self.assertEqual(transform_set_id, 1)
# db
result = TransformSetModel().select(1)
self.assertEqual(result, (1, 'face', 1, None))
result = TransformModel().list(1)
self.assertEqual(len(result), 3)
t = list(result[0])
json.loads(t.pop(3))
self.assertEqual(t, [1, 1, 1, 0])
t = list(result[1])
json.loads(t.pop(3))
self.assertEqual(t, [2, 1, 3, 0])
t = list(result[2])
json.loads(t.pop(3))
self.assertEqual(t, [3, 1, 5, 0])
# files
p1 = TransformSetSubDir.path(1)
# transforms
self.assertTrue(os.path.isfile(TransformFile.path(p1, 1, 'jpg')))
self.assertTrue(os.path.isfile(TransformFile.path(p1, 2, 'jpg')))
self.assertTrue(os.path.isfile(TransformFile.path(p1, 3, 'jpg')))
def test_select(self):
with deepstar_path():
FrameSetModel().insert(None)
FrameModel().insert(1, 0)
TransformSetModel().insert('test', 1)
transform_model = TransformModel()
transform_model.insert(1, 1, '{}', 0)
result = transform_model.select(1)
self.assertEqual(result, (1, 1, 1, '{}', 0))
def test_fk_transform_sets_on_delete_cascade(self):
with deepstar_path():
FrameSetModel().insert(None)
FrameModel().insert(1, 0)
transform_set_model = TransformSetModel()
transform_set_model.insert('test', 1)
transform_model = TransformModel()
transform_model.insert(1, 1, '{}', 0)
self.assertEqual(transform_model.list(1), [(1, 1, 1, '{}', 0)])
transform_set_model.delete(1)
self.assertIsNone(transform_model.list(1))
def test_select_extract_one(self):
with deepstar_path():
video_0001 = os.path.dirname(os.path.realpath(__file__)) + '/../../support/video_0001.mp4' # noqa
shutil.copyfile(video_0001, VideoFile.path('video_0001.mp4'))
VideoModel().insert('test', 'video_0001.mp4')
args = ['main.py', 'select', 'videos', '1', 'extract']
opts = {}
route_handler = VideoCommandLineRouteHandler()
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_id=1, video_id=1')
# db
result = FrameSetModel().select(1)
self.assertEqual(result, (1, 1))
result = FrameModel().list(1)
self.assertEqual(len(result), 5)
self.assertEqual(result[0], (1, 1, 0))
self.assertEqual(result[1], (2, 1, 0))
self.assertEqual(result[2], (3, 1, 0))
self.assertEqual(result[3], (4, 1, 0))
self.assertEqual(result[4], (5, 1, 0))
# files
p1 = FrameSetSubDir.path(1)
# frames
self.assertTrue(os.path.isfile(FrameFile.path(p1, 1, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 2, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 3, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 4, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 5, 'jpg')))
# thumbnails
self.assertTrue(os.path.isfile(FrameFile.path(p1, 1, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 2, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 3, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 4, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 5, 'jpg', '192x192'))) # noqa
def test_count(self):
with deepstar_path():
VideoModel().insert('test1', 'test2')
FrameSetModel().insert(1)
frame_model = FrameModel()
frame_model.insert(1, 0)
frame_model.insert(1, 0)
frame_model.insert(1, 0)
self.assertEqual(frame_model.count(1), 3)
def test_count_rejected_false(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)
self.assertEqual(frame_model.count(1, rejected=False), 2)
def mock_transform_set(self):
image_0006 = os.path.dirname(os.path.realpath(
__file__)) + '/../../support/image_0006.jpg' # noqa
FrameSetModel().insert('1')
FrameModel().insert(1, 0)
p1 = FrameSetSubDir.path(1)
os.mkdir(p1)
shutil.copy(image_0006, FrameFile.path(p1, 1, 'jpg'))
MTCNNFrameSetSelectExtractPlugin().frame_set_select_extract(1, {})
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 test_insert_images(self):
with deepstar_path():
image_0001 = os.path.dirname(os.path.realpath(__file__)) + '/../../support/image_0001.jpg' # noqa
image_0007 = os.path.dirname(os.path.realpath(__file__)) + '/../../support/image_0007.png' # noqa
tmpdir = os.environ['DEEPSTAR_PATH'] + '/test'
os.mkdir(tmpdir)
shutil.copy(image_0001, os.path.join(tmpdir, 'image_0001.jpg'))
shutil.copy(image_0001, os.path.join(tmpdir, 'image_0002.jpg'))
shutil.copy(image_0007, os.path.join(tmpdir, 'image_0003.png'))
shutil.copy(image_0007, os.path.join(tmpdir, 'image_0004.png'))
args = ['main.py', 'insert', 'frame_sets', 'images', tmpdir]
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_id=1, fk_videos=None')
# db
result = FrameSetModel().select(1)
self.assertEqual(result, (1, None))
result = FrameModel().list(1)
self.assertEqual(len(result), 4)
self.assertEqual(result[0], (1, 1, 0))
self.assertEqual(result[1], (2, 1, 0))
self.assertEqual(result[2], (3, 1, 0))
self.assertEqual(result[3], (4, 1, 0))
# files
p1 = FrameSetSubDir.path(1)
# frames
self.assertTrue(os.path.isfile(FrameFile.path(p1, 1, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 2, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 3, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 4, 'jpg')))
def get(self, frame_set_id):
if 'offset' in request.args:
offset = int(request.args['offset'])
else:
offset = 0
if 'length' in request.args:
length = int(request.args['length'])
else:
length = 100
result = FrameModel().list(frame_set_id,
length=length,
offset=offset)
if result is None:
abort(404)
return jsonify(result)
def test_update(self):
with deepstar_path():
FrameSetModel().insert(None)
FrameModel().insert(1, 0)
TransformSetModel().insert('test', 1)
transform_model = TransformModel()
transform_model.insert(1, 1, 'test1', 0)
result = transform_model.select(1)
self.assertEqual(result, (1, 1, 1, 'test1', 0))
result = transform_model.update(1, 'test2', 1)
self.assertTrue(result)
result = transform_model.select(1)
self.assertEqual(result, (1, 1, 1, 'test2', 1))
def test_count(self):
with deepstar_path():
FrameSetModel().insert(None)
frame_model = FrameModel()
frame_model.insert(1, 0)
frame_model.insert(1, 0)
frame_model.insert(1, 0)
TransformSetModel().insert('test', 1)
transform_model = TransformModel()
transform_model.insert(1, 1, '{}', 0)
transform_model.insert(1, 2, '{}', 0)
transform_model.insert(1, 3, '{}', 0)
self.assertEqual(frame_model.count(1), 3)
def frame_set_select_extract(self, frame_set_id, opts):
"""
This method extracts faces from each frame in a frame set.
:param int frame_set_id: The frame set ID.
:param dict opts: The dict of opts.
:rtype: int
"""
detector = MTCNN()
offset_percent = 0.2
min_confidence = 0.9
debug_ = True if 'debug' in opts else False
frame_set_path = FrameSetSubDir.path(frame_set_id)
transform_set_id = TransformSetModel().insert(self.name, frame_set_id)
transform_set_path = TransformSetSubDir.path(transform_set_id)
os.makedirs(transform_set_path)
length = int(os.environ.get('MODEL_LIST_LENGTH', '100'))
offset = 0
while True:
result = FrameModel().list(frame_set_id,
length=length,
offset=offset,
rejected=False)
if not result:
break
for frame_id, _, rejected in result:
self._extract_faces(frame_set_path, frame_id,
transform_set_path, transform_set_id,
detector, offset_percent, min_confidence,
debug_)
offset += length
return transform_set_id
def test_count_rejected_false(self):
with deepstar_path():
FrameSetModel().insert(None)
frame_model = FrameModel()
frame_model.insert(1, 0)
frame_model.insert(1, 0)
frame_model.insert(1, 0)
TransformSetModel().insert('test', 1)
transform_model = TransformModel()
transform_model.insert(1, 1, '{}', 0)
transform_model.insert(1, 2, '{}', 1)
transform_model.insert(1, 3, '{}', 0)
self.assertEqual(transform_model.count(1, rejected=False), 2)
def test_select_export_dir_one_rejected(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])
FrameModel().update(5, rejected=1)
tmpdir = os.environ['DEEPSTAR_PATH'] + '/test'
os.mkdir(tmpdir)
args = ['main.py', 'select', 'frame_sets', '1', 'export', 'dir', tmpdir] # noqa
opts = {}
route_handler = FrameSetCommandLineRouteHandler()
try:
sys.stdout = StringIO()
with mock.patch.dict(os.environ, {'MODEL_LIST_LENGTH': '2'}):
route_handler.handle(args, opts)
actual = sys.stdout.getvalue().strip()
finally:
sys.stdout = sys.__stdout__
# stdout
self.assertEqual(actual, f'4 frames were successfully exported to {tmpdir}') # noqa
# frames
self.assertTrue(os.path.isfile(FrameFile.path(tmpdir, 1, 'jpg'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(tmpdir, 2, 'jpg'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(tmpdir, 3, 'jpg'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(tmpdir, 4, 'jpg'))) # noqa
def test_video_select_extract_sub_sample(self):
with deepstar_path():
video_0001 = os.path.dirname(os.path.realpath(
__file__)) + '/../../support/video_0001.mp4' # noqa
shutil.copyfile(video_0001, VideoFile.path('video_0001.mp4'))
VideoModel().insert('test', 'video_0001.mp4')
DefaultVideoSelectExtractPlugin().video_select_extract(
1, sub_sample=2) # noqa
# db
result = FrameModel().list(1)
self.assertEqual(len(result), 3)
self.assertEqual(result[0], (1, 1, 0))
self.assertEqual(result[1], (2, 1, 0))
self.assertEqual(result[2], (3, 1, 0))
# files
p1 = FrameSetSubDir.path(1)
# frames
self.assertTrue(os.path.isfile(FrameFile.path(p1, 1, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 2, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 3, 'jpg')))
# thumbnails
self.assertTrue(
os.path.isfile(FrameFile.path(p1, 1, 'jpg',
'192x192'))) # noqa
self.assertTrue(
os.path.isfile(FrameFile.path(p1, 2, 'jpg',
'192x192'))) # noqa
self.assertTrue(
os.path.isfile(FrameFile.path(p1, 3, 'jpg',
'192x192'))) # noqa
def test_list(self):
with deepstar_path():
FrameSetModel().insert(None)
FrameModel().insert(1, 0)
TransformSetModel().insert('test', 1)
transform_model = TransformModel()
transform_model.insert(1, 1, '{}', 0)
transform_model.insert(1, 1, '{}', 1)
transform_model.insert(1, 1, '{}', 0)
result = transform_model.list(1)
self.assertEqual(len(result), 3)
self.assertEqual(result[0], (1, 1, 1, '{}', 0))
self.assertEqual(result[1], (2, 1, 1, '{}', 1))
self.assertEqual(result[2], (3, 1, 1, '{}', 0))
result = transform_model.list(1, length=2)
self.assertEqual(len(result), 2)
self.assertEqual(result[0], (1, 1, 1, '{}', 0))
self.assertEqual(result[1], (2, 1, 1, '{}', 1))
result = transform_model.list(1, offset=1)
self.assertEqual(len(result), 2)
self.assertEqual(result[0], (2, 1, 1, '{}', 1))
self.assertEqual(result[1], (3, 1, 1, '{}', 0))
result = transform_model.list(1, length=1, offset=1)
self.assertEqual(len(result), 1)
self.assertEqual(result[0], (2, 1, 1, '{}', 1))
result = transform_model.list(1, rejected=False)
self.assertEqual(len(result), 2)
self.assertEqual(result[0], (1, 1, 1, '{}', 0))
self.assertEqual(result[1], (3, 1, 1, '{}', 0))
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)
def test_select_merge_rejected(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.insert_file(video_0001)
route_handler.insert_file(video_0001)
route_handler.select_extract([1, 2, 3])
FrameModel().update(15, rejected=1)
args = ['main.py', 'select', 'frame_sets', '1-2,3', 'merge']
opts = {}
route_handler = FrameSetCommandLineRouteHandler()
try:
sys.stdout = StringIO()
with mock.patch.dict(os.environ, {'MODEL_LIST_LENGTH': '4'}):
route_handler.handle(args, opts)
actual = sys.stdout.getvalue().strip()
finally:
sys.stdout = sys.__stdout__
# stdout
self.assertEqual(actual, 'frame_set_id=4, fk_videos=None')
# db
result = FrameSetModel().select(4)
self.assertEqual(result, (4, None))
result = FrameModel().list(4)
self.assertEqual(len(result), 14)
self.assertEqual(result[0], (16, 4, 0))
self.assertEqual(result[1], (17, 4, 0))
self.assertEqual(result[2], (18, 4, 0))
self.assertEqual(result[3], (19, 4, 0))
self.assertEqual(result[4], (20, 4, 0))
self.assertEqual(result[5], (21, 4, 0))
self.assertEqual(result[6], (22, 4, 0))
self.assertEqual(result[7], (23, 4, 0))
self.assertEqual(result[8], (24, 4, 0))
self.assertEqual(result[9], (25, 4, 0))
self.assertEqual(result[10], (26, 4, 0))
self.assertEqual(result[11], (27, 4, 0))
self.assertEqual(result[12], (28, 4, 0))
self.assertEqual(result[13], (29, 4, 0))
# files
p1 = FrameSetSubDir.path(4)
# frames
self.assertTrue(os.path.isfile(FrameFile.path(p1, 16, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 17, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 18, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 19, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 20, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 21, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 22, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 23, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 24, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 25, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 26, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 27, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 28, 'jpg')))
self.assertTrue(os.path.isfile(FrameFile.path(p1, 29, 'jpg')))
# thumbnails
self.assertTrue(os.path.isfile(FrameFile.path(p1, 16, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 17, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 18, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 19, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 20, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 21, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 22, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 23, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 24, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 25, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 26, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 27, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 28, 'jpg', '192x192'))) # noqa
self.assertTrue(os.path.isfile(FrameFile.path(p1, 29, 'jpg', '192x192'))) # noqa
def video_select_extract(self, video_id, sub_sample=1, max_sample=0):
"""
This method extracts frames and thumbnails from a video to a frame set.
:param int video_id: The video ID.
:param int sub_sample: Sample frames at a rate of 1 sample per
sub_sample rate. For example, if sub_sample is 30, then sample
frames at a rate of 1 frame per 30 frames. The default value of 1
indicates to sample every frame.
:param int max_sample: Sample up to maximum count of frames. For
example, sample up to 1000 total frames (and then cease to sample).
The default value of 0 indicates that there is no maximum count of
frames.
:raises: CommandLineRouteHandlerError
:rtype: int
"""
result = VideoModel().select(video_id)
p1 = VideoFile.path(result[2])
vc = cv2.VideoCapture(p1)
try:
if not vc.isOpened():
raise CommandLineRouteHandlerError(
f'OpenCV VideoCapture isOpened returned false for {p1}')
frame_set_id = FrameSetModel().insert(video_id)
p2 = FrameSetSubDir.path(frame_set_id)
os.makedirs(p2)
frame_model = FrameModel()
sub_sample_ = sub_sample
max_sample_ = 0
while True:
ret, frame = vc.read()
if not ret:
break
# sub_sample
if sub_sample_ != sub_sample:
sub_sample_ += 1
continue
else:
sub_sample_ = 1
frame_id = frame_model.insert(frame_set_id, 0)
p3 = FrameFile.path(p2, frame_id, 'jpg')
cv2.imwrite(p3, frame, [cv2.IMWRITE_JPEG_QUALITY, 100])
# imutils.resize preserves aspect ratio.
thumbnail = imutils.resize(frame, width=192, height=192)
p4 = FrameFile.path(p2, frame_id, 'jpg', '192x192')
# can adjust jpeg quality thus impacting file size
cv2.imwrite(p4, thumbnail, [cv2.IMWRITE_JPEG_QUALITY, 100])
debug(
f'Frame with ID {frame_id:08d} and thumbnail extracted '
f'to {p3} and {p4}', 4)
# max_sample
max_sample_ += 1
if max_sample > 0:
if max_sample_ == max_sample:
break
finally:
vc.release()
return frame_set_id
