图像分类识别——（5）完整的项目代码
环境配置代码
1 #安装库 2 !pip install numpy pandas matplotlib requests tqdm opencv-python pillow gc -i https://pypi.tuna.tsinghua.edu.cn/simple 3 # 下载安装Pytorch 4 !pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 5 # 安装mmcv -full库，对视频文件进行逐帧分割保存并再串成一个视频 6 !pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.10.0/index.html 7 # 下载中文字体文件 8 !wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/dataset/SimHei.ttf 9 # 下载 ImageNet 1000类别信息 10 #下载预训练模型，也可换成自己训练好的模型 11 !wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/dataset/meta_data/imagenet_class_index.csv 12 #创建目录 13 import os 14 # 存放测试图片 15 os.mkdir('test_img') 16 17 # 存放结果文件 18 os.mkdir('output') 19 20 # 下载测试图像文件 至 test_img 文件夹 21 22 !wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/watermelon1.jpg -O test_img/watermelon1.jpg 23 !wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/banana1.jpg -O test_img/banana1.jpg 24 !wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/cat1.jpg -O test_img/cat1.jpg 25 26 # 哈士奇，来源：https://www.pexels.com/zh-cn/photo/2853130/ 27 !wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/husky1.jpeg -O test_img/husky1.jpeg 28 29 # 猫狗，来源：https://unsplash.com/photos/ouo1hbizWwo 30 !wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/cat_dog.jpg -O test_img/cat_dog.jpg 31 32 !wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/video_2.mp4 -O test_img/video_2.mp4
单张图片的简单识别实现代码
1 # 导入操作系统库 2 import os 3 # 导入opencv库 4 import cv2 5 # 导入数据分析处理库 6 import pandas as pd 7 import numpy as np 8 # 导入pytorch 9 import torch 10 # 导入数据可视化库用来画图 11 import matplotlib.pyplot as plt 12 %matplotlib inline 13 # 获取设备信息，有 GPU 就用 GPU，没有就用 CPU 14 15 device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 16 #输出当前使用的计算设备 17 print('device', device) 18 from torchvision import models 19 # 载入预训练图像分类模型 20 # 这里使用的是最基础的resnet18模型 21 model = models.resnet18(pretrained=True) 22 # 也可换成下面更精准的模型，有152层，数字代表层数 23 # model = models.resnet152(pretrained=True) 24 # 把模型调整到评估模式 25 model = model.eval() 26 # 把模型放到刚才找到的计算设备中 27 model = model.to(device) 28 29 from torchvision import transforms 30 31 # 测试集图像预处理-RCTN：缩放裁剪、转 Tensor、归一化。为固定的代码模板 32 # Resize:缩放，把任何一张图像缩放成256*256方框 33 # CenterCrop 裁剪，把上步缩放的图片裁剪成224的小方框 34 # ToTensor：转换，把上一步的小方框图像转换成pytroch的张量Tensor 35 # Normalize：归一化，对图像色彩的三通道进行归一化，减去均值，除以标准差。是公认的固定参数 36 test_transform = transforms.Compose([transforms.Resize(256), 37 transforms.CenterCrop(224), 38 transforms.ToTensor(), 39 transforms.Normalize( 40 mean=[0.485, 0.456, 0.406], 41 std=[0.229, 0.224, 0.225]) 42 ]) 43 44 45 # 从之前准备的图片文件夹 test_img 里选择一张图片加载 46 img_path = 'test_img/banana1.jpg' 47 48 # img_path = 'test_img/cat_dog.jpg' 49 50 # 用 pillow 图像处理工具载入，pytroch的预处理函数只接受pillow格式的图像 51 from PIL import Image 52 # 用Image.open载入图像 53 img_pil = Image.open(img_path) 54 55 # 查询图片效果 56 img_pil 57 58 np.array(img_pil).shape 59 60 61 62 # 把img_pil传入给transform图像预处理函数 63 input_img = test_transform(img_pil) # 预处理 64 # 查询处理后的图像尺寸 65 input_img.shape 66 # 用 unsqueeze(0) 函数，给图像增加一个维度，因为pytroch模型输入模型必须整理成这样的维度 67 input_img = input_img.unsqueeze(0).to(device) 68 # 再次查询看下效果 69 input_img.shape 70 # 执行前向预测，得到所有类别的 logit 预测分数 71 pred_logits = model(input_img) 72 pred_logits.shape 73 # 显示这一千个类别的分数，这里是科学计数法 74 pred_logits 75 76 import torch.nn.functional as F 77 # 对 logit 分数做 softmax 归一化运算，把所有类别的分类转成0到1之间的概率，且所有类别求和为1 78 pred_softmax = F.softmax(pred_logits, dim=1) 79 pred_softmax.shape 80 81 # 再次查看，即显示概率 82 pred_softmax 83 84 85 86 # 指定图像大小 87 plt.figure(figsize=(8,4)) 88 # 横轴是1000个类别 89 x = range(1000) 90 # 纵轴是每一个类别的置信度 91 y = pred_softmax.cpu().detach().numpy()[0] 92 # 画柱状图 93 ax = plt.bar(x, y, alpha=0.5, width=0.3, color='yellow', edgecolor='red', lw=3) 94 # y轴取值范围 95 plt.ylim([0, 1.0]) 96 # plt.bar_label(ax, fmt='%.2f', fontsize=15) # 置信度数值 97 98 # 横轴与纵轴的图像的说明，字体的大小 99 plt.xlabel('Class', fontsize=20) 100 plt.ylabel('Confidence', fontsize=20) 101 # 坐标文字大小 102 plt.tick_params(labelsize=16) 103 # 图的标题 104 plt.title(img_path, fontsize=25) 105 # 展示图 106 plt.show() 107 108 109 # 设置取置信度最大的 n 个结果 110 n = 10 111 top_n = torch.topk(pred_softmax, n) 112 # 显示的是前 n 个类别的置信度和类别id 113 top_n 114 # 解析出类别 115 pred_ids = top_n[1].cpu().detach().numpy().squeeze() 116 # 显示解析出的类别id 117 pred_ids 118 # 解析出置信度，对应每个类别 119 confs = top_n[0].cpu().detach().numpy().squeeze() 120 # 显示置信度 121 confs 122 # 把1000类别字典的文件读入 123 df = pd.read_csv('imagenet_class_index.csv') 124 # 显示字典格式 125 df 126 127 128 # 把表格信息转成python的字典，键值对形式。 129 # 如果指定的是文件的class就是英文，chinese列就是中文 130 idx_to_labels = {} 131 for idx, row in df.iterrows(): 132 idx_to_labels[row['ID']] = [row['wordnet'], row['class']] 133 134 # 显示字典样式 135 idx_to_labels 136 137 138 # 用 opencv 载入原图 139 img_bgr = cv2.imread(img_path) 140 141 # 遍及前 n 个预测结果，获取每个预测结果的名称与置信度，并生成字符串 142 for i in range(n): 143 class_name = idx_to_labels[pred_ids[i]][1] # 获取类别名称 144 confidence = confs[i] * 100 # 获取置信度 145 text = '{:<15} {:>.4f}'.format(class_name, confidence) 146 print(text) 147 148 # 将文字写到图片上 149 # !图片，添加的文字，左上角坐标，字体，字号，bgr颜色，线宽 150 img_bgr = cv2.putText(img_bgr, text, (25, 50 + 40 * i), cv2.FONT_HERSHEY_SIMPLEX, 1.25, (0, 0, 255), 3) 151 152 153 # 保存图像 154 cv2.imwrite('output/img_pred.jpg', img_bgr) 155 # 载入预测结果图像 156 img_pred = Image.open('output/img_pred.jpg') 157 img_pred 158 159 160 161 # 图像和柱状图一起显示 162 fig = plt.figure(figsize=(18,6)) 163 164 # 绘制左图-预测图 165 ax1 = plt.subplot(1,2,1) 166 ax1.imshow(img_pred) 167 ax1.axis('off') 168 169 # 绘制右图-柱状图 170 ax2 = plt.subplot(1,2,2) 171 x = df['ID'] 172 y = pred_softmax.cpu().detach().numpy()[0] 173 ax2.bar(x, y, alpha=0.5, width=0.3, color='yellow', edgecolor='red', lw=3) 174 175 plt.ylim([0, 1.0]) # y轴取值范围 176 plt.title('{} Classification'.format(img_path), fontsize=30) 177 plt.xlabel('Class', fontsize=20) 178 plt.ylabel('Confidence', fontsize=20) 179 ax2.tick_params(labelsize=16) # 坐标文字大小 180 181 plt.tight_layout() 182 fig.savefig('output/预测图+柱状图.jpg') 183 184 # 预测结果表格输出 185 pred_df = pd.DataFrame() # 预测结果表格 186 for i in range(n): 187 class_name = idx_to_labels[pred_ids[i]][1] # 获取类别名称 188 label_idx = int(pred_ids[i]) # 获取类别号 189 wordnet = idx_to_labels[pred_ids[i]][0] # 获取 WordNet 190 confidence = confs[i] * 100 # 获取置信度 191 pred_df = pred_df.append({'Class':class_name, 'Class_ID':label_idx, 'Confidence(%)':confidence, 'WordNet':wordnet}, ignore_index=True) # 预测结果表格添加一行 192 display(pred_df) # 展示预测结果表格
单张图片的简单识别实现代码的中文版文字显示
1 1 # 设置matplotlib中文字体 2 2 # # windows操作系统 3 3 # plt.rcParams['font.sans-serif']=['SimHei'] # 用来正常显示中文标签 4 4 # plt.rcParams['axes.unicode_minus']=False # 用来正常显示负号 5 5 6 6 # Mac操作系统，参考 https://www.ngui.cc/51cto/show-727683.html 7 7 # 下载 simhei.ttf 字体文件 8 8 # !wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/dataset/SimHei.ttf 9 9 10 10 # Linux操作系统，例如 云GPU平台：https://featurize.cn/?s=d7ce99f842414bfcaea5662a97581bd1 11 11 # 如果遇到 SSL 相关报错，重新运行本代码块即可 12 12 !wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/dataset/SimHei.ttf -O /environment/miniconda3/lib/python3.7/site-packages/matplotlib/mpl-data/fonts/ttf/SimHei.ttf 13 13 !rm -rf /home/featurize/.cache/matplotlib 14 14 15 15 import matplotlib 16 16 matplotlib.rc("font",family='SimHei') # 中文字体 17 17 18 18 19 19 # 对于pillow在图像上写中文的字体，也要进行导入 20 20 from PIL import ImageFont, ImageDraw 21 21 # 导入中文字体，指定字号 22 22 font = ImageFont.truetype('SimHei.ttf', 32) 23 23 24 24 25 25 # 导入工具包 26 26 import os 27 27 28 28 import cv2 29 29 from PIL import Image, ImageFont, ImageDraw 30 30 31 31 import numpy as np 32 32 import pandas as pd 33 33 import matplotlib.pyplot as plt 34 34 %matplotlib inline 35 35 36 36 import torch 37 37 from torchvision import models 38 38 import torch.nn.functional as F 39 39 from torchvision import transforms 40 40 41 41 # 有 GPU 就用 GPU，没有就用 CPU 42 42 device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 43 43 print('device', device) 44 44 45 45 46 46 # 载入预训练图像分类模型 47 47 48 48 model = models.resnet18(pretrained=True) 49 49 50 50 # model = models.resnet152(pretrained=True) 51 51 52 52 model = model.eval() 53 53 model = model.to(device) 54 54 55 55 56 56 57 57 # 载入ImageNet 1000图像分类中文标签 58 58 # 这里的标签就换成了中文chinese 59 59 df = pd.read_csv('imagenet_class_index.csv') 60 60 idx_to_labels = {} 61 61 for idx, row in df.iterrows(): 62 62 idx_to_labels[row['ID']] = [row['wordnet'], row['Chinese']] 63 63 64 64 idx_to_labels 65 65 66 66 67 67 68 68 # 图像预处理 69 69 # 测试集图像预处理-RCTN：缩放裁剪、转 Tensor、归一化 70 70 test_transform = transforms.Compose([transforms.Resize(256), 71 71 transforms.CenterCrop(224), 72 72 transforms.ToTensor(), 73 73 transforms.Normalize( 74 74 mean=[0.485, 0.456, 0.406], 75 75 std=[0.229, 0.224, 0.225]) 76 76 ]) 77 77 78 78 79 79 80 80 # 载入一张测试图像 81 81 # img_path = 'test_img/banana1.jpg' 82 82 # img_path = 'test_img/husky1.jpeg' 83 83 # img_path = 'test_img/watermelon1.jpg' 84 84 img_path = 'test_img/cat_dog.jpg' 85 85 86 86 img_pil = Image.open(img_path) # 用 pillow 载入 87 87 88 88 img_pil 89 89 90 90 # 执行图像分类预测 91 91 input_img = test_transform(img_pil).unsqueeze(0).to(device) # 预处理 92 92 pred_logits = model(input_img) # 执行前向预测，得到所有类别的 logit 预测分数 93 93 pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算 94 94 95 95 96 96 # 各类别置信度柱状图 97 97 plt.figure(figsize=(8,4)) 98 98 99 99 x = range(1000) 100 100 y = pred_softmax.cpu().detach().numpy()[0] 101 101 102 102 ax = plt.bar(x, y, alpha=0.5, width=0.3, color='yellow', edgecolor='red', lw=3) 103 103 plt.ylim([0, 1.0]) # y轴取值范围 104 104 # plt.bar_label(ax, fmt='%.2f', fontsize=15) # 置信度数值 105 105 106 106 plt.title(img_path, fontsize=30) 107 107 plt.xlabel('类别', fontsize=20) 108 108 plt.ylabel('置信度', fontsize=20) 109 109 plt.tick_params(labelsize=16) # 坐标文字大小 110 110 111 111 plt.show() 112 112 113 113 114 114 115 115 116 116 # 取置信度最大的 n 个结果 117 117 n = 10 118 118 top_n = torch.topk(pred_softmax, n) # 取置信度最大的 n 个结果 119 119 pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析出类别 120 120 confs = top_n[0].cpu().detach().numpy().squeeze() # 解析出置信度 121 121 122 122 123 123 124 124 125 125 # 图像分类结果写在原图上 126 126 # opencv不能写中文，这里用pil来写中文 127 127 draw = ImageDraw.Draw(img_pil) 128 128 for i in range(n): 129 129 class_name = idx_to_labels[pred_ids[i]][1] # 获取类别名称 130 130 confidence = confs[i] * 100 # 获取置信度 131 131 text = '{:<15} {:>.4f}'.format(class_name, confidence) 132 132 print(text) 133 133 134 134 # 文字坐标，中文字符串，字体，rgba颜色 135 135 draw.text((50, 100 + 50 * i), text, font=font, fill=(255, 0, 0, 1)) 136 136 137 137 138 138 139 139 img_pil 140 140 141 141 142 142 143 143 144 144 # 保存图像 145 145 img_pil.save('output/img_pred.jpg') 146 146 147 147 148 148 149 149 150 150 151 151 # 图像和柱状图一起显示 152 152 fig = plt.figure(figsize=(18,6)) 153 153 154 154 # 绘制左图-预测图 155 155 ax1 = plt.subplot(1,2,1) 156 156 ax1.imshow(img_pil) 157 157 ax1.axis('off') 158 158 159 159 # 绘制右图-柱状图 160 160 ax2 = plt.subplot(1,2,2) 161 161 x = df['ID'] 162 162 y = pred_softmax.cpu().detach().numpy()[0] 163 163 ax2.bar(x, y, alpha=0.5, width=0.3, color='yellow', edgecolor='red', lw=3) 164 164 plt.ylim([0, 1.0]) # y轴取值范围 165 165 plt.xlabel('类别', fontsize=20) 166 166 plt.ylabel('置信度', fontsize=20) 167 167 ax2.tick_params(labelsize=16) # 坐标文字大小 168 168 169 169 plt.title('{} 图像分类预测结果'.format(img_path), fontsize=30) 170 170 171 171 plt.tight_layout() 172 172 fig.savefig('output/预测图+柱状图.jpg') 173 173 174 174 175 175 176 176 177 177 # 预测结果表格输出 178 178 pred_df = pd.DataFrame() # 预测结果表格 179 179 for i in range(n): 180 180 class_name = idx_to_labels[pred_ids[i]][1] # 获取类别名称 181 181 label_idx = int(pred_ids[i]) # 获取类别号 182 182 wordnet = idx_to_labels[pred_ids[i]][0] # 获取 WordNet 183 183 confidence = confs[i] * 100 # 获取置信度 184 184 pred_df = pred_df.append({'Class':class_name, 'Class_ID':label_idx, 'Confidence(%)':confidence, 'WordNet':wordnet}, ignore_index=True) # 预测结果表格添加一行 185 185 display(pred_df) # 展示预测结果表格
视频文件的处理
1 # 导入工具包，操作系统相关的包 2 import os 3 import time 4 import shutil 5 import tempfile 6 from tqdm import tqdm 7 8 # 导入图像处理相关的包 9 import cv2 10 from PIL import Image 11 12 # 导入与可视化相关的包 13 import numpy as np 14 import pandas as pd 15 import matplotlib.pyplot as plt 16 %matplotlib inline 17 plt.rcParams['axes.unicode_minus']=False # 用来正常显示负号 18 plt.rcParams['font.sans-serif']=['SimHei'] # 用来正常显示中文标签 19 import gc 20 21 # 导入pytorch相关的包 22 import torch 23 import torch.nn.functional as F 24 from torchvision import models 25 26 27 28 import mmcv 29 30 # 获取设备 31 # 有 GPU 就用 GPU，没有就用 CPU 32 device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 33 print('device:', device) 34 35 36 37 #这次要生成的是视频文件，所以不需要显示每一帧的绘图 38 # 后端绘图，不显示，只保存 39 import matplotlib 40 matplotlib.use('Agg') 41 42 43 # 载入预训练图像分类模型 44 model = models.resnet18(pretrained=True) 45 model = model.eval() 46 model = model.to(device) 47 48 49 # 载入ImageNet 1000图像分类标签 50 df = pd.read_csv('imagenet_class_index.csv') 51 idx_to_labels = {} 52 for idx, row in df.iterrows(): 53 idx_to_labels[row['ID']] = [row['wordnet'], row['class']] 54 55 idx_to_labels 56 57 58 59 # 图像预处理 60 from torchvision import transforms 61 62 # 测试集图像预处理-RCTN：缩放裁剪、转 Tensor、归一化 63 test_transform = transforms.Compose([transforms.Resize(256), 64 transforms.CenterCrop(224), 65 transforms.ToTensor(), 66 transforms.Normalize( 67 mean=[0.485, 0.456, 0.406], 68 std=[0.229, 0.224, 0.225]) 69 ]) 70 71 72 73 # 图像分类预测函数 74 def pred_single_frame(img, n=5): 75 ''' 76 输入摄像头画面bgr-array，输出前n个图像分类预测结果的图像bgr-array 77 ''' 78 img_bgr = img 79 img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # BGR 转 RGB 80 img_pil = Image.fromarray(img_rgb) # array 转 pil 81 input_img = test_transform(img_pil).unsqueeze(0).to(device) # 预处理 82 pred_logits = model(input_img) # 执行前向预测，得到所有类别的 logit 预测分数 83 pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算 84 85 top_n = torch.topk(pred_softmax, n) # 取置信度最大的 n 个结果 86 pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析出类别 87 confs = top_n[0].cpu().detach().numpy().squeeze() # 解析出置信度 88 89 # 在图像上写字 90 for i in range(n): 91 class_name = idx_to_labels[pred_ids[i]][1] # 获取类别名称 92 confidence = confs[i] * 100 # 获取置信度 93 text = '{:<15} {:>.4f}'.format(class_name, confidence) 94 95 # !图片，添加的文字，左上角坐标，字体，字号，bgr颜色，线宽 96 img_bgr = cv2.putText(img_bgr, text, (25, 50 + 40 * i), cv2.FONT_HERSHEY_SIMPLEX, 1.25, (0, 0, 255), 3) 97 98 return img_bgr, pred_softmax 99 100 101 102 # 输入输出视频路径 103 input_video = 'test_img/video_3.mp4' 104 105 106 # 创建临时文件夹，存放每帧结果,文件名是现在的时间。用来存储每一帧的预测结果 107 temp_out_dir = time.strftime('%Y%m%d%H%M%S') 108 os.mkdir(temp_out_dir) 109 print('创建文件夹 {} 用于存放每帧预测结果'.format(temp_out_dir)) 110 111 112 # 读入待预测视频 113 imgs = mmcv.VideoReader(input_video) 114 115 prog_bar = mmcv.ProgressBar(len(imgs)) 116 117 # 对视频逐帧处理 118 for frame_id, img in enumerate(imgs): 119 120 ## 处理单帧画面 121 img, pred_softmax = pred_single_frame(img, n=5) 122 123 # 将处理后的该帧画面图像文件，保存至 /tmp 目录下 124 cv2.imwrite(f'{temp_out_dir}/{frame_id:06d}.jpg', img) 125 126 prog_bar.update() # 更新进度条 127 128 # 把每一帧串成视频文件 129 mmcv.frames2video(temp_out_dir, 'output/output_pred.mp4', fps=imgs.fps, fourcc='mp4v') 130 131 shutil.rmtree(temp_out_dir) # 删除存放每帧画面的临时文件夹 132 print('删除临时文件夹', temp_out_dir)
