mat标注文件转COCO格式标注文件（shanghaiTech数据集）

johnhamtom

已于 2023-01-11 15:40:45 修改

阅读量794

点赞数 4

文章标签：计算机视觉图像处理

于 2020-02-24 17:19:01 首次发布

本文链接：https://blog.csdn.net/johnhamtom/article/details/104481186

版权

1.数据集介绍

shanghaiTech数据集是一个用于行人计数的数据集，里面包含了大量的人群图片，它由两部分组成：part_A_final文件夹和 part_B_final文件夹，它们内部都由 train_data 和 test_data 两部分组成，意思为训练集和测试集，对于train_data 和 test_data 其中包括 images 和 ground_truth，images储存着图片文件、ground_truth里储存着对应的标注文件，一张图片对应一个标注文件。
shanghaiTech结构如下所示：
在这里插入图片描述
标注文件为mat文件，标注方式为点标注，标注位置为人头中心。

2.mat标注文件转COCO格式标注文件（shanghaiTech数据集）

参考CSRNet的官方源码(https://github.com/leeyeehoo/CSRNet-pytorch)中的make_dataset.py可以得知，读取mat标注文件并获取其点坐标的代码如下：

mat = io.loadmat(img_path.replace('.jpg','.mat').replace('images','ground_truth').replace('IMG_','GT_IMG_'))
gt = mat["image_info"][0,0][0,0][0] #gt为用列表存储的点坐标对

知道如何获取mat文件的点坐标后，参考网络上的YOLO格式文件转化为COCO格式文件的脚本就可以写出mat标注文件转COCO格式标注文件的脚本了。
定义Mat2COCO函数，传入的参数为图片地址、shanghaiTech根目录地址和COCO文件保存地址，然后初始化变量记录标注文件地址、图片地址、图片名、标注点类别名和字典型变量dataset用于记录COCO信息，代码如下：

def Mat2COCO(path_sets,root,save_path):
    assert os.path.exists(path_sets)
    LabelDir = path_sets.replace('images','ground_truth') # 标签位置
    ImageDir = path_sets                                  # 图片位置
    with open(os.path.join(root, 'classes.txt')) as f:
        classes = f.read().strip().split()                      # 类别
    indexes = os.listdir(ImageDir)                       # 图片名
    dataset = {'categories': [], 'annotations': [], 'images': [],'licenses':[]} #字典型变量用于记录COCO信息

读取categories（类别）信息并写入dataset变量中，licenses信息一般用不上，这里就进行了简单的置空操作，代码如下：

    for i, cls in enumerate(classes, 0):
        print(i)
        print(cls)
        dataset['categories'].append({'id': i+1, 'name': cls, 'supercategory': 'person'})

    dataset['licenses'].append({
        'name': "",
        'id': 0,
        'url': "",
    })

读取图片信息写入dataset变量的images里，读取标注信息写入dataset变量的annotations里，代码里的注释已经解释的比较清楚了，这里再做几点说明，（1）labelList 里存储的是多个坐标对，格式类似[ [坐标x1，坐标y1] , [坐标x2，坐标y2 ] …]；（2）由于shanghaiTech是点标注且只有person类别的点标注，所以类别id：cls_id置为了固定值，area的值和bbox里的后两个值设置为了0；（3）keypoints记录的就是关键点坐标（x,y）及可见性信息（0不可见，1被遮挡，2完全可见）由于shanghaiTech里并没有提供可见性信息，我这里就将其统一设置为了2。

    ann_id_cnt = 0
    for k, index in enumerate(tqdm(indexes)):
        # jpg 格式的图片。
        matFile = index.replace('.jpg','.mat').replace('IMG_','GT_IMG_')
        # 读取图像的宽和高
        im = cv2.imread(os.path.join(ImageDir, index))
        height, width, _ = im.shape
        # 添加图像的信息
        dataset['images'].append({'file_name': index,
                                    'id': k+1,
                                    'width': width,
                                    'height': height,
                                  'license': 0,
                                  'date_captured': 0,
                                  })
        if not os.path.exists(os.path.join(LabelDir, matFile)):
            # 如没标签，跳过，只保留图片信息。
            print("无标签："+os.path.join(LabelDir, matFile))
            continue
        print(os.path.join(LabelDir, matFile))

        # 读取标注点坐标信息
        mat = io.loadmat(os.path.join(LabelDir, matFile))
        labelList = mat["image_info"][0, 0][0, 0][0]
        
        for label in labelList:
            x = float(label[0])
            y = float(label[1])
            cls_id = 1
            width = 0.0
            height = 0.0
            dataset['annotations'].append({
                'area': width * height,
                'bbox': [x, y, width, height],
                'category_id': cls_id,
                'id': ann_id_cnt+1,
                'image_id': k+1,
                'iscrowd': 0,
                'num_keypoints': 1,
                'segmentation': [],
                'attributes': {"occluded": False, "track_id": 0, "keyframe": True},
                'keypoints':[x,y,2]
            })
            ann_id_cnt += 1

保存dataset变量为json文件，代码如下：

    folder = os.path.join(root, 'annotations')
  if not os.path.exists(folder):
      os.makedirs(folder)     #如果没有annotations文件夹就建立一个

  json_name = os.path.join(root, 'annotations/{}'.format(save_path))
  with open(json_name, 'w') as f:
      json.dump(dataset, f)
      print('Save annotation to {}'.format(json_name))

最后在主函数中调用Mat2COCO实现建立所需的COCO文件，如下的代码展示了对part_A和part_B的训练集和测试集建立了对应的COCO格式文件的代码逻辑。

  # 设置根目录地址
 root = r'.\Shanghai'
 # 根据根目录地址合成图片地址
 part_A_train = os.path.join(root,'part_A_final','train_data','images')
 part_A_test = os.path.join(root,'part_A_final','test_data','images')
 part_B_train = os.path.join(root,'part_B_final','train_data','images')
 part_B_test = os.path.join(root,'part_B_final','test_data','images')
 path_sets = [part_A_train,part_A_test,part_B_train,part_B_test]
 save_path = ['part_A_train','part_A_test','part_B_train','part_B_test']
 i=0
 for path in path_sets:
     Mat2COCO(path,root,save_path[i])
     i=i+1

综上，完整代码如下：

import scipy.io as io
import json
from image import *
from tqdm import tqdm

def Mat2COCO(path_sets,root,save_path):
    assert os.path.exists(path_sets)
    LabelDir = path_sets.replace('images','ground_truth') # 标签位置
    ImageDir = path_sets                                  # 图片位置
    with open(os.path.join(root, 'classes.txt')) as f:
        classes = f.read().strip().split()                      # 类别
    indexes = os.listdir(ImageDir)                       # 图片名
    dataset = {'categories': [], 'annotations': [], 'images': [],'licenses':[]} #字典型变量用于记录COCO信息

    for i, cls in enumerate(classes, 0):
        print(i)
        print(cls)
        dataset['categories'].append({'id': i+1, 'name': cls, 'supercategory': 'person'})

    dataset['licenses'].append({
        'name': "",
        'id': 0,
        'url': "",
    })
    # 标注的id
    ann_id_cnt = 0
    for k, index in enumerate(tqdm(indexes)):
        # jpg 格式的图片。
        matFile = index.replace('.jpg','.mat').replace('IMG_','GT_IMG_')
        # 读取图像的宽和高
        im = cv2.imread(os.path.join(ImageDir, index))
        height, width, _ = im.shape
        # 添加图像的信息
        dataset['images'].append({'file_name': index,
                                    'id': k+1,
                                    'width': width,
                                    'height': height,
                                  'license': 0,
                                  'date_captured': 0,
                                  })
        if not os.path.exists(os.path.join(LabelDir, matFile)):
            # 如没标签，跳过，只保留图片信息。
            print("无标签："+os.path.join(LabelDir, matFile))
            continue
        print(os.path.join(LabelDir, matFile))

        # 读取标注点坐标信息
        mat = io.loadmat(os.path.join(LabelDir, matFile))
        labelList = mat["image_info"][0, 0][0, 0][0]

        for label in labelList:
            x = float(label[0])
            y = float(label[1])
            cls_id = 1
            width = 0.0
            height = 0.0
            dataset['annotations'].append({
                'area': width * height,
                'bbox': [x, y, width, height],
                'category_id': cls_id,
                'id': ann_id_cnt+1,
                'image_id': k+1,
                'iscrowd': 0,
                'num_keypoints': 1,
                'segmentation': [],
                'attributes': {"occluded": False, "track_id": 0, "keyframe": True},
                'keypoints':[x,y,2]
            })
            ann_id_cnt += 1


    folder = os.path.join(root, 'annotations')
    if not os.path.exists(folder):
        os.makedirs(folder)     #如果没有annotations文件夹就建立一个

    json_name = os.path.join(root, 'annotations/{}'.format(save_path))
    with open(json_name, 'w') as f:
        json.dump(dataset, f)
        print('Save annotation to {}'.format(json_name))


if __name__ == "__main__":
    # 设置根目录地址
    root = r'.\Shanghai'
    # 根据根目录地址合成图片地址
    part_A_train = os.path.join(root,'part_A_final','train_data','images')
    part_A_test = os.path.join(root,'part_A_final','test_data','images')
    part_B_train = os.path.join(root,'part_B_final','train_data','images')
    part_B_test = os.path.join(root,'part_B_final','test_data','images')
    path_sets = [part_A_train,part_A_test,part_B_train,part_B_test]
    save_path = ['part_A_train','part_A_test','part_B_train','part_B_test']
    i=0
    for path in path_sets:
        Mat2COCO(path,root,save_path[i])
        i=i+1

转换结果如下：
在这里插入图片描述

3.COCO可视化

为了验证转换结果是否正确，参考博客COCO数据集可视化程序(包括bbox和segmentation)
（https: // blog.csdn.net / wtandyn / article / details / 109751015）的代码，进行了可视化验证，即将标注结果展示在图片上，黄点为标注点，展示如下图所示：
在这里插入图片描述

从展示结果上来看，转换生成的COCO文件应该是没有什么问题的。