65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
import queue
|
|
import threading
|
|
import logging
|
|
import time
|
|
import json
|
|
import base64
|
|
import os
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import requests
|
|
|
|
from edge import imgproc
|
|
from edge import settings
|
|
|
|
compute_queue = queue.Queue()
|
|
|
|
|
|
def load_models(base='.'):
|
|
logging.info('Loading models...')
|
|
proto_path = os.path.join(base, 'models/deploy.prototxt')
|
|
model_path = os.path.join(
|
|
base, 'models/res10_300x300_ssd_iter_140000.caffemodel')
|
|
detector = cv2.dnn.readNetFromCaffe(proto_path, model_path)
|
|
embedder = cv2.dnn.readNetFromTorch(
|
|
os.path.join(base, 'models/nn4.small2.v1.t7'))
|
|
return (detector, embedder)
|
|
|
|
|
|
def process_images():
|
|
logging.info('Processing thread started')
|
|
detector, embedder = load_models()
|
|
while True:
|
|
host, device, jpeg_image = compute_queue.get()
|
|
logging.info('Received image from %s (%s)' % (device, host))
|
|
np_img = np.asarray(bytearray(jpeg_image), dtype=np.uint8)
|
|
image = cv2.imdecode(np_img, cv2.IMREAD_COLOR)
|
|
|
|
vectors = []
|
|
for box in imgproc.detect_faces(image, detector, 0.5):
|
|
embedding = imgproc.vectorize_face(image, box, embedder)
|
|
vectors.append(
|
|
base64.b64encode(embedding.tobytes()).decode('utf8'))
|
|
|
|
logging.info('Found %d face(s)', len(vectors))
|
|
|
|
if vectors:
|
|
logging.info('Forwarding faces to cloud')
|
|
payload = json.dumps({
|
|
'faces': vectors
|
|
})
|
|
logging.info('Sending request to cloud')
|
|
try:
|
|
response = requests.post(
|
|
settings.FACES_URL % device, data=payload)
|
|
except Exception as e:
|
|
logging.warning('Error while sending request: %s', str(e))
|
|
|
|
|
|
def start_processing(nb_thread=4):
|
|
logging.info('Starting processing threads...')
|
|
for i in range(nb_thread):
|
|
thread = threading.Thread(target=process_images, daemon=True)
|
|
thread.start()
|