Add cleanup thread, use systemd directories

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)