File size: 1,236 Bytes
91e3dad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
import os
import numpy as np
from PIL import Image
import sys
class Results(object):
def __init__(self, root_dir):
self.root_dir = root_dir
def _read_mask(self, sequence, frame_id):
try:
mask_path = os.path.join(self.root_dir, sequence, f'{frame_id}.png')
return np.array(Image.open(mask_path))
except IOError as err:
sys.stdout.write(sequence + " frame %s not found!\n" % frame_id)
sys.stdout.write("The frames have to be indexed PNG files placed inside the corespondent sequence "
"folder.\nThe indexes have to match with the initial frame.\n")
sys.stderr.write("IOError: " + err.strerror + "\n")
sys.exit()
def read_masks(self, sequence, masks_id):
mask_0 = self._read_mask(sequence, masks_id[0])
masks = np.zeros((len(masks_id), *mask_0.shape))
for ii, m in enumerate(masks_id):
masks[ii, ...] = self._read_mask(sequence, m)
num_objects = int(np.max(masks))
tmp = np.ones((num_objects, *masks.shape))
tmp = tmp * np.arange(1, num_objects + 1)[:, None, None, None]
masks = (tmp == masks[None, ...]) > 0
return masks
|