Update app.py

app.py

@@ -4,6 +4,8 @@ import numpy as np
 from tensorflow.keras.applications.efficientnet import preprocess_input
 from PIL import Image
 from tensorflow.keras.layers import InputLayer
+# Import các lớp Augmentation mà mô hình có thể đã lưu
+from tensorflow.keras.layers import RandomFlip, RandomRotation, RandomZoom, RandomContrast
 
 # 1. Định nghĩa hằng số
 MODEL_PATH = "best_model.h5" # Tên file mô hình đã upload (ĐÃ SỬA: Đồng bộ với tên file trên GitHub)
@@ -11,21 +13,32 @@ IMG_SIZE = (224, 224)
 CLASS_NAMES = ['bad', 'good', 'very_good'] 
 
 # =================================================================
-# KHẮC PHỤC LỖI DESERIALIZATION (batch_shape)
-# Nếu mô hình được lưu bằng Keras cũ, ta cần định nghĩa lại InputLayer 
-# để chấp nhận tham số batch_shape khi tải.
+# KHẮC PHỤC LỖI DESERIALIZATION CHO CÁC LỚP AUGMENTATION VÀ INPUT LAYER
+# Mục đích: Bỏ qua tham số 'batch_shape' và 'data_format' không tương thích.
+
 class FixedInputLayer(InputLayer):
-    """
-    Tạo InputLayer cố định để xử lý lỗi "Unrecognized keyword arguments: ['batch_shape']"
-    khi tải mô hình đã lưu bằng TF/Keras cũ hơn.
-    """
+    """Xử lý lỗi batch_shape."""
     def __init__(self, **kwargs):
         if 'batch_shape' in kwargs:
-            # Chuyển batch_shape thành input_shape
             kwargs['input_shape'] = kwargs['batch_shape'][1:]
             del kwargs['batch_shape']
         super().__init__(**kwargs)
 
+def fix_augmentation_layer(LayerClass):
+    """Tạo lớp cố định để loại bỏ tham số 'data_format'."""
+    class FixedLayer(LayerClass):
+        def __init__(self, **kwargs):
+            if 'data_format' in kwargs:
+                del kwargs['data_format']
+            super().__init__(**kwargs)
+    return FixedLayer
+
+# Tạo các lớp cố định cho Augmentation
+FixedRandomFlip = fix_augmentation_layer(RandomFlip)
+FixedRandomRotation = fix_augmentation_layer(RandomRotation)
+FixedRandomZoom = fix_augmentation_layer(RandomZoom)
+FixedRandomContrast = fix_augmentation_layer(RandomContrast)
+
 # =================================================================
 
 # 2. Tải mô hình
@@ -38,7 +51,12 @@ try:
     model = tf.keras.models.load_model(
         MODEL_PATH, 
         custom_objects={
-            'InputLayer': FixedInputLayer # Sử dụng lớp FixedInputLayer để thay thế
+            'InputLayer': FixedInputLayer,        # Fix InputLayer
+            'RandomFlip': FixedRandomFlip,        # Fix RandomFlip
+            'RandomRotation': FixedRandomRotation, # Fix RandomRotation
+            'RandomZoom': FixedRandomZoom,         # Fix RandomZoom
+            'RandomContrast': FixedRandomContrast # Fix RandomContrast
+            # Bổ sung các lớp Augmentation khác nếu cần
         }
     )
     print("✅ Mô hình đã được tải thành công.")