68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
|
# u2net_saliency_only.py
|
|||
|
|
import torch
|
|||
|
|
import cv2
|
|||
|
|
import torchvision.transforms as transforms
|
|||
|
|
import numpy as np
|
|||
|
|
from PIL import Image
|
|||
|
|
from model.u2net import U2NETP
|
|||
|
|
import os
|
|||
|
|
# 缓存模型，避免重复加载
|
|||
|
|
_model = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_u2net_model(model_path="./saved_models/u2netp.pth"):
|
|||
|
|
global _model
|
|||
|
|
if _model is None:
|
|||
|
|
_model = U2NETP(3, 1)
|
|||
|
|
_model.load_state_dict(torch.load(model_path))
|
|||
|
|
_model.cuda()
|
|||
|
|
_model.eval()
|
|||
|
|
return _model
|
|||
|
|
|
|||
|
|
|
|||
|
|
def preprocess(image):
|
|||
|
|
transform = transforms.Compose([
|
|||
|
|
transforms.Resize((320, 320)),
|
|||
|
|
transforms.ToTensor(),
|
|||
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406],
|
|||
|
|
std=[0.229, 0.224, 0.225])
|
|||
|
|
])
|
|||
|
|
return transform(image).unsqueeze(0)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def postprocess(output, original_size):
|
|||
|
|
output = output.squeeze().cpu().detach().numpy()
|
|||
|
|
output = (output * 255).astype(np.uint8)
|
|||
|
|
output = cv2.resize(output, original_size)
|
|||
|
|
return output
|
|||
|
|
|
|||
|
|
|
|||
|
|
def generate_saliency_map(input_image_path, output_image_path, model_path="./saved_models/u2netp.pth"):
|
|||
|
|
net = load_u2net_model(model_path)
|
|||
|
|
|
|||
|
|
original_image = Image.open(input_image_path).convert("RGB")
|
|||
|
|
original_size = original_image.size
|
|||
|
|
input_tensor = preprocess(original_image).cuda()
|
|||
|
|
|
|||
|
|
output = net(input_tensor)[0]
|
|||
|
|
saliency_map = postprocess(output, original_size)
|
|||
|
|
|
|||
|
|
cv2.imwrite(output_image_path, saliency_map)
|
|||
|
|
print(f"显著图已保存至: {output_image_path}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
triplets = [
|
|||
|
|
# (标签, 原图路径, 显著 / 掩模 路径)
|
|||
|
|
('front', './image/front.jpg', './saliency/front.png'), # 正面
|
|||
|
|
('rear', './image/rear.jpg', './saliency/rear.png'), # 后面
|
|||
|
|
('side', './image/side.jpg', './saliency/side.png'), # 侧面（做圆检测）
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
thresh_dir = './thresh'
|
|||
|
|
os.makedirs(thresh_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
# # ======================= 生成显著性图 （可以注释掉，在u2net_saliency生成）=======================
|
|||
|
|
for tag, image_path, saliency_path in triplets:
|
|||
|
|
print(f"处理 {tag} 图像中...")
|
|||
|
|
generate_saliency_map(image_path, saliency_path)
|