142 lines
4.3 KiB
Python
142 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
|
|
from detector import Detector
|
|
from fastapi import FastAPI, HTTPException, Request, UploadFile, File
|
|
from fastapi.responses import FileResponse
|
|
from os import listdir, remove
|
|
from os.path import getatime, splitext, basename, isfile
|
|
from pathlib import Path
|
|
from threading import Thread
|
|
from time import sleep
|
|
import hashlib
|
|
import os
|
|
import re
|
|
import tempfile
|
|
|
|
|
|
def get_envvar(name):
|
|
return os.environ.get(name)
|
|
|
|
def get_envvar_or_fail(name):
|
|
result = get_envvar(name)
|
|
if result:
|
|
return result
|
|
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 = 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,
|
|
yolo_config,
|
|
yolo_labels,
|
|
outgoing_dir)
|
|
|
|
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,
|
|
"confidence": d.confidence,
|
|
"box": {
|
|
"x": d.box[0],
|
|
"y": d.box[1],
|
|
"width": d.box[2],
|
|
"height": d.box[3],
|
|
},
|
|
}
|
|
|
|
def result_to_dict(res, base_url):
|
|
return {
|
|
"labels": list(map(lambda d: d.label, res.detections)),
|
|
"detections": list(map(detection_to_dict, res.detections)),
|
|
"output": base_url + basename(res.outfile),
|
|
}
|
|
|
|
# @app.put("/images")
|
|
# def analyze_image(file: UploadFile(), request: Request):
|
|
# print("Initiating file receipt, url: " + str(request.url))
|
|
# base_url = re.sub(r'\/images\/$', '/analyzed_images/', str(request.url))
|
|
# infile = open(incoming_dir / file.filename)
|
|
# file_hash = hashlib.sha256()
|
|
# with open(infile, mode="wb") as f:
|
|
# chunk = f.read(buffer_size)
|
|
# while chunk:
|
|
# print("writing chunk")
|
|
# file_hash.update(chunk)
|
|
# infile.write(chunk)
|
|
# chunk=f.read(buffer_size)
|
|
# print("saving complete")
|
|
# print("analyzing image")
|
|
# result = detector.detect_objects(infile, file_hash.hexdigest() + ".png")
|
|
# print("image analyzed")
|
|
# remove(infile)
|
|
# return result_to_dict(result, base_url)
|
|
|
|
@app.post("/images")
|
|
def analyze_image(request: Request, image: UploadFile):
|
|
print("Initiating file receipt, url: " + str(request.url))
|
|
base_url = re.sub(r'\/images\/?$', '/analyzed_images/', str(request.url))
|
|
infile = incoming_dir / image.filename
|
|
file_hash = hashlib.sha256()
|
|
with open(infile, "wb") as f:
|
|
chunk = image.file.read(buffer_size)
|
|
while chunk:
|
|
file_hash.update(chunk)
|
|
print ("writing " + str(buffer_size) + " bytes")
|
|
f.write(chunk)
|
|
chunk = image.file.read(buffer_size)
|
|
print("save complete")
|
|
print("analyzing image")
|
|
result = detector.detect_objects(
|
|
infile,
|
|
str(outgoing_dir / (file_hash.hexdigest() + ".png")))
|
|
print("image analyzed")
|
|
remove(infile)
|
|
return result_to_dict(result, base_url)
|
|
|
|
@app.get("/analyzed_images/{image_name}", response_class=FileResponse)
|
|
def get_analyzed_image(image_name: str):
|
|
filename = outgoing_dir / image_name
|
|
if isfile(filename):
|
|
return str(filename)
|
|
else:
|
|
raise HTTPException(status_code=404, detail="file not found: " + str(filename))
|