Update tensorflow

app.py

@@ -13,6 +13,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torchvision import transforms
+from tensorflow.keras import layers, models
 
 app = Flask(__name__)
 
@@ -124,15 +125,43 @@ def load_pytorch():
     return _pytorch_model
 
 
+
+def build_cnn_tf(num_classes: int = 6, input_shape: tuple = (228, 228, 3)):
+    return models.Sequential([
+        layers.Input(shape=input_shape),
+
+        layers.Conv2D(32, (5, 5), activation="relu"),
+        layers.MaxPooling2D(2, 2),
+
+        layers.Conv2D(32, (5, 5), activation="relu"),
+        layers.MaxPooling2D(2, 2),
+
+        layers.Conv2D(32, (3, 3), activation="relu"),
+        layers.MaxPooling2D(2, 2),
+
+        layers.Conv2D(64, (3, 3), activation="relu"),
+        layers.MaxPooling2D(2, 2),
+
+        layers.Conv2D(64, (3, 3), activation="relu"),
+        layers.MaxPooling2D(2, 2),
+
+        layers.Flatten(),
+        layers.Dense(1024, activation="relu"),
+        layers.Dropout(0.20),
+        layers.Dense(124, activation="relu"),
+        layers.Dropout(0.20),
+        layers.Dense(num_classes, activation="softmax"),
+    ])
+
+
 def load_tensorflow():
     global _tf_model
     if _tf_model is not None:
         return _tf_model
 
-    import tensorflow as tf
-
     model = build_cnn_tf(num_classes=6, input_shape=(228, 228, 3))
     model.load_weights("parfait_model.keras")
+
     _tf_model = model
     return _tf_model
     
@@ -205,4 +234,4 @@ def predict():
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 5000))
-    app.run(host="0.0.0.0", port=port, debug=False)
+    app.run(host="0.0.0.0", port=port, debug=False)