diff --git a/objectifier-module.nix b/objectifier-module.nix index 372f1f6..db0b11f 100644 --- a/objectifier-module.nix +++ b/objectifier-module.nix @@ -35,6 +35,21 @@ in { "List of IP addresses on which to listen for incoming requests."; default = [ "127.0.0.1" ]; }; + + cleanup = { + max_file_age = mkOption { + type = int; + description = + "Maximum age of a file (in seconds), after which it will be removed."; + default = 60 * 60 * 8; # 8 hours + }; + + delay = { + type = int; + description = "Time between cleanup sweeps."; + default = 60 * 10; # 10 minutes + }; + }; }; config = mkIf cfg.enable { @@ -48,6 +63,8 @@ in { OBJECTIFIER_YOLOV3_WEIGHTS = "${pkgs.yolov3-data}/yolov3.weights"; OBJECTIFIER_YOLOV3_LABELS = "${pkgs.yolov3-data}/labels"; OBJECTIFIER_BUFFER_SIZE = "524288"; + OBJECTIFIER_CLEANUP_MAX_AGE = toString cfg.cleanup.max_file_age; + OBJECTIFIER_CLEANUP_DELAY = toString cfg.cleanup.delay; }; serviceConfig = { # PrivateUsers = true; diff --git a/src/objectifier.py b/src/objectifier.py index 7e02872..070c724 100644 --- a/src/objectifier.py +++ b/src/objectifier.py @@ -7,9 +7,10 @@ import tempfile from pathlib import Path import hashlib import os +from os import listdir, remove +from os.path import getatime, splitext +from time import sleep -incoming_dir = Path(tempfile.mkdtemp()) -outgoing_dir = Path(tempfile.mkdtemp()) def get_envvar(name): return os.environ.get(name) @@ -21,10 +22,21 @@ def get_envvar_or_fail(name): else: raise EnvironmentError('Missing required environment variable: ' + name) +def to_int(input_int): + if input_int: + return int(input_int) + else: + return None + yolo_config = get_envvar_or_fail('OBJECTIFIER_YOLOV3_CONFIG') yolo_weights = get_envvar_or_fail('OBJECTIFIER_YOLOV3_WEIGHTS') yolo_labels = get_envvar_or_fail('OBJECTIFIER_YOLOV3_LABELS') -buffer_size = get_envvar('OBJECTIFIER_BUFFER_SIZE') or 524288 +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')) + +incoming_dir = Path(get_envvar_or_fail('CACHE_DIRECTORY')) +outgoing_dir = Path(get_envvar_or_fail('STATE_DIRECTORY')) detector = Detector( yolo_weights, @@ -36,6 +48,27 @@ app = FastAPI() analyzed_images = {} +img_formats = [ "png", "jpg", "gif", "bmp" ] + +def cleanup_old_files(path, extensions, max_age): + for filename in listdir(path): + if (splitext(filename) in extensions) and (getatime(filename) - time.time() > max_age): + print("removing old output file: " + filename) + remove(filename) + +def run_cleanup_thread(path, extensions, age, delay): + while True: + cleanup_old_files(path, extensions, age) + sleep(delay) + +cleanup_thread = Thread( + target=run_cleanup_thread, + args=(outgoing_dir, img_formats, max_file_age, file_cleanup_delay)) + +cleanup_thread.daemon = True + +cleanup_thread.start() + def detection_to_dict(d): return { "label": d.label, @@ -67,6 +100,7 @@ async def analyze_image(file: UploadFile, request: Request): infile.write(chunk) chunk=f.read(buffer_size) result = detector.detect_objects(infile, file_hash.hexdigest() + ".png") + remove(infile) return result_to_dict(result, base_url) @app.get("/analyzed_images/${image_name}", response_class=FileResponse) diff --git a/yolo.org b/yolo.org new file mode 100644 index 0000000..9084893 --- /dev/null +++ b/yolo.org @@ -0,0 +1,7 @@ +#+title: Yolo + +* TODO + +** [ ] Periodically clean up files from the output directory + +** [ ] Share a common output directory between the various threads