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)