MobilePose-Pi/dataset_factory.py at master · savageyusuff/MobilePose-Pi · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
from dataloader import *

ROOT_DIR = "./pose_dataset/mpii"  # root dir to the dataset

DEBUG_MODE = False

def get_transform(modeltype, input_size):
    """
    :param modeltype: "resnet" / "mobilenet"
    :param input_size:
    :return:
    """
    if modeltype == "resnet":
        return Rescale((input_size, input_size))
    elif modeltype == "mobilenet":
        return Wrap((input_size, input_size))
    elif modeltype == "shufflenet":
		return Wrap((input_size, input_size))
    else:
        raise ValueError("modeltype is not wrong")


class DatasetFactory:
    def __init__(self):
        pass

    @staticmethod
    def get_train_dataset(modeltype, input_size, debug=DEBUG_MODE):
        """
        :param modeltype: "resnet" / "mobilenet"
        :return: type: PoseDataset
        Example:
        DataFactory.get_train_dataset("resnet", 224)
        In debug mode, it will return a small dataset
        """
        csv_name = "train_joints.csv"  #target
        if debug:
            csv_name = "train_joints-500.csv"

        return PoseDataset(csv_file=os.path.join(ROOT_DIR, csv_name),
                           transform=transforms.Compose([
                               Augmentation(),
                               # Rescale((inputsize, inputsize)),  # for resnet18
                               # # Wrap((inputsize,inputsize)),# for mobilenetv2
                               get_transform(modeltype, input_size),
                               Expansion(),
                               # Guass(),
                               ToTensor()
                           ]))

    @staticmethod
    def get_test_dataset(modeltype, input_size, debug=DEBUG_MODE):
        """
        :param modeltype: resnet / mobilenet
        :return: type: PoseDataset
        Example:
        DataFactory.get_test_dataset("resnet", 224)
        In debug mode, it will return a small dataset
        """
        csv_name = "test_joints.csv"
        if debug:
            csv_name = "test_joints-500.csv"
        return PoseDataset(
            csv_file=os.path.join(ROOT_DIR, csv_name),
            transform=transforms.Compose([
                # Rescale((inputsize, inputsize)),  # for resnet18
                # # Wrap((inputsize, inputsize)),# for mobilenetv2
                get_transform(modeltype, input_size),
                Expansion(),
                # Guass(),
                ToTensor()
            ]))