Klette_Computer_Vision_2019/part5/codes/segImgPython.py - main - Leylin/ComputerVision

Klette_Computer_Vision_2019/part5/codes/segImgPython.py
75 строк3 KB
Art_Lun
Клетте глава 5
21 дек 2025, 01:55
21 дек 2025, 01:555d76b5d
import cv2
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt

def segment_image_python(image_path, num_clusters=3):
    # Загрузка изображения
    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # Изменение формы для кластеризации
    pixels = image_rgb.reshape((-1, 3))
    pixels = np.float32(pixels)
    
    # Применение K-means кластеризации
    kmeans = KMeans(n_clusters=num_clusters, random_state=42, n_init=10)
    labels = kmeans.fit_predict(pixels)
    
    # Создание сегментированного изображения
    centers = np.uint8(kmeans.cluster_centers_)
    segmented_data = centers[labels]
    segmented_image = segmented_data.reshape(image_rgb.shape)
    
    # Отображение результатов
    fig, axes = plt.subplots(1, 2, figsize=(12, 6))
    axes[0].imshow(image_rgb)
    axes[0].set_title('Original Image')
    axes[0].axis('off')
    
    axes[1].imshow(segmented_image)
    axes[1].set_title(f'Segmented Image (K={num_clusters})')
    axes[1].axis('off')
    
    plt.tight_layout()
    plt.show()
    
    return segmented_image

# Альтернативный метод: пороговая сегментация
def threshold_segmentation_python(image_path):
    # Загрузка изображения в оттенках серого
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    
    # Автоматическое определение порога (метод Оцу)
    _, binary = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    
    # Поиск контуров
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Создание цветного изображения с контурами
    color_image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    cv2.drawContours(color_image, contours, -1, (0, 255, 0), 2)
    
    # Отображение
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    axes[0].imshow(image, cmap='gray')
    axes[0].set_title('Original')
    axes[0].axis('off')
    
    axes[1].imshow(binary, cmap='gray')
    axes[1].set_title('Binary')
    axes[1].axis('off')
    
    axes[2].imshow(color_image)
    axes[2].set_title('Contours')
    axes[2].axis('off')
    
    plt.tight_layout()
    plt.show()
    
    return binary

# Использование
# segmented = segment_image_python('path/to/your/image.jpg', num_clusters=4)
# binary_seg = threshold_segmentation_python('path/to/your/image.jpg')