Some further optimizations.

2024-11-29 18:22:48 +01:00
parent ad1ef42b99
commit d50364cda7
3 changed files with 63 additions and 11 deletions
@@ -44,14 +44,12 @@ class SimCC(BaseModel):
        self.scale = 0

    def preprocess(self, image: np.ndarray):
-        tensor, self.dx, self.dy, self.scale = image, 0, 0, 1
-        tensor = tensor.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))
        return tensor

    def postprocess(self, tensor: List[np.ndarray]):
-        kpts = tensor[0][0]
-        scores = np.expand_dims(tensor[1][0], axis=-1)
-        keypoints = np.concatenate([kpts, scores], axis=-1)
-
+        keypoints = np.concatenate(
+            [tensor[0][0], np.expand_dims(tensor[1][0], axis=-1)], axis=-1
+        )
        return keypoints