plato-learn 1.1__py3-none-any.whl

plato/datasources/datalib/refer_utils/referitgame_utils.py

@@ -0,0 +1,237 @@
+import sys
+import os.path as osp
+import json
+import logging
+import pickle
+import time
+import itertools
+import skimage.io as io
+
+from pprint import pprint
+import numpy as np
+
+
+class REFER:
+    def __init__(
+        self,
+        data_root,  # the root where the data is stored
+        image_dataroot,  # the path of the source images
+        dataset="refcoco",
+        splitBy="unc",
+    ):
+        # provide data_root folder which contains refclef, refcoco, refcoco+ and refcocog
+        # also provide dataset name and splitBy information
+        # e.g., dataset = 'refcoco', splitBy = 'unc'
+        print("loading dataset %s into memory..." % dataset)
+
+        self.ROOT_DIR = osp.abspath(osp.dirname(__file__))
+        self.DATA_DIR = osp.join(data_root, dataset)
+        if dataset in ["refcoco", "refcoco+", "refcocog"]:
+            self.IMAGE_DIR = image_dataroot
+        else:
+            print("No refer dataset is called [%s]" % dataset)
+            sys.exit()
+
+        # load refs from data/datasetrefs(dataset).json
+        tic = time.perf_counter()
+        ref_file = osp.join(self.DATA_DIR, "refs(" + splitBy + ").p")
+        self.data = {}
+        self.data["dataset"] = dataset
+        self.data["refs"] = pickle.load(open(ref_file, "rb"))
+
+        # load annotations from data/dataset/instances.json
+        instances_file = osp.join(self.DATA_DIR, "instances.json")
+        instances = json.load(open(instances_file, "r"))
+        self.data["images"] = instances["images"]
+        self.data["annotations"] = instances["annotations"]
+        self.data["categories"] = instances["categories"]
+
+        # create index
+        self.createIndex()
+        print("DONE (t=%.2fs)" % (time.perf_counter() - tic))
+
+    def createIndex(self):
+        # create sets of mapping
+        # 1)  Refs:         {ref_id: ref}
+        # 2)  Anns:         {ann_id: ann}
+        # 3)  Imgs:         {image_id: image}
+        # 4)  Cats:         {category_id: category_name}
+        # 5)  Sents:        {sent_id: sent}
+        # 6)  imgToRefs:    {image_id: refs}
+        # 7)  imgToAnns:    {image_id: anns}
+        # 8)  refToAnn:     {ref_id: ann}
+        # 9)  annToRef:     {ann_id: ref}
+        # 10) catToRefs:    {category_id: refs}
+        # 11) sentToRef:    {sent_id: ref}
+        # 12) sentToTokens: {sent_id: tokens}
+        print("creating index...")
+        # fetch info from instances
+        Anns, Imgs, Cats, imgToAnns = {}, {}, {}, {}
+        for ann in self.data["annotations"]:
+            Anns[ann["id"]] = ann
+            imgToAnns[ann["image_id"]] = imgToAnns.get(ann["image_id"], []) + [ann]
+        for img in self.data["images"]:
+            img["file_name"] = img["file_name"].strip().split("_")[-1]
+            Imgs[img["id"]] = img
+        for cat in self.data["categories"]:
+            Cats[cat["id"]] = cat["name"]
+
+        # fetch info from refs
+        Refs, imgToRefs, refToAnn, annToRef, catToRefs = {}, {}, {}, {}, {}
+        Sents, sentToRef, sentToTokens = {}, {}, {}
+        for ref in self.data["refs"]:
+            # ids
+            ref_id = ref["ref_id"]
+            ann_id = ref["ann_id"]
+            category_id = ref["category_id"]
+            image_id = ref["image_id"]
+
+            # add mapping related to ref
+            Refs[ref_id] = ref
+            imgToRefs[image_id] = imgToRefs.get(image_id, []) + [ref]
+            catToRefs[category_id] = catToRefs.get(category_id, []) + [ref]
+            refToAnn[ref_id] = Anns[ann_id]
+            annToRef[ann_id] = ref
+
+            # add mapping of sent
+            for sent in ref["sentences"]:
+                Sents[sent["sent_id"]] = sent
+                sentToRef[sent["sent_id"]] = ref
+                sentToTokens[sent["sent_id"]] = sent["tokens"]
+
+        # create class members
+        self.Refs = Refs
+        self.Anns = Anns
+        self.Imgs = Imgs
+        self.Cats = Cats
+        self.Sents = Sents
+        self.imgToRefs = imgToRefs
+        self.imgToAnns = imgToAnns
+        self.refToAnn = refToAnn
+        self.annToRef = annToRef
+        self.catToRefs = catToRefs
+        self.sentToRef = sentToRef
+        self.sentToTokens = sentToTokens
+        print("index created.")
+
+    def getRefIds(self, image_ids=[], cat_ids=[], ref_ids=[], split=""):
+        image_ids = image_ids if type(image_ids) == list else [image_ids]
+        cat_ids = cat_ids if type(cat_ids) == list else [cat_ids]
+        ref_ids = ref_ids if type(ref_ids) == list else [ref_ids]
+
+        if len(image_ids) == len(cat_ids) == len(ref_ids) == len(split) == 0:
+            refs = self.data["refs"]
+        else:
+            if not len(image_ids) == 0:
+                refs = [self.imgToRefs[image_id] for image_id in image_ids]
+            else:
+                refs = self.data["refs"]
+            if not len(cat_ids) == 0:
+                refs = [ref for ref in refs if ref["category_id"] in cat_ids]
+            if not len(ref_ids) == 0:
+                refs = [ref for ref in refs if ref["ref_id"] in ref_ids]
+            if not len(split) == 0:
+                if split in ["testA", "testB", "testC"]:
+                    refs = [
+                        ref for ref in refs if split[-1] in ref["split"]
+                    ]  # we also consider testAB, testBC, ...
+                elif split in ["testAB", "testBC", "testAC"]:
+                    refs = [
+                        ref for ref in refs if ref["split"] == split
+                    ]  # rarely used I guess...
+                elif split == "test":
+                    refs = [ref for ref in refs if "test" in ref["split"]]
+                elif split == "train" or split == "val":
+                    refs = [ref for ref in refs if ref["split"] == split]
+                else:
+                    print("No such split [%s]" % split)
+                    sys.exit()
+        ref_ids = [ref["ref_id"] for ref in refs]
+        return ref_ids
+
+    def getAnnIds(self, image_ids=[], cat_ids=[], ref_ids=[]):
+        image_ids = image_ids if type(image_ids) == list else [image_ids]
+        cat_ids = cat_ids if type(cat_ids) == list else [cat_ids]
+        ref_ids = ref_ids if type(ref_ids) == list else [ref_ids]
+
+        if len(image_ids) == len(cat_ids) == len(ref_ids) == 0:
+            ann_ids = [ann["id"] for ann in self.data["annotations"]]
+        else:
+            if not len(image_ids) == 0:
+                lists = [
+                    self.imgToAnns[image_id]
+                    for image_id in image_ids
+                    if image_id in self.imgToAnns
+                ]  # list of [anns]
+                anns = list(itertools.chain.from_iterable(lists))
+            else:
+                anns = self.data["annotations"]
+            if not len(cat_ids) == 0:
+                anns = [ann for ann in anns if ann["category_id"] in cat_ids]
+            ann_ids = [ann["id"] for ann in anns]
+            if not len(ref_ids) == 0:
+                ids = set(ann_ids).intersection(
+                    set([self.Refs[ref_id]["ann_id"] for ref_id in ref_ids])
+                )
+        return ann_ids
+
+    def getImgIds(self, ref_ids=[]):
+        ref_ids = ref_ids if type(ref_ids) == list else [ref_ids]
+
+        if not len(ref_ids) == 0:
+            image_ids = list(set([self.Refs[ref_id]["image_id"] for ref_id in ref_ids]))
+        else:
+            image_ids = self.Imgs.keys()
+        return image_ids
+
+    def getCatIds(self):
+        return self.Cats.keys()
+
+    def loadRefs(self, ref_ids=[]):
+        if type(ref_ids) == list:
+            return [self.Refs[ref_id] for ref_id in ref_ids]
+        elif type(ref_ids) == int:
+            return [self.Refs[ref_ids]]
+
+    def loadAnns(self, ann_ids=[]):
+        if type(ann_ids) == list:
+            return [self.Anns[ann_id] for ann_id in ann_ids]
+        elif type(ann_ids) == int or type(ann_ids) == unicode:
+            return [self.Anns[ann_ids]]
+
+    def loadImgs(self, image_ids=[]):
+        if type(image_ids) == list:
+            return [self.Imgs[image_id] for image_id in image_ids]
+        elif type(image_ids) == int:
+            return [self.Imgs[image_ids]]
+
+    def loadImgspath(self, image_ids=[]):
+        if type(image_ids) == list:
+            return [
+                osp.join(self.IMAGE_DIR, self.Imgs[image_id]["file_name"])
+                for image_id in image_ids
+            ]
+        elif type(image_ids) == int:
+            return [osp.join(self.IMAGE_DIR, self.Imgs[image_ids]["file_name"])]
+
+    def loadImgsData(self, image_ids=[]):  # the image_ids is obtained from refs
+        if type(image_ids) == list:
+            return [
+                io.imread(osp.join(self.IMAGE_DIR, self.Imgs[image_id]["file_name"]))
+                for image_id in image_ids
+            ]
+        elif type(image_ids) == int:
+            return [
+                io.imread(osp.join(self.IMAGE_DIR, self.Imgs[image_ids]["file_name"]))
+            ]
+
+    def loadCats(self, cat_ids=[]):
+        if type(cat_ids) == list:
+            return [self.Cats[cat_id] for cat_id in cat_ids]
+        elif type(cat_ids) == int:
+            return [self.Cats[cat_ids]]
+
+    def getRefBox(self, ref_id):
+        ref = self.Refs[ref_id]
+        ann = self.refToAnn[ref_id]
+        return ann["bbox"]  # [x, y, w, h]

