anon5's picture
Rename a0p.py to app.py
85b3f26 verified
raw
history blame
800 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)