Update app.py
Browse files
app.py
CHANGED
@@ -12,7 +12,7 @@ labels = dataset['train'].features['labels'].names
|
|
12 |
example_imgs = ["example_0.jpg", "example_1.jpg"]
|
13 |
|
14 |
def classify(im):
|
15 |
-
features =
|
16 |
logits = model(features["pixel_values"])[-1]
|
17 |
probability = torch.nn.functional.softmax(logits, dim=-1)
|
18 |
probs = probability[0].detach().numpy()
|
|
|
12 |
example_imgs = ["example_0.jpg", "example_1.jpg"]
|
13 |
|
14 |
def classify(im):
|
15 |
+
features = extractor(im, return_tensors='pt')
|
16 |
logits = model(features["pixel_values"])[-1]
|
17 |
probability = torch.nn.functional.softmax(logits, dim=-1)
|
18 |
probs = probability[0].detach().numpy()
|