VideoAnomolyDetection/model.py at master · Colprit/VideoAnomolyDetection · 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
import keras
from keras.layers import (
    Conv2DTranspose,
    ConvLSTM2D,
    TimeDistributed,
    Conv2D,
)
from keras.models import Sequential, load_model
from keras.layers import LayerNormalization

from clip_generator import DataGenerator
from config import *

def get_model(train=True):
    """
    Parameters
    ----------
    reload_model : bool
        Load saved model or retrain it
    """

    if not train:
        return load_model(
            MODEL_PATH,
            custom_objects={'LayerNormalization': LayerNormalization}
        )

    training_generator = DataGenerator(
        DATASET_PATH,
        CLIP_LEN,
        STRIDE,
        DIM,
        BATCH_SIZE,
        N_CHANNELS,
        SHUFFLE
    )

    seq = Sequential()
    seq.add(TimeDistributed(
        Conv2D(16, (11,11), strides=4, padding="same"),
        batch_input_shape=(None, *DIM, N_CHANNELS)
    ))
    seq.add(LayerNormalization())
    seq.add(TimeDistributed(Conv2D(8, (8,8), strides=2, padding="same")))
    seq.add(LayerNormalization())
    ######
    seq.add(ConvLSTM2D(8, (3,3), padding="same", return_sequences=True))
    seq.add(LayerNormalization())
    seq.add(ConvLSTM2D(4, (3,3), padding="same", return_sequences=True))
    seq.add(LayerNormalization())
    seq.add(ConvLSTM2D(8, (3,3), padding="same", return_sequences=True))
    seq.add(LayerNormalization())
    ######
    seq.add(TimeDistributed(Conv2DTranspose(8, (8,8), strides=2, padding="same")))
    seq.add(LayerNormalization())
    seq.add(TimeDistributed(Conv2DTranspose(16, (11, 11), strides=4, padding="same")))
    seq.add(LayerNormalization())
    seq.add(TimeDistributed(Conv2D(1, (11,11), activation="sigmoid", padding="same")))

    print(seq.summary())

    seq.compile(
        loss='mse',
        optimizer=keras.optimizers.Adam(lr=1e-4, decay=1e-5, epsilon=1e-6)
    )
    seq.fit(
        x=training_generator,
        epochs=EPOCHS,
        verbose=True,
        workers=0,
        use_multiprocessing=False
    )
    seq.save(MODEL_PATH)

    return seq