需求说明
手头上有一两万张图片,其中大多数是猫,但是还有不少乱七八糟的图片混入其中。目的是将非猫图片删掉,同时对是猫的图片在名称上带上猫的品种
图片概括如下:
paddlehub
调研市场开源产品后,还是选择百度的paddlehub来实现,效果也比较好。
环境搭建
本次使用python是3.8版本,使用虚拟环境安装包,python安装略。
1、虚拟环境创建
创建虚拟环境的目录假设在 D:\venv\ai
# 进入上述目录后执行
python -m venv .
cd D:\venv\ai\Scripts
activate.bat # 进入虚拟环境
2、安装包
pip install paddlehub -i https://pypi.douban.com/simple
# 注意,安装paddlehub不会自动安装paddlepaddle,运行时会报错,no module named paddle
pip install paddlepaddle -i https://pypi.douban.com/simple
代码编写
import os
import pathlib
import paddlehub as hub
import cv2
def all_path(dirname):
"""
根据指定目录,返回该目录下所有文件的绝对路径,返回值类型是 WindowsPath(不同系统会不同)
"""
r = []
for filepath in pathlib.Path(dirname).glob('**/*'):
r.append(filepath.absolute())
return r
n = 50 # 指定一次识别多少张图片,根据图片和内存大小设置合理数值
img_paths = all_path("D:\imgs\cat\pc") # 获取指定目录内所有图片路径
paths = []
# 切分图片到多个列表中
for i in range(0, len(img_paths) // n + 1):
paths.append(img_paths[i * n: (i + 1) * n])
# 加载训练好的模型,第一次执行时会自动下载模型
module = hub.Module(name="resnet50_vd_animals")
# 另一个模型,两个模型均可
# module = hub.Module(name="mobilenet_v2_animals")
for path in paths:
print(path)
np_images = [cv2.imread(image_path.__str__()) for image_path in path]
results = module.classification(images=np_images)
for result in zip(path, results):
print(result[0], result[1]) # D:\imgs\cat\pc\111.jpeg {'中国冠毛犬': 0.19150744378566742}
file_path = result[0].__str__()
cat_name = list(result[1].keys())[0]
if '猫' in cat_name: # 如果结果中含猫,保留并改名,否则删掉(图片中没有熊猫,如果有可作过滤)
filename = file_path.split('\\')[-1]
new_filename = cat_name + '__cat__' + filename
new_path = file_path.replace(filename, new_filename)
print(new_path)
# 重命名,猫的类型+原文件名为新文件名
os.rename(file_path, new_path)
else:
# 删掉非猫图片
print('删除', result[0], result[1])
os.remove(str(result[0]))
参考链接
https://www.paddlepaddle.org.cn/hub/scene/animals
https://blog.csdn.net/qq_53016081/article/details/119719496