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

import matplotlib.pyplot as plt

from pathlib import Path
from PIL import Image

import tensorflow as tf
12
from tensorflow import keras
josh's avatar
josh committed
13 14 15 16 17
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
josh's avatar
josh committed
18
from sklearn.model_selection import train_test_split
josh's avatar
josh committed
19 20 21 22


class HDRModel:

23 24

    def __init__(self, shape=None, loadDir=None):
josh's avatar
josh committed
25 26 27
        self.featureModel   = None
        self.rfClassifer    = None
        self.trainTestSplit = 0.2
josh's avatar
josh committed
28 29 30 31

        self.inputChannels  = 3
        self.outputChannels = 1

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
        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
63 64


josh's avatar
josh committed
65
    def fit(self, xTrain, yTrain, xTest=None, yTest=None):
josh's avatar
josh committed
66 67 68 69
        # Get features from VGG16 built model
        features = self.__getFeatures(xTrain)

        # Reshape input for random forest
josh's avatar
test  
josh committed
70
        # TO DO: sklearn split data to train and test
josh's avatar
josh committed
71 72
        xForest = features.reshape(-1, features.shape[-1])
        yForest = yTrain.reshape(-1)
josh's avatar
josh committed
73

josh's avatar
josh committed
74 75
        dataset = pd.DataFrame(xForest)
        dataset["Label"] = yForest
josh's avatar
josh committed
76

josh's avatar
josh committed
77 78 79
        print("BEFORE DROP VALUE 0")
        print(dataset["Label"].unique())
        print(dataset["Label"].value_counts())
josh's avatar
josh committed
80

josh's avatar
josh committed
81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104
        #dataset = dataset[dataset["Label"] != 0]

        print("AFTER DROP VALUE 0")
        print(dataset["Label"].unique())
        print(dataset["Label"].value_counts())

        newX = dataset.drop(labels=["Label"], axis=1)
        newY = dataset["Label"]

        #xTrain, xTest, yTrain, yTest = train_test_split(xForest, yForest, test_size=self.trainTestSplit)
        self.rfClassifer.fit(newX, newY)
        print("NEW DATA")
        if xTest is not None and yTest is not None:
            return self.rfClassifer.score(xTest, yTest)
        else:
            return -1


    def score(self, X, Y):
        features = self.__getFeatures(X)
        xForest = features.reshape(-1, features.shape[-1])
        yForest = Y.reshape(-1)

        return self.rfClassifer.score(xForest, yForest)
josh's avatar
josh committed
105 106 107 108 109 110 111 112 113


    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)
114
        
josh's avatar
josh committed
115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132
        # 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
133 134
        self.featureModel.summary()
        self.rfClassifer  = RandomForestClassifier(n_estimators=50, random_state=42, verbose=1, n_jobs=3)
josh's avatar
josh committed
135 136