If my image contains multiple people, how can I get the mask of a single person? #109
Description
Now I can get a mulitple masks, can I get a single person mask?
this is my code:
for i, image_name in tqdm(enumerate(image_names), total=len(image_names)):
print(image_name)
image_path = os.path.join(input_dir, image_name)
result = inference_model(model, image_path)
print("predict picture")
output_file = os.path.join(output_root, os.path.basename(image_path).replace('.jpg', '.png').replace('.jpeg', '.png').replace('.png', '.npy'))
output_seg_file = os.path.join(output_root, os.path.basename(image_path).replace('.jpg', '.png').replace('.jpeg', '.png').replace('.png', '_seg.npy'))
image = cv2.imread(image_path)
pred_sem_seg = result.pred_sem_seg.data[0].cpu().numpy() # H x W. seg ids.
mask = (pred_sem_seg > 0)
random_number = random.randint(1000, 9999)
mask_image_path = os.path.join(output_root, f"{random_number}_mask.png")
cv2.imwrite(mask_image_path, (mask * 255).astype(np.uint8))
np.save(output_file, mask)
np.save(output_seg_file, pred_sem_seg)
