teeth-segmentation/loader.py at main · mdciri/teeth-segmentation · 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

import numpy as np
from torch.utils.data import Dataset
from PIL import Image


class TuftsDataset(Dataset):
    """ Tufts Dataset
    """

    def __init__(
        self,
        data_list,
        masking = False,
        transform = None,
    ):
        super(TuftsDataset).__init__()

        self.data_list = data_list
        self.masking = masking
        self.transform = transform

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, id):

        if self.masking == True:
            msk = np.asarray(Image.open(self.data_list[id]["msk"]).convert("1"), dtype="float32")
            seg = np.asarray(Image.open(self.data_list[id]["seg"]).convert("1"), dtype="float32")
            img = np.asarray(Image.open(self.data_list[id]["img"]).convert("L"), dtype="float32") * msk
        else:
            seg = np.asarray(Image.open(self.data_list[id]["seg"]).convert("1"), dtype="float32")
            img = np.asarray(Image.open(self.data_list[id]["img"]).convert("L"), dtype="float32")

        data = {
            "img": np.expand_dims(img, axis=0),
            "seg": np.expand_dims(seg, axis=0)
        }

        if self.transform is not None:
            data = self.transform(data)

        return data