Как построить модель для классификации изображений на бинарной маске с квадратами, кругами и треугольниками?

Вопрос о классификации изображений с геометрическими фигурами — это интересная задача компьютерного зрения. Рассмотрю несколько подходов от простого к сложному.

Задача

Дано:

• Бинарные маски (черно-белые изображения)
• Каждая маска содержит ОДНУ фигуру из трёх типов: квадрат, круг или треугольник
• Нужно классифицировать тип фигуры

Подход 1: Классические признаки + Machine Learning

Для простых геометрических фигур можно использовать традиционные методы компьютерного зрения:

import cv2
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from scipy import ndimage

def extract_shape_features(mask):
    """Извлекаем признаки из маски"""
    
    # Находим контур
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    if len(contours) == 0:
        return None
    
    cnt = contours[0]
    
    # Признак 1: Периметр
    perimeter = cv2.arcLength(cnt, True)
    
    # Признак 2: Площадь
    area = cv2.contourArea(cnt)
    
    # Признак 3: Компактность (изопериметрическое отношение)
    # Для круга это близко к 1, для других фигур меньше
    compactness = (4 * np.pi * area) / (perimeter ** 2)
    
    # Признак 4: Эксцентриситет (через момент инерции)
    mu = cv2.moments(cnt)
    hu_moments = cv2.HuMoments(cnt).flatten()  # 7 инвариантных моментов Ху
    
    # Признак 5: Количество углов
    # Аппроксимируем контур полигоном
    epsilon = 0.02 * perimeter
    approx = cv2.approxPolyDP(cnt, epsilon, True)
    n_vertices = len(approx)
    
    # Признак 6: Коэффициент заполнения (солидность)
    # hull = cv2.convexHull(cnt)
    # hull_area = cv2.contourArea(hull)
    # solidity = area / hull_area
    
    # Признак 7: Вытянутость (ratio of principal axes)
    # (x, y), (MA, ma), angle = cv2.fitEllipse(cnt)
    # if ma > 0:
    #     elongation = MA / ma
    # else:
    #     elongation = 1
    
    features = np.concatenate([
        [perimeter, area, compactness, n_vertices],
        hu_moments
    ])
    
    return features

# Подготовка данных
X = []
y = []

# Загружаем изображения
for label, shape_name in enumerate(["square", "circle", "triangle"]):
    image_paths = glob.glob(f"data/{shape_name}/*.png")
    for path in image_paths:
        mask = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        features = extract_shape_features(mask)
        if features is not None:
            X.append(features)
            y.append(label)

X = np.array(X)
y = np.array(y)

# Обучение
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

accuracy = model.score(X_test, y_test)
print(f"Accuracy: {accuracy:.4f}")

# Важность признаков
feature_importance = model.feature_importances_
print(f"\nВажность признаков:\nПериметр: {feature_importance[0]:.4f}")
print(f"Площадь: {feature_importance[1]:.4f}")
print(f"Компактность: {feature_importance[2]:.4f}")
print(f"N вершин: {feature_importance[3]:.4f}")

Практический пример: использование количества вершин

Для идеальных фигур это почти детерминированное правило:

def classify_shape_by_vertices(mask):
    """Простая классификация по количеству углов"""
    
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if len(contours) == 0:
        return None
    
    cnt = contours[0]
    perimeter = cv2.arcLength(cnt, True)
    epsilon = 0.02 * perimeter
    approx = cv2.approxPolyDP(cnt, epsilon, True)
    n_vertices = len(approx)
    
    # Правила классификации
    if n_vertices == 4:
        return "square"
    elif n_vertices >= 8:  # Круг аппроксимируется многоугольником с много сторон
        return "circle"
    elif n_vertices == 3:
        return "triangle"
    else:
        return "unknown"

# Тестирование
mask = cv2.imread("test_image.png", cv2.IMREAD_GRAYSCALE)
predicted = classify_shape_by_vertices(mask)
print(f"Predicted: {predicted}")

Подход 2: Deep Learning с CNN

Для более сложных случаев (шум, деформация) используйте нейросети:

import tensorflow as tf
from tensorflow import keras
import numpy as np
from sklearn.model_selection import train_test_split

# Подготовка данных
X = np.load("images.npy")  # (N, 224, 224, 1) - чёрно-белые изображения
y = np.load("labels.npy")  # (N,) - 0: square, 1: circle, 2: triangle

# Нормализация
X = X / 255.0

