Prototype references

references/classification/presets.py

@@ -1,5 +1,5 @@
 import torch
-from torchvision.transforms import autoaugment, transforms
+from torchvision.prototype import transforms
 from torchvision.transforms.functional import InterpolationMode
 
 
@@ -17,22 +17,24 @@ def __init__(
         augmix_severity=3,
         random_erase_prob=0.0,
     ):
-        trans = [transforms.RandomResizedCrop(crop_size, interpolation=interpolation)]
+        trans = [
+            transforms.ToImageTensor(),
+            transforms.RandomResizedCrop(crop_size, interpolation=interpolation, antialias=True),
+        ]
         if hflip_prob > 0:
-            trans.append(transforms.RandomHorizontalFlip(hflip_prob))
+            trans.append(transforms.RandomHorizontalFlip(p=hflip_prob))
         if auto_augment_policy is not None:
             if auto_augment_policy == "ra":
-                trans.append(autoaugment.RandAugment(interpolation=interpolation, magnitude=ra_magnitude))
+                trans.append(transforms.RandAugment(interpolation=interpolation, magnitude=ra_magnitude))
             elif auto_augment_policy == "ta_wide":
-                trans.append(autoaugment.TrivialAugmentWide(interpolation=interpolation))
+                trans.append(transforms.TrivialAugmentWide(interpolation=interpolation))
             elif auto_augment_policy == "augmix":
-                trans.append(autoaugment.AugMix(interpolation=interpolation, severity=augmix_severity))
+                trans.append(transforms.AugMix(interpolation=interpolation, severity=augmix_severity))
             else:
