Fixed running custom onnx models.

2024-11-29 15:18:57 +01:00
parent f6d13ea5a7
commit 93d4611a91
6 changed files with 158 additions and 32 deletions
@@ -44,10 +44,12 @@ class SimCC(BaseModel):
        self.scale = 0

    def preprocess(self, image: np.ndarray):
-        th, tw = self.input_shape[2:]
-        image, self.dx, self.dy, self.scale = letterbox(image, (tw, th))
-        tensor = (image - np.array((103.53, 116.28, 123.675))) / np.array((57.375, 57.12, 58.395))
-        tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2)).astype(np.float32)
+        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)
+        )
        return tensor

    def postprocess(self, tensor: List[np.ndarray]):