From e6bd5af6a8e306a1cdef63402a77a980a04ad6e1 Mon Sep 17 00:00:00 2001
From: Grafting Rayman <156515434+GraftingRayman@users.noreply.github.com>
Date: Fri, 17 Jan 2025 11:06:44 +0000
Subject: Add files via upload

---
 scripts/r_masking/segs.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)
 create mode 100644 scripts/r_masking/segs.py

(limited to 'scripts/r_masking/segs.py')

diff --git a/scripts/r_masking/segs.py b/scripts/r_masking/segs.py
new file mode 100644
index 0000000..60c22d7
--- /dev/null
+++ b/scripts/r_masking/segs.py
@@ -0,0 +1,22 @@
+def filter(segs, labels):
+    labels = set([label.strip() for label in labels])
+
+    if 'all' in labels:
+        return (segs, (segs[0], []), )
+    else:
+        res_segs = []
+        remained_segs = []
+
+        for x in segs[1]:
+            if x.label in labels:
+                res_segs.append(x)
+            elif 'eyes' in labels and x.label in ['left_eye', 'right_eye']:
+                res_segs.append(x)
+            elif 'eyebrows' in labels and x.label in ['left_eyebrow', 'right_eyebrow']:
+                res_segs.append(x)
+            elif 'pupils' in labels and x.label in ['left_pupil', 'right_pupil']:
+                res_segs.append(x)
+            else:
+                remained_segs.append(x)
+
+    return ((segs[0], res_segs), (segs[0], remained_segs), )
-- 
cgit v1.2.3