davis.patch

index d831be6..8891b88 100644
--- a/davis2017/davis.py
+++ b/davis2017/davis.py
@@ -8,10 +8,11 @@ from PIL import Image
 class DAVIS(object):
     SUBSET_OPTIONS = ['train', 'val', 'test-dev', 'test-challenge']
     TASKS = ['semi-supervised', 'unsupervised']
+    YEARS = ['2016', '2017', '2019']
     DATASET_WEB = 'https://davischallenge.org/davis2017/code.html'
     VOID_LABEL = 255
 
-    def __init__(self, root, task='unsupervised', subset='val', sequences='all', resolution='480p', codalab=False):
+    def __init__(self, root, task='unsupervised', subset='val', sequences='all', resolution='480p', codalab=False, year='2017'):
         """
         Class to read the DAVIS dataset
         :param root: Path to the DAVIS folder that contains JPEGImages, Annotations, etc. folders.
@@ -24,6 +25,8 @@ class DAVIS(object):
             raise ValueError(f'Subset should be in {self.SUBSET_OPTIONS}')
         if task not in self.TASKS:
             raise ValueError(f'The only tasks that are supported are {self.TASKS}')
+        if year not in self.YEARS:
+            raise ValueError(f'Year should be one of the following {self.YEARS}')
 
         self.task = task
         self.subset = subset
@@ -31,8 +34,12 @@ class DAVIS(object):
         self.img_path = os.path.join(self.root, 'JPEGImages', resolution)
         annotations_folder = 'Annotations' if task == 'semi-supervised' else 'Annotations_unsupervised'
         self.mask_path = os.path.join(self.root, annotations_folder, resolution)
-        year = '2019' if task == 'unsupervised' and (subset == 'test-dev' or subset == 'test-challenge') else '2017'
-        self.imagesets_path = os.path.join(self.root, 'ImageSets', year)
+
+        self.year = year
+        if self.year == '2019' and not (task == 'unsupervised' and (subset == 'test-dev' or subset == 'test-challenge')):
+            raise ValueError("Set 'task' to 'unsupervised' and subset to 'test-dev' or 'test-challenge'")
+
+        self.imagesets_path = os.path.join(self.root, 'ImageSets', self.year)
 
         self._check_directories()
 
@@ -95,6 +102,10 @@ class DAVIS(object):
             tmp = tmp * np.arange(1, num_objects + 1)[:, None, None, None]
             masks = (tmp == masks[None, ...])
             masks = masks > 0
+        else:
+            # for single object evaluation (e.g. DAVIS2016)
+            masks = np.expand_dims(masks, axis=0)
+            masks = masks > 0
         return masks, masks_void, masks_id
 
     def get_sequences(self):
diff --git a/davis2017/evaluation.py b/davis2017/evaluation.py
index 7bfb80f..eae777c 100644
--- a/davis2017/evaluation.py
+++ b/davis2017/evaluation.py
@@ -12,7 +12,7 @@ from scipy.optimize import linear_sum_assignment
 
 
 class DAVISEvaluation(object):
-    def __init__(self, davis_root, task, gt_set, sequences='all', codalab=False):
+    def __init__(self, davis_root, task, gt_set, sequences='all', codalab=False, year='2017'):
         """
         Class to evaluate DAVIS sequences from a certain set and for a certain task
         :param davis_root: Path to the DAVIS folder that contains JPEGImages, Annotations, etc. folders.
@@ -22,7 +22,8 @@ class DAVISEvaluation(object):
         """
         self.davis_root = davis_root
         self.task = task
-        self.dataset = DAVIS(root=davis_root, task=task, subset=gt_set, sequences=sequences, codalab=codalab)
+        self.year = year
+        self.dataset = DAVIS(root=davis_root, task=task, subset=gt_set, sequences=sequences, codalab=codalab, year=self.year)
 
     @staticmethod
     def _evaluate_semisupervised(all_gt_masks, all_res_masks, all_void_masks, metric):
@@ -77,10 +78,12 @@ class DAVISEvaluation(object):
         if 'F' in metric:
             metrics_res['F'] = {"M": [], "R": [], "D": [], "M_per_object": {}}
 
+        separate_objects_masks = self.year != '2016'
+
         # Sweep all sequences
         results = Results(root_dir=res_path)
         for seq in tqdm(list(self.dataset.get_sequences())):
-            all_gt_masks, all_void_masks, all_masks_id = self.dataset.get_all_masks(seq, True)
+            all_gt_masks, all_void_masks, all_masks_id = self.dataset.get_all_masks(seq, separate_objects_masks)
             if self.task == 'semi-supervised':
                 all_gt_masks, all_masks_id = all_gt_masks[:, 1:-1, :, :], all_masks_id[1:-1]
             all_res_masks = results.read_masks(seq, all_masks_id)
diff --git a/evaluation_method.py b/evaluation_method.py
index 04f67d1..d364f81 100644
--- a/evaluation_method.py
+++ b/evaluation_method.py
@@ -20,6 +20,8 @@ parser.add_argument('--task', type=str, help='Task to evaluate the results', def
                     choices=['semi-supervised', 'unsupervised'])
 parser.add_argument('--results_path', type=str, help='Path to the folder containing the sequences folders',
                     required=True)
+parser.add_argument("--year", type=str, help="Davis dataset year (default: 2017)", default='2017',
+                    choices=['2016', '2017', '2019'])
 args, _ = parser.parse_known_args()
 csv_name_global = f'global_results-{args.set}.csv'
 csv_name_per_sequence = f'per-sequence_results-{args.set}.csv'
@@ -34,7 +36,7 @@ if os.path.exists(csv_name_global_path) and os.path.exists(csv_name_per_sequence
 else:
     print(f'Evaluating sequences for the {args.task} task...')
     # Create dataset and evaluate
-    dataset_eval = DAVISEvaluation(davis_root=args.davis_path, task=args.task, gt_set=args.set)
+    dataset_eval = DAVISEvaluation(davis_root=args.davis_path, task=args.task, gt_set=args.set, year=args.year)
     metrics_res = dataset_eval.evaluate(args.results_path)
     J, F = metrics_res['J'], metrics_res['F']