OVIS数据集代码解析

OVIS数据集格式整体和COCO类似，但是是以video的形式存储的，对应的解析代码见：https://github.com/qjy981010/cocoapi/blob/main/PythonAPI/pycocotools/ovis.py。

由于OVIS仅train提供了标注，因此，这里均以train进行说明。

1.init函数

创建一些字典，包括dataset、anns、cats、vids以及vidToAnns、catToVids。

其中dataset是原始标注数据，由json文件读取得到。
这里对其进行详细说明

dataset包括5个key：

• info: 主要是存储一些meta信息，如下：

‘description’: ‘OVIS’,
‘url’: ‘http://songbai.site/ovis/’,
‘version’: ‘1.0’,
‘year’: 2021,
‘contributor’: ‘youku’,
‘date_created’: ‘2021-01-01’

• licenses: 存储licenses信息。

• categories: 类别信息。OVIS一共包括25个类别，每个类别包括三个信息：supercategory表示该类别的超类，OVIS都设置为object，id表示该类别的id，而name则是该类别的字符串名称。

  具体类别及对应实例数目如下：
  
  可以看出，person是主要类别。

• videos: 视频详细信息。训练集一共包含607段视频，每段视频的信息以一个字典构成，因此，videos是一个长为607的列表。

  每个视频信息字典包含6个基本属性：

  • width: 视频的宽
  • height: 视频的长
  • license: license id
  • file_names: 对应图片的相对路径，例如’85aa3b0e/img_0000001.jpg’，说明该图片位于85aa3b0e视频文件夹下，图片名称为img_0000001.jpg
  • id: 视频id
  • length: 该视频采样得到的图片数目

• annotations: 该视频中所有标注实例的信息。这里和coco存在一定差异。由于OVIS是一个视频数据集，这里的实例是针对视频进行描述的。例如一个视频假设有500帧，某一个人物在这个视频中出现了，那么该人物就是一个实例，其在该视频中500帧图片中的实例id都是一致的。这里的标注就是采用以上设置给出的。

  每个实例标注由一个字典表示。每个字典有11个基本属性：

  • length: 数值全为1，不太清楚具体含义
  • category_id: 该实例的类别id
  • video_id: 该实例对应所在的视频id
  • iscrowd: 这里的iscrowd应该和coco不太一样，不用于区分segmentation的类别。（在coco中，iscrowd=0时意味着segmentation的格式为polygon，iscrowd=1时则意味着segmentation的格式为rle）
  • id: 该实例对应的id
  • height: 视频帧的长
  • width: 视频帧的宽
  • segmentations: 该实例的mask标注信息，每一帧都会给出对应的标注，如果当前帧没有该实例，则为None。这里的标注信息以字典形式给出：
    • size: 该标注对应的[height, width]，用于rle解码
    • counts: rle编码
  • bboxes: 该实例的bbox标注信息，每一帧都会给出对应的标注。这里的标注以四元组 (x, y, h, w) 形式给出，其中 (x,y) 是左上角坐标。
  • areas: 该实例对应的bbox的面积，每一帧都给出。
  • occulusion: 该实例的遮挡情况，分为三类no occlusion, slight occlusion 以及 severe occlusion，具体说明可参考原论文。

其余字典（anns、cats、vids以及vidToAnns、catToVids）则是通过createIndex进行填充。

2. createIndex

该函数主要是用于填充在init中创建的一些属性字典。

for ann in self.dataset['annotations']:vidToAnns[ann['video_id']].append(ann)anns[ann['id']] = ann

vidToAnns: key为视频的id，而value则为一个列表，里面存储了视频id为key的所有实例信息。
anns: key为实例的id，而value即为该实例。

for vid in self.dataset['videos']:vids[vid['id']] = vid

vids: key为视频id，而value即为该视频信息。

for cat in self.dataset['categories']:cats[cat['id']] = cat

cats: key为类别id，而value即为该类别信息。

for ann in self.dataset['annotations']:catToVids[ann['category_id']].append(ann['video_id'])

catToVids: key是类别id，而value则为一个列表，里面存储了类别id为key的所有视频id。

3. 一些常用函数的使用

3.1 getAnnIds

