随机森林-语义分割

语义分割大家应该都听说过，一般我们都是用深度学习来实现。比如比较经典的U-NET、DeepLab等。但是用随机森林来实现语义分割，大家可能还是第一次听说。下面的例子是随机森林实现语义分割的一个简单的实现，利用像素的颜色信息作为特征信息，掩码作为类别信息。如果想改进分割效果，可以自己去添加其它特征，比如CLBP纹理特征之类的。

import os  
import numpy as np  
import matplotlib.pyplot as plt  
from sklearn.ensemble import RandomForestClassifier  
from sklearn.model_selection import train_test_split  
from sklearn.metrics import classification_report  
from PIL import Image  
from matplotlib.colors import ListedColormap  

# 函数：加载数据集  
def load_dataset(images_dir, annotations_dir):  
    images = []  
    annotations = []  
    unique_labels = set()  

    for image_name in os.listdir(annotations_dir):  
        # 加载图像  
        image_path = os.path.join(images_dir, image_name)  
        image = Image.open(image_path)  
        images.append(np.array(image))  

        # 加载对应的标注  
        annotation_name = image_name.replace('.png', '.png')  # 假设标注文件是PNG格式  
        annotation_path = os.path.join(annotations_dir, annotation_name)  
        annotation = Image.open(annotation_path)  
        annotation_array = np.array(annotation)  
        annotations.append(annotation_array)  

        # 更新标签集合  
        unique_labels.update(np.unique(annotation_array))  

    num_classes = len(unique_labels)  
    return np.array(images), np.array(annotations), num_classes  

target_size = (1216, 480)  

# 添加数据预处理步骤  
def preprocess_images(images, target_size=target_size):  
    images_resized = np.array([np.array(Image.fromarray(img).resize(target_size)) for img in images])  
    images_normalized = images_resized / 255.0  
    return images_normalized  

def preprocess_annotations(annotations, target_size=target_size):  
    annotations_resized = np.array([np.array(Image.fromarray(ann).resize(target_size, Image.NEAREST)) for ann in annotations])  
    return annotations_resized  

# 函数：训练随机森林模型  
def train_random_forest(x_train, y_train, num_classes):  
    # 将图像数据展平为二维数组  
    n_samples, h, w, c = x_train.shape  
    x_train_flat = x_train.reshape(-1, c)  # 每个像素的RGB值作为特征  
    y_train_flat = y_train.flatten()  # 将标签展平  

    # 创建随机森林分类器  
    rf_classifier = RandomForestClassifier(n_estimators=300, random_state=42)  
    rf_classifier.fit(x_train_flat, y_train_flat)  

    return rf_classifier  

# 函数：进行预测  
def predict_with_rf(model, x):  
    n_samples, h, w, c = x.shape  
    x_flat = x.reshape(-1, c)  # 展平为二维数组  
    predictions_flat = model.predict(x_flat)  # 进行预测  
    predictions = predictions_flat.reshape(h, w)  # 重塑为图像形状  
    return predictions  

# 颜色定义  
colors = [(0, 0, 0),    # black  
          (255, 0, 0),  # red  
          (0, 255, 0),  # green  
          (0, 0, 255),  # blue  
          (255, 255, 0),  # yellow  
          (0, 255, 255),  # cyan  
          (255, 0, 255),  # magenta  
          (255, 128, 0),  # orange  
          (128, 255, 0),  # light green  
          (255, 0, 128),  # pink  
          (128, 0, 255),  # purple  
          (0, 255, 128),  # teal  
          (0, 128, 255),  # sky blue  
          (255, 128, 128),  # coral  
          (128, 255, 128),  # light green  
          (128, 128, 255),  # light blue  
          (255, 255, 128),  # light yellow  
          (128, 255, 255),  # aqua  
          (255, 128, 255)   # light pink  
          ]  

colormap = ListedColormap(colors)  

# 函数：可视化预测结果  
def visualize_prediction(predictions, save_path):  
    pred_colored = np.zeros((predictions.shape[0], predictions.shape[1], 3), dtype=np.uint8)  

    for label, color in enumerate(colors):  
        pred_colored[predictions == label] = color  

    pred_image = Image.fromarray(pred_colored)  
    pred_image.save(save_path)  
    print(f'Saved prediction to {save_path}')  

# 主程序  
images_dir = 'D:/images_dir'  
annotations_dir = 'D:/annotations_dir'  
x_train, y_train, num_classes = load_dataset(images_dir, annotations_dir)  

# 数据预处理  
x_train_preprocessed = preprocess_images(x_train)  
y_train_preprocessed = preprocess_annotations(y_train)  

# 训练随机森林模型  
rf_model = train_random_forest(x_train_preprocessed, y_train_preprocessed, num_classes)  

# 进行预测  
for i in range(len(x_train)):  
    predictions = predict_with_rf(rf_model, x_train_preprocessed[i:i+1])  # 逐张预测  
    save_path = os.path.join('D:/predictions', f'prediction_{i}.png')  
    visualize_prediction(predictions, save_path)  

print('所有预测已完成。')