Update utils_mask.py

Combines upper_clothes, skirt, and dress parts to simulate saree

utils_mask.py

@@ -100,6 +100,15 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
                              (parse_array == 14).astype(np.float32) + \
                              (parse_array == 15).astype(np.float32)
         parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
+    elif category == 'saree':
+    # Include upper clothes, dress, and skirt for saree draping
+    parse_mask = (parse_array == 4).astype(np.float32) + \
+                 (parse_array == 5).astype(np.float32) + \
+                 (parse_array == 7).astype(np.float32)
+    parser_mask_fixed += (parse_array == label_map["pants"]).astype(np.float32)
+    parser_mask_fixed += arms_left + arms_right
+    parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
+    
     else:
         raise NotImplementedError
 
@@ -112,7 +121,7 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
     im_arms_right = Image.new('L', (width, height))
     arms_draw_left = ImageDraw.Draw(im_arms_left)
     arms_draw_right = ImageDraw.Draw(im_arms_right)
-    if category == 'dresses' or category == 'upper_body':
+    if category in ['dresses', 'upper_body', 'saree']:
         shoulder_right = np.multiply(tuple(pose_data[2][:2]), height / 512.0)
         shoulder_left = np.multiply(tuple(pose_data[5][:2]), height / 512.0)
         elbow_right = np.multiply(tuple(pose_data[3][:2]), height / 512.0)
@@ -145,7 +154,7 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
 
     parser_mask_fixed = np.logical_or(parser_mask_fixed, parse_head)
     parse_mask = cv2.dilate(parse_mask, np.ones((5, 5), np.uint16), iterations=5)
-    if category == 'dresses' or category == 'upper_body':
+    if category in ['dresses', 'upper_body', 'saree']:
         neck_mask = (parse_array == 18).astype(np.float32)
         neck_mask = cv2.dilate(neck_mask, np.ones((5, 5), np.uint16), iterations=1)
         neck_mask = np.logical_and(neck_mask, np.logical_not(parse_head))