initial commit

src/app.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+
+import os
+import uuid
+import json
+import time
+import cv2
+import queue
+from pathlib import Path as P
+from PIL import Image
+from flask import Flask, request, send_from_directory
+from flask_cors import CORS
+from threading import Thread
+
+app = Flask(__name__)
+CORS(app)
+
+QUEUE_STATE_IN_QUEUE = 'IN_QUEUE'
+QUEUE_STATE_PROCESSING = 'PROCESSING'
+QUEUE_STATE_READY = 'READY'
+
+q = queue.Queue(maxsize=1000)
+state = {}
+root = P("./q").absolute()
+root.mkdir(parents=True, exist_ok=True)
+
+
+def process():
+    print("importing ultralytics")
+    from ultralytics import YOLO
+    print("creating model")
+    model = YOLO(os.environ.get("MODEL_PATH", "/best.pt"))
+    print("created model")
+    last = 0
+    while True:
+        if time.time() < last + 10:
+            print("aa")
+            time.sleep(0.5)
+            continue
+
+        print("BB")
+
+        print("waiting for/getting job")
+        job = q.get()
+        filepath = root.joinpath(job)
+        print("got job")
+        state[job] = {
+            'state': QUEUE_STATE_PROCESSING,
+        }
+        print(f"processing job: {job=}")
+        img = cv2.imread(str(filepath))
+        results = model.predict(img)
+
+        for r in results:
+            im_array = r.plot()
+            im = Image.fromarray(im_array[..., ::-1])
+            im.save(str(filepath.with_suffix(".png")))
+            state[job] = {
+                'state': QUEUE_STATE_READY,
+            }
+
+        filepath.unlink()
+
+
+@app.route("/submit", methods=['POST'])
+def render():
+    id = str(uuid.uuid4())
+    request.files['image'].save(root.joinpath(id))
+    q.put(id)
+    state[id] = {'state': QUEUE_STATE_IN_QUEUE}
+    return json.dumps({'id': id})
+
+
+@app.route("/state/<id>")
+def get_state(id):
+    return json.dumps(state[id])
+
+
+@app.route("/result/<i>")
+def get_result(i):
+    print(root)
+    print(i + '.png')
+    return send_from_directory(root, i + '.png')
+
+
+if __name__ == '__main__':
+    p = Thread(target=process)
+    p.start()
+    app.run(debug=True)
+    p.join()