urandu/face_detect_api

View on GitHub
api/tasks/image.py

Summary

Maintainability
A
0 mins
Test Coverage
# The future is now!
import uuid

import requests
from django.core.files.storage import default_storage
from mtcnn.mtcnn import MTCNN
from numpy import asarray
from api.models.face import Face
from api.models.image import Image as Image_object
from PIL import Image as PImage
from api.celery_app import app
import numpy as np
import json
import cv2
import io
from django.core.files.uploadedfile import InMemoryUploadedFile
from django.conf import settings


@app.task(bind=True, name='detect_faces')
def detect_faces(self, *args, **kwargs):
    image_id = kwargs.get("image_id")
    image_object = Image_object.objects.get(image_id=image_id)
    filename = image_object.name
    image = PImage.open(default_storage.open(filename))
    image = image.convert('RGB')
    pixels = asarray(image)

    detector = MTCNN()

    # detect faces in the image
    results = detector.detect_faces(pixels)

    detected_faces = list()
    for result in results:

        # only detect faces with a confidence of 94% and above
        if result['confidence'] > 0.94:
            face_object = Face()
            face_id = str(uuid.uuid4())
            face_object.face_id = face_id
            face_object.image_id = image_id
            face_object.confidence = result['confidence']
            face_object.box = result['box']
            face_object.keypoints = result['keypoints']
            face_object.save()
            detected_faces.append(face_id)

    return detected_faces



@app.task(bind=True, name='detect_faces_callback')
def detect_faces_callback(self, *args, **kwargs):
    image_id = kwargs.get("image_id")
    image_object = Image_object.objects.get(image_id=image_id)

    filename = image_object.name
    output_filename = "detected_faces/" + image_object.name
    faces_on_image = Face.objects.filter(image_id=image_id)
    image = PImage.open(default_storage.open(filename))
    image = np.array(image)
    image = image.copy()
    faces_dict = list()
    for face in faces_on_image:
        faces_dict.append({
            "confidence":face.confidence,
            "box":face.box,
            "keypoints":face.keypoints
        })
        box = json.loads(face.box)

        x1, y1, width, height = box
        x1, y1 = abs(x1), abs(y1)
        x2, y2 = x1 + width, y1 + height
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 5)
        cv2.putText(image,
                    "P: " + "{0:.4f}".format(float(face.confidence)),
                    (x1, (y2 + 25)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)

    pil_im = PImage.fromarray(image)
    silver_bullet = io.BytesIO()
    pil_im.save(silver_bullet, format="JPEG")

    image_file = InMemoryUploadedFile(silver_bullet, None, output_filename, 'image/jpeg',
                                      len(silver_bullet.getvalue()), None)

    default_storage.save(output_filename, image_file)

    callback = dict({
        "image_id":image_id,
        "request_id":image_object.request_id,
        "faces":faces_dict,
        "output_image_url":"{host}/api/image/?image_id={image_id}".format(host=settings.API_HOST, image_id=image_id)
    })
    image_object.status = "completed"
    image_object.save()

    try:
        requests.post(url=image_object.callback_url, data=json.dumps(callback))
    except Exception as e:
        # log exception
        pass

    return image_id