Prototype references

references/classification/presets.py

@@ -1,6 +1,5 @@
 import torch
-from torchvision.transforms import autoaugment, transforms
-from torchvision.transforms.functional import InterpolationMode
+import torchvision.transforms.v2 as transforms
 
 
 class ClassificationPresetTrain:
@@ -10,29 +9,44 @@ def __init__(
         crop_size,
         mean=(0.485, 0.456, 0.406),
         std=(0.229, 0.224, 0.225),
-        interpolation=InterpolationMode.BILINEAR,
+        interpolation=transforms.InterpolationMode.BILINEAR,
         hflip_prob=0.5,
         auto_augment_policy=None,
         ra_magnitude=9,
         augmix_severity=3,
         random_erase_prob=0.0,
+        backend="pil",
     ):
-        trans = [transforms.RandomResizedCrop(crop_size, interpolation=interpolation)]
+        trans = []
+
+        backend = backend.lower()
+        if backend == "datapoint":
+            trans.append(transforms.ToImageTensor())
+        elif backend == "tensor":
+            trans.append(transforms.PILToTensor())
+        else:
+            assert backend == "pil"
+
+        trans.append(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))
+
+        if backend == "pil":
+            # Note: we could also just use pure tensors?
+            trans.append(transforms.ToImageTensor())
+
         trans.extend(
             [
-                transforms.PILToTensor(),
                 transforms.ConvertImageDtype(torch.float),
                 transforms.Normalize(mean=mean, std=std),
             ]
@@ -54,18 +68,33 @@ def __init__(
         resize_size=256,
         mean=(0.485, 0.456, 0.406),
         std=(0.229, 0.224, 0.225),
-        interpolation=InterpolationMode.BILINEAR,
+        interpolation=transforms.InterpolationMode.BILINEAR,
+        backend="pil",
     ):
+        trans = []
 
-        self.transforms = transforms.Compose(
-            [
-                transforms.Resize(resize_size, interpolation=interpolation),
-                transforms.CenterCrop(crop_size),
-                transforms.PILToTensor(),
-                transforms.ConvertImageDtype(torch.float),
-                transforms.Normalize(mean=mean, std=std),
-            ]
-        )
+        backend = backend.lower()
+        if backend == "datapoint":
+            trans.append(transforms.ToImageTensor())
+        elif backend == "tensor":
+            trans.append(transforms.PILToTensor())
+        else:
+            assert backend == "pil"
+
+        trans += [
+            transforms.Resize(resize_size, interpolation=interpolation, antialias=True),
+            transforms.CenterCrop(crop_size),
+        ]
+
+        if backend == "pil":
+            trans.append(transforms.ToImageTensor())
+
+        trans += [
+            transforms.ConvertImageDtype(torch.float),
+            transforms.Normalize(mean=mean, std=std),
+        ]
+
+        self.transforms = transforms.Compose(trans)
 
     def __call__(self, img):
         return self.transforms(img)

references/classification/train.py

@@ -4,15 +4,17 @@
 import warnings
 
 import presets
+from sampler import RASampler
+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.transforms.functional import InterpolationMode
+from torchvision.transforms.v2 import InterpolationMode
 
 
 def train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, args, model_ema=None, scaler=None):
@@ -143,6 +145,7 @@ def load_data(traindir, valdir, args):
                 random_erase_prob=random_erase_prob,
                 ra_magnitude=ra_magnitude,
                 augmix_severity=augmix_severity,
+                backend=args.backend,
             ),
         )
         if args.cache_dataset:
@@ -163,12 +166,15 @@ def load_data(traindir, valdir, args):
             preprocessing = weights.transforms()
         else:
             preprocessing = presets.ClassificationPresetEval(
-                crop_size=val_crop_size, resize_size=val_resize_size, interpolation=interpolation
+                crop_size=val_crop_size,
+                resize_size=val_resize_size,
+                interpolation=interpolation,
+                backend=args.backend,
             )
 
         dataset_test = torchvision.datasets.ImageFolder(
             valdir,
-            preprocessing,
+            transform=preprocessing,
         )
         if args.cache_dataset:
             print(f"Saving dataset_test to {cache_path}")
@@ -507,6 +513,7 @@ def get_args_parser(add_help=True):
         "--ra-reps", default=3, type=int, help="number of repetitions for Repeated Augmentation (default: 3)"
     )
     parser.add_argument("--weights", default=None, type=str, help="the weights enum name to load")
