From 00cd32dfdcf20892e34444a5da6dc2413d29d431 Mon Sep 17 00:00:00 2001 From: 283375 Date: Wed, 20 Mar 2024 15:53:10 +0800 Subject: [PATCH] feat: ETERNAL rating class support --- src/arcaea_offline_ocr/device/ocr.py | 1 + src/arcaea_offline_ocr/device/rois/masker/auto.py | 14 ++++++++++++++ .../device/rois/masker/common.py | 4 ++++ 3 files changed, 19 insertions(+) diff --git a/src/arcaea_offline_ocr/device/ocr.py b/src/arcaea_offline_ocr/device/ocr.py index f4c637c..85f74ee 100644 --- a/src/arcaea_offline_ocr/device/ocr.py +++ b/src/arcaea_offline_ocr/device/ocr.py @@ -80,6 +80,7 @@ class DeviceOcr: self.masker.rating_class_prs(roi), self.masker.rating_class_ftr(roi), self.masker.rating_class_byd(roi), + self.masker.rating_class_etr(roi), ] return max(enumerate(results), key=lambda i: np.count_nonzero(i[1]))[0] diff --git a/src/arcaea_offline_ocr/device/rois/masker/auto.py b/src/arcaea_offline_ocr/device/rois/masker/auto.py index 29cb4ab..f3894a4 100644 --- a/src/arcaea_offline_ocr/device/rois/masker/auto.py +++ b/src/arcaea_offline_ocr/device/rois/masker/auto.py @@ -34,6 +34,9 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto): BYD_HSV_MIN = np.array([170, 50, 50], np.uint8) BYD_HSV_MAX = np.array([179, 210, 198], np.uint8) + ETR_HSV_MIN = np.array([130, 60, 80], np.uint8) + ETR_HSV_MAX = np.array([140, 145, 180], np.uint8) + TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8) TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8) @@ -87,6 +90,10 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto): def rating_class_byd(cls, roi_bgr: Mat) -> Mat: return cls.mask_bgr_in_hsv(roi_bgr, cls.BYD_HSV_MIN, cls.BYD_HSV_MAX) + @classmethod + def rating_class_etr(cls, roi_bgr: Mat) -> Mat: + return cls.mask_bgr_in_hsv(roi_bgr, cls.ETR_HSV_MIN, cls.ETR_HSV_MAX) + @classmethod def max_recall(cls, roi_bgr: Mat) -> Mat: return cls.gray(roi_bgr) @@ -135,6 +142,9 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto): BYD_HSV_MIN = np.array([170, 50, 50], np.uint8) BYD_HSV_MAX = np.array([179, 210, 198], np.uint8) + ETR_HSV_MIN = np.array([130, 60, 80], np.uint8) + ETR_HSV_MAX = np.array([140, 145, 180], np.uint8) + MAX_RECALL_HSV_MIN = np.array([125, 0, 0], np.uint8) MAX_RECALL_HSV_MAX = np.array([145, 100, 150], np.uint8) @@ -186,6 +196,10 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto): def rating_class_byd(cls, roi_bgr: Mat) -> Mat: return cls.mask_bgr_in_hsv(roi_bgr, cls.BYD_HSV_MIN, cls.BYD_HSV_MAX) + @classmethod + def rating_class_etr(cls, roi_bgr: Mat) -> Mat: + return cls.mask_bgr_in_hsv(roi_bgr, cls.ETR_HSV_MIN, cls.ETR_HSV_MAX) + @classmethod def max_recall(cls, roi_bgr: Mat) -> Mat: return cls.mask_bgr_in_hsv( diff --git a/src/arcaea_offline_ocr/device/rois/masker/common.py b/src/arcaea_offline_ocr/device/rois/masker/common.py index cb1223f..cc63b88 100644 --- a/src/arcaea_offline_ocr/device/rois/masker/common.py +++ b/src/arcaea_offline_ocr/device/rois/masker/common.py @@ -34,6 +34,10 @@ class DeviceRoisMasker: def rating_class_byd(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() + @classmethod + def rating_class_etr(cls, roi_bgr: Mat) -> Mat: + raise NotImplementedError() + @classmethod def max_recall(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError()