From bouding boxes to masks

Asked by At

I have a csv dataframe as follow:

    filename  width  height  class  xmin  ymin  xmax  ymax
0     1.jpg   2048    1251    1      706   513   743   562
1    10.jpg   1600     980    1      715   157   733   181
2    11.jpg   2828    1828    1      460  1530   482  1557
3    12.jpg   1276    1754    1      846   517   878   563
19   10.jpg   1600     980    1      428    83   483   145

I would like to get the masks for every image. I've succeded to get them if there is only one box for each image, however some images have multiple bouding boxes (example 10.jpg). How can I be able to add that bounding box to the mask ?

So far my code is as follow ( works good if image has 1 row ):

for idimage in annotations['filename']:
    img = cv2.imread('images/'+idimage)

    x1 = annotations[annotations['filename'] == idimage]['xmin'][0]
    y1 = annotations[annotations['filename'] == idimage]['ymin'][0]
    x2 = annotations[annotations['filename'] == idimage]['xmax'][0]
    y2 = annotations[annotations['filename'] == idimage]['ymax'][0]

    mask = np.zeros((img.shape[0],img.shape[1])).astype('uint8')
    mask[y1:y2, x1:x2] = 1

    mask = cv2.imwrite('mask/'+idimage,mask)

Thank you !

1 Answers

Valentino On Best Solutions

Actually, this is wrong:

I've succeded to get them if there is only one box for each image

Your code works only for the first row, because you request index 0. All the other rows fails because the dataframes remember their original index.

In this case, groupby does the trick.

for fn, subdf in annotations.groupby('filename'):
    img = cv2.imread('images/'+fn)
    mask = np.zeros((img.shape[0],img.shape[1])).astype('uint8')
    for _, row in subdf.iterrows():
        mask[row['ymin']:row['ymax'], row['xmin']:row['xmax']] = 1

    cv2.imwrite('mask/'+fn, mask)

Here groupby allows to iterate over a series of subdataframes with the same 'filename'.
Then in a nested loop iterrows is used to iterate over each row of each subdataframe in order to extract the value and build the mask.
As you see, the mask is buildt each iteration of the outer loop, leaving the inner loop to "paint" different rectangles of the mask, one rectangle for each row of the subdataframe.


A similar but slightly faster solution for the inner loop, instead of iterrows is:

for x1, y1, x2, y2 in zip(subdf['xmin'], subdf['ymin'], subdf['xmax'], subdf['ymax']):
    mask[y1:y2, x1:x2] = 1

If you have a large amount of rows may be useful.