waspinator / pycococreator

Helper functions to create COCO datasets
Apache License 2.0
775 stars 179 forks source link

IndexError: list index out of range #36

Open anmspro opened 4 years ago

anmspro commented 4 years ago

Full error:

Traceback (most recent call last):
  File "segment.py", line 108, in <module>
    class_id = [x['id'] for x in CATEGORIES if x['name'] in annotation_filename][0]
IndexError: list index out of range

Here's the code:

INFO = {
    "description": "Fashion Dataset",
    "url": "https://github.com/waspinator/pycococreator",
    "version": "0.1.0",
    "year": 2020,
    "contributor": "Abu Noman Md. Sakib",
    "date_created": datetime.datetime.utcnow().isoformat(' ')
}

LICENSES = [
    {
        "id": 1,
        "name": "GB",
        "url": "GB"
    }
]

CATEGORIES = [
    {
        'id': 1,
        'name': 'mask',
        'supercategory': 'fashion',
    }
]

coco_output = {
    "info": INFO,
    "licenses": LICENSES,
    "categories": CATEGORIES,
    "images": [],
    "annotations": []
}

image_id = 1
segmentation_id = 1

ROOT_DIR = "train"
IMAGE_DIR = "train/images"
ANNOTATION_DIR = "train/annotations"

image_files = [f for f in listdir(IMAGE_DIR) if isfile(join(IMAGE_DIR, f))]
annotation_files = [f for f in listdir(ANNOTATION_DIR) if isfile(join(ANNOTATION_DIR, f))]

# go through each image
for image_filename in image_files:
    image = Image.open(IMAGE_DIR + '/' + image_filename)
    image_info = pycococreatortools.create_image_info(image_id, os.path.basename(image_filename), image.size)

    coco_output["images"].append(image_info)

# go through each associated annotation
for annotation_filename in annotation_files:

    print(annotation_filename)
    class_id = [x['id'] for x in CATEGORIES if x['name'] in annotation_filename][0]
    category_info = {'id': class_id, 'is_crowd': 'crowd' in image_filename}
    binary_mask = np.asarray(Image.open(annotation_filename).convert('1')).astype(np.uint8)

    annotation_info = pycococreatortools.create_annotation_info(segmentation_id, image_id, category_info, binary_mask, image.size, tolerance=2)

    if annotation_info is not None:
        coco_output["annotations"].append(annotation_info)

    segmentation_id = segmentation_id + 1

    image_id = image_id + 1

with open('train/images.json', 'w') as output_json_file:
    json.dump(coco_output, output_json_file)

I still don't know if the code works or not. Please help!

kp421 commented 4 years ago

Did you manage to solve this?

kp421 commented 4 years ago

I just managed to solve it - the name for one of the categories was misspelled.