Add files via upload

Author: 1038lab

AILab_BiRefNet.py

@@ -350,8 +350,8 @@ def INPUT_TYPES(s):
             }
         }
 
-    RETURN_TYPES = ("IMAGE", "MASK")
-    RETURN_NAMES = ("IMAGE", "MASK")
+    RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
+    RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
     FUNCTION = "process_image"
     CATEGORY = "🧪AILab/🧽RMBG"
 
@@ -447,7 +447,16 @@ def process_image(self, image, model, **params):
 
                 processed_masks.append(pil2tensor(mask))
 
-            return (torch.cat(processed_images, dim=0), torch.cat(processed_masks, dim=0))
+            # Create mask image for visualization
+            mask_images = []
+            for mask_tensor in processed_masks:
+                # Convert mask to RGB image format for visualization
+                mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
+                mask_images.append(mask_image)
+            
+            mask_image_output = torch.cat(mask_images, dim=0)
+
+            return (torch.cat(processed_images, dim=0), torch.cat(processed_masks, dim=0), mask_image_output)
 
         except Exception as e:
             handle_model_error(f"Error in image processing: {str(e)}")

AILab_BodySegment.py

@@ -85,8 +85,8 @@ def INPUT_TYPES(cls):
             },
         }
 
-    RETURN_TYPES = ("IMAGE", "MASK")
-    RETURN_NAMES = ("IMAGE", "MASK")
+    RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
+    RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
     FUNCTION = "segment_body"
     CATEGORY = "🧪AILab/🧽RMBG"
 
@@ -217,11 +217,20 @@ def segment_body(self, images, mask_blur=0, mask_offset=0, background_color="Alp
                 batch_tensor.append(result_image)
                 batch_masks.append(pil2tensor(mask_image))
 
+            # Create mask image for visualization
+            mask_images = []
+            for mask_tensor in batch_masks:
+                # Convert mask to RGB image format for visualization
+                mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
+                mask_images.append(mask_image)
+            
+            mask_image_output = torch.cat(mask_images, dim=0)
+            
             # Prepare final output
             batch_tensor = torch.cat(batch_tensor, dim=0)
             batch_masks = torch.cat(batch_masks, dim=0)
 
-            return (batch_tensor, batch_masks)
+            return (batch_tensor, batch_masks, mask_image_output)
 
         except Exception as e:
             self.clear_model()

AILab_ClothSegment.py

@@ -85,8 +85,8 @@ def INPUT_TYPES(cls):
             },
         }
 
-    RETURN_TYPES = ("IMAGE", "MASK")
-    RETURN_NAMES = ("IMAGE", "MASK")
+    RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
+    RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
     FUNCTION = "segment_clothes"
     CATEGORY = "🧪AILab/🧽RMBG"
 
@@ -253,11 +253,20 @@ def segment_clothes(self, images, process_res=1024, mask_blur=0, mask_offset=0,
                     batch_tensor.append(result_image)
                     batch_masks.append(pil2tensor(mask_image))
 
+            # Create mask image for visualization
+            mask_images = []
+            for mask_tensor in batch_masks:
+                # Convert mask to RGB image format for visualization
+                mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
+                mask_images.append(mask_image)
+            
+            mask_image_output = torch.cat(mask_images, dim=0)
+            
             # Prepare final output
             batch_tensor = torch.cat(batch_tensor, dim=0)
             batch_masks = torch.cat(batch_masks, dim=0)
 
-            return (batch_tensor, batch_masks)
+            return (batch_tensor, batch_masks, mask_image_output)
 
         except Exception as e:
             self.clear_model()