anon5's picture
Update app.py
a8373a3 verified
raw
history blame
837 Bytes
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
url = request.args.get('url')
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)