Read metadata for actual dataset in Visualizer, if available

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/516

# Context

D2go allows for training with more than one datasets, and as long as the categories are consistent, the IDs do not necessarily have to correspond to each other between annotations of two different data sets.

It is still loaded correctly to the data loader, and the training works as expected.

# Problem

However, I observed weird mis-labelleing issues in the Visualizer for Tensorboard. Originally I thought this was a data/conversion issue, but upon inspecting the logs I see that the data is loaded correctly. See example below.

{F924075931}

"Plant" labelled as "Refrigerator", "Floor" labelled as "Lamp"

{F924078113}

... but the loaded annotations doesn't actually contain any samples of "Refrigerator".

The reason is that the Visualizer always loads the metadata (and thus the labels) from the first train data set, but the order of the categories between the data sets may not be consistent, but still be a valid training run.

# Fix

If there is a data set name associated with the data to visualize, use that to fetch the metadata, and the correct labels, otherwise default to the first data set (current situation).

Reviewed By: wat3rBro

Differential Revision:
D44495363

Privacy Context Container: L1127277

fbshipit-source-id: 37b940d393aa794cd2f39aabdc66c6d23abd8000

d2go/utils/visualization.py

@@ -60,7 +60,12 @@ class VisualizerWrapper(object):
 
         img = per_image["image"].permute(1, 2, 0).cpu().detach().numpy()
         img = utils.convert_image_to_rgb(img, cfg.INPUT.FORMAT)
-        metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0])
+
+        if "dataset_name" in input_dict:
+            metadata = MetadataCatalog.get(input_dict["dataset_name"])
+        else:
+            metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0])
+
         scale = 2.0
         visualizer = self.visualizer(img, metadata=metadata, scale=scale)