Spaces:
No application file
No application file
| import cv2 | |
| import pickle | |
| import numpy as np | |
| from pathlib import Path | |
| from tqdm import tqdm | |
| from isegm.data.datasets import LvisDataset, CocoDataset | |
| from isegm.utils.misc import get_bbox_from_mask, get_bbox_iou | |
| from scripts.annotations_conversion.common import get_masks_hierarchy, get_iou, encode_masks | |
| def create_annotations(lvis_path: Path, coco_path: Path, dataset_split='train', min_object_area=80): | |
| lvis_dataset = LvisDataset(lvis_path, split=dataset_split) | |
| lvis_samples = lvis_dataset.dataset_samples | |
| lvis_annotations = lvis_dataset.annotations | |
| coco_dataset = CocoDataset(coco_path, split=dataset_split + '2017') | |
| coco_lvis_mapping = [] | |
| lvis_images = {x['coco_url'].split('/')[-1].split('.')[0]: lvis_indx | |
| for lvis_indx, x in enumerate(lvis_samples)} | |
| for indx, coco_sample in enumerate(coco_dataset.dataset_samples): | |
| lvis_indx = lvis_images.get(coco_sample['file_name'].split('.')[0], None) | |
| if lvis_indx is not None: | |
| coco_lvis_mapping.append((indx, lvis_indx)) | |
| output_masks_path = lvis_path / dataset_split / 'masks' | |
| output_masks_path.mkdir(parents=True, exist_ok=True) | |
| hlvis_annotation = dict() | |
| for coco_indx, lvis_indx in tqdm(coco_lvis_mapping): | |
| coco_sample = get_coco_sample(coco_dataset, coco_indx) | |
| lvis_info = lvis_samples[lvis_indx] | |
| lvis_annotation = lvis_annotations[lvis_info['id']] | |
| empty_mask = np.zeros((lvis_info['height'], lvis_info['width'])) | |
| image_name = lvis_info['coco_url'].split('/')[-1].split('.')[0] | |
| lvis_masks = [] | |
| lvis_bboxes = [] | |
| for obj_annotation in lvis_annotation: | |
| obj_mask = lvis_dataset.get_mask_from_polygon(obj_annotation, empty_mask) | |
| obj_mask = obj_mask == 1 | |
| if obj_mask.sum() >= min_object_area: | |
| lvis_masks.append(obj_mask) | |
| lvis_bboxes.append(get_bbox_from_mask(obj_mask)) | |
| coco_bboxes = [] | |
| coco_masks = [] | |
| for inst_id in coco_sample['instances_info'].keys(): | |
| obj_mask = coco_sample['instances_mask'] == inst_id | |
| if obj_mask.sum() >= min_object_area: | |
| coco_masks.append(obj_mask) | |
| coco_bboxes.append(get_bbox_from_mask(obj_mask)) | |
| masks = [] | |
| for coco_j, coco_bbox in enumerate(coco_bboxes): | |
| for lvis_i, lvis_bbox in enumerate(lvis_bboxes): | |
| if get_bbox_iou(lvis_bbox, coco_bbox) > 0.70 and \ | |
| get_iou(lvis_masks[lvis_i], coco_masks[coco_j]) > 0.70: | |
| break | |
| else: | |
| masks.append(coco_masks[coco_j]) | |
| for ti, (lvis_mask, lvis_bbox) in enumerate(zip(lvis_masks, lvis_bboxes)): | |
| for tj_mask, tj_bbox in zip(lvis_masks[ti + 1:], lvis_bboxes[ti + 1:]): | |
| bbox_iou = get_bbox_iou(lvis_bbox, tj_bbox) | |
| if bbox_iou > 0.7 and get_iou(lvis_mask, tj_mask) > 0.85: | |
| break | |
| else: | |
| masks.append(lvis_mask) | |
| masks_meta = [(get_bbox_from_mask(x), x.sum()) for x in masks] | |
| if not masks: | |
| continue | |
| hierarchy = get_masks_hierarchy(masks, masks_meta) | |
| for obj_id, obj_info in list(hierarchy.items()): | |
| if obj_info['parent'] is None and len(obj_info['children']) == 0: | |
| hierarchy[obj_id] = None | |
| merged_mask = np.max(masks, axis=0) | |
| num_instance_masks = len(masks) | |
| for obj_id in coco_sample['semantic_info'].keys(): | |
| obj_mask = coco_sample['semantic_map'] == obj_id | |
| obj_mask = np.logical_and(obj_mask, np.logical_not(merged_mask)) | |
| if obj_mask.sum() > 500: | |
| masks.append(obj_mask) | |
| hlvis_annotation[image_name] = { | |
| 'num_instance_masks': num_instance_masks, | |
| 'hierarchy': hierarchy | |
| } | |
| with open(output_masks_path / f'{image_name}.pickle', 'wb') as f: | |
| pickle.dump(encode_masks(masks), f) | |
| with open(lvis_path / dataset_split / 'hannotation.pickle', 'wb') as f: | |
| pickle.dump(hlvis_annotation, f, protocol=pickle.HIGHEST_PROTOCOL) | |
| def get_coco_sample(dataset, index): | |
| dataset_sample = dataset.dataset_samples[index] | |
| image_path = dataset.images_path / dataset.get_image_name(dataset_sample['file_name']) | |
| label_path = dataset.labels_path / dataset_sample['file_name'] | |
| image = cv2.imread(str(image_path)) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| label = cv2.imread(str(label_path), cv2.IMREAD_UNCHANGED).astype(np.int32) | |
| label = 256 * 256 * label[:, :, 0] + 256 * label[:, :, 1] + label[:, :, 2] | |
| instance_map = np.full_like(label, 0) | |
| semantic_map = np.full_like(label, 0) | |
| semantic_info = dict() | |
| instances_info = dict() | |
| for segment in dataset_sample['segments_info']: | |
| class_id = segment['category_id'] | |
| obj_id = segment['id'] | |
| if class_id not in dataset._things_labels_set: | |
| semantic_map[label == obj_id] = obj_id | |
| semantic_info[obj_id] = {'ignore': False} | |
| continue | |
| instance_map[label == obj_id] = obj_id | |
| ignore = segment['iscrowd'] == 1 | |
| instances_info[obj_id] = { | |
| 'ignore': ignore | |
| } | |
| sample = { | |
| 'image': image, | |
| 'instances_mask': instance_map, | |
| 'instances_info': instances_info, | |
| 'semantic_map': semantic_map, | |
| 'semantic_info': semantic_info | |
| } | |
| return sample | |