diff --git a/clashroyalebuildabot/detectors/number_detector.py b/clashroyalebuildabot/detectors/number_detector.py index 373a953..5770a7d 100644 --- a/clashroyalebuildabot/detectors/number_detector.py +++ b/clashroyalebuildabot/detectors/number_detector.py @@ -70,11 +70,7 @@ def _post_process(self, pred): return clean_pred def _preprocess(self, image): - image = self.resize(image) - image = np.array(image, dtype=np.float32) - image, padding = self.pad(image) - image = image.transpose(2, 0, 1) - image /= 255 + image, padding = self.resize_pad_transpose_and_scale(image) return image, padding def run(self, image): @@ -86,8 +82,7 @@ def run(self, image): crops.append(crop) paddings.append(padding) - crops = np.array(crops, dtype=np.float16) - preds = self._infer(crops).astype(np.float32) + preds = self._infer(crops) for i, padding in enumerate(paddings): preds[i] = self.fix_bboxes( diff --git a/clashroyalebuildabot/detectors/onnx_detector.py b/clashroyalebuildabot/detectors/onnx_detector.py index bba967a..f3bf924 100644 --- a/clashroyalebuildabot/detectors/onnx_detector.py +++ b/clashroyalebuildabot/detectors/onnx_detector.py @@ -49,6 +49,14 @@ def pad(self, x): ) return x, padding + def resize_pad_transpose_and_scale(self, image): + image = self.resize(image) + image = np.array(image, dtype=np.float16) + image, padding = self.pad(image) + image = image.transpose(2, 0, 1) + image /= 255 + return image, padding + def fix_bboxes(self, x, width, height, padding): x[:, [0, 2]] -= padding[0] x[:, [1, 3]] -= padding[2] diff --git a/clashroyalebuildabot/detectors/unit_detector.py b/clashroyalebuildabot/detectors/unit_detector.py index 1fd1bbe..9d77986 100644 --- a/clashroyalebuildabot/detectors/unit_detector.py +++ b/clashroyalebuildabot/detectors/unit_detector.py @@ -67,11 +67,8 @@ def _preprocess(self, image): self.UNIT_Y_END * image.height, ) ) - image = self.resize(image) - image = np.array(image, dtype=np.float32) - image, padding = self.pad(image) - image = np.expand_dims(image.transpose(2, 0, 1), axis=0) - image /= 255 + image, padding = self.resize_pad_transpose_and_scale(image) + image = np.expand_dims(image, axis=0) return image, padding def _post_process(self, pred, height, image): @@ -99,7 +96,7 @@ def _post_process(self, pred, height, image): def run(self, image): height, width = image.height, image.width np_image, padding = self._preprocess(image) - pred = self._infer(np_image.astype(np.float16)).astype(np.float32)[0] + pred = self._infer(np_image)[0] pred = pred[pred[:, 4] > self.MIN_CONF] pred = self.fix_bboxes(pred, width, height, padding) pred = self._post_process(pred, height, image)