Face identification and recognition scalable server with multiple face directories.
https://github.com/ehp/faceserver
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
87 lines
2.8 KiB
87 lines
2.8 KiB
import numpy as np |
|
import torch |
|
import argparse |
|
import json |
|
from PIL import Image, ImageDraw |
|
|
|
from identification.dataloader import Normalizer, Resizer |
|
from torchvision import transforms |
|
|
|
|
|
def fan_detect(model, img_data, threshold=0.9, max_detections=100, is_cuda=True): |
|
input_data = {'img': img_data, 'annot': np.zeros((0, 5)), 'scale': 1} |
|
transform = transforms.Compose([Resizer(), Normalizer()]) |
|
transformed = transform(input_data) |
|
|
|
model.eval() |
|
with torch.no_grad(): |
|
img_data = transformed['img'].permute(2, 0, 1).float().unsqueeze(dim=0) |
|
if is_cuda: |
|
img_data = img_data.cuda() |
|
scores, labels, boxes = model(img_data) |
|
if scores is None: |
|
return np.empty((0,0)), np.empty((0,0)) |
|
|
|
scores = scores.cpu().numpy() |
|
scale = transformed['scale'] |
|
boxes = boxes.cpu().numpy() / scale |
|
|
|
indices = np.where(scores > threshold)[0] |
|
scores = scores[indices] |
|
scores_sort = np.argsort(-scores)[:max_detections] |
|
image_boxes = boxes[indices[scores_sort], :] |
|
|
|
return image_boxes, scores[:max_detections] |
|
|
|
|
|
def img_rectangles(img, output_path, boxes=None): |
|
if boxes is not None: |
|
draw = ImageDraw.Draw(img) |
|
for arr in boxes: |
|
draw.rectangle(((arr[0], arr[1]), (arr[2], arr[3])), outline="black", width=1) |
|
|
|
img.save(output_path) |
|
|
|
|
|
def load_model(model_path, is_cuda=True): |
|
# load possible cuda model as cpu |
|
model = torch.load(model_path, map_location=lambda storage, location: storage) |
|
if is_cuda: |
|
model = model.cuda() |
|
|
|
model.anchors.is_cuda=is_cuda |
|
|
|
return model |
|
|
|
|
|
def load_image(filepath): |
|
img = Image.open(filepath) |
|
img = img.convert(mode="RGB") |
|
return img |
|
|
|
|
|
def main(args=None): |
|
parser = argparse.ArgumentParser(description='Simple training script for training a RetinaNet network.') |
|
|
|
parser.add_argument('--model', help='Path to model') |
|
parser.add_argument('--image', help='Path to image') |
|
parser.add_argument('--rect', help='Output image with rectangles') |
|
parser.add_argument('--threshold', help='Probability threshold (default 0.9)', type=float, default=0.9) |
|
parser.add_argument('--force-cpu', help='Force CPU for detection (default false)', dest='force_cpu', |
|
default=False, action='store_true') |
|
|
|
parser = parser.parse_args(args) |
|
|
|
is_cuda = torch.cuda.is_available() and not parser.force_cpu |
|
|
|
model = load_model(parser.model, is_cuda=is_cuda) |
|
img = load_image(parser.image) |
|
boxes, scores = fan_detect(model, img, threshold=parser.threshold, is_cuda=is_cuda) |
|
print(json.dumps({'boxes': boxes.tolist(), 'scores': scores})) |
|
if parser.rect: |
|
img = load_image(parser.image) |
|
img_rectangles(img, parser.rect, boxes) |
|
|
|
|
|
if __name__ == '__main__': |
|
main()
|
|
|