Create app.py

app.py

@@ -0,0 +1,36 @@
+import gradio as gr
+import timm
+import hyenapixel.models
+import torch
+import numpy as np
+from PIL import Image
+ 
+with open("imagenet.txt") as file:
+    class_names = [line.rstrip() for line in file]
+ 
+def predict(model_name, image):
+    model = timm.create_model(model_name, pretrained=True)
+    model.eval()
+    image_size = 224
+    if "_384" in model_name:
+        image_size = 384
+    transform = timm.data.create_transform(image_size)
+    input_tensor = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        output = model(input_tensor)
+    output_np = output[0].numpy()
+    class_ind = np.argmax(output_np)
+    return class_names[class_ind]
+ 
+interface = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Dropdown(label="Select Model", value="hb_former_b36", choices=["hpx_former_s18", "hpx_former_s18_384", "hb_former_s18", "c_hpx_former_s18", "hpx_a_former_s18", "hb_a_former_s18", "hpx_former_b36", "hb_former_b36"]),
+        gr.Image(type="pil", label="Upload Image")
+    ],
+    outputs=gr.Textbox(label="Predicted Class"),
+    title="Image Classification",
+    description="Choose a model and upload an image to predict the class."
+)
+ 
+interface.launch()