Spaces:
Running
Running
| 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) |