Pytorch框架下的maskrcnn实例分割（含COCO格式转换）

人生忽如寄-

已于 2024-09-19 23:50:36 修改

阅读量1.3k

点赞数 13

文章标签： pytorch 人工智能 python 计算机视觉深度学习

于 2024-08-31 14:52:24 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_48679164/article/details/141730568

版权

一、数据准备

labelme2coco: 将LabelMe注解转换为COCO格式

1. 安装labelme2coco

pip install -U labelme2coco

2. 修改文件名
由于我的数据集图片名称中带有_w后缀labelmecoco无法识别转换，因此需要编写个rename.py脚本修改带_w的文件名，另外修改json文件名的同时要修改json文件内容，将imagePath一同修改

# #把图片文件名中_w部分删去并在前面+1
# import os
#
# # 指定要重命名文件的文件夹路径
# folder_path = "C:/.../maskrcnn-benchmark-main/maskrcnn-benchmark-main/data/json"
#
# # 遍历文件夹中的所有文件
# for filename in os.listdir(folder_path):
#     # 检查文件名中是否包含"_w"
#     if "_w" in filename:
#         # 获取文件的扩展名（如.jpg）
#         file_extension = os.path.splitext(filename)[1]
#         # 去掉文件名中的"_w"
#         new_filename = filename.replace("_w", "")
#         # 在文件名前加上"1"
#         new_filename = "1" + new_filename
#         # 完整的新文件名
#         new_filename_with_extension = new_filename + file_extension
#         # 获取完整的旧文件路径和新文件路径
#         old_file_path = os.path.join(folder_path, filename)
#         new_file_path = os.path.join(folder_path, new_filename_with_extension)
#         # 重命名文件
#         os.rename(old_file_path, new_file_path)
#         print(f"文件已重命名: {old_file_path} -> {new_file_path}")
# #
# ##修复双重扩展名问题；
# import os
#
# # 指定要修复文件扩展名的文件夹路径
# folder_path = "C:/.../maskrcnn-benchmark-main/maskrcnn-benchmark-main/data/json"
#
# # 遍历文件夹中的所有文件
# for filename in os.listdir(folder_path):
#     # 检查文件名中是否包含两个扩展名，如.jpg.jpg或.json.json
#     if filename.endswith(".jpg.jpg") or filename.endswith(".json.json"):
#         # 去掉多余的扩展名
#         new_filename = filename.replace(".jpg.jpg", ".jpg").replace(".json.json", ".json")
#         # 获取完整的旧文件路径和新文件路径
#         old_file_path = os.path.join(folder_path, filename)
#         new_file_path = os.path.join(folder_path, new_filename)
#         # 重命名文件
#         os.rename(old_file_path, new_file_path)
#         print(f"文件已修复: {old_file_path} -> {new_file_path}")
#遍历文件夹中的所有 JSON 文件，并将 imagePath 字段的值修改为与 JSON 文件名相对应的值：
import os
import json

# 指定包含 JSON 文件的文件夹路径
folder_path = "C:/.../maskrcnn-benchmark-main/maskrcnn-benchmark-main/data/json"

# 遍历文件夹中的所有文件
for filename in os.listdir(folder_path):
    if filename.endswith(".json"):
        # 获取完整的 JSON 文件路径
        file_path = os.path.join(folder_path, filename)

        # 读取 JSON 文件内容
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)

        # 获取新的 imagePath 值
        new_image_path = filename.replace(".json", ".jpg")

        # 修改 JSON 中的 imagePath 值
        if 'imagePath' in data:
            old_image_path = data['imagePath']
            data['imagePath'] = new_image_path
            print(f"修改 {file_path} 中的 imagePath: {old_image_path} -> {new_image_path}")

        # 将修改后的内容写回 JSON 文件
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=4, ensure_ascii=False)

print("所有 JSON 文件中的 imagePath 字段已修改完成。")

3. 运用labelme2coco转换
将文件名修改好后，运用labelme2coco转换成coco格式：

