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_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 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_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_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 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 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 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_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_path_with_res(self):
self.assertEqual(FrameFile.path('test', 1, 'jpg', '192'), 'test/00000001-192.jpg') # noqa
def test_path(self):
self.assertEqual(FrameFile.path('test', 1, 'jpg'), 'test/00000001.jpg')
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
def _extract_faces(self, frame_set_path, frame_id, transform_set_path,
transform_set_id, detector, offset_percent,
min_confidence, debug_):
"""
This method extracts faces from a frame.
:param str frame_set_path:The frame set path.
:param int frame_id: The frame ID.
:param str transform_set_path: The transform set path.
:param transform_set_id: The transform set ID.
:param MTCNN detector: The detector to use to detect faces.
:param float offset_percent:
:param float min_confidence: The minimum confidence value required to
accept/reject a detected face.
:param bool debug_: True if should place markers on landmarks else
False if should not.
:rtype: None
"""
frame_path = FrameFile.path(frame_set_path, frame_id, 'jpg')
img = cv2.imread(frame_path)
img_height, img_width = img.shape[:2]
results = detector.detect_faces(img)
for r in results:
if r['confidence'] < min_confidence:
continue
x, y, width, height = r['box']
adjusted_x = int(max(0, x - (0.5 * width * offset_percent)))
adjusted_y = int(max(0, y - (0.5 * height * offset_percent)))
t = x + width + (0.5 * width * offset_percent)
adjusted_right_x = int(min(img_width, t))
t = y + height + (0.5 * height * offset_percent)
adjusted_bottom_y = int(min(img_height, t))
metadata = {
'face': {
k: [v[0] - adjusted_x, v[1] - adjusted_y]
for k, v in r['keypoints'].items()
}
}
transform_id = TransformModel().insert(transform_set_id, frame_id,
json.dumps(metadata), 0)
face_crop = img[adjusted_y:adjusted_bottom_y,
adjusted_x:adjusted_right_x]
output_path = TransformFile.path(transform_set_path, transform_id,
'jpg')
if debug_ is True:
for _, v in metadata['face'].items():
cv2.drawMarker(face_crop,
tuple(v), (0, 0, 255),
markerType=cv2.MARKER_DIAMOND,
markerSize=15,
thickness=2)
cv2.imwrite(output_path, face_crop,
[cv2.IMWRITE_JPEG_QUALITY, 100])
debug(
f'Transform with ID {transform_id:08d} at {output_path} '
f'extracted from frame with ID {frame_id:08d} at '
f'{frame_path}', 4)