文章描述了一个将公司标注平台的图片分类结果从扁平化存储转换为树状结构的过程，以便于更高效地处理和展示。作者使用了`CustomTree`类重写了`to_dict`方法，移除了children和data字段，只保留叶节点的数据。
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'
}
}
"""
大功告成！
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
