labelme2coco.py 脚本增强

前言

最近在做一个实例分割的项目,用使用 Segment anything 的脚本先处理一遍图片,然后用labelme人工审核和调整。(labelme 新版本已经加入了 Segment anything,可以AI 抠图,可惜用的是onnxruntime,于是我折腾一番改成了onnxruntime-gpu。最后,个人感觉用cuda扣的比cpu好。)

扣完图按labelme项目中Readme.md指导的,接着生成coco数据集格式的数据集。结果,发现原脚本功能不能满足我的需求:

  • 它不划分测试集和训练集,coco2017是分train2017和val2017,不划分就得改训练脚本。
  • Segment anything 抠图时有些粘连部分没处理好,人工检查没发现,转化时生成的Visualization图片能更好的发现这些,于是需要更新,原脚本没有更新功能,只能从头开始生成全部信息。
  • 工业项目的像素不像coco数据集那么低,训练时要均衡显存和分割效果,要进行一定的缩放。当然可以放在训练前做预处理(既要reszie图片又要resize boxes 和 masks),但是你会发现gpu工作效率下降(gpu在等待cpu处理)。

增强

新加入参数

python 复制代码
    parser.add_argument("--split_rate", help="val file rate", default=0.1)
    parser.add_argument("--scale_rate", help="val file rate", type=float, default=1.0)
    parser.add_argument("--update_dir", help="update input annotated directory,"
                                             "It will check json with input_dir. If changed generate new",
                        type=str, default=None)
  • split_rate 训练集测试集分割比例
  • scale_rate 缩放比例(实现方式是长宽乘以浮点数后强转int,不是四舍五入)
  • update_dir 指定一个annotations目录(也就是新版本的json目录),比对老目录的json,改变的重新生成并覆盖老目录。

划分数据集

分割数据集实现:

python 复制代码
    # 划分 train 和 val
    split_rate = args.split_rate
    indices = list(range(len(label_files)))
    random.shuffle(indices)
    train_size = int((1.0 - split_rate) * len(label_files))
    train_indices = indices[:train_size]
    val_indices = indices[train_size:]

这就是一个随机实现,主要差别是后面,划分后其实是2个非常独立的任务,放到2个进程去加速生成。

python 复制代码
    print('启动2个进程分别执行')
    pool = multiprocessing.Pool()
    pool.apply_async(gen_coco_dateset, args=(
        train_indices, label_files,
        train_path, copy.deepcopy(data),
        class_name_to_id, train_ann_file,
        args.noviz, train_visual_path, args.scale_rate
    ))
    pool.apply_async(gen_coco_dateset, args=(
        val_indices, label_files,
        val_path, copy.deepcopy(data),
        class_name_to_id, val_ann_file,
        args.noviz, val_visual_path, args.scale_rate
    ))
    pool.close()
    pool.join()
    print('所有进程执行完毕')

更新

更新主要流程:

  1. 遍历更新文件目录,区别新增和更新(新增我没有实现)
  2. 判断是否需要更新
  3. 判断在训练集更新还是验证集更新
  4. 最后处理完coco集的json后,覆盖老目录的json文件

更新时最主要要调整的是,原来最后生成单个json文件的 annotations 。

处理时

python 复制代码
# 删除所有 annotations 中关于此图的所有标记信息
data['annotations'] = [item for item in data['annotations'] if item['image_id'] != image_id]

最后要写入时,重新整理

python 复制代码
# 重新排序 data['annotations']
for i, item in enumerate(train_data['annotations']):
	item['id'] = i
for i, item in enumerate(val_data['annotations']):
	item['id'] = i

完整脚本

我fork的labelme项目原labelme2coco.py同级labeme2coco_v2.py

解释一番是希望,能读者能修改我的脚本来满足自己的需求。

相关推荐
Teacher.chenchong24 分钟前
现代R语言机器学习:Tidymodel/Tidyverse语法+回归/树模型/集成学习/SVM/深度学习/降维/聚类分类与科研绘图可视化
机器学习·回归·r语言
AndrewHZ26 分钟前
【图像处理基石】如何入门色彩评估?
图像处理·人工智能·深度学习·色彩科学·hvs·色彩评估·颜色工程
TomatoSCI26 分钟前
聚类的可视化选择:PCA / t-SNE丨TomatoSCI分析日记
人工智能·机器学习
大咖分享课28 分钟前
深度剖析:最新发布的ChatGPT Agent 技术架构与应用场景
人工智能·openai·智能助手·ai代理·chatgpt agent·自主任务执行
lucky_lyovo38 分钟前
卷积神经网络--网络性能提升
人工智能·神经网络·cnn
liliangcsdn42 分钟前
smolagents - 如何在mac用agents做简单算术题
人工智能·macos·prompt
nju_spy1 小时前
周志华《机器学习导论》第8章 集成学习 Ensemble Learning
人工智能·随机森林·机器学习·集成学习·boosting·bagging·南京大学
星座5281 小时前
基于现代R语言【Tidyverse、Tidymodel】的机器学习方法与案例分析
机器学习·r语言·tidyverse·tidymodel
静心问道1 小时前
TrOCR: 基于Transformer的光学字符识别方法,使用预训练模型
人工智能·深度学习·transformer·多模态
说私域1 小时前
基于开源AI大模型、AI智能名片与S2B2C商城小程序源码的用户价值引导与核心用户沉淀策略研究
人工智能·开源