La compatibilidad con el formato COCO es un requisito fundamental para entrenar la mayoría de las redes neuronales modernas de detección y segmentación. En este artículo, se detalla el proceso técnico para transformar un conjunto de datos de imágenes de segmentación con anotaciones en color RGB al formato JSON estándar de COCO.

1. Instalación de dependencias

Entes de comenzar, es necesario instalar las herramiantas oficiales de COCO y el paquete pycococreator. Asegúrese de tener cython instalado previamente en su entorno virtual.

# Para entornos Linux/macOS
pip install cython
pip install "git+https://github.com/waleedka/cocoapi.git#egg=pycocotools&subdirectory=PythonAPI"
pip install git+https://github.com/waspinator/coco.git@2.1.0

# Para entornos Windows
pip install cython
pip install "git+https://github.com/philferriere/cocoapi.git#egg=pycocotools&subdirectory=PythonAPI"
pip install git+https://github.com/waspinator/coco.git@2.1.0

2. Estructura de directorios

Organice los archivos del conjunto de datos siguiendo una estructura jerárquica clara. Esto facilitará el emparejamiento automático entre las imágenes originales y sus respectivas máscaras durante el procesamiento.

dataset_root/
│
├── labels/                  # Imágenes de segmentación RGB originales
│   ├── img_001.png
│   └── img_002.png
│
└── train/
    ├── annotations/         # Máscaras binarias generadas (salida del paso 3)
    │   ├── img_001_object_0.png
    │   └── img_002_object_0.png
    │
    └── images/              # Imágenes originales para el entrenamiento
        ├── img_001.png
        └── img_002.png

3. Conversión de anotaciones RGB a máscaras binarias

Las redes de segmentación requieren máscaras binarias (blanco y negro) para cada instancia de objeto. El sigueinte script utiliza numpy para identificar los colores únicos en la imagen RGB y generar una máscara individual por cada color, omitiendo el fondo negro. El formato de nombre de archivo resultante será <id_imagen>_<clase>_<id_instancia>.png.

import cv2
import numpy as np
from pathlib import Path

def extract_binary_masks(rgb_image_path, output_directory):
    image_id = Path(rgb_image_path).stem
    image_data = cv2.imread(str(rgb_image_path))
    
    # Aplanar la imagen a una lista de píxeles para encontrar colores únicos de forma eficiente
    pixels = image_data.reshape(-1, 3)
    unique_colors = np.unique(pixels, axis=0)
    
    instance_counter = 0
    white_color = np.array([255, 255, 255], dtype=np.uint8)
    
    for color in unique_colors:
        # Ignorar el color de fondo (asumiendo que es negro puro)
        if np.array_equal(color, [0, 0, 0]):
            continue
            
        # Generar máscara booleana para el color actual
        color_mask = np.all(image_data == color, axis=-1)
        
        # Crear imagen de máscara de 3 canales
        binary_mask = np.zeros_like(image_data)
        binary_mask[color_mask] = white_color
        
        # Guardar la máscara en disco
        mask_filename = f"{image_id}_object_{instance_counter}.png"
        save_path = Path(output_directory) / mask_filename
        cv2.imwrite(str(save_path), binary_mask)
        instance_counter += 1

# Ejecución del proceso por lotes
labels_dir = Path('./dataset_root/labels')
masks_dir = Path('./dataset_root/train/annotations')
masks_dir.mkdir(parents=True, exist_ok=True)

for label_file in labels_dir.glob('*.png'):
    extract_binary_masks(label_file, masks_dir)

4. Generación del archivo JSON de COCO

Una vez obtenidas las máscaras binarias, se utiliza pycococreatortools para compilar la estructura JSON final. Este script recorre el directorio de imágenes, busca sus máscaras asociadas y calcula los polígonos de segmentación y las cajas delimitadoras (bounding boxes) requeridas por la API de COCO.

import datetime
import json
from pathlib import Path
from PIL import Image
import numpy as np
from pycococreatortools import pycococreatortools

BASE_PATH = Path('/ruta/completa/dataset_root')
IMAGES_PATH = BASE_PATH / 'train' / 'images'
MASKS_PATH = BASE_PATH / 'train' / 'annotations'
OUTPUT_FILE = BASE_PATH / 'instances_train.json'

DATASET_METADATA = {
    "description": "Conjunto de datos de segmentación personalizado",
    "version": "1.0.0",
    "year": 2023,
    "date_created": datetime.datetime.utcnow().isoformat(' ')
}

CATEGORIES_CONFIG = [
    {"id": 1, "name": "objeto_objetivo", "supercategory": "entidad"}
]

def compile_coco_json():
    coco_structure = {
        "info": DATASET_METADATA,
        "licenses": [],
        "categories": CATEGORIES_CONFIG,
        "images": [],
        "annotations": []
    }
    
    current_img_id = 1
    current_ann_id = 1
    
    valid_extensions = {'.jpg', '.jpeg', '.png'}
    image_files = sorted([f for f in IMAGES_PATH.iterdir() if f.suffix.lower() in valid_extensions])
    
    for img_path in image_files:
        pil_image = Image.open(img_path)
        img_metadata = pycococreatortools.create_image_info(
            current_img_id, img_path.name, pil_image.size
        )
        coco_structure["images"].append(img_metadata)
        
        # Buscar máscaras correspondientes a la imagen actual mediante patrones
        mask_pattern = f"{img_path.stem}_*.png"
        associated_masks = list(MASKS_PATH.glob(mask_pattern))
        
        for mask_path in associated_masks:
            target_category_id = CATEGORIES_CONFIG[0]['id']
            category_metadata = {'id': target_category_id, 'is_crowd': 0}
            
            # Cargar máscara y convertir a binario estricto
            mask_array = np.asarray(Image.open(mask_path).convert('1')).astype(np.uint8)
            
            ann_metadata = pycococreatortools.create_annotation_info(
                current_ann_id, current_img_id, category_metadata, mask_array,
                pil_image.size, tolerance=2
            )
            
            if ann_metadata is not None:
                coco_structure["annotations"].append(ann_metadata)
                current_ann_id += 1
                
        current_img_id += 1
        
    with open(OUTPUT_FILE, 'w') as json_out:
        json.dump(coco_structure, json_out)

if __name__ == "__main__":
    compile_coco_json()