U2-Net 预测函数
包含单个图片检测以及视频检测
import os
import time
import cv2
import numpy as np
import matplotlib.pyplot as plt
import torch
import time
import subprocess
from torchvision.transforms import transforms
from src import u2net_full
os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'
def time_synchronized():
torch.cuda.synchronize() if torch.cuda.is_available() else None
return time.time()
# 获取GPU相关信息
def get_gpu_info():
try:
cmd_out = subprocess.check_output('nvidia-smi --query-gpu=name,memory.used,memory.total --format=csv,noheader',
shell=True)
gpu_info = cmd_out.decode().strip().split('\n')
gpu_info = [info.split(', ') for info in gpu_info]
return gpu_info
except subprocess.CalledProcessError as e:
print("Error while invoking nvidia-smi: ", e)
return None
# 打印 GPU 型号及占用情况
def print_gpu_usage():
gpu_info = get_gpu_info()
if gpu_info:
total_memory = 0
used_memory = 0
for name, used, total in gpu_info:
used_memory += int(used.strip().split()[0])
total_memory += int(total.strip().split()[0])
memory_usage_percent = round(used_memory / total_memory * 100, 2)
print(f"GPU: {name.strip()}, Memory used: {used.strip()}, Memory total: {total.strip()}"
f", Memory usage: {memory_usage_percent}%")
# 将原图像与分割后的图像混合
def Image_Blend(src, res):
info = res.shape
height = info[0]
width = info[1]
dst = np.zeros((height, width, 3), np.uint8)
# 分割后的图换色
mask = ~(res == [0, 0, 0]).all(axis=2)
res[mask] = [0, 0, 255]
dst = res
# 2.图像混合
img = cv2.addWeighted(src, 0.8, dst, 0.2, 0, dtype=cv2.CV_8UC3)
return img
# 读取 Gpu 信息
def gpu_info() -> str:
info = ''
for id in range(torch.cuda.device_count()):
p = torch.cuda.get_device_properties(id)
info += f'CUDA:{id} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n'
return info[:-1]
# 单单检测图片
def pic_predict(threshold, device, data_transform, origin_img, model):
h, w = origin_img.shape[:2]
img = data_transform(origin_img)
img = torch.unsqueeze(img, 0).to(device) # [C, H, W] -> [1, C, H, W]
with torch.no_grad():
# init model
img_height, img_width = img.shape[-2:]
init_img = torch.zeros((1, 3, img_height, img_width), device=device)
model(init_img)
# 推理
pred = model(img)
pred = torch.squeeze(pred).to("cpu").numpy() # [1, 1, H, W] -> [H, W]
pred = cv2.resize(pred, dsize=(w, h), interpolation=cv2.INTER_LINEAR)
pred_mask = np.where(pred > threshold, 1, 0)
origin_img = np.array(origin_img, dtype=np.uint8)
seg_img = origin_img * pred_mask[..., None]
img_res = Image_Blend(origin_img,seg_img)
cv2.imwrite("result/pred_result11.png", cv2.cvtColor(img_res.astype(np.uint8), cv2.COLOR_RGB2BGR))
# 视频检测
def video_pre(threshold, device, data_transform, origin_img, model):
h, w = origin_img.shape[:2]
img = data_transform(origin_img)
img = torch.unsqueeze(img, 0).to(device) # [C, H, W] -> [1, C, H, W]
with torch.no_grad():
# init model
img_height, img_width = img.shape[-2:]
init_img = torch.zeros((1, 3, img_height, img_width), device=device)
model(init_img)
# 推理
pred = model(img)
# 打印GPU占用信息
print_gpu_usage()
pred = torch.squeeze(pred).to("cpu").numpy() # [1, 1, H, W] -> [H, W]
pred = cv2.resize(pred, dsize=(w, h), interpolation=cv2.INTER_LINEAR)
pred_mask = np.where(pred > threshold, 1, 0)
origin_img = np.array(origin_img, dtype=np.uint8)
seg_img = origin_img * pred_mask[..., None]
img_res = Image_Blend(origin_img, seg_img)
return img_res
def main():
weights_path = "model_best.pth"
img_path = "test/video.mp4"
threshold = 0.5
# 判断图片路径是否正确
assert os.path.exists(img_path), f"image file {img_path} dose not exists."
# 判断 Gpu 是否可用
if torch.cuda.is_available():
print(gpu_info())
# 设置硬件 根据Gpu 是否可用，来选择用GPU 还是 CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Resize(320),
transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
])
# 载入模型
model = u2net_full()
weights = torch.load(weights_path, map_location='cpu')
if "model" in weights:
model.load_state_dict(weights["model"])
else:
model.load_state_dict(weights)
model.to(device)
model.eval()
str = os.path.splitext(img_path)[-1]
if str == ".mp4":
print("视频读入")
# 获取视频部分
cap = cv2.VideoCapture(img_path)
i = 0
# 2、获取图像的属性（宽和高），并将其转化为整数
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))
# 3、创建保存视频的对象，设置编码格式、帧率、图像的宽高等
out = cv2.VideoWriter('result/OutPut2.avi', cv2.VideoWriter_fourcc(*'FFV1'), 30,
(frame_width, frame_height))
while (cap.isOpened()):
# 4、获取每一帧图像
ret, frame = cap.read()
img = frame
i += 1
start_time = time.time() # 开始处理一帧图片的时间
img_res = video_pre(threshold, device, data_transform, img, model)
# 5、将每一帧图像写入到输出文件中
if ret == True:
out.write(img_res)
else:
break
end_time = time.time()
cost_time = end_time - start_time
print("检测第 {} 帧花了 {:.8f}s 。".format(i, cost_time))
cap.release()
out.release()
cv2.destroyAllWindows()
elif str == ".jpg":
print("图片读入")
start_time = time.time() # 开始处理一帧图片的时间
origin_img = cv2.cvtColor(cv2.imread(img_path, flags=cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)
pic_predict(threshold, device, data_transform, origin_img, model)
end_time = time.time()
cost_time = end_time - start_time
print("检测一张图片花了 {:.8f}s 。".format(cost_time))
else:
print("请重新读入图片或者视频")
if __name__ == '__main__':
main()
