From cb961ed0d6896309336159bda34ed2f75fb60a84 Mon Sep 17 00:00:00 2001
From: carp <25677564+carp@users.noreply.github.com>
Date: Mon, 13 Jul 2020 13:27:55 -0400
Subject: add face detection
---
anime-face-detector/nms_wrapper.py | 29 +++++++++++++++++++++++++++++
1 file changed, 29 insertions(+)
create mode 100644 anime-face-detector/nms_wrapper.py
(limited to 'anime-face-detector/nms_wrapper.py')
diff --git a/anime-face-detector/nms_wrapper.py b/anime-face-detector/nms_wrapper.py
new file mode 100644
index 0000000..ca900e8
--- /dev/null
+++ b/anime-face-detector/nms_wrapper.py
@@ -0,0 +1,29 @@
+from enum import Enum
+
+
+class NMSType(Enum):
+ PY_NMS = 1
+ CPU_NMS = 2
+ GPU_NMS = 3
+
+
+default_nms_type = NMSType.PY_NMS
+
+
+class NMSWrapper:
+ def __init__(self, nms_type=default_nms_type):
+ assert type(nms_type) == NMSType
+ if nms_type == NMSType.PY_NMS:
+ from nms.py_cpu_nms import py_cpu_nms
+ self._nms = py_cpu_nms
+ elif nms_type == NMSType.CPU_NMS:
+ from nms.cpu_nms import cpu_nms
+ self._nms = cpu_nms
+ elif nms_type == NMSType.GPU_NMS:
+ from nms.gpu_nms import gpu_nms
+ self._nms = gpu_nms
+ else:
+ raise ValueError('current nms type is not implemented yet')
+
+ def __call__(self, *args, **kwargs):
+ return self._nms(*args, **kwargs)
--
cgit v1.2.3