diff --git a/AILab_FashionSegment.py b/AILab_FashionSegment.py index 93856e7..8e05bbe 100644 --- a/AILab_FashionSegment.py +++ b/AILab_FashionSegment.py @@ -236,8 +236,10 @@ class FashionSegmentClothing: except Exception as e: return False, f"Error downloading model files: {str(e)}" - def segment_fashion(self, images, accessories_options, process_res=512, mask_blur=0, mask_offset=0, + def segment_fashion(self, images, accessories_options=None, process_res=512, mask_blur=0, mask_offset=0, background="Alpha", background_color="#222222", invert_output=False, **class_selections): + if accessories_options is None: + accessories_options = [] try: # Check and download model cache_status, message = self.check_model_cache() @@ -360,7 +362,7 @@ class FashionSegmentClothing: self.clear_model() raise RuntimeError(f"Error in fashion segmentation: {str(e)}") finally: - if not self.model.training: + if self.model is not None and not self.model.training: self.clear_model() def __del__(self): diff --git a/AILab_SegmentV2.py b/AILab_SegmentV2.py index f998f53..c5edd12 100644 --- a/AILab_SegmentV2.py +++ b/AILab_SegmentV2.py @@ -16,7 +16,6 @@ from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection from AILab_ImageMaskTools import pil2tensor, tensor2pil -# SAM model definitions (6 models) SAM_MODELS = { "sam_vit_h (2.56GB)": { "model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_vit_h.pth", @@ -50,7 +49,6 @@ SAM_MODELS = { } } -# GroundingDINO model definitions (2 models) DINO_MODELS = { "GroundingDINO_SwinT_OGC (694MB)": { "config_url": "https://huggingface.co/1038lab/GroundingDINO/resolve/main/GroundingDINO_SwinT_OGC.cfg.py", @@ -97,7 +95,6 @@ def apply_background_color(image: Image.Image, mask_image: Image.Image, background_color: str = "#222222") -> Image.Image: rgba_image = image.copy().convert('RGBA') rgba_image.putalpha(mask_image.convert('L')) - if background == "Color": def hex_to_rgba(hex_color): hex_color = hex_color.lstrip('#') @@ -172,25 +169,18 @@ class SegmentV2: background_color="#222222", invert_output=False): device = "cuda" if torch.cuda.is_available() else "cpu" - # 处理批量图像 batch_size = image.shape[0] if len(image.shape) == 4 else 1 if len(image.shape) == 3: image = image.unsqueeze(0) - result_images = [] result_masks = [] result_mask_images = [] - for b in range(batch_size): img_pil = tensor2pil(image[b]) img_np = np.array(img_pil.convert("RGB")) - - # Load GroundingDINO config and weights dino_info = DINO_MODELS[dino_model] config_path = get_or_download_model_file(dino_info["config_filename"], dino_info["config_url"], "grounding-dino") weights_path = get_or_download_model_file(dino_info["model_filename"], dino_info["model_url"], "grounding-dino") - - # Load and cache GroundingDINO model dino_key = (config_path, weights_path, device) if dino_key not in self.dino_model_cache: args = SLConfig.fromfile(config_path) @@ -201,12 +191,8 @@ class SegmentV2: model.to(device) self.dino_model_cache[dino_key] = model dino = self.dino_model_cache[dino_key] - - # Download/check SAM weights sam_info = SAM_MODELS[sam_model] sam_ckpt_path = get_or_download_model_file(sam_info["filename"], sam_info["model_url"], "SAM") - - # Load SAM model (cache to avoid reloading) sam_key = (sam_info["model_type"], sam_ckpt_path, device) if sam_key not in self.sam_model_cache: try: @@ -230,8 +216,6 @@ class SegmentV2: else: raise e predictor = self.sam_model_cache[sam_key] - - # Preprocess image for DINO from groundingdino.datasets.transforms import Compose, RandomResize, ToTensor, Normalize transform = Compose([ RandomResize([800], max_size=1333), @@ -240,21 +224,13 @@ class SegmentV2: ]) image_tensor, _ = transform(img_pil.convert("RGB"), None) image_tensor = image_tensor.unsqueeze(0).to(device) - - # Prepare text prompt text_prompt = prompt if prompt.endswith(".") else prompt + "." - - # Forward pass with torch.no_grad(): outputs = dino(image_tensor, captions=[text_prompt]) logits = outputs["pred_logits"].sigmoid()[0] boxes = outputs["pred_boxes"][0] - - # Filter boxes by threshold filt_mask = logits.max(dim=1)[0] > threshold boxes_filt = boxes[filt_mask] - - # Handle case with no detected boxes if boxes_filt.shape[0] == 0: width, height = img_pil.size empty_mask = torch.zeros((1, height, width), dtype=torch.float32, device="cpu") @@ -264,14 +240,10 @@ class SegmentV2: result_masks.append(empty_mask) result_mask_images.append(empty_mask_rgb) continue - - # Convert boxes to xyxy H, W = img_pil.size[1], img_pil.size[0] boxes_xyxy = box_ops.box_cxcywh_to_xyxy(boxes_filt) boxes_xyxy = boxes_xyxy * torch.tensor([W, H, W, H], dtype=torch.float32, device=boxes_xyxy.device) boxes_xyxy = boxes_xyxy.cpu().numpy() - - # Use SAM to get masks for each box predictor.set_image(img_np) boxes_tensor = torch.tensor(boxes_xyxy, dtype=torch.float32, device=predictor.device) transformed_boxes = predictor.transform.apply_boxes_torch(boxes_tensor, img_np.shape[:2]) @@ -281,38 +253,27 @@ class SegmentV2: boxes=transformed_boxes, multimask_output=False ) - - # Combine all masks into one - combined_mask = torch.max(masks, dim=0)[0] # Take maximum across all masks + combined_mask = torch.max(masks, dim=0)[0] mask = combined_mask.float().cpu().numpy() mask = mask.squeeze(0) mask = (mask * 255).astype(np.uint8) mask_pil = Image.fromarray(mask, mode="L") - - # Process mask and apply background mask_image = process_mask(mask_pil, invert_output, mask_blur, mask_offset) result_image = apply_background_color(img_pil, mask_image, background, background_color) if background == "Color": result_image = result_image.convert("RGB") else: result_image = result_image.convert("RGBA") - - # Convert to tensors mask_tensor = torch.from_numpy(np.array(mask_image).astype(np.float32) / 255.0).unsqueeze(0) mask_image_vis = mask_tensor.reshape((-1, 1, mask_image.height, mask_image.width)).movedim(1, -1).expand(-1, -1, -1, 3) - result_images.append(pil2tensor(result_image)) result_masks.append(mask_tensor) result_mask_images.append(mask_image_vis) - - # 如果没有成功处理任何图像，返回空结果 if len(result_images) == 0: width, height = tensor2pil(image[0]).size empty_mask = torch.zeros((batch_size, 1, height, width), dtype=torch.float32, device="cpu") empty_mask_rgb = empty_mask.reshape((-1, 1, height, width)).movedim(1, -1).expand(-1, -1, -1, 3) return (image, empty_mask, empty_mask_rgb) - - # 合并所有批次的结果 return (torch.cat(result_images, dim=0), torch.cat(result_masks, dim=0), torch.cat(result_mask_images, dim=0)) @@ -324,4 +285,3 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "SegmentV2": "Segmentation V2 (RMBG)", } -