labelme2coco path/to/labelme/dir

dir为labelme输出的文件夹包含jpg+json
labelme2coco.py的地址为：C:\Users\llxc120401\anaconda3\envs\maskrcnn\Lib\site-packages\labelme2coco\labelme2coco.py

4. 划分数据集

 import json
import os
import random
import shutil

def split_coco_dataset(json_file, img_dir, output_dir, train_ratio=0.8, val_ratio=0.1, test_ratio=0.1):
    with open(json_file, 'r') as f:
        coco_data = json.load(f)

    images = coco_data['images']
    annotations = coco_data['annotations']

    # 随机打乱图像顺序
    random.shuffle(images)

    # 计算每个子集的大小
    total_images = len(images)
    train_count = int(total_images * train_ratio)
    val_count = int(total_images * val_ratio)
    test_count = total_images - train_count - val_count

    # 划分图像
    train_images = images[:train_count]
    val_images = images[train_count:train_count + val_count]
    test_images = images[train_count + val_count:]

    # 创建图像ID集合以快速查找
    train_image_ids = {img['id'] for img in train_images}
    val_image_ids = {img['id'] for img in val_images}
    test_image_ids = {img['id'] for img in test_images}

    # 划分注释
    train_annotations = [ann for ann in annotations if ann['image_id'] in train_image_ids]
    val_annotations = [ann for ann in annotations if ann['image_id'] in val_image_ids]
    test_annotations = [ann for ann in annotations if ann['image_id'] in test_image_ids]

    # 创建新的COCO数据结构
    coco_template = {
        "info": coco_data.get('info', {}),
        "licenses": coco_data.get('licenses', []),
        "categories": coco_data.get('categories', []),
    }

    train_data = coco_template.copy()
    train_data['images'] = train_images
    train_data['annotations'] = train_annotations

    val_data = coco_template.copy()
    val_data['images'] = val_images
    val_data['annotations'] = val_annotations

    test_data = coco_template.copy()
    test_data['images'] = test_images
    test_data['annotations'] = test_annotations

    # 保存分割后的数据集
    os.makedirs(output_dir, exist_ok=True)

    with open(os.path.join(output_dir, 'train_dataset.json'), 'w') as f:
        json.dump(train_data, f, indent=4)

    with open(os.path.join(output_dir, 'val_dataset.json'), 'w') as f:
        json.dump(val_data, f, indent=4)

    with open(os.path.join(output_dir, 'test_dataset.json'), 'w') as f:
        json.dump(test_data, f, indent=4)

    # 创建图像存储目录
    train_img_dir = os.path.join(output_dir, 'train_images')
    val_img_dir = os.path.join(output_dir, 'val_images')
    test_img_dir = os.path.join(output_dir, 'test_images')

    os.makedirs(train_img_dir, exist_ok=True)
    os.makedirs(val_img_dir, exist_ok=True)
    os.makedirs(test_img_dir, exist_ok=True)

    # 移动图像文件到相应目录


    def copy_images(images, target_dir):
        for img in images:
            img_filename = img['file_name']
            src_path = os.path.join(img_dir, img_filename)
            dst_path = os.path.join(target_dir, img_filename)

            if not os.path.exists(src_path):
                print(f"未找到文件，跳过: {src_path}")
                continue
            shutil.copy(src_path, dst_path)

    copy_images(train_images, train_img_dir)
    copy_images(val_images, val_img_dir)
    copy_images(test_images, test_img_dir)

    print(f"数据集已划分并保存到 {output_dir} 中。")

# 使用你的数据集路径
json_file = "C:/.../maskrcnn-benchmark-main/maskrcnn-benchmark-main/runs/labelme2coco/dataset.json"
img_dir = "C:/.../maskrcnn-benchmark-main/maskrcnn-benchmark-main/data/pic"
output_dir = "C:/.../maskrcnn-benchmark-main/maskrcnn-benchmark-main/split_data"

