from datasets import Dataset, Image as DatasetsImage
import os
from PIL import Image
import torch
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
from torchvision import transforms
import numpy as np

# Fonction build_dataset (inchangée)
def build_dataset(base_path):
    image_paths = []
    labels = []
    categories = []

    for category in os.listdir(base_path):
        category_path = os.path.join(base_path, category)
        if not os.path.isdir(category_path):
            continue

        category_images_path = os.path.join(category_path, "Data", "Images")

        # Images normales
        normal_path = os.path.join(category_images_path, "Normal")
        if os.path.exists(normal_path):
            for img in os.listdir(normal_path):
                image_paths.append(os.path.join(normal_path, img))
                labels.append(0)  # Normal
                categories.append(category)

        # Images anomalies
        anomaly_path = os.path.join(category_images_path, "Anomaly")
        if os.path.exists(anomaly_path):
            for img in os.listdir(anomaly_path):
                image_paths.append(os.path.join(anomaly_path, img))
                labels.append(1)  # Anomalie
                categories.append(category)

    return Dataset.from_dict({
        "image_path": image_paths,
        "label": labels,
        "category": categories
    })

feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
model.to(device)
model.eval()

# Fonction load_and_preprocess_image (inchangée)
def load_and_preprocess_image(img_path):
    """
    Charge et prétraite une image en utilisant le feature extractor.
    """
    image = Image.open(img_path).convert("RGB")
    inputs = feature_extractor(images=image, return_tensors="pt")
    return inputs["pixel_values"]

# Fonction extract_features (inchangée)
def extract_features(image_paths):
    """
    Extrait des caractéristiques pour chaque image de la liste image_paths.
    Ici, on utilise les logits du modèle comme représentation.
    """
    features = []
    with torch.no_grad():
        for img_path in image_paths:
            image_tensor = load_and_preprocess_image(img_path).to(device)
            outputs = model(pixel_values=image_tensor)
            # Dans ce cas, nous utilisons les logits comme vecteur de caractéristiques.
            feature_vector = outputs.logits.cpu().numpy()
            features.append(feature_vector)
    return np.vstack(features)