Update app.py

app.py

@@ -125,6 +125,57 @@ id2label= {1: '动物皮',
  65: '纱窗',
  0: '背景/未知'}
 
+id2material={1: 'Animal skin',
+ 2: 'Bone/teeth/horn',
+ 3: 'Brickwork',
+ 4: 'Cardboard/Paper',
+ 6: 'Ceiling tile',
+ 7: 'Ceramic',
+ 8: 'Chalkboard/blackboard',
+ 10: 'Concrete',
+ 13: 'Fabric/cloth',
+ 15: 'Fire',
+ 16: 'Foliage',
+ 17: 'Food',
+ 18: 'Fur',
+ 19: 'Gemstone/quartz',
+ 20: 'Glass',
+ 21: 'Hair',
+ 23: 'Ice',
+ 24: 'Leather',
+ 26: 'Metal',
+ 27: 'Mirror',
+ 29: 'Paint/plaster',
+ 32: 'Photograph/painting',
+ 33: 'Plastic, clear',
+ 34: 'Plastic, non-clear',
+ 35: 'Rubber/latex',
+ 36: 'Sand',
+ 37: 'Skin/lips',
+ 38: 'Sky',
+ 39: 'Snow',
+ 41: 'Soil/mud',
+ 43: 'natural stone',
+ 44: 'polished stone & engineered stone',
+ 46: 'Tile',
+ 47: 'Wallpaper',
+ 48: 'Water',
+ 49: 'Wax',
+ 50: 'Whiteboard',
+ 52: 'Wood',
+ 53: 'tree',
+ 56: 'Asphalt',
+ 57: 'enamel',
+ 58: 'Rammed earth',
+ 59: 'composite decorative board',
+ 60: 'Cement',
+ 61: 'Pottery',
+ 62: 'Roofing waterproof material',
+ 63: 'Metal mesh window (perspective)',
+ 64: 'carved brick',
+ 65: 'window screen',
+ 0: 'background'}
+
 model_save_path ='jinfengxie/BFM_segformer0821'
 model = TFSegformerForSemanticSegmentation.from_pretrained(model_save_path)
 
@@ -294,9 +345,9 @@ def process_image(image_path):
     #image = Image.open(image_path)
     one_mask,color_mask, counts_dict = predict_and_visualize(image_path)
     colors_per_class=ext_colors(image_path,one_mask,n_clusters=4)
-    colors_per_label = {id2label[key]: value for key, value in colors_per_class.items()}
+    colors_per_label = {id2material[key]: value for key, value in colors_per_class.items()}
     # 定义一个列表，包含需要从字典中删除的键
-    labels_to_remove = ['天空', '背景/未知','玻璃','树木','水','透明塑料']
+    labels_to_remove = ['sky', 'background','Glass','tree','water','Plastic, clear']
     # 使用字典推导式删除列表中的键
     colors_per_label = {key: value for key, value in colors_per_label.items() if key not in labels_to_remove}
     palette_image = plot_material_color_palette_grid(colors_per_label)
@@ -304,6 +355,7 @@ def process_image(image_path):
     # 将结果转化为图片展示
     plt.figure(figsize=(5, 5))
     plt.imshow(color_mask)
+    plt.tight_layout()
     plt.axis('off')
     color_mask_img = plt_to_image()
     counts_dict2={id2label[key]: value for key, value in counts_dict.items()}