def test_select_extract_many(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'))
shutil.copyfile(video_0001, VideoFile.path('video_0002.mp4'))
shutil.copyfile(video_0001, VideoFile.path('video_0003.mp4'))
VideoModel().insert('test', 'video_0001.mp4')
VideoModel().insert('test', 'video_0002.mp4')
VideoModel().insert('test', 'video_0003.mp4')
args = ['main.py', 'select', 'videos', '1-2,3', '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
expected = textwrap.dedent('''
frame_set_id=1, video_id=1
frame_set_id=2, video_id=2
frame_set_id=3, video_id=3
''').strip()
self.assertEqual(actual, expected)
def test_transform_set_select_curate_min_size_fails_due_to_missing_required_option(
self): # noqa
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
self.mock_transform_set()
with self.assertRaises(ValueError):
try:
MinSizeTransformSetSelectCuratePlugin(
).transform_set_select_curate(1, {}) # noqa
except ValueError as e:
self.assertEqual(
str(e),
'The min-size option is required but was not supplied'
) # noqa
raise e
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)
video_model = VideoModel()
old = video_model.select(1)
self.assertIsNotNone(old)
args = ['main.py', 'delete', 'videos', '1']
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, 'Video 1 was successfully deleted')
# db + files
new = video_model.select(1)
self.assertIsNone(new)
self.assertFalse(os.path.exists(VideoFile.path(old[2])))
def test_select_deploy(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')
class TestPlugin:
def video_select_deploy(self, video_id, opts):
pass
args = ['main.py', 'select', 'videos', '1', 'deploy', 'test']
opts = {}
route_handler = VideoCommandLineRouteHandler()
try:
sys.stdout = StringIO()
with mock.patch.dict(Plugin._map, {'video_select_deploy': {'test': TestPlugin}}): # noqa
route_handler.handle(args, opts)
actual = sys.stdout.getvalue().strip()
finally:
sys.stdout = sys.__stdout__
# stdout
self.assertEqual(actual, 'success')
def test_insert_youtube_url_file(self):
with deepstar_path():
video_0001 = os.path.dirname(os.path.realpath(__file__)) + '/../../support/video_0001.mp4' # noqa
class TestVideoCommandLineRouteHandler_(VideoCommandLineRouteHandler): # noqa
def uuid(self):
return '12345678-1234-1234-1234-123456789012'
args = ['main.py', 'insert', 'videos', 'youtube', 'https://www.youtube.com/watch?v=test'] # noqa
opts = {'file': video_0001}
route_handler = TestVideoCommandLineRouteHandler_()
try:
sys.stdout = StringIO()
route_handler.handle(args, opts)
actual = sys.stdout.getvalue().strip()
finally:
sys.stdout = sys.__stdout__
# stdout
self.assertIsNotNone(re.match('^video_id=1, uri=.*video_0001.mp4, filename=12345678-1234-1234-1234-123456789012.mp4, description=None$', actual)) # noqa
# db + files
result = VideoModel().select(1)
self.assertEqual(len(result), 4)
self.assertEqual(result[0], 1)
self.assertIsNotNone(re.match('^.*video_0001.mp4', result[1]))
UUID(result[2].rstrip('.mp4'), version=4)
self.assertTrue(os.path.isfile(VideoFile.path(result[2])))
self.assertEqual(result[3], None)
def test_insert_file(self):
with deepstar_path():
video_0001 = os.path.dirname(os.path.realpath(__file__)) + '/../../support/video_0001.mp4' # noqa
class TestVideoCommandLineRouteHandler_(VideoCommandLineRouteHandler): # noqa
def uuid(self):
return '12345678-1234-1234-1234-123456789012'
args = ['main.py', 'insert', 'videos', 'file', video_0001]
opts = {}
route_handler = TestVideoCommandLineRouteHandler_()
try:
sys.stdout = StringIO()
route_handler.handle(args, opts)
actual = sys.stdout.getvalue().strip()
finally:
sys.stdout = sys.__stdout__
# stdout
expected = 'video_id=1, uri=' + video_0001 + ', filename=12345678-1234-1234-1234-123456789012.mp4, description=None' # noqa
self.assertEqual(actual, expected)
# db
result = VideoModel().select(1)
self.assertEqual(len(result), 4)
self.assertEqual(result[0], 1)
self.assertEqual(result[1], video_0001)
UUID(result[2].rstrip('.mp4'), version=4)
self.assertEqual(result[3], None)
# files
self.assertTrue(os.path.isfile(VideoFile.path(result[2])))
def test_transform_set_select_extract_max_size(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
self.mock_transform_set()
with mock.patch.dict(os.environ, {'MODEL_LIST_LENGTH': '2'}):
transform_set_id = MaxSizeTransformSetSelectExtractPlugin(
).transform_set_select_extract(1, {}) # noqa
self.assertEqual(transform_set_id, 2)
# db
result = TransformSetModel().select(2)
self.assertEqual(result, (2, 'max_size', 1, 1))
result = TransformModel().list(2)
self.assertEqual(len(result), 5)
t = list(result[0])
json.loads(t.pop(3))
self.assertEqual(t, [6, 2, 1, 0])
t = list(result[1])
json.loads(t.pop(3))
self.assertEqual(t, [7, 2, 2, 0])
t = list(result[2])
json.loads(t.pop(3))
self.assertEqual(t, [8, 2, 3, 0])
t = list(result[3])
json.loads(t.pop(3))
self.assertEqual(t, [9, 2, 4, 0])
t = list(result[4])
json.loads(t.pop(3))
self.assertEqual(t, [10, 2, 5, 0])
# files
p1 = TransformSetSubDir.path(2)
# transforms
self.assertEqual(
cv2.imread(TransformFile.path(p1, 6, 'jpg')).shape[0],
299) # noqa
self.assertEqual(
cv2.imread(TransformFile.path(p1, 7, 'jpg')).shape[0],
299) # noqa
self.assertEqual(
cv2.imread(TransformFile.path(p1, 8, 'jpg')).shape[0],
299) # noqa
self.assertEqual(
cv2.imread(TransformFile.path(p1, 9, 'jpg')).shape[0],
299) # noqa
self.assertEqual(
cv2.imread(TransformFile.path(p1, 10, 'jpg')).shape[0],
299) # noqa
def delete(self, video_ids):
"""
This method deletes a video from the video collection.
:param list(ints) video_ids: The video IDs.
:raises: CommandLineRouteHandlerError
:rtype: None
"""
video_model = VideoModel()
for video_id in video_ids:
result = video_model.select(video_id)
if result is None:
raise CommandLineRouteHandlerError(
f'Video with ID {video_id:08d} not found')
for video_id in video_ids:
result = video_model.select(video_id)
path = VideoFile.path(result[2])
video_model.delete(video_id)
os.remove(path)
debug(f'Video {video_id} was successfully deleted', 3)
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_transform_set_select_extract_pad_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
self.mock_transform_set()
MaxSizeTransformSetSelectExtractPlugin(
).transform_set_select_extract(1, {}) # noqa
transform_model = TransformModel()
transform_model.update(7, rejected=1)
transform_model.update(9, rejected=1)
with mock.patch.dict(os.environ, {'MODEL_LIST_LENGTH': '2'}):
transform_set_id = PadTransformSetSelectExtractPlugin(
).transform_set_select_extract(2, {}) # noqa
self.assertEqual(transform_set_id, 3)
# db
result = TransformSetModel().select(3)
self.assertEqual(result, (3, 'pad', 1, 2))
result = TransformModel().list(3)
self.assertEqual(len(result), 3)
t = list(result[0])
json.loads(t.pop(3))
self.assertEqual(t, [11, 3, 1, 0])
t = list(result[1])
json.loads(t.pop(3))
self.assertEqual(t, [12, 3, 3, 0])
t = list(result[2])
json.loads(t.pop(3))
self.assertEqual(t, [13, 3, 5, 0])
# files
p1 = TransformSetSubDir.path(3)
# transforms
self.assertIsInstance(cv2.imread(TransformFile.path(p1, 11,
'jpg')),
np.ndarray) # noqa
self.assertIsInstance(cv2.imread(TransformFile.path(p1, 12,
'jpg')),
np.ndarray) # noqa
self.assertIsInstance(cv2.imread(TransformFile.path(p1, 13,
'jpg')),
np.ndarray) # noqa
def test_insert_youtube_file_description(self):
with deepstar_path():
video_0001 = os.path.dirname(os.path.realpath(__file__)) + '/../../support/video_0001.mp4' # noqa
class TestVideoCommandLineRouteHandler_(VideoCommandLineRouteHandler): # noqa
def youtube_context(self, url):
return None
def youtube_streams(self, context):
return [
['22', 'video/mp4', '720p']
]
def uuid(self):
return '12345678-1234-1234-1234-123456789012'
def youtube_download(self, context, stream, output_path, filename): # noqa
filename = filename + '.mp4'
shutil.copyfile(video_0001, output_path + '/' + filename)
return filename
path = os.environ['DEEPSTAR_PATH'] + '/test.txt'
with open(path, 'w') as file_:
file_.write('https://www.youtube.com/watch?v=test1, test')
args = ['main.py', 'insert', 'videos', 'youtube', path]
opts = {}
route_handler = TestVideoCommandLineRouteHandler_()
try:
sys.stdout = StringIO()
route_handler.handle(args, opts)
actual = sys.stdout.getvalue().strip()
finally:
sys.stdout = sys.__stdout__
# stdout
expected = \
'video_id=1, uri=https://www.youtube.com/watch?v=test1, ' \
'filename=12345678-1234-1234-1234-123456789012.mp4, ' \
'description=test' \
self.assertEqual(actual, expected)
# db + files
result = VideoModel().select(1)
self.assertEqual(result[0], 1)
self.assertEqual(result[1], 'https://www.youtube.com/watch?v=test1') # noqa
UUID(result[2].rstrip('.mp4'), version=4)
self.assertTrue(os.path.isfile(VideoFile.path(result[2])))
self.assertEqual(result[3], 'test')
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_frame_set_select_extract_face(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
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), 5)
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, 2, 0])
t = list(result[2])
json.loads(t.pop(3))
self.assertEqual(t, [3, 1, 3, 0])
t = list(result[3])
json.loads(t.pop(3))
self.assertEqual(t, [4, 1, 4, 0])
t = list(result[4])
json.loads(t.pop(3))
self.assertEqual(t, [5, 1, 5, 0])
# files
p1 = TransformSetSubDir.path(1)
# transforms
self.assertIsInstance(cv2.imread(TransformFile.path(p1, 1, 'jpg')), np.ndarray) # noqa
self.assertIsInstance(cv2.imread(TransformFile.path(p1, 2, 'jpg')), np.ndarray) # noqa
self.assertIsInstance(cv2.imread(TransformFile.path(p1, 3, 'jpg')), np.ndarray) # noqa
self.assertIsInstance(cv2.imread(TransformFile.path(p1, 4, 'jpg')), np.ndarray) # noqa
self.assertIsInstance(cv2.imread(TransformFile.path(p1, 5, 'jpg')), np.ndarray) # noqa
def test_insert_vimeo_url(self):
with deepstar_path():
video_0001 = os.path.dirname(os.path.realpath(__file__)) + '/../../support/video_0001.mp4' # noqa
class TestVideoCommandLineRouteHandler_(VideoCommandLineRouteHandler): # noqa
def vimeo_context(self, url):
return None
def vimeo_title(self, video):
return 'Unittest video'
def vimeo_stream(self, video, quality):
return None
def uuid(self):
return '12345678-1234-1234-1234-123456789012'
def vimeo_stream_download(self, stream, output_path, filename):
shutil.copyfile(video_0001,
os.path.join(output_path, filename))
args = ['main.py', 'insert', 'videos', 'vimeo', 'https://vimeo.com/test'] # noqa
opts = {}
route_handler = TestVideoCommandLineRouteHandler_()
try:
sys.stdout = StringIO()
route_handler.handle(args, opts)
actual = sys.stdout.getvalue().strip()
finally:
sys.stdout = sys.__stdout__
# stdout
self.assertEqual(actual, 'video_id=1, uri=https://vimeo.com/test, filename=12345678-1234-1234-1234-123456789012.mp4, description=Unittest video') # noqa
# db + files
result = VideoModel().select(1)
self.assertEqual(len(result), 4)
self.assertEqual(result[0], 1)
self.assertEqual(result[1], 'https://vimeo.com/test') # noqa
UUID(result[2].rstrip('.mp4'), version=4)
self.assertTrue(os.path.isfile(VideoFile.path(result[2])))
self.assertEqual(result[3], 'Unittest video')
def test_insert_image_frame_count(self):
with deepstar_path():
image_0001 = os.path.dirname(os.path.realpath(__file__)) + '/../../support/image_0001.jpg' # noqa
args = ['main.py', 'insert', 'videos', 'image', image_0001]
opts = {'frame-count': '5'}
with mock.patch.dict(os.environ, {'DEBUG_LEVEL': '0'}):
VideoCommandLineRouteHandler().handle(args, opts)
# files
video_path = VideoFile.path(VideoModel().select(1)[2])
vc = cv2.VideoCapture(video_path)
try:
self.assertTrue(vc.isOpened())
self.assertEqual(vc.get(cv2.CAP_PROP_FRAME_COUNT), 5)
finally:
vc.release()
def test_insert_image(self):
with deepstar_path():
image_0001 = os.path.dirname(os.path.realpath(__file__)) + '/../../support/image_0001.jpg' # noqa
args = ['main.py', 'insert', 'videos', 'image', image_0001]
opts = {}
class TestVideoCommandLineRouteHandler_(VideoCommandLineRouteHandler): # noqa
def uuid(self):
return '12345678-1234-1234-1234-123456789012'
route_handler = TestVideoCommandLineRouteHandler_()
try:
sys.stdout = StringIO()
route_handler.handle(args, opts)
actual = sys.stdout.getvalue().strip()
finally:
sys.stdout = sys.__stdout__
# stdout
self.assertIsNotNone(re.match('^video_id=1, uri=.*video.mp4, filename=12345678-1234-1234-1234-123456789012.mp4, description=None$', actual)) # noqa
# db
result = VideoModel().select(1)
self.assertEqual(len(result), 4)
self.assertEqual(result[0], 1)
self.assertIsNotNone(re.match('^.*video.mp4$', result[1]))
self.assertEqual(result[2], '12345678-1234-1234-1234-123456789012.mp4') # noqa
self.assertEqual(result[3], None)
# files
video_path = VideoFile.path(result[2])
vc = cv2.VideoCapture(video_path)
try:
self.assertTrue(vc.isOpened())
self.assertEqual(vc.get(cv2.CAP_PROP_FRAME_COUNT), 1)
finally:
vc.release()
def select_detect(self, video_id, model_name, opts):
"""
This method runs a deepfake detection model over a video.
:param int video_id: The video ID.
:param str model_name: The name of the detection model to run.
:param opts dict: The dict of options.
"""
video_model = VideoModel()
result = video_model.select(video_id)
if result is None:
raise CommandLineRouteHandlerError(
f'Video with ID {video_id:08d} not found')
plugin = Plugin.get('video_select_detect', model_name)
if plugin is None:
raise CommandLineRouteHandlerError(
f"'{model_name}' is not a valid video select detect plugin name")
video_path = VideoFile.path(result[2])
plugin().video_select_detect(video_path, **opts)
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_transform_set_select_curate_min_size(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
self.mock_transform_set()
MinSizeTransformSetSelectCuratePlugin(
).transform_set_select_curate(1, {'min-size': '300'}) # noqa
# db
result = TransformSetModel().select(1)
self.assertEqual(result, (1, 'face', 1, None))
result = TransformModel().list(1)
self.assertEqual(len(result), 5)
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, 2, 1])
t = list(result[2])
json.loads(t.pop(3))
self.assertEqual(t, [3, 1, 3, 0])
t = list(result[3])
json.loads(t.pop(3))
self.assertEqual(t, [4, 1, 4, 1])
t = list(result[4])
json.loads(t.pop(3))
self.assertEqual(t, [5, 1, 5, 0])
def test_transform_set_select_extract_pad_size(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
self.mock_transform_set()
MaxSizeTransformSetSelectExtractPlugin(
).transform_set_select_extract(1, {'max-size': '200'}) # noqa
with mock.patch.dict(os.environ, {'MODEL_LIST_LENGTH': '2'}):
PadTransformSetSelectExtractPlugin(
).transform_set_select_extract(2, {'size': '200'}) # noqa
# db
result = TransformSetModel().select(3)
self.assertEqual(result, (3, 'pad', 1, 2))
result = TransformModel().list(3)
self.assertEqual(len(result), 5)
t = list(result[0])
json.loads(t.pop(3))
self.assertEqual(t, [11, 3, 1, 0])
t = list(result[1])
json.loads(t.pop(3))
self.assertEqual(t, [12, 3, 2, 0])
t = list(result[2])
json.loads(t.pop(3))
self.assertEqual(t, [13, 3, 3, 0])
t = list(result[3])
json.loads(t.pop(3))
self.assertEqual(t, [14, 3, 4, 0])
t = list(result[4])
json.loads(t.pop(3))
self.assertEqual(t, [15, 3, 5, 0])
# files
p1 = TransformSetSubDir.path(3)
# transforms
self.assertEqual(
cv2.imread(TransformFile.path(p1, 11, 'jpg')).shape[:2],
(200, 200)) # noqa
self.assertEqual(
cv2.imread(TransformFile.path(p1, 12, 'jpg')).shape[:2],
(200, 200)) # noqa
self.assertEqual(
cv2.imread(TransformFile.path(p1, 13, 'jpg')).shape[:2],
(200, 200)) # noqa
self.assertEqual(
cv2.imread(TransformFile.path(p1, 14, 'jpg')).shape[:2],
(200, 200)) # noqa
self.assertEqual(
cv2.imread(TransformFile.path(p1, 15, 'jpg')).shape[:2],
(200, 200)) # 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 test_path(self):
with mock.patch.dict(os.environ, {'DEEPSTAR_PATH': 'test'}):
self.assertEqual(
VideoFile.path('test'),
os.path.realpath('test/files/videos/test')) # noqa