From 5c261fe7a88cede95bfee38d67c8633abf62dfbc Mon Sep 17 00:00:00 2001
From: changhy666 <57101883+changhy666@users.noreply.github.com>
Date: Thu, 14 Oct 2021 21:03:22 +0800
Subject: [PATCH] Update aux_functions.py

support that pattern is a dir.
---
 utils/aux_functions.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/utils/aux_functions.py b/utils/aux_functions.py
index df556ba..5c218e3 100644
--- a/utils/aux_functions.py
+++ b/utils/aux_functions.py
@@ -318,7 +318,17 @@ def mask_face(image, face_location, six_points, angle, args, type="surgical"):
     # Process the mask if necessary
     if args.pattern:
         # Apply pattern to mask
-        img = texture_the_mask(img, args.pattern, args.pattern_weight)
+        if os.path.isfile(args.pattern):
+            img = texture_the_mask(img, args.pattern, args.pattern_weight)
+        else:
+            patterns_list = []
+            for root, dirs, files in os.walk(args.pattern):
+                if len(files) != 0 and len(dirs) == 0:
+                    for item in files:
+                        pattern_path = os.path.join(root, item)
+                        patterns_list.append(pattern_path)
+            choice_pattern = random.choice(patterns_list)
+            img = texture_the_mask(img, choice_pattern, args.pattern_weight)
 
     if args.color:
         # Apply color to mask