app.py

from PIL import Image, ImageDraw, ImageFont
from skimage.measure import label, regionprops
import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
from tensorflow.keras.preprocessing.image import array_to_img
import json
import os
from transformers import AutoModel
from transformers import TFSegformerForSemanticSegmentation
import matplotlib.pyplot as plt
import matplotlib
import matplotlib.font_manager as fm
from sklearn.cluster import KMeans
from skimage import color
import io
import pandas as pd

# Set the font to support Chinese characters
#font_path = 'simhei.ttf'
#font_prop = fm.FontProperties(fname=font_path)
#matplotlib.rcParams['font.family'] = font_prop.get_name()
#matplotlib.rcParams['font.family'] = 'Droid Sans Fallback'

id2color= {1: [209, 35, 69],
 2: [216, 208, 246],
 3: [172, 196, 170],
 4: [178, 80, 80],
 6: [89, 89, 89],
 7: [160, 146, 229],
 8: [18, 17, 20],
 10: [190, 209, 189],
 13: [37, 12, 156],
 15: [250, 50, 83],
 16: [61, 245, 61],
 17: [230, 203, 104],
 18: [125, 104, 227],
 19: [228, 225, 249],
 20: [51, 221, 255],
 21: [95, 95, 95],
 23: [156, 239, 255],
 24: [153, 102, 51],
 26: [0, 0, 226],
 27: [254, 242, 208],
 29: [89, 134, 179],
 32: [255, 0, 204],
 33: [170, 240, 209],
 34: [140, 120, 240],
 35: [118, 255, 166],
 36: [250, 250, 55],
 37: [243, 232, 208],
 38: [1, 118, 141],
 39: [243, 241, 255],
 41: [158, 108, 4],
 43: [132, 0, 0],
 44: [245, 147, 49],
 46: [240, 120, 240],
 47: [149, 83, 203],
 48: [52, 209, 183],
 49: [200, 101, 0],
 50: [65, 112, 192],
 52: [255, 204, 51],
 53: [36, 179, 83],
 56: [90, 98, 89],
 57: [255, 191, 0],
 58: [204, 153, 51],
 59: [31, 73, 125],
 60: [155, 149, 205],
 61: [154, 150, 169],
 62: [128, 128, 128],
 63: [163, 160, 172],
 64: [255, 106, 77],
 65: [115, 51, 128],
 0: [10, 9, 10]}

id2label= {1: '动物皮',
 2: '骨/牙/角',
 3: '砖块',
 4: '纸板/纸',
 6: '天花板瓦片',
 7: '瓷',
 8: '黑板',
 10: '混凝土',
 13: '织物/布/地毯',
 15: '火',
 16: '树叶',
 17: '食物',
 18: '毛皮',
 19: '宝石/石英',
 20: '玻璃',
 21: '毛发',
 23: '冰',
 24: '皮革',
 26: '金属',
 27: '镜子',
 29: '油漆/抹灰/石膏',
 32: '照片/绘画/布面招牌',
 33: '透明塑料',
 34: '非透明塑料',
 35: '橡胶/乳胶',
 36: '沙',
 37: '皮肤/嘴唇',
 38: '天空',
 39: '雪',
 41: '土壤/泥土',
 43: '天然石材',
 44: '抛光石材',
 46: '片状地砖/石地砖/瓷地砖',
 47: '壁纸',
 48: '水',
 49: '蜡',
 50: '白板',
 52: '木材',
 53: '树木',
 56: '沥青',
 57: '珐琅/琉璃',
 58: '夯土',
 59: '塑钢复合装饰板',
 60: '水泥',
 61: '陶',
 62: '屋顶防水卷材',
 63: '金属网窗（远景）',
 64: '砖雕',
 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)

def predict_and_visualize(image):
    #image = Image.open(image_path)
    image_np = np.array(image)
    height,width,_=image_np.shape
    maxhl=max(height,width)
    image = tf.convert_to_tensor(image_np, dtype=tf.float32) 
    if maxhl>1500:
        if maxhl==height:
            image=tf.image.resize(image,(1500,int(1500*width/height)))
        if maxhl==width:
            image=tf.image.resize(image,(int(1500*height/width),1500))
    #image = tf.image.resize_with_pad(image, 1500, 1500)
    image = tf.cast(image, tf.float32) / 255.0
    image = tf.transpose(image, perm=[2, 0, 1])
    images= tf.expand_dims(image, axis=0)
    # 进行预测
    preds = model.predict(images).logits
    pred_mask = tf.argmax(preds, axis=1)
    pred_mask = tf.expand_dims(pred_mask, axis=-1)
    pred_mask = pred_mask[0]  # 取出批处理的第一个结果
    pred_mask=tf.image.resize(pred_mask,(height,width),method='nearest')
    pred_mask=tf.squeeze(pred_mask)
    print(pred_mask.shape)
    #pred_mask = pred_mask[:,:,0] .numpy() # 取出批处理的第一个结果
    #print(pred_mask.shape)
    unique, counts = np.unique(pred_mask, return_counts=True)
    counts_dict = dict(zip(unique, counts))
    # 转换预测掩码为颜色图像
    color_mask = np.zeros((height,width, 3))
    label_positions = {}
    for key, value in id2color.items():
        #print("mask shape",mask.shape)
        color_mask[pred_mask == key] = np.array(value)   # 颜色值需要被标准化到[0,1]
        indices = np.where(pred_mask == key)
        if indices[0].size > 0:
            # 计算标签的位置为当前类别像素的中心点
            label_positions[key] = (np.mean(indices[1]), np.mean(indices[0]))        
   
    color_mask = color_mask.astype(np.uint8)
    result_image = Image.fromarray(color_mask)
    draw = ImageDraw.Draw(result_image)
    font = ImageFont.truetype("arial.ttf", int(height/30))  # 尝试加载Arial字体，大小为12

    for key, position in label_positions.items():
        if key in id2label:
            # 绘制文本，您可能需要调整文本位置和字体大小
            material=id2material[key]
            draw.text((position[0], position[1]), str(material), font=font, fill='white')

    return pred_mask,result_image,counts_dict


