text_plain_output_1.png
from joblib import Parallel, delayed

from tqdm.notebook import tqdm

import json

import numpy as np

import os

import pandas as pd

thingClasses = ['Aortic enlargement', 'Atelectasis', 'Calcification', 'Cardiomegaly', 'Consolidation', 'ILD', 'Infiltration', 'Lung Opacity', 'Nodule/Mass', 'Other lesion', 'Pleural effusion', 'Pleural thickening', 'Pneumothorax', 'Pulmonary fibrosis', 'No finding']
cfgDict = {'dicomPath': None, 'orgDataPath': '../input/sartorius-cell-instance-segmentation/', 'newDataPath': None, 'cachePath': './', 'trainDataName': 'vinbigdataTrain', 'validDataName': 'vinbigdataValid', 'sampleSize': 1000, 'imSize': 256, 'modelName': 'COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml', 'debug': False, 'outdir': './results/', 'logFile': 'log.txt', 'splitMode': True, 'seed': 111, 'device': 'cuda', 'iter': 1000, 'ims_per_batch': 16, 'roi_batch_size_per_image': 512, 'eval_period': 20, 'lr_scheduler_name': 'WarmupCosineLR', 'base_lr': 0.001, 'checkpoint_period': 500, 'num_workers': 4, 'score_thresh_test': 0.05, 'augKwargs': {'RandomFlip': {'prob': 0.5}, 'RandomRotation': {'angle': [0, 360]}}}

def rle2mask(rle, h, w):
    rleArray = np.fromiter(rle.split(), dtype=np.uint)
    rleArray = rleArray.reshape((-1, 2)).T
    rleArray[0] = rleArray[0] - 1
    starts, lenghts = rleArray
    rleArray = np.concatenate([np.arange(s, s + l, dtype=np.uint) for s, l in zip(starts, lenghts)])
    mask = np.zeros(h * w, dtype=np.uint8)
    mask[rleArray] = 1
    mask = mask.reshape((h, w))
    mask = np.asfortranarray(mask)
    return mask

def mask2annotation(idx, row, catIds):
    mask = rle2mask(row['annotation'], row['height'], row['width'])
    rle = maskUtils.encode(mask)
    rle['counts'] = rle['counts'].decode('utf-8')
    area = maskUtils.area(rle).item()
    bbox = maskUtils.toBbox(rle).astype(int).tolist()
    annotation = {'segmentation': rle, 'bbox': bbox, 'area': area, 'image_id': row['id'], 'category_id': catIds[row['cell_type']], 'iscrowd': 0, 'id': idx}
    return annotation

def df2COCO(cfg, df, workers=4):
    catIds = {name: id + 1 for id, name in enumerate(df.cell_type.unique())}
    cats = [{'name': name, 'id': id} for name, id in catIds.items()]
    images = [{'id': id, 'width': row.width, 'height': row.height, 'file_name': f'train/{id}.png'} for id, row in df.groupby('id').agg('first').iterrows()]
    annotations = Parallel(n_jobs=workers)((delayed(mask2annotation)(idx, row, catIds) for idx, row in tqdm(df.iterrows(), total=len(df))))
    return {'categories': cats, 'images': images, 'annotations': annotations}

df = pd.read_csv(os.path.join(cfgDict['orgDataPath'], 'train.csv'))
df.head()
root = df2COCO(cfgDict, df[:cfgDict['sampleSize']])
with open('annotations_train.json', 'w', encoding='utf-8') as f:
    json.dump(root, f, ensure_ascii=True, indent=4)