File size: 983 Bytes
4f6b78d |
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 33 34 |
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):
# mask_path = os.path.join(self.root_dir, sequence, f'{frame_id}.png')
mask_path = os.path.join(self.root_dir, sequence, f'{int(frame_id[-4:]) -1 :05d}.png')
obj = Image.open(mask_path)
obj_mode = obj.mode
if obj_mode == 'LA':
obj = np.array(obj)[:,:,0]
else:
obj = np.array(obj)
return obj
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 = 1
tmp = np.ones((num_objects, *masks.shape)) * 255
masks = (tmp == masks[None, ...]) > 0
return masks
|