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.
This commit is contained in:
parent
d4de83bf4b
commit
ed962240b3
|
|
@ -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;
|
||||
|
|
|
|||
123
src/detector.py
123
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))
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 + ":")
|
||||
|
|
|
|||
Loading…
Reference in New Issue