# One-hot encoding
y_categorical = keras.utils.to_categorical(y, num_classes=3)

X_train, X_test, y_train, y_test = train_test_split(
    X, y_categorical, test_size=0.2, random_state=42
)

# Построение CNN
model = keras.Sequential([
    keras.layers.Conv2D(32, (3, 3), activation="relu", input_shape=(224, 224, 1)),
    keras.layers.MaxPooling2D((2, 2)),
    
    keras.layers.Conv2D(64, (3, 3), activation="relu"),
    keras.layers.MaxPooling2D((2, 2)),
    
    keras.layers.Conv2D(64, (3, 3), activation="relu"),
    
    keras.layers.Flatten(),
    keras.layers.Dense(64, activation="relu"),
    keras.layers.Dropout(0.5),
    keras.layers.Dense(3, activation="softmax")
])

model.compile(
    optimizer="adam",
    loss="categorical_crossentropy",
    metrics=["accuracy"]
)

# Обучение
history = model.fit(
    X_train, y_train,
    epochs=20,
    batch_size=32,
    validation_split=0.2,
    verbose=1
)

# Оценка
test_loss, test_accuracy = model.evaluate(X_test, y_test)
print(f"Test Accuracy: {test_accuracy:.4f}")

# Предсказание
y_pred_proba = model.predict(X_test)
y_pred = np.argmax(y_pred_proba, axis=1)

print(f"\nПримеры предсказаний:\n")
for i in range(5):
    shapes = ["square", "circle", "triangle"]
    print(f"Истина: {shapes[np.argmax(y_test[i])]}, "
          f"Предсказание: {shapes[y_pred[i]]}, "
          f"Уверенность: {y_pred_proba[i].max():.4f}")

Подход 3: Transfer Learning (эффективнее)

Для маленьких датасетов используйте предобученную модель:

from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.preprocessing import image

# Загружаем предобученную модель (без top слоёв)
base_model = MobileNetV2(
    input_shape=(224, 224, 3),  # Если нужно 3 канала
    include_top=False,
    weights="imagenet"
)

# Замораживаем веса базовой модели
base_model.trainable = False

# Добавляем свои слои
model = keras.Sequential([
    keras.layers.Input(shape=(224, 224, 3)),
    keras.layers.Lambda(lambda x: tf.image.grayscale_to_rgb(x[:, :, :, :1])),  # Если бинарное
    base_model,
    keras.layers.GlobalAveragePooling2D(),
    keras.layers.Dense(128, activation="relu"),
    keras.layers.Dropout(0.5),
    keras.layers.Dense(3, activation="softmax")
])

model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=1e-4),
    loss="categorical_crossentropy",
    metrics=["accuracy"]
)

model.fit(X_train, y_train, epochs=10, batch_size=32, validation_split=0.2)

Полный пайплайн с аугментацией

from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Аугментация для обучения
train_datagen = ImageDataGenerator(
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    zoom_range=0.2,
    fill_mode="nearest"
)

test_datagen = ImageDataGenerator()  # Только нормализация

# Обучение с аугментацией
batch_size = 32
epochs = 20

model.fit(
    train_datagen.flow(X_train, y_train, batch_size=batch_size),
    steps_per_epoch=len(X_train) // batch_size,
    epochs=epochs,
    validation_data=(X_test, y_test),
    verbose=1
)

Выбор подхода

Мой рекомендуемый подход

# Для вашего случая (геометрические фигуры на маске):

1. Сначала попробуйте классический подход (признаки + RF)
   - Вероятно, даст 95%+ accuracy
   - Быстро обучается
   - Легко интерпретировать

2. Если нужна лучшая точность:
   - Используйте CNN или Transfer Learning
   - Добавьте Data Augmentation
   - Настройте гиперпараметры

3. Для production:
   - Выбирайте скорость предсказания
   - Классический подход: ~1 ms на изображение
   - CNN: ~10-50 ms в зависимости от модели

Проверка качества

from sklearn.metrics import classification_report, confusion_matrix

y_true = [0, 1, 2, 0, 1, 2]
y_pred = [0, 1, 2, 0, 2, 2]

print(classification_report(y_true, y_pred, 
                          target_names=["square", "circle", "triangle"]))

print(f"\nConfusion Matrix:\n{confusion_matrix(y_true, y_pred)}")

Для вашей задачи я рекомендую начать с классических признаков — это 80/20 solution, который даст хороший результат быстро и с минимальным кодом.