Rearrange category mapping of COCO dictionary. Can also be used to filter some of the categories. Args: desired_name2id: Desired category name to id mapping, e.g. {"big_vehicle": 1, "car": 2, "human": 3}. coco_dict: COCO formatted dictionary. Returns:
(desired_name2id: dict, coco_dict: dict)
| 1725 | |
| 1726 | |
| 1727 | def update_categories(desired_name2id: dict, coco_dict: dict) -> dict: |
| 1728 | """Rearrange category mapping of COCO dictionary. |
| 1729 | |
| 1730 | Can also be used to filter some of the categories. |
| 1731 | |
| 1732 | Args: |
| 1733 | desired_name2id: Desired category name to id mapping, |
| 1734 | e.g. {"big_vehicle": 1, "car": 2, "human": 3}. |
| 1735 | coco_dict: COCO formatted dictionary. |
| 1736 | |
| 1737 | Returns: |
| 1738 | coco_target: COCO dict with updated/filtered categories. |
| 1739 | """ |
| 1740 | # so that original variable doesn't get affected |
| 1741 | coco_source = copy.deepcopy(coco_dict) |
| 1742 | |
| 1743 | # init target coco dict |
| 1744 | coco_target: dict = {"images": [], "annotations": [], "categories": []} |
| 1745 | |
| 1746 | # init vars |
| 1747 | currentid2desiredid_mapping: dict = {} |
| 1748 | # create category id mapping (currentid2desiredid_mapping) |
| 1749 | for category in coco_source["categories"]: |
| 1750 | current_category_id = category["id"] |
| 1751 | current_category_name = category["name"] |
| 1752 | if current_category_name in desired_name2id.keys(): |
| 1753 | currentid2desiredid_mapping[current_category_id] = desired_name2id[current_category_name] |
| 1754 | else: |
| 1755 | # ignore categories that are not included in desired_name2id |
| 1756 | currentid2desiredid_mapping[current_category_id] = -1 |
| 1757 | |
| 1758 | # update annotations |
| 1759 | for annotation in coco_source["annotations"]: |
| 1760 | current_category_id = annotation["category_id"] |
| 1761 | desired_category_id = currentid2desiredid_mapping[current_category_id] |
| 1762 | # append annotations with category id present in desired_name2id |
| 1763 | if desired_category_id != -1: |
| 1764 | # update cetegory id |
| 1765 | annotation["category_id"] = desired_category_id |
| 1766 | # append updated annotation to target coco dict |
| 1767 | coco_target["annotations"].append(annotation) |
| 1768 | |
| 1769 | # create desired categories |
| 1770 | categories = [] |
| 1771 | for name in desired_name2id.keys(): |
| 1772 | category = {} |
| 1773 | category["name"] = category["supercategory"] = name |
| 1774 | category["id"] = desired_name2id[name] |
| 1775 | categories.append(category) |
| 1776 | |
| 1777 | # update categories |
| 1778 | coco_target["categories"] = categories |
| 1779 | |
| 1780 | # update images |
| 1781 | coco_target["images"] = coco_source["images"] |
| 1782 | |
| 1783 | return coco_target |
| 1784 |