dev(narugo): complete visual part for dwpose (87daa6c7) · Commits · git-mirror / Imgutils

imgutils/pose/format.py

+21 −1

Original line number	Diff line number	Diff line
		from enum import IntEnum, unique

		import numpy as np

		OP18_BODY_MIN = 0
		OP18_BODY_MAX = 17

		@@ -50,7 +52,7 @@ class OpenPose18(IntEnum):

		class OP18KeyPointSet:
		def __init__(self, all_):
		self.all = all_
		self.all: np.ndarray = all_

		@property
		def body(self):
		@@ -79,3 +81,21 @@ class OP18KeyPointSet:
		@property
		def right_hand(self):
		return self.all[OP18_RIGHT_HAND_MIN:OP18_RIGHT_HAND_MAX + 1]

		def __mul__(self, multiplier):
		if isinstance(multiplier, (float, int)):
		new_all = self.all.copy()
		new_all[:, 0] *= multiplier
		new_all[:, 1] *= multiplier
		return OP18KeyPointSet(new_all)
		else:
		raise TypeError(f'Invalid type of multiplier - {multiplier!r}.')

		def __truediv__(self, divisor):
		if isinstance(divisor, (float, int)):
		new_all = self.all.copy()
		new_all[:, 0] /= divisor
		new_all[:, 1] /= divisor
		return OP18KeyPointSet(new_all)
		else:
		raise TypeError(f'Invalid type of divisor - {divisor!r}.')

imgutils/pose/visual.py

+71 −2

Original line number	Diff line number	Diff line
		from typing import List
		import colorsys
		from typing import List, Optional

		import numpy as np
		from PIL import ImageDraw, Image

		from imgutils.data import ImageTyping, load_image
		@@ -26,12 +28,79 @@ def _op18_body(keypoints: OP18KeyPointSet, draw: ImageDraw.ImageDraw, threshold:
		draw.line([(x0, y0), (x1, y1)], width=5, fill=color)


		_OP18_HAND_EDGES = [
		[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10],
		[10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]
		]


		# noinspection DuplicatedCode
		def _op18_hands(keypoints: OP18KeyPointSet, draw: ImageDraw.ImageDraw, threshold: float = 0.3):
		for hands in [keypoints.left_hand, keypoints.right_hand]:
		for ie, (e0, e1) in enumerate(_OP18_HAND_EDGES):
		x0, y0, s0 = hands[e0]
		x1, y1, s1 = hands[e1]
		if s0 >= threshold and s1 >= threshold:
		color = (np.array(colorsys.hsv_to_rgb(ie / float(len(_OP18_HAND_EDGES)), 1.0, 1.0)) * 255) \
		.astype(np.uint8)
		draw.line([(x0, y0), (x1, y1)], width=3, fill=tuple(color.tolist()))


		_OP18_FOOT_EDGES = [(2, 0), (2, 1)]


		# noinspection DuplicatedCode
		def _op18_feet(keypoints: OP18KeyPointSet, draw: ImageDraw.ImageDraw, threshold: float = 0.3):
		for hands in [keypoints.left_foot, keypoints.right_foot]:
		for ie, (e0, e1) in enumerate(_OP18_FOOT_EDGES):
		x0, y0, s0 = hands[e0]
		x1, y1, s1 = hands[e1]
		if s0 >= threshold and s1 >= threshold:
		color = (np.array(colorsys.hsv_to_rgb(ie / float(len(_OP18_FOOT_EDGES)), 1.0, 1.0)) * 255) \
		.astype(np.uint8)
		draw.line([(x0, y0), (x1, y1)], width=3, fill=tuple(color.tolist()))


		_FACE_POINT_SIZE = 4.5


		def _op18_face(keypoints: OP18KeyPointSet, draw: ImageDraw.ImageDraw, threshold: float = 0.3):
		for x0, y0, s0 in keypoints.face:
		if s0 >= threshold:
		draw.ellipse(
		(
		x0 - _FACE_POINT_SIZE / 2, y0 - _FACE_POINT_SIZE / 2,
		x0 + _FACE_POINT_SIZE / 2, y0 + _FACE_POINT_SIZE / 2,
		),
		fill='#00ff00', outline='#00ff00',
		)


		def op18_visualize(image: ImageTyping, keypoints_list: List[OP18KeyPointSet], threshold: float = 0.3,
		draw_body: bool = True) -> Image.Image:
		min_edge_size: Optional[int] = 512, draw_body: bool = True, draw_hands: bool = True,
		draw_feet: bool = True, draw_face: bool = True) -> Image.Image:
		image = load_image(image, force_background='white', mode='RGB')
		if min_edge_size is not None and min(image.width, image.height) > min_edge_size:
		r = min(image.width, image.height) / min_edge_size
		else:
		r = None

		if r is not None:
		new_width = int(image.width / r)
		new_height = int(image.height / r)
		image = image.resize((new_width, new_height))

		draw = ImageDraw.Draw(image)
		for keypoints in keypoints_list:
		if r is not None:
		keypoints /= r
		if draw_body:
		_op18_body(keypoints, draw, threshold)
		if draw_hands:
		_op18_hands(keypoints, draw, threshold)
		if draw_feet:
		_op18_feet(keypoints, draw, threshold)
		if draw_face:
		_op18_face(keypoints, draw, threshold)

		return image

test/pose/init.py

0 → 100644

+0 −0

Empty file added.