原文:
towardsdatascience.com/how-many-cars-are-in-this-aerial-imagery-lets-count-them-with-yolov8-from-scratch-7c24a3919d21
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/8b40564aff7e1a1eda4beb40ba637034.png
YOLO 在空中图像中检测到的车辆,图片由作者提供
目录
🌟 简介
🚀设置 Google Colab
🛰️在哪里找到免费空中图像
🔄将 GeoTIFF 转换为 JPG
📐将图像分割成子图像
🏷️在子图像中标记车辆
🛠️ 准备训练和测试数据库
🧠 训练 YOLOv8
🖼️ 在整个图像上部署 YOLOv8
📄 结论
📚 参考资料
🌟 简介
几乎一个月前,我在一家知名报纸上读到了一篇分析每家特斯拉库存中汽车数量的文章,该文章使用了超高分辨率的图像(分辨率低于 20 厘米)。文章旨在估计特斯拉在第二季度交付的汽车数量。阅读后,我认为这可以是一个将目标检测算法应用于检测和计数空中图像中汽车,评估汽车公司库存状况的绝佳应用。
受此想法的启发,我决定准备一个教程,演示您如何使用您的标记数据集训练一个目标检测算法,并将其从 A 到 Z 应用于空中图像。在这个教程中,您将学习如何找到免费空中图像,将 GeoTIFF 文件转换为 JPG 文件而不丢失分辨率,标记您的数据集,训练您的模型,并将其部署到整个图像上。如果您想在 Medium 的一篇文章中学习所有这些步骤,这个故事就是为您准备的。
🚀 在 Google Colab 中设置(可选)
在下载图像和运行脚本之前,您需要设置您将工作的环境。我测试了我的脚本在 Google Colab 上,因为任何拥有 Gmail 账户的人都可以访问并运行它。然而,如果您使用的是 VS Code 或 Anaconda 等不同的环境,您仍然可以按照以下步骤创建必要的输入和输出文件夹。
首先,请确保您已经在您的系统上安装了以下包:
pip install supervision ultralytics rasterio接下来,我们需要以下文件夹:“raw_inputs”,“images”,“labels”,“sub_images”,并且在“images”和“labels”文件夹中,我们需要“train”和“test”子文件夹。在其他部分,您将学习每个文件夹中会保存哪些文件。以下代码将在您的目录中创建这些文件夹:
importos folders=['raw_inputs','images/train','images/test','labels/train','labels/test','sub_images']forfolderinfolders:os.makedirs(folder,exist_ok=True)🛰️在哪里找到免费空中图像
我们需要遵循本教程步骤的唯一输入是高分辨率航空影像,它允许我们清晰地看到对象。你可以在网上找到许多来源,但我在这个教程中使用的是 OpenAerialMap,它为不同地点提供无人机影像。访问网站,点击 “Start Exploring”,并下载航空影像。
我使用的图像名为 “Petre Bagrationi Street – Batumi”,在 CC-BY 4.0 许可下可用。这个许可允许我们在任何媒介或格式下为任何目的复制和重新分发材料。请随意探索其他地点并下载包含你想要检测的对象的航空影像。这张图像是在大约一个月前在一个城市区域上空拍摄的。你可以在开源软件如 QGIS 中打开下载的 geotiff 文件:
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/4e7da5bd8bef3195cb6acb4ac75d53a4.png
在 CC-BY 4.0 许可下从openaerialmap.org/下载的航空影像,由作者在 QGIS 中可视化
下载图像后,将其上传到“raw_inputs”文件夹。
🔄将 GeoTIFF 转换为 JPG
当你从 OpenAerialMap 网站下载航空影像时,图像的格式是 GeoTIFF。然而,为了使用 YOLO 模型等对象检测算法——我们在这个教程中使用的是这个模型——我们需要将其从 GeoTIFF 转换为 JPG。以下脚本将读取 GeoTIFF 文件并将其导出为 YOLO 算法可以读取的 JPG 文件。它将 JPG 文件保存在 "raw_inputs"文件夹中:
importrasteriofromrasterio.plotimportreshape_as_imagefromPILimportImageimportnumpyasnpimportmatplotlib.pyplotaspltwithrasterio.open("raw_inputs/drone.tif")assrc:red=src.read(1)green=src.read(2)blue=src.read(3)rgb_composite_n=np.dstack((red,green,blue))plt.imshow(rgb_composite_n)plt.axis('off')plt.margins(0,0)plt.gca().xaxis.set_major_locator(plt.NullLocator())plt.gca().yaxis.set_major_locator(plt.NullLocator())dpi=2400plt.savefig(f"raw_inputs/drone.jpg",dpi=dpi,bbox_inches='tight',pad_inches=0)在脚本中,只有一个名为 dpi 的参数,它会影响导出图像的质量。根据图像的分辨率和飞行高度,这个数值可能会变化。在我的情况下,我通过使用2400 dpi的分辨率在导出的图像中达到了足够的分辨率,但你可以根据你的图像调整它,直到你对输出结果的分辨率满意。
输出结果将包括:
在 CC-BY 4.0 许可下从openaerialmap.org/下载的航空影像,转换为 JPG 并由作者可视化
在 CC-BY 4.0 许可下从openaerialmap.org/下载的航空影像,转换为 JPG 并由作者可视化
📐将图像分割成子图像
要训练任何目标检测算法以在图像中检测特定对象,我们需要两组数据库:训练集和测试集。在这个例子中,因为我们有一个大图像,我们需要将这个图像分割成几个子图像,然后将这些子图像分类到训练组和测试组。以下代码将转换后的图像分割成子图像,并将它们存储在**"sub_images"文件夹中。它只有一个参数,指定了我们在 x 轴和 y 轴上想要将图像分割成多少个段。在这个模板中,我使用15**,这意味着图像被分割成一个15 x 15的子图像网格:
fromPILimportImage parameter=15# Read the imageimg=Image.open('raw_inputs/drone.jpg')# Get Dims of the imagewidth,height=img.size sub_width=width//parameter sub_height=height//parameter# Loop through each sub-image and save itforiinrange(parameter):forjinrange(parameter):left=j*sub_width top=i*sub_height right=left+sub_width bottom=top+sub_height sub_img=img.crop((left,top,right,bottom))sub_img.save(f'sub_images/sub_image_{i}_{j}.jpg')到目前为止,您应该在**"raw_inputs"文件夹中有drone.tif和drone.jpg**,在**"sub_images"文件夹中有225 张图像**。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/b6b8e71879b87c7e6c3cf804ede75d37.png
主目录中保存的无人机和子图像,图像由作者提供
让我们看看保存在"sub_images"文件夹中的一些图像:
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/386d29327958b87dd48d42d4099fc4c2.png
从原始图像创建用于训练和测试的子图像,图像由作者提供。
如您所见,我们将原始图像裁剪成更小的方形图像。
🏷️在子图像中标记车辆
在将这些子图像分割成训练集和测试集之前,我们需要在每个子图像中标记所有车辆。为了标记这些图像,我们将使用MakeSense.AI,这是一个用于标记图像的开源 Web 应用程序。要使用这个 Web 应用程序,您需要下载 sub_images 文件夹并将其上传到 MakeSense.AI Web 应用程序。以下代码将压缩 sub_images 文件夹并将其下载到您的本地计算机:
importshutil folder_name='sub_images'# Create a zip fileshutil.make_archive(folder_name,'zip',folder_name)# Download the zip filefromgoogle.colabimportfiles files.download(f'{folder_name}.zip')下载后,解压文件夹并将图像上传到 makesense.ai。然后您需要在每张图像中围绕每辆车绘制多边形并标记为“汽车”。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/7a62ddc8e6b5cb5444494c866aeb1096.png
使用makesense.ai标记车辆,图像由作者提供
我应该提到,在这个例子中的 225 个子图像(15 x 15 网格)中,你不需要上传所有图像。相反,你可以选择包含各种汽车的子图像。例如,由于我正在处理无人机图像,边缘图像的质量通常由于晕影问题而不太好,所以我将它们从我的列表中移除。此外,在分割原始图像后,一些子图像完全为黑色,因为它们与图像的角落相关。我也将它们移除了。还有,如果有不包含任何车辆的子图像,我也将它们移除了。拥有干净且高质量的子图像可以提高你标注图像的速度,并提高目标检测算法的准确性。以下是我从 “sub_images” 文件夹中移除的图像示例:
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/528af83412dcd8e526742ff147b40cf2.png
从 “sub_images” 文件夹中移除的图像示例,图片由作者提供。
当你完成所有图像中车辆的标注后,你需要从 MakeSense.AI 导出标签。然而,由于我们使用多边形工具进行图像标注,不幸的是,MakeSense.AI 只允许我们以 COCO JSON 格式导出它们,这与 YOLO 格式不兼容。因此,我们需要将其转换为与 YOLO 兼容的格式。
要这样做,首先在 MakeSense.AI 中点击“操作”选项卡,然后点击“导出标注”并选择“单个 COCO JSON 格式文件”下载文件。接下来,将其上传到“raw_inputs”文件夹,并运行以下代码将此 JSON 文件转换为 YOLO 可读的 TXT 文件:
importjsonimportos# Function to convert COCO bounding box to YOLO formatdefconvert_bbox(size,box):dw=1\./size[0]dh=1\./size[1]x=(box[0]+box[2]/2.0)*dw y=(box[1]+box[3]/2.0)*dh w=box[2]*dw h=box[3]*dhreturn(x,y,w,h)defconvert_coco_to_yolo(coco_json_path,output_dir):withopen(coco_json_path)asf:data=json.load(f)images={img['id']:imgforimgindata['images']}categories={cat['id']:catforcatindata['categories']}# Create a mapping from category ids to their corresponding YOLO class idscategory_mapping={cat_id:idxforidx,cat_idinenumerate(categories.keys())}forannotationindata['annotations']:image_id=annotation['image_id']category_id=annotation['category_id']bbox=annotation['bbox']image_info=images[image_id]image_width=image_info['width']image_height=image_info['height']yolo_bbox=convert_bbox((image_width,image_height),bbox)# Create or append to the corresponding YOLO format file for this imageyolo_file_path=os.path.join(output_dir,f"{image_info['file_name'].split('.')[0]}.txt")withopen(yolo_file_path,'a')asyolo_file:yolo_file.write(f"{category_mapping[category_id]}{' '.join(map(str,yolo_bbox))}n")coco_json_path='raw_inputs/labels_my-project-name_2024-07-04-01-18-59.json'output_dir='raw_inputs'convert_coco_to_yolo(coco_json_path,output_dir)如果代码运行成功,你将在 “raw_inputs” 文件夹中看到这些 TXT 文件:
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/dc4e998509113b430151ca66816dbbeb.png
从 COCO JSON 格式生成的 TXT 文件,包含每张图像中车辆的边界框信息,图片由作者提供。
🛠️ 准备训练和测试数据库
现在我们有了子图像以及标签 TXT 文件(每个汽车的边界框),我们可以将它们分成训练和测试数据集。我们之前创建了两个文件夹,images 和 labels,每个文件夹都包含 train 和 test 子文件夹,但它们目前都是空的。以下函数将随机从 sub_images 文件夹中选择 75% 的图像及其相应的 TXT 文件,并将它们分别放置在 images/train 和 labels/train 中。剩余的 25% 将放置在 images/test 和 labels/test 中:
importosimportrandomimportshutildefsplit_files(source_folder,labels_train_folder,labels_test_folder,images_train_folder,images_test_folder,train_ratio=0.75):# List all txt files starting with sub_image_txt_files=[fforfinos.listdir(source_folder)iff.startswith('sub_image_')andf.endswith('.txt')]# Shuffle the listrandom.shuffle(txt_files)split_index=int(len(txt_files)*train_ratio)# Split into training and testing setstrain_files=txt_files[:split_index]test_files=txt_files[split_index:]fortxt_fileintrain_files:# Move txt fileshutil.move(os.path.join(source_folder,txt_file),os.path.join(labels_train_folder,txt_file))# Replace .txt with .jpg and move themjpg_file=txt_file.replace('.txt','.jpg')ifos.path.exists(os.path.join(source_folder_image,jpg_file)):shutil.copy(os.path.join(source_folder_image,jpg_file),os.path.join(images_train_folder,jpg_file))fortxt_fileintest_files:# Move txt fileshutil.move(os.path.join(source_folder,txt_file),os.path.join(labels_test_folder,txt_file))# Replace .txt with .jpg and move themjpg_file=txt_file.replace('.txt','.jpg')ifos.path.exists(os.path.join(source_folder_image,jpg_file)):shutil.copy(os.path.join(source_folder_image,jpg_file),os.path.join(images_test_folder,jpg_file))source_folder='raw_inputs'source_folder_image='sub_images'labels_train_folder='labels/train'labels_test_folder='labels/test'images_train_folder='images/train'images_test_folder='images/test'split_files(source_folder,labels_train_folder,labels_test_folder,images_train_folder,images_test_folder,train_ratio=0.75)如果代码运行成功,你将在你的文件夹中有这些文件:
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/8c95d5163fdd3465c601dda01229d123.png
作者存储在 images/train、images/test、labels/train 和 labels/test 中的文件快照
🧠 训练 YOLOv8
下一步是创建 YOLO 算法的配置文件,并根据训练和测试目录中的子图像及其标签来训练模型。
以下代码创建用于训练 YOLO 物体检测模型的配置文件。此配置文件指定了训练和测试数据集的路径、类别数量以及这些类别的名称:
importyaml data={'train':'/content/images/train','val':'/content/images/test','nc':1,'names':['car']}# Write the yaml filewithopen('yolo_config.yaml','w')asfile:yaml.dump(data,file).yaml 文件将由 YOLO 训练算法用来理解如何找到图像以及如何解释标签。
现在,是时候使用配置文件来训练 YOLO 模型了:
fromultralyticsimportYOLOimportyamlimporttorchfromPILimportImageimportosimportcv2importtime# Choose the configyaml_filename='yolo_config.yaml'# Create Yolo modelmodel=YOLO('yolov8m.yaml').load('yolov8m.pt')# Train the modelmodel.train(data=f'{yaml_filename}',epochs=30)代码通过导入库和指定配置文件来设置环境。然后,它使用 YAML 配置文件创建 YOLO 模型并加载预训练的权重。最后,模型使用指定数据配置进行一定数量的周期训练。在这个过程中,YOLO 模型从提供的训练数据中学习,以检测和分类新图像中的对象。
**提示:**如果你想加快学习过程,请确保你在 Google Colab 环境中选择一个 GPU(点击更改运行时间并选择 T4 GPU):
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/81208ba181f996fa949ccc90da3e8657.png
在 Google Colab 环境中选择 T4 GPU,图片由作者提供
如果这一步执行成功,YOLO 将在你的目录中创建一个名为 “runs” 的新文件夹,其中包含这些文件:
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/96de6c27d0e3290cf63340c7935c1bd6.png
运行目录的快照,图片由作者提供
在这些文件中,best.pt 包含了训练过程中的最佳权重。我们将使用此文件来检测我们航空影像中的所有车辆。
🖼️ 在整个图像上部署 YOLOv8
在这一步,我们希望将训练好的模型部署到整个航空影像上,以检测和计数所有车辆并报告结果。然而,将模型应用于整个图像并不高效。相反,我们需要为所有子图像部署模型,然后将结果合并在一起。以下代码将使用训练好的模型来检测每个子图像中的所有车辆,计数车辆,并将总数保存为 TXT 文件:
# Load the trained weightsmodel=YOLO('runs/detect/train/weights/best.pt')# Run predictionim_dir='sub_images/'results=model.predict(source=im_dir,save=True,conf=0.9)total_cars=0# Iterate over the results to count the number of carsforresultinresults:# Count the number of detectionstotal_cars+=len(result.boxes)print(f'Total number of cars detected:{total_cars}')withopen('total_cars.txt','w')asf:f.write(f'Total number of cars detected:{total_cars}n')一旦运行代码,YOLO 将在 runs 目录下创建另一个名为 “predicts” 的文件夹,其中包含每个子图像中每个车辆的边界框和置信度分数:
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/e1fd8f4ee5e2b94ee2c962d315fb31c6.png
预测文件夹的快照,图片由作者提供
如果你打开它们,你会看到边界框以及置信度分数:
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/16909d41a8e86407fdae534619547c37.png
YOLO 在其中一个子图像中检测到的车辆,图片由作者提供。
此外,您还可以查看 total_cars.txt 文件,以了解 YOLO 在所有子图像中检测到的车辆总数。我的总数是185!
现在在每个子图像中检测到车辆后,我们可以运行以下脚本将它们合并在一起,以查看原始航拍图像中检测到的车辆:
fromPILimportImage# Set the number of rows and columnsparameter=15# Create a new image with the size of the sub-imagessub_image_size=Image.open('runs/detect/predict/sub_image_0_0.jpg').size image_width=sub_image_size[0]*parameter image_height=sub_image_size[1]*parameter image=Image.new('RGB',(image_width,image_height))# Loop through the sub-images and merge themforiinrange(parameter):forjinrange(parameter):sub_image=Image.open(f'runs/detect/predict/sub_image_{i}_{j}.jpg')x=j*sub_image_size[0]y=i*sub_image_size[1]image.paste(sub_image,(x,y))# Save itimage.save('combined_image.jpg')输出结果将如下:
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/724ee2c92ab7954a6a6595df50f9da0f.png
YOLO 在全部航拍图像中检测到的车辆,图片由作者提供
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/93177e2ad544c684335ffb550eedbbbd.png
YOLO 在全部航拍图像中检测到的车辆(在几个地方进行了放大),图片由作者提供。
📄 结论
在本教程中,我介绍了如何找到免费的航拍图像,将其从 GeoTIFF 转换为 JPG,将其分割成子图像,然后将子图像分类为训练集和测试集。我还使用训练集对基于预训练模型的 YOLO 模型进行了微调,并将其最终部署在全部图像集上。学习这项技术有助于读者将其应用于库存管理、交通流量分析或停车需求评估等实际应用。希望您喜欢阅读,如有任何问题、评论或反馈,请随时联系。
📚 参考文献
OpenAerialMap
Piotr Skalski.Make Sense. 2019. 可在:github.com/SkalskiP/make-sense/
Glenn Jocher, Ayush Chaurasia, Jing Qiu.Ultralytics YOLO. 版本 8.0.0, Ultralytics, 2023. 可在:ultralytics.com. 仓库:github.com/ultralytics/ultralytics. 发布于 2023–01–10. ORCIDs: Glenn Jocher, Ayush Chaurasia, Jing Qiu.
📱 在其他平台上与我联系,获取更多精彩内容!LinkedIn,ResearchGate,Github, 和Twitter.
通过以下链接可以找到相关帖子:
哪种版本的 Segment Anything Model (SAM)可以检测更多对象?(Python)
创建用于卫星图像可视化的 Streamlit 应用:一步步指南
从太空观看风暴:创建惊人视图的 Python 脚本