Доброго времени суток! Я первый раз столкнулся с нейронными сетями и мне требуется ваша помощь.
У меня есть следующая задача: Найти на фотографии определенные объекты, после чего получить маску этих объектов на фото.
С какой проблемой я столкнулся: После обучения нейронной сети я получаю сегментацию и могу сохранить ее в файл. Но на многих картинках сегменты не захватывают весь объект. Иногда на 5-10 пикселей маска меньше или больше в некоторых местах. А так же имеет очень не ровные края.
Какой способ борьбы с этим я вижу: С недостаточным размером маски: просто увеличить маску. Но не понимаю как это сделать правильно. Просто увеличить маску до размера большего чем само изображение а потом просто кропнуть?
Сглаживание краев: Вот тут у меня вообще тупик. сегментация имеет множество точек (больше 20) и я не могу придумать алгоритм для сглаживания....
Мой код
from ultralytics import YOLO
import matplotlib.pyplot as plt
import cv2
import torch
model = YOLO('runs/segment/train3/weights/best.pt')
img = cv2.imread('test_img/manual_test6.jpg')
print(img.shape)
results = model(
img, imgsz=640,
iou=1, conf=0.8,
verbose=False,
retina_masks=True
)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
for result in results:
boxes = result.boxes.data
mask = torch.any(boxes, dim=0).int() * 255
annotated_frame = results[0].plot()
annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
plt.imshow(annotated_frame)
plt.show()
