| import os | |
| import cv2 | |
| import torch | |
| from tqdm import tqdm | |
| from torch.utils.data import DataLoader | |
| from segmentation_models_pytorch.base.modules import Activation | |
| from SemanticModel.data_loader import SegmentationDataset | |
| from SemanticModel.metrics import compute_mean_iou | |
| from SemanticModel.image_preprocessing import get_validation_augmentations | |
| def evaluate_model(model_config, data_path, image_size=None): | |
| """Evaluates model performance on a dataset.""" | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| classes = ['background'] + model_config.classes if model_config.background_flag else model_config.classes | |
| data_path = os.path.realpath(data_path) | |
| image_subdir = os.path.join(data_path, 'Images') | |
| mask_subdir = os.path.join(data_path, 'Masks') | |
| if not all(os.path.exists(d) for d in [image_subdir, mask_subdir]): | |
| raise Exception("Missing required subdirectories: 'Images' and 'Masks'") | |
| if not image_size: | |
| sample_image = cv2.imread(os.path.join(image_subdir, os.listdir(image_subdir)[0])) | |
| height, width = sample_image.shape[:2] | |
| image_size = max(height, width) | |
| evaluation_dataset = SegmentationDataset( | |
| data_path, | |
| classes=classes, | |
| augmentation=get_validation_augmentations( | |
| im_width=image_size, | |
| im_height=image_size, | |
| fixed_size=False | |
| ), | |
| preprocessing=model_config.preprocessing | |
| ) | |
| evaluation_loader = DataLoader( | |
| evaluation_dataset, | |
| batch_size=1, | |
| shuffle=False, | |
| num_workers=2 | |
| ) | |
| model = model_config.model.to(device) | |
| model.eval() | |
| requires_sigmoid = False | |
| if model_config.n_classes == 1: | |
| current_activation = _check_activation_function(model) | |
| if current_activation != 'Sigmoid': | |
| requires_sigmoid = True | |
| predictions = [] | |
| ground_truth = [] | |
| print("Evaluating model performance...") | |
| with torch.no_grad(): | |
| for images, masks in tqdm(evaluation_loader): | |
| images = images.to(device) | |
| masks = masks.to(device) | |
| outputs = model.forward(images) | |
| if model_config.n_classes > 1: | |
| predictions.extend([p.cpu().argmax(dim=0) for p in outputs]) | |
| ground_truth.extend([gt.cpu().argmax(dim=0) for gt in masks]) | |
| else: | |
| if requires_sigmoid: | |
| predictions.extend([ | |
| (torch.sigmoid(p) > 0.5).float().squeeze().cpu() | |
| for p in outputs | |
| ]) | |
| else: | |
| predictions.extend([ | |
| (p > 0.5).float().squeeze().cpu() | |
| for p in outputs | |
| ]) | |
| ground_truth.extend([gt.cpu().squeeze() for gt in masks]) | |
| metrics = compute_mean_iou( | |
| predictions, | |
| ground_truth, | |
| num_labels=len(classes), | |
| ignore_index=255 | |
| ) | |
| print("\nEvaluation Results:") | |
| print(f"Mean IoU: {metrics['mean_iou']:.3f}") | |
| print("\nPer-class IoU:") | |
| for idx, iou in enumerate(metrics['per_category_iou']): | |
| print(f"{classes[idx]}: {iou:.3f}") | |
| return metrics | |
| def _check_activation_function(model): | |
| """Checks the activation function used in model's segmentation head.""" | |
| from segmentation_models_pytorch.base.modules import Activation | |
| activation_functions = [] | |
| for _, module in model.segmentation_head.named_children(): | |
| if isinstance(module, Activation): | |
| activation_functions.append(type(module.activation).__name__) | |
| return activation_functions[-1] if activation_functions else None |