同济子豪兄模板 半天搞定图像分类

‘’'import cv2
import numpy as np
import time
from tqdm import tqdm

视频逐帧处理代码模板

不需修改任何代码，只需定义process_frame函数即可

def generate_video(input_path=‘videos/robot.mp4’):
filehead = input_path.split(‘/’)[-1]
output_path = “out-” + filehead

print('视频开始处理',input_path)# 获取视频总帧数
cap = cv2.VideoCapture(input_path)
frame_count = 0
while(cap.isOpened()):success, frame = cap.read()frame_count += 1if not success:break
cap.release()
print('视频总帧数为',frame_count)# cv2.namedWindow('Crack Detection and Measurement Video Processing')
cap = cv2.VideoCapture(input_path)
frame_size = (cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT))# fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))
# fourcc = cv2.VideoWriter_fourcc(*'XVID')
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
fps = cap.get(cv2.CAP_PROP_FPS)out = cv2.VideoWriter(output_path, fourcc, fps, (int(frame_size[0]), int(frame_size[1])))# 进度条绑定视频总帧数
with tqdm(total=frame_count-1) as pbar:try:while(cap.isOpened()):success, frame = cap.read()if not success:break# 处理帧# frame_path = './temp_frame.png'# cv2.imwrite(frame_path, frame)try:frame = process_frame(frame)except:print('报错！', error)passif success == True:# cv2.imshow('Video Processing', frame)out.write(frame)# 进度条更新一帧pbar.update(1)# if cv2.waitKey(1) & 0xFF == ord('q'):# breakexcept:print('中途中断')passcv2.destroyAllWindows()
out.release()
cap.release()
print('视频已保存', output_path)

处理帧函数

def process_frame(img_bgr):

'''
输入摄像头拍摄画面bgr-array，输出图像分类预测结果bgr-array
'''# 记录该帧开始处理的时间
start_time = time.time()## 画面转成 RGB 的 Pillow 格式
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) # BGR转RGB
img_pil = Image.fromarray(img_rgb) # array 转 PIL## 预处理
input_img = test_transform(img_pil) # 预处理
input_tensor = input_img.unsqueeze(0).numpy()## onnx runtime 预测
ort_inputs = {'input': input_tensor} # onnx runtime 输入
pred_logits = ort_session.run(['output'], ort_inputs)[0] # onnx runtime 输出
pred_logits = torch.tensor(pred_logits)
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算## 解析top-n预测结果的类别和置信度
top_n = torch.topk(pred_softmax, 3) # 取置信度最大的 n 个结果
pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析预测类别
confs = top_n[0].cpu().detach().numpy().squeeze() # 解析置信度# 在图像上写英文
for i in range(len(confs)):pred_class = idx_to_labels[pred_ids[i]]# 写字：图片，添加的文字，左上角坐标，字体，字体大小，颜色，线宽，线型text = '{:<12} {:>.2f}'.format(pred_class, confs[i])img_bgr = cv2.putText(img_bgr, text, (50, 160 + 40 * i), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)# 记录该帧处理完毕的时间
end_time = time.time()
# 计算每秒处理图像帧数FPS
FPS = 1/(end_time - start_time)  
# 图片，添加的文字，左上角坐标，字体，字体大小，颜色，线宽，线型
img_bgr = cv2.putText(img_bgr, 'FPS  '+str(int(FPS)), (50, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2, cv2.LINE_AA)return img_bgr

while True:
img_bgr = cap.read()

if img_bgr is None:continueimg_bgr = process_frame(img_bgr)cvs.imshow(img_bgr)

‘’’

demo见B站分享： Aidlux两天搞定图像分类，半天就可以 |【aidlux实现图像分类～还是蛮好用的，可以推荐-哔哩哔哩】 https://b23.tv/codx3GL