Skip to content

Rework predictable classes to work with batch viz

Jon Crall requested to merge dev/class_head_fix-tweaks into dev/class_head_fix

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:

image

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:

image

More work can be done here to allow mapping classes to something other than background, but that is the default.

Edited by Jon Crall

Merge request reports