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
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 ,and
这些函数的作用都相近 ， 都是根据提供对应的id列表来获取对应的信息 ， 包括标注、类别以及视频等 。
3.5 以及
这两个函数一般使用 ， 用于将ann转化为对应的 mask ， 而则是用于对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. 解析代码
这里提供一份用于提取某个特定类别数据的代码 ， 例如以下代码可以提取这个类别：
import cv2import os.pathimport numpy as npfrom tqdm import tqdm# https://github.com/qjy981010/cocoapi/blob/main/PythonAPI/pycocotools/ovis.pyfrom ovis import OVIS# a tool function for parallel workfrom 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)