| 1527 | return seg, bc, rgb, im.shape[2:] |
| 1528 | |
| 1529 | class ClassifierSegRunner: |
| 1530 | def __init__(self, dataset, recover_image=None): |
| 1531 | # The dataset contains explicit segmentations |
| 1532 | if recover_image is None: |
| 1533 | recover_image = reverse_normalize_from_transform(dataset) |
| 1534 | self.recover_image = recover_image |
| 1535 | self.dataset = dataset |
| 1536 | def get_label_and_category_names(self): |
| 1537 | catnames = self.dataset.categories |
| 1538 | label_and_cat_names = [(readable(label), |
| 1539 | catnames[self.dataset.label_category[i]]) |
| 1540 | for i, label in enumerate(self.dataset.labels)] |
| 1541 | return label_and_cat_names, catnames |
| 1542 | def run_and_segment_batch(self, batch, model, |
| 1543 | want_bincount=False, want_rgb=False): |
| 1544 | ''' |
| 1545 | Runs the dissected model on one batch of the dataset, and |
| 1546 | returns a multilabel semantic segmentation for the data. |
| 1547 | Given a batch of size (n, c, y, x) the segmentation should |
| 1548 | be a (long integer) tensor of size (n, d, y//r, x//r) where |
| 1549 | d is the maximum number of simultaneous labels given to a pixel, |
| 1550 | and where r is some (optional) resolution reduction factor. |
| 1551 | In the segmentation returned, the label `0` is reserved for |
| 1552 | the background "no-label". |
| 1553 | |
| 1554 | In addition to the segmentation, bc, rgb, and shape are returned |
| 1555 | where bc is a per-image bincount counting returned label pixels, |
| 1556 | rgb is a viewable (n, y, x, rgb) byte image tensor for the data |
| 1557 | for visualizations (reversing normalizations, for example), and |
| 1558 | shape is the (y, x) size of the data. If want_bincount or |
| 1559 | want_rgb are False, those return values may be None. |
| 1560 | ''' |
| 1561 | im, seg, bc = batch |
| 1562 | device = next(model.parameters()).device |
| 1563 | if want_rgb: |
| 1564 | rgb = self.recover_image(im.clone() |
| 1565 | ).permute(0, 2, 3, 1).mul_(255).clamp(0, 255).byte() |
| 1566 | else: |
| 1567 | rgb = None |
| 1568 | # Run the model. |
| 1569 | model(im.to(device)) |
| 1570 | return seg, bc, rgb, im.shape[2:] |
| 1571 | |
| 1572 | class GeneratorSegRunner: |
| 1573 | def __init__(self, segmenter): |