图像分类识别——（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

View Code

单张图片的简单识别实现代码

  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) # 展示预测结果表格

View Code

单张图片的简单识别实现代码的中文版文字显示

  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) # 展示预测结果表格

View Code

视频文件的处理

  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)