def ext_colors(image_path,mask,n_clusters=4):
    #image = Image.open(image_path)

    # 将图像和掩码转换为numpy数组
    image_np = np.array(image_path)
    mask_np = np.array(mask)

    # 获取掩码中的唯一类别
    unique_classes = np.unique(mask_np)

    # 为每个类别提取颜色
    colors_per_class = {}
    for cls in unique_classes:
        # 提取当前类别的像素点
        indices = np.where(mask_np == cls)
        #print(indices)
        pixels = image_np[indices]

        # 使用K-means聚类来找到主要颜色
        kmeans = KMeans(n_clusters=n_clusters,n_init=10)
        kmeans.fit(pixels)
        dominant_colors = kmeans.cluster_centers_

        # 将颜色存储为整数值
        dominant_colors = dominant_colors.astype(int)
        
        # 保存颜色
        colors_per_class[cls] = dominant_colors


    return colors_per_class

def plot_material_color_palette_grid(material_dict, materials_per_row=4):
    # Calculate total number of color rows and header rows needed
    total_rows = sum((len(colors) + 1) for colors in material_dict.values())  # +1 for the header row per material
    num_materials = len(material_dict)
    grid_rows = (num_materials + materials_per_row - 1) // materials_per_row
    total_grid_rows = 0
    for i in range(grid_rows):
        row_materials = list(material_dict.keys())[i * materials_per_row:(i + 1) * materials_per_row]
        row_height = max(len(material_dict[mat]) for mat in row_materials if mat in material_dict) + 1
        total_grid_rows += row_height

    # Set dimensions and spacing
    block_width = 1
    block_height = 0.5
    text_gap = 0.2
    row_gap = 0.2
    column_gap = 1.5  # Gap between material columns within the same row

    # Calculate figure width and height dynamically
    fig_width = materials_per_row * (block_width + text_gap + column_gap)
    fig_height = total_grid_rows * (block_height + row_gap)

    # Create a figure and a set of subplots
    fig, ax = plt.subplots(figsize=(fig_width, fig_height))

    # Set the title of the figure
    #ax.set_title('Material Color Palette Grid')

    # Remove axes
    ax.axis('off')

    # Reverse the Y-axis to top-align the origin
    ax.invert_yaxis()

    current_row = 0  # Tracker for the current row position in the grid
    for i in range(grid_rows):
        row_materials = list(material_dict.keys())[i * materials_per_row:(i + 1) * materials_per_row]
        max_row_height = max(len(material_dict[mat]) for mat in row_materials if mat in material_dict) + 1
        for j, material in enumerate(row_materials):
            if material not in material_dict:
                continue
            colors = material_dict[material]
            # Add a header for each material class
            ax.text(j * (block_width + text_gap + column_gap), current_row * (block_height + row_gap)+0.5,
                    material, va='center', fontsize=12, fontweight='bold', ha='left')
            material_row_start = current_row
            for k, color in enumerate(colors):
                # Normalize the RGB values to [0, 1] for Matplotlib
                normalized_color = np.array(color) / 255.0
                y_pos = (material_row_start + 1 + k) * (block_height + row_gap)
                # Draw a rectangle for each color
                ax.add_patch(plt.Rectangle((j * (block_width + text_gap + column_gap), y_pos),
                                           block_width, block_height, color=normalized_color))
                # Annotate the RGB values to the right of each color block
                ax.text(j * (block_width + text_gap + column_gap) + block_width + text_gap, y_pos + block_height / 2,
                        str(color), va='center', fontsize=10)
        current_row += max_row_height

    # Adjust plot limits
    ax.set_xlim(0, fig_width)
    ax.set_ylim(current_row * (block_height + row_gap), 0)

    # 保存到内存，而不是显示图像
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)
    img = Image.open(buf)
    return img

# 将matplotlib图转换为图像
def plt_to_image():
    buf = io.BytesIO()
    plt.savefig(buf, format='png',dpi=300)
    plt.close()
    buf.seek(0)
    img = Image.open(buf)
    return img

def calculate_slice_statistics(one_mask, slice_size=256):
    """计算每个切片的材质占比"""
    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=256):    
    fig, axs = plt.subplots(nrows=len(top_slices), ncols=3, figsize=(15, 5 * len(top_slices)))
    image=Image.fromarray(image)
    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']
    # 使用字典推导式删除列表中的键
    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)

    # 将结果转化为图片展示
    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()}
    counts_df = pd.DataFrame(list(counts_dict2.items()), columns=['类别', '计数'])
    # 计算总计数
    total_count = counts_df['计数'].sum()
    # 计算每个类别的百分比
    counts_df['百分比'] = (counts_df['计数'] / total_count * 100).round(2)
    # 重新命名 DataFrame 为 percentage_df 以清楚表达其内容
    percentage_df = counts_df.rename(columns={'计数': 'pixels', '百分比': 'percentage (%)'})

    slice_size = 128
    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, slice_image, percentage_df


iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(),
    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="Building Facade Material Segmentation",
    description="Upload an image to segment material masks, and get color palettes."
)

iface.launch()