|
from pathlib import Path |
|
from typing import Optional, Union |
|
|
|
from PIL import Image, ImageDraw |
|
from torch import ge |
|
|
|
from internals.util.commons import download_file, download_image, safe_index |
|
from internals.util.config import get_root_dir |
|
from models.pose.body import Body |
|
|
|
|
|
class PoseDetector: |
|
__pose_model = ( |
|
"https://comic-assets.s3.ap-south-1.amazonaws.com/models/body_pose_model.pth" |
|
) |
|
|
|
__loaded = False |
|
|
|
def load(self): |
|
if self.__loaded: |
|
return |
|
|
|
pose_path = Path.home() / ".cache" / self.__pose_model.split("/")[-1] |
|
|
|
download_file(self.__pose_model, pose_path) |
|
|
|
self.body_estimation = Body(str(pose_path)) |
|
|
|
self.__loaded = True |
|
|
|
def transform( |
|
self, |
|
image: Union[str, Image.Image], |
|
width: int, |
|
height: int, |
|
client_coordinates: Optional[dict], |
|
) -> Image.Image: |
|
"Infer pose coordinates from image, map head and body coordinates to infered ones, create pose" |
|
self.load() |
|
if type(image) is str: |
|
image = download_image(image) |
|
|
|
infer_coordinates = self.infer(image, width, height) |
|
candidate_list = self.make_pose_from_subset( |
|
infer_coordinates["candidate"], infer_coordinates["subset"] |
|
) |
|
|
|
infer_coordinates["candidate"] = candidate_list[0] |
|
|
|
if client_coordinates and client_coordinates["candidate"]: |
|
client_coordinates = self.resize_coordinates( |
|
client_coordinates, 384, 384, width, height |
|
) |
|
infer_coordinates = self.map_coordinates( |
|
client_coordinates, infer_coordinates |
|
) |
|
|
|
print(infer_coordinates) |
|
|
|
return self.create_pose(infer_coordinates, width, height) |
|
|
|
def resize_coordinates( |
|
self, data: dict, ori_width, ori_height, new_width, new_height |
|
): |
|
points = data["candidate"] |
|
new_points = [] |
|
|
|
if new_width > new_height: |
|
ori_min = min(ori_width, ori_height) |
|
new_min = min(new_width, new_height) |
|
else: |
|
ori_min = max(ori_width, ori_height) |
|
new_min = max(new_width, new_height) |
|
|
|
for _, pair in enumerate(points): |
|
x = pair[0] * new_min / ori_min |
|
y = pair[1] * new_min / ori_min |
|
new_points.append([x, y]) |
|
|
|
return {"candidate": new_points, "subset": data["subset"]} |
|
|
|
def create_pose(self, data: dict, width: int, height: int) -> Image.Image: |
|
image = Image.new("RGBA", (width, height), "black") |
|
draw = ImageDraw.Draw(image) |
|
|
|
points: list = data["candidate"] |
|
for pair in self.__pose_logical_map: |
|
xy = safe_index(points, pair[0] - 1) |
|
x1y1 = safe_index(points, pair[1] - 1) |
|
|
|
if xy and x1y1: |
|
draw.line( |
|
(xy[0], xy[1], x1y1[0], x1y1[1]), |
|
fill=pair[2], |
|
width=4, |
|
) |
|
|
|
for i, point in enumerate(points): |
|
x = safe_index(point, 0) |
|
y = safe_index(point, 1) |
|
if x and y: |
|
draw.ellipse((x - 3, y - 3, x + 3, y + 3), fill=self.__points_color[i]) |
|
|
|
return image |
|
|
|
def infer(self, image: Union[str, Image.Image], width, height) -> dict: |
|
self.load() |
|
candidate = [] |
|
subset = [] |
|
|
|
if type(image) == str: |
|
image = download_image(image) |
|
|
|
image = image.resize((width, height)) |
|
|
|
candidate, subset = self.body_estimation.__call__(image) |
|
candidate = candidate.tolist() |
|
subset = subset.tolist() |
|
|
|
candidate = [item[:2] for item in candidate] |
|
|
|
return {"candidate": candidate, "subset": subset} |
|
|
|
def map_coordinates( |
|
self, client_coordinates: dict, infer_coordinates: dict |
|
) -> dict: |
|
client_points = client_coordinates["candidate"] |
|
infer_points = infer_coordinates["candidate"] |
|
|
|
c_neck = client_points[1] |
|
i_neck = infer_points[1] |
|
|
|
dx = i_neck[0] - c_neck[0] |
|
dy = i_neck[1] - c_neck[1] |
|
|
|
|
|
for i in range(len(client_points)): |
|
point = client_points[i - 1] |
|
infer_points[i - 1] = [point[0] + dx, point[1] + dy] |
|
|
|
return {"candidate": infer_points, "subset": infer_coordinates["subset"]} |
|
|
|
def make_pose_from_subset(self, candidate, subset): |
|
"Maps pose coordinates for subset" |
|
|
|
def make_pose_from_subset_item(candidate, subset_item): |
|
pose = [] |
|
for j in range(18): |
|
i = int(subset_item[j]) |
|
pose.append( |
|
None |
|
if i < 0 or not safe_index(candidate, i) |
|
else list(map(lambda x: x, candidate[i])) |
|
) |
|
return pose |
|
|
|
return list( |
|
map( |
|
lambda subset_item: make_pose_from_subset_item(candidate, subset_item), |
|
subset, |
|
) |
|
) |
|
|
|
def __convert_keypoints(self, keypoints): |
|
return [keypoints[i] for i in self.__kim] |
|
|
|
__kim = [0, 17, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3] |
|
__pose_logical_map = [ |
|
[2, 3, (255, 0, 0, 153)], |
|
[3, 4, (255, 85, 0, 153)], |
|
[4, 5, (255, 170, 0, 153)], |
|
[2, 6, (255, 255, 0, 153)], |
|
[6, 7, (170, 255, 0, 153)], |
|
[7, 8, (85, 255, 0, 153)], |
|
[2, 9, (0, 255, 0, 153)], |
|
[9, 10, (0, 255, 85, 153)], |
|
[10, 11, (0, 255, 170, 153)], |
|
[2, 12, (0, 255, 255, 153)], |
|
[12, 13, (0, 170, 255, 153)], |
|
[13, 14, (0, 85, 255, 153)], |
|
[2, 1, (0, 0, 255, 153)], |
|
[1, 15, (85, 0, 255, 153)], |
|
[15, 17, (170, 0, 255, 153)], |
|
[1, 16, (255, 0, 255, 153)], |
|
[16, 18, (255, 0, 170, 153)], |
|
] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
__points_color = [ |
|
"#ff0000", |
|
"#ff5600", |
|
"#ffaa01", |
|
"#ffff00", |
|
"#aaff03", |
|
"#53ff00", |
|
"#03ff00", |
|
"#03ff55", |
|
"#03ffaa", |
|
"#03ffff", |
|
"#05aaff", |
|
"#0055ff", |
|
"#0000ff", |
|
"#5500ff", |
|
"#aa00ff", |
|
"#ff00aa", |
|
"#ff00ff", |
|
"#ff0055", |
|
] |
|
|
