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)