|
@@ -0,0 +1,59 @@
|
|
|
|
|
+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)
|