+    parser.add_argument("--backend", default="PIL", type=str, help="PIL, tensor or datapoint - case insensitive")
     return parser
 
 

references/detection/coco_utils.py

@@ -1,30 +1,12 @@
-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.datasets import wrap_dataset_for_transforms_v2
 
 
 def convert_coco_poly_to_mask(segmentations, height, width):
@@ -44,8 +26,10 @@ def convert_coco_poly_to_mask(segmentations, height, width):
     return masks
 
 
+# TODO: Is this still needed?
 class ConvertCocoPolysToMask:
-    def __call__(self, image, target):
+    def __call__(self, sample):
+        image, target = sample
         w, h = image.size
 
         image_id = target["image_id"]
@@ -126,10 +110,10 @@ def _has_valid_annotation(anno):
             return True
         return False
 
-    if not isinstance(dataset, torchvision.datasets.CocoDetection):
-        raise TypeError(
-            f"This function expects dataset of type torchvision.datasets.CocoDetection, instead  got {type(dataset)}"
-        )
+    # if not isinstance(dataset, torchvision.datasets.CocoDetection):
+    #     raise TypeError(
+    #         f"This function expects dataset of type torchvision.datasets.CocoDetection, instead  got {type(dataset)}"
+    #     )
     ids = []
     for ds_idx, img_id in enumerate(dataset.ids):
         ann_ids = dataset.coco.getAnnIds(imgIds=img_id, iscrowd=None)
@@ -201,50 +185,35 @@ def get_coco_api_from_dataset(dataset):
             break
         if isinstance(dataset, torch.utils.data.Subset):
             dataset = dataset.dataset
-    if isinstance(dataset, torchvision.datasets.CocoDetection):
+    # TODO: hmmmmm
+    if isinstance(dataset, torchvision.datasets.CocoDetection) or isinstance(
+        getattr(dataset, "_dataset", None), torchvision.datasets.CocoDetection
+    ):
         return dataset.coco
     return convert_to_coco_api(dataset)
 
 
-class CocoDetection(torchvision.datasets.CocoDetection):
-    def __init__(self, img_folder, ann_file, transforms):
-        super().__init__(img_folder, ann_file)
-        self._transforms = transforms
-
-    def __getitem__(self, idx):
-        img, target = super().__getitem__(idx)
-        image_id = self.ids[idx]
-        target = dict(image_id=image_id, annotations=target)
-        if self._transforms is not None:
-            img, target = self._transforms(img, target)
-        return img, target
+# TODO: Maybe not critical but the wrapper doesn't work on sub-classes
 
 
 def get_coco(root, image_set, transforms, mode="instances"):
     anno_file_template = "{}_{}2017.json"
     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()]
-
-    if transforms is not None:
-        t.append(transforms)
-    transforms = T.Compose(t)
-
     img_folder, ann_file = PATHS[image_set]
     img_folder = os.path.join(root, img_folder)
     ann_file = os.path.join(root, ann_file)
 
-    dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
+    dataset = torchvision.datasets.CocoDetection(img_folder, ann_file, transforms=transforms)
+    # TODO: need to update target_keys to handle masks for segmentation!
+    dataset = wrap_dataset_for_transforms_v2(dataset, target_keys={"boxes", "labels", "image_id"})
 
     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,9 @@ 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]
+        for t in targets:
+            assert t["labels"].shape[0] == t["boxes"].shape[0], f"{t['labels'].shape} {t['boxes'].shape}"
+        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 +99,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/group_by_aspect_ratio.py

@@ -164,7 +164,10 @@ def compute_aspect_ratios(dataset, indices=None):
     if hasattr(dataset, "get_height_and_width"):
         return _compute_aspect_ratios_custom_dataset(dataset, indices)
 
-    if isinstance(dataset, torchvision.datasets.CocoDetection):
+    # TODO: hmmmmm
+    if isinstance(dataset, torchvision.datasets.CocoDetection) or isinstance(
+        getattr(dataset, "_dataset", None), torchvision.datasets.CocoDetection
+    ):
         return _compute_aspect_ratios_coco_dataset(dataset, indices)
 
     if isinstance(dataset, torchvision.datasets.VOCDetection):