От коробочек до масок

Question

У меня есть csv dataframe:

    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

Я хотел бы получить маски для каждого изображения.Мне удалось получить их, если для каждого изображения есть только один блок, однако на некоторых изображениях есть несколько блоков (например, 10.jpg).Как я могу добавить эту ограничивающую рамку к маске?

Пока мой код выглядит следующим образом (хорошо работает, если изображение имеет 1 строку):

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)

Спасибо!

Answer 1

На самом деле, это неправильно:

Мне удалось получить их, если для каждого изображения есть только одна коробка

Ваш код работает только для первой строки, потому что вы запрашиваете индекс 0. Все остальные строки не работают, потому что кадры данных запоминают свой исходный индекс.

В этом случае groupby добивается цели.

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)

Здесь groupby позволяет перебирать серию подкадров с одинаковыми 'filename'.
Затем во вложенном цикле iterrows используется для итерации по каждой строке каждого подкадра с целью извлечения значения и построения маски.
Как видите, маска строится на каждой итерации внешнего цикла, оставляя внутренний цикл для «рисования» различных прямоугольников маски, по одному прямоугольнику для каждой строки подкадра.

EDIT

Аналогичное, но немного более быстрое решение для внутреннего цикла вместо iterrows:

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

Если у вас большое количество строк, может быть полезно.