From ed962240b3e9d39d315fe8502d0937ce97a3ea35 Mon Sep 17 00:00:00 2001 From: niten Date: Wed, 15 Mar 2023 12:31:01 -0700 Subject: [PATCH] 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 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 | 14 +++++ src/detector.py | 123 +++++++++++++++++++++++++++-------------- src/objectifier.py | 14 +++-- src/yolo-cli.py | 9 +-- 4 files changed, 111 insertions(+), 49 deletions(-) diff --git a/objectifier-module.nix b/objectifier-module.nix index 10f2a0a..0f54046 100644 --- a/objectifier-module.nix +++ b/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; diff --git a/src/detector.py b/src/detector.py index 2020fb5..462fae7 100644 --- a/src/detector.py +++ b/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): - self.net = cv.dnn.readNet(weights, cfg) + def __init__(self, weights, cfg, classes, tempdir, pool_size, confidence=0.7): + 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) - outs = self.net.forward(self.output_layer) + with self.nets.reserve(timeout) as net: + 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 = [] - confidences = [] - boxes = [] - detections = [] + class_ids = [] + confidences = [] + boxes = [] + detections = [] - for out in outs: - for detection in out: - scores = detection[5:] - class_id = np.argmax(scores) - confidence = scores[class_id] - if confidence > self.minimum_confidence: - center_x = int(detection[0] * width) - center_y = int(detection[1] * height) - w = int(detection[2] * width) - h = int(detection[3] * height) - x = int(center_x - w/2) - y = int(center_y - h/2) - boxes.append([x, y, w, h]) - confidences.append(float(confidence)) - class_ids.append(class_id) + for out in outs: + for detection in out: + scores = detection[5:] + class_id = np.argmax(scores) + confidence = scores[class_id] + if confidence > self.minimum_confidence: + center_x = int(detection[0] * width) + center_y = int(detection[1] * height) + w = int(detection[2] * width) + h = int(detection[3] * height) + x = int(center_x - w/2) + y = int(center_y - h/2) + boxes.append([x, y, w, h]) + confidences.append(float(confidence)) + 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: - label = str(self.classes[class_ids[i]]) - box = [int(n) for n in boxes[i]] - detections.append(Detection(label, confidences[i], box)) + for i in indexes: + label = str(self.classes[class_ids[i]]) + box = [int(n) for n in boxes[i]] + detections.append(Detection(label, confidences[i], box)) - font = cv.FONT_HERSHEY_PLAIN - marked = cv.imread(str(filename)) - for detection in detections: - x, y, w, h = detection.box - 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) + font = cv.FONT_HERSHEY_PLAIN + marked = cv.imread(str(filename)) + for detection in detections: + x, y, w, h = detection.box + 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) - out_file = output_filename if output_filename else self.output_filename(filename) - cv.imwrite(out_file, marked) - return AnalyzedImage(filename, detections, str(out_file)) + out_file = output_filename if output_filename else self.output_filename(filename) + cv.imwrite(out_file, marked) + return AnalyzedImage(filename, detections, str(out_file)) diff --git a/src/objectifier.py b/src/objectifier.py index d435b2f..d364259 100644 --- a/src/objectifier.py +++ b/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, - yolo_config, - yolo_labels, - outgoing_dir) + weights=yolo_weights, + cfg=yolo_config, + classes=yolo_labels, + 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) diff --git a/src/yolo-cli.py b/src/yolo-cli.py index f475cbf..11b9e4f 100644 --- a/src/yolo-cli.py +++ b/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, - args.yolo_config, - classes, - outgoing_dir) + weights=args.yolo_weights, + cfg=args.yolo_config, + classes=classes, + tempdir=outgoing_dir, + pool_size=1) for filename in args.files: print(filename + ":")