1. 测试数据
公司标注平台某张图片的classification
导出结果,也就是图片分类结果。
这个结果应该是树状结构,只是存储时为了方便以平铺形式存储。但是这样使用时就不太方便,所以要复原为树状结构。
clf = [
{
"classificationId": 831139,
"classificationVersion": 1,
"id": "de4f2a89-376e-4b26-a1a4-9987fb60d74b",
"values": [
{
"alias": "物体计数",
"attributeVersion": 1,
"id": "6c020d84-e486-4300-8c96-b1ff9542811c",
"isLeaf": False,
"name": "object_counts",
"value": "object_counts"
},
{
"alias": "拖车数量",
"attributeVersion": 1,
"id": "f8120dc7-45fc-4bcc-8020-9369edfab72a",
"isLeaf": True,
"name": "trailer_count",
"pid": "6c020d84-e486-4300-8c96-b1ff9542811c",
"value": "1"
},
{
"alias": "拖车头数量",
"attributeVersion": 1,
"id": "d01f6289-5ba7-4980-99d1-3644d1883be2",
"isLeaf": True,
"name": "tractor_count",
"pid": "6c020d84-e486-4300-8c96-b1ff9542811c",
"value": "1"
},
{
"alias": "货车数量",
"attributeVersion": 1,
"id": "00129c8e-ce66-4bef-841f-6d0a885c49e9",
"isLeaf": True,
"name": "truck_count",
"pid": "6c020d84-e486-4300-8c96-b1ff9542811c",
"value": "0"
}
]
},
{
"classificationId": 831140,
"classificationVersion": 1,
"id": "4475dd3f-a5a6-40ee-899a-645df2aa0b32",
"values": [
{
"alias": "天气",
"attributeVersion": 1,
"id": "3f6075ed-9ddd-4ccb-a4d6-86bc60e410e5",
"isLeaf": True,
"name": "environment_conditions",
"value": "sunny"
}
]
},
{
"classificationId": 831141,
"classificationVersion": 1,
"id": "f40b57e0-7eb5-4604-96da-6dc009edcdfe",
"values": [
{
"alias": "环境",
"attributeVersion": 1,
"id": "0f5ea0c5-4e7f-45c8-966d-ef46b5d42f85",
"isLeaf": True,
"name": "scene_type",
"value": "other"
}
]
},
{
"classificationId": 831142,
"classificationVersion": 1,
"id": "c29afde4-7c2e-49eb-b8e1-a10e091f579e",
"values": [
{
"alias": "图片质量",
"attributeVersion": 1,
"id": "306949de-022d-4aa7-9b4d-782dc752874a",
"isLeaf": False,
"name": "image_quality",
"value": "image_quality"
},
{
"alias": "图片破损",
"attributeVersion": 1,
"id": "2ba006ac-dc21-4dbe-a4f7-4115811cd1ae",
"isLeaf": True,
"name": "image_broken",
"pid": "306949de-022d-4aa7-9b4d-782dc752874a",
"value": "no"
},
{
"alias": "镜头尘土",
"attributeVersion": 1,
"id": "e05205d9-0135-43de-8551-ef3653ea5f10",
"isLeaf": True,
"name": "dirt/dust_on_image",
"pid": "306949de-022d-4aa7-9b4d-782dc752874a",
"value": "no"
},
{
"alias": "镜头水渍",
"attributeVersion": 1,
"id": "07a8405d-f8c2-4f5b-be2b-d9db74fc1dc8",
"isLeaf": True,
"name": "water_on_lense",
"pid": "306949de-022d-4aa7-9b4d-782dc752874a",
"value": "no"
},
{
"alias": "图片过曝",
"attributeVersion": 1,
"id": "b436d1c4-3061-4235-9c9a-303687c00c10",
"isLeaf": True,
"name": "overexposure_of_image",
"pid": "306949de-022d-4aa7-9b4d-782dc752874a",
"value": "no"
}
]
},
{
"classificationId": 831143,
"classificationVersion": 1,
"id": "1b923564-f887-441f-98ca-e01968da4d4d",
"values": [
{
"alias": "生物",
"attributeVersion": 1,
"id": "e47c2003-542e-40d5-881e-3600a80cee97",
"isLeaf": False,
"name": "mammals",
"value": "mammals"
},
{
"alias": "人类",
"attributeVersion": 1,
"id": "55de181d-c835-4f9e-beed-01ea071f2980",
"isLeaf": True,
"name": "person",
"pid": "e47c2003-542e-40d5-881e-3600a80cee97",
"value": "no"
},
{
"alias": "动物",
"attributeVersion": 1,
"id": "bbe48866-3aa0-4852-a0dc-f75102354eac",
"isLeaf": True,
"name": "animal",
"pid": "e47c2003-542e-40d5-881e-3600a80cee97",
"value": "no"
}
]
},
{
"classificationId": 831144,
"classificationVersion": 2,
"id": "0e13d564-15c2-461a-b55f-bdf282f516ab",
"values": [
{
"alias": "材料",
"attributeVersion": 2,
"id": "fecbe77c-58f9-4879-bb40-dfed6b4181c0",
"isLeaf": False,
"name": "material",
"value": "material"
},
{
"alias": "材料在拖车内",
"attributeVersion": 2,
"id": "8364e0f6-ce41-4c9f-951e-3c4c036045b8",
"isLeaf": True,
"name": "material_in_trailer",
"pid": "fecbe77c-58f9-4879-bb40-dfed6b4181c0",
"value": "no"
},
{
"alias": "材料在空中",
"attributeVersion": 2,
"id": "3a85701d-db8d-41dd-83fe-90d92cfd7617",
"isLeaf": True,
"name": "material_in_air",
"pid": "fecbe77c-58f9-4879-bb40-dfed6b4181c0",
"value": "no"
}
]
}
]
2. 原始的to_dict
源码:
class Tree(object):
...
def to_dict(self, nid=None, key=None, sort=True, reverse=False, with_data=False):
"""Transform the whole tree into a dict."""
nid = self.root if (nid is None) else nid
ntag = self[nid].tag
tree_dict = {ntag: {"children": []}}
if with_data:
tree_dict[ntag]["data"] = self[nid].data
if self[nid].expanded:
queue = [self[i] for i in self[nid].successors(self._identifier)]
key = (lambda x: x) if (key is None) else key
if sort:
queue.sort(key=key, reverse=reverse)
for elem in queue:
tree_dict[ntag]["children"].append(
self.to_dict(
elem.identifier, with_data=with_data, sort=sort, reverse=reverse
)
)
if len(tree_dict[ntag]["children"]) == 0:
tree_dict = (
self[nid].tag if not with_data else {ntag: {"data": self[nid].data}}
)
return tree_dict
效果:
from treelib import Node, Tree
tree = Tree()
tree.create_node("Root", "root")
for x in clf:
for v in x["values"]:
tree.create_node(
tag=v["name"],
identifier=v["id"],
parent=v.get("pid") or "root",
data=v["value"]
)
print(tree.to_dict(with_data=True))
"""
{
'Root': {
'children': [
{'environment_conditions': {'data': 'sunny'}},
{
'image_quality': {
'children': [
{'dirt/dust_on_image': {'data': 'no'}},
{'image_broken': {'data': 'no'}},
{'overexposure_of_image': {'data': 'no'}},
{'water_on_lense': {'data': 'no'}}
],
'data': 'image_quality'
}
},
{
'mammals': {
'children': [{'animal': {'data': 'no'}}, {'person': {'data': 'no'}}],
'data': 'mammals'
}
},
{
'material': {
'children': [{'material_in_air': {'data': 'no'}}, {'material_in_trailer': {'data': 'no'}}],
'data': 'material'
}
},
{
'object_counts': {
'children': [
{'tractor_count': {'data': '1'}},
{'trailer_count': {'data': '1'}},
{'truck_count': {'data': '0'}}
],
'data': 'object_counts'
}
},
{'scene_type': {'data': 'other'}}
],
'data': None
}
}
"""
可以看出,多了很多children
和data
字段,而且children
是以列表形式存储,没有整合到一个字典里。
2. 重写to_dict
本次需求是,不要children
和data
字段,如果不是叶子节点,就一直把tag
存为值,data
不管,如果到了叶子节点,再把data
作为值存储。
重写:
from treelib import Node, Tree
class CustomTree(Tree):
def to_dict(self, nid=None, key=None, sort=True, reverse=False, with_data=False):
"""Transform the whole tree into a dict."""
nid = self.root if (nid is None) else nid
ntag = self[nid].tag
tree_dict = {ntag: {}}
queue = [self[i] for i in self[nid].successors(self._identifier)]
key = (lambda x: x) if (key is None) else key
if sort:
queue.sort(key=key, reverse=reverse)
for elem in queue:
tree_dict[ntag].update(
self.to_dict(
elem.identifier, with_data=with_data, sort=sort, reverse=reverse
)
)
if len(tree_dict[ntag]) == 0:
tree_dict = {ntag: None} if not with_data else {ntag: self[nid].data}
return tree_dict
效果:
tree = CustomTree()
tree.create_node("Root", "root")
for x in clf:
for v in x["values"]:
tree.create_node(
tag=v["name"],
identifier=v["id"],
parent=v.get("pid") or "root",
data=v["value"]
)
print(tree.to_dict(with_data=True))
"""
{
'Root': {
'environment_conditions': 'sunny',
'image_quality': {
'dirt/dust_on_image': 'no',
'image_broken': 'no',
'overexposure_of_image': 'no',
'water_on_lense': 'no'
},
'mammals': {'animal': 'no', 'person': 'no'},
'material': {'material_in_air': 'no', 'material_in_trailer': 'no'},
'object_counts': {'tractor_count': '1', 'trailer_count': '1', 'truck_count': '0'},
'scene_type': 'other'
}
}
"""
大功告成!