crosstyan/RapidPoseTriangulation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Updated easypose scripts.

Browse Source
This commit is contained in:
Daniel
2024-12-06 17:35:49 +01:00
parent 108937d96c
commit ee8b9bafb3
4 changed files with 12 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
extras/easypose/base_model.py
View File

@ -29,10 +29,14 @@ class BaseModel(ABC):
self.input_shape = self.session.get_inputs()[0].shape self.input_shape = self.session.get_inputs()[0].shape
input_type = self.session.get_inputs()[0].type input_type = self.session.get_inputs()[0].type
if input_type == 'tensor(float16)': if input_type == 'tensor(float32)':
self.input_type = np.float16
else:
self.input_type = np.float32 self.input_type = np.float32
elif input_type == 'tensor(float16)':
self.input_type = np.float16
elif input_type == 'tensor(uint8)':
self.input_type = np.uint8
else:
raise ValueError('Unknown input type: ', input_type)
if warmup > 0: if warmup > 0:
self.warmup(warmup) self.warmup(warmup)

4
extras/easypose/detection.py
View File

@ -20,12 +20,12 @@ class RTMDet(BaseModel):
self.scale = 0 self.scale = 0
def preprocess(self, image: np.ndarray): def preprocess(self, image: np.ndarray):
th, tw = self.input_shape[2:] th, tw = self.input_shape[1:3]
image, self.dx, self.dy, self.scale = letterbox( image, self.dx, self.dy, self.scale = letterbox(
image, (tw, th), fill_value=114 image, (tw, th), fill_value=114
) )
tensor = np.asarray(image).astype(self.input_type, copy=False)[..., ::-1] tensor = np.asarray(image).astype(self.input_type, copy=False)[..., ::-1]
tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2)) tensor = np.expand_dims(tensor, axis=0)
return tensor return tensor
def postprocess(self, tensor: List[np.ndarray]): def postprocess(self, tensor: List[np.ndarray]):

5
extras/easypose/pose.py
View File

@ -39,13 +39,10 @@ class Heatmap(BaseModel):
class SimCC(BaseModel): class SimCC(BaseModel):
def __init__(self, model_path: str, device: str = 'CUDA', warmup: int = 30): def __init__(self, model_path: str, device: str = 'CUDA', warmup: int = 30):
super(SimCC, self).__init__(model_path, device, warmup) super(SimCC, self).__init__(model_path, device, warmup)
self.dx = 0
self.dy = 0
self.scale = 0
def preprocess(self, image: np.ndarray): def preprocess(self, image: np.ndarray):
tensor = np.asarray(image).astype(self.input_type, copy=False) tensor = np.asarray(image).astype(self.input_type, copy=False)
tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2)) tensor = np.expand_dims(tensor, axis=0)
return tensor return tensor
def postprocess(self, tensor: List[np.ndarray]): def postprocess(self, tensor: List[np.ndarray]):

7
scripts/utils_2d_pose_ep.py → extras/easypose/utils_2d_pose_ep.py
View File

@ -14,13 +14,10 @@ filepath = os.path.dirname(os.path.realpath(__file__)) + "/"
def load_model(): def load_model():
print("Loading mmpose model ...") print("Loading mmpose model ...")
# model = ep.TopDown("rtmpose_m", "SimCC", "rtmdet_s")
model = ep.TopDown( model = ep.TopDown(
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_with-norm.onnx", "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_fp16_extra-steps.onnx",
# "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_fp16_with-norm.onnx",
"SimCC", "SimCC",
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_with-norm.onnx", "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_fp16_extra-steps.onnx",
# "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_fp16_with-norm.onnx",
conf_threshold=0.3, conf_threshold=0.3,
iou_threshold=0.3, iou_threshold=0.3,
warmup=10, warmup=10,