33 lines
1021 B
Python
33 lines
1021 B
Python
from PIL import Image, ImageDraw
|
|
|
|
from matplotlib import pyplot
|
|
from retinaface.pre_trained_models import get_model
|
|
from torch import save
|
|
from retinaface.predict_single import Model
|
|
from retinaface.network import RetinaFace
|
|
from numpy import array
|
|
|
|
im_path = "道重さゆみ=sayumimichishige-blog=12719215891-10.jpg"
|
|
model: Model = get_model(model_name='resnet50_2020-07-20', max_size=1024, device='cuda')
|
|
image = array(Image.open(im_path))
|
|
faces = model.predict_jsons(image, confidence_threshold=0.4, nms_threshold=0.4)
|
|
|
|
print(faces)
|
|
# print(dfs)
|
|
image = Image.open(im_path)
|
|
draw = ImageDraw.Draw(image)
|
|
if faces.__len__() == 1:
|
|
if faces[0]['score'] == -1:
|
|
exit()
|
|
for face in faces:
|
|
print(face['bbox'])
|
|
x, y, w, h = face['bbox']
|
|
print(x, y, w, h)
|
|
draw.rectangle((x, y, w, h))
|
|
for order, landmark in enumerate(face['landmarks'], start=1):
|
|
draw.point(tuple(landmark))
|
|
draw.text(tuple(landmark), text=str(order))
|
|
#
|
|
image.save("dest.png")
|
|
# RetinaFace.detect_faces(im_path)
|