from __future__ import absolute_import, division import json import os import tempfile import time import unittest from functools import wraps import numpy as np import pandas as pd import pytest from davisinteractive.common import Path, patch from davisinteractive.connector.local import LocalConnector from davisinteractive.dataset import Davis from davisinteractive.session import DavisInteractiveSession from davisinteractive.utils.scribbles import annotated_frames, is_empty EMPTY_SCRIBBLE = { 'scribbles': [[] for _ in range(69)], 'sequence': 'test-sequence' } def dataset(subset, **samples): def decorator(func): @wraps(func) def wrapper(*args, **kwargs): sequence_set = Davis.sets[subset] new_set = list(samples.keys()) new_set.sort() # Necessary as the keys order is not garanteed original = Davis.dataset.copy() for seq in samples: Davis.dataset[seq].update(samples[seq]) Davis.sets[subset] = new_set result = func(*args, **kwargs) # Recover orignal state at Davis class Davis.sets[subset] = sequence_set Davis.dataset = original return result return wrapper return decorator class TestDavisInteractiveSession(unittest.TestCase): @patch.object(Davis, 'check_files', return_value=None) def test_subset(self, mock_davis): davis_root = '/tmp/DAVIS' session = DavisInteractiveSession(subset='val', davis_root=davis_root) session.__enter__() session = DavisInteractiveSession(subset='train', davis_root=davis_root) session.__enter__() session1 = DavisInteractiveSession( subset='test-dev', davis_root=davis_root) session2 = DavisInteractiveSession(subset='xxxx', davis_root=davis_root) with pytest.raises(ValueError): session1.__enter__() with pytest.raises(ValueError): session2.__enter__() assert mock_davis.call_count == 2 @patch.object( LocalConnector, 'post_predicted_masks', return_value=EMPTY_SCRIBBLE) @patch.object(LocalConnector, 'get_report', return_value=pd.DataFrame()) @patch.object(LocalConnector, 'get_scribble', return_value=EMPTY_SCRIBBLE) @patch.object( LocalConnector, 'get_samples', return_value=([('bear', 2), ('bear', 1)], 5, None)) def test_interactions_limit(self, mock_start_session, mock_get_scribble, mock_get_report, mock_submit_masks): davis_root = '/tmp/DAVIS' with DavisInteractiveSession( davis_root=davis_root, max_nb_interactions=5, report_save_dir=tempfile.mkdtemp(), max_time=None) as session: assert mock_start_session.call_count == 1 for i in range(7): assert session.next() seq, scribbles, new_seq = session.get_scribbles() assert seq == 'bear' assert is_empty(scribbles) if i % 5 == 0: assert new_seq, i session.submit_masks(None) assert mock_get_scribble.call_count == 2 assert mock_submit_masks.call_count == 7 @dataset('train', bear={'num_frames': 2, 'num_scribbles': 1}) @patch.object(Davis, '_download_scribbles', return_value=None) def test_integration_single(self, mock_davis): dataset_dir = Path(__file__).parent.joinpath('test_data', 'DAVIS') tmp_dir = Path(tempfile.mkdtemp()) with DavisInteractiveSession( davis_root=dataset_dir, subset='train', max_nb_interactions=5, report_save_dir=tmp_dir, max_time=None) as session: count = 0 temp_csv = tmp_dir / ("%s.tmp.csv" % session.report_name) final_csv = tmp_dir / ("%s.csv" % session.report_name) while session.next(): assert not final_csv.exists() assert temp_csv.exists() df = pd.read_csv(temp_csv, index_col=0) assert df.shape == (count * 2, 10) seq, scribble, new_seq = session.get_scribbles(only_last=True) assert new_seq == (count == 0) assert seq == 'bear' if count == 0: with dataset_dir.joinpath('Scribbles', 'bear', '001.json').open() as fp: sc = json.load(fp) assert sc == scribble else: assert annotated_frames(scribble) == [1] assert not is_empty(scribble) assert len(scribble['scribbles']) == 2 assert len(scribble['scribbles'][1]) > 0 assert len(scribble['scribbles'][0]) == 0 # Simulate model predicting masks pred_masks = np.zeros((2, 480, 854)) session.submit_masks( pred_masks, next_scribble_frame_candidates=[1]) if count > 0: assert df.sequence.unique() == ['bear'] assert np.all(df.interaction.unique() == [i + 1 for i in range(count)]) assert np.all(df.object_id.unique() == [1]) count += 1 assert count == 5 assert final_csv.exists() assert not temp_csv.exists() assert mock_davis.call_count == 0 @dataset('train', bear={'num_frames': 2, 'num_scribbles': 2}) @patch.object(Davis, '_download_scribbles', return_value=None) def test_integration_multiple(self, mock_davis): dataset_dir = Path(__file__).parent.joinpath('test_data', 'DAVIS') with DavisInteractiveSession( davis_root=dataset_dir, subset='train', max_nb_interactions=5, report_save_dir=tempfile.mkdtemp(), max_time=None) as session: count = 0 while session.next(): seq, scribble, new_seq = session.get_scribbles() assert new_seq == (count == 0 or count == 5) assert seq == 'bear' if count == 0: with dataset_dir.joinpath('Scribbles', 'bear', '001.json').open() as fp: sc = json.load(fp) assert sc == scribble if count == 5: with dataset_dir.joinpath('Scribbles', 'bear', '002.json').open() as fp: sc = json.load(fp) assert sc == scribble assert not is_empty(scribble) # Simulate model predicting masks pred_masks = np.zeros((2, 480, 854)) session.submit_masks(pred_masks) count += 1 assert count == 10 assert mock_davis.call_count == 0 @dataset( 'train', bear={ 'num_frames': 2, 'num_scribbles': 2 }, tennis={ 'num_frames': 2, 'num_scribbles': 1 }) @patch.object(Davis, '_download_scribbles', return_value=None) def test_integration_multiple_sequences(self, mock_davis): dataset_dir = Path(__file__).parent.joinpath('test_data', 'DAVIS') with DavisInteractiveSession( davis_root=dataset_dir, subset='train', max_nb_interactions=4, report_save_dir=tempfile.mkdtemp(), max_time=None) as session: count = 0 for seq, scribble, new_seq in session.scribbles_iterator(): assert new_seq == (count == 0 or count == 4 or count == 8) if count < 8: assert seq == 'bear', count else: assert seq == 'tennis', count if count == 0: with dataset_dir.joinpath('Scribbles', 'bear', '001.json').open() as fp: sc = json.load(fp) assert sc == scribble if count == 4: with dataset_dir.joinpath('Scribbles', 'bear', '002.json').open() as fp: sc = json.load(fp) assert sc == scribble if count == 8: with dataset_dir.joinpath('Scribbles', 'tennis', '001.json').open() as fp: sc = json.load(fp) assert sc == scribble assert not is_empty(scribble) # Simulate model predicting masks pred_masks = np.ones((2, 480, 854)) session.submit_masks(pred_masks) count += 1 assert count == 12 df = session.get_report() assert mock_davis.call_count == 0 assert df.shape == (2 * 4 * 2 * 1 + 4 * 2 * 2, 10) global_summary_file = os.path.join(tempfile.mkdtemp(), 'summary.json') summary = session.get_global_summary() self.assertFalse(os.path.exists(global_summary_file)) self.assertTrue('auc' in summary) self.assertTrue('metric_at_threshold' in summary) self.assertEqual(summary['metric_at_threshold']['threshold'], 60) np.testing.assert_almost_equal(summary['metric_at_threshold']['J_AND_F'], 0.035155) self.assertTrue('curve' in summary) curve = summary['curve'] self.assertEqual(len(curve['J_AND_F']), 6) self.assertEqual(len(curve['time']), 6) summary = session.get_global_summary(save_file=global_summary_file) self.assertTrue(os.path.exists(global_summary_file)) @dataset( 'train', bear={ 'num_frames': 2, 'num_scribbles': 2 }, tennis={ 'num_frames': 2, 'num_scribbles': 1 }) @patch.object(Davis, '_download_scribbles', return_value=None) def test_integration_multiple_sequences_metric(self, mock_davis): dataset_dir = Path(__file__).parent.joinpath('test_data', 'DAVIS') with DavisInteractiveSession( davis_root=dataset_dir, subset='train', max_nb_interactions=4, report_save_dir=tempfile.mkdtemp(), metric_to_optimize='J', max_time=None) as session: count = 0 for seq, scribble, new_seq in session.scribbles_iterator(): assert new_seq == (count == 0 or count == 4 or count == 8) if count < 8: assert seq == 'bear', count else: assert seq == 'tennis', count if count == 0: with dataset_dir.joinpath('Scribbles', 'bear', '001.json').open() as fp: sc = json.load(fp) assert sc == scribble if count == 4: with dataset_dir.joinpath('Scribbles', 'bear', '002.json').open() as fp: sc = json.load(fp) assert sc == scribble if count == 8: with dataset_dir.joinpath('Scribbles', 'tennis', '001.json').open() as fp: sc = json.load(fp) assert sc == scribble assert not is_empty(scribble) # Simulate model predicting masks pred_masks = np.ones((2, 480, 854)) session.submit_masks(pred_masks) count += 1 assert count == 12 df = session.get_report() assert mock_davis.call_count == 0 assert df.shape == (2 * 4 * 2 * 1 + 4 * 2 * 2, 10) global_summary_file = os.path.join(tempfile.mkdtemp(), 'summary.json') summary = session.get_global_summary() self.assertFalse(os.path.exists(global_summary_file)) self.assertTrue('auc' in summary) self.assertTrue('metric_at_threshold' in summary) self.assertEqual(summary['metric_at_threshold']['threshold'], 60) np.testing.assert_almost_equal(summary['metric_at_threshold']['J'], 0.07031) self.assertTrue('curve' in summary) curve = summary['curve'] self.assertEqual(len(curve['J']), 6) self.assertEqual(len(curve['time']), 6) summary = session.get_global_summary(save_file=global_summary_file) self.assertTrue(os.path.exists(global_summary_file)) @dataset('train', blackswan={'num_frames': 6, 'num_scribbles': 1}) @patch.object(Davis, '_download_scribbles', return_value=None) def test_integration_single_only_last(self, mock_davis): dataset_dir = Path(__file__).parent.joinpath('test_data', 'DAVIS') with DavisInteractiveSession( davis_root=dataset_dir, subset='train', max_nb_interactions=4, report_save_dir=tempfile.mkdtemp(), max_time=None) as session: count = 0 annotated_frames_list = [] while session.next(): seq, scribble, new_seq = session.get_scribbles(only_last=True) assert new_seq == (count == 0) assert seq == 'blackswan' if count == 0: with dataset_dir.joinpath('Scribbles', 'blackswan', '001.json').open() as fp: sc = json.load(fp) assert sc == scribble else: assert len(annotated_frames(scribble)) == 1 a_fr = annotated_frames(scribble)[0] assert a_fr not in annotated_frames_list annotated_frames_list.append(a_fr) assert not is_empty(scribble) # Simulate model predicting masks pred_masks = np.zeros((6, 480, 854)) session.submit_masks(pred_masks) count += 1 assert count == 4 assert mock_davis.call_count == 0 @dataset('train', bear={'num_frames': 2, 'num_scribbles': 1}) @patch.object(Davis, '_download_scribbles', return_value=None) def test_integration_single_timeout(self, mock_davis): dataset_dir = Path(__file__).parent.joinpath('test_data', 'DAVIS') with DavisInteractiveSession( davis_root=dataset_dir, subset='train', max_nb_interactions=None, max_time=1, report_save_dir=tempfile.mkdtemp()) as session: count = 0 while session.next(): seq, scribble, new_seq = session.get_scribbles(only_last=True) assert new_seq == (count == 0) assert seq == 'bear' with dataset_dir.joinpath('Scribbles', 'bear', '001.json').open() as fp: sc = json.load(fp) assert sc == scribble assert not is_empty(scribble) # Simulate model predicting masks pred_masks = np.zeros((2, 480, 854)) time.sleep(1.2) session.submit_masks(pred_masks) count += 1 assert count == 1 assert mock_davis.call_count == 0 @dataset('train', bear={'num_frames': 2, 'num_scribbles': 1}) @patch.object(Davis, '_download_scribbles', return_value=None) def test_report_folder_creation(self, mock_davis): dataset_dir = Path(__file__).parent.joinpath('test_data', 'DAVIS') tmp_dir = Path(tempfile.mkdtemp()) / 'test' assert not tmp_dir.exists() session = DavisInteractiveSession( davis_root=dataset_dir, subset='train', report_save_dir=tmp_dir) assert tmp_dir.exists() assert mock_davis.call_count == 0 @dataset( 'train', bear={ 'num_frames': 2, 'num_scribbles': 2 }, tennis={ 'num_frames': 2, 'num_scribbles': 1 }) @patch.object(Davis, '_download_scribbles', return_value=None) def test_shuffle(self, mock_davis): dataset_dir = Path(__file__).parent.joinpath('test_data', 'DAVIS') with DavisInteractiveSession( davis_root=dataset_dir, subset='train', shuffle=True, report_save_dir=tempfile.mkdtemp()) as session: assert ('bear', 1) in session.samples assert ('bear', 2) in session.samples assert ('tennis', 1) in session.samples assert mock_davis.call_count == 0