mevCJ commited on
Commit
f62d400
1 Parent(s): dc8b3bf

Segment upper clothes only

Browse files
Files changed (1) hide show
  1. process.py +3 -3
process.py CHANGED
@@ -124,9 +124,9 @@ def generate_mask(input_image, net, palette, device = 'cpu'):
124
  classes_to_save = []
125
 
126
  # Check which classes are present in the image
127
- for cls in range(1, 4): # Exclude background class (0)
128
- if np.any(output_arr == cls):
129
- classes_to_save.append(cls)
130
 
131
  # Save alpha masks
132
  for cls in classes_to_save:
 
124
  classes_to_save = []
125
 
126
  # Check which classes are present in the image
127
+ #for cls in range(1, 4): # Exclude background class (0)
128
+ if np.any(output_arr == 1):
129
+ classes_to_save.append(1)
130
 
131
  # Save alpha masks
132
  for cls in classes_to_save: