Klette_Computer_Vision_2019/part10/python/viola_jones_detector.py - main - pentryyy/ComputerVision

Klette_Computer_Vision_2019/part10/python/viola_jones_detector.py
581 строка23 KB
SergDSV0
part10
24 ноя 2025, 22:33
24 ноя 2025, 22:337a320ae
import os
import numpy as np
from PIL import Image
import pickle
from tqdm import tqdm
from collections import namedtuple
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import random

# ----------------------------------------
# Интегральные изображения и вспом.
# ----------------------------------------

def integral_image(img: np.ndarray) -> np.ndarray:
    """Возвращает интегральное изображение (int64) размера HxW"""
    # img: HxW grayscale (uint8 или float)
    ii = np.cumsum(np.cumsum(img.astype(np.int64), axis=0), axis=1)
    return ii


def rect_sum(ii: np.ndarray, x: int, y: int, w: int, h: int) -> int:
    """Сумма значений в прямоугольнике (x,y) ширина w, высота h, используя ii.
    (x,y) — левый верхний пиксель, 0-based."""
    x1, y1 = x, y
    x2, y2 = x + w - 1, y + h - 1
    A = ii[y1 - 1, x1 - 1] if (x1 > 0 and y1 > 0) else 0
    B = ii[y1 - 1, x2] if (y1 > 0) else 0 
    C = ii[y2, x1 - 1] if (x1 > 0) else 0
    D = ii[y2, x2]
    return int(D - B - C + A)

# ----------------------------------------
# Описание и генерация Хаар-признаков
# ----------------------------------------

HaarFeature = namedtuple('HaarFeature', ['type', 'x', 'y', 'w', 'h'])
# type: 'two-vertical', 'two-horizontal', 'three-horizontal', 'three-vertical'


def generate_haar_features(window_w: int, window_h: int, min_w=1, min_h=1):
    """Генерирует список Хаар-признаков для заданного размера окна.
    Параметры min_w/min_h задают минимальные размеры базовых зон.
    Для производительности можно уменьшить плотность (шаги > 1)."""
    feats = []
    # two-vertical (left vs right)
    for w in range(min_w*2, window_w+1, min_w):
        for h in range(min_h, window_h+1, min_h):
            if w % 2 != 0:
                continue
            for x in range(0, window_w - w + 1, min_w):
                for y in range(0, window_h - h + 1, min_h):
                    feats.append(HaarFeature('two-vertical', x, y, w, h))
    # two-horizontal (top vs bottom)
    for w in range(min_w, window_w+1, min_w):
        for h in range(min_h*2, window_h+1, min_h):
            if h % 2 != 0:
                continue
            for x in range(0, window_w - w + 1, min_w):
                for y in range(0, window_h - h + 1, min_h):
                    feats.append(HaarFeature('two-horizontal', x, y, w, h))
    # three-horizontal (left, center, right)
    for w in range(min_w*3, window_w+1, min_w):
        for h in range(min_h, window_h+1, min_h):
            if w % 3 != 0:
                continue
            for x in range(0, window_w - w + 1, min_w):
                for y in range(0, window_h - h + 1, min_h):
                    feats.append(HaarFeature('three-horizontal', x, y, w, h))
    # three-vertical (top, mid, bottom)
    for w in range(min_w, window_w+1, min_w):
        for h in range(min_h*3, window_h+1, min_h):
            if h % 3 != 0:
                continue
            for x in range(0, window_w - w + 1, min_w):
                for y in range(0, window_h - h + 1, min_h):
                    feats.append(HaarFeature('three-vertical', x, y, w, h))

    return feats


def visualize_haar_features(features, window_size=(24, 24), num_features=12):
    """Визуализация случайных Хаар-признаков"""
    fig, axes = plt.subplots(3, 4, figsize=(15, 12))
    axes = axes.ravel()
    
    # Выбираем случайные признаки для визуализации
    selected_features = random.sample(features, min(num_features, len(features)))
    
    for i, feat in enumerate(selected_features):
        if i >= len(axes):
            break
            
        # Создаем пустое изображение
        img = np.ones((window_size[1], window_size[0]))
        
        # Рисуем признак
        x, y, w, h = feat.x, feat.y, feat.w, feat.h
        
        if feat.type == 'two-vertical':
            half = w // 2
            # Левая часть - темная
            img[y:y+h, x:x+half] = 0
            # Правая часть - светлая
            img[y:y+h, x+half:x+w] = 0.5
            
        elif feat.type == 'two-horizontal':
            half = h // 2
            # Верхняя часть - темная
            img[y:y+half, x:x+w] = 0
            # Нижняя часть - светлая
            img[y+half:y+h, x:x+w] = 0.5
            
        elif feat.type == 'three-horizontal':
            third = w // 3
            # Левая - темная
            img[y:y+h, x:x+third] = 0
            # Центр - светлая
            img[y:y+h, x+third:x+2*third] = 0.5
            # Правая - темная
            img[y:y+h, x+2*third:x+w] = 0
            
        elif feat.type == 'three-vertical':
            third = h // 3
            # Верхняя - темная
            img[y:y+third, x:x+w] = 0
            # Центр - светлая
            img[y+third:y+2*third, x:x+w] = 0.5
            # Нижняя - темная
            img[y+2*third:y+h, x:x+w] = 0
        
        axes[i].imshow(img, cmap='gray', vmin=0, vmax=1)
        axes[i].set_title(f'{feat.type}\n({w}x{h}) at ({x},{y})')
        axes[i].axis('off')
    
    plt.tight_layout()
    plt.show()


def feature_value(ii: np.ndarray, feat: HaarFeature) -> int:
    """Вычисляет значение одного Хаар-признака для данного интегрального изображения окна.
    Ожидается, что ii относится к окну (т.е. интеграл по тому же окну).
    Для эффективности при работе со множеством окон нужно применять векторизацию."""
    x, y, w, h = feat.x, feat.y, feat.w, feat.h
    if feat.type == 'two-vertical':
        half = w // 2
        left = rect_sum(ii, x, y, half, h)
        right = rect_sum(ii, x + half, y, half, h)
        return left - right
    elif feat.type == 'two-horizontal':
        half = h // 2
        top = rect_sum(ii, x, y, w, half)
        bot = rect_sum(ii, x, y + half, w, half)
        return top - bot
    elif feat.type == 'three-horizontal':
        third = w // 3
        left = rect_sum(ii, x, y, third, h)
        mid = rect_sum(ii, x + third, y, third, h)
        right = rect_sum(ii, x + 2*third, y, third, h)
        return left - mid + right
    elif feat.type == 'three-vertical':
        third = h // 3
        top = rect_sum(ii, x, y, w, third)
        mid = rect_sum(ii, x, y + third, w, third)
        bot = rect_sum(ii, x, y + 2*third, w, third)
        return top - mid + bot
    else:
        raise ValueError('Unknown feature type')

# ----------------------------------------
# Векторизированная матрица признаков
# ----------------------------------------

def compute_feature_matrix(integral_images, features):
    """Возвращает матрицу размера (N_images x N_features) значений признаков.
    integral_images: list/array интегральных изображений (все одного размера).
    features: список HaarFeature.
    WARNING: может потребовать много памяти для больших наборов.
    """
    N = len(integral_images)
    F = len(features)
    X = np.zeros((N, F), dtype=np.int32)
    for j, ii in enumerate(integral_images):
        row = np.zeros(F, dtype=np.int32)
        for k, f in enumerate(features):
            row[k] = feature_value(ii, f)
        X[j] = row
    return X

# ----------------------------------------
# AdaBoost: слабые классификаторы — пороги по одной признаковой координате
# ----------------------------------------

class WeakClassifier:
    def __init__(self, feature_index, threshold, polarity):
        self.feature_index = feature_index
        self.threshold = threshold
        self.polarity = polarity  # 1 or -1

    def predict_array(self, X):
        # X: NxF
        feat = X[:, self.feature_index]
        preds = np.ones(feat.shape, dtype=int)
        if self.polarity == 1:
            preds[feat < self.threshold] = -1
        else:
            preds[feat >= self.threshold] = -1
        return preds

    def predict(self, x_row):
        v = x_row[self.feature_index]
        if self.polarity == 1:
            return 1 if v >= self.threshold else -1
        else:
            return 1 if v < self.threshold else -1


class AdaBoostStump:
    def __init__(self, T=20):
        self.T = T
        self.alphas = []
        self.clfs = []
        self.errors = []

    def fit(self, X, y):
        # X: NxF, y: N (labels in {1,-1})
        N, F = X.shape
        w = np.ones(N) / N
        for t in range(self.T):
            # Поиск лучшего слабого классификатора (feature + threshold + polarity)
            best_err = np.inf
            best_clf = None
            # перебираем признаки
            for j in range(F):
                feat_vals = X[:, j]
                thresholds = np.unique(feat_vals)
                # Для ускорения можно брать квантильные пороги
                # но здесь берем уникальные значения (медленно при больших данных)
                for thr in thresholds:
                    for polarity in (1, -1):
                        # предсказания
                        if polarity == 1:
                            preds = np.where(feat_vals >= thr, 1, -1)
                        else:
                            preds = np.where(feat_vals < thr, 1, -1)
                        err = np.sum(w * (preds != y))
                        if err < best_err:
                            best_err = err
                            best_clf = WeakClassifier(j, thr, polarity)
            # Если ошибка >= 0.5 — останавливаемся
            if best_err >= 0.5:
                break
            # вычисляем alpha
            eps = 1e-12
            alpha = 0.5 * np.log((1 - best_err + eps) / (best_err + eps))
            preds = best_clf.predict_array(X)
            # обновляем веса
            w = w * np.exp(-alpha * y * preds)
            w /= np.sum(w)
            # сохраняем
            self.alphas.append(alpha)
            self.clfs.append(best_clf)
            self.errors.append(best_err)
            print(f"Iter {t+1}: feature={best_clf.feature_index}, thr={best_clf.threshold}, pol={best_clf.polarity}, err={best_err:.4f}, alpha={alpha:.4f}")
        return self

    def predict_scores(self, X):
        if not self.clfs:
            return np.zeros(X.shape[0])
        S = np.zeros(X.shape[0])
        for a, clf in zip(self.alphas, self.clfs):
            S += a * clf.predict_array(X)
        return S

    def predict(self, X):
        return np.sign(self.predict_scores(X)).astype(int)

# ----------------------------------------
# Non-Maximum Suppression (NMS)
# ----------------------------------------

def non_max_suppression(boxes, scores, iou_threshold=0.5):
    # boxes: Nx4 [x,y,w,h]
    if len(boxes) == 0:
        return []
    boxes = np.array(boxes, dtype=float)
    scores = np.array(scores, dtype=float)
    x1 = boxes[:,0]
    y1 = boxes[:,1]
    x2 = x1 + boxes[:,2]
    y2 = y1 + boxes[:,3]
    areas = (x2 - x1) * (y2 - y1)
    order = scores.argsort()[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        w = np.maximum(0.0, xx2 - xx1)
        h = np.maximum(0.0, yy2 - yy1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr <= iou_threshold)[0]
        order = order[inds + 1]
    return keep

# ----------------------------------------
# Обучение/детекция — утилиты
# ----------------------------------------

def load_images_as_gray(folder, target_size=None):
    files = [os.path.join(folder, f) for f in os.listdir(folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))]
    imgs = []
    for p in files:
        im = Image.open(p).convert('L')
        if target_size is not None:
            im = im.resize(target_size, Image.BILINEAR)
        imgs.append(np.array(im, dtype=np.uint8))
    return imgs


def windows_from_image(img, win_w, win_h, step=8, scale=1.0):
    # генерируем (x,y,w,h,window_image) в списке
    H, W = img.shape
    res = []
    for y in range(0, H - win_h + 1, step):
        for x in range(0, W - win_w + 1, step):
            window = img[y:y+win_h, x:x+win_w]
            res.append((x, y, win_w, win_h, window))
    return res

# ----------------------------------------
# Визуализация процесса обучения
# ----------------------------------------

