Computer_Vision/convert-to-coco-format-checkpoint.ipynb at 5c00e9ffe705296c92963cde33978d49d66290f1

from torch_snippets import *

from sklearn.model_selection import train_test_split
df = pd.read_csv('df.csv')
trn_df, val_df = train_test_split(df, random_state=10)

df_mini = df[df.ImageID.isin(df.ImageID.unique()[:500].tolist())]
trn_df_mini, val_df_mini = train_test_split(df_mini, random_state=10)

len(df)
df.head()

	ImageID	Source	LabelName	Confidence	XMin	XMax	YMin	YMax	IsOccluded	IsTruncated	...	XClick1X	XClick2X	XClick3X	XClick4X	XClick1Y	XClick2Y	XClick3Y	XClick4Y
0	0000599864fd15b3	xclick	Bus	1	0.343750	0.908750	0.156162	0.650047	1	0	...	0.421875	0.343750	0.795000	0.908750	0.156162	0.512700	0.650047	0.457197
1	00006bdb1eb5cd74	xclick	Truck	1	0.276667	0.697500	0.141604	0.437343	1	0	...	0.299167	0.276667	0.697500	0.659167	0.141604	0.241855	0.352130	0.437343
2	00006bdb1eb5cd74	xclick	Truck	1	0.702500	0.999167	0.204261	0.409774	1	1	...	0.849167	0.702500	0.906667	0.999167	0.204261	0.398496	0.409774	0.295739
3	00010bf498b64bab	xclick	Bus	1	0.156250	0.371250	0.269188	0.705228	0	0	...	0.274375	0.371250	0.311875	0.156250	0.269188	0.493882	0.705228	0.521691
4	00013f14dd4e168f	xclick	Bus	1	0.287500	0.999375	0.194184	0.999062	0	1	...	0.920000	0.999375	0.648750	0.287500	0.194184	0.303940	0.999062	0.523452

5 rows × 21 columns

val_df_mini.LabelName.unique()

array(['Bus', 'Truck'], dtype=object)

categories = [{'id': 1, 'name': 'Bus', 'supercategory': 'none'}, {'id': 2, 'name': 'Truck', 'supercategory': 'none'}]
category_ids = {'Bus': 1, 'Truck': 2}

def get_image_infos(df):
    image_infos = []
    for image in df['ImageID'].unique():
        info = {}
        info['file_name'] = image+'.jpg'
        im = read(f'images/{image}.jpg')
        h, w = im.shape
        info['height'], info['width'] = h, w
        info['id'] = len(image_infos)+1
        image_infos.append(info)
    return image_infos
def get_annotations_for_image_info(image_info, df):
    imginfo = image_info['file_name'].split('.')[0]
    image_id = image_info['id']
    h, w = image_info['height'], image_info['width']
    _df = df[df['ImageID'] == imginfo]
    annotations = []
    for ix, row in _df.iterrows():
        annot = {}
        row = row.squeeze()
        x,y,X,Y = row.XMin,row.YMin,row.XMax,row.YMax
        x,y,X,Y = x*w,y*h,X*w,Y*h
        x,y,X,Y = [int(i) for i in [x,y,X,Y]]
        annot['bbox'] = [x,y,X-x,Y-y]
        annot['ignore'] = '0'
        annot['category_id'] = category_ids[row.LabelName]
        annot['area'] = (X-x)*(Y-y)
        annot['iscrowd'] = 0
        annot['segmentation'] = [[x,y,x,Y,X,Y,X,y]]
        annot['image_id'] = image_id
        annotations.append(annot)
    return annotations
def get_annotations_for_all_image_infos(image_infos, df):
    ANNOTATIONS = []
    for image_info in image_infos:
        annotations = get_annotations_for_image_info(image_info, df)
        for annot in annotations:
            annot['id'] = len(ANNOTATIONS) + 1
            ANNOTATIONS.append(annot)
    return ANNOTATIONS
def get_coco_annotations(df):
    image_infos = get_image_infos(df)
    annotations = get_annotations_for_all_image_infos(image_infos, df)
    data_in_coco_format = {}
    data_in_coco_format['annotations'] = annotations
    data_in_coco_format['categories'] = categories
    data_in_coco_format['images'] = image_infos
    data_in_coco_format['type'] = 'instances'
    return data_in_coco_format

import json
!mkdir annotations
json.dump(get_coco_annotations(trn_df), open('annotations/open_images_train_coco_format.json','w'),indent=4)
json.dump(get_coco_annotations(val_df), open('annotations/open_images_val_coco_format.json','w'),indent=4)

Load and test

category_ids = {'Bus': 0, 'Truck': 1}
def get_dump(df):
    data = []
    for imageID in df.ImageID.unique():
        filename = imageID+'.jpg'
        width, height = read(f'images/{filename}').shape
        _df = df[df.ImageID==imageID]
        bbs = _df['XMin,YMin,XMax,YMax'.split(',')].values
        bbs = bbs*np.array([width,height,width,height])
        # bbs[:,2:] = bbs[:,2:] - bbs[:,:2]
        labels = np.array([category_ids[ix] for ix in _df['LabelName']])
        if len(labels) == 0: continue
        annot = dict(
            filename=filename,
            width=width, height=height,
            ann=dict(
                bboxes=bbs.astype(np.float32),
                labels=labels.astype(np.int64)
            )
        )
        data.append(annot)
    return data

import pickle
pickle.dump(get_dump(trn_df_mini), open('pickles/2-mini-open-images-bus-truck-train.pickle', 'wb'))
pickle.dump(get_dump(val_df_mini), open('pickles/2-mini-open-images-bus-truck-val.pickle', 'wb'))

# pickle.dump(get_dump(trn_df_mini), open('mini-open-images-bus-truck-train.pickle', 'wb'))
# pickle.dump(get_dump(val_df_mini), open('mini-open-images-bus-truck-val.pickle', 'wb'))

from mmcv import load
data = load('open-images-bus-truck-train.pickle')
data[:3]

[{'filename': '02b4c34bd52f1217.jpg',
  'width': 194,
  'height': 256,
  'ann': {'bboxes': array([[  0.      , 126.014465, 159.44376 , 105.36499 ]], dtype=float32),
   'labels': array([1])}},
 {'filename': '9b92c6f48307ece3.jpg',
  'width': 194,
  'height': 256,
  'ann': {'bboxes': array([[ 19.885   ,  25.621248, 154.9575  , 209.41594 ],
          [  1.455   ,  55.47725 ,  99.78875 , 139.32826 ]], dtype=float32),
   'labels': array([1, 1])}},
 {'filename': '8ee9ac71d758e982.jpg',
  'width': 144,
  'height': 256,
  'ann': {'bboxes': array([[  8.73    ,  17.92    ,  66.24    , 200.53325 ],
          [ 65.34    ,  15.928832,  78.03    , 200.81792 ]], dtype=float32),
   'labels': array([1, 1])}}]

16 KiB Raw Blame History Unescape Escape

Load and test

16 KiB

Raw Blame History