739 lines
29 KiB
Python
739 lines
29 KiB
Python
import os
|
||
import json
|
||
import tempfile
|
||
import numpy as np
|
||
import torch
|
||
import matplotlib.pyplot as plt
|
||
import segmentation_models_pytorch as smp
|
||
import mlflow
|
||
import re
|
||
import ast
|
||
import time
|
||
|
||
# Visualisation
|
||
import matplotlib.patches as mpatches
|
||
from matplotlib.colors import ListedColormap
|
||
|
||
# Preprocessing/Images
|
||
import albumentations as A
|
||
from albumentations.pytorch import ToTensorV2
|
||
|
||
|
||
def load_cityscapes_config(config_path="../cityscapes_config.json", verbose=True):
|
||
"""
|
||
Charge la configuration Cityscapes depuis un fichier JSON.
|
||
INCHANGÉ - Compatible PyTorch
|
||
"""
|
||
with open(config_path, "r", encoding="utf-8") as f:
|
||
config = json.load(f)
|
||
|
||
class_groups = config["class_groups"]
|
||
cityscapes_mapping = {int(k): v for k, v in config["cityscapes_mapping"].items()}
|
||
id_to_label = {int(k): v for k, v in config["id_to_label"].items()}
|
||
group_colors = np.array(config["group_colors"], dtype=np.uint8)
|
||
group_names = config["group_names"]
|
||
|
||
id_to_group = np.zeros(256, dtype=np.uint8)
|
||
for id_, group_id in cityscapes_mapping.items():
|
||
id_to_group[id_] = group_id
|
||
|
||
num_classes = len(group_names)
|
||
|
||
if verbose:
|
||
print("== Groupes de classes Cityscapes ==")
|
||
for group, classes in class_groups.items():
|
||
print(f" - {group} : {classes}")
|
||
|
||
print("\n== Mapping ID Cityscapes → Groupe ==")
|
||
for id_ in sorted(cityscapes_mapping.keys()):
|
||
label = id_to_label.get(id_, "Label inconnu")
|
||
group_id = cityscapes_mapping[id_]
|
||
print(f" ID {id_:2d} ('{label}') → groupe '{group_names[group_id]}' ({group_id})")
|
||
|
||
print("\n== Liste des groupes et couleurs associées ==")
|
||
for i, (name, color) in enumerate(zip(group_names, group_colors)):
|
||
print(f" Groupe {i}: {name} - Couleur RGB : {color}")
|
||
|
||
print(f"\nNombre total de groupes/classes : {num_classes}")
|
||
|
||
return {
|
||
"class_groups": class_groups,
|
||
"cityscapes_mapping": cityscapes_mapping,
|
||
"id_to_label": id_to_label,
|
||
"group_colors": group_colors,
|
||
"group_names": group_names,
|
||
"id_to_group": id_to_group,
|
||
"num_classes": num_classes
|
||
}
|
||
|
||
def map_mask_ids(mask, id_to_group_mapping):
|
||
"""
|
||
Applique le mapping des IDs Cityscapes vers les groupes
|
||
ADAPTÉ PYTORCH : Utilise torch au lieu de tf
|
||
"""
|
||
if isinstance(mask, np.ndarray):
|
||
mask = torch.from_numpy(mask)
|
||
|
||
if isinstance(id_to_group_mapping, np.ndarray):
|
||
if len(id_to_group_mapping) != 256:
|
||
raise ValueError(f"id_to_group_mapping doit avoir 256 éléments, reçu: {len(id_to_group_mapping)}")
|
||
mapping_tensor = torch.tensor(id_to_group_mapping, dtype=torch.long)
|
||
else:
|
||
mapping_tensor = torch.tensor(list(id_to_group_mapping), dtype=torch.long)
|
||
|
||
# S'assurer que les valeurs du masque sont dans la plage valide [0, 255]
|
||
mask_clipped = torch.clamp(mask, 0, 255)
|
||
mapped_mask = mapping_tensor[mask_clipped.long()]
|
||
return mapped_mask
|
||
|
||
def colorize_mask(mask, group_colors):
|
||
"""
|
||
Applique la colorisation du masque avec les couleurs définies
|
||
INCHANGÉ - Compatible PyTorch (prend numpy en entrée)
|
||
"""
|
||
if isinstance(mask, torch.Tensor):
|
||
mask = mask.numpy()
|
||
|
||
colored_mask = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
|
||
num_classes = len(group_colors)
|
||
|
||
for group_id in range(num_classes):
|
||
colored_mask[mask == group_id] = group_colors[group_id]
|
||
|
||
return colored_mask
|
||
|
||
# def get_preprocessing_fn(encoder_name=None):
|
||
# """
|
||
# Retourne la fonction de prétraitement adaptée au backbone utilisé.
|
||
# ADAPTÉ PYTORCH : utilise smp.encoders.get_preprocessing_params
|
||
# """
|
||
# if encoder_name is None:
|
||
# print("ℹ️ Modèle vanilla - normalisation simple")
|
||
# return lambda x: x / 255.0
|
||
# else:
|
||
# print(f"ℹ️ Préprocessing du backbone {encoder_name}")
|
||
# try:
|
||
# # PyTorch segmentation_models utilise get_preprocessing_params
|
||
# preprocess_params = smp.encoders.get_preprocessing_params(encoder_name)
|
||
# mean = preprocess_params['mean']
|
||
# std = preprocess_params['std']
|
||
|
||
# def preprocess_fn(x):
|
||
# # x doit être dans [0, 1] puis normalisé
|
||
# x = x / 255.0 if x.max() > 1.0 else x
|
||
# # Normalisation ImageNet ou spécifique au backbone
|
||
# if isinstance(x, torch.Tensor):
|
||
# for i in range(3): # RGB
|
||
# x[:, i, :, :] = (x[:, i, :, :] - mean[i]) / std[i]
|
||
# else: # numpy
|
||
# for i in range(3):
|
||
# x[:, :, i] = (x[:, :, i] - mean[i]) / std[i]
|
||
# return x
|
||
|
||
# return preprocess_fn
|
||
# except Exception as e:
|
||
# print(f"⚠️ Erreur preprocessing {encoder_name}: {e}, fallback simple")
|
||
# return lambda x: x / 255.0
|
||
|
||
def get_preprocessing_fn(img_size=(512, 512), encoder=None):
|
||
"""
|
||
Retourne une fonction qui prétraite à la fois les images et les masques
|
||
pour l'inférence, en reproduisant exactement le pipeline CityscapesDataset.
|
||
"""
|
||
|
||
if encoder is None:
|
||
mean, std = [0.0, 0.0, 0.0], [1.0, 1.0, 1.0]
|
||
else:
|
||
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
|
||
|
||
# Pipeline unifié (gérera image ET masque ensemble)
|
||
transform = A.Compose([
|
||
A.Resize(height=img_size[0], width=img_size[1]),
|
||
A.Normalize(mean=mean, std=std),
|
||
ToTensorV2()
|
||
])
|
||
|
||
def preprocess(image_np, mask_np=None):
|
||
if mask_np is not None:
|
||
# Traitement simultané image + masque
|
||
transformed = transform(image=image_np, mask=mask_np)
|
||
return transformed['image'], transformed['mask'].long()
|
||
else:
|
||
# Image seule
|
||
transformed = transform(image=image_np)
|
||
return transformed['image'], None
|
||
|
||
return preprocess
|
||
|
||
# Version simplifiée de preprocess_image_and_mask utilisant la nouvelle get_preprocessing_fn
|
||
def preprocess_image_and_mask(image_path, mask_path=None, img_size=(512, 512),
|
||
encoder_name=None, id_to_group_mapping=None):
|
||
"""
|
||
Version simplifiée utilisant get_preprocessing_fn pour la cohérence avec CityscapesDataset
|
||
"""
|
||
import numpy as np
|
||
from PIL import Image
|
||
import torch
|
||
|
||
print(f"🔄 Préprocessing de l'image: {os.path.basename(image_path)}")
|
||
|
||
# Chargement des données
|
||
image_pil = Image.open(image_path).convert('RGB')
|
||
image_np = np.array(image_pil)
|
||
img_display = np.array(image_pil.resize(img_size, Image.BILINEAR)) # Pour affichage
|
||
|
||
mask_np = None
|
||
if mask_path and os.path.exists(mask_path):
|
||
print(f"🔄 Préprocessing du masque: {os.path.basename(mask_path)}")
|
||
mask_pil = Image.open(mask_path).convert('L')
|
||
mask_np = np.array(mask_pil)
|
||
print(f" Valeurs uniques avant mapping: {np.unique(mask_np)[:10]}...")
|
||
|
||
# Utiliser notre preprocessing unifié
|
||
preprocess_fn = get_preprocessing_fn(img_size=img_size, encoder=encoder_name)
|
||
img_preprocessed, mask_preprocessed = preprocess_fn(image_np, mask_np)
|
||
|
||
# Appliquer le mapping si fourni ET si on a un masque
|
||
if mask_preprocessed is not None and id_to_group_mapping is not None:
|
||
mask_preprocessed = map_mask_ids(mask_preprocessed, id_to_group_mapping)
|
||
print(f" Valeurs uniques après mapping: {torch.unique(mask_preprocessed)}")
|
||
|
||
return img_preprocessed, img_display, mask_preprocessed
|
||
|
||
|
||
# def preprocess_image_and_mask(image_path, mask_path=None, img_size=(512, 512), # ← CHANGÉ: 224→512
|
||
# encoder_name=None, id_to_group_mapping=None):
|
||
# """
|
||
# Préprocesse une image (et optionnellement un masque) avec le même pipeline que l'entraînement
|
||
# ADAPTÉ PYTORCH : Utilise PIL + torch au lieu de tf
|
||
# """
|
||
|
||
# print(f"🔄 Préprocessing de l'image: {os.path.basename(image_path)}")
|
||
|
||
# # Chargement et preprocessing de l'image
|
||
# img_pil = Image.open(image_path).convert('RGB')
|
||
# img_original = img_pil.resize(img_size, Image.BILINEAR)
|
||
|
||
# # Conversion en tensor et normalisation
|
||
# img_tensor = transforms.ToTensor()(img_original) # [3, H, W] dans [0, 1]
|
||
|
||
# # Appliquer le preprocessing du backbone
|
||
# preprocessing_fn = get_preprocessing_fn(encoder_name)
|
||
# if encoder_name is not None:
|
||
# # Pour les backbones pré-entraînés, on applique leur normalisation
|
||
# img_preprocessed = preprocessing_fn(img_tensor.permute(1, 2, 0).numpy())
|
||
# img_preprocessed = torch.from_numpy(img_preprocessed).permute(2, 0, 1).float()
|
||
# else:
|
||
# # Simple normalisation [0, 1]
|
||
# img_preprocessed = img_tensor.float()
|
||
|
||
# # Preprocessing du masque si fourni
|
||
# mask_preprocessed = None
|
||
# if mask_path and os.path.exists(mask_path):
|
||
# print(f"🔄 Préprocessing du masque: {os.path.basename(mask_path)}")
|
||
|
||
# mask_pil = Image.open(mask_path).convert('L') # Grayscale
|
||
# mask_resized = mask_pil.resize(img_size, Image.NEAREST)
|
||
# mask_tensor = torch.from_numpy(np.array(mask_resized)).long()
|
||
|
||
# print(f" Valeurs uniques avant mapping: {torch.unique(mask_tensor)[:10]}...")
|
||
|
||
# # Appliquer le mapping si fourni
|
||
# if id_to_group_mapping is not None:
|
||
# mask_tensor = map_mask_ids(mask_tensor, id_to_group_mapping)
|
||
# print(f" Valeurs uniques après mapping: {torch.unique(mask_tensor)}")
|
||
|
||
# mask_preprocessed = mask_tensor
|
||
|
||
# # img_original en numpy pour affichage
|
||
# img_display = np.array(img_original)
|
||
|
||
# return img_preprocessed, img_display, mask_preprocessed
|
||
|
||
# def run_inference_and_visualize(image_path, model, encoder_name, img_size, mapping_config,
|
||
# mask_path=None, save_dir=None, show_stats=True):
|
||
# """
|
||
# ADAPTÉ PYTORCH : Inférence avec modèle PyTorch
|
||
# Returns: (predicted_mask, inference_time_seconds)
|
||
# """
|
||
|
||
# print(f"=== INFÉRENCE PYTORCH SUR {os.path.basename(image_path)} ===")
|
||
|
||
# img_preprocessed, img_display, mask_gt = preprocess_image_and_mask(
|
||
# image_path=image_path,
|
||
# mask_path=mask_path,
|
||
# img_size=img_size or (512, 512),
|
||
# encoder_name=encoder_name,
|
||
# id_to_group_mapping=mapping_config.get('id_to_group')
|
||
# )
|
||
|
||
# print("🔮 Exécution de l'inférence PyTorch...")
|
||
# model.eval()
|
||
# device = next(model.parameters()).device
|
||
|
||
# with torch.no_grad():
|
||
# # Ajouter dimension batch
|
||
# img_batch = img_preprocessed.unsqueeze(0).to(device) # [1, 3, H, W]
|
||
|
||
# # ⏱️ MESURE DU TEMPS D'INFÉRENCE
|
||
# start_time = time.time()
|
||
# predictions = model(img_batch) # [1, num_classes, H, W]
|
||
# end_time = time.time()
|
||
|
||
# inference_time = end_time - start_time
|
||
|
||
# predicted_mask = torch.argmax(predictions[0], dim=0).cpu().numpy() # [H, W]
|
||
|
||
# print(f"✅ Inférence terminée - Shape: {predicted_mask.shape}")
|
||
# print(f"⏱️ Temps d'inférence: {inference_time*1000:.1f} ms ({inference_time:.4f} s)")
|
||
|
||
# # Visualisation
|
||
# num_plots = 2 if mask_gt is None else 4
|
||
# fig, axes = plt.subplots(1, num_plots, figsize=(5*num_plots, 5))
|
||
# if num_plots == 1:
|
||
# axes = [axes]
|
||
|
||
# axes[0].imshow(img_display)
|
||
# axes[0].set_title("Image Originale")
|
||
# axes[0].axis('off')
|
||
|
||
# predicted_colored = colorize_mask(predicted_mask, mapping_config['group_colors'])
|
||
# axes[1].imshow(predicted_colored)
|
||
# axes[1].set_title("Masque Prédit")
|
||
# axes[1].axis('off')
|
||
|
||
# if mask_gt is not None:
|
||
# mask_gt_np = mask_gt.numpy() if isinstance(mask_gt, torch.Tensor) else mask_gt
|
||
# gt_colored = colorize_mask(mask_gt_np, mapping_config['group_colors'])
|
||
# axes[2].imshow(gt_colored)
|
||
# axes[2].set_title("Masque de Vérité")
|
||
# axes[2].axis('off')
|
||
|
||
# # Visualisation des erreurs
|
||
# diff_mask = (mask_gt_np != predicted_mask).astype(float)
|
||
# axes[3].imshow(img_display)
|
||
# axes[3].imshow(diff_mask, alpha=0.5, cmap='Reds')
|
||
# axes[3].set_title("Erreurs (en rouge)")
|
||
# axes[3].axis('off')
|
||
|
||
# colors_normalized = mapping_config['group_colors'] / 255.0
|
||
# custom_cmap = ListedColormap(colors_normalized)
|
||
|
||
# legend_elements = []
|
||
# for i, (name, color) in enumerate(zip(mapping_config['group_names'], colors_normalized)):
|
||
# legend_elements.append(mpatches.Patch(color=color, label=f'{i}: {name}'))
|
||
|
||
# fig.legend(handles=legend_elements,
|
||
# loc='center left',
|
||
# bbox_to_anchor=(1.01, 0.5),
|
||
# fontsize=9,
|
||
# frameon=True,
|
||
# fancybox=True,
|
||
# shadow=True)
|
||
|
||
# plt.tight_layout()
|
||
# fig.subplots_adjust(right=0.95)
|
||
|
||
# if save_dir:
|
||
# os.makedirs(save_dir, exist_ok=True)
|
||
# img_name = os.path.splitext(os.path.basename(image_path))[0]
|
||
# fig_path = os.path.join(save_dir, f'{img_name}_inference.png')
|
||
# fig.savefig(fig_path, dpi=150, bbox_inches='tight')
|
||
# print(f"💾 Résultat sauvegardé: {fig_path}")
|
||
|
||
# plt.show()
|
||
|
||
# # Statistiques
|
||
# if show_stats:
|
||
# print("\n📊 Statistiques du masque prédit:")
|
||
# unique_classes, counts = np.unique(predicted_mask, return_counts=True)
|
||
# for class_id, count in zip(unique_classes, counts):
|
||
# if class_id < len(mapping_config['group_names']):
|
||
# percentage = (count / predicted_mask.size) * 100
|
||
# print(f" {mapping_config['group_names'][class_id]}: {count} pixels ({percentage:.1f}%)")
|
||
|
||
# if mask_gt is not None:
|
||
# mask_gt_np = mask_gt.numpy() if isinstance(mask_gt, torch.Tensor) else mask_gt
|
||
# print("\n📊 Statistiques du masque de vérité:")
|
||
# unique_classes, counts = np.unique(mask_gt_np, return_counts=True)
|
||
# for class_id, count in zip(unique_classes, counts):
|
||
# if class_id < len(mapping_config['group_names']):
|
||
# percentage = (count / mask_gt_np.size) * 100
|
||
# print(f" {mapping_config['group_names'][class_id]}: {count} pixels ({percentage:.1f}%)")
|
||
|
||
# return predicted_mask, inference_time
|
||
|
||
|
||
# Version mise à jour de run_inference_and_visualize
|
||
def run_inference_and_visualize(image_path, model, encoder_name, img_size, mapping_config,
|
||
mask_path=None, save_dir=None, show_stats=True, show_plot=True):
|
||
"""
|
||
Version simplifiée utilisant preprocess_image_and_mask
|
||
"""
|
||
print(f"=== INFÉRENCE PYTORCH SUR {os.path.basename(image_path)} ===")
|
||
|
||
# Utiliser la version simplifiée du preprocessing
|
||
img_preprocessed, img_display, mask_gt = preprocess_image_and_mask(
|
||
image_path=image_path,
|
||
mask_path=mask_path,
|
||
img_size=img_size or (512, 512),
|
||
encoder_name=encoder_name,
|
||
id_to_group_mapping=mapping_config.get('id_to_group')
|
||
)
|
||
|
||
print("🔮 Exécution de l'inférence PyTorch...")
|
||
model.eval()
|
||
device = next(model.parameters()).device
|
||
|
||
with torch.no_grad():
|
||
# Ajouter dimension batch
|
||
img_batch = img_preprocessed.unsqueeze(0).to(device) # [1, 3, H, W]
|
||
|
||
# ⏱️ MESURE DU TEMPS D'INFÉRENCE
|
||
start_time = time.time()
|
||
predictions = model(img_batch) # [1, num_classes, H, W]
|
||
end_time = time.time()
|
||
|
||
inference_time = end_time - start_time
|
||
|
||
predicted_mask = torch.argmax(predictions[0], dim=0).cpu().numpy() # [H, W]
|
||
|
||
print(f"✅ Inférence terminée - Shape: {predicted_mask.shape}")
|
||
print(f"⏱️ Temps d'inférence: {inference_time*1000:.1f} ms ({inference_time:.4f} s)")
|
||
|
||
# [Le reste du code de visualisation reste identique]
|
||
# Visualisation
|
||
num_plots = 2 if mask_gt is None else 4
|
||
fig, axes = plt.subplots(1, num_plots, figsize=(5*num_plots, 5))
|
||
if num_plots == 1:
|
||
axes = [axes]
|
||
|
||
axes[0].imshow(img_display)
|
||
axes[0].set_title("Image Originale")
|
||
axes[0].axis('off')
|
||
|
||
predicted_colored = colorize_mask(predicted_mask, mapping_config['group_colors'])
|
||
axes[1].imshow(predicted_colored)
|
||
axes[1].set_title("Masque Prédit")
|
||
axes[1].axis('off')
|
||
|
||
if mask_gt is not None:
|
||
mask_gt_np = mask_gt.numpy() if isinstance(mask_gt, torch.Tensor) else mask_gt
|
||
gt_colored = colorize_mask(mask_gt_np, mapping_config['group_colors'])
|
||
axes[2].imshow(gt_colored)
|
||
axes[2].set_title("Masque de Vérité")
|
||
axes[2].axis('off')
|
||
|
||
# Visualisation des erreurs
|
||
diff_mask = (mask_gt_np != predicted_mask).astype(float)
|
||
axes[3].imshow(img_display)
|
||
axes[3].imshow(diff_mask, alpha=0.5, cmap='Reds')
|
||
axes[3].set_title("Erreurs (en rouge)")
|
||
axes[3].axis('off')
|
||
|
||
colors_normalized = mapping_config['group_colors'] / 255.0
|
||
custom_cmap = ListedColormap(colors_normalized)
|
||
|
||
legend_elements = []
|
||
for i, (name, color) in enumerate(zip(mapping_config['group_names'], colors_normalized)):
|
||
legend_elements.append(mpatches.Patch(color=color, label=f'{i}: {name}'))
|
||
|
||
fig.legend(handles=legend_elements,
|
||
loc='center left',
|
||
bbox_to_anchor=(1.01, 0.5),
|
||
fontsize=9,
|
||
frameon=True,
|
||
fancybox=True,
|
||
shadow=True)
|
||
|
||
plt.tight_layout()
|
||
fig.subplots_adjust(right=0.95)
|
||
|
||
if save_dir:
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
img_name = os.path.splitext(os.path.basename(image_path))[0]
|
||
fig_path = os.path.join(save_dir, f'{img_name}_inference.png')
|
||
fig.savefig(fig_path, dpi=150, bbox_inches='tight')
|
||
print(f"💾 Résultat sauvegardé: {fig_path}")
|
||
|
||
|
||
# Statistiques
|
||
if show_stats:
|
||
print("\n📊 Statistiques du masque prédit:")
|
||
unique_classes, counts = np.unique(predicted_mask, return_counts=True)
|
||
for class_id, count in zip(unique_classes, counts):
|
||
if class_id < len(mapping_config['group_names']):
|
||
percentage = (count / predicted_mask.size) * 100
|
||
print(f" {mapping_config['group_names'][class_id]}: {count} pixels ({percentage:.1f}%)")
|
||
|
||
if mask_gt is not None:
|
||
mask_gt_np = mask_gt.numpy() if isinstance(mask_gt, torch.Tensor) else mask_gt
|
||
print("\n📊 Statistiques du masque de vérité:")
|
||
unique_classes, counts = np.unique(mask_gt_np, return_counts=True)
|
||
for class_id, count in zip(unique_classes, counts):
|
||
if class_id < len(mapping_config['group_names']):
|
||
percentage = (count / mask_gt_np.size) * 100
|
||
print(f" {mapping_config['group_names'][class_id]}: {count} pixels ({percentage:.1f}%)")
|
||
|
||
if show_plot:
|
||
plt.show()
|
||
return predicted_mask, inference_time
|
||
else:
|
||
return predicted_mask, inference_time, fig
|
||
|
||
|
||
def load_model(run_id=None, experiment_name="OC Projet 9",
|
||
metric="test_mean_iou", top_n=1):
|
||
"""
|
||
Charge un modèle PyTorch depuis MLflow avec plusieurs méthodes de sélection
|
||
|
||
Args:
|
||
run_id (str, optional): Run ID spécifique (prioritaire si fourni)
|
||
experiment_name (str): Nom de l'expérience MLflow
|
||
metric (str): Métrique pour classer les runs
|
||
top_n (int): Position dans le classement (1=meilleur, 2=deuxième, etc.)
|
||
|
||
Returns:
|
||
tuple: (model, source_run_id, encoder_name, script_img_size)
|
||
"""
|
||
|
||
print(f"\n🔍 Chargement d'un modèle PyTorch depuis MLflow...")
|
||
|
||
client = mlflow.tracking.MlflowClient()
|
||
|
||
# 1. Sélection du run selon les paramètres
|
||
if run_id:
|
||
print(f"🎯 Mode run_id spécifique: {run_id}")
|
||
try:
|
||
run = mlflow.get_run(run_id)
|
||
source_run_id = run_id
|
||
print(f"✅ Run trouvé: {run.data.tags.get('mlflow.runName', 'Unknown')}")
|
||
except Exception as e:
|
||
raise RuntimeError(f"❌ Run {run_id} introuvable: {e}")
|
||
|
||
else:
|
||
print(f"🏆 Mode classement: top_{top_n} de '{metric}' dans '{experiment_name}'")
|
||
try:
|
||
# Récupérer l'expérience
|
||
experiment = client.get_experiment_by_name(experiment_name)
|
||
if experiment is None:
|
||
raise RuntimeError(f"❌ Expérience '{experiment_name}' non trouvée")
|
||
|
||
# Rechercher les runs avec la métrique, triés par ordre décroissant
|
||
filter_string = f"status = 'FINISHED' and metrics.{metric} >= 0"
|
||
order_by = f"metrics.{metric} DESC"
|
||
|
||
runs_df = mlflow.search_runs(
|
||
experiment_ids=[experiment.experiment_id],
|
||
filter_string=filter_string,
|
||
order_by=[order_by],
|
||
max_results=top_n
|
||
)
|
||
|
||
if runs_df.empty:
|
||
raise RuntimeError(f"❌ Aucun run avec '{metric}' trouvé dans '{experiment_name}'")
|
||
|
||
if len(runs_df) < top_n:
|
||
raise RuntimeError(f"❌ Seulement {len(runs_df)} runs trouvés, impossible d'obtenir le top_{top_n}")
|
||
|
||
# Prendre le run à la position top_n
|
||
selected_run = runs_df.iloc[top_n - 1]
|
||
source_run_id = selected_run['run_id']
|
||
metric_value = selected_run[f'metrics.{metric}']
|
||
run_name = selected_run.get('tags.mlflow.runName', 'Unknown')
|
||
|
||
print(f"✅ Run sélectionné (position #{top_n}):")
|
||
print(f" • Nom: {run_name}")
|
||
print(f" • Run ID: {source_run_id}")
|
||
print(f" • {metric}: {metric_value:.4f}")
|
||
|
||
# Afficher le contexte
|
||
print(f"📊 Contexte du classement:")
|
||
for i, (_, row) in enumerate(runs_df.head(min(5, len(runs_df))).iterrows()):
|
||
indicator = "👑" if i == (top_n - 1) else f"{i+1}."
|
||
name = row.get('tags.mlflow.runName', 'Unknown')[:30]
|
||
value = row[f'metrics.{metric}']
|
||
print(f" {indicator} {name}: {value:.4f}")
|
||
|
||
run = mlflow.get_run(source_run_id)
|
||
|
||
except Exception as e:
|
||
raise RuntimeError(f"❌ Erreur lors de la sélection: {e}")
|
||
|
||
# 2-5. [Reste du code identique - extraction params, détection variante, chargement...]
|
||
try:
|
||
params = run.data.params
|
||
|
||
encoder_name = params.get("encoder_name", None)
|
||
architecture = params.get("architecture", "FPN")
|
||
num_classes = int(params.get("num_classes", 8))
|
||
script_img_size_raw = params.get("script_img_size", None)
|
||
run_name = run.data.tags.get('mlflow.runName', 'Unknown')
|
||
|
||
# Parser script_img_size
|
||
script_img_size = None
|
||
if script_img_size_raw:
|
||
try:
|
||
script_img_size = ast.literal_eval(script_img_size_raw)
|
||
if isinstance(script_img_size, (tuple, list)) and len(script_img_size) == 2:
|
||
script_img_size = tuple(script_img_size)
|
||
except:
|
||
script_img_size = None
|
||
|
||
print(f"\n📋 Paramètres du modèle:")
|
||
print(f"🏗️ Architecture: {architecture}")
|
||
print(f"🔧 Encodeur: {encoder_name}")
|
||
print(f"📊 Nombre de classes: {num_classes}")
|
||
print(f"🖼️ Script img size: {script_img_size}")
|
||
|
||
except Exception as e:
|
||
raise RuntimeError(f"❌ Erreur lors de la récupération des paramètres: {e}")
|
||
|
||
# Détecter la variante SegFormer
|
||
def detect_segformer_variant_from_name(run_name, params):
|
||
segformer_pattern = r'SegFormer[_-]?b(\d+)'
|
||
match = re.search(segformer_pattern, run_name, re.IGNORECASE)
|
||
if match:
|
||
variant_num = match.group(1)
|
||
detected = f"mit_b{variant_num}"
|
||
print(f"🔍 Variante détectée: 'SegFormer_b{variant_num}' → {detected}")
|
||
return detected
|
||
|
||
if encoder_name and encoder_name.startswith("mit_b"):
|
||
print(f"🔍 Variante depuis paramètres: {encoder_name}")
|
||
return encoder_name
|
||
|
||
b_pattern = r'\bb(\d+)\b'
|
||
match = re.search(b_pattern, run_name, re.IGNORECASE)
|
||
if match:
|
||
variant_num = match.group(1)
|
||
detected = f"mit_b{variant_num}"
|
||
print(f"🔍 Pattern détecté: 'b{variant_num}' → {detected}")
|
||
return detected
|
||
|
||
print(f"🔍 Aucune variante détectée, utilisation de mit_b0")
|
||
return "mit_b0"
|
||
|
||
detected_variant = detect_segformer_variant_from_name(run_name, params)
|
||
|
||
# Télécharger et charger le modèle
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
try:
|
||
artifacts = client.list_artifacts(source_run_id)
|
||
model_artifact_candidates = [a.path for a in artifacts if a.path.endswith('.pth')]
|
||
|
||
if not model_artifact_candidates:
|
||
raise RuntimeError("❌ Aucun fichier .pth trouvé")
|
||
|
||
model_artifact = model_artifact_candidates[0]
|
||
print(f"⬇️ Téléchargement de {model_artifact}...")
|
||
|
||
local_path = mlflow.artifacts.download_artifacts(
|
||
run_id=source_run_id,
|
||
artifact_path=model_artifact,
|
||
dst_path=tmpdir
|
||
)
|
||
|
||
checkpoint = torch.load(local_path, map_location='cpu')
|
||
if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
|
||
state_dict = checkpoint['state_dict']
|
||
else:
|
||
state_dict = checkpoint
|
||
|
||
except Exception as e:
|
||
raise RuntimeError(f"❌ Erreur lors du téléchargement: {e}")
|
||
|
||
# Créer le modèle
|
||
try:
|
||
is_transformer = any('patch_embed' in key for key in state_dict.keys())
|
||
|
||
if is_transformer:
|
||
print(f"🔧 Création SegFormer ({detected_variant})")
|
||
model = smp.Segformer(encoder_name=detected_variant, classes=num_classes)
|
||
encoder_name = detected_variant
|
||
else:
|
||
print("🔧 Création modèle CNN")
|
||
if encoder_name is None or encoder_name == "None":
|
||
encoder_name = "efficientnet-b0"
|
||
|
||
if architecture.lower() == "fpn":
|
||
model = smp.FPN(encoder_name=encoder_name, classes=num_classes)
|
||
elif architecture.lower() == "unet":
|
||
model = smp.Unet(encoder_name=encoder_name, classes=num_classes)
|
||
elif architecture.lower() == "deeplabv3plus":
|
||
model = smp.DeepLabV3Plus(encoder_name=encoder_name, classes=num_classes)
|
||
else:
|
||
raise ValueError(f"Architecture {architecture} non supportée")
|
||
|
||
print("🔄 Chargement des poids...")
|
||
model.load_state_dict(state_dict, strict=True)
|
||
model.eval()
|
||
print("✅ Modèle chargé avec succès!")
|
||
|
||
except Exception as e:
|
||
try:
|
||
model.load_state_dict(state_dict, strict=False)
|
||
model.eval()
|
||
print("✅ Chargement partiel réussi!")
|
||
except Exception as e2:
|
||
raise RuntimeError(f"❌ Erreur lors du chargement: {e2}")
|
||
|
||
print(f"🎉 Modèle chargé depuis le run {source_run_id}")
|
||
return model, source_run_id, encoder_name, script_img_size
|
||
|
||
|
||
def load_best_model_from_registry(model_name):
|
||
"""
|
||
Charge le meilleur modèle depuis le MLflow Model Registry
|
||
(Version refactorisée - wrapper autour de load_model)
|
||
|
||
Args:
|
||
model_name (str): Nom du modèle dans le Registry
|
||
|
||
Returns:
|
||
tuple: (model, source_run_id, encoder_name, script_img_size)
|
||
"""
|
||
print(f"🏛️ Récupération du modèle '{model_name}' depuis le Registry...")
|
||
|
||
client = mlflow.tracking.MlflowClient()
|
||
|
||
try:
|
||
# Récupérer la version en Production
|
||
versions = client.search_model_versions(f"name='{model_name}'")
|
||
prod_versions = [v for v in versions if v.current_stage == "Production"]
|
||
|
||
if not prod_versions:
|
||
raise RuntimeError(f"❌ Aucune version en Production trouvée pour {model_name}")
|
||
|
||
latest_prod = max(prod_versions, key=lambda v: int(v.version))
|
||
print(f"📦 Version Registry: {latest_prod.version} (Stage: {latest_prod.current_stage})")
|
||
|
||
# Récupérer le source_run_id depuis le tag du Registry
|
||
source_run_id = latest_prod.tags.get("source_run_id")
|
||
if not source_run_id:
|
||
raise RuntimeError(f"❌ Tag 'source_run_id' absent pour la version {latest_prod.version}")
|
||
|
||
print(f"🔗 Source run détecté: {source_run_id}")
|
||
|
||
except Exception as e:
|
||
raise RuntimeError(f"❌ Erreur lors de l'accès au Registry: {e}")
|
||
|
||
# Déléguer le chargement à load_model avec le run_id trouvé
|
||
print("➡️ Délégation vers load_model...")
|
||
return load_model(run_id=source_run_id)
|
||
|
||
|
||
# Fonctions de convenance
|
||
def load_best_model(experiment_name="OC Projet 9", metric="test_mean_iou"):
|
||
"""Charge le meilleur modèle selon une métrique"""
|
||
return load_model(experiment_name=experiment_name, metric=metric, top_n=1)
|
||
|
||
def load_second_best_model(experiment_name="OC Projet 9", metric="test_mean_iou"):
|
||
"""Charge le deuxième meilleur modèle selon une métrique"""
|
||
return load_model(experiment_name=experiment_name, metric=metric, top_n=2)
|
||
|
||
def load_model_by_run_id(run_id):
|
||
"""Charge un modèle par run_id spécifique"""
|
||
return load_model(run_id=run_id)
|
||
|
||
|