volo-api/main.py

import os
import shutil
import tempfile
import uuid
from datetime import time
from uuid import UUID

import cv2
from flask import Flask, jsonify, request
from flasgger import Swagger
import torch
from werkzeug.utils import secure_filename

app = Flask(__name__)
swagger = Swagger(app)

@app.route('/classify', methods=["POST"])
def classify():
    """Classify an image or video
    ---
    parameters:
      - name: file
        in: formData
        description: The uploaded file data
        required: true
        type: file
      - name: confidence
        in: query
        type: float
        required: false
        default: 0.75
    responses:
      200:
        description: A list of entities found in the source
    """
    file = request.files.getlist('file')[0]
    filename = secure_filename(str(uuid.uuid1()) + '-' + file.filename)
    try:
        os.mkdir('tmp')
    except:
        shutil.rmtree('tmp')
        os.mkdir('tmp')
        pass
    img = os.path.join('tmp', filename)
    file.save(img)
    print(file.name)
    print(img)
    model = torch.hub.load('ultralytics/yolov5', 'yolov5x6')  # or yolov5m, yolov5l, yolov5x, custom

    # Inference
    model.conf = float(request.args.get('confidence'))
    model.iou = 0.5  # NMS IoU threshold (0-1)
    # Results
    results = model(img)
    # results.json()  # or .show(), .save(), .crop(), .pandas(), etc.
    results.show()
    return results.pandas().xyxy[0].to_json()

app.run(debug=True)