face3d0725 / FaceExtraction

MIT License
36 stars 7 forks source link

Unable to get results #5

Closed angelabr closed 1 year ago

angelabr commented 1 year ago

I'm loading the pretrained model and executing:

data = load_data("my_image.png")
data = data.to(DEVICE)
with torch.no_grad():
    pred = model(data)

# m = nn.Softmax(dim=2)
# pred = m(pred)

pred_mask = (pred > 0).type(torch.int8)
pred_mask = pred.squeeze().cpu().numpy()
cv2.imwrite("a.png", pred_mask)

It doesn't matter the image, the result is always black, could you point me in the right direction to obtain the mask? Thank you