plato/datasources/datalib/tiny_data_tools.py

@@ -0,0 +1,81 @@
+"""
+This includes tools for creating the processing tiny dataset.
+
+The tiny dataset is sampled from the whole dataset.
+
+"""
+
+import os
+import collections
+import logging
+
+import pandas as pd
+import numpy as np
+
+
+def create_tiny_kinetics_anno(kinetics_annotation_files_info, num_samples, random_seed):
+    """Creating the annotation files for tiny kinetics.
+
+    Args:
+        kinetics_annotation_files_info (dict): a dict contains the annotation
+         files for different splits. e.g., {"train": xxx, "val": xxx}.
+        num_fist_videos (int): the number of samples utilized to create
+         tiny dataset.
+        random_seed (int): the random seed for sampling samples.
+    """
+    train_anno_file_path = kinetics_annotation_files_info["train"]
+    test_anno_file_path = kinetics_annotation_files_info["test"]
+    val_anno_file_path = kinetics_annotation_files_info["val"]
+    np.random.seed(random_seed)
+
+    train_anno_df = pd.read_csv(train_anno_file_path)
+    train_selected_samples_df = train_anno_df.iloc[:num_samples]
+
+    train_selected_classes = train_selected_samples_df["label"].tolist()
+
+    # select from test/val anno files based on the train classes
+    def select_anchored_samples(src_anno_df, anchor_classes):
+        selected_df = None
+        counted_classes = collections.Counter(anchor_classes)
+        for cls in list(counted_classes.keys()):
+            num_cls = counted_classes[cls]
+            cls_anno_df = src_anno_df[src_anno_df["label"] == cls]
+            len_cls_df = len(cls_anno_df)
+            selected_samples_idx = np.random.choice(
+                list(range(len_cls_df)), size=num_cls
+            )
+            cls_selected_df = cls_anno_df.iloc[selected_samples_idx]
+            if selected_df is None:
+                selected_df = cls_selected_df
+            else:
+                selected_df = pd.concat([selected_df, cls_selected_df])
+        return selected_df
+
+    def save_anno_df(anno_file_path, selected_anno_df):
+        anno_file_base_dir = os.path.dirname(anno_file_path)
+        anno_file_base_name, ext_type = os.path.basename(anno_file_path).split(".")
+        tiny_anno_file_path = os.path.join(
+            anno_file_base_dir, anno_file_base_name + "_tiny." + ext_type
+        )
+
+        if os.path.exists(tiny_anno_file_path):
+            logging.info(
+                "Annotation file for tiny data exists: %s, Using it directly",
+                tiny_anno_file_path,
+            )
+        else:
+            selected_anno_df.to_csv(path_or_buf=tiny_anno_file_path, index=False)
+
+    test_anno_df = pd.read_csv(test_anno_file_path)
+    test_selected_df = test_anno_df.iloc[:num_samples]
+
+    val_anno_df = pd.read_csv(val_anno_file_path)
+    val_selected_df = select_anchored_samples(
+        src_anno_df=val_anno_df, anchor_classes=train_selected_classes
+    )
+
+    save_anno_df(
+        anno_file_path=train_anno_file_path, selected_anno_df=train_selected_samples_df
+    )
+    save_anno_df(anno_file_path=test_anno_file_path, selected_anno_df=test_selected_df)
+    save_anno_df(anno_file_path=val_anno_file_path, selected_anno_df=val_selected_df)

