init

.gitignore

@@ -0,0 +1,2 @@
+tmp/
+yolo*.pt

.idea/.gitignore

@@ -0,0 +1,3 @@
+# Default ignored files
+/shelf/
+/workspace.xml

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,12 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyPep8NamingInspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
+      <option name="ignoredErrors">
+        <list>
+          <option value="N803" />
+        </list>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.9 (yolo-api)" project-jdk-type="Python SDK" />
+  <component name="PyCharmProfessionalAdvertiser">
+    <option name="shown" value="true" />
+  </component>
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/yolo-api.iml" filepath="$PROJECT_DIR$/.idea/yolo-api.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="$PROJECT_DIR$" vcs="Git" />
+  </component>
+</project>

.idea/yolo-api.iml

@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$">
+      <excludeFolder url="file://$MODULE_DIR$/venv" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

DOCKERFILE

@@ -0,0 +1,19 @@
+FROM ubuntu:20.04
+
+MAINTANER Torsten Simon "simon@edu-sharing.net"
+
+RUN apt-get update -y && \
+    apt-get install -y python-pip python-dev
+
+# We copy just the requirements.txt first to leverage Docker cache
+COPY ./requirements.txt /app/requirements.txt
+
+WORKDIR /app
+
+RUN pip install -r requirements.txt
+
+COPY . /app
+
+ENTRYPOINT [ "python" ]
+
+CMD [ "main.py" ]

main.py

@@ -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)

requirements.txt

@@ -0,0 +1,25 @@
+# required
+Flask>=0.10
+PyYAML>=3.0
+jsonschema>=3.0.1
+six>=1.10.0
+mistune
+werkzeug
+torch
+
+
+# Base ----------------------------------------
+matplotlib>=3.2.2
+numpy>=1.18.5
+opencv-python>=4.1.2
+Pillow>=7.1.2
+PyYAML>=5.3.1
+requests>=2.23.0
+scipy>=1.4.1
+torch>=1.7.0
+torchvision>=0.8.1
+tqdm>=4.41.0
+
+# Plotting ------------------------------------
+pandas>=1.1.4
+seaborn>=0.11.0