update

fastreid/data/datasets/pes3d.py

+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['PeS3D',]
+
+
+@DATASET_REGISTRY.register()
+class PeS3D(ImageDataset):
+    """3Dpes
+    """
+    dataset_dir = "3DPeS"
+    dataset_name = "pes3d"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+
+        pid_list = os.listdir(train_path)
+        for pid_dir in pid_list:
+            pid = self.dataset_name + "_" + pid_dir
+            img_list = glob(os.path.join(train_path, pid_dir,  "*.bmp"))
+            for img_path in img_list:
+                camid = self.dataset_name + "_cam0"
+                data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/pku.py

+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['PKU', ]
+
+
+@DATASET_REGISTRY.register()
+class PKU(ImageDataset):
+    """PKU
+    """
+    dataset_dir = "PKUv1a_128x48"
+    dataset_name = 'pku'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        img_paths = glob(os.path.join(train_path, "*.png"))
+
+        for img_path in img_paths:
+            split_path = img_path.split('/')
+            img_info = split_path[-1].split('_')
+            pid = self.dataset_name + "_" + img_info[0]
+            camid = self.dataset_name + "_" + img_info[1]
+            data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/prai.py

+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['PRAI', ]
+
+
+@DATASET_REGISTRY.register()
+class PRAI(ImageDataset):
+    """PRAI
+    """
+    dataset_dir = "PRAI-1581"
+    dataset_name = 'prai'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir, 'images')
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        img_paths = glob(os.path.join(train_path, "*.jpg"))
+        for img_path in img_paths:
+            split_path = img_path.split('/')
+            img_info = split_path[-1].split('_')
+            pid = self.dataset_name + "_" + img_info[0]
+            camid = self.dataset_name + "_" + img_info[1]
+            data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/prid.py

+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['PRID', ]
+
+
+@DATASET_REGISTRY.register()
+class PRID(ImageDataset):
+    """PRID
+    """
+    dataset_dir = "prid_2011"
+    dataset_name = 'prid'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir, 'slim_train')
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        for root, dirs, files in os.walk(train_path):
+            for img_name in filter(lambda x: x.endswith('.png'), files):
+                img_path = os.path.join(root, img_name)
+                pid = self.dataset_name + '_' + root.split('/')[-1].split('_')[1]
+                camid = self.dataset_name + '_' + img_name.split('_')[0]
+                data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/saivt.py

+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['SAIVT', ]
+
+
+@DATASET_REGISTRY.register()
+class SAIVT(ImageDataset):
+    """SAIVT
+    """
+    dataset_dir = "SAIVT-SoftBio"
+    dataset_name = "saivt"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+
+        pid_path = os.path.join(train_path, "cropped_images")
+        pid_list = os.listdir(pid_path)
+
+        for pid_name in pid_list:
+            pid = self.dataset_name + '_' + pid_name
+            img_list = glob(os.path.join(pid_path, pid_name, "*.jpeg"))
+            for img_path in img_list:
+                img_name = os.path.basename(img_path)
+                camid = self.dataset_name + '_' + img_name.split('-')[2]
+                data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/sensereid.py

+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['SenseReID', ]
+
+
+@DATASET_REGISTRY.register()
+class SenseReID(ImageDataset):
+    """Sense reid
+    """
+    dataset_dir = "SenseReID"
+    dataset_name = "senseid"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        file_path_list = ['test_gallery', 'test_prob']
+
+        for file_path in file_path_list:
+            sub_file = os.path.join(train_path, file_path)
+            img_name = glob(os.path.join(sub_file, "*.jpg"))
+            for img_path in img_name:
+                img_name = img_path.split('/')[-1]
+                img_info = img_name.split('_')
+                pid = self.dataset_name + "_" + img_info[0]
+                camid = self.dataset_name + "_" + img_info[1].split('.')[0]
+                data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/shinpuhkan.py

+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['Shinpuhkan', ]
+
+
+@DATASET_REGISTRY.register()
+class Shinpuhkan(ImageDataset):
+    """shinpuhkan
+    """
+    dataset_dir = "shinpuhkan"
+    dataset_name = 'shinpuhkan'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+
+        for root, dirs, files in os.walk(train_path):
+            img_names = list(filter(lambda x: x.endswith(".jpg"), files))
+            # fmt: off
+            if len(img_names) == 0: continue
+            # fmt: on
+            for img_name in img_names:
+                img_path = os.path.join(root, img_name)
+                split_path = img_name.split('_')
+                pid = self.dataset_name + "_" + split_path[0]
+                camid = self.dataset_name + "_" + split_path[2]
+                data.append((img_path, pid, camid))
+
+        return data

fastreid/data/datasets/sysu_mm.py

+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['SYSU_mm', ]
+
+
+@DATASET_REGISTRY.register()
+class SYSU_mm(ImageDataset):
+    """sysu mm
+    """
+    dataset_dir = "SYSU-MM01"
+    dataset_name = "sysumm01"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+
+        file_path_list = ['cam1', 'cam2', 'cam4', 'cam5']
+
+        for file_path in file_path_list:
+            camid = self.dataset_name + "_" + file_path
+            pid_list = os.listdir(os.path.join(train_path, file_path))
+            for pid_dir in pid_list:
+                pid = self.dataset_name + "_" + pid_dir
+                img_list = glob(os.path.join(train_path, file_path, pid_dir, "*.jpg"))
+                for img_path in img_list:
+                    data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/thermalworld.py

+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['Thermalworld', ]
+
+
+@DATASET_REGISTRY.register()
+class Thermalworld(ImageDataset):
+    """thermal world
+    """
+    dataset_dir = "thermalworld_rgb"
+    dataset_name = "thermalworld"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        pid_list = os.listdir(train_path)
+        for pid_dir in pid_list:
+            pid = self.dataset_name + "_" + pid_dir
+            img_list = glob(os.path.join(train_path, pid_dir, "*.jpg"))
+            for img_path in img_list:
+                camid = self.dataset_name + "_cam0"
+                data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/vehicleid.py