def getAnnIds(self, vidIds=[], catIds=[], areaRng=[], iscrowd=None):"""Get ann ids that satisfy given filter conditions. default skips that filter:param vidIds  (int array)     : get anns for given vidscatIds  (int array)     : get anns for given catsareaRng (float array)   : get anns for given area range (e.g. [0 inf])iscrowd (boolean)       : get anns for given crowd label (False or True):return: ids (int array)       : integer array of ann ids"""

这个函数的作用在于在给定视频id列表、类别id列表、面积约束范围以及是否包括crowd实例的情况下，输出对应满足条件的ann id。
其大致逻辑是：

1. 根据给定的视频id，通过vidToAnns来获取所有的实例信息。
2. 根据类别id、面积约束范围以及是否包含crowd实例来对得到的实例信息进行筛选。
3. 最终返回所有实例对应的id。

3.2 getCatIds

def getCatIds(self, catNms=[], supNms=[], catIds=[]):"""filtering parameters. default skips that filter.:param catNms (str array)  : get cats for given cat names:param supNms (str array)  : get cats for given supercategory names:param catIds (int array)  : get cats for given cat ids:return: ids (int array)   : integer array of cat ids"""

这个函数的作用在于在给定类别名称列表、类超类类别名称列表以及类别id列表的情况下，输出对应满足条件的类别id。

3.3 getVidIds

def getVidIds(self, vidIds=[], catIds=[]):'''Get vid ids that satisfy given filter conditions.:param vidIds (int array) : get vids for given ids:param catIds (int array) : get vids with all given cats:return: ids (int array)  : integer array of vid ids'''

这个函数的作用在于在给定视频id列表以及类别id列表的情况下，输出对应满足条件的类别id。

3.4 loadAnns, loadCats and loadVids

这些函数的作用都相近，都是根据提供对应的id列表来获取对应的信息，包括标注、类别以及视频等。

3.5 annToRLE以及annToMask

这两个函数一般使用annToMask，用于将ann转化为对应的binary mask，而annToRLE则是用于对ann进行转化。

def annToRLE(self, ann, frameId):"""Convert annotation which can be polygons, uncompressed RLE to RLE.:return: binary mask (numpy 2D array)"""t = self.vids[ann['video_id']]h, w = t['height'], t['width']segm = ann['segmentations'][frameId]if type(segm) == list:# polygon -- a single object might consist of multiple parts# we merge all parts into one mask rle coderles = maskUtils.frPyObjects(segm, h, w)rle = maskUtils.merge(rles)elif type(segm['counts']) == list:# uncompressed RLErle = maskUtils.frPyObjects(segm, h, w)else:# rlerle = segmreturn rledef annToMask(self, ann, frameId):"""Convert annotation which can be polygons, uncompressed RLE, or RLE to binary mask.:return: binary mask (numpy 2D array)"""rle = self.annToRLE(ann, frameId)m = maskUtils.decode(rle)return m

4. 解析代码

这里提供一份用于提取某个特定类别数据的代码，例如以下代码可以提取Person这个类别：

