0

I am trying to create black and white masks from images from COCO Dataset, so I can train my u2net model (salient object detection). However, I am facing this error:

  File "D:\downloads\U-2-Net-master\U-2-Net-master\mask.py", line 25, in <module>
    mask |= coco.annToMask(anns[i])
ValueError: operands could not be broadcast together with shapes (479,640) (375,500) (479,640)

Here my code:

from pycocotools.coco import COCO
import os
from matplotlib import image
from pathlib import Path
from PIL import Image
img_dir = "./coco/coco-tr/train2017"
annFile = "./coco/coco-tr/annotations/instances_train2017.json"
coco=COCO(annFile)
# Get category IDs and annotation IDs
catIds = coco.getCatIds(catNms=['dog'])
img_dir=os.path.join("train_data", "dogs", "images"+os.sep)

annsIds = coco.getAnnIds(catIds=catIds,iscrowd=False)

# Create folders named after annotation categories
for cat in catIds:

    Path(os.path.join("train_data","dogs","masks",coco.loadCats(cat)[0]['name'])).mkdir(parents=True, exist_ok=True)
    
    anns=coco.loadAnns(annsIds)
        
    mask = coco.annToMask(anns[0])
    for i in range(len(anns)):

        mask |= coco.annToMask(anns[i])

    mask = mask * 255
    mask = Image.fromarray(mask)
    file_path = os.path.join("train_data","dogs","masks",coco.loadCats(cat)[0]['name'],coco.loadImgs(anns[0]['image_id'])[0]['file_name'])
    image.imsave(file_path, mask, cmap="gray")

Thanks for any helpful advice

codingIsLove
  • 101
  • 1
  • 8

0 Answers0