Update app.py

app.py

@@ -1,40 +1,63 @@
-import os
-import tensorflow as tf
-import numpy as np
-import pandas as pd
-import matplotlib.pyplot as plt
-from tensorflow import keras
+import sys
 
-import requests
-import PIL
-import io
 import matplotlib.pyplot as plt
+import PIL
+from PIL import Image
+import json
+
+import torch
+import torchvision
+import torchvision.transforms as T
 
 
-from keras_cv_attention_models import convnext
+from timm import create_model
 
 import gradio as gr
 
-mm = convnext.ConvNeXtXlarge()
+
+model_name = "convnext_xlarge_in22k"
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+# create a ConvNeXt model : https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/convnext.py
+model = create_model(model_name, pretrained=True).to(device)
+
+# Define transforms for test
+from timm.data.constants import \
+    IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+
+NORMALIZE_MEAN = IMAGENET_DEFAULT_MEAN
+NORMALIZE_STD = IMAGENET_DEFAULT_STD
+SIZE = 256
+
+# Here we resize smaller edge to 256, no center cropping
+transforms = [
+              T.Resize(SIZE, interpolation=T.InterpolationMode.BICUBIC),
+              T.ToTensor(),
+              T.Normalize(NORMALIZE_MEAN, NORMALIZE_STD),
+              ]
+
+transforms = T.Compose(transforms)
+
+os.system("wget https://dl.fbaipublicfiles.com/convnext/label_to_words.json")
+imagenet_labels = json.load(open('label_to_words.json'))
 
 def inference(img):
-    img = img
-    imm = keras.applications.imagenet_utils.preprocess_input(img, mode='torch')
-    image_input = tf.expand_dims(tf.image.resize(imm, mm.input_shape[1:3]), 0)
-    
-    pred = mm(image_input)
-    pred_np = pred.numpy()
-    
-    pred_names = keras.applications.imagenet_utils.decode_predictions(pred.numpy())[0]
+    img_tensor = transforms(img).unsqueeze(0).to(device)
+    # inference
+    output = torch.softmax(model(img_tensor), dim=1)
+    top5 = torch.topk(output, k=5)
+    top5_prob = top5.values[0]
+    top5_indices = top5.indices[0]
     
     result = {}
-    
+
     for i in range(5):
-        result[pred_names[i][1]] = pred_names[i][2].item()
+        labels = imagenet_labels[str(int(top5_indices[i]))]
+        prob = "{:.2f}%".format(float(top5_prob[i])*100)
+        results[labels] = prob
     
     return result
 
-inputs = gr.inputs.Image(type='numpy')
+inputs = gr.inputs.Image(type='pil')
 outputs = gr.outputs.Label(type="confidences",num_top_classes=5)
 
 title = "ConvNeXt"