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
from sklearn.cluster import KMeans
from skimage import color
import io
import pandas as pd

# Set the font to support Chinese characters
matplotlib.rcParams['font.family'] = 'SimHei'

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: '背景/未知'}

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==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/40))  # 尝试加载Arial字体，大小为12

    for key, position in label_positions.items():
        if key in id2label:
            # 绘制文本，您可能需要调整文本位置和字体大小
            draw.text((position[0], position[1]), str(key), 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')
    plt.close()
    buf.seek(0)
    img = Image.open(buf)
    return img


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()}
    # 定义一个列表，包含需要从字典中删除的键
    labels_to_remove = ['天空', '背景/未知','玻璃','树木','水','透明塑料']
    # 使用字典推导式删除列表中的键
    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.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={'计数': '数量', '百分比': '占比 (%)'})

    return color_mask_img, palette_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.DataFrame()
    ],
    title="Image Processing for Mask Visualization and Color Extraction",
    description="Upload an image to extract color masks, class counts, and color palettes."
)

iface.launch()