nodes/LK-168-comfyui_imgutils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix seg_bbox; improve README; change some display of node

Browse Source
This commit is contained in:
LK-168
2025-06-30 15:21:48 +08:00
parent 4c3019145d
commit bba945b1aa
10 changed files with 1075 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+25 -7
View File
@@ -6,7 +6,14 @@
https://github.com/deepghs/imgutils
模型文件全部由该库提供和管理下载，会自动下载到
`HF_HOME` 环境变量指定的目录下
因此需要在 ComfyUI 的启动项中设置该环境变量
**因此需要在 ComfyUI 的启动项中设置该环境变量**
如果你是官方包，可以在 .\run_nvidia_gpu.bat 的开头中添加：
```bat
set "CURDIR=%cd%"
set HF_HOME="%CURDIR%\(你想要的字目录)"
```
秋叶包可以不用设，默认是下载到 `.cache\huggingface\hub\` 下面
Segment-Anything 模型需要手动下载到
@@ -20,7 +27,8 @@ Segment-Anything 模型需要手动下载到
节点介绍：
#### 检测节点
#### BBox节点
![alt text](md_img/bbox_nodes.png)
- Imgutils Generic Detector
- 支持多种 `imgutils` 提供的多种基于anime的检测模型
- `detection_type`：检测类型
@@ -32,14 +40,23 @@ Segment-Anything 模型需要手动下载到
- Mask to BBox 、 BBox to Mask
- 用于 `Mask` 和 `BBox` 之间的转换
- BBoxFilter
- 用于过滤 `BBox`，可以根据置信度、面积和 标签进行过滤
- `labels`的值依照的是 `Imgutils Generic Detector` 输出的图片中bbox上标注的标签，可以输入多个，用逗号分隔
- 用于过滤 `BBox`，可以根据置信度、面积和标签进行过滤
- `labels`的值依照的是 `Imgutils Generic Detector` 的输出 `image with boxes` 中bbox上标注的标签，可以输入多个，用逗号分隔
#### segment-anything节点
基本上抄 Impact-Pack 的
基本上抄的 Impact-Pack
![alt text](md_img/detailer_example.png)
- SAMPredictorNode
- SAMLoader for SAMPredictorNode
#### segment 节点
![alt text](md_img/segment_nodes.png)
- Imgutils Auto Segmenter
- 仅能对图片进行 **前景和背景** 的分割
@@ -56,6 +73,8 @@ Segment-Anything 模型需要手动下载到
#### Mask处理节点
收录一些常用的 Mask 处理节点
![alt text](md_img/mask_nodes.png)
- Mask Morphology:
- 提供了常用的形态学操作：膨胀、腐蚀、开运算和闭运算
-
@@ -69,8 +88,7 @@ Segment-Anything 模型需要手动下载到
- Mask Info:
- 显示mask的统计信息，如形状、覆盖率、范围和均值
- MaskHelperLK:
- 如果你忘记了以上`Mask`处理节点的功能，可以使用这个节点查看，因为我知道这个东西只看节点名字很难知道效果
- 如果你忘记了以上`Mask`处理节点的功能，可以使用这个节点查看，因为我知道以上四个节点只看节点名字很难知道效果
懒了，今天先写到这，传播民主与自由去了。
example_workflow/detailer_example.json
+997
View File
@@ -0,0 +1,997 @@
{
"id": "924f8d9b-3083-4698-8424-5d1ff6375ff9",
"revision": 0,
"last_node_id": 23,
"last_link_id": 23,
"nodes": [
{
"id": 18,
"type": "MaskPreview+",
"pos": [
1430,
580
],
"size": [
230,
246
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"localized_name": "mask",
"name": "mask",
"type": "MASK",
"link": 15
}
],
"outputs": [],
"title": "BBox result",
"properties": {
"cnr_id": "comfyui_essentials",
"ver": "1.1.0",
"widget_ue_connectable": {},
"Node name for S&R": "MaskPreview+"
},
"widgets_values": []
},
{
"id": 19,
"type": "MaskPreview+",
"pos": [
1670,
580
],
"size": [
230,
246
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"localized_name": "mask",
"name": "mask",
"type": "MASK",
"link": 16
}
],
"outputs": [],
"title": "Segmented result",
"properties": {
"cnr_id": "comfyui_essentials",
"ver": "1.1.0",
"widget_ue_connectable": {},
"Node name for S&R": "MaskPreview+"
},
"widgets_values": []
},
{
"id": 16,
"type": "SAMLoaderLK",
"pos": [
1500,
450
],
"size": [
261.6646423339844,
82
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [
{
"localized_name": "model_name",
"name": "model_name",
"type": "COMBO",
"widget": {
"name": "model_name"
},
"link": null
},
{
"localized_name": "device_mode",
"name": "device_mode",
"type": "COMBO",
"widget": {
"name": "device_mode"
},
"link": null
}
],
"outputs": [
{
"localized_name": "SAM_MODEL",
"name": "SAM_MODEL",
"type": "SAM_MODEL",
"links": [
13
]
}
],
"properties": {
"aux_id": "LK-168/comfyui_imgutils",
"ver": "4c3019145dd6a17c7c82f4b1d8867b3b8a0077b6",
"widget_ue_connectable": {},
"Node name for S&R": "SAMLoaderLK"
},
"widgets_values": [
"sam_vit_b.pth",
"AUTO"
]
},
{
"id": 12,
"type": "ImgutilsGenericDetector",
"pos": [
1500,
200
],
"size": [
270,
218
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"type": "IMAGE",
"link": 10
},
{
"localized_name": "detection_type",
"name": "detection_type",
"type": "COMBO",
"widget": {
"name": "detection_type"
},
"link": null
},
{
"localized_name": "conf_threshold",
"name": "conf_threshold",
"type": "FLOAT",
"widget": {
"name": "conf_threshold"
},
"link": null
},
{
"localized_name": "iou_threshold",
"name": "iou_threshold",
"type": "FLOAT",
"widget": {
"name": "iou_threshold"
},
"link": null
},
{
"localized_name": "draw_boxes",
"name": "draw_boxes",
"type": "BOOLEAN",
"widget": {
"name": "draw_boxes"
},
"link": null
},
{
"localized_name": "level",
"name": "level",
"type": "STRING",
"widget": {
"name": "level"
},
"link": null
},
{
"localized_name": "version",
"name": "version",
"type": "STRING",
"widget": {
"name": "version"
},
"link": null
}
],
"outputs": [
{
"localized_name": "image_with_boxes",
"name": "image_with_boxes",
"type": "IMAGE",
"links": [
11
]
},
{
"localized_name": "detection_mask",
"name": "detection_mask",
"type": "MASK",
"links": [
15
]
},
{
"localized_name": "detection_bbox",
"name": "detection_bbox",
"type": "BBOX",
"links": [
12
]
}
],
"properties": {
"aux_id": "LK-168/comfyui_imgutils",
"ver": "4c3019145dd6a17c7c82f4b1d8867b3b8a0077b6",
"widget_ue_connectable": {},
"Node name for S&R": "ImgutilsGenericDetector"
},
"widgets_values": [
"Hand Detection",
0.5,
0.7,
true,
"s",
""
]
},
{
"id": 21,
"type": "MaskToSEGS",
"pos": [
2150,
200
],
"size": [
270,
154
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"localized_name": "mask",
"name": "mask",
"type": "MASK",
"link": 18
},
{
"localized_name": "combined",
"name": "combined",
"type": "BOOLEAN",
"widget": {
"name": "combined"
},
"link": null
},
{
"localized_name": "crop_factor",
"name": "crop_factor",
"type": "FLOAT",
"widget": {
"name": "crop_factor"
},
"link": null
},
{
"localized_name": "bbox_fill",
"name": "bbox_fill",
"type": "BOOLEAN",
"widget": {
"name": "bbox_fill"
},
"link": null
},
{
"localized_name": "drop_size",
"name": "drop_size",
"type": "INT",
"widget": {
"name": "drop_size"
},
"link": null
},
{
"localized_name": "contour_fill",
"name": "contour_fill",
"type": "BOOLEAN",
"widget": {
"name": "contour_fill"
},
"link": null
}
],
"outputs": [
{
"localized_name": "SEGS",
"name": "SEGS",
"type": "SEGS",
"links": [
19
]
}
],
"properties": {
"cnr_id": "comfyui-impact-pack",
"ver": "8.17.1",
"widget_ue_connectable": {},
"Node name for S&R": "MaskToSEGS"
},
"widgets_values": [
false,
3,
false,
10,
false
]
},
{
"id": 13,
"type": "LoadImage",
"pos": [
1210,
200
],
"size": [
274.080078125,
314
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"type": "COMBO",
"widget": {
"name": "image"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "IMAGEUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "IMAGE",
"name": "IMAGE",
"type": "IMAGE",
"links": [
10,
20,
21
]
},
{
"localized_name": "MASK",
"name": "MASK",
"type": "MASK",
"links": []
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.41",
"widget_ue_connectable": {},
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"1067_969609715223178_ST_F_00001_.png",
"image"
]
},
{
"id": 14,
"type": "PreviewImage",
"pos": [
1920,
420
],
"size": [
300,
420
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"localized_name": "images",
"name": "images",
"type": "IMAGE",
"link": 11
}
],
"outputs": [],
"title": "Image with boxes",
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.41",
"widget_ue_connectable": {},
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 15,
"type": "SAMPredictorNode",
"pos": [
1850,
200
],
"size": [
270,
170
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"localized_name": "sam_model",
"name": "sam_model",
"type": "SAM_MODEL",
"link": 13
},
{
"localized_name": "image",
"name": "image",
"type": "IMAGE",
"link": 20
},
{
"localized_name": "bbox",
"name": "bbox",
"type": "BBOX",
"link": 12
},
{
"localized_name": "threshold",
"name": "threshold",
"type": "FLOAT",
"widget": {
"name": "threshold"
},
"link": null
},
{
"localized_name": "points_method",
"name": "points_method",
"type": "COMBO",
"widget": {
"name": "points_method"
},
"link": null
},
{
"localized_name": "merge_options",
"name": "merge_options",
"type": "COMBO",
"widget": {
"name": "merge_options"
},
"link": null
},
{
"localized_name": "crop_factor",
"name": "crop_factor",
"shape": 7,
"type": "FLOAT",
"widget": {
"name": "crop_factor"
},
"link": null
}
],
"outputs": [
{
"localized_name": "MASK",
"name": "MASK",
"type": "MASK",
"links": [
16,
18,
23
]
},
{
"localized_name": "SEGS",
"name": "SEGS",
"type": "SEGS",
"links": null
}
],
"properties": {
"aux_id": "LK-168/comfyui_imgutils",
"ver": "4c3019145dd6a17c7c82f4b1d8867b3b8a0077b6",
"widget_ue_connectable": {},
"Node name for S&R": "SAMPredictorNode"
},
"widgets_values": [
0.4,
"None",
"Merge All",
3
]
},
{
"id": 23,
"type": "ImageAndMaskPreview",
"pos": [
2240,
410
],
"size": [
370,
430
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"label": "image",
"localized_name": "image",
"name": "image",
"shape": 7,
"type": "IMAGE",
"link": 21
},
{
"label": "mask",
"localized_name": "mask",
"name": "mask",
"shape": 7,
"type": "MASK",
"link": 23
},
{
"localized_name": "mask_opacity",
"name": "mask_opacity",
"type": "FLOAT",
"widget": {
"name": "mask_opacity"
},
"link": null
},
{
"localized_name": "mask_color",
"name": "mask_color",
"type": "STRING",
"widget": {
"name": "mask_color"
},
"link": null
},
{
"localized_name": "pass_through",
"name": "pass_through",
"type": "BOOLEAN",
"widget": {
"name": "pass_through"
},
"link": null
}
],
"outputs": [
{
"label": "composite",
"localized_name": "composite",
"name": "composite",
"type": "IMAGE",
"links": null
}
],
"properties": {
"cnr_id": "comfyui-kjnodes",
"ver": "8950c5fe673f30b7bedee6650ed88e998b6caf27",
"widget_ue_connectable": {},
"Node name for S&R": "ImageAndMaskPreview"
},
"widgets_values": [
1,
"255, 255, 255",
false
]
},
{
"id": 22,
"type": "DetailerForEach",
"pos": [
2650,
190
],
"size": [
400,
680
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"type": "IMAGE",
"link": null
},
{
"localized_name": "segs",
"name": "segs",
"type": "SEGS",
"link": 19
},
{
"localized_name": "model",
"name": "model",
"type": "MODEL",
"link": null
},
{
"localized_name": "clip",
"name": "clip",
"type": "CLIP",
"link": null
},
{
"localized_name": "vae",
"name": "vae",
"type": "VAE",
"link": null
},
{
"localized_name": "positive",
"name": "positive",
"type": "CONDITIONING",
"link": null
},
{
"localized_name": "negative",
"name": "negative",
"type": "CONDITIONING",
"link": null
},
{
"localized_name": "detailer_hook",
"name": "detailer_hook",
"shape": 7,
"type": "DETAILER_HOOK",
"link": null
},
{
"localized_name": "scheduler_func_opt",
"name": "scheduler_func_opt",
"shape": 7,
"type": "SCHEDULER_FUNC",
"link": null
},
{
"localized_name": "guide_size",
"name": "guide_size",
"type": "FLOAT",
"widget": {
"name": "guide_size"
},
"link": null
},
{
"localized_name": "guide_size_for",
"name": "guide_size_for",
"type": "BOOLEAN",
"widget": {
"name": "guide_size_for"
},
"link": null
},
{
"localized_name": "max_size",
"name": "max_size",
"type": "FLOAT",
"widget": {
"name": "max_size"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
},
{
"localized_name": "steps",
"name": "steps",
"type": "INT",
"widget": {
"name": "steps"
},
"link": null
},
{
"localized_name": "cfg",
"name": "cfg",
"type": "FLOAT",
"widget": {
"name": "cfg"
},
"link": null
},
{
"localized_name": "sampler_name",
"name": "sampler_name",
"type": "COMBO",
"widget": {
"name": "sampler_name"
},
"link": null
},
{
"localized_name": "scheduler",
"name": "scheduler",
"type": "COMBO",
"widget": {
"name": "scheduler"
},
"link": null
},
{
"localized_name": "denoise",
"name": "denoise",
"type": "FLOAT",
"widget": {
"name": "denoise"
},
"link": null
},
{
"localized_name": "feather",
"name": "feather",
"type": "INT",
"widget": {
"name": "feather"
},
"link": null
},
{
"localized_name": "noise_mask",
"name": "noise_mask",
"type": "BOOLEAN",
"widget": {
"name": "noise_mask"
},
"link": null
},
{
"localized_name": "force_inpaint",
"name": "force_inpaint",
"type": "BOOLEAN",
"widget": {
"name": "force_inpaint"
},
"link": null
},
{
"localized_name": "wildcard",
"name": "wildcard",
"type": "STRING",
"widget": {
"name": "wildcard"
},
"link": null
},
{
"localized_name": "cycle",
"name": "cycle",
"type": "INT",
"widget": {
"name": "cycle"
},
"link": null
},
{
"localized_name": "inpaint_model",
"name": "inpaint_model",
"shape": 7,
"type": "BOOLEAN",
"widget": {
"name": "inpaint_model"
},
"link": null
},
{
"localized_name": "noise_mask_feather",
"name": "noise_mask_feather",
"shape": 7,
"type": "INT",
"widget": {
"name": "noise_mask_feather"
},
"link": null
},
{
"localized_name": "tiled_encode",
"name": "tiled_encode",
"shape": 7,
"type": "BOOLEAN",
"widget": {
"name": "tiled_encode"
},
"link": null
},
{
"localized_name": "tiled_decode",
"name": "tiled_decode",
"shape": 7,
"type": "BOOLEAN",
"widget": {
"name": "tiled_decode"
},
"link": null
}
],
"outputs": [
{
"localized_name": "IMAGE",
"name": "IMAGE",
"type": "IMAGE",
"links": null
}
],
"properties": {
"cnr_id": "comfyui-impact-pack",
"ver": "8.17.1",
"widget_ue_connectable": {},
"Node name for S&R": "DetailerForEach"
},
"widgets_values": [
512,
true,
1024,
121335371194066,
"randomize",
20,
8,
"euler",
"normal",
0.5,
5,
true,
true,
"",
1,
false,
20,
false,
false,
[
false,
true
]
]
}
],
"links": [
[
10,
13,
0,
12,
0,
"IMAGE"
],
[
11,
12,
0,
14,
0,
"IMAGE"
],
[
12,
12,
2,
15,
2,
"BBOX"
],
[
13,
16,
0,
15,
0,
"SAM_MODEL"
],
[
15,
12,
1,
18,
0,
"MASK"
],
[
16,
15,
0,
19,
0,
"MASK"
],
[
18,
15,
0,
21,
0,
"MASK"
],
[
19,
21,
0,
22,
1,
"SEGS"
],
[
20,
13,
0,
15,
1,
"IMAGE"
],
[
21,
13,
0,
23,
0,
"IMAGE"
],
[
23,
15,
0,
23,
1,
"MASK"
]
],
"groups": [],
"config": {},
"extra": {
"workspace_info": {
"id": "M1pq72Nrhan7TltzaAQUJ"
},
"ue_links": [],
"links_added_by_ue": [],
"ds": {
"scale": 0.7266303355934305,
"offset": [
-564.8769777196603,
274.40995013281713
]
}
},
"version": 0.4
}
md_img/bbox_nodes.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 111 KiB

md_img/detailer_example.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 978 KiB

md_img/mask_nodes.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 86 KiB

md_img/segment_nodes.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 865 KiB

utils/detect.py
+2 -2
View File
@@ -108,13 +108,13 @@ class ImgutilsGenericDetector:
"draw_boxes": ("BOOLEAN", {"default": True}),
# Model specific parameters, always present for simplicity in this generic node
"level": ("STRING", {"default": "s", "options": ["n", "s"]}),
"version": ("STRING", {"default": "v1.1", "options": ["v0", "v1", "v1.1"]}),
"version": ("STRING", {"default": "", "options": ["v0", "v1", "v1.1"]}),
# "model_name": ("STRING", {"default": ""}),
},
}
RETURN_TYPES = ("IMAGE", "MASK",BBOX, )
RETURN_NAMES = ("image_with_boxes", "detection_mask","detection_results",)
RETURN_NAMES = ("image_with_boxes", "detection_mask","detection_bbox",)
FUNCTION = "detect"
def detect(self, image, detection_type, conf_threshold,
utils/mask.py
+25 -3
View File
@@ -276,13 +276,17 @@ class MaskInfoNode:
}
CATEGORY = "imgutils/mask"
RETURN_TYPES = ("MASK", "STRING")
RETURN_NAMES = ("mask", "info")
RETURN_TYPES = ("MASK", "FLOAT", "STRING", "FLOAT", "STRING")
RETURN_NAMES = ("mask", "coverage", "value_range", "mean_value", "detailed_info")
FUNCTION = "get_mask_info"
def get_mask_info(self, mask):
batch_size = mask.shape[0]
info_list = []
total_coverage = 0.0
total_mean = 0.0
overall_min = 1.0
overall_max = 0.0
for i in range(batch_size):
current_mask = mask[i].squeeze().cpu().numpy()
@@ -297,12 +301,30 @@ class MaskInfoNode:
max_val = np.max(current_mask)
mean_val = np.mean(current_mask)
# Update overall statistics
total_coverage += coverage
total_mean += mean_val
overall_min = min(overall_min, min_val)
overall_max = max(overall_max, max_val)
info = f"Batch {i}: Shape={shape}, Coverage={coverage:.1f}%, "
info += f"Range=[{min_val:.3f}, {max_val:.3f}], Mean={mean_val:.3f}"
info_list.append(info)
# Calculate averages
avg_coverage = total_coverage / batch_size
avg_mean = total_mean / batch_size
value_range = overall_max - overall_min
# Create detailed info
combined_info = "\n".join(info_list)
return (mask, combined_info)
if batch_size > 1:
combined_info += f"\n\nOverall Statistics:"
combined_info += f"\nAverage Coverage: {avg_coverage:.1f}%"
combined_info += f"\nOverall Range: [{overall_min:.3f}, {overall_max:.3f}]"
combined_info += f"\nAverage Mean: {avg_mean:.3f}"
return (mask, avg_coverage, f"[{overall_min:.3f}, {overall_max:.3f}]", avg_mean, combined_info)
class MaskHelperLK:
"""Helper class for above mask operations"""
utils/sam.py
+1 -1
View File
@@ -342,7 +342,7 @@ NODE_CLASS_MAPPINGS = {
}
NODE_DISPLAY_NAME_MAPPINGS = {
"SAMLoaderLK": "SAM Loader LK",
"SAMLoaderLK": "SAM Loader for SAMPredictorNode",
"SAMPredictorNode": "SAM Predictor",
}
utils/segment.py
+25 -11
View File
@@ -104,7 +104,7 @@ class ImgutilsBBoxSegmenter(object):
original_image_pil = Image.fromarray(img_np_255).convert("RGB")
bbox_mask_np_0_1 = bbox_mask[0].cpu().numpy() # MASK is [B, H, W], values 0-1
bbox_mask_np = (bbox_mask_np_0_1 > 0.5).astype(np.float32)
bbox_mask_np = (bbox_mask_np_0_1 > 0.01).astype(np.float32)
# If bbox_mask is empty, return original image and empty mask
if not np.any(bbox_mask_np):
@@ -112,7 +112,17 @@ class ImgutilsBBoxSegmenter(object):
empty_mask_tensor = torch.zeros((1, original_image_pil.height, original_image_pil.width), dtype=torch.float32)
return (image, empty_mask_tensor)
# Get the main subject's mask from imgutils
# Find bounding box coordinates from the mask
bbox_indices = np.where(bbox_mask_np > 0)
if len(bbox_indices[0]) == 0:
print("BBox mask has no valid pixels. Returning original image and empty mask.")
empty_mask_tensor = torch.zeros((1, original_image_pil.height, original_image_pil.width), dtype=torch.float32)
return (image, empty_mask_tensor)
min_y, max_y = bbox_indices[0].min(), bbox_indices[0].max() + 1
min_x, max_x = bbox_indices[1].min(), bbox_indices[1].max() + 1
# Get the main subject's mask from imgutils (full image)
seg_mask_np_full = get_isnetis_mask(original_image_pil, scale=scale)
# Normalize mask to 0-1 float32
if seg_mask_np_full.dtype == np.uint8:
@@ -123,17 +133,16 @@ class ImgutilsBBoxSegmenter(object):
# Combine imgutils mask with provided bbox_mask
final_seg_mask_np = seg_mask_np_full * bbox_mask_np
# Get the segmented image based on mode
# Get the segmented image based on mode (full image)
if segment_mode == "rgba_transparent":
_, seg_image_pil_full = segment_rgba_with_isnetis(original_image_pil, scale=scale)
if seg_image_pil_full.mode != 'RGBA':
seg_image_pil_full = seg_image_pil_full.convert('RGBA')
# Apply the combined mask to the alpha channel of the segmented image
# Create an alpha channel from the final_seg_mask_np
alpha_from_mask = Image.fromarray((final_seg_mask_np * 255).astype(np.uint8), mode='L')
seg_image_pil_full.putalpha(alpha_from_mask)
output_image_pil = seg_image_pil_full
full_output_image_pil = seg_image_pil_full
else: # RGB modes
bg_color_val = 0 if segment_mode == "rgb_black_bg" else 255
@@ -149,20 +158,25 @@ class ImgutilsBBoxSegmenter(object):
seg_image_pil_full = seg_image_pil_full.convert('RGB')
# Paste the imgutils segmented image onto the background using the combined mask
# Use the final_seg_mask_np to control what gets pasted
paste_mask_pil = Image.fromarray((final_seg_mask_np * 255).astype(np.uint8), mode='L')
background_image.paste(seg_image_pil_full, (0, 0), paste_mask_pil)
output_image_pil = background_image
full_output_image_pil = background_image
# Crop the segmented image to the bbox region
cropped_output_image_pil = full_output_image_pil.crop((min_x, min_y, max_x, max_y))
# Crop the mask to the bbox region as well
cropped_seg_mask_np = final_seg_mask_np[min_y:max_y, min_x:max_x]
# Format output
if output_image_pil:
output_image_np_255 = np.array(output_image_pil.convert("RGB"))
if cropped_output_image_pil:
output_image_np_255 = np.array(cropped_output_image_pil.convert("RGB"))
output_image_tensor = torch.from_numpy(output_image_np_255.astype(np.float32) / 255.0).unsqueeze(0)
else: # Fallback
output_image_tensor = image
# Final mask is the combined mask
output_mask_tensor = torch.from_numpy(final_seg_mask_np).unsqueeze(0)
# Output the cropped mask
output_mask_tensor = torch.from_numpy(cropped_seg_mask_np).unsqueeze(0)
return (output_image_tensor, output_mask_tensor)