Update app.py

app.py

@@ -2,10 +2,8 @@ import os
 import cv2
 import numpy as np
 import pickle
-import tempfile
-import traceback
-from pathlib import Path
 from PIL import Image
+import matplotlib.pyplot as plt
 import tensorflow as tf
 from tensorflow.keras import layers
 from tensorflow.keras.models import load_model, Model
@@ -13,24 +11,20 @@ from tensorflow.keras.applications import EfficientNetV2B0
 from tensorflow.keras.applications.efficientnet import preprocess_input as efficientnet_preprocess
 from tensorflow.keras.preprocessing.sequence import pad_sequences
 from tensorflow.keras.preprocessing.image import img_to_array
-from huggingface_hub import hf_hub_download
-import gradio as gr
+from tqdm import tqdm
+import random
+from tensorflow.keras.preprocessing.sequence import pad_sequences
 
+import tempfile
+import traceback
+from pathlib import Path
+from huggingface_hub import hf_hub_download
 
-MODEL_REPO = "slyviee/img_cap"
+import gradio as gr
+from PIL import Image
+import pickle
 
-# Khởi tạo tài nguyên toàn cục khi app start
-model_path = hf_hub_download(repo_id=MODEL_REPO, filename="best_model.keras")
-tokenizer_path = hf_hub_download(repo_id=MODEL_REPO, filename="tokenizer.pkl")
-config_path = hf_hub_download(repo_id=MODEL_REPO, filename="model_config.pkl")
 
-model = None
-tokenizer = None
-max_length = None
-vocab_size = None
-extractor = None
-ready = False
-startup_error = ""
 
 # -----------------------------
 # Custom attention layers
@@ -64,6 +58,10 @@ class ChannelAttention(layers.Layer):
         config.update({'ratio': self.ratio})
         return config
 
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
 
 
 class SpatialAttention(layers.Layer):
@@ -81,6 +79,11 @@ class SpatialAttention(layers.Layer):
     def get_config(self):
         return super(SpatialAttention, self).get_config()
 
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
+
 
 # -----------------------------
 # Load model + tokenizer
@@ -92,6 +95,7 @@ def load_caption_model(model_path):
         'SpatialAttention': SpatialAttention
     }
     model = load_model(model_path, custom_objects=custom_objects)
+    print("✅ Đã load model thành công!")
     return model
 
 
@@ -115,6 +119,7 @@ def load_feature_extractor():
 def extract_features_from_image(image_path, extractor):
     image = cv2.imread(image_path)
     if image is None:
+        print(f"❌ Không đọc được ảnh: {image_path}")
         return None
     image = cv2.resize(image, (224, 224))
     image = img_to_array(image)
@@ -143,9 +148,24 @@ def generate_caption(model, tokenizer, image_features, max_length):
 
 
 # -----------------------------
-# App initialization
+# Chạy test
 # -----------------------------
 
+MODEL_REPO = "slyviee/img_cap"
+
+# Khởi tạo tài nguyên toàn cục khi app start
+model_path = hf_hub_download(repo_id=MODEL_REPO, filename="best_model.keras")
+tokenizer_path = hf_hub_download(repo_id=MODEL_REPO, filename="tokenizer.pkl")
+config_path = hf_hub_download(repo_id=MODEL_REPO, filename="model_config.pkl")
+
+model = None
+tokenizer = None
+max_length = None
+vocab_size = None
+extractor = None
+ready = False
+startup_error = ""
+
 
 def _startup():
     global model, tokenizer, max_length, vocab_size, extractor, ready, startup_error
@@ -157,15 +177,23 @@ def _startup():
             ready = False
             return
 
+        print("🔄 Đang tải model...")
         model = load_caption_model(model_path)
+        print("✅ Model đã được tải.")
+
+        print("🔄 Đang tải tokenizer và config...")
         tokenizer, max_length, vocab_size = load_tokenizer_and_config(tokenizer_path, config_path)
+        print("✅ Tokenizer và config đã được tải.")
+
+        print("🔄 Đang tải feature extractor...")
         extractor = load_feature_extractor()
+        print("✅ Feature extractor đã được tải.")
+
         ready = True
     except Exception as e:
         startup_error = f"Khởi tạo lỗi: {e}\n{traceback.format_exc()}"
         ready = False
 
-_startup()
 
 def predict(pil_image: Image.Image):
     if not ready:
@@ -185,7 +213,7 @@ def predict(pil_image: Image.Image):
         caption = generate_caption(model, tokenizer, features, max_length)
         return caption
     except Exception as e:
-        return f"Lỗi suy luận: {e}"
+        return f"Lỗi trong quá trình dự đoán: {e}\n{traceback.format_exc()}"
 
 DESCRIPTION = (
     "Upload ảnh và nhận caption sinh ra bởi mô hình. "
@@ -200,5 +228,6 @@ demo = gr.Interface(
     allow_flagging="never",
 )
 
-if __name__ == "__main__":
+if __name__ == '__main__':
+    _startup()
     demo.launch()