Spaces:
Running
Running
File size: 2,665 Bytes
6462c1c fe6cf06 2a161ab 3be26bb fe6cf06 6462c1c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | 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) |