LookingForSeagrassSemanticSegmentation/data.py at master · EnviewFulda/LookingForSeagrassSemanticSegmentation · 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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
# dataset provider for classifications

import os
import cv2
import numpy as np
from PIL import Image
import sys
import json
from time import sleep
from IPython.display import display


class Data:

    # loads the all file names or the serialized numpy object
    def loadDataset(self):
        self.pathImages = {
            "train": self.config["path"],
            "trainLabel": self.config["path"],
            "test" : self.config["path"],
            "testLabel": self.config["path"],
            "validation": self.config["path"],
            "validationLabel": self.config["path"]
        }

        if self.config["preProcessedPath"] != "":
            self.pathImages["train"] += self.config["preProcessedPath"]+self.config["images"]
            self.pathImages["trainLabel"] += self.config["preProcessedPath"]+self.config["labels"]
        else:
            self.pathImages["train"] += self.config["images"]
            self.pathImages["trainLabel"] += self.config["labels"]
            self.pathImages["test"] += self.config["images"]
            self.pathImages["testLabel"] += self.config["labels"]
            self.pathImages["validation"] += self.config["images"]
            self.pathImages["validationLabel"] += self.config["labels"]

        if not self.config["serializedObject"]:

            if self.config["name"] == "Seagrass":
                jsonData = json.load(open(self.config["path"]+"train.json"))
                trainData = list(map(lambda i:os.path.basename(i["image"]) if i["depth"] <= float(self.config["depth"]) else None, jsonData))
                labelData = list(map(lambda i:os.path.basename(i["ground-truth"]) if i["depth"] <= float(self.config["depth"]) else None, jsonData))

                jsonData = json.load(open(self.config["path"]+"test.json"))
                testData = list(map(lambda i:os.path.basename(i["image"]) if i["depth"] <= float(self.config["depth"]) else None, jsonData))
                testLabelData = list(map(lambda i:os.path.basename(i["ground-truth"]) if i["depth"] <= float(self.config["depth"]) else None, jsonData))

                jsonData = json.load(open(self.config["path"]+"validate.json"))
                validateData = list(map(lambda i:os.path.basename(i["image"]) if i["depth"] <= float(self.config["depth"]) else None, jsonData))
                validateLabelData = list(map(lambda i:os.path.basename(i["ground-truth"]) if i["depth"] <= float(self.config["depth"]) else None, jsonData))

                self.imageData = {
                    "train": list(filter(lambda i:i != None, trainData)),
                    "trainLabel": list(filter(lambda i:i != None, labelData)),
                    "test": list(filter(lambda i:i != None, testData)),
                    "testLabel": list(filter(lambda i:i != None, testLabelData)),
                    "validation": list(filter(lambda i:i != None, validateData)),
                    "validationLabel": list(filter(lambda i:i != None, validateLabelData))
                }


            else:
                # sort data because os.listdir selects files in arbitrary order
                trainDataFiles = os.listdir(self.pathImages["train"])
                trainLabelDataFiles = os.listdir(self.pathImages["trainLabel"])
                trainDataFiles.sort()
                trainLabelDataFiles.sort()


                trainElements = int(self.config["trainSize"]*self.config["size"])
                testElements = int(self.config["testSize"]*self.config["size"])

                self.imageData = {
                    "train": trainDataFiles[:trainElements],
                    "trainLabel": trainLabelDataFiles[:trainElements],
                    "test": trainDataFiles[trainElements:trainElements+testElements],
                    "testLabel": trainLabelDataFiles[trainElements:trainElements+testElements],
                    "validation": trainDataFiles[trainElements+testElements if testElements > 0 else trainElements:],
                    "validationLabel": trainLabelDataFiles[trainElements+testElements if testElements > 0 else trainElements:],
                }


            self.config["trainSize"] = len(self.imageData["train"])
            self.config["testSize"] = len(self.imageData["test"])
            self.config["validationSize"] = len(self.imageData["validation"])

            print("trainSize: ", self.config["trainSize"], " Testsize: ", self.config["testSize"], "Validationsize: ", self.config["validationSize"])

        else:

            self.imageData = np.load(self.config["path"]+self.config["fileName"]+".npy")
            self.config["trainSize"] = len(self.imageData.item().get("train"))
            self.config["testSize"] = len(self.imageData.item().get("test"))
            self.config["validationSize"] = len(self.imageData.item().get("validation"))
            print("Finished loading dataset...")


    # string configPath - path of the json file which describes the dataset
    def __init__(self, configPath):

        try:
            self.config = json.load(open(configPath))
        except:
            raise "Wrong path for data config file given!"


        self.loadDataset()

    # gets a value from the config file with its given name
    def getConfig(self, name):
        return self.config[name]


    def getImageTuple(self, imageFilename, labelFilename):

        img = cv2.imread(self.pathImages["train"]+imageFilename.decode())
        labelImg = cv2.imread(self.pathImages["trainLabel"]+labelFilename.decode())
        if self.config["downsize"]:
            img = cv2.resize(img, (self.config["x"], self.config["y"]), interpolation=cv2.INTER_AREA)
            labelImg = cv2.resize(labelImg, (self.config["x"], self.config["y"]), interpolation=cv2.INTER_AREA)

        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        labelImg = cv2.cvtColor(labelImg, cv2.COLOR_BGR2RGB)


        # exterminate conversion errors by opencv
        labelImg[(labelImg <= 127).all(-1)] = [0,0,0]
        labelImg[(labelImg >= 128).all(-1)] = [255,255,255]

        for rgbIdx, rgbV in enumerate(self.config["ClassToRGB"]):
            labelImg[(labelImg == rgbV).all(-1)] = rgbIdx


        labelImg = labelImg[:,:,0]
        img = ((img - img.mean()) / img.std()).astype(np.float32)

        return img, labelImg