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)