Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
-import tensorflow as tf
-from transformers import AutoFeatureExtractor
+import torch
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 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 = "hoangthan/image-classification"
+model_name = "microsoft/quickdraw"
 logging.info("Loading image processor and model...")
 feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
-model = tf.keras.models.load_model('https://huggingface.co/hoangthan/image-classification/resolve/main/tf_model.h5')
+model = AutoModelForImageClassification.from_pretrained(model_name)
 
 # 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]
 
-        # 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]
+        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]
         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.")