От боксов до масок
У меня есть фрейм данных csv следующим образом:
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)
Спасибо !
Ответы 1
На самом деле это неправильно:
I've succeded to get them if there is only one box for each image
Ваш код работает только для первой строки, потому что вы запрашиваете индекс 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 используется для перебора каждой строки каждого подкадра данных, чтобы извлечь значение и построить маску.
Как видите, маска строится на каждой итерации внешнего цикла, оставляя внутреннему циклу «рисовать» разные прямоугольники маски, по одному прямоугольнику для каждой строки подкадра данных.
РЕДАКТИРОВАТЬ
Аналогичное, но немного более быстрое решение для внутреннего цикла вместо iterrows:
for x1, y1, x2, y2 in zip(subdf['xmin'], subdf['ymin'], subdf['xmax'], subdf['ymax']):
mask[y1:y2, x1:x2] = 1
Если у вас есть большое количество строк, может быть полезно.