Update app.py

app.py

@@ -1,139 +1,51 @@
 import gradio as gr
 import tensorflow as tf
 import numpy as np
-from tensorflow.keras.applications.efficientnet import preprocess_input
 from PIL import Image
-from tensorflow.keras.layers import InputLayer, RandomFlip, RandomRotation, RandomZoom, RandomContrast
-from tensorflow.keras.layers import RandomWidth, RandomHeight 
+from tensorflow.keras.applications.efficientnet import preprocess_input
 
-# 1. Định nghĩa hằng số
-MODEL_PATH = "best_model.h5"
+# Đường dẫn model SavedModel
+MODEL_PATH = "exported_model"
 IMG_SIZE = (224, 224)
-CLASS_NAMES = ['bad', 'good', 'very_good'] 
-
-# =================================================================
-# KHẮC PHỤC LỖI TƯƠNG THÍCH PHIÊN BẢN (KERAS 3 -> KERAS 2)
-# =================================================================
-
-# 1. Mock DTypePolicy (Xử lý lỗi: Unknown dtype policy, Attribute name & compute_dtype)
-class MockDTypePolicy:
-    """
-    Lớp giả lập để thay thế DTypePolicy của Keras 3.
-    Giúp tránh lỗi deserialization khi chạy trên môi trường cũ.
-    """
-    def __init__(self, **kwargs):
-        # SỬA LỖI: Thêm đầy đủ các thuộc tính mà Keras 3 yêu cầu
-        self.name = kwargs.get("name", "float32")
-        self.compute_dtype = kwargs.get("compute_dtype", "float32")
-        self.variable_dtype = kwargs.get("variable_dtype", "float32")
-    
-    @classmethod
-    def from_config(cls, config):
-        return cls(**config)
-    
-    def get_config(self):
-        return {
-            "name": self.name,
-            "compute_dtype": self.compute_dtype,
-            "variable_dtype": self.variable_dtype
-        }
-
-# 2. Xử lý InputLayer (Xử lý lỗi: batch_shape)
-class FixedInputLayer(InputLayer):
-    def __init__(self, **kwargs):
-        if 'batch_shape' in kwargs:
-            kwargs['input_shape'] = kwargs['batch_shape'][1:]
-            del kwargs['batch_shape']
-        # Xóa dtype nếu nó là dạng dictionary (config của Keras 3)
-        if 'dtype' in kwargs and isinstance(kwargs['dtype'], dict):
-            del kwargs['dtype']
-        super().__init__(**kwargs)
+CLASS_NAMES = ['bad', 'good', 'very_good']
 
-# 3. Xử lý Augmentation Layers (Xử lý lỗi: data_format, dtype & value_range)
-def fix_augmentation_layer(LayerClass):
-    class FixedLayer(LayerClass):
-        def __init__(self, **kwargs):
-            # Danh sách các tham số gây lỗi tương thích giữa Keras 3 và 2
-            ignore_keys = ['data_format', 'dtype', 'value_range']
-            
-            for key in ignore_keys:
-                if key in kwargs:
-                    del kwargs[key]
-                
-            super().__init__(**kwargs)
-    return FixedLayer
+# Load model
+model = tf.saved_model.load(MODEL_PATH)
 
-# 4. Đăng ký tất cả Custom Objects
-CUSTOM_OBJECTS = {
-    'InputLayer': FixedInputLayer,
-    # Đăng ký lớp giả lập DTypePolicy
-    'DTypePolicy': MockDTypePolicy, 
-    # Augmentation Layers đã được vá lỗi
-    'RandomFlip': fix_augmentation_layer(RandomFlip),
-    'RandomRotation': fix_augmentation_layer(RandomRotation),
-    'RandomZoom': fix_augmentation_layer(RandomZoom),
-    'RandomContrast': fix_augmentation_layer(RandomContrast),
-    'RandomWidth': fix_augmentation_layer(RandomWidth),
-    'RandomHeight': fix_augmentation_layer(RandomHeight)
-}
-
-# =================================================================
-
-# 2. Tải mô hình
-model = None
-try:
-    # Tắt log TensorFlow
-    tf.get_logger().setLevel('ERROR') 
-    
-    # Tải mô hình với danh sách custom objects đầy đủ
-    model = tf.keras.models.load_model(
-        MODEL_PATH, 
-        custom_objects=CUSTOM_OBJECTS
-    )
-    print("✅ Mô hình đã được tải thành công.")
-except Exception as e:
-    print(f"❌ Lỗi tải mô hình: {e}")
-    model = None
+infer = model.signatures["serving_default"]
 
 def predict_guava_quality(img_input):
-    if model is None:
-        return "❌ Lỗi: Không thể tải mô hình.", 0.0
-
     if img_input is None:
-        return "❌ Vui lòng tải ảnh lên.", 0.0
+        return "❌ Vui lòng tải ảnh", 0.0
+
+    # Convert image
+    img = Image.fromarray(img_input).convert("RGB")
+    img = img.resize(IMG_SIZE)
+
+    arr = np.array(img).astype("float32")
+    arr = preprocess_input(arr)
+    arr = np.expand_dims(arr, axis=0)
 
-    try:
-        # Chuyển đổi ảnh
-        img_pil = Image.fromarray(img_input).convert("RGB")
-        img_resized = img_pil.resize(IMG_SIZE)
-        
-        # Preprocess
-        arr = np.array(img_resized).astype("float32")
-        arr = preprocess_input(arr)
-        arr = np.expand_dims(arr, 0)
+    # TensorFlow serving
+    outputs = infer(tf.constant(arr))
+    preds = list(outputs.values())[0].numpy()[0]
 
-        # Dự đoán
-        preds = model.predict(arr)[0]
-        idx = np.argmax(preds)
-        confidence = preds[idx]
-        label = CLASS_NAMES[idx]
+    idx = np.argmax(preds)
+    confidence = preds[idx]
+    label = CLASS_NAMES[idx]
 
-        return f"✅ Kết quả: {label}", float(confidence)
-    except Exception as e:
-        return f"❌ Lỗi xử lý ảnh: {str(e)}", 0.0
+    return f"✅ Kết quả: {label}", float(confidence)
 
-# 3. Giao diện Gradio
 demo = gr.Interface(
     fn=predict_guava_quality,
     inputs=gr.Image(type="numpy", label="Tải ảnh Quả Ổi"),
     outputs=[
         gr.Textbox(label="Dự đoán"),
-        gr.Number(label="Độ tin cậy (%)", precision=2)
+        gr.Number(label="Độ tin cậy (%)", precision=4)
     ],
-    title="Phân loại Chất lượng Quả Ổi (EfficientNetB0)",
-    description="Tải lên ảnh quả ổi để phân loại thành: Hàng xuất khẩu (very_good), Hàng nội địa (good), hoặc Loại bỏ (bad)."
+    title="Phân loại chất lượng Ổi",
+    description="Model EfficientNetB0 | very_good / good / bad"
 )
 
-# 4. Chạy App
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()