def load(self): video = cv2.VideoCapture(self.video_metadata.file) video_start = self.offset if self.offset else 0 video.set(cv2.CAP_PROP_POS_FRAMES, video_start) _, frame = video.read() frame_ind = video_start - 1 info = None if frame is not None: (height, width, channels) = frame.shape info = FrameInfo(height, width, channels, ColorSpace.BGR) frames = [] while frame is not None: frame_ind += 1 eva_frame = Frame(frame_ind, frame, info) if self.skip_frames > 0 and frame_ind % self.skip_frames != 0: _, frame = video.read() continue frames.append(eva_frame) if self.limit and frame_ind >= self.limit: return FrameBatch(frames, info) if len(frames) % self.batch_size == 0: yield FrameBatch(frames, info) frames = [] _, frame = video.read() if frames: return FrameBatch(frames, info)
def test_should_return_only_frames_satisfy_predicate(self): frame_1 = Frame(1, np.ones((1, 1)), None) frame_2 = Frame(1, 2 * np.ones((1, 1)), None) frame_3 = Frame(1, 3 * np.ones((1, 1)), None) outcome_1 = Prediction(frame_1, ["car", "bus"], [0.5, 0.6]) outcome_2 = Prediction(frame_2, ["bus"], [0.5, 0.6]) outcome_3 = Prediction(frame_3, ["car", "train"], [0.5, 0.6]) batch = FrameBatch( frames=[ frame_1, frame_2, frame_3, ], info=None, outcomes={"test": [outcome_1, outcome_2, outcome_3]}) predicate = Predicate(name="test", predicate=lambda prediction: prediction.eq("car") and not prediction.eq("bus")) predicate_executor = PredicateExecutor(None, predicate=predicate) expected = FrameBatch(frames=[frame_3], info=None, outcomes={"test": [outcome_3]}) filtered = predicate_executor.execute(batch) self.assertEqual(expected, filtered)
def test_frames_as_numpy_array_should_frames_as_numpy_array(self): batch = FrameBatch(frames=[ Frame(1, np.ones((1, 1)), None), Frame(1, np.ones((1, 1)), None) ], info=None) expected = list(np.ones((2, 1, 1))) actual = list(batch.frames_as_numpy_array()) self.assertEqual(expected, actual)
def test_return_only_frames_specified_in_the_indices(self): batch = FrameBatch(frames=[ Frame(1, np.ones((1, 1)), None), Frame(1, np.ones((1, 1)), None) ], info=None) expected = FrameBatch(frames=[Frame(1, np.ones((1, 1)), None)], info=None) output = batch[[0]] self.assertEqual(expected, output)
def test_slicing_should_word_for_negative_stop_value(self): batch = FrameBatch(frames=[ Frame(1, np.ones((1, 1)), None), Frame(1, 2 * np.ones((1, 1)), None) ], info=None, outcomes={'test': [[None], [None]]}) expected = FrameBatch(frames=[Frame(1, np.ones((1, 1)), None)], info=None, outcomes={'test': [[None]]}) self.assertEqual(expected, batch[:-1])
def test_fetching_frames_by_index_should_also_return_outcomes(self): batch = FrameBatch(frames=[ Frame(1, np.ones((1, 1)), None), Frame(1, np.ones((1, 1)), None) ], info=None, outcomes={'test': [[None], [None]]}) expected = FrameBatch(frames=[Frame(1, np.ones((1, 1)), None)], info=None, outcomes={'test': [[None]]}) self.assertEqual(expected, batch[[0]])
def test_slicing_on_batched_should_return_new_batch_frame(self): batch = FrameBatch(frames=[ Frame(1, np.ones((1, 1)), None), Frame(1, 2 * np.ones((1, 1)), None) ], info=None, outcomes={'test': [[None], [None]]}) expected = FrameBatch(frames=[Frame(1, np.ones((1, 1)), None)], info=None, outcomes={'test': [[None]]}) self.assertEqual(batch, batch[:]) self.assertEqual(expected, batch[:-1])
def test_slicing_should_work_with_skip_value(self): batch = FrameBatch(frames=[ Frame(1, np.ones((1, 1)), None), Frame(1, 2 * np.ones((1, 1)), None), Frame(1, np.ones((1, 1)), None) ], info=None, outcomes={'test': [[None], [None], [None]]}) expected = FrameBatch(frames=[ Frame(1, np.ones((1, 1)), None), Frame(1, np.ones((1, 1)), None) ], info=None, outcomes={'test': [[None], [None]]}) self.assertEqual(expected, batch[::2])
def test_should_return_list_of_predictions_for_each_frame_in_batch(self): batch = FrameBatch(frames=[Frame(1, np.ones((1, 1)), None)], info=None) predictions = [['A', 'B']] scores = [[1, 1]] expected = [Prediction(batch.frames[0], predictions[0], scores[0])] actual = Prediction.predictions_from_batch_and_lists( batch, predictions, scores) self.assertEqual(expected, actual)
def test_should_check_if_batch_frames_equivalent_to_number_of_scores(self): batch = FrameBatch(frames=[Frame(1, np.ones((1, 1)), None)], info=None) predictions = [['A', 'B']] scores = [] self.assertRaises( AssertionError, lambda x=None: Prediction.predictions_from_batch_and_lists( batch, predictions, scores))
def execute(self, batch: FrameBatch): predictions = batch.get_outcomes_for(self.predicate.name) required_frame_ids = [] for i, prediction in enumerate(predictions): if self.predicate(prediction): required_frame_ids.append(i) return batch[required_frame_ids]
def classify(self, batch: FrameBatch) -> List[Prediction]: frames = batch.frames_as_numpy_array() (prediction_classes, prediction_scores, prediction_boxes) = self._get_predictions(frames) return Prediction.predictions_from_batch_and_lists( batch, prediction_classes, prediction_scores, boxes=prediction_boxes)
def test_should_return_batches_equivalent_to_number_of_frames(self): frame_dog = Frame(1, self._load_image( os.path.join(self.base_path, 'data', 'dog.jpeg')), None) frame_dog_cat = Frame(1, self._load_image( os.path.join(self.base_path, 'data', 'dog_cat.jpg')), None) frame_batch = FrameBatch([frame_dog, frame_dog_cat], None) detector = FastRCNNObjectDetector() result = detector.classify(frame_batch) self.assertEqual(["dog"], result[0].labels) self.assertEqual(["cat", "dog"], result[1].labels)
def test_should_return_batches_equivalent_to_number_of_frames(self): frame_dog = Frame(1, self._load_image( os.path.join(self.base_path, 'data', 'dog.jpeg')), None) frame_dog_cat = Frame(1, self._load_image( os.path.join(self.base_path, 'data', 'dog_cat.jpg')), None) frame_batch = FrameBatch([frame_dog, frame_dog_cat], None) estimator = DepthEstimator() result = estimator.process_frames(frame_batch) print(result) print(result[0].frame) print(result[0].depth) print(result[0].segm)
def load_data(self): """ method to load images and depth base_path: file location root path num_images: number of images to load, image names have to be in format 0.png, 1.png ... num_images.png both for depth and RGB images returns Framebatch object containing all the frames Depth images in numpy array """ images = self._load_data_from_folder(self.images_folder_path) depths = self._load_data_from_folder(self.depths_folder_path, True) return FrameBatch(images, None), np.array(depths)
def test_get_outcome_from_non_existing_udf_name_returns_empty_list(self): batch = FrameBatch(frames=[Frame(1, np.ones((1, 1)), None)], info=None) self.assertEqual([], batch.get_outcomes_for('test'))
def test_set_outcomes_method_should_set_the_predictions_with_udf_name( self): batch = FrameBatch(frames=[Frame(1, np.ones((1, 1)), None)], info=None) batch.set_outcomes('test', [None]) self.assertEqual([None], batch.get_outcomes_for('test'))
def process_frames(self, batch: FrameBatch) -> List[DepthEstimationResult]: frames = batch.frames_as_numpy_array() (segms, depths) = self._get_depth_estimation(frames) return DepthEstimationResult.depth_estimation_result_from_batch_and_lists( batch, segms, depths)