diff --git a/extras/easypose/detection.py b/extras/easypose/detection.py index a57c524..0c29a2c 100644 --- a/extras/easypose/detection.py +++ b/extras/easypose/detection.py @@ -19,25 +19,26 @@ class RTMDet(BaseModel): self.dy = 0 self.scale = 0 + norm_mean = -1 * np.array([123.675, 116.28, 103.53]) + norm_std = 1.0 / np.array([58.395, 57.12, 57.375]) + self.norm_mean = np.reshape(norm_mean, (1, 1, 3)).astype(np.float32) + self.norm_std = np.reshape(norm_std, (1, 1, 3)).astype(np.float32) + def preprocess(self, image: np.ndarray): th, tw = self.input_shape[2:] tensor, self.dx, self.dy, self.scale = letterbox( image, (tw, th), fill_value=114 ) - tensor -= np.array((123.675, 116.28, 103.53)) - tensor /= np.array((58.395, 57.12, 57.375)) + tensor = tensor.astype(np.float32, copy=False) + tensor += self.norm_mean + tensor *= self.norm_std tensor = tensor[..., ::-1] - tensor = ( - np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2)).astype(np.float32) - ) + tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2)) return tensor def postprocess(self, tensor: List[np.ndarray]): - boxes = tensor[0] - classes = tensor[1] - boxes = np.squeeze(boxes, axis=0) - classes = np.squeeze(classes, axis=0) - classes = np.expand_dims(classes, axis=-1) + boxes = np.squeeze(tensor[0], axis=0) + classes = np.expand_dims(np.squeeze(tensor[1], axis=0), axis=-1) boxes = np.concatenate([boxes, classes], axis=-1) boxes = nms(boxes, self.iou_threshold, self.conf_threshold) diff --git a/extras/easypose/pipeline.py b/extras/easypose/pipeline.py index 31bf8db..333e46f 100644 --- a/extras/easypose/pipeline.py +++ b/extras/easypose/pipeline.py @@ -199,7 +199,8 @@ class TopDown: # See: /mmpose/models/pose_estimators/topdown.py - add_pred_to_datasample() th, tw = region.shape[:2] bw, bh = [p.box[2] - p.box[0], p.box[3] - p.box[1]] - kp[:, :2] = kp[:, :2] / np.array([tw, th]) * np.array([bw, bh]) + kp[:, :2] /= np.array([tw, th]) + kp[:, :2] *= np.array([bw, bh]) kp[:, :2] += np.array([p.box[0] + bw / 2, p.box[1] + bh / 2]) kp[:, :2] -= 0.5 * np.array([bw, bh]) diff --git a/extras/easypose/pose.py b/extras/easypose/pose.py index be8a29c..677493b 100644 --- a/extras/easypose/pose.py +++ b/extras/easypose/pose.py @@ -43,13 +43,17 @@ class SimCC(BaseModel): self.dy = 0 self.scale = 0 + norm_mean = -1 * np.array([123.675, 116.28, 103.53]) + norm_std = 1.0 / np.array([58.395, 57.12, 57.375]) + self.norm_mean = np.reshape(norm_mean, (1, 1, 3)).astype(np.float32) + self.norm_std = np.reshape(norm_std, (1, 1, 3)).astype(np.float32) + def preprocess(self, image: np.ndarray): tensor, self.dx, self.dy, self.scale = image, 0, 0, 1 - tensor -= np.array((123.675, 116.28, 103.53)) - tensor /= np.array((58.395, 57.12, 57.375)) - tensor = ( - np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2)).astype(np.float32) - ) + tensor = tensor.astype(np.float32, copy=False) + tensor += self.norm_mean + tensor *= self.norm_std + tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2)) return tensor def postprocess(self, tensor: List[np.ndarray]): diff --git a/extras/easypose/utils.py b/extras/easypose/utils.py index bb42582..b81ca8c 100644 --- a/extras/easypose/utils.py +++ b/extras/easypose/utils.py @@ -11,13 +11,10 @@ def letterbox(img: np.ndarray, target_size: Sequence[int], fill_value: int = 128 scale = min(tw / w, th / h) nw, nh = int(w * scale), int(h * scale) - - resized_img = cv2.resize(img, (nw, nh)) + dx, dy = (tw - nw) // 2, (th - nh) // 2 canvas = np.full((th, tw, img.shape[2]), fill_value, dtype=img.dtype) - - dx, dy = (tw - nw) // 2, (th - nh) // 2 - canvas[dy:dy + nh, dx:dx + nw, :] = resized_img + canvas[dy:dy + nh, dx:dx + nw, :] = cv2.resize(img, (nw, nh)) return canvas, dx, dy, scale