model.py 3.09 KB
Newer Older
josh's avatar
josh committed
1
import numpy as np
2
3
import pickle
import os
josh's avatar
josh committed
4
5
6
7
8
9
10

import matplotlib.pyplot as plt

from pathlib import Path
from PIL import Image

import tensorflow as tf
11
from tensorflow import keras
josh's avatar
josh committed
12
13
14
15
16
17
18
19
20
from tensorflow.keras.layers import Input, Conv2D
from tensorflow.keras.models import Model
from tensorflow.keras.applications.vgg16 import VGG16

from sklearn.ensemble import RandomForestClassifier


class HDRModel:

21
22

    def __init__(self, shape=None, loadDir=None):
josh's avatar
josh committed
23
24
25
26
27
28
        self.featureModel = None
        self.rfClassifer  = None

        self.inputChannels  = 3
        self.outputChannels = 1

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
        self.__fmSaveDir  = "featureModel"
        self.__fmSaveFile = "fm.h5"
        self.__rfSaveDir  = "rfClassifer"
        self.__rfSaveFile = "classifier.sav"

        if loadDir is not None:
            self.load(loadDir)
        elif shape is not None:
            self.__initModels(shape)


    def save(self, saveDir):
        # Save tf model
        fmDir = saveDir.joinpath(self.__fmSaveDir)
        os.mkdir(fmDir)
        self.featureModel.save(str(fmDir.joinpath(self.__fmSaveFile)))
        # Save sklearn random forest
        classifierDir = saveDir.joinpath(self.__rfSaveDir)
        os.mkdir(classifierDir)
        pickle.dump(self.rfClassifer, open(classifierDir.joinpath(self.__rfSaveFile), "wb"))


    def load(self, modelDir):
        # Load tf model
        fmPath = modelDir.joinpath(self.__fmSaveDir).joinpath(self.__fmSaveFile)
        self.featureModel = keras.models.load_model(str(fmPath))
        self.featureModel.summary()

        # Load sklean model
        rfPath = modelDir.joinpath(self.__rfSaveDir).joinpath(self.__rfSaveFile)
        self.rfClassifer = pickle.load(open(rfPath, "rb"))
josh's avatar
josh committed
60
61
62
63
64
65
66


    def fit(self, xTrain, yTrain):
        # Get features from VGG16 built model
        features = self.__getFeatures(xTrain)

        # Reshape input for random forest
josh's avatar
test    
josh committed
67
        # TO DO: sklearn split data to train and test
josh's avatar
josh committed
68
69
70
71
72
73
74
75
76
77
78
79
80
        xForestTrain = features.reshape(-1, features.shape[-1])
        yForestTrain = yTrain.reshape(-1)

        self.rfClassifer.fit(xForestTrain, yForestTrain)


    def predict(self, xInput):
        '''
        Supports only a single input: shape=(x, y, z)
        '''
        # Extract features from input
        xExpand  = np.expand_dims(xInput, axis=0)
        features = self.__getFeatures(xExpand)
81
        
josh's avatar
josh committed
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
        # Reshape input for random forest
        forestInput = features.reshape(-1, features.shape[-1])
        forestOutput = self.rfClassifer.predict(forestInput)
        
        # Reshape output back to image shape
        outputShape = (xInput.shape[0], xInput.shape[1])
        finalOutput = forestOutput.reshape(outputShape)
        return finalOutput


    def __getFeatures(self, xInput):
        return self.featureModel.predict(xInput)
    

    def __initModels(self, shape):
        self.inputChannels = shape[-1]
        vgg = VGG16(weights="imagenet", include_top=False, input_shape=shape)
        self.featureModel = Model(inputs=vgg.input, outputs=vgg.get_layer("block1_conv2").output)
josh's avatar
josh committed
100
        self.rfClassifer  = RandomForestClassifier(n_estimators=50, random_state=42, verbose=1, n_jobs=-1)
josh's avatar
josh committed
101
102