Rework predictable classes to work with batch viz
This makes modifications to dev/class_head_fix
(!14 (merged)) so the visualized batch will reflect the predictable classes.
Note the new draw_item
doctest:
>>> from geowatch.tasks.fusion.datamodules.kwcoco_dataset import * # NOQA
>>> import geowatch
>>> anchors = np.array([[0.1, 0.1]])
>>> coco_dset = geowatch.coerce_kwcoco('vidshapes1', num_frames=4, num_tracks=40, anchors=anchors)
>>> self = KWCocoVideoDataset(coco_dset, time_dims=4, window_dims=(300, 300))
>>> self._notify_about_tasks(predictable_classes=['star', 'eff', 'background'])
>>> self.requested_tasks['change'] = False
>>> from geowatch.tasks.fusion import utils
>>> utils.category_tree_ensure_color(self.predictable_classes)
>>> index = self.new_sample_grid['targets'][self.new_sample_grid['positives_indexes'][0]]
>>> item = self[index]
>>> canvas = self.draw_item(item, draw_weights=False)
>>> # xdoctest: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl()
>>> label_to_color = {
>>> node: data['color']
>>> for node, data in self.predictable_classes.graph.nodes.items()}
>>> label_to_color = ub.sorted_keys(label_to_color)
>>> legend_img = utils._memo_legend(label_to_color)
>>> legend_img = kwimage.imresize(legend_img, scale=4.0)
>>> show_canvas = kwimage.stack_images([canvas, legend_img], axis=1)
>>> kwplot.imshow(show_canvas)
>>> kwplot.show_if_requested()
Which produces:
You'll note that the "superstars" visible in the image and saliency head have moved to "background" in the class head. If I remove "eff" from the list of predictable classes and rerun then the image becomes:
More work can be done here to allow mapping classes to something other than background, but that is the default.
Edited by Jon Crall