File size: 1,002 Bytes
32580f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from flask import Flask, request, jsonify
from transformers import ViTFeatureExtractor, ViTForImageClassification
from PIL import Image
import requests

print("Loading models...")

feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch32-384')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch32-384')

print("Starting webapp...")

app = Flask(__name__)

print("Ready")

@app.route("/")
def hello_world():
    global feature_extractor, model
    
    url = request.args.get('url')
    
    if url is None:
        return jsonify(error="Url is required", url=None, classes=[])
    
    image = Image.open(requests.get(url, stream=True).raw)
    
    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits

    # model predicts one of the 1000 ImageNet classes
    predicted_class_idx = logits.argmax(-1).item()

    return jsonify(url=url, classes=model.config.id2label[predicted_class_idx])