+# encoding: utf-8
+"""
+@author:  Jinkai Zheng
+@contact: 1315673509@qq.com
+"""
+
+import os.path as osp
+import random
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class VehicleID(ImageDataset):
+    """VehicleID.
+
+    Reference:
+        Liu et al. Deep relative distance learning: Tell the difference between similar vehicles. CVPR 2016.
+
+    URL: `<https://pkuml.org/resources/pku-vehicleid.html>`_
+
+    Train dataset statistics:
+        - identities: 13164.
+        - images: 113346.
+    """
+    dataset_dir = "vehicleid"
+    dataset_name = "vehicleid"
+
+    def __init__(self, root='datasets', test_list='', **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+
+        self.image_dir = osp.join(self.dataset_dir, 'image')
+        self.train_list = osp.join(self.dataset_dir, 'train_test_split/train_list.txt')
+        if test_list:
+            self.test_list = test_list
+        else:
+            self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_13164.txt')
+
+        required_files = [
+            self.dataset_dir,
+            self.image_dir,
+            self.train_list,
+            self.test_list,
+        ]
+        self.check_before_run(required_files)
+
+        train = self.process_dir(self.train_list, is_train=True)
+        query, gallery = self.process_dir(self.test_list, is_train=False)
+
+        super(VehicleID, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, list_file, is_train=True):
+        img_list_lines = open(list_file, 'r').readlines()
+
+        dataset = []
+        for idx, line in enumerate(img_list_lines):
+            line = line.strip()
+            vid = int(line.split(' ')[1])
+            imgid = line.split(' ')[0]
+            img_path = osp.join(self.image_dir, f"{imgid}.jpg")
+            imgid = int(imgid)
+            if is_train:
+                vid = f"{self.dataset_name}_{vid}"
+                imgid = f"{self.dataset_name}_{imgid}"
+            dataset.append((img_path, vid, imgid))
+
+        if is_train: return dataset
+        else:
+            random.shuffle(dataset)
+            vid_container = set()
+            query = []
+            gallery = []
+            for sample in dataset:
+                if sample[1] not in vid_container:
+                    vid_container.add(sample[1])
+                    gallery.append(sample)
+                else:
+                    query.append(sample)
+
+            return query, gallery
+
+
+@DATASET_REGISTRY.register()
+class SmallVehicleID(VehicleID):
+    """VehicleID.
+    Small test dataset statistics:
+        - identities: 800.
+        - images: 6493.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.test_list = osp.join(dataset_dir, 'train_test_split/test_list_800.txt')
+
+        super(SmallVehicleID, self).__init__(root, self.test_list, **kwargs)
+
+
+@DATASET_REGISTRY.register()
+class MediumVehicleID(VehicleID):
+    """VehicleID.
+    Medium test dataset statistics:
+        - identities: 1600.
+        - images: 13377.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.test_list = osp.join(dataset_dir, 'train_test_split/test_list_1600.txt')
+
+        super(MediumVehicleID, self).__init__(root, self.test_list, **kwargs)
+
+
+@DATASET_REGISTRY.register()
+class LargeVehicleID(VehicleID):
+    """VehicleID.
+    Large test dataset statistics:
+        - identities: 2400.
+        - images: 19777.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.test_list = osp.join(dataset_dir, 'train_test_split/test_list_2400.txt')
+
+        super(LargeVehicleID, self).__init__(root, self.test_list, **kwargs)

fastreid/data/datasets/veri.py

+# encoding: utf-8
+"""
+@author:  Jinkai Zheng
+@contact: 1315673509@qq.com
+"""
+
+import glob
+import os.path as osp
+import re
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class VeRi(ImageDataset):
+    """VeRi.
+
+    Reference:
+        Xinchen Liu et al. A Deep Learning based Approach for Progressive Vehicle Re-Identification. ECCV 2016.
+        Xinchen Liu et al. PROVID: Progressive and Multimodal Vehicle Reidentification for Large-Scale Urban Surveillance. IEEE TMM 2018.
+
+    URL: `<https://vehiclereid.github.io/VeRi/>`_
+
+    Dataset statistics:
+        - identities: 775.
+        - images: 37778 (train) + 1678 (query) + 11579 (gallery).
+    """
+    dataset_dir = "veri"
+    dataset_name = "veri"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+
+        self.train_dir = osp.join(self.dataset_dir, 'image_train')
+        self.query_dir = osp.join(self.dataset_dir, 'image_query')
+        self.gallery_dir = osp.join(self.dataset_dir, 'image_test')
+
+        required_files = [
+            self.dataset_dir,
+            self.train_dir,
+            self.query_dir,
+            self.gallery_dir,
+        ]
+        self.check_before_run(required_files)
+
+        train = self.process_dir(self.train_dir)
+        query = self.process_dir(self.query_dir, is_train=False)
+        gallery = self.process_dir(self.gallery_dir, is_train=False)
+
+        super(VeRi, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, dir_path, is_train=True):
+        img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
+        pattern = re.compile(r'([\d]+)_c(\d\d\d)')
+
+        data = []
+        for img_path in img_paths:
+            pid, camid = map(int, pattern.search(img_path).groups())
+            if pid == -1: continue  # junk images are just ignored
+            assert 0 <= pid <= 776
+            assert 1 <= camid <= 20
+            camid -= 1  # index starts from 0
+            if is_train:
+                pid = self.dataset_name + "_" + str(pid)
+                camid = self.dataset_name + "_" + str(camid)
+            data.append((img_path, pid, camid))
+
+        return data

fastreid/data/datasets/veriwild.py

+# encoding: utf-8
+"""
+@author:  Jinkai Zheng
+@contact: 1315673509@qq.com
+"""
+
+import os.path as osp
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class VeRiWild(ImageDataset):
+    """VeRi-Wild.
+
+    Reference:
+        Lou et al. A Large-Scale Dataset for Vehicle Re-Identification in the Wild. CVPR 2019.
+
+    URL: `<https://github.com/PKU-IMRE/VERI-Wild>`_
+
+    Train dataset statistics:
+        - identities: 30671.
+        - images: 277797.
+    """
+    dataset_dir = "VERI-Wild"
+    dataset_name = "veriwild"
+
+    def __init__(self, root='datasets', query_list='', gallery_list='', **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+
+        self.image_dir = osp.join(self.dataset_dir, 'images')
+        self.train_list = osp.join(self.dataset_dir, 'train_test_split/train_list.txt')
+        self.vehicle_info = osp.join(self.dataset_dir, 'train_test_split/vehicle_info.txt')
+        if query_list and gallery_list:
+            self.query_list = query_list
+            self.gallery_list = gallery_list
+        else:
+            self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_10000_query.txt')
+            self.gallery_list = osp.join(self.dataset_dir, 'train_test_split/test_10000.txt')
+
+        required_files = [
+            self.image_dir,
+            self.train_list,
+            self.query_list,
+            self.gallery_list,
+            self.vehicle_info,
+        ]
+        self.check_before_run(required_files)
+
+        self.imgid2vid, self.imgid2camid, self.imgid2imgpath = self.process_vehicle(self.vehicle_info)
+
+        train = self.process_dir(self.train_list)
+        query = self.process_dir(self.query_list, is_train=False)
+        gallery = self.process_dir(self.gallery_list, is_train=False)
+
+        super(VeRiWild, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, img_list, is_train=True):
+        img_list_lines = open(img_list, 'r').readlines()
+
+        dataset = []
+        for idx, line in enumerate(img_list_lines):
+            line = line.strip()
+            vid = int(line.split('/')[0])
+            imgid = line.split('/')[1].split('.')[0]
+            camid = int(self.imgid2camid[imgid])
+            if is_train:
+                vid = f"{self.dataset_name}_{vid}"
+                camid = f"{self.dataset_name}_{camid}"
+            dataset.append((self.imgid2imgpath[imgid], vid, camid))
+
+        assert len(dataset) == len(img_list_lines)
+        return dataset
+
+    def process_vehicle(self, vehicle_info):
+        imgid2vid = {}
+        imgid2camid = {}
+        imgid2imgpath = {}
+        vehicle_info_lines = open(vehicle_info, 'r').readlines()
+
+        for idx, line in enumerate(vehicle_info_lines[1:]):
+            vid = line.strip().split('/')[0]
+            imgid = line.strip().split(';')[0].split('/')[1]
+            camid = line.strip().split(';')[1]
+            img_path = osp.join(self.image_dir, vid, imgid + '.jpg')
+            imgid2vid[imgid] = vid
+            imgid2camid[imgid] = camid
+            imgid2imgpath[imgid] = img_path
+
+        assert len(imgid2vid) == len(vehicle_info_lines) - 1
+        return imgid2vid, imgid2camid, imgid2imgpath
+
+
+@DATASET_REGISTRY.register()
+class SmallVeRiWild(VeRiWild):
+    """VeRi-Wild.
+    Small test dataset statistics:
+        - identities: 3000.
+        - images: 41861.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.query_list = osp.join(dataset_dir, 'train_test_split/test_3000_query.txt')
+        self.gallery_list = osp.join(dataset_dir, 'train_test_split/test_3000.txt')
+
+        super(SmallVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)
+
+
+@DATASET_REGISTRY.register()
+class MediumVeRiWild(VeRiWild):
+    """VeRi-Wild.
+    Medium test dataset statistics:
+        - identities: 5000.
+        - images: 69389.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.query_list = osp.join(dataset_dir, 'train_test_split/test_5000_query.txt')
+        self.gallery_list = osp.join(dataset_dir, 'train_test_split/test_5000.txt')
+
+        super(MediumVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)
+
+
+@DATASET_REGISTRY.register()
+class LargeVeRiWild(VeRiWild):
+    """VeRi-Wild.
+    Large test dataset statistics:
+        - identities: 10000.
+        - images: 138517.
+    """
+
+    def __init__(self, root='datasets', **kwargs):
+        dataset_dir = osp.join(root, self.dataset_dir)
+        self.query_list = osp.join(dataset_dir, 'train_test_split/test_10000_query.txt')
+        self.gallery_list = osp.join(dataset_dir, 'train_test_split/test_10000.txt')
+
+        super(LargeVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)

fastreid/data/datasets/viper.py

+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['VIPeR', ]
+
+
+@DATASET_REGISTRY.register()
+class VIPeR(ImageDataset):
+    dataset_dir = "VIPeR"
+    dataset_name = "viper"
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir)
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+
+        file_path_list = ['cam_a', 'cam_b']
+
+        for file_path in file_path_list:
+            camid = self.dataset_name + "_" + file_path
+            img_list = glob(os.path.join(train_path, file_path, "*.bmp"))
+            for img_path in img_list:
+                img_name = img_path.split('/')[-1]
+                pid = self.dataset_name + "_" + img_name.split('_')[0]
+                data.append([img_path, pid, camid])
+
+        return data

fastreid/data/datasets/wildtracker.py

+# encoding: utf-8
+"""
+@author:  wangguanan
+@contact: guan.wang0706@gmail.com
+"""
+
+import glob
+import os
+
+from .bases import ImageDataset
+from ..datasets import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class WildTrackCrop(ImageDataset):
+    """WildTrack.
+    Reference:
+        WILDTRACK: A Multi-camera HD Dataset for Dense Unscripted Pedestrian Detection
+            T. Chavdarova; P. Baqué; A. Maksai; S. Bouquet; C. Jose et al.
+    URL: `<https://www.epfl.ch/labs/cvlab/data/data-wildtrack/>`_
+    Dataset statistics:
+        - identities: 313
+        - images: 33979 (train only)
+        - cameras: 7
+    Args:
+        data_path(str): path to WildTrackCrop dataset
+        combineall(bool): combine train and test sets as train set if True
+    """
+    dataset_url = None
+    dataset_dir = 'Wildtrack_crop_dataset'
+    dataset_name = 'wildtrack'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.dataset_dir = os.path.join(self.root, self.dataset_dir)
+
+        self.train_dir = os.path.join(self.dataset_dir, "crop")
+
+        train = self.process_dir(self.train_dir)
+        query = []
+        gallery = []
+
+        super(WildTrackCrop, self).__init__(train, query, gallery, **kwargs)
+
+    def process_dir(self, dir_path):
+        r"""
+        :param dir_path: directory path saving images
+        Returns
+            data(list) = [img_path, pid, camid]
+        """
+        data = []
+        for dir_name in os.listdir(dir_path):
+            img_lists = glob.glob(os.path.join(dir_path, dir_name, "*.png"))
+            for img_path in img_lists:
+                pid = self.dataset_name + "_" + dir_name
+                camid = img_path.split('/')[-1].split('_')[0]
+                camid = self.dataset_name + "_" + camid
+                data.append([img_path, pid, camid])
+        return data

fastreid/data/samplers/__init__.py

+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+from .triplet_sampler import BalancedIdentitySampler, NaiveIdentitySampler, SetReWeightSampler
+from .data_sampler import TrainingSampler, InferenceSampler
+from .imbalance_sampler import ImbalancedDatasetSampler
+
+__all__ = [
+    "BalancedIdentitySampler",
+    "NaiveIdentitySampler",
+    "SetReWeightSampler",
+    "TrainingSampler",
+    "InferenceSampler",
+    "ImbalancedDatasetSampler",
+]

fastreid/data/samplers/data_sampler.py

+# encoding: utf-8
+"""
+@author:  l1aoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+import itertools
+from typing import Optional
+
+import numpy as np
+from torch.utils.data import Sampler
+
+from fastreid.utils import comm
+
+
+class TrainingSampler(Sampler):
+    """
+    In training, we only care about the "infinite stream" of training data.
+    So this sampler produces an infinite stream of indices and
+    all workers cooperate to correctly shuffle the indices and sample different indices.
+    The samplers in each worker effectively produces `indices[worker_id::num_workers]`
+    where `indices` is an infinite stream of indices consisting of
+    `shuffle(range(size)) + shuffle(range(size)) + ...` (if shuffle is True)
+    or `range(size) + range(size) + ...` (if shuffle is False)
+    """
+
+    def __init__(self, size: int, shuffle: bool = True, seed: Optional[int] = None):
+        """
+        Args:
+            size (int): the total number of data of the underlying dataset to sample from
+            shuffle (bool): whether to shuffle the indices or not
+            seed (int): the initial seed of the shuffle. Must be the same
+                across all workers. If None, will use a random seed shared
+                among workers (require synchronization among all workers).
+        """
+        self._size = size
+        assert size > 0
+        self._shuffle = shuffle
+        if seed is None:
+            seed = comm.shared_random_seed()
+        self._seed = int(seed)
+
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(self._infinite_indices(), start, None, self._world_size)
+
+    def _infinite_indices(self):
+        np.random.seed(self._seed)
+        while True:
+            if self._shuffle:
+                yield from np.random.permutation(self._size)
+            else:
+                yield from np.arange(self._size)
+
+
+class InferenceSampler(Sampler):
+    """
+    Produce indices for inference.
+    Inference needs to run on the __exact__ set of samples,
+    therefore when the total number of samples is not divisible by the number of workers,
+    this sampler produces different number of samples on different workers.
+    """
+
+    def __init__(self, size: int):
+        """
+        Args:
+            size (int): the total number of data of the underlying dataset to sample from
+        """
+        self._size = size
+        assert size > 0
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+
+        shard_size = (self._size - 1) // self._world_size + 1
+        begin = shard_size * self._rank
+        end = min(shard_size * (self._rank + 1), self._size)
+        self._local_indices = range(begin, end)
+
+    def __iter__(self):
+        yield from self._local_indices
+
+    def __len__(self):
+        return len(self._local_indices)

fastreid/data/samplers/imbalance_sampler.py

+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+# based on:
+# https://github.com/ufoym/imbalanced-dataset-sampler/blob/master/torchsampler/imbalanced.py
+
+
+import itertools
+from typing import Optional, List, Callable
+
+import numpy as np
+import torch
+from torch.utils.data.sampler import Sampler
+
+from fastreid.utils import comm
+
+
+class ImbalancedDatasetSampler(Sampler):
+    """Samples elements randomly from a given list of indices for imbalanced dataset
+    Arguments:
+        data_source: a list of data items
+        size: number of samples to draw
+    """
+
+    def __init__(self, data_source: List, size: int = None, seed: Optional[int] = None,
+                 callback_get_label: Callable = None):
+        self.data_source = data_source
+        # consider all elements in the dataset
+        self.indices = list(range(len(data_source)))
+        # if num_samples is not provided, draw `len(indices)` samples in each iteration
+        self._size = len(self.indices) if size is None else size
+        self.callback_get_label = callback_get_label
+
+        # distribution of classes in the dataset
+        label_to_count = {}
+        for idx in self.indices:
+            label = self._get_label(data_source, idx)
+            label_to_count[label] = label_to_count.get(label, 0) + 1
+
+        # weight for each sample
+        weights = [1.0 / label_to_count[self._get_label(data_source, idx)] for idx in self.indices]
+        self.weights = torch.DoubleTensor(weights)
+
+        if seed is None:
+            seed = comm.shared_random_seed()
+        self._seed = int(seed)
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+
+    def _get_label(self, dataset, idx):
+        if self.callback_get_label:
+            return self.callback_get_label(dataset, idx)
+        else:
+            return dataset[idx][1]
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(self._infinite_indices(), start, None, self._world_size)
+
+    def _infinite_indices(self):
+        np.random.seed(self._seed)
+        while True:
+            for i in torch.multinomial(self.weights, self._size, replacement=True):
+                yield self.indices[i]

fastreid/data/samplers/triplet_sampler.py

+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: liaoxingyu2@jd.com
+"""
+
+import copy
+import itertools
+from collections import defaultdict
+from typing import Optional, List
+
+import numpy as np
+from torch.utils.data.sampler import Sampler
+
+from fastreid.utils import comm
+
+
+def no_index(a, b):
+    assert isinstance(a, list)
+    return [i for i, j in enumerate(a) if j != b]
+
+
+def reorder_index(batch_indices, world_size):
+    r"""Reorder indices of samples to align with DataParallel training.
+    In this order, each process will contain all images for one ID, triplet loss
+    can be computed within each process, and BatchNorm will get a stable result.
+    Args:
+        batch_indices: A batched indices generated by sampler
+        world_size: number of process
+    Returns:
+
+    """
+    mini_batchsize = len(batch_indices) // world_size
+    reorder_indices = []
+    for i in range(0, mini_batchsize):
+        for j in range(0, world_size):
+            reorder_indices.append(batch_indices[i + j * mini_batchsize])
+    return reorder_indices
+
+
+class BalancedIdentitySampler(Sampler):
+    def __init__(self, data_source: List, mini_batch_size: int, num_instances: int, seed: Optional[int] = None):
+        self.data_source = data_source
+        self.num_instances = num_instances
+        self.num_pids_per_batch = mini_batch_size // self.num_instances
+
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+        self.batch_size = mini_batch_size * self._world_size
+
+        self.index_pid = dict()
+        self.pid_cam = defaultdict(list)
+        self.pid_index = defaultdict(list)
+
+        for index, info in enumerate(data_source):
+            pid = info[1]
+            camid = info[2]
+            self.index_pid[index] = pid
+            self.pid_cam[pid].append(camid)
+            self.pid_index[pid].append(index)
+
+        self.pids = sorted(list(self.pid_index.keys()))
+        self.num_identities = len(self.pids)
+
+        if seed is None:
+            seed = comm.shared_random_seed()
+        self._seed = int(seed)
+
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(self._infinite_indices(), start, None, self._world_size)
+
+    def _infinite_indices(self):
+        np.random.seed(self._seed)
+        while True:
+            # Shuffle identity list
+            identities = np.random.permutation(self.num_identities)
+
+            # If remaining identities cannot be enough for a batch,
+            # just drop the remaining parts
+            drop_indices = self.num_identities % (self.num_pids_per_batch * self._world_size)
+            if drop_indices: identities = identities[:-drop_indices]
+
+            batch_indices = []
+            for kid in identities:
+                i = np.random.choice(self.pid_index[self.pids[kid]])
+                _, i_pid, i_cam = self.data_source[i]
+                batch_indices.append(i)
+                pid_i = self.index_pid[i]
+                cams = self.pid_cam[pid_i]
+                index = self.pid_index[pid_i]
+                select_cams = no_index(cams, i_cam)
+
+                if select_cams:
+                    if len(select_cams) >= self.num_instances:
+                        cam_indexes = np.random.choice(select_cams, size=self.num_instances - 1, replace=False)
+                    else:
+                        cam_indexes = np.random.choice(select_cams, size=self.num_instances - 1, replace=True)
+                    for kk in cam_indexes:
+                        batch_indices.append(index[kk])
+                else:
+                    select_indexes = no_index(index, i)
+                    if not select_indexes:
+                        # Only one image for this identity
+                        ind_indexes = [0] * (self.num_instances - 1)
+                    elif len(select_indexes) >= self.num_instances:
+                        ind_indexes = np.random.choice(select_indexes, size=self.num_instances - 1, replace=False)
+                    else:
+                        ind_indexes = np.random.choice(select_indexes, size=self.num_instances - 1, replace=True)
+
+                    for kk in ind_indexes:
+                        batch_indices.append(index[kk])
+
+                if len(batch_indices) == self.batch_size:
+                    yield from reorder_index(batch_indices, self._world_size)
+                    batch_indices = []
+
+
+class SetReWeightSampler(Sampler):
+    def __init__(self, data_source: str, mini_batch_size: int, num_instances: int, set_weight: list,
+                 seed: Optional[int] = None):
+        self.data_source = data_source
+        self.num_instances = num_instances
+        self.num_pids_per_batch = mini_batch_size // self.num_instances
+
+        self.set_weight = set_weight
+
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+        self.batch_size = mini_batch_size * self._world_size
+
+        assert self.batch_size % (sum(self.set_weight) * self.num_instances) == 0 and \
+               self.batch_size > sum(
+            self.set_weight) * self.num_instances, "Batch size must be divisible by the sum set weight"
+
+        self.index_pid = dict()
+        self.pid_cam = defaultdict(list)
+        self.pid_index = defaultdict(list)
+
+        self.cam_pid = defaultdict(list)
+
+        for index, info in enumerate(data_source):
+            pid = info[1]
+            camid = info[2]
+            self.index_pid[index] = pid
+            self.pid_cam[pid].append(camid)
+            self.pid_index[pid].append(index)
+            self.cam_pid[camid].append(pid)
+
+        # Get sampler prob for each cam
+        self.set_pid_prob = defaultdict(list)
+        for camid, pid_list in self.cam_pid.items():
+            index_per_pid = []
+            for pid in pid_list:
+                index_per_pid.append(len(self.pid_index[pid]))
+            cam_image_number = sum(index_per_pid)
+            prob = [i / cam_image_number for i in index_per_pid]
+            self.set_pid_prob[camid] = prob
+
+        self.pids = sorted(list(self.pid_index.keys()))
+        self.num_identities = len(self.pids)
+
+        if seed is None:
+            seed = comm.shared_random_seed()
+        self._seed = int(seed)
+
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(self._infinite_indices(), start, None, self._world_size)
+
+    def _infinite_indices(self):
+        np.random.seed(self._seed)
+        while True:
+            batch_indices = []
+            for camid in range(len(self.cam_pid.keys())):
+                select_pids = np.random.choice(self.cam_pid[camid], size=self.set_weight[camid], replace=False,
+                                               p=self.set_pid_prob[camid])
+                for pid in select_pids:
+                    index_list = self.pid_index[pid]
+                    if len(index_list) > self.num_instances:
+                        select_indexs = np.random.choice(index_list, size=self.num_instances, replace=False)
+                    else:
+                        select_indexs = np.random.choice(index_list, size=self.num_instances, replace=True)
+
+                    batch_indices += select_indexs
+            np.random.shuffle(batch_indices)
+
+            if len(batch_indices) == self.batch_size:
+                yield from reorder_index(batch_indices, self._world_size)
+
+
+class NaiveIdentitySampler(Sampler):
+    """
+    Randomly sample N identities, then for each identity,
+    randomly sample K instances, therefore batch size is N*K.
+    Args:
+    - data_source (list): list of (img_path, pid, camid).
+    - num_instances (int): number of instances per identity in a batch.
+    - batch_size (int): number of examples in a batch.
+    """
+
+    def __init__(self, data_source: str, mini_batch_size: int, num_instances: int, seed: Optional[int] = None):
+        self.data_source = data_source
+        self.num_instances = num_instances
+        self.num_pids_per_batch = mini_batch_size // self.num_instances
+
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+        self.batch_size = mini_batch_size * self._world_size
+
+        self.pid_index = defaultdict(list)
+
+        for index, info in enumerate(data_source):
+            pid = info[1]
+            self.pid_index[pid].append(index)
+
+        self.pids = sorted(list(self.pid_index.keys()))
+        self.num_identities = len(self.pids)
+
+        if seed is None:
+            seed = comm.shared_random_seed()
+        self._seed = int(seed)
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(self._infinite_indices(), start, None, self._world_size)
+
+    def _infinite_indices(self):
+        np.random.seed(self._seed)
+        while True:
+            avl_pids = copy.deepcopy(self.pids)
+            batch_idxs_dict = {}
+
+            batch_indices = []
+            while len(avl_pids) >= self.num_pids_per_batch:
+                selected_pids = np.random.choice(avl_pids, self.num_pids_per_batch, replace=False).tolist()
+                for pid in selected_pids:
+                    # Register pid in batch_idxs_dict if not
+                    if pid not in batch_idxs_dict:
+                        idxs = copy.deepcopy(self.pid_index[pid])
+                        if len(idxs) < self.num_instances:
+                            idxs = np.random.choice(idxs, size=self.num_instances, replace=True).tolist()
+                        np.random.shuffle(idxs)
+                        batch_idxs_dict[pid] = idxs
+
+                    avl_idxs = batch_idxs_dict[pid]
+                    for _ in range(self.num_instances):
+                        batch_indices.append(avl_idxs.pop(0))
+
+                    if len(avl_idxs) < self.num_instances: avl_pids.remove(pid)
+
+                if len(batch_indices) == self.batch_size:
+                    yield from reorder_index(batch_indices, self._world_size)
+                    batch_indices = []

fastreid/data/transforms/__init__.py

+# encoding: utf-8
+"""
+@author:  sherlock
+@contact: sherlockliao01@gmail.com
+"""
+
+from .autoaugment import AutoAugment
+from .build import build_transforms
+from .transforms import *
+
+__all__ = [k for k in globals().keys() if not k.startswith("_")]

fastreid/data/transforms/autoaugment.py

+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+""" AutoAugment, RandAugment, and AugMix for PyTorch
+This code implements the searched ImageNet policies with various tweaks and improvements and
+does not include any of the search code.
+AA and RA Implementation adapted from:
+    https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py
+AugMix adapted from:
+    https://github.com/google-research/augmix
+Papers:
+    AutoAugment: Learning Augmentation Policies from Data - https://arxiv.org/abs/1805.09501
+    Learning Data Augmentation Strategies for Object Detection - https://arxiv.org/abs/1906.11172
+    RandAugment: Practical automated data augmentation... - https://arxiv.org/abs/1909.13719
+    AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty - https://arxiv.org/abs/1912.02781
+Hacked together by Ross Wightman
+"""
+import math
+import random
+import re
+
+import PIL
+import numpy as np
+from PIL import Image, ImageOps, ImageEnhance
+
+_PIL_VER = tuple([int(x) for x in PIL.__version__.split('.')[:2]])
+
+_FILL = (128, 128, 128)
+
+# This signifies the max integer that the controller RNN could predict for the
+# augmentation scheme.
+_MAX_LEVEL = 10.
+
+_HPARAMS_DEFAULT = dict(
+    translate_const=57,
+    img_mean=_FILL,
+)
+
+_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)
+
+
+def _interpolation(kwargs):
+    interpolation = kwargs.pop('resample', Image.BILINEAR)
+    if isinstance(interpolation, (list, tuple)):
+        return random.choice(interpolation)
+    else:
+        return interpolation
+
+
+def _check_args_tf(kwargs):
+    if 'fillcolor' in kwargs and _PIL_VER < (5, 0):
+        kwargs.pop('fillcolor')
+    kwargs['resample'] = _interpolation(kwargs)
+
+
+def shear_x(img, factor, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0), **kwargs)
+
+
+def shear_y(img, factor, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0), **kwargs)
+
+
+def translate_x_rel(img, pct, **kwargs):
+    pixels = pct * img.size[0]
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)
+
+
+def translate_y_rel(img, pct, **kwargs):
+    pixels = pct * img.size[1]
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)
+
+
+def translate_x_abs(img, pixels, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)
+
+
+def translate_y_abs(img, pixels, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)
+
+
+def rotate(img, degrees, **kwargs):
+    _check_args_tf(kwargs)
+    if _PIL_VER >= (5, 2):
+        return img.rotate(degrees, **kwargs)
+    elif _PIL_VER >= (5, 0):
+        w, h = img.size
+        post_trans = (0, 0)
+        rotn_center = (w / 2.0, h / 2.0)
+        angle = -math.radians(degrees)
+        matrix = [
+            round(math.cos(angle), 15),
+            round(math.sin(angle), 15),
+            0.0,
+            round(-math.sin(angle), 15),
+            round(math.cos(angle), 15),
+            0.0,
+        ]
+
+        def transform(x, y, matrix):
+            (a, b, c, d, e, f) = matrix
+            return a * x + b * y + c, d * x + e * y + f
+
+        matrix[2], matrix[5] = transform(
+            -rotn_center[0] - post_trans[0], -rotn_center[1] - post_trans[1], matrix
+        )
+        matrix[2] += rotn_center[0]
+        matrix[5] += rotn_center[1]
+        return img.transform(img.size, Image.AFFINE, matrix, **kwargs)
+    else:
+        return img.rotate(degrees, resample=kwargs['resample'])
+
+
+def auto_contrast(img, **__):
+    return ImageOps.autocontrast(img)
+
+
+def invert(img, **__):
+    return ImageOps.invert(img)
+
+
+def equalize(img, **__):
+    return ImageOps.equalize(img)
+
+
+def solarize(img, thresh, **__):
+    return ImageOps.solarize(img, thresh)
+
+
+def solarize_add(img, add, thresh=128, **__):
+    lut = []
+    for i in range(256):
+        if i < thresh:
+            lut.append(min(255, i + add))
+        else:
+            lut.append(i)
+    if img.mode in ("L", "RGB"):
+        if img.mode == "RGB" and len(lut) == 256:
+            lut = lut + lut + lut
+        return img.point(lut)
+    else:
+        return img
+
+
+def posterize(img, bits_to_keep, **__):
+    if bits_to_keep >= 8:
+        return img
+    return ImageOps.posterize(img, bits_to_keep)
+
+
+def contrast(img, factor, **__):
+    return ImageEnhance.Contrast(img).enhance(factor)
+
+
+def color(img, factor, **__):
+    return ImageEnhance.Color(img).enhance(factor)
+
+
+def brightness(img, factor, **__):
+    return ImageEnhance.Brightness(img).enhance(factor)
+
+
+def sharpness(img, factor, **__):
+    return ImageEnhance.Sharpness(img).enhance(factor)
+
+
+def _randomly_negate(v):
+    """With 50% prob, negate the value"""
+    return -v if random.random() > 0.5 else v
+
+
+def _rotate_level_to_arg(level, _hparams):
+    # range [-30, 30]
+    level = (level / _MAX_LEVEL) * 30.
+    level = _randomly_negate(level)
+    return level,
+
+
+def _enhance_level_to_arg(level, _hparams):
+    # range [0.1, 1.9]
+    return (level / _MAX_LEVEL) * 1.8 + 0.1,
+
+
+def _enhance_increasing_level_to_arg(level, _hparams):
+    # the 'no change' level is 1.0, moving away from that towards 0. or 2.0 increases the enhancement blend
+    # range [0.1, 1.9]
+    level = (level / _MAX_LEVEL) * .9
+    level = 1.0 + _randomly_negate(level)
+    return level,
+
+
+def _shear_level_to_arg(level, _hparams):
+    # range [-0.3, 0.3]
+    level = (level / _MAX_LEVEL) * 0.3
+    level = _randomly_negate(level)
+    return level,
+
+
+def _translate_abs_level_to_arg(level, hparams):
+    translate_const = hparams['translate_const']
+    level = (level / _MAX_LEVEL) * float(translate_const)
+    level = _randomly_negate(level)
+    return level,
+
+
+def _translate_rel_level_to_arg(level, hparams):
+    # default range [-0.45, 0.45]
+    translate_pct = hparams.get('translate_pct', 0.45)
+    level = (level / _MAX_LEVEL) * translate_pct
+    level = _randomly_negate(level)
+    return level,
+
+
+def _posterize_level_to_arg(level, _hparams):
+    # As per Tensorflow TPU EfficientNet impl
+    # range [0, 4], 'keep 0 up to 4 MSB of original image'
+    # intensity/severity of augmentation decreases with level
+    return int((level / _MAX_LEVEL) * 4),
+
+
+def _posterize_increasing_level_to_arg(level, hparams):
+    # As per Tensorflow models research and UDA impl
+    # range [4, 0], 'keep 4 down to 0 MSB of original image',
+    # intensity/severity of augmentation increases with level
+    return 4 - _posterize_level_to_arg(level, hparams)[0],
+
+
+def _posterize_original_level_to_arg(level, _hparams):
+    # As per original AutoAugment paper description
+    # range [4, 8], 'keep 4 up to 8 MSB of image'
+    # intensity/severity of augmentation decreases with level
+    return int((level / _MAX_LEVEL) * 4) + 4,
+
+
+def _solarize_level_to_arg(level, _hparams):
+    # range [0, 256]
+    # intensity/severity of augmentation decreases with level
+    return int((level / _MAX_LEVEL) * 256),
+
+
+def _solarize_increasing_level_to_arg(level, _hparams):
+    # range [0, 256]
+    # intensity/severity of augmentation increases with level
+    return 256 - _solarize_level_to_arg(level, _hparams)[0],
+
+
+def _solarize_add_level_to_arg(level, _hparams):
+    # range [0, 110]
+    return int((level / _MAX_LEVEL) * 110),
+
+
+LEVEL_TO_ARG = {
+    'AutoContrast': None,
+    'Equalize': None,
+    'Invert': None,
+    'Rotate': _rotate_level_to_arg,
+    # There are several variations of the posterize level scaling in various Tensorflow/Google repositories/papers
+    'Posterize': _posterize_level_to_arg,
+    'PosterizeIncreasing': _posterize_increasing_level_to_arg,
+    'PosterizeOriginal': _posterize_original_level_to_arg,
+    'Solarize': _solarize_level_to_arg,
+    'SolarizeIncreasing': _solarize_increasing_level_to_arg,
+    'SolarizeAdd': _solarize_add_level_to_arg,
+    'Color': _enhance_level_to_arg,
+    'ColorIncreasing': _enhance_increasing_level_to_arg,
+    'Contrast': _enhance_level_to_arg,
+    'ContrastIncreasing': _enhance_increasing_level_to_arg,
+    'Brightness': _enhance_level_to_arg,
+    'BrightnessIncreasing': _enhance_increasing_level_to_arg,
+    'Sharpness': _enhance_level_to_arg,
+    'SharpnessIncreasing': _enhance_increasing_level_to_arg,
+    'ShearX': _shear_level_to_arg,
+    'ShearY': _shear_level_to_arg,
+    'TranslateX': _translate_abs_level_to_arg,
+    'TranslateY': _translate_abs_level_to_arg,
+    'TranslateXRel': _translate_rel_level_to_arg,
+    'TranslateYRel': _translate_rel_level_to_arg,
+}
+
+NAME_TO_OP = {
+    'AutoContrast': auto_contrast,
+    'Equalize': equalize,
+    'Invert': invert,
+    'Rotate': rotate,
+    'Posterize': posterize,
+    'PosterizeIncreasing': posterize,
+    'PosterizeOriginal': posterize,
+    'Solarize': solarize,
+    'SolarizeIncreasing': solarize,
+    'SolarizeAdd': solarize_add,
+    'Color': color,
+    'ColorIncreasing': color,
+    'Contrast': contrast,
+    'ContrastIncreasing': contrast,
+    'Brightness': brightness,
+    'BrightnessIncreasing': brightness,
+    'Sharpness': sharpness,
+    'SharpnessIncreasing': sharpness,
+    'ShearX': shear_x,
+    'ShearY': shear_y,
+    'TranslateX': translate_x_abs,
+    'TranslateY': translate_y_abs,
+    'TranslateXRel': translate_x_rel,
+    'TranslateYRel': translate_y_rel,
+}
+
+
+class AugmentOp:
+
+    def __init__(self, name, prob=0.5, magnitude=10, hparams=None):
+        hparams = hparams or _HPARAMS_DEFAULT
+        self.aug_fn = NAME_TO_OP[name]
+        self.level_fn = LEVEL_TO_ARG[name]
+        self.prob = prob
+        self.magnitude = magnitude
+        self.hparams = hparams.copy()
+        self.kwargs = dict(
+            fillcolor=hparams['img_mean'] if 'img_mean' in hparams else _FILL,
+            resample=hparams['interpolation'] if 'interpolation' in hparams else _RANDOM_INTERPOLATION,
+        )
+
+        # If magnitude_std is > 0, we introduce some randomness
+        # in the usually fixed policy and sample magnitude from a normal distribution
+        # with mean `magnitude` and std-dev of `magnitude_std`.
+        # NOTE This is my own hack, being tested, not in papers or reference impls.
+        self.magnitude_std = self.hparams.get('magnitude_std', 0)
+
+    def __call__(self, img):
+        if self.prob < 1.0 and random.random() > self.prob:
+            return img
+        magnitude = self.magnitude
+        if self.magnitude_std and self.magnitude_std > 0:
+            magnitude = random.gauss(magnitude, self.magnitude_std)
+        magnitude = min(_MAX_LEVEL, max(0, magnitude))  # clip to valid range
+        level_args = self.level_fn(magnitude, self.hparams) if self.level_fn is not None else tuple()
+        return self.aug_fn(img, *level_args, **self.kwargs)
+
+
+def auto_augment_policy_v0(hparams):
+    # ImageNet v0 policy from TPU EfficientNet impl, cannot find a paper reference.
+    policy = [
+        [('Equalize', 0.8, 1), ('ShearY', 0.8, 4)],
+        [('Color', 0.4, 9), ('Equalize', 0.6, 3)],
+        [('Color', 0.4, 1), ('Rotate', 0.6, 8)],
+        [('Solarize', 0.8, 3), ('Equalize', 0.4, 7)],
+        [('Solarize', 0.4, 2), ('Solarize', 0.6, 2)],
+        [('Color', 0.2, 0), ('Equalize', 0.8, 8)],
+        [('Equalize', 0.4, 8), ('SolarizeAdd', 0.8, 3)],
+        [('ShearX', 0.2, 9), ('Rotate', 0.6, 8)],
+        [('Color', 0.6, 1), ('Equalize', 1.0, 2)],
+        [('Invert', 0.4, 9), ('Rotate', 0.6, 0)],
+        [('Equalize', 1.0, 9), ('ShearY', 0.6, 3)],
+        [('Color', 0.4, 7), ('Equalize', 0.6, 0)],
+        [('Posterize', 0.4, 6), ('AutoContrast', 0.4, 7)],
+        [('Solarize', 0.6, 8), ('Color', 0.6, 9)],
+        [('Solarize', 0.2, 4), ('Rotate', 0.8, 9)],
+        [('Rotate', 1.0, 7), ('TranslateYRel', 0.8, 9)],
+        [('ShearX', 0.0, 0), ('Solarize', 0.8, 4)],
+        [('ShearY', 0.8, 0), ('Color', 0.6, 4)],
+        [('Color', 1.0, 0), ('Rotate', 0.6, 2)],
+        [('Equalize', 0.8, 4), ('Equalize', 0.0, 8)],
+        [('Equalize', 1.0, 4), ('AutoContrast', 0.6, 2)],
+        [('ShearY', 0.4, 7), ('SolarizeAdd', 0.6, 7)],
+        [('Posterize', 0.8, 2), ('Solarize', 0.6, 10)],  # This results in black image with Tpu posterize
+        [('Solarize', 0.6, 8), ('Equalize', 0.6, 1)],
+        [('Color', 0.8, 6), ('Rotate', 0.4, 5)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy_v0r(hparams):
+    # ImageNet v0 policy from TPU EfficientNet impl, with variation of Posterize used
+    # in Google research implementation (number of bits discarded increases with magnitude)
+    policy = [
+        [('Equalize', 0.8, 1), ('ShearY', 0.8, 4)],
+        [('Color', 0.4, 9), ('Equalize', 0.6, 3)],
+        [('Color', 0.4, 1), ('Rotate', 0.6, 8)],
+        [('Solarize', 0.8, 3), ('Equalize', 0.4, 7)],
+        [('Solarize', 0.4, 2), ('Solarize', 0.6, 2)],
+        [('Color', 0.2, 0), ('Equalize', 0.8, 8)],
+        [('Equalize', 0.4, 8), ('SolarizeAdd', 0.8, 3)],
+        [('ShearX', 0.2, 9), ('Rotate', 0.6, 8)],
+        [('Color', 0.6, 1), ('Equalize', 1.0, 2)],
+        [('Invert', 0.4, 9), ('Rotate', 0.6, 0)],
+        [('Equalize', 1.0, 9), ('ShearY', 0.6, 3)],
+        [('Color', 0.4, 7), ('Equalize', 0.6, 0)],
+        [('PosterizeIncreasing', 0.4, 6), ('AutoContrast', 0.4, 7)],
+        [('Solarize', 0.6, 8), ('Color', 0.6, 9)],
+        [('Solarize', 0.2, 4), ('Rotate', 0.8, 9)],
+        [('Rotate', 1.0, 7), ('TranslateYRel', 0.8, 9)],
+        [('ShearX', 0.0, 0), ('Solarize', 0.8, 4)],
+        [('ShearY', 0.8, 0), ('Color', 0.6, 4)],
+        [('Color', 1.0, 0), ('Rotate', 0.6, 2)],
+        [('Equalize', 0.8, 4), ('Equalize', 0.0, 8)],
+        [('Equalize', 1.0, 4), ('AutoContrast', 0.6, 2)],
+        [('ShearY', 0.4, 7), ('SolarizeAdd', 0.6, 7)],
+        [('PosterizeIncreasing', 0.8, 2), ('Solarize', 0.6, 10)],
+        [('Solarize', 0.6, 8), ('Equalize', 0.6, 1)],
+        [('Color', 0.8, 6), ('Rotate', 0.4, 5)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy_original(hparams):
+    # ImageNet policy from https://arxiv.org/abs/1805.09501
+    policy = [
+        [('PosterizeOriginal', 0.4, 8), ('Rotate', 0.6, 9)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+        [('PosterizeOriginal', 0.6, 7), ('PosterizeOriginal', 0.6, 6)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Equalize', 0.4, 4), ('Rotate', 0.8, 8)],
+        [('Solarize', 0.6, 3), ('Equalize', 0.6, 7)],
+        [('PosterizeOriginal', 0.8, 5), ('Equalize', 1.0, 2)],
+        [('Rotate', 0.2, 3), ('Solarize', 0.6, 8)],
+        [('Equalize', 0.6, 8), ('PosterizeOriginal', 0.4, 6)],
+        [('Rotate', 0.8, 8), ('Color', 0.4, 0)],
+        [('Rotate', 0.4, 9), ('Equalize', 0.6, 2)],
+        [('Equalize', 0.0, 7), ('Equalize', 0.8, 8)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Rotate', 0.8, 8), ('Color', 1.0, 2)],
+        [('Color', 0.8, 8), ('Solarize', 0.8, 7)],
+        [('Sharpness', 0.4, 7), ('Invert', 0.6, 8)],
+        [('ShearX', 0.6, 5), ('Equalize', 1.0, 9)],
+        [('Color', 0.4, 0), ('Equalize', 0.6, 3)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy_originalr(hparams):
+    # ImageNet policy from https://arxiv.org/abs/1805.09501 with research posterize variation
+    policy = [
+        [('PosterizeIncreasing', 0.4, 8), ('Rotate', 0.6, 9)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+        [('PosterizeIncreasing', 0.6, 7), ('PosterizeIncreasing', 0.6, 6)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Equalize', 0.4, 4), ('Rotate', 0.8, 8)],
+        [('Solarize', 0.6, 3), ('Equalize', 0.6, 7)],
+        [('PosterizeIncreasing', 0.8, 5), ('Equalize', 1.0, 2)],
+        [('Rotate', 0.2, 3), ('Solarize', 0.6, 8)],
+        [('Equalize', 0.6, 8), ('PosterizeIncreasing', 0.4, 6)],
+        [('Rotate', 0.8, 8), ('Color', 0.4, 0)],
+        [('Rotate', 0.4, 9), ('Equalize', 0.6, 2)],
+        [('Equalize', 0.0, 7), ('Equalize', 0.8, 8)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Rotate', 0.8, 8), ('Color', 1.0, 2)],
+        [('Color', 0.8, 8), ('Solarize', 0.8, 7)],
+        [('Sharpness', 0.4, 7), ('Invert', 0.6, 8)],
+        [('ShearX', 0.6, 5), ('Equalize', 1.0, 9)],
+        [('Color', 0.4, 0), ('Equalize', 0.6, 3)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy(name="original"):
+    hparams = _HPARAMS_DEFAULT
+    if name == 'original':
+        return auto_augment_policy_original(hparams)
+    elif name == 'originalr':
+        return auto_augment_policy_originalr(hparams)
+    elif name == 'v0':
+        return auto_augment_policy_v0(hparams)
+    elif name == 'v0r':
+        return auto_augment_policy_v0r(hparams)
+    else:
+        assert False, 'Unknown AA policy (%s)' % name
+
+
+class AutoAugment:
+
+    def __init__(self):
+        self.policy = auto_augment_policy()
+
+    def __call__(self, img):
+        sub_policy = random.choice(self.policy)
+        for op in sub_policy:
+            img = op(img)
+        return img
+
+
+def auto_augment_transform(config_str, hparams):
+    """
+    Create a AutoAugment transform
+    :param config_str: String defining configuration of auto augmentation. Consists of multiple sections separated by
+    dashes ('-'). The first section defines the AutoAugment policy (one of 'v0', 'v0r', 'original', 'originalr').
+    The remaining sections, not order sepecific determine
+        'mstd' -  float std deviation of magnitude noise applied
+    Ex 'original-mstd0.5' results in AutoAugment with original policy, magnitude_std 0.5
+    :param hparams: Other hparams (kwargs) for the AutoAugmentation scheme
+    :return: A PyTorch compatible Transform
+    """
+    config = config_str.split('-')
+    policy_name = config[0]
+    config = config[1:]
+    for c in config:
+        cs = re.split(r'(\d.*)', c)
+        if len(cs) < 2:
+            continue
+        key, val = cs[:2]
+        if key == 'mstd':
+            # noise param injected via hparams for now
+            hparams.setdefault('magnitude_std', float(val))
+        else:
+            assert False, 'Unknown AutoAugment config section'
+    aa_policy = auto_augment_policy(policy_name)
+    return AutoAugment(aa_policy)
+
+
+_RAND_TRANSFORMS = [
+    'AutoContrast',
+    'Equalize',
+    'Invert',
+    'Rotate',
+    'Posterize',
+    'Solarize',
+    'SolarizeAdd',
+    'Color',
+    'Contrast',
+    'Brightness',
+    'Sharpness',
+    'ShearX',
+    'ShearY',
+    'TranslateXRel',
+    'TranslateYRel',
+    # 'Cutout'  # NOTE I've implement this as random erasing separately
+]
+
+_RAND_INCREASING_TRANSFORMS = [
+    'AutoContrast',
+    'Equalize',
+    'Invert',
+    'Rotate',
+    'PosterizeIncreasing',
+    'SolarizeIncreasing',
+    'SolarizeAdd',
+    'ColorIncreasing',
+    'ContrastIncreasing',
+    'BrightnessIncreasing',
+    'SharpnessIncreasing',
+    'ShearX',
+    'ShearY',
+    'TranslateXRel',
+    'TranslateYRel',
+    # 'Cutout'  # NOTE I've implement this as random erasing separately
+]
+
+# These experimental weights are based loosely on the relative improvements mentioned in paper.
+# They may not result in increased performance, but could likely be tuned to so.
+_RAND_CHOICE_WEIGHTS_0 = {
+    'Rotate': 0.3,
+    'ShearX': 0.2,
+    'ShearY': 0.2,
+    'TranslateXRel': 0.1,
+    'TranslateYRel': 0.1,
+    'Color': .025,
+    'Sharpness': 0.025,
+    'AutoContrast': 0.025,
+    'Solarize': .005,
+    'SolarizeAdd': .005,
+    'Contrast': .005,
+    'Brightness': .005,
+    'Equalize': .005,
+    'Posterize': 0,
+    'Invert': 0,
+}
+
+
+def _select_rand_weights(weight_idx=0, transforms=None):
+    transforms = transforms or _RAND_TRANSFORMS
+    assert weight_idx == 0  # only one set of weights currently
+    rand_weights = _RAND_CHOICE_WEIGHTS_0
+    probs = [rand_weights[k] for k in transforms]
+    probs /= np.sum(probs)
+    return probs
+
+
+def rand_augment_ops(magnitude=10, hparams=None, transforms=None):
+    hparams = hparams or _HPARAMS_DEFAULT
+    transforms = transforms or _RAND_TRANSFORMS
+    return [AugmentOp(
+        name, prob=0.5, magnitude=magnitude, hparams=hparams) for name in transforms]
+
+
+class RandAugment:
+    def __init__(self, ops, num_layers=2, choice_weights=None):
+        self.ops = ops
+        self.num_layers = num_layers
+        self.choice_weights = choice_weights
+
+    def __call__(self, img):
+        # no replacement when using weighted choice
+        ops = np.random.choice(
+            self.ops, self.num_layers, replace=self.choice_weights is None, p=self.choice_weights)
+        for op in ops:
+            img = op(img)
+        return img
+
+
+def rand_augment_transform(config_str, hparams):
+    """
+    Create a RandAugment transform
+    :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by
+    dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining
+    sections, not order sepecific determine
+        'm' - integer magnitude of rand augment
+        'n' - integer num layers (number of transform ops selected per image)
+        'w' - integer probabiliy weight index (index of a set of weights to influence choice of op)
+        'mstd' -  float std deviation of magnitude noise applied
+        'inc' - integer (bool), use augmentations that increase in severity with magnitude (default: 0)
+    Ex 'rand-m9-n3-mstd0.5' results in RandAugment with magnitude 9, num_layers 3, magnitude_std 0.5
+    'rand-mstd1-w0' results in magnitude_std 1.0, weights 0, default magnitude of 10 and num_layers 2
+    :param hparams: Other hparams (kwargs) for the RandAugmentation scheme
+    :return: A PyTorch compatible Transform
+    """
+    magnitude = _MAX_LEVEL  # default to _MAX_LEVEL for magnitude (currently 10)
+    num_layers = 2  # default to 2 ops per image
+    weight_idx = None  # default to no probability weights for op choice
+    transforms = _RAND_TRANSFORMS
+    config = config_str.split('-')
+    assert config[0] == 'rand'
+    config = config[1:]
+    for c in config:
+        cs = re.split(r'(\d.*)', c)
+        if len(cs) < 2:
+            continue
+        key, val = cs[:2]
+        if key == 'mstd':
+            # noise param injected via hparams for now
+            hparams.setdefault('magnitude_std', float(val))
+        elif key == 'inc':
+            if bool(val):
+                transforms = _RAND_INCREASING_TRANSFORMS
+        elif key == 'm':
+            magnitude = int(val)
+        elif key == 'n':
+            num_layers = int(val)
+        elif key == 'w':
+            weight_idx = int(val)
+        else:
+            assert False, 'Unknown RandAugment config section'
+    ra_ops = rand_augment_ops(magnitude=magnitude, hparams=hparams, transforms=transforms)
+    choice_weights = None if weight_idx is None else _select_rand_weights(weight_idx)
+    return RandAugment(ra_ops, num_layers, choice_weights=choice_weights)
+
+
+_AUGMIX_TRANSFORMS = [
+    'AutoContrast',
+    'ColorIncreasing',  # not in paper
+    'ContrastIncreasing',  # not in paper
+    'BrightnessIncreasing',  # not in paper
+    'SharpnessIncreasing',  # not in paper
+    'Equalize',
+    'Rotate',
+    'PosterizeIncreasing',
+    'SolarizeIncreasing',
+    'ShearX',
+    'ShearY',
+    'TranslateXRel',
+    'TranslateYRel',
+]
+
+
+def augmix_ops(magnitude=10, hparams=None, transforms=None):
+    hparams = hparams or _HPARAMS_DEFAULT
+    transforms = transforms or _AUGMIX_TRANSFORMS
+    return [AugmentOp(
+        name, prob=1.0, magnitude=magnitude, hparams=hparams) for name in transforms]
+
+
+class AugMixAugment:
+    """ AugMix Transform
+    Adapted and improved from impl here: https://github.com/google-research/augmix/blob/master/imagenet.py
+    From paper: 'AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty -
+    https://arxiv.org/abs/1912.02781
+    """
+
+    def __init__(self, ops, alpha=1., width=3, depth=-1, blended=False):
+        self.ops = ops
+        self.alpha = alpha
+        self.width = width
+        self.depth = depth
+        self.blended = blended  # blended mode is faster but not well tested
+
+    def _calc_blended_weights(self, ws, m):
+        ws = ws * m
+        cump = 1.
+        rws = []
+        for w in ws[::-1]:
+            alpha = w / cump
+            cump *= (1 - alpha)
+            rws.append(alpha)
+        return np.array(rws[::-1], dtype=np.float32)
+
+    def _apply_blended(self, img, mixing_weights, m):
+        # This is my first crack and implementing a slightly faster mixed augmentation. Instead
+        # of accumulating the mix for each chain in a Numpy array and then blending with original,
+        # it recomputes the blending coefficients and applies one PIL image blend per chain.
+        # TODO the results appear in the right ballpark but they differ by more than rounding.
+        img_orig = img.copy()
+        ws = self._calc_blended_weights(mixing_weights, m)
+        for w in ws:
+            depth = self.depth if self.depth > 0 else np.random.randint(1, 4)
+            ops = np.random.choice(self.ops, depth, replace=True)
+            img_aug = img_orig  # no ops are in-place, deep copy not necessary
+            for op in ops:
+                img_aug = op(img_aug)
+            img = Image.blend(img, img_aug, w)
+        return img
+
+    def _apply_basic(self, img, mixing_weights, m):
+        # This is a literal adaptation of the paper/official implementation without normalizations and
+        # PIL <-> Numpy conversions between every op. It is still quite CPU compute heavy compared to the
+        # typical augmentation transforms, could use a GPU / Kornia implementation.
+        img_shape = img.size[0], img.size[1], len(img.getbands())
+        mixed = np.zeros(img_shape, dtype=np.float32)
+        for mw in mixing_weights:
+            depth = self.depth if self.depth > 0 else np.random.randint(1, 4)
+            ops = np.random.choice(self.ops, depth, replace=True)
+            img_aug = img  # no ops are in-place, deep copy not necessary
+            for op in ops:
+                img_aug = op(img_aug)
+            mixed += mw * np.asarray(img_aug, dtype=np.float32)
+        np.clip(mixed, 0, 255., out=mixed)
+        mixed = Image.fromarray(mixed.astype(np.uint8))
+        return Image.blend(img, mixed, m)
+
+    def __call__(self, img):
+        mixing_weights = np.float32(np.random.dirichlet([self.alpha] * self.width))
+        m = np.float32(np.random.beta(self.alpha, self.alpha))
+        if self.blended:
+            mixed = self._apply_blended(img, mixing_weights, m)
+        else:
+            mixed = self._apply_basic(img, mixing_weights, m)
+        return mixed
+
+
+def augment_and_mix_transform(config_str, hparams):
+    """ Create AugMix PyTorch transform
+    :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by
+    dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining
+    sections, not order sepecific determine
+        'm' - integer magnitude (severity) of augmentation mix (default: 3)
+        'w' - integer width of augmentation chain (default: 3)
+        'd' - integer depth of augmentation chain (-1 is random [1, 3], default: -1)
+        'b' - integer (bool), blend each branch of chain into end result without a final blend, less CPU (default: 0)
+        'mstd' -  float std deviation of magnitude noise applied (default: 0)
+    Ex 'augmix-m5-w4-d2' results in AugMix with severity 5, chain width 4, chain depth 2
+    :param hparams: Other hparams (kwargs) for the Augmentation transforms
+    :return: A PyTorch compatible Transform
+    """
+    magnitude = 3
+    width = 3
+    depth = -1
+    alpha = 1.
+    blended = False
+    config = config_str.split('-')
+    assert config[0] == 'augmix'
+    config = config[1:]
+    for c in config:
+        cs = re.split(r'(\d.*)', c)
+        if len(cs) < 2:
+            continue
+        key, val = cs[:2]
+        if key == 'mstd':
+            # noise param injected via hparams for now
+            hparams.setdefault('magnitude_std', float(val))
+        elif key == 'm':
+            magnitude = int(val)
+        elif key == 'w':
+            width = int(val)
+        elif key == 'd':
+            depth = int(val)
+        elif key == 'a':
+            alpha = float(val)
+        elif key == 'b':
+            blended = bool(val)
+        else:
+            assert False, 'Unknown AugMix config section'
+    ops = augmix_ops(magnitude=magnitude, hparams=hparams)
+    return AugMixAugment(ops, alpha=alpha, width=width, depth=depth, blended=blended)