Segmentation 문제

아래 그림에서 Segmentation 모델을 활용해 검출된 모든 클래스를 출력하고, 배경을 제외한 각 클래스마다 이미지 내에서 차지하는 비율을 출력

• 모델은 반드시 torchvision에서 제공하는 fcn_resnet50 Pretrained Model을 사용할 것
• 비율은 %로 표시

model = torch.hub.load('pytorch/vision:v0.10.0', 'fcn_resnet50', pretrained=True)
model.eval()

# 코드 작성
seg_image_url = "http://host.robots.ox.ac.uk/pascal/VOC/voc2007/segexamples/images/006585.jpg"
res = request.urlopen(seg_image_url).read()
img = Image.open(io.BytesIO(res))

preprocess = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

img = img.convert("RGB")

origin_image_array = np.array(img)
input_img_tensor = preprocess(img)
input_img_batch = input_img_tensor.unsqueeze(0)
input_img_batch = input_img_batch.to('cuda')
model.to('cuda')

with torch.no_grad():
  output = model(input_img_batch)

class_feature = output['out'].squeeze(0).argmax(0)
print('classes:', torch.unique(class_feature))

pixels = output['out'].squeeze(0).shape[1]*output['out'].squeeze(0).shape[2]
class_array = class_feature.byte().cpu().numpy()

print('배경', (np.where(class_array == 0)[0].shape[0] / pixels * 100), '%')
print('오토바이', (np.where(class_array == 14)[0].shape[0] / pixels * 100), '%')
print('사람', (np.where(class_array == 15)[0].shape[0] / pixels * 100), '%')

plt.imshow(img)
plt.show()

#결과:
classes: tensor([ 0, 14, 15], device='cuda:0')
배경 87.91117824773413 %
오토바이 8.636858006042296 %
사람 3.4519637462235653 %