Update app.py

app.py

@@ -342,14 +342,84 @@ def plt_to_image():
     img = Image.open(buf)
     return img
 
+def calculate_slice_statistics(one_mask, slice_size=128):
+    """计算每个切片的材质占比"""
+    num_rows, num_cols = one_mask.shape[0] // slice_size, one_mask.shape[1] // slice_size
+    slice_stats = {}
+    
+    for i in range(num_rows):
+        for j in range(num_cols):
+            slice_mask = one_mask[i*slice_size:(i+1)*slice_size, j*slice_size:(j+1)*slice_size]
+            unique, counts = np.unique(slice_mask, return_counts=True)
+            total_pixels = counts.sum()
+            slice_stats[(i, j)] = {k: v / total_pixels for k, v in zip(unique, counts)}
+            
+    return slice_stats
+
+def find_top_slices(slice_stats, exclusion_list, min_percent=0.7, min_slices=1, top_k=3):
+    """找出每个类材质占比最高的前三个切片，加入新的筛选条件"""
+    from collections import defaultdict
+    import heapq
+
+    top_slices = defaultdict(list)
+
+    for slice_pos, stats in slice_stats.items():
+        for material_id, percent in stats.items():
+            # 第一个判断：材质是否在排除列表中
+            if material_id in exclusion_list:
+                continue
+            # 第二个判断：材质占比是否至少为70%
+            if percent < min_percent:
+                continue
+
+            # 将符合条件的切片添加到堆中
+            if len(top_slices[material_id]) < top_k:
+                heapq.heappush(top_slices[material_id], (percent, slice_pos))
+            else:
+                heapq.heappushpop(top_slices[material_id], (percent, slice_pos))
+
+    # 过滤出符合第三个条件的材质
+    valid_top_slices = {}
+    for material_id, slices in top_slices.items():
+        if len(slices) > min_slices:  # 至少有超过一个切片
+            valid_top_slices[material_id] = sorted(slices, reverse=True)
+
+    return valid_top_slices
+
+def extract_and_visualize_top_slices(image, top_slices, slice_size=128):
+    """从原始图像中提取并可视化顶部切片"""
+    import matplotlib.pyplot as plt
+    
+    fig, axs = plt.subplots(nrows=len(top_slices), ncols=3, figsize=(15, 5 * len(top_slices)))
+    
+    if len(top_slices) == 1:
+        axs = [axs]
+    
+    for idx, (material_id, slices) in enumerate(top_slices.items()):
+        for col, (_, pos) in enumerate(slices):
+            i, j = pos
+            img_slice = image.crop((j * slice_size, i * slice_size, (j + 1) * slice_size, (i + 1) * slice_size))
+            axs[idx][col].imshow(img_slice)
+            axs[idx][col].set_title(f'Material {id2material[material_id]} - Slice {pos}')
+            axs[idx][col].axis('off')
+    
+    plt.tight_layout()
+    # 保存到内存，而不是显示图像
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    plt.close()
+    buf.seek(0)
+    img = Image.open(buf)
+    return img
 
+# main program
 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 = {id2material[key]: value for key, value in colors_per_class.items()}
     # 定义一个列表，包含需要从字典中删除的键
-    labels_to_remove = ['sky', 'background','Glass','tree','water','Plastic, clear']
+    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)
@@ -369,7 +439,14 @@ def process_image(image_path):
     # 重新命名 DataFrame 为 percentage_df 以清楚表达其内容
     percentage_df = counts_df.rename(columns={'计数': '数量', '百分比': '占比 (%)'})
 
-    return color_mask_img, palette_image, percentage_df
+    slice_size = 64
+    exclusion_list = [38]
+    slice_stats = calculate_slice_statistics(one_mask, slice_size=slice_size)
+    top_slices = find_top_slices(slice_stats, exclusion_list=exclusion_list, min_percent=0.5, min_slices=1)
+    slice_image=extract_and_visualize_top_slices(image_path, top_slices, slice_size=slice_size)
+    
+    return color_mask_img, palette_image, percentage_df,slice_image
+
 
 iface = gr.Interface(
     fn=process_image,
@@ -377,6 +454,7 @@ iface = gr.Interface(
     outputs=[
         gr.Image(type="pil", label="Color Mask"),
         gr.Image(type="pil", label="Color Palette"),
+        gr.Image(type='pil', label='Texture Slices')
         gr.DataFrame()
     ],
     title="Image Processing for Mask Visualization and Color Extraction",