diff --git a/py/image_auto_crop_v2.py b/py/image_auto_crop_v2.py index c0605a6..d62d747 100644 --- a/py/image_auto_crop_v2.py +++ b/py/image_auto_crop_v2.py @@ -113,8 +113,17 @@ def image_auto_crop_v2(self, image, fill_background, background_color, aspect_ra previous_dino_model = grounding_dino_model item = _image.convert('RGBA') boxes = groundingdino_predict(DINO_MODEL, item, sam_prompt, sam_threshold) - (_, _mask) = sam_segment(SAM_MODEL, item, boxes) - _mask = mask2image(_mask[0]) + sam_result = sam_segment(SAM_MODEL, item, boxes) + if sam_result is not None: + (_, _mask) = sam_result + _mask = mask2image(_mask[0]) + else: + # don't do any cropping + _mask = Image.new('RGBA', _image.size, color="#FFFFFF") + ret_images.append(pil2tensor(_image)) + ret_box_previews.append(pil2tensor(_mask)) + ret_masks.append(image2mask(_mask)) + continue else: _mask = RMBG(_image) if ultra_detail_range: @@ -250,4 +259,4 @@ def image_auto_crop_v2(self, image, fill_background, background_color, aspect_ra NODE_DISPLAY_NAME_MAPPINGS = { "LayerUtility: ImageAutoCrop V2": "LayerUtility: ImageAutoCrop V2(Advance)" -} \ No newline at end of file +}