Spaces:
No application file
No application file
File size: 2,778 Bytes
907462b | 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | from pathlib import Path
import pickle
import random
import numpy as np
import json
import cv2
from copy import deepcopy
from isegm.data.base import ISDataset
from isegm.data.sample import DSample
class CocoLvisDataset(ISDataset):
def __init__(self, dataset_path, split='train', stuff_prob=0.0,
allow_list_name=None, anno_file='hannotation.pickle', **kwargs):
super(CocoLvisDataset, self).__init__(**kwargs)
dataset_path = Path(dataset_path)
self._split_path = dataset_path / split
self.split = split
self._images_path = self._split_path / 'images'
self._masks_path = self._split_path / 'masks'
self.stuff_prob = stuff_prob
with open(self._split_path / anno_file, 'rb') as f:
self.dataset_samples = sorted(pickle.load(f).items())
if allow_list_name is not None:
allow_list_path = self._split_path / allow_list_name
with open(allow_list_path, 'r') as f:
allow_images_ids = json.load(f)
allow_images_ids = set(allow_images_ids)
self.dataset_samples = [sample for sample in self.dataset_samples
if sample[0] in allow_images_ids]
def get_sample(self, index) -> DSample:
image_id, sample = self.dataset_samples[index]
image_path = self._images_path / f'{image_id}.jpg'
image = cv2.imread(str(image_path))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
packed_masks_path = self._masks_path / f'{image_id}.pickle'
with open(packed_masks_path, 'rb') as f:
encoded_layers, objs_mapping = pickle.load(f)
layers = [cv2.imdecode(x, cv2.IMREAD_UNCHANGED) for x in encoded_layers]
layers = np.stack(layers, axis=2)
instances_info = deepcopy(sample['hierarchy'])
for inst_id, inst_info in list(instances_info.items()):
if inst_info is None:
inst_info = {'children': [], 'parent': None, 'node_level': 0}
instances_info[inst_id] = inst_info
inst_info['mapping'] = objs_mapping[inst_id]
if self.stuff_prob > 0 and random.random() < self.stuff_prob:
for inst_id in range(sample['num_instance_masks'], len(objs_mapping)):
instances_info[inst_id] = {
'mapping': objs_mapping[inst_id],
'parent': None,
'children': []
}
else:
for inst_id in range(sample['num_instance_masks'], len(objs_mapping)):
layer_indx, mask_id = objs_mapping[inst_id]
layers[:, :, layer_indx][layers[:, :, layer_indx] == mask_id] = 0
return DSample(image, layers, objects=instances_info)
|