欢迎访问我的网站,希望内容对您有用,感兴趣的可以加入我们的社群。

COCO 格式转 YOLO

YOLO 迷途小书童 4年前 (2021-11-18) 4401次浏览 1个评论

环境

  • window 10 64bit
  • coco yolo

前言

前文 MS COCO数据集 介绍过了 COCO 数据集,COCO 是将所有图片对应的标注信息写在了一个 json 文件里,如下

coco to yolo

因此要将 coco 格式的数据集转换成 yolo 格式,本质上就是去解析这个 json 文件,将对应图片的标注信息提取出来,写入 txt 文件中

实操

这里还是使用我们熟悉的 roboflow 平台上的口罩数据集,下载地址是 https://public.roboflow.com/object-detection/mask-wearing/4

下载的格式选择 COCO

coco to yolo

解压后,进入 train 文件夹,可以看到 json 文件和图片文件

coco to yolo

在同级目录创建 python 脚本 coco2yolo.py,输入如下代码,少量注释嵌入到代码中了

  1. import os
  2. import json
  3. from tqdm import tqdm
  4. import argparse
  5. def convert(size, box):
  6. '''
  7. size: 图片的宽和高(w,h)
  8. box格式: x,y,w,h
  9. 返回值:x_center/image_width y_center/image_height width/image_width height/image_height
  10. '''
  11. dw = 1. / (size[0])
  12. dh = 1. / (size[1])
  13. x = box[0] + box[2] / 2.0
  14. y = box[1] + box[3] / 2.0
  15. w = box[2]
  16. h = box[3]
  17. x = x * dw
  18. w = w * dw
  19. y = y * dh
  20. h = h * dh
  21. return (x, y, w, h)
  22. if __name__ == '__main__':
  23. parser = argparse.ArgumentParser()
  24. parser.add_argument('--json_file', default='test.json',
  25. type=str, help="coco file path")
  26. parser.add_argument('--save_dir', default='labels', type=str,
  27. help="where to save .txt labels")
  28. arg = parser.parse_args()
  29. data = json.load(open(arg.json_file, 'r'))
  30. # 如果存放txt文件夹不存在,则创建
  31. if not os.path.exists(arg.save_dir):
  32. os.makedirs(arg.save_dir)
  33. id_map = {}
  34. # 解析目标类别,也就是 categories 字段,并将类别写入文件 classes.txt 中
  35. with open(os.path.join(arg.save_dir, 'classes.txt'), 'w') as f:
  36. for i, category in enumerate(data['categories']):
  37. f.write(f"{category['name']}\n")
  38. id_map[category['id']] = i
  39. for img in tqdm(data['images']):
  40. # 解析 images 字段,分别取出图片文件名、图片的宽和高、图片id
  41. filename = img["file_name"]
  42. img_width = img["width"]
  43. img_height = img["height"]
  44. img_id = img["id"]
  45. head, tail = os.path.splitext(filename)
  46. # txt文件名,与对应图片名只有后缀名不一样
  47. txt_name = head + ".txt"
  48. f_txt = open(os.path.join(arg.save_dir, txt_name), 'w')
  49. for ann in data['annotations']:
  50. if ann['image_id'] == img_id:
  51. box = convert((img_width, img_height), ann["bbox"])
  52. # 写入txt,共5个字段
  53. f_txt.write("%s %s %s %s %s\n" % (
  54. id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))
  55. f_txt.close()

执行上述代码

  1. python coco2yolo.py --json_file _annotations.coco.json --save_dir labels

完成后,还是在同级目录,就能看到 labels 文件夹及里面的 txt 标注文件了

coco to yolo

关联阅读

喜欢 (1)

您必须 登录 才能发表评论!

(1)个小伙伴在吐槽
  1. 报错了个KeyError:'categories'
    匿名2022-04-14 21:19