segmentation.py

import cv2
import numpy as np
from sklearn.cluster import KMeans

def segment_image(image_path, k=2, resize_shape=(256, 256)):
    # Charger l'image
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convertir en RGB
    img = cv2.resize(img, resize_shape)  # Redimensionner l'image

    # Normaliser l'image
    img_normalized = img / 255.0

    # Réduire le bruit
    img_denoised = cv2.GaussianBlur(img_normalized, (5, 5), 0)

    # Aplatir l'image pour l'appliquer à K-means
    img_flat = img_denoised.reshape(-1, 3)

    # Appliquer K-means
    kmeans = KMeans(n_clusters=k, random_state=0).fit(img_flat)
    labels = kmeans.labels_

    # Réorganiser les labels en image
    segmented_img = labels.reshape(resize_shape)

    # Convertir l'image segmentée en niveaux de gris
    segmented_img_gray = (segmented_img * 255).astype(np.uint8)

    return img, segmented_img_gray