Use a pool of detectors.

When more than a single request was coming in at a time, the server
would crash when trying to write the output file. Each worker had a
single detector with one DN network, so if more than <WORKER#> requests
came in at once, one network would be used multiple times--and it's not
threadsafe.

I've added a queue of nets to pull from, so each request has exclusive
access to a specific network.

objectifier-module.nix

@@ -36,6 +36,18 @@ in {
       default = [ "127.0.0.1" ];
     };
 
+    pool-size = mkOption {
+      type = int;
+      description = "Number of nets to initialize.";
+      default = 5;
+    };
+
+    detection-timeout = mkOption {
+      type = int;
+      description = "Time in seconds to allow for detection to start.";
+      default = 5;
+    };
+
     cleanup = {
       max_file_age = mkOption {
         type = int;
@@ -63,6 +75,8 @@ in {
         OBJECTIFIER_BUFFER_SIZE = "524288";
         OBJECTIFIER_CLEANUP_MAX_AGE = toString cfg.cleanup.max_file_age;
         OBJECTIFIER_CLEANUP_DELAY = toString cfg.cleanup.delay;
+        OBJECTIFIER_TIMEOUT = toString cfg.detection-timeout;
+        OBJECTIFIER_POOL_SIZE = toString cfg.pool-size;
       };
       serviceConfig = {
         PrivateUsers = true;

src/detector.py

@@ -6,9 +6,11 @@ import sys
 import shutil
 from os import path
 import hashlib
+from contextlib import contextmanager
 
 import tempfile
 from pathlib import Path
+from queue import Queue, Empty, Full
 
 class Detection:
     """Represents an object dectected in an image."""
@@ -25,14 +27,52 @@ class AnalyzedImage:
         self.detections = detections
         self.outfile = outfile
 
+class ResourcePoolError(Exception):
+    """Base class for Pool errors."""
+
+class ResourcePoolTimeout(Exception):
+    """Timed out while waiting to resource to become available."""
+
+class ResourcePoolFull(Exception):
+    """Pool is full."""
+
+class ResourcePool:
+    """A pool to store shared resources."""
+
+    def __init__(self, pool_size, factory):
+        self._pool = Queue(pool_size)
+        for _ in range(pool_size):
+            self.__put(factory())
+
+    def __get(self, timeout):
+        try:
+            return self._pool.get(timeout=timeout)
+        except Empty:
+            raise ResourcePoolTimeout()
+
+    def __put(self, resource):
+        try:
+            return self._pool.put_nowait(resource)
+        except Full:
+            raise ResourcePoolFull()
+
+    @contextmanager
+    def reserve(self, timeout):
+        resource = self.__get(timeout)
+        try:
+            yield resource
+        finally:
+            self.__put(resource)
+
+def build_net(weights, cfg):
+    return cv.dnn.readNet(weights, cfg)
+
 class Detector:
     """Detects objects in images, returning an AnalyzedImage."""
 
-    def __init__(self, weights, cfg, classes, tempdir, confidence=0.6):
+    def __init__(self, weights, cfg, classes, tempdir, pool_size, confidence=0.7):
-        self.net = cv.dnn.readNet(weights, cfg)
+        self.nets = ResourcePool(pool_size, lambda: build_net(weights, cfg))
         self.classes = classes
-        self.layer_names = self.net.getLayerNames()
-        self.output_layer = [self.layer_names[i - 1] for i in self.net.getUnconnectedOutLayers()]
         self.tmpdir = tempdir
         self.minimum_confidence = confidence
 
@@ -40,48 +80,51 @@ class Detector:
         simple_name = path.splitext(path.basename(filename))[0]
         return str(self.tmpdir / (simple_name + ".png"))
 
-    def detect_objects(self, filename, output_filename=None):
+    def detect_objects(self, filename, timeout=5, output_filename=None):
         img = cv.imread(str(filename))
         height, width, channel = img.shape
         blob = cv.dnn.blobFromImage(img, 0.00392, (416, 416), (0,0,0), True, crop=False)
-        self.net.setInput(blob)
+        with self.nets.reserve(timeout) as net:
-        outs = self.net.forward(self.output_layer)
+            layer_names = net.getLayerNames()
+            output_layer = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
+            net.setInput(blob)
+            outs = net.forward(output_layer)
 
-        class_ids = []
+            class_ids = []
-        confidences = []
+            confidences = []
-        boxes = []
+            boxes = []
-        detections = []
+            detections = []
 
-        for out in outs:
+            for out in outs:
-            for detection in out:
+                for detection in out:
-                scores = detection[5:]
+                    scores = detection[5:]
-                class_id = np.argmax(scores)
+                    class_id = np.argmax(scores)
-                confidence = scores[class_id]
+                    confidence = scores[class_id]
-                if confidence > self.minimum_confidence:
+                    if confidence > self.minimum_confidence:
-                    center_x = int(detection[0] * width)
+                        center_x = int(detection[0] * width)
-                    center_y = int(detection[1] * height)
+                        center_y = int(detection[1] * height)
-                    w = int(detection[2] * width)
+                        w = int(detection[2] * width)
-                    h = int(detection[3] * height)
+                        h = int(detection[3] * height)
-                    x = int(center_x - w/2)
+                        x = int(center_x - w/2)
-                    y = int(center_y - h/2)
+                        y = int(center_y - h/2)
-                    boxes.append([x, y, w, h])
+                        boxes.append([x, y, w, h])
-                    confidences.append(float(confidence))
+                        confidences.append(float(confidence))
-                    class_ids.append(class_id)
+                        class_ids.append(class_id)
 
-        indexes = cv.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
+            indexes = cv.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
 
-        for i in indexes:
+            for i in indexes:
-            label = str(self.classes[class_ids[i]])
+                label = str(self.classes[class_ids[i]])
-            box = [int(n) for n in boxes[i]]
+                box = [int(n) for n in boxes[i]]
-            detections.append(Detection(label, confidences[i], box))
+                detections.append(Detection(label, confidences[i], box))
 
-        font = cv.FONT_HERSHEY_PLAIN
+            font = cv.FONT_HERSHEY_PLAIN
-        marked = cv.imread(str(filename))
+            marked = cv.imread(str(filename))
-        for detection in detections:
+            for detection in detections:
-            x, y, w, h = detection.box
+                x, y, w, h = detection.box
-            cv.rectangle(marked, (x,y), (x + w, y + h), (255,255,255,0), 2)
+                cv.rectangle(marked, (x,y), (x + w, y + h), (255,255,255,0), 2)
-            cv.putText(marked, detection.label, (x,y+30), font, 3, (255,255,255,0), 1)
+                cv.putText(marked, detection.label, (x,y+30), font, 3, (255,255,255,0), 1)
 
-        out_file = output_filename if output_filename else self.output_filename(filename)
+            out_file = output_filename if output_filename else self.output_filename(filename)
-        cv.imwrite(out_file, marked)
+            cv.imwrite(out_file, marked)
-        return AnalyzedImage(filename, detections, str(out_file))
+            return AnalyzedImage(filename, detections, str(out_file))

src/objectifier.py

@@ -36,6 +36,8 @@ yolo_labels_file = get_envvar_or_fail('OBJECTIFIER_YOLOV3_LABELS')
 buffer_size = to_int(get_envvar('OBJECTIFIER_BUFFER_SIZE')) or 524288
 max_file_age = to_int(get_envvar('OBJECTIFIER_CLEANUP_MAX_AGE'))
 file_cleanup_delay = to_int(get_envvar('OBJECTIFIER_CLEANUP_DELAY'))
+detection_timeout = to_int(get_envvar('OBJECTIFIER_TIMEOUT')) or 5
+pool_size = to_int(get_envvar('OBJECTIFIER_POOL_SIZE')) or 10
 
 yolo_labels = []
 with open(yolo_labels_file, "r") as f:
@@ -45,10 +47,11 @@ incoming_dir = Path(get_envvar_or_fail('CACHE_DIRECTORY'))
 outgoing_dir = Path(get_envvar_or_fail('STATE_DIRECTORY'))
 
 detector = Detector(
-    yolo_weights,
+    weights=yolo_weights,
-    yolo_config,
+    cfg=yolo_config,
-    yolo_labels,
+    classes=yolo_labels,
-    outgoing_dir)
+    tempdir=outgoing_dir,
+    pool_size=pool_size)
 
 app = FastAPI()
 
@@ -108,7 +111,8 @@ def analyze_image(request: Request, image: UploadFile):
             chunk = image.file.read(buffer_size)
     result = detector.detect_objects(
         infile,
-        str(outgoing_dir / (file_hash.hexdigest() + ".png")))
+        str(outgoing_dir / (file_hash.hexdigest() + ".png")),
+        detection_timeout)
     remove(infile)
     return result_to_dict(result, base_url)
 

src/yolo-cli.py

@@ -27,10 +27,11 @@ with open(args.yolo_labels) as f:
     classes = [line.strip() for line in f.readlines()]
 
 detector = Detector(
-    args.yolo_weights,
+    weights=args.yolo_weights,
-    args.yolo_config,
+    cfg=args.yolo_config,
-    classes,
+    classes=classes,
-    outgoing_dir)
+    tempdir=outgoing_dir,
+    pool_size=1)
 
 for filename in args.files:
     print(filename + ":")