anon5's picture
Update app.py
ff692be verified
raw
history blame
802 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
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)