import cv2
import os.path
import numpy as np
from tqdm import tqdm# https://github.com/qjy981010/cocoapi/blob/main/PythonAPI/pycocotools/ovis.py
from ovis import OVIS
# a tool function for parallel work
from parallel_work import parallel_workPALETTE = [(220, 20, 60), (119, 11, 32), (0, 0, 142), (0, 0, 230), (106, 0, 228),(0, 60, 100), (0, 80, 100), (0, 0, 70), (0, 0, 192), (250, 170, 30),(100, 170, 30), (220, 220, 0), (175, 116, 175), (250, 0, 30),(165, 42, 42), (255, 77, 255), (0, 226, 252), (182, 182, 255),(0, 82, 0), (120, 166, 157), (110, 76, 0), (174, 57, 255),(199, 100, 0), (72, 0, 118), (255, 179, 240), (0, 125, 92),(209, 0, 151), (188, 208, 182), (0, 220, 176), (255, 99, 164),(92, 0, 73), (133, 129, 255), (78, 180, 255), (0, 228, 0),(174, 255, 243), (45, 89, 255), (134, 134, 103), (145, 148, 174),(255, 208, 186), (197, 226, 255), (171, 134, 1), (109, 63, 54),(207, 138, 255), (151, 0, 95), (9, 80, 61), (84, 105, 51),(74, 65, 105), (166, 196, 102), (208, 195, 210), (255, 109, 65),(0, 143, 149), (179, 0, 194), (209, 99, 106), (5, 121, 0),(227, 255, 205), (147, 186, 208), (153, 69, 1), (3, 95, 161),(163, 255, 0), (119, 0, 170), (0, 182, 199), (0, 165, 120),(183, 130, 88), (95, 32, 0), (130, 114, 135), (110, 129, 133),(166, 74, 118), (219, 142, 185), (79, 210, 114), (178, 90, 62),(65, 70, 15), (127, 167, 115), (59, 105, 106), (142, 108, 45),(196, 172, 0), (95, 54, 80), (128, 76, 255), (201, 57, 1),(246, 0, 122), (191, 162, 208)]def load_mask_from_person_video_ids(person_video_ids, ovis, person_cat_id, target_dir):for video_id in person_video_ids:video_info = ovis.vids[video_id]height, width, length = video_info['height'], video_info['width'], video_info['length']video_anns = ovis.vidToAnns[video_id]person_anns = [anns for anns in video_anns if anns['category_id'] == person_cat_id]for frame_id in tqdm(range(length)):frame_mask = np.zeros((height, width), dtype=np.uint8)for ann in person_anns:if ann['segmentations'][frame_id] is None:continuemask = ovis.annToMask(ann, frame_id)mask_id = max(np.unique(frame_mask)) + 1frame_mask = np.where(mask > 0, mask_id, frame_mask)if len(np.unique(frame_mask)) > 0:file_path = video_info['file_names'][frame_id]video_name, frame_name = file_path.split('/')os.makedirs(os.path.join(target_dir, video_name), exist_ok=True)cv2.imwrite(os.path.join(target_dir, video_name, frame_name.replace('jpg', 'png')), frame_mask)def load_mask_from_ovis(ovis_root):subsets = ['train']for subset in subsets:target_dir = os.path.join(ovis_root, 'labels', subset)os.makedirs(target_dir, exist_ok=True)json_file = os.path.join(ovis_root, f'annotations_{subset}.json')ovis = OVIS(json_file)# get person cat idperson_cat_id = ovis.getCatIds(catNms=['Person'])[0]print(person_cat_id)# get person video idsperson_video_ids = ovis.getVidIds(catIds=person_cat_id)print(person_video_ids)parallel_work(person_video_ids, load_mask_from_person_video_ids, ovis, person_cat_id, target_dir)def visualization_label_files(label_files, target_dir):for root, file in label_files:video_name = root.split('/')[-1]os.makedirs(os.path.join(target_dir, video_name), exist_ok=True)raw_mask = cv2.imread(os.path.join(root, file), 0)h, w = raw_mask.shape[:2]color_mask = np.zeros((h, w, 3))instance_ids = np.unique(raw_mask)[1:]for instance_id in instance_ids:ci = int(instance_id) % len(PALETTE)color_mask[:, :, 0][raw_mask == instance_id] = PALETTE[ci][0]color_mask[:, :, 1][raw_mask == instance_id] = PALETTE[ci][1]color_mask[:, :, 2][raw_mask == instance_id] = PALETTE[ci][2]cv2.imwrite(os.path.join(target_dir, video_name, file), color_mask)def visualization(ovis_root):subsets = ['train']for subset in subsets:label_dir = os.path.join(ovis_root, 'labels', subset)target_dir = os.path.join(ovis_root, 'visualizations', subset)label_files = []for root, dirs, files in os.walk(label_dir):for file in tqdm(files, total=len(files)):if file.endswith('png'):label_files.append((root, file))parallel_work(label_files, visualization_label_files, target_dir)if __name__ == '__main__':ovis_root = '/data/data/segmenation/instance_segmentation/OVIS'load_mask_from_ovis(ovis_root)visualization(ovis_root)

parallel_work是一个工具函数：

import numpy as np
import multiprocessingdef parallel_work(datas, worker, *args):n_cpu = multiprocessing.cpu_count()print('n_cpu', n_cpu)num_data = len(datas)stride = int(np.ceil(num_data / n_cpu))processes = []for i in range(n_cpu):end = min((i + 1) * stride, num_data)worker_data = datas[i * stride:end]process = multiprocessing.Process(target=worker,args=(worker_data, *args),)processes.append(process)process.start()for process in processes:process.join()