xuebinqin / U-2-Net

The code for our newly accepted paper in Pattern Recognition 2020: "U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection."
Apache License 2.0
8.34k stars 1.43k forks source link

how get image matting #329

Open dh12306 opened 1 year ago

dh12306 commented 1 year ago

用测试脚本 得到的png如下图,使用pil读取的时候是 rgb模型,没有alpha通道 image

怎么通过原图获取想要的抠图区域呢,实现 origin.paste(mask, (0,0), mask= a) @xuebinqin image

xuebinqin commented 1 year ago

You can comcatenate the rgb with the prediction to form a 4-channel map, then save it to png. Be aware of that both rgb and the predicted map should be in the same scale e.g., [0,255].

On Mon., Sep. 5, 2022, 2:35 a.m. dh12306, @.***> wrote:

when I use u2net_test.py ,I get white-black image,like this: [image: image] https://user-images.githubusercontent.com/20471681/188417685-247568a7-6686-421f-984f-cb23ba6129f2.png

but how can I get the real matting img? origin pic is: [image: image] https://user-images.githubusercontent.com/20471681/188418035-6970e813-9b6a-4558-9448-b0a15e45db24.png

the png result is RGB mode, no alpha ,how can I use pil to get the real pic where I want?

— Reply to this email directly, view it on GitHub https://github.com/xuebinqin/U-2-Net/issues/329, or unsubscribe https://github.com/notifications/unsubscribe-auth/ADSGORN3THYI3JOTZJSVWWTV4W5HDANCNFSM6AAAAAAQEZ55FE . You are receiving this because you are subscribed to this thread.Message ID: @.***>

shuaimoumou commented 1 month ago

def save_output(image_name, pred, d_dir): predict = pred predict = predict.squeeze() predict = predict.unsqueeze(0) # 添加批处理维度 predict_np = predict.cpu().data.numpy()

# 读取原始图像
image = Image.open(image_name)

# 将预测结果上采样到原始图像大小
resized_pred = TF.resize(predict, image.size[::-1], interpolation=Image.BILINEAR)  # 需要反转图像大小顺序
resized_pred = resized_pred.squeeze()  # 移除批处理维度

# 转换为 [0, 255] 范围的 numpy 数组
resized_pred = (resized_pred * 255).byte().cpu().numpy()

# 创建4通道图像
rgba = np.zeros((image.size[1], image.size[0], 4), dtype=np.uint8)

# 将RGB图像复制到RGBA图像的前三个通道
rgba[:, :, :3] = np.array(image)

# 将预测结果放置到第四个通道
rgba[:, :, 3] = resized_pred

# 保存为PNG
img_name = image_name.split(os.sep)[-1]
img_name = img_name.split('.')[0] + '.png'
Image.fromarray(rgba).save(os.path.join(d_dir, img_name))

print(f"Saved {img_name} successfully.")