包含单个图片检测以及视频检测
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()