Spaces:
Build error
Build error
Commit
·
daa014f
1
Parent(s):
be2b94f
init imagesense app
Browse files- app.py +33 -0
- requirements.txt +2 -0
app.py
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import requests
|
| 2 |
+
import gradio as gr
|
| 3 |
+
import torch
|
| 4 |
+
|
| 5 |
+
from timm import create_model
|
| 6 |
+
from timm.data import resolve_data_config
|
| 7 |
+
from timm.data.transforms_factory import create_transform
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
IMAGENET_1K_URL = 'https://storage.googleapis.com/bit_models/ilsvrc2012_1k_wordnet_lemmas.txt'
|
| 11 |
+
IMAGENET_1K_LABELS = requests.get(IMAGENET_1K_URL).text.strip().split('\n')
|
| 12 |
+
|
| 13 |
+
model = create_model('restnet50', pretrained=True)
|
| 14 |
+
|
| 15 |
+
transform = create_transform(
|
| 16 |
+
**resolve_data_config({}, model=model)['test_time_augmentation'][0])
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
model.eval()
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def predict(image):
|
| 23 |
+
img = image.convert('RGB')
|
| 24 |
+
transformed_image = transform(img).unsqueeze(0)
|
| 25 |
+
with torch.no_grad():
|
| 26 |
+
out = model(transformed_image)
|
| 27 |
+
probabilities = torch.nn.functional.softmax(out[0], dim=0)
|
| 28 |
+
values, indices = torch.topk(probabilities, k=5)
|
| 29 |
+
return {IMAGENET_1K_LABELS[i]: v.item() for i, v in zip(indices, values)}
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
gr.Interface(predict, gr.inputs.Image(type='pil'),
|
| 33 |
+
output='label').launch()
|
requirements.txt
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
timm
|
| 2 |
+
gradio
|