Ejemplo n.º 1
0
    def test_background_masks_dont_introduce_instances_but_cover_others(self):
        dataset = Dataset.from_iterable([
            DatasetItem(1, image=np.zeros((4, 1, 1)), annotations=[
                Mask([1, 1, 1, 1], label=1, attributes={'z_order': 1}),
                Mask([0, 0, 1, 1], label=2, attributes={'z_order': 2}),
                Mask([0, 0, 1, 1], label=0, attributes={'z_order': 3}),
            ])
        ], categories=['background', 'a', 'b'])

        with TestDir() as test_dir:
            VocConverter.convert(dataset, test_dir, apply_colormap=False)

            cls_mask = load_mask(
                osp.join(test_dir, 'SegmentationClass', '1.png'))
            inst_mask = load_mask(
                osp.join(test_dir, 'SegmentationObject', '1.png'))
            self.assertTrue(np.array_equal([0, 1], np.unique(cls_mask)))
            self.assertTrue(np.array_equal([0, 1], np.unique(inst_mask)))
Ejemplo n.º 2
0
    def _parse_annotations(cls, xml_root, dataset_root, categories):
        def parse_attributes(attr_str):
            parsed = []
            if not attr_str:
                return parsed

            for attr in [a.strip() for a in attr_str.split(',') if a.strip()]:
                if '=' in attr:
                    name, value = attr.split('=', maxsplit=1)
                    if value.lower() in {'true', 'false'}:
                        value = value.lower() == 'true'
                    else:
                        try:
                            value = float(value)
                        except ValueError:
                            pass
                    parsed.append((name, value))
                else:
                    parsed.append((attr, True))

            return parsed

        label_cat = categories[AnnotationType.label]
        def get_label_id(label):
            if not label:
                return None
            idx, _ = label_cat.find(label)
            if idx is None:
                idx = label_cat.add(label)
            return idx

        image_annotations = []

        parsed_annotations = dict()
        group_assignments = dict()
        root_annotations = set()
        for obj_elem in xml_root.iter('object'):
            obj_id = int(obj_elem.find('id').text)

            ann_items = []

            label = get_label_id(obj_elem.find('name').text)

            attributes = []
            attributes_elem = obj_elem.find('attributes')
            if attributes_elem is not None and attributes_elem.text:
                attributes = parse_attributes(attributes_elem.text)

            occluded = False
            occluded_elem = obj_elem.find('occluded')
            if occluded_elem is not None and occluded_elem.text:
                occluded = (occluded_elem.text == 'yes')
            attributes.append(('occluded', occluded))

            deleted = False
            deleted_elem = obj_elem.find('deleted')
            if deleted_elem is not None and deleted_elem.text:
                deleted = bool(int(deleted_elem.text))

            user = ''

            poly_elem = obj_elem.find('polygon')
            segm_elem = obj_elem.find('segm')
            type_elem = obj_elem.find('type') # the only value is 'bounding_box'
            if poly_elem is not None:
                user_elem = poly_elem.find('username')
                if user_elem is not None and user_elem.text:
                    user = user_elem.text
                attributes.append(('username', user))

                points = []
                for point_elem in poly_elem.iter('pt'):
                    x = float(point_elem.find('x').text)
                    y = float(point_elem.find('y').text)
                    points.append(x)
                    points.append(y)

                if type_elem is not None and type_elem.text == 'bounding_box':
                    xmin = min(points[::2])
                    xmax = max(points[::2])
                    ymin = min(points[1::2])
                    ymax = max(points[1::2])
                    ann_items.append(Bbox(xmin, ymin, xmax - xmin, ymax - ymin,
                        label=label, attributes=attributes, id=obj_id,
                    ))
                else:
                    ann_items.append(Polygon(points,
                        label=label, attributes=attributes, id=obj_id,
                    ))
            elif segm_elem is not None:
                user_elem = segm_elem.find('username')
                if user_elem is not None and user_elem.text:
                    user = user_elem.text
                attributes.append(('username', user))

                mask_path = osp.join(dataset_root, LabelMePath.MASKS_DIR,
                    segm_elem.find('mask').text)
                if not osp.isfile(mask_path):
                    raise Exception("Can't find mask at '%s'" % mask_path)
                mask = load_mask(mask_path)
                mask = np.any(mask, axis=2)
                ann_items.append(Mask(image=mask, label=label, id=obj_id,
                    attributes=attributes))

            if not deleted:
                parsed_annotations[obj_id] = ann_items

            # Find parents and children
            parts_elem = obj_elem.find('parts')
            if parts_elem is not None:
                children_ids = []
                hasparts_elem = parts_elem.find('hasparts')
                if hasparts_elem is not None and hasparts_elem.text:
                    children_ids = [int(c) for c in hasparts_elem.text.split(',')]

                parent_ids = []
                ispartof_elem = parts_elem.find('ispartof')
                if ispartof_elem is not None and ispartof_elem.text:
                    parent_ids = [int(c) for c in ispartof_elem.text.split(',')]

                if children_ids and not parent_ids and hasparts_elem.text:
                    root_annotations.add(obj_id)
                group_assignments[obj_id] = [None, children_ids]

        # assign single group to all grouped annotations
        current_group_id = 0
        annotations_to_visit = list(root_annotations)
        while annotations_to_visit:
            ann_id = annotations_to_visit.pop()
            ann_assignment = group_assignments[ann_id]
            group_id, children_ids = ann_assignment
            if group_id:
                continue

            if ann_id in root_annotations:
                current_group_id += 1 # start a new group

            group_id = current_group_id
            ann_assignment[0] = group_id

            # continue with children
            annotations_to_visit.extend(children_ids)

        assert current_group_id == len(root_annotations)

        for ann_id, ann_items in parsed_annotations.items():
            group_id = 0
            if ann_id in group_assignments:
                ann_assignment = group_assignments[ann_id]
                group_id = ann_assignment[0]

            for ann_item in ann_items:
                if group_id:
                    ann_item.group = group_id

                image_annotations.append(ann_item)

        return image_annotations