Vision-AI) coco dataset을 mask로 변환하기_coco2mask

이전 포스팅에서는 mask를 coco dataset으로 변환했습니다.

https://mightytedkim.tistory.com/215

Vision-AI) mask 이미지 coco datset으로 변환하기_mask2coco

 

이번에는 반대로 mask를 coco로 변환해보겠습니다.

 

이제 확인해봐야겠죠.

이전 포스팅에서 mask와 coco를 시가화하는 글을 썼는데 이걸 이용해볼게요

1. https://mightytedkim.tistory.com/212

Vision-AI) coco dataset 시각화 하기 (코드 포함)

 

2. https://mightytedkim.tistory.com/213

Vision-AI) Mask 정리해서 시각화하기(one-hot, show_channel)

 

2개가 비슷한 것을 확인할 수 있습니다 :)

coco 시각화 + mask 시각화

 

코드는 아래와 같습니다

class Preprocess(object):
    def __init__(self, coco_path, output_mask_path):
        print("Initializing PreProcess")
        self.annotation_file_path = os.path.join(coco_path, "annotations", [f for f in os.listdir(os.path.join(coco_path, "annotations")) if f.endswith('.json')][0])
        self.images_dir = os.path.join(coco_path, "images")
        self.masks_output_dir = os.path.join(output_mask_path)
        print(self.annotation_file_path)
        print(self.images_dir)
        print(self.masks_output_dir)

        self.annotations = self.load_coco_annotations()
        self.categories = {category['id']: category for category in self.annotations['categories']}
        self.category_names = list({category['name']: category for category in self.annotations['categories']}.keys())
        
    def load_coco_annotations(self):
        with open(self.annotation_file_path, 'r') as file:
            return json.load(file)

    def installPip(self):
        # This method should be run outside of the class in your environment setup
        pass

    def getCatNames(self):
        return self.category_names

    def convertCocoToMasked(self):
        import matplotlib.pyplot as plt
        from collections import Counter

        if os.path.isdir(self.masks_output_dir):
            print('delete mask since exists')
            shutil.rmtree(self.masks_output_dir)
        os.makedirs(self.masks_output_dir, exist_ok=True)

        with open(self.annotation_file_path, 'r') as f:
            data = json.load(f)

        # Assign unique grayscale values for each category
        category_ids = {cat['id']: idx + 1 for idx, cat in enumerate(data['categories'])}  # Start at 1 to preserve 0 for background
        coco_category_names = {cat['name']: idx + 1 for idx, cat in enumerate(data['categories'])}  # Start at 1 to preserve 0 for background

        for image_info in tqdm(data['images'], desc="convertCocoToMasked"):
            image_id = image_info['id']
            image_file_name = os.path.splitext(image_info['file_name'])[0]
            image_path = os.path.join(self.images_dir, image_info['file_name'])
            mask_path = os.path.join(self.masks_output_dir, image_file_name + '.png')

            image = Image.open(image_path)
            mask = Image.new('L', image.size, 0)
            draw = ImageDraw.Draw(mask)

            # mask_array = np.zeros(image.size[::-1], dtype=np.uint8)  # Create an array for the mask
            for annotation in data['annotations']:
                if annotation['image_id'] == image_id and 'segmentation' in annotation:
                    category_value = category_ids[annotation['category_id']]
                    for segmentation in annotation['segmentation']:
                        if isinstance(segmentation, list):
                            polygon = [(segmentation[i], segmentation[i + 1]) for i in range(0, len(segmentation), 2)]
                            draw.polygon(polygon, fill=category_value,outline=None)

            # mask = Image.fromarray(mask_array)
            mask.save(mask_path, 'PNG')

 

mask와 coco를 함께 사용하는 분들께 도움이 되었으면 좋겠습니다.