plato/datasources/datalib/video_transform.py

@@ -0,0 +1,79 @@
+"""
+The transformers for the video classification task
+"""
+
+import torch
+import torch.nn as nn
+from torchvision.transforms import transforms
+
+
+class ConvertBHWCtoBCHW(nn.Module):
+    """Convert tensor from (B, H, W, C) to (B, C, H, W)"""
+
+    def forward(self, vid: torch.Tensor) -> torch.Tensor:
+        """Change the order of tensor dims"""
+        return vid.permute(0, 3, 1, 2)
+
+
+class ConvertBCHWtoCBHW(nn.Module):
+    """Convert tensor from (B, C, H, W) to (C, B, H, W)"""
+
+    def forward(self, vid: torch.Tensor) -> torch.Tensor:
+        """Change the order of tensor dims"""
+        return vid.permute(1, 0, 2, 3)
+
+
+class VideoClassificationTrainTransformer:
+    """The transformer to process the data for video classification model training"""
+
+    def __init__(
+        self,
+        resize_size,
+        crop_size,
+        mean=(0.43216, 0.394666, 0.37645),
+        std=(0.22803, 0.22145, 0.216989),
+        hflip_prob=0.5,
+    ):
+        trans = [
+            ConvertBHWCtoBCHW(),
+            transforms.ConvertImageDtype(torch.float32),
+            transforms.Resize(resize_size),
+        ]
+        if hflip_prob > 0:
+            trans.append(transforms.RandomHorizontalFlip(hflip_prob))
+        trans.extend(
+            [
+                transforms.Normalize(mean=mean, std=std),
+                transforms.RandomCrop(crop_size),
+                ConvertBCHWtoCBHW(),
+            ]
+        )
+        self.transforms = transforms.Compose(trans)
+
+    def __call__(self, video_data):
+        return self.transforms(video_data)
+
+
+class VideoClassificationEvalTransformer:
+    """The transformer to process the data for video classification model evaluation"""
+
+    def __init__(
+        self,
+        resize_size,
+        crop_size,
+        mean=(0.43216, 0.394666, 0.37645),
+        std=(0.22803, 0.22145, 0.216989),
+    ):
+        self.transforms = transforms.Compose(
+            [
+                ConvertBHWCtoBCHW(),
+                transforms.ConvertImageDtype(torch.float32),
+                transforms.Resize(resize_size),
+                transforms.Normalize(mean=mean, std=std),
+                transforms.CenterCrop(crop_size),
+                ConvertBCHWtoCBHW(),
+            ]
+        )
+
+    def __call__(self, video_data):
+        return self.transforms(video_data)

