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

coco 이미지를 mask로 변환하는 법이 있다면, 그 반대도 있겠죠.

 

이번에는 mask 이미지를 이용해서 coco datset을 만들어보겠습니다.

실행하고 나면 파일이 생성된 것을 확인할 수 있습니다.

사용한 변수는 아래와 같습니다

여기서 주의할 점은 pixel 0의 경우 일반적으로 category_list에서 unlabeled 가 없는 값을 넣어줍니다.

만약 넣으면 json 파일이 기가 단위로 나올 수 있어요.

 

잘 생성되었는지 확인해봐야겠죠?

전에 포스팅했던 글의 class를 써보도록 하겠습니다.

잘 나오는 것을 확인했습니다 :)

https://mightytedkim.tistory.com/212

 

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

 

코드는 아래와 같습니다.

class COCODatasetCreator:
    def __init__(self, input_img_dir, input_mask_dir, output_mask_dir, cat_names):
        self.input_img_dir = input_img_dir
        self.cat_names = cat_names #self.read_category_names_from_coco_json(input_coco_json_path)
        self.input_mask_dir = input_mask_dir
        print(self.input_mask_dir)
        self.output_mask_dir = output_mask_dir
        self.coco_images_dir = os.path.join(output_mask_dir, 'images')
        self.annotations_dir = os.path.join(output_mask_dir, 'annotations')
        self.annotations_path = os.path.join(self.annotations_dir, 'annotations.json')
    def is_image_file(self, filename):
        valid_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"]
        return any(filename.lower().endswith(ext) for ext in valid_extensions)
        
    def read_category_names_from_coco_json(self, coco_json_path):
        """Read category names from a COCO JSON file."""
        with open(coco_json_path, 'r') as f:
            data = json.load(f)
        return [cat['name'] for cat in data['categories']]

    def create_sub_masks(self, mask_image, width, height):
        """Generate a sub-mask for each category ID in the image."""
        sub_masks = {}
        unique_ids = np.unique(mask_image)
        for category_id in unique_ids:
            if category_id == 0 or category_id == 255:  # Assuming 0 is the background, 255 may be an error or unused
                continue
            if category_id - 1 >= len(self.cat_names):
                print(f"Skipping category ID {category_id}, which is out of bounds.")
                continue
            sub_mask = np.zeros((height, width), dtype=np.uint8)
            sub_mask[mask_image == category_id] = 1
            #sub_masks[self.cat_names[category_id - 1]] = sub_mask
            sub_masks[self.cat_names[category_id ]] = sub_mask
        return sub_masks

    def create_annotations(self, sub_masks, image_id):
        """Create annotations from sub-masks."""
        annotations = []
        annotation_id = 1
        for category_name, sub_mask in sub_masks.items():
            contours, _ = cv2.findContours(sub_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
            for contour in contours:
                if len(contour) < 3:
                    continue
                segmentation = contour.flatten().tolist()
                if len(segmentation) < 6:
                    continue  # Need at least 3 points (x, y)
                rle = maskUtils.frPyObjects([segmentation], sub_mask.shape[0], sub_mask.shape[1])
                area = maskUtils.area(rle)
                bbox = maskUtils.toBbox(rle)[0].tolist()
                annotations.append({
                    'id': annotation_id,
                    'image_id': image_id,
                    # 'category_id': self.cat_names.index(category_name),
                    'category_id': self.cat_names.index(category_name) + 1,
                    'segmentation': [segmentation],
                    'area': int(area),
                    'bbox': bbox,
                    'iscrowd': 0
                })
                annotation_id += 1
        return annotations

    def create_coco_dataset(self):
        """Create COCO dataset by processing images and masks."""
        os.makedirs(self.coco_images_dir, exist_ok=True)
        os.makedirs(self.annotations_dir, exist_ok=True)

        coco_format = {
            "images": [],
            "annotations": [],
            "categories": [{"id": idx, "name": name, "supercategory": "none"} for idx, name in enumerate(self.cat_names)],
            "licenses": [{"id": 1, "name": "Unknown", "url": "unknown"}]
        }

        image_id = 1
        for filename in tqdm(os.listdir(self.input_img_dir), desc="create_coco_dataset"):
            if self.is_image_file(filename):
                original_image_path = os.path.join(self.input_img_dir, filename)
                target_image_path = os.path.join(self.coco_images_dir, filename)
                shutil.copy2(original_image_path, target_image_path)

                image = cv2.imread(original_image_path)
                height, width = image.shape[:2]

                coco_format['images'].append({
                    "id": image_id,
                    "width": width,
                    "height": height,
                    "file_name": filename,
                    "license": 1,
                    "flickr_url": "",
                    "coco_url": "",
                    "date_captured": ""
                })

                mask_path = os.path.join(self.input_mask_dir, os.path.splitext(filename)[0] + '.png')
                mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
                if mask is not None:
                    sub_masks = self.create_sub_masks(mask, width, height)
                    annotations = self.create_annotations(sub_masks, image_id)
                    coco_format['annotations'].extend(annotations)

                image_id += 1

        with open(self.annotations_path, 'w') as f:
            json.dump(coco_format, f, indent=4)
        print(f'self.annotations_path:{self.coco_images_dir}')
        print(f'self.annotations_path:{self.annotations_path}')

 

semantic segmentation과 coco를 함께 사용할 분들께 도움이 되었으면 좋겠습니다.