update
continuous-integration/drone/push Build is passing
Details
continuous-integration/drone/push Build is passing
Details
This commit is contained in:
parent
f362c10892
commit
6e2c6260fd
|
|
@ -3,11 +3,11 @@ from math import ceil
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
from numpy import array
|
from numpy import array, transpose
|
||||||
from retinaface.pre_trained_models import get_model
|
from retinaface.pre_trained_models import get_model
|
||||||
from retinaface.predict_single import Model
|
from retinaface.predict_single import Model
|
||||||
from retinaface.network import RetinaFace
|
from retinaface.network import RetinaFace
|
||||||
from torch import jit, randn, no_grad, Tensor, int64, tensor, onnx
|
from torch import jit, randn, no_grad, Tensor, int64, tensor, onnx, from_numpy
|
||||||
import albumentations as A
|
import albumentations as A
|
||||||
from torchinfo import summary
|
from torchinfo import summary
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
@ -21,13 +21,12 @@ from torchvision.utils import _log_api_usage_once
|
||||||
# model.eval()
|
# model.eval()
|
||||||
|
|
||||||
|
|
||||||
image = Image.open(
|
image = Image.open(r"C:\Users\tomokazu\すぐ消す\野中美希.jpg").convert(mode="RGB")
|
||||||
fp="/home/tomokazu/PycharmProjects/helloproject-ai/data/blog_images"
|
image_arr = from_numpy(np.array(object=image, dtype=np.float32)).unsqueeze(0).permute(0, 3, 1, 2)
|
||||||
"/稲場愛香/稲場愛香=juicejuice-official=12737097989-2.jpg").convert(mode="RGB")
|
|
||||||
image_arr = array(image)
|
|
||||||
max_size = 512
|
max_size = 512
|
||||||
|
|
||||||
example_input = randn(size=[1, 3, 256, 256]).float().cuda()
|
example_input = randn(size=[10, 3, 256, 256]).float()
|
||||||
|
|
||||||
retina_model = RetinaFace(
|
retina_model = RetinaFace(
|
||||||
name="Resnet50",
|
name="Resnet50",
|
||||||
|
|
@ -35,267 +34,31 @@ retina_model = RetinaFace(
|
||||||
return_layers={"layer2": 1, "layer3": 2, "layer4": 3},
|
return_layers={"layer2": 1, "layer3": 2, "layer4": 3},
|
||||||
in_channels=256,
|
in_channels=256,
|
||||||
out_channels=256,
|
out_channels=256,
|
||||||
).cuda()
|
).eval()
|
||||||
|
|
||||||
|
print(image_arr.size())
|
||||||
# onnx.export(
|
|
||||||
# model=retina_model, args=example_input, export_params=True, verbose=False, input_names=["input"],
|
torch.onnx.export(
|
||||||
# output_names=["bbox", "confidence", "landmark"],
|
model=retina_model,
|
||||||
# dynamic_axes={"input": {
|
args=example_input,
|
||||||
# 0: "batch_size",
|
f="retinaface.onnx",
|
||||||
# 2: "height",
|
input_names=["input"],
|
||||||
# 3: "width"
|
output_names=["bbox", "confidence", "landmark"],
|
||||||
# }, "bbox": {1: "bbox"}, "confidence": {1: "confidence"}, "landmark": {1: "landmark"}}, opset_version=16,
|
dynamic_axes={"input": {
|
||||||
# f="retinaface.onnx"
|
0: "batch_size",
|
||||||
# )
|
2: "height",
|
||||||
|
3: "width"
|
||||||
def pad_to_size(
|
},
|
||||||
target_size: Tuple[int, int],
|
"bbox": {0: "batch_size", 1: "length"},
|
||||||
image: np.array,
|
"confidence": {0: "batch_size", 1: "length"},
|
||||||
bboxes: Optional[np.ndarray] = None,
|
"landmark": {0: "batch_size", 1: "length"},
|
||||||
keypoints: Optional[np.ndarray] = None,
|
},
|
||||||
) -> Dict[str, Union[np.ndarray, Tuple[int, int, int, int]]]:
|
|
||||||
target_height, target_width = target_size
|
)
|
||||||
|
with no_grad():
|
||||||
image_height, image_width = image.shape[:2]
|
bbox_regressions, classifications, ldm_regressions = retina_model(image_arr)
|
||||||
|
print(bbox_regressions)
|
||||||
if target_width < image_width:
|
print(bbox_regressions.data[0][:, :2])
|
||||||
raise ValueError(f"Target width should bigger than image_width" f"We got {target_width} {image_width}")
|
print(bbox_regressions.size())
|
||||||
|
print(classifications.size())
|
||||||
if target_height < image_height:
|
print(ldm_regressions.size())
|
||||||
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)
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue