302 lines
8.9 KiB
Python
302 lines
8.9 KiB
Python
from itertools import product
|
|
from math import ceil
|
|
|
|
import torch
|
|
from PIL import Image
|
|
from numpy import array
|
|
from retinaface.pre_trained_models import get_model
|
|
from retinaface.predict_single import Model
|
|
from retinaface.network import RetinaFace
|
|
from torch import jit, randn, no_grad, Tensor, int64, tensor, onnx
|
|
import albumentations as A
|
|
from torchinfo import summary
|
|
import numpy as np
|
|
from typing import Dict, List, Optional, Tuple, Union
|
|
import cv2
|
|
from torch.nn import functional as F
|
|
from torchvision.extension import _assert_has_ops
|
|
from torchvision.utils import _log_api_usage_once
|
|
|
|
# model: Model = get_model(model_name='resnet50_2020-07-20', max_size=512, device='cuda')
|
|
# model.eval()
|
|
|
|
|
|
image = Image.open(
|
|
fp="/home/tomokazu/PycharmProjects/helloproject-ai/data/blog_images"
|
|
"/稲場愛香/稲場愛香=juicejuice-official=12737097989-2.jpg").convert(mode="RGB")
|
|
image_arr = array(image)
|
|
max_size = 512
|
|
|
|
example_input = randn(size=[1, 3, 256, 256]).float().cuda()
|
|
|
|
retina_model = RetinaFace(
|
|
name="Resnet50",
|
|
pretrained=False,
|
|
return_layers={"layer2": 1, "layer3": 2, "layer4": 3},
|
|
in_channels=256,
|
|
out_channels=256,
|
|
).cuda()
|
|
|
|
|
|
# onnx.export(
|
|
# model=retina_model, args=example_input, export_params=True, verbose=False, input_names=["input"],
|
|
# output_names=["bbox", "confidence", "landmark"],
|
|
# dynamic_axes={"input": {
|
|
# 0: "batch_size",
|
|
# 2: "height",
|
|
# 3: "width"
|
|
# }, "bbox": {1: "bbox"}, "confidence": {1: "confidence"}, "landmark": {1: "landmark"}}, opset_version=16,
|
|
# f="retinaface.onnx"
|
|
# )
|
|
|
|
def pad_to_size(
|
|
target_size: Tuple[int, int],
|
|
image: np.array,
|
|
bboxes: Optional[np.ndarray] = None,
|
|
keypoints: Optional[np.ndarray] = None,
|
|
) -> Dict[str, Union[np.ndarray, Tuple[int, int, int, int]]]:
|
|
target_height, target_width = target_size
|
|
|
|
image_height, image_width = image.shape[:2]
|
|
|
|
if target_width < image_width:
|
|
raise ValueError(f"Target width should bigger than image_width" f"We got {target_width} {image_width}")
|
|
|
|
if target_height < image_height:
|
|
raise ValueError(f"Target height should bigger than image_height" f"We got {target_height} {image_height}")
|
|
|
|
if image_height == target_height:
|
|
y_min_pad = 0
|
|
y_max_pad = 0
|
|
else:
|
|
y_pad = target_height - image_height
|
|
y_min_pad = y_pad // 2
|
|
y_max_pad = y_pad - y_min_pad
|
|
|
|
if image_width == target_width:
|
|
x_min_pad = 0
|
|
x_max_pad = 0
|
|
else:
|
|
x_pad = target_width - image_width
|
|
x_min_pad = x_pad // 2
|
|
x_max_pad = x_pad - x_min_pad
|
|
|
|
result = {
|
|
"pads": (x_min_pad, y_min_pad, x_max_pad, y_max_pad),
|
|
"image": cv2.copyMakeBorder(image, y_min_pad, y_max_pad, x_min_pad, x_max_pad, cv2.BORDER_CONSTANT),
|
|
}
|
|
|
|
if bboxes is not None:
|
|
bboxes[:, 0] += x_min_pad
|
|
bboxes[:, 1] += y_min_pad
|
|
bboxes[:, 2] += x_min_pad
|
|
bboxes[:, 3] += y_min_pad
|
|
|
|
result["bboxes"] = bboxes
|
|
|
|
if keypoints is not None:
|
|
keypoints[:, 0] += x_min_pad
|
|
keypoints[:, 1] += y_min_pad
|
|
|
|
result["keypoints"] = keypoints
|
|
|
|
return result
|
|
|
|
|
|
def tensor_from_rgb_image(image: np.ndarray) -> torch.Tensor:
|
|
image = np.transpose(image, (2, 0, 1))
|
|
return torch.from_numpy(image)
|
|
|
|
|
|
def priorbox(min_sizes, steps, clip, image_size):
|
|
feature_maps = [[ceil(image_size[0] / step), ceil(image_size[1] / step)] for step in steps]
|
|
|
|
anchors = []
|
|
for k, f in enumerate(feature_maps):
|
|
t_min_sizes = min_sizes[k]
|
|
for i, j in product(range(f[0]), range(f[1])):
|
|
for min_size in t_min_sizes:
|
|
s_kx = min_size / image_size[1]
|
|
s_ky = min_size / image_size[0]
|
|
dense_cx = [x * steps[k] / image_size[1] for x in [j + 0.5]]
|
|
dense_cy = [y * steps[k] / image_size[0] for y in [i + 0.5]]
|
|
for cy, cx in product(dense_cy, dense_cx):
|
|
anchors += [cx, cy, s_kx, s_ky]
|
|
|
|
# back to torch land
|
|
output = torch.Tensor(anchors).view(-1, 4)
|
|
if clip:
|
|
output.clamp_(max=1, min=0)
|
|
return output
|
|
|
|
|
|
def decode(
|
|
loc: torch.Tensor, priors: torch.Tensor, variances: Union[List[float], Tuple[float, float]]
|
|
) -> torch.Tensor:
|
|
boxes = torch.cat(
|
|
(
|
|
priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
|
|
priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1]),
|
|
),
|
|
1,
|
|
)
|
|
boxes[:, :2] -= boxes[:, 2:] / 2
|
|
boxes[:, 2:] += boxes[:, :2]
|
|
return boxes
|
|
|
|
|
|
def decode_landm(
|
|
pre: torch.Tensor, priors: torch.Tensor, variances: Union[List[float], Tuple[float, float]]
|
|
) -> torch.Tensor:
|
|
return torch.cat(
|
|
(
|
|
priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
|
|
priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
|
|
priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
|
|
priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
|
|
priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
|
|
),
|
|
dim=1,
|
|
)
|
|
|
|
|
|
def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(nms)
|
|
_assert_has_ops()
|
|
return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
|
|
|
|
|
|
def unpad_from_size(
|
|
pads: Tuple[int, int, int, int],
|
|
image: Optional[np.array] = None,
|
|
bboxes: Optional[np.ndarray] = None,
|
|
keypoints: Optional[np.ndarray] = None,
|
|
) -> Dict[str, np.ndarray]:
|
|
x_min_pad, y_min_pad, x_max_pad, y_max_pad = pads
|
|
|
|
result = {}
|
|
|
|
if image is not None:
|
|
height, width = image.shape[:2]
|
|
result["image"] = image[y_min_pad: height - y_max_pad, x_min_pad: width - x_max_pad]
|
|
|
|
if bboxes is not None:
|
|
bboxes[:, 0] -= x_min_pad
|
|
bboxes[:, 1] -= y_min_pad
|
|
bboxes[:, 2] -= x_min_pad
|
|
bboxes[:, 3] -= y_min_pad
|
|
|
|
result["bboxes"] = bboxes
|
|
|
|
if keypoints is not None:
|
|
keypoints[:, 0] -= x_min_pad
|
|
keypoints[:, 1] -= y_min_pad
|
|
|
|
result["keypoints"] = keypoints
|
|
|
|
return result
|
|
|
|
|
|
device = "cuda"
|
|
transform = A.Compose([A.LongestMaxSize(max_size=max_size, p=1), A.Normalize(p=1)])
|
|
variance = [0.1, 0.2]
|
|
nms_threshold = .4
|
|
confidence_threshold = .7
|
|
_priorbox = priorbox(
|
|
min_sizes=[[16, 32], [64, 128], [256, 512]],
|
|
steps=[8, 16, 32],
|
|
clip=False,
|
|
image_size=(max_size, max_size),
|
|
).to(device)
|
|
original_height, original_width = image_arr.shape[:2]
|
|
|
|
scale_landmarks = torch.from_numpy(np.tile([max_size, max_size], 5)).to(device)
|
|
scale_bboxes = torch.from_numpy(np.tile([max_size, max_size], 2)).to(device)
|
|
|
|
transformed_image = transform(image=image_arr)["image"]
|
|
|
|
paded = pad_to_size(target_size=(max_size, max_size), image=transformed_image)
|
|
|
|
pads = paded["pads"]
|
|
|
|
torched_image = tensor_from_rgb_image(paded["image"]).to(device)
|
|
|
|
|
|
# loc, conf, land = retina_model(torched_image.unsqueeze(0))
|
|
|
|
|
|
def infer(loc, conf, land):
|
|
conf = F.softmax(conf, dim=-1)
|
|
|
|
annotations = []
|
|
|
|
boxes = decode(loc.data[0], _priorbox, variance)
|
|
|
|
boxes *= scale_bboxes
|
|
scores = conf[0][:, 1]
|
|
|
|
landmarks = decode_landm(land.data[0], _priorbox, variance)
|
|
landmarks *= scale_landmarks
|
|
|
|
# ignore low scores
|
|
valid_index = torch.where(scores > confidence_threshold)[0]
|
|
boxes = boxes[valid_index]
|
|
landmarks = landmarks[valid_index]
|
|
scores = scores[valid_index]
|
|
|
|
# Sort from high to low
|
|
order = scores.argsort(descending=True)
|
|
boxes = boxes[order]
|
|
landmarks = landmarks[order]
|
|
scores = scores[order]
|
|
|
|
# do NMS
|
|
keep = nms(boxes, scores, nms_threshold)
|
|
boxes = boxes[keep, :].int()
|
|
|
|
if boxes.shape[0] == 0:
|
|
return [{"bbox": [], "score": -1, "landmarks": []}]
|
|
|
|
landmarks = landmarks[keep]
|
|
|
|
scores = scores[keep].cpu().detach().numpy().astype(np.float64)
|
|
boxes = boxes.cpu().numpy()
|
|
landmarks = landmarks.cpu().numpy()
|
|
landmarks = landmarks.reshape([-1, 2])
|
|
|
|
unpadded = unpad_from_size(pads, bboxes=boxes, keypoints=landmarks)
|
|
|
|
resize_coeff = max(original_height, original_width) / max_size
|
|
|
|
boxes = (unpadded["bboxes"] * resize_coeff).astype(int)
|
|
landmarks = (unpadded["keypoints"].reshape(-1, 10) * resize_coeff).astype(int)
|
|
|
|
for box_id, bbox in enumerate(boxes):
|
|
x_min, y_min, x_max, y_max = bbox
|
|
|
|
x_min = np.clip(x_min, 0, original_width - 1)
|
|
x_max = np.clip(x_max, x_min + 1, original_width - 1)
|
|
|
|
if x_min >= x_max:
|
|
continue
|
|
|
|
y_min = np.clip(y_min, 0, original_height - 1)
|
|
y_max = np.clip(y_max, y_min + 1, original_height - 1)
|
|
|
|
if y_min >= y_max:
|
|
continue
|
|
|
|
annotations += [
|
|
{
|
|
"bbox": bbox.tolist(),
|
|
"score": scores[box_id],
|
|
"landmarks": landmarks[box_id].reshape(-1, 2).tolist(),
|
|
}
|
|
]
|
|
return annotations
|
|
|
|
|
|
ans = infer(*retina_model(torched_image.unsqueeze(0)))
|
|
print(ans)
|