def visualize_training_process(clf, features, window_size=(24, 24)):
    """Визуализация лучших признаков, найденных во время обучения"""
    if not clf.clfs:
        print("Нет обученных классификаторов для визуализации")
        return
    
    num_clfs = len(clf.clfs)
    cols = 4
    rows = min(3, (num_clfs + cols - 1) // cols)
    
    fig, axes = plt.subplots(rows, cols, figsize=(15, 4 * rows))
    if rows == 1:
        axes = [axes] if cols == 1 else axes
    else:
        axes = axes.ravel()
    
    for i in range(min(num_clfs, len(axes))):
        clf_i = clf.clfs[i]
        feat = features[clf_i.feature_index]
        alpha = clf.alphas[i]
        error = clf.errors[i]
        
        # Создаем визуализацию признака
        img = np.ones((window_size[1], window_size[0]))
        x, y, w, h = feat.x, feat.y, feat.w, feat.h
        
        if feat.type == 'two-vertical':
            half = w // 2
            img[y:y+h, x:x+half] = 0
            img[y:y+h, x+half:x+w] = 0.5
        elif feat.type == 'two-horizontal':
            half = h // 2
            img[y:y+half, x:x+w] = 0
            img[y+half:y+h, x:x+w] = 0.5
        elif feat.type == 'three-horizontal':
            third = w // 3
            img[y:y+h, x:x+third] = 0
            img[y:y+h, x+third:x+2*third] = 0.5
            img[y:y+h, x+2*third:x+w] = 0
        elif feat.type == 'three-vertical':
            third = h // 3
            img[y:y+third, x:x+w] = 0
            img[y+third:y+2*third, x:x+w] = 0.5
            img[y+2*third:y+h, x:x+w] = 0
        
        axes[i].imshow(img, cmap='gray', vmin=0, vmax=1)
        axes[i].set_title(f'#{i+1}: {feat.type}\nα={alpha:.3f}, err={error:.3f}\npol={clf_i.polarity}')
        axes[i].axis('off')
    
    # Скрываем пустые subplots
    for i in range(min(num_clfs, len(axes)), len(axes)):
        axes[i].axis('off')
    
    plt.tight_layout()
    plt.show()
    
    # Визуализация ошибок обучения
    if clf.errors:
        plt.figure(figsize=(10, 4))
        plt.subplot(1, 2, 1)
        plt.plot(range(1, len(clf.errors) + 1), clf.errors, 'bo-')
        plt.xlabel('Итерация')
        plt.ylabel('Ошибка')
        plt.title('Ошибки слабых классификаторов')
        plt.grid(True)
        
        plt.subplot(1, 2, 2)
        plt.plot(range(1, len(clf.alphas) + 1), clf.alphas, 'ro-')
        plt.xlabel('Итерация')
        plt.ylabel('Вес (α)')
        plt.title('Веса классификаторов')
        plt.grid(True)
        
        plt.tight_layout()
        plt.show()

def visualize_detections(image, detections, scores=None, original_boxes=None):
    """Визуализация детекций на изображении"""
    fig, axes = plt.subplots(1, 2, figsize=(15, 6))
    
    # Оригинальное изображение с детекциями после NMS
    axes[0].imshow(image, cmap='gray')
    axes[0].set_title('Детекции после NMS')
    
    for i, det in enumerate(detections):
        x, y, w, h = det
        rect = patches.Rectangle((x, y), w, h, linewidth=2, edgecolor='r', facecolor='none')
        axes[0].add_patch(rect)
        if scores is not None and i < len(scores):
            axes[0].text(x, y-5, f'{scores[i]:.2f}', color='r', fontsize=8)
    
    # Все кандидаты до NMS (если предоставлены)
    if original_boxes is not None:
        axes[1].imshow(image, cmap='gray')
        axes[1].set_title('Все кандидаты до NMS')
        for det in original_boxes:
            x, y, w, h = det
            rect = patches.Rectangle((x, y), w, h, linewidth=1, edgecolor='b', facecolor='none', alpha=0.3)
            axes[1].add_patch(rect)
    
    for ax in axes:
        ax.axis('off')
    
    plt.tight_layout()
    plt.show()

# ----------------------------------------
# Демонстрация работы в Colab/Jupyter
# ----------------------------------------

def train_demo():
    """Демонстрация обучения на маленьком наборе данных"""
    print("=== Демонстрация обучения Viola-Jones ===")
    
    # Создаем синтетические данные для демонстрации
    WIN_W, WIN_H = 24, 24  # меньший размер для демонстрации
    
    # Создаем простые "лица" (темная верхняя часть, светлая нижняя)
    print("Создание синтетических данных...")
    pos_imgs = []
    for i in range(50):
        img = np.random.randint(100, 200, (WIN_H, WIN_W), dtype=np.uint8)
        # Создаем простой паттерн "глаза" (темные области сверху)
        img[:WIN_H//3, :] = np.random.randint(0, 50, (WIN_H//3, WIN_W))
        pos_imgs.append(img)
    
    # Создаем негативные изображения (случайный шум)
    neg_imgs = [np.random.randint(0, 255, (WIN_H, WIN_W), dtype=np.uint8) for _ in range(50)]
    
    # Визуализация примеров данных
    fig, axes = plt.subplots(2, 4, figsize=(12, 6))
    for i in range(4):
        axes[0, i].imshow(pos_imgs[i], cmap='gray')
        axes[0, i].set_title(f'Положительный пример {i+1}')
        axes[0, i].axis('off')
        
        axes[1, i].imshow(neg_imgs[i], cmap='gray')
        axes[1, i].set_title(f'Отрицательный пример {i+1}')
        axes[1, i].axis('off')
    
    plt.tight_layout()
    plt.show()
    
    all_imgs = pos_imgs + neg_imgs
    labels = np.array([1]*len(pos_imgs) + [-1]*len(neg_imgs))
    
    print("Вычисление интегральных изображений...")
    iis = [integral_image(im) for im in all_imgs]
    
    print("Генерация признаков Хаара...")
    feats = generate_haar_features(WIN_W, WIN_H, min_w=4, min_h=4)
    print(f"Сгенерировано {len(feats)} признаков")
    
    # Визуализация случайных признаков Хаара
    visualize_haar_features(feats, (WIN_W, WIN_H))
    
    print("Вычисление матрицы признаков...")
    X = compute_feature_matrix(iis, feats)
    
    print("Обучение AdaBoost...")
    clf = AdaBoostStump(T=10)  # меньше итераций для демонстрации
    clf.fit(X, labels)
    
    # Визуализация процесса обучения
    visualize_training_process(clf, feats, (WIN_W, WIN_H))
    
    # Сохраняем модель
    model = {
        'window': (WIN_W, WIN_H),
        'features': feats,
        'clf': clf
    }
    
    with open('vj_demo_model.pkl', 'wb') as f:
        pickle.dump(model, f)
    
    print("Модель сохранена как 'vj_demo_model.pkl'")
    return model

def detect_demo(model_path='vj_demo_model.pkl'):
    """Демонстрация детекции"""
    print("=== Демонстрация детекции ===")
    
    if not os.path.exists(model_path):
        print("Модель не найдена. Сначала запустите обучение.")
        return
    
    with open(model_path, 'rb') as f:
        model = pickle.load(f)
    
    WIN_W, WIN_H = model['window']
    feats = model['features']
    clf = model['clf']
    
    # Создаем тестовое изображение с несколькими "лицами"
    test_img = np.random.randint(100, 200, (300, 300), dtype=np.uint8)
    
    # Добавляем несколько "лиц" в разных местах
    faces_positions = [(50, 50), (150, 80), (80, 180)]
    for face_x, face_y in faces_positions:
        test_img[face_y:face_y+WIN_H, face_x:face_x+WIN_W] = np.random.randint(100, 200, (WIN_H, WIN_W))
        test_img[face_y:face_y+WIN_H//3, face_x:face_x+WIN_W] = np.random.randint(0, 50, (WIN_H//3, WIN_W))
    
    print("Поиск объектов...")
    cand_boxes = []
    cand_scores = []
    
    for (x, y, w, h, window) in windows_from_image(test_img, WIN_W, WIN_H, step=12):
        ii = integral_image(window)
        vals = np.array([feature_value(ii, f) for f in feats])
        score = clf.predict_scores(vals.reshape(1, -1))[0]
        if score > 0:
            cand_boxes.append([x, y, w, h])
            cand_scores.append(score)
    
    keep = non_max_suppression(cand_boxes, cand_scores, iou_threshold=0.3)
    detections = [cand_boxes[i] for i in keep]
    detection_scores = [cand_scores[i] for i in keep]
    
    print(f"Найдено детекций: {len(detections)}")
    for i, det in enumerate(detections):
        print(f"Детекция {i+1}: x={det[0]}, y={det[1]}, w={det[2]}, h={det[3]}, score={detection_scores[i]:.2f}")
    
    # Визуализация детекций
    visualize_detections(test_img, detections, detection_scores, cand_boxes)
    
    return detections, detection_scores, cand_boxes

# ----------------------------------------
# Основной код для Colab/Jupyter
# ----------------------------------------

if __name__ == '__main__':
    # В Colab/Jupyter просто вызываем функции напрямую
    print("Viola-Jones реализация для Colab/Jupyter")
    
    # Раскомментируйте нужную демонстрацию:
    
    # Демонстрация обучения
    model = train_demo()
    
    # Демонстрация детекции
    # detections, scores, all_boxes = detect_demo()