plato/datasources/emnist.py

@@ -0,0 +1,64 @@
+"""
+The Extended MNIST (EMNIST) dataset from the torchvision package.
+"""
+
+from torchvision import datasets, transforms
+
+from plato.config import Config
+from plato.datasources import base
+
+
+class DataSource(base.DataSource):
+    """The EMNIST dataset."""
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        _path = Config().params["data_path"]
+
+        train_transform = (
+            kwargs["train_transform"]
+            if "train_transform" in kwargs
+            else (
+                transforms.Compose(
+                    [
+                        transforms.RandomHorizontalFlip(),
+                        transforms.RandomAffine(
+                            degrees=10, translate=(0.2, 0.2), scale=(0.8, 1.2)
+                        ),
+                        transforms.ToTensor(),
+                        transforms.Normalize(mean=[0.5], std=[0.5]),
+                    ]
+                )
+            )
+        )
+
+        test_transform = (
+            kwargs["test_transform"]
+            if "test_transform" in kwargs
+            else (
+                transforms.Compose(
+                    [transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5])]
+                )
+            )
+        )
+
+        self.trainset = datasets.EMNIST(
+            root=_path,
+            split="balanced",
+            train=True,
+            download=True,
+            transform=train_transform,
+        )
+        self.testset = datasets.EMNIST(
+            root=_path,
+            split="balanced",
+            train=False,
+            download=True,
+            transform=test_transform,
+        )
+
+    def num_train_examples(self):
+        return 112800
+
+    def num_test_examples(self):
+        return 18800

