| 1503 | return data.mul(self.stdev.to(device)).add_(self.mean.to(device)) |
| 1504 | |
| 1505 | class ImageOnlySegRunner: |
| 1506 | def __init__(self, dataset, recover_image=None): |
| 1507 | if recover_image is None: |
| 1508 | recover_image = reverse_normalize_from_transform(dataset) |
| 1509 | self.recover_image = recover_image |
| 1510 | self.dataset = dataset |
| 1511 | def get_label_and_category_names(self): |
| 1512 | return [('-', '-')], ['-'] |
| 1513 | def run_and_segment_batch(self, batch, model, |
| 1514 | want_bincount=False, want_rgb=False): |
| 1515 | [im] = batch |
| 1516 | device = next(model.parameters()).device |
| 1517 | if want_rgb: |
| 1518 | rgb = self.recover_image(im.clone() |
| 1519 | ).permute(0, 2, 3, 1).mul_(255).clamp(0, 255).byte() |
| 1520 | else: |
| 1521 | rgb = None |
| 1522 | # Stubs for seg and bc |
| 1523 | seg = torch.zeros(im.shape[0], 1, 1, 1, dtype=torch.long) |
| 1524 | bc = torch.ones(im.shape[0], 1, dtype=torch.long) |
| 1525 | # Run the model. |
| 1526 | model(im.to(device)) |
| 1527 | return seg, bc, rgb, im.shape[2:] |
| 1528 | |
| 1529 | class ClassifierSegRunner: |
| 1530 | def __init__(self, dataset, recover_image=None): |