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