plato/datasources/fashion_mnist.py

@@ -0,0 +1,41 @@
+"""
+The FashionMNIST dataset from the torchvision package.
+"""
+
+from torchvision import datasets, transforms
+
+from plato.config import Config
+from plato.datasources import base
+
+
+class DataSource(base.DataSource):
+    """The FashionMNIST dataset."""
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        _path = Config().params["data_path"]
+
+        train_transform = (
+            kwargs["train_transform"]
+            if "train_transform" in kwargs
+            else (
+                transforms.Compose(
+                    [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+                )
+            )
+        )
+        test_transform = train_transform
+
+        self.trainset = datasets.FashionMNIST(
+            root=_path, train=True, download=True, transform=train_transform
+        )
+
+        self.testset = datasets.FashionMNIST(
+            root=_path, train=False, download=True, transform=test_transform
+        )
+
+    def num_train_examples(self):
+        return 60000
+
+    def num_test_examples(self):
+        return 10000

plato/datasources/feature.py

@@ -0,0 +1,24 @@
+"""
+The feature dataset server received from clients.
+"""
+
+from itertools import chain
+from plato.datasources import base
+
+
+class DataSource(base.DataSource):
+    """The feature dataset."""
+
+    def __init__(self, features, **kwargs):
+        super().__init__()
+
+        # Faster way to deep flatten a list of lists compared to list comprehension
+        self.feature_dataset = list(chain.from_iterable(features))
+        self.trainset = self.feature_dataset
+        self.testset = []
+
+    def __len__(self):
+        return len(self.trainset)
+
+    def __getitem__(self, item):
+        return self.trainset[item]

plato/datasources/feature_dataset.py

@@ -0,0 +1,15 @@
+import torch
+
+
+class FeatureDataset(torch.utils.data.Dataset):
+    """Used to prepare a feature dataset for a DataLoader in PyTorch."""
+
+    def __init__(self, dataset):
+        self.dataset = dataset
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, item):
+        image, label = self.dataset[item]
+        return image, label

plato/datasources/femnist.py

