File size: 800 Bytes
79e23f0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from flask import Flask, request, jsonify
from transformers import pipeline
from PIL import Image
import requests
import logging
import torch
from transformers import pipeline

print("Loading model...")

#torch.backends.cuda.matmul.allow_tf32 = True
#torch.backends.cudnn.allow_tf32 = True

pipe = pipeline("image-classification", model="./checkpoint-600")

print("Starting webapp...")

app = Flask(__name__)
log = logging.getLogger('werkzeug')
log.disabled = True
app.logger.disabled = True

print("Ready")

@app.route("/")
def classify_image():
    global pipe
    
    if url is None:
        return jsonify(error="Url is required", url=None, label=None)
    
    image = Image.open(requests.get(url, stream=True).raw)

    output = pipe(images=[image])

    return jsonify(url=url, result=output)