Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
-import torch
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+import tensorflow as tf
+from transformers import AutoFeatureExtractor
 from PIL import Image
 import numpy as np
 import logging
@@ -9,10 +9,10 @@ import logging
 logging.basicConfig(level=logging.DEBUG)
 
 # Load the pre-trained model and feature extractor
-model_name = "google/vit-base-patch16-224"
+model_name = "hoangthan/image-classification"
 logging.info("Loading image processor and model...")
 feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
-model = AutoModelForImageClassification.from_pretrained(model_name)
+model = tf.keras.models.load_model('https://huggingface.co/hoangthan/image-classification/resolve/main/tf_model.h5')
 
 # Define the prediction function
 def predict(image):
@@ -29,15 +29,15 @@ def predict(image):
         logging.debug("Converting NumPy array to PIL image...")
         image = Image.fromarray(image, 'RGBA').convert('RGB')
         logging.debug("Image converted successfully.")
+        
+        # Process the image for the model
+        inputs = feature_extractor(images=image, return_tensors="np")
+        pixel_values = inputs['pixel_values'][0]
 
-        logging.info("Processing image...")
-        inputs = feature_extractor(images=image, return_tensors="pt")
-        outputs = model(**inputs)
-        logits = outputs.logits
-        probs = torch.nn.functional.softmax(logits, dim=-1)
-        top_probs, top_idxs = probs.topk(3, dim=-1)
-        top_probs = top_probs.detach().numpy()[0]
-        top_idxs = top_idxs.detach().numpy()[0]
+        # Predict using the model
+        preds = model.predict(np.expand_dims(pixel_values, axis=0))
+        top_probs = tf.nn.softmax(preds[0])
+        top_idxs = np.argsort(-top_probs)[:3]
         top_classes = [model.config.id2label[idx] for idx in top_idxs]
         result = {top_classes[i]: float(top_probs[i]) for i in range(3)}
         logging.info("Prediction successful.")