split_coco_dataset(json_file, img_dir, output_dir)

5. 修改maskrcnn-benchmark/paths_catalogs.py文件中的数据地址

class DatasetCatalog(object):
    DATA_DIR = "datasets"
    DATASETS = {
        "my_dataset_train": {
            "img_dir": "C:/Users/.../maskrcnn-benchmark-main/maskrcnn-benchmark-main/split_data/train_images",
            "ann_file": "C:/Users/.../maskrcnn-benchmark-main/maskrcnn-benchmark-main/split_data/train_dataset.json"
        },
        "my_dataset_val": {
            "img_dir": "C:/Users/.../maskrcnn-benchmark-main/maskrcnn-benchmark-main/split_data/val_images",
            "ann_file": "C:/Users/.../maskrcnn-benchmark-main/maskrcnn-benchmark-main/split_data/val_dataset.json"

二、模型训练

1. 环境搭建

conda create --name maskrcnn_benchmark -y
conda activate maskrcnn_benchmark

conda install ipython
pip install ninja yacs cython matplotlib tqdm opencv-python
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

#we give the instructions for CUDA 11.8
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

git clone https://github.com/cocodataset/cocoapi.git
cd cocoapi/PythonAPI
python setup.py build_ext install

git clone https://github.com/NVIDIA/apex.git
cd apex
python setup.py install --cuda_ext --cpp_ext

python setup.py build develop

2. 关于报错

一般会涉及电脑是否安装了Visual Studio\Git\APEX，以及CUDA版本和VS版本的兼容问题

可通过nvcc --version以及Visual Studio Installer查看两者的版本
Visual Studio下载地址:https://visualstudio.microsoft.com/zh-hans/downloads/，下来找到用于Visual Studio的工具选择下载Visual Studio生成工具，安装时记得勾选Desktop development with C++。
如果Git代码下不下来则需手动下载：https://git-scm.com/download/win，下载完成后需要在高级系统设置-环境变量-系统变量，点击“编辑”-“新建” ，添加地址：C:\Program Files\Git\cmd;C:\Program Files\Git\bin, 然后重新打开终端即可继续别的操作

3. 模型训练

Pre-trained models, baselines and comparison with Detectron and mmdetection
can be found in MODEL_ZOO.md

运行代码：

python train_net.py --config-file "C:\Users\...\maskrcnn-benchmark-main\maskrcnn-benchmark-main\configs\e2e_mask_rcnn_R_50_FPN_1x.yaml"

整体逻辑：

1 修改文件名_w改为数字重新用labelme2coco转换coco格式
2.使用datasplit.py拆分训练测试验证
3.修改catalogs.py数据集地址
4.修改相关代码device为cpu，注释掉gpu相关函数
5.python train_net.py --config-file “C:\Users\llxc120401\Desktop\硕士学位论文\image segment\maskrcnn-benchmark-main\maskrcnn-benchmark-main\configs\e2e_mask_rcnn_R_50_FPN_1x.yaml”

小技巧

1. 搜索项目中的某代码在什么脚本文件中
打开cmd,定位到项目文件夹，findstr /s /i “amp.scale_loss” *.py
2. 定位脚本文件中的代码
打开该脚本文件，快捷键shift+F,，即可打开搜索栏
3. 文件目录过长，代码报错

win+R打开“运行”，输入 “regedit”,导航至 "HKEY_LOCAL_MACHINE/SYSTEM\CurrentControlSet\Control\FileSystem\LongPathsEnabled”,双击，将数值数据改为1，重启电脑即可。
若还是报错，则cmd，以管理员权限打开，输入“mklink /D C:\shortpath “C:\Users…\maskrcnn-benchmark-main\maskrcnn-benchmark-main””，之后“cd C:\shortpath"即可

参考博文：
5. 基于TensorFlow和Keras框架的maskrcnn
6. 基于TensorFlow的mask r-cnn(cpu版）
7. Mask R-CNN 1.0 2.0 2.1版本源码