-                aa_policy = autoaugment.AutoAugmentPolicy(auto_augment_policy)
-                trans.append(autoaugment.AutoAugment(policy=aa_policy, interpolation=interpolation))
+                aa_policy = transforms.AutoAugmentPolicy(auto_augment_policy)
+                trans.append(transforms.AutoAugment(policy=aa_policy, interpolation=interpolation))
         trans.extend(
             [
-                transforms.PILToTensor(),
                 transforms.ConvertImageDtype(torch.float),
                 transforms.Normalize(mean=mean, std=std),
             ]
@@ -59,9 +61,9 @@ def __init__(
 
         self.transforms = transforms.Compose(
             [
-                transforms.Resize(resize_size, interpolation=interpolation),
+                transforms.ToImageTensor(),
+                transforms.Resize(resize_size, interpolation=interpolation, antialias=True),
                 transforms.CenterCrop(crop_size),
-                transforms.PILToTensor(),
                 transforms.ConvertImageDtype(torch.float),
                 transforms.Normalize(mean=mean, std=std),
             ]

references/classification/train.py

@@ -4,14 +4,17 @@
 import warnings
 
 import presets
+from sampler import RASampler
+from transforms import WrapIntoFeatures
+import utils  # usort: skip
+
 import torch
 import torch.utils.data
 import torchvision
-import transforms
-import utils
-from sampler import RASampler
+
 from torch import nn
 from torch.utils.data.dataloader import default_collate
+from torchvision.prototype import features, transforms
 from torchvision.transforms.functional import InterpolationMode
 
 
@@ -144,6 +147,7 @@ def load_data(traindir, valdir, args):
                 ra_magnitude=ra_magnitude,
                 augmix_severity=augmix_severity,
             ),
+            target_transform=lambda target: features.Label(target),
         )
         if args.cache_dataset:
             print(f"Saving dataset_train to {cache_path}")
@@ -168,7 +172,8 @@ def load_data(traindir, valdir, args):
 
         dataset_test = torchvision.datasets.ImageFolder(
             valdir,
-            preprocessing,
+            transform=preprocessing,
+            target_transform=lambda target: features.Label(target),
         )
         if args.cache_dataset:
             print(f"Saving dataset_test to {cache_path}")
@@ -210,16 +215,23 @@ def main(args):
 
     collate_fn = None
     num_classes = len(dataset.classes)
-    mixup_transforms = []
+    mixup_or_cutmix = []
     if args.mixup_alpha > 0.0:
-        mixup_transforms.append(transforms.RandomMixup(num_classes, p=1.0, alpha=args.mixup_alpha))
+        mixup_or_cutmix.append(transforms.RandomMixup(alpha=args.mixup_alpha, p=1.0))
     if args.cutmix_alpha > 0.0:
-        mixup_transforms.append(transforms.RandomCutmix(num_classes, p=1.0, alpha=args.cutmix_alpha))
-    if mixup_transforms:
-        mixupcutmix = torchvision.transforms.RandomChoice(mixup_transforms)
+        mixup_or_cutmix.append(transforms.RandomCutmix(alpha=args.cutmix_alpha, p=1.0))
+    if mixup_or_cutmix:
+        batch_transform = transforms.Compose(
+            [
+                WrapIntoFeatures(),
+                transforms.LabelToOneHot(num_categories=num_classes),
+                transforms.ToDtype({features.OneHotLabel: torch.float, features.Image: None}),
+                transforms.RandomChoice(mixup_or_cutmix),
+            ]
+        )
 
         def collate_fn(batch):
-            return mixupcutmix(*default_collate(batch))
+            return batch_transform(*default_collate(batch))
 
     data_loader = torch.utils.data.DataLoader(
         dataset,

references/classification/transforms.py

@@ -3,9 +3,20 @@
 
 import torch
 from torch import Tensor
+from torchvision.prototype import features
+from torchvision.prototype.transforms import functional as PF
 from torchvision.transforms import functional as F
 
 
+class WrapIntoFeatures(torch.nn.Module):
+    def forward(self, sample):
+        image, target = sample
+        return PF.to_image_tensor(image), features.Label(target)
+
+
+# Original Transforms can be removed:
+
+
 class RandomMixup(torch.nn.Module):
     """Randomly apply Mixup to the provided batch and targets.
     The class implements the data augmentations as described in the paper

references/detection/coco_utils.py

@@ -1,30 +1,13 @@
-import copy
 import os
 
 import torch
 import torch.utils.data
 import torchvision
-import transforms as T
+
 from pycocotools import mask as coco_mask
 from pycocotools.coco import COCO
-
-
-class FilterAndRemapCocoCategories:
-    def __init__(self, categories, remap=True):
-        self.categories = categories
-        self.remap = remap
-
-    def __call__(self, image, target):
-        anno = target["annotations"]
-        anno = [obj for obj in anno if obj["category_id"] in self.categories]
-        if not self.remap:
-            target["annotations"] = anno
-            return image, target
-        anno = copy.deepcopy(anno)
-        for obj in anno:
-            obj["category_id"] = self.categories.index(obj["category_id"])
-        target["annotations"] = anno
-        return image, target
+from torchvision.prototype import features, transforms as T
+from torchvision.prototype.transforms import functional as F
 
 
 def convert_coco_poly_to_mask(segmentations, height, width):
@@ -45,7 +28,8 @@ def convert_coco_poly_to_mask(segmentations, height, width):
 
 
 class ConvertCocoPolysToMask:
-    def __call__(self, image, target):
+    def __call__(self, sample):
+        image, target = sample
         w, h = image.size
 
         image_id = target["image_id"]
@@ -100,6 +84,27 @@ def __call__(self, image, target):
         return image, target
 
 
+class WrapIntoFeatures:
+    def __call__(self, sample):
+        image, target = sample
+
+        wrapped_target = dict(
+            boxes=features.BoundingBox(
+                target["boxes"],
+                format=features.BoundingBoxFormat.XYXY,
+                spatial_size=(image.height, image.width),
+            ),
+            # TODO: add categories
+            labels=features.Label(target["labels"], categories=None),
+            masks=features.Mask(target["masks"]),
+            image_id=int(target["image_id"]),
+            area=target["area"].tolist(),
+            iscrowd=target["iscrowd"].bool().tolist(),
+        )
+
+        return F.to_image_tensor(image), wrapped_target
+
+
 def _coco_remove_images_without_annotations(dataset, cat_list=None):
     def _has_only_empty_bbox(anno):
         return all(any(o <= 1 for o in obj["bbox"][2:]) for obj in anno)
@@ -225,10 +230,12 @@ def get_coco(root, image_set, transforms, mode="instances"):
     PATHS = {
         "train": ("train2017", os.path.join("annotations", anno_file_template.format(mode, "train"))),
         "val": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val"))),
-        # "train": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val")))
     }
 
-    t = [ConvertCocoPolysToMask()]
+    t = [
+        ConvertCocoPolysToMask(),
+        WrapIntoFeatures(),
+    ]
 
     if transforms is not None:
         t.append(transforms)
@@ -243,8 +250,6 @@ def get_coco(root, image_set, transforms, mode="instances"):
     if image_set == "train":
         dataset = _coco_remove_images_without_annotations(dataset)
 
-    # dataset = torch.utils.data.Subset(dataset, [i for i in range(500)])
-
     return dataset
 
 

references/detection/engine.py

@@ -26,7 +26,7 @@ def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq, sc
 
     for images, targets in metric_logger.log_every(data_loader, print_freq, header):
         images = list(image.to(device) for image in images)
-        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
+        targets = [{k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in t.items()} for t in targets]
         with torch.cuda.amp.autocast(enabled=scaler is not None):
             loss_dict = model(images, targets)
             losses = sum(loss for loss in loss_dict.values())
@@ -97,7 +97,7 @@ def evaluate(model, data_loader, device):
         outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
         model_time = time.time() - model_time
 
-        res = {target["image_id"].item(): output for target, output in zip(targets, outputs)}
+        res = {target["image_id"]: output for target, output in zip(targets, outputs)}
         evaluator_time = time.time()
         coco_evaluator.update(res)
         evaluator_time = time.time() - evaluator_time

references/detection/presets.py

@@ -1,73 +1,56 @@
+from collections import defaultdict
+
 import torch
-import transforms as T
+from torchvision.prototype import features, transforms as T
 
 
-class DetectionPresetTrain:
+class DetectionPresetTrain(T.Compose):
     def __init__(self, *, data_augmentation, hflip_prob=0.5, mean=(123.0, 117.0, 104.0)):
         if data_augmentation == "hflip":
-            self.transforms = T.Compose(
-                [
-                    T.RandomHorizontalFlip(p=hflip_prob),
-                    T.PILToTensor(),
-                    T.ConvertImageDtype(torch.float),
-                ]
-            )
+            transforms = [
+                T.RandomHorizontalFlip(p=hflip_prob),
+                T.ConvertImageDtype(torch.float),
+            ]
         elif data_augmentation == "lsj":
-            self.transforms = T.Compose(
-                [
-                    T.ScaleJitter(target_size=(1024, 1024)),
-                    T.FixedSizeCrop(size=(1024, 1024), fill=mean),
-                    T.RandomHorizontalFlip(p=hflip_prob),
-                    T.PILToTensor(),
-                    T.ConvertImageDtype(torch.float),
-                ]
-            )
+            transforms = [
+                T.ScaleJitter(target_size=(1024, 1024), antialias=True),
+                T.FixedSizeCrop(size=(1024, 1024), fill=defaultdict(lambda: mean, {features.Mask: 0})),
+                T.RandomHorizontalFlip(p=hflip_prob),
+                T.ConvertImageDtype(torch.float),
+            ]
         elif data_augmentation == "multiscale":
-            self.transforms = T.Compose(
-                [
-                    T.RandomShortestSize(
-                        min_size=(480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=1333
-                    ),
-                    T.RandomHorizontalFlip(p=hflip_prob),
-                    T.PILToTensor(),
-                    T.ConvertImageDtype(torch.float),
-                ]
-            )
+            transforms = [
+                T.RandomShortestSize(
+                    min_size=(480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=1333, antialias=True
+                ),
+                T.RandomHorizontalFlip(p=hflip_prob),
+                T.ConvertImageDtype(torch.float),
+            ]
         elif data_augmentation == "ssd":
-            self.transforms = T.Compose(
-                [
-                    T.RandomPhotometricDistort(),
-                    T.RandomZoomOut(fill=list(mean)),
-                    T.RandomIoUCrop(),
-                    T.RandomHorizontalFlip(p=hflip_prob),
-                    T.PILToTensor(),
-                    T.ConvertImageDtype(torch.float),
-                ]
-            )
+            transforms = [
+                T.RandomPhotometricDistort(),
+                T.RandomZoomOut(fill=defaultdict(lambda: mean, {features.Mask: 0})),
+                T.RandomIoUCrop(),
+                T.RandomHorizontalFlip(p=hflip_prob),
+                T.ConvertImageDtype(torch.float),
+            ]
         elif data_augmentation == "ssdlite":
-            self.transforms = T.Compose(
-                [
-                    T.RandomIoUCrop(),
-                    T.RandomHorizontalFlip(p=hflip_prob),
-                    T.PILToTensor(),
-                    T.ConvertImageDtype(torch.float),
-                ]
-            )
+            transforms = [
+                T.RandomIoUCrop(),
+                T.RandomHorizontalFlip(p=hflip_prob),
+                T.ConvertImageDtype(torch.float),
+            ]
         else:
             raise ValueError(f'Unknown data augmentation policy "{data_augmentation}"')
 
-    def __call__(self, img, target):
-        return self.transforms(img, target)
+        super().__init__(transforms)
 
 
-class DetectionPresetEval:
+class DetectionPresetEval(T.Compose):
     def __init__(self):
-        self.transforms = T.Compose(
+        super().__init__(
             [
-                T.PILToTensor(),
+                T.ToImageTensor(),
                 T.ConvertImageDtype(torch.float),
             ]
         )
-
-    def __call__(self, img, target):
-        return self.transforms(img, target)

references/detection/train.py

@@ -31,12 +31,12 @@
 from coco_utils import get_coco, get_coco_kp
 from engine import evaluate, train_one_epoch
 from group_by_aspect_ratio import create_aspect_ratio_groups, GroupedBatchSampler
+from torchvision.prototype import transforms as T
 from torchvision.transforms import InterpolationMode
-from transforms import SimpleCopyPaste
 
 
 def copypaste_collate_fn(batch):
-    copypaste = SimpleCopyPaste(blending=True, resize_interpolation=InterpolationMode.BILINEAR)
+    copypaste = T.SimpleCopyPaste(blending=True, resize_interpolation=InterpolationMode.BILINEAR, antialias=True)
     return copypaste(*utils.collate_fn(batch))
 
 

references/detection/transforms.py

@@ -1,3 +1,4 @@
+# Original Transforms can be removed:
 from typing import Dict, List, Optional, Tuple, Union
 
 import torch