Spaces:
Running
Running
Segment upper clothes only
Browse files- 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 |
-
|
129 |
-
|
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:
|