@@ -0,0 +1,141 @@
+"""
+The Federated EMNIST dataset.
+
+The Federated EMNIST dataset originates from the EMNIST dataset, which contains
+817851 images, each of which is a 28x28 greyscale image in 1 out of 62 classes.
+The difference between the Federated EMNIST dataset and its original counterpart
+is that this dataset is already partitioned by the client ID, using the data
+provider IDs included in the original EMNIST dataset. As a result of this
+partitioning, there are 3597 clients in total, each of which has 227.37 images
+on average (std is 88.84). For each client, 90% data samples are used for
+training, while the remaining samples are used for testing.
+
+Reference:
+
+G. Cohen, S. Afshar, J. Tapson, and A. Van Schaik, "EMNIST: Extending MNIST to
+handwritten letters," in the 2017 International Joint Conference on Neural
+Networks (IJCNN).
+
+"""
+
+import json
+import logging
+import os
+
+import numpy as np
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+from plato.config import Config
+from plato.datasources import base
+
+
+class CustomDictDataset(Dataset):
+    """Custom dataset from a dictionary with support of transforms."""
+
+    def __init__(self, loaded_data, transform=None):
+        """Initializing the custom dataset."""
+        super().__init__()
+        self.loaded_data = loaded_data
+        self.transform = transform
+
+    def __getitem__(self, index):
+        sample = self.loaded_data["x"][index]
+        target = self.loaded_data["y"][index]
+
+        if self.transform:
+            sample = self.transform(sample)
+
+        return sample, target
+
+    def __len__(self):
+        return len(self.loaded_data["y"])
+
+
+class ReshapeListTransform:
+    """The transform that reshapes an image."""
+
+    def __init__(self, new_shape):
+        self.new_shape = new_shape
+
+    def __call__(self, img):
+        return np.array(img, dtype=np.float32).reshape(self.new_shape)
+
+
+class DataSource(base.DataSource):
+    """The FEMNIST dataset."""
+
+    def __init__(self, client_id=0, **kwargs):
+        super().__init__()
+        self.trainset = None
+        self.testset = None
+
+        root_path = os.path.join(
+            Config().params["data_path"], "FEMNIST", "packaged_data"
+        )
+        if client_id == 0:
+            # If we are on the federated learning server
+            data_dir = os.path.join(root_path, "test")
+            data_url = (
+                "http://iqua.ece.toronto.edu/baochun/FEMNIST/test/"
+                + str(client_id)
+                + ".zip"
+            )
+        else:
+            data_dir = os.path.join(root_path, "train")
+            data_url = (
+                "http://iqua.ece.toronto.edu/baochun/FEMNIST/train/"
+                + str(client_id)
+                + ".zip"
+            )
+
+        if not os.path.exists(os.path.join(data_dir, str(client_id))):
+            logging.info(
+                "Downloading the Federated EMNIST dataset "
+                "with the client datasets pre-partitioned. This may take a while.",
+            )
+            self.download(url=data_url, data_path=data_dir)
+
+        loaded_data = DataSource.read_data(
+            file_path=os.path.join(data_dir, str(client_id), "data.json")
+        )
+
+        train_transform = (
+            kwargs["train_transform"]
+            if "train_transform" in kwargs
+            else (
+                transforms.Compose(
+                    [
+                        ReshapeListTransform((28, 28, 1)),
+                        transforms.ToPILImage(),
+                        transforms.RandomCrop(
+                            28, padding=2, padding_mode="constant", fill=1.0
+                        ),
+                        transforms.RandomResizedCrop(
+                            28, scale=(0.8, 1.2), ratio=(4.0 / 5.0, 5.0 / 4.0)
+                        ),
+                        transforms.RandomRotation(5, fill=1.0),
+                        transforms.ToTensor(),
+                        transforms.Normalize(0.9637, 0.1597),
+                    ]
+                )
+            )
+        )
+
+        dataset = CustomDictDataset(loaded_data=loaded_data, transform=train_transform)
+
+        self.testset = dataset
+        self.trainset = dataset
+
+    @staticmethod
+    def read_data(file_path):
+        """Reading the dataset specific to a client_id."""
+        with open(file_path, "r", encoding="utf-8") as fin:
+            loaded_data = json.load(fin)
+        return loaded_data
+
+    def num_train_examples(self):
+        return len(self.trainset)
+
+    def num_test_examples(self):
+        return len(self.testset)