[python/coco] 원하는 클래스만 필터링하여 사용하기(Filtering sub classes on COCO Json file)

COCO 형태의 train, val, test JSON파일 안에는 여러 클래스가 있다.

이 중 내가 특정 경우에 일부 클래스만 학습하고자 할 때 해당 Json파일을 수정할 수 있도록 해주는 오픈소스가 있어 공유한다.

 

https://github.com/immersive-limit/coco-manager

GitHub - immersive-limit/coco-manager: Scripts to manage COCO datasets

 

간단하게 원하는 input file, output file, 원하는 카테고리들을 나열해주면 해당 카테고리를 제외한 나머지 Annotation을 삭제해준다.

 

근데, 내 경우엔 이 제외 항목들을 다 한 클래스로 덮어씌우고 싶어서 아래와 같은 코드를 _filter_annotations()함수에 추가해서 사용해줬다.

 

    def _filter_annotations(self):
        """ Create new collection of annotations matching category ids
            Keep track of image ids matching annotations
        """
        self.new_segmentations = []
        self.new_image_ids = set()
        for image_id, segmentation_list in self.segmentations.items():
            for segmentation in segmentation_list:
                original_seg_cat = segmentation['category_id']
                if original_seg_cat in self.new_category_map.keys():
                    new_segmentation = dict(segmentation)
                    new_segmentation['category_id'] = self.new_category_map[original_seg_cat]
                    self.new_segmentations.append(new_segmentation)
                    self.new_image_ids.add(image_id)
                ##### 하단이 추가된 코드 ####
                else:# 해당하지 않는 클래스 인 경우엔 Text로 변경
                    new_segmentation = dict(segmentation)
                    new_segmentation['category_id'] = self.new_category_map[10] # 10이 Text
                    self.new_segmentations.append(new_segmentation)
                    self.new_image_ids.add(image_id)

 

결론은 너무 편한 오픈소스.