DeepShield/main.py at main · TanmayT134/DeepShield · 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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
import os
import tensorflow as tf

# ==================================================
# TensorFlow Setup
# ==================================================

from src.utils.tensorflow_setup import setup_tensorflow

setup_tensorflow()

# ==================================================
# Configuration
# ==================================================

from config.config import *

# ==================================================
# Dataset
# ==================================================

from src.datasets.dataloader import create_dataset
from src.utils.dataset_statistics import print_dataset_statistics

print("\nLoading datasets...\n")

train_dataset = create_dataset(
    TRAIN_PATH,
    training=True
)

validation_dataset = create_dataset(
    VALIDATION_PATH,
    training=False
)

test_dataset = create_dataset(
    TEST_PATH,
    training=False
)

print("\nDataset Statistics\n")

print_dataset_statistics(
    train_dataset,
    "Train"
)

print_dataset_statistics(
    validation_dataset,
    "Validation"
)

print_dataset_statistics(
    test_dataset,
    "Test"
)

# ==================================================
# Model
# ==================================================

from src.models.cnn import DeepShieldCNN
from src.utils.save_model_summary import save_model_summary

CHECKPOINT_PATH = "models/checkpoints/checkpoint.keras"

if os.path.exists(CHECKPOINT_PATH):

    print("\nResuming from latest checkpoint...\n")

    model = tf.keras.models.load_model(
        CHECKPOINT_PATH
    )

elif os.path.exists(MODEL_PATH):

    print("\nLoading best model...\n")

    model = tf.keras.models.load_model(
        MODEL_PATH
    )

else:

    print("\nBuilding new model...\n")

    cnn = DeepShieldCNN()

    model = cnn.build()

save_model_summary(model)

# ==================================================
# Training
# ==================================================

from src.training.trainer import Trainer

trainer = Trainer(model)

trainer.compile()

history = trainer.train(

    train_dataset,

    validation_dataset

)

# ==================================================
# Save Training History
# ==================================================

from src.utils.save_history import save_history

save_history(history)

# ==================================================
# Plot Training Curves
# ==================================================

from src.utils.history_plotter import HistoryPlotter

plotter = HistoryPlotter(history)

plotter.plot()

# ==================================================
# Evaluation
# ==================================================

from src.evaluation.evaluate import Evaluator

evaluator = Evaluator(model)

evaluator.evaluate(test_dataset)

# ==================================================
# Save Final Model
# ==================================================

os.makedirs(
    "models/final",
    exist_ok=True
)

FINAL_MODEL_PATH = os.path.join(
    "models",
    "final",
    "final_cnn.keras"
)

model.save(
    FINAL_MODEL_PATH
)

print("\nFinal model saved successfully.")

print(f"Location : {FINAL_MODEL_PATH}")

print("\nDeepShield Pipeline Completed Successfully.")