File size: 1,213 Bytes
023c2b7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import tempfile
import gradio as gr
from autodistill_fastvit import FASTVIT_IMAGENET_1K_CLASSES, FastViT
from PIL import Image
base_model = FastViT(None)
def infer(image):
with tempfile.NamedTemporaryFile(suffix=".jpg") as temp:
image = Image.fromarray(image.astype("uint8"), "RGB")
image.save(temp.name)
predictions = base_model.predict(temp.name, confidence=0.1)
labels = [FASTVIT_IMAGENET_1K_CLASSES[i] for i in predictions.class_id.tolist()]
confidences = predictions.confidence.tolist()
# divide by 100 to convert to percentage
confidences = [c / 100 for c in confidences]
return {
k: v
for k, v in zip(labels, confidences)
}
iface = gr.Interface(
fn=infer,
inputs="image",
outputs="label",
allow_flagging=False,
title="FastViT",
description="[FastViT](https://github.com/apple/ml-fastvit) is a fast Vision Transformer developed by Apple. FastViT was trained on the ImageNet-1k dataset.\n\nUse the space below to test FastViT on your own images.\n\nThis space uses [Autodistill FastViT](https://github.com/autodistill/autodistill-fastvit) for inference.",
)
iface.launch()
|