import os import argparse import evaluate import numpy as np import torch import csv import pandas as pd from tqdm.auto import tqdm import cv2 import re from lib.utils_segfly import InferenceDataset, InferenceDatasetThermal, Timing, ID2COLOR, mask2label from safetensors.torch import load_file from transformers import AutoImageProcessor from lib.firefly_rgb import FireflyForSemanticSegmentationRGB, FireflyConfigRGB from lib.firefly_thermal import FireflyForSemanticSegmentationThermal, FireflyConfigThermal timing = Timing() @timing def segment_image(image, _model, _device, target_size, image_processor): with torch.no_grad(): pixel_values = image_processor(image, return_tensors="pt").pixel_values.to(_device, dtype=torch.float32) outputs = _model(pixel_values) if hasattr(outputs, "logits"): outputs = outputs.logits upsampled_logits = torch.nn.functional.interpolate( outputs.float(), size=target_size, mode="bilinear", align_corners=False ) return upsampled_logits.argmax(dim=1).detach().cpu().numpy() def run_inference(image, _model, _device, target_size, image_processor): torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() predicted = segment_image(image, _model, _device, target_size, image_processor) mem_used = torch.cuda.max_memory_allocated() / 1024**2 return predicted, mem_used def get_image_path_info(dataset, idx): path = "" for attr in ['images', 'image_paths', 'filepaths', 'samples', 'img_files', 'data']: if hasattr(dataset, attr): val = getattr(dataset, attr) if isinstance(val, (list, tuple, np.ndarray)) and len(val) > idx: path = str(val[idx]) break subfolder_name = f"pred_img_{idx:04d}" if path: scene_match = re.search(r'(scene_\d+)', path, re.IGNORECASE) alt_match = re.search(r'(\d+m)', path, re.IGNORECASE) if scene_match or alt_match: scene_str = scene_match.group(1) if scene_match else "scene_unk" alt_str = alt_match.group(1) if alt_match else "unk_m" subfolder_name = f"pred_{scene_str}_{alt_str}_{idx:04d}" actual_path = path if path and not os.path.exists(path): if hasattr(dataset, 'root_dir') and os.path.exists(os.path.join(dataset.root_dir, path)): actual_path = os.path.join(dataset.root_dir, path) elif hasattr(dataset, '_root_dir') and os.path.exists(os.path.join(dataset._root_dir, path)): actual_path = os.path.join(dataset._root_dir, path) return subfolder_name, actual_path def generate_metrics_for_dataset(loader, _model, _device, _metric, config, image_processor, visualize=False): mem_usage = [] global_gt_counts = np.zeros(config["num_classes"], dtype=np.int64) ignore_label = config.get("ignore_label", 255) if visualize: vis_dir = os.path.join(config.get("output_dir", "./"), "visualizations") os.makedirs(vis_dir, exist_ok=True) palette = ID2COLOR csv_path = os.path.join(config.get("output_dir", "./"), "per_image_iou.csv") print("Running inference and accumulating batches...") pbar = tqdm(total=len(loader)) with open(csv_path, mode="w", newline="", encoding="utf-8") as csv_file: csv_writer = csv.writer(csv_file) csv_writer.writerow(["Image_Index", "mIoU"]) for idx, (images, masks) in enumerate(loader): image = images.squeeze(0) mask = masks.squeeze(0) if len(mask.shape) == 2: target_h, target_w = mask.shape else: target_h, target_w = mask.shape[-2:] target_size = (target_h, target_w) prediction, mem = run_inference(image, _model, _device, target_size, image_processor) if isinstance(mask, torch.Tensor): mask_np = mask.detach().cpu().numpy().astype(int) else: mask_np = np.array(mask).astype(int) if getattr(image_processor, 'do_reduce_labels', False): mask_np = mask_np.copy() mask_np[mask_np == 0] = 255 mask_np = mask_np - 1 mask_np[mask_np == 254] = 255 valid_pixels = mask_np[mask_np != ignore_label] valid_pixels = valid_pixels[valid_pixels >= 0] counts = np.bincount(valid_pixels.flatten(), minlength=config["num_classes"]) global_gt_counts += counts[:config["num_classes"]] pred_np = prediction.squeeze() if isinstance(pred_np, torch.Tensor): pred_np = pred_np.cpu().numpy() valid = (mask_np != ignore_label) & (mask_np >= 0) pred_valid = pred_np[valid] mask_valid = mask_np[valid] if len(mask_valid) > 0: classes_in_img = np.unique(np.concatenate([mask_valid, pred_valid])) else: classes_in_img = [] image_ious = [] for c in classes_in_img: intersection = np.sum((pred_valid == c) & (mask_valid == c)) union = np.sum((pred_valid == c) | (mask_valid == c)) if union > 0: image_ious.append(intersection / union) else: image_ious.append(0.0) image_miou = np.mean(image_ious) if len(image_ious) > 0 else 0.0 csv_writer.writerow([idx, f"{image_miou:.4f}"]) csv_file.flush() if visualize: subfolder_name, actual_image_path = get_image_path_info(loader.dataset, idx) true_img_bgr = None if actual_image_path and os.path.exists(actual_image_path): true_img_bgr = cv2.imread(actual_image_path) if true_img_bgr is not None: orig_h, orig_w = true_img_bgr.shape[:2] img_bgr = true_img_bgr else: print(f"\n[Warning] Could not load raw image from disk for idx {idx}. Falling back to tensor size.") if isinstance(image, torch.Tensor): img_vis = image.cpu().numpy() else: img_vis = np.array(image) if img_vis.ndim == 3: if img_vis.shape[0] in [1, 3]: img_vis = np.transpose(img_vis, (1, 2, 0)) elif img_vis.shape[1] in [1, 3]: img_vis = np.transpose(img_vis, (0, 2, 1)) if img_vis.dtype.kind == 'f' and img_vis.max() <= 1.0: img_vis = img_vis * 255.0 img_vis = np.ascontiguousarray(img_vis).astype(np.uint8) if img_vis.ndim == 2: img_vis = cv2.cvtColor(img_vis, cv2.COLOR_GRAY2RGB) elif img_vis.ndim == 3 and img_vis.shape[2] == 1: img_vis = cv2.cvtColor(img_vis[:, :, 0], cv2.COLOR_GRAY2RGB) elif img_vis.ndim == 3 and img_vis.shape[2] > 3: img_vis = img_vis[:, :, :3] orig_h, orig_w = img_vis.shape[:2] img_bgr = cv2.cvtColor(img_vis, cv2.COLOR_RGB2BGR) gt_color = mask2label(mask_np, palette) pred_color = mask2label(pred_np, palette) gt_color[mask_np == ignore_label] = [0, 0, 0] try: if gt_color.shape[:2] != (orig_h, orig_w): gt_color = cv2.resize(gt_color, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST) if pred_color.shape[:2] != (orig_h, orig_w): pred_color = cv2.resize(pred_color, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST) except cv2.error as e: print(f"\n[Error] OpenCV failed to resize masks. Original mask shape: {gt_color.shape[:2]}, Target image size: ({orig_w}, {orig_h})") raise e instance_dir = os.path.join(vis_dir, subfolder_name) os.makedirs(instance_dir, exist_ok=True) gt_bgr = cv2.cvtColor(gt_color, cv2.COLOR_RGB2BGR) pred_bgr = cv2.cvtColor(pred_color, cv2.COLOR_RGB2BGR) cv2.imwrite(os.path.join(instance_dir, "image.png"), img_bgr) cv2.imwrite(os.path.join(instance_dir, "gt.png"), gt_bgr) cv2.imwrite(os.path.join(instance_dir, "pred.png"), pred_bgr) if len(prediction.shape) == 2: prediction = prediction[None, :, :] if len(mask.shape) == 2: mask = mask[None, :, :] _metric.add_batch( predictions=prediction, references=mask ) mem_usage.append(mem) pbar.update() pbar.close() print("Computing global metrics...") results = _metric.compute( num_labels=config["num_classes"], ignore_index=config.get("ignore_label"), reduce_labels=image_processor.do_reduce_labels, ) return results, mem_usage, global_gt_counts if __name__ == "__main__": parser = argparse.ArgumentParser(description="Evaluation script for Firefly RGB / Thermal models") parser.add_argument("--data_dir", type=str, default="./data", help="Path to the root directory of the dataset.") parser.add_argument("--weights_path", type=str, default="", help="Path to the .safetensors or .bin or .pth trained weights file.") parser.add_argument("--class_dict_path", type=str, default="./classes_segfly.csv", help="Path to the class dictionary CSV file.") parser.add_argument("--modality", type=str, default="rgb", choices=["rgb", "thermal"], help="Modality of the dataset (rgb or thermal).") parser.add_argument("--output_dir", type=str, default="./eval_output", help="Directory where evaluation results and visualizations will be saved.") parser.add_argument("--visualize", action="store_true", help="Generate and save side-by-side visualizations of image, ground truth, and prediction.") parser.add_argument("--image_size", type=int, default=640, help="Image size for inference.") parser.add_argument("--dinov3_repo_dir", type=str, default="./dinov3", help="Path to the local dinov3 repository.") args = parser.parse_args() model_type = "thermal" if args.modality.lower() == "thermal" else "rgb" if not args.weights_path: if args.modality == "rgb": args.weights_path = "./Firefly_RGB/model.safetensors" elif args.modality == "thermal": args.weights_path = "./Firefly_Thermal/model.safetensors" print(f"No weights path provided, defaulting to: {args.weights_path}") os.makedirs(args.output_dir, exist_ok=True) device = torch.device("cuda") df = pd.read_csv(args.class_dict_path) df = df.iloc[1:].reset_index(drop=True) classes = df["name"] id2label = classes.to_dict() label2id = {v: k for k, v in id2label.items()} num_classes = len(classes) print(f"Loaded {num_classes} classes from {args.class_dict_path}") print(label2id) config_dict = { "num_classes": num_classes, "ignore_label": 255, "class_dict_path": args.class_dict_path, "image_size": args.image_size, "output_dir": args.output_dir } print("Loading dataset...") if args.modality.lower() == "thermal": val_set = InferenceDatasetThermal( _root_dir=args.data_dir, config=config_dict, ) else: val_set = InferenceDataset( _root_dir=args.data_dir, config=config_dict, ) print(f"Number of images in validation set: {len(val_set)}") weights_dir = os.path.dirname(args.weights_path) if os.path.isfile(args.weights_path) else args.weights_path image_processor_loaded = False if weights_dir and os.path.exists(os.path.join(weights_dir, "preprocessor_config.json")): try: print(f"Loading image processor config from {weights_dir}...") image_processor = AutoImageProcessor.from_pretrained(weights_dir) image_processor_loaded = True except Exception as e: print(f"Warning: Failed to load image processor from {weights_dir}: {e}") if not image_processor_loaded: print("Falling back to standard image processor initialization...") image_processor_args = { "size": {"height": args.image_size, "width": args.image_size}, "crop_size": {"height": args.image_size, "width": args.image_size}, "image_mean": [0.485, 0.456, 0.406], "image_std": [0.229, 0.224, 0.225], "do_center_crop": False, "do_normalize": True, "do_resize": True, "do_rescale": True, "do_reduce_labels": True, } try: image_processor = AutoImageProcessor.from_pretrained("nvidia/mit-b3", **image_processor_args) except: image_processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512", **image_processor_args) print(f"Initializing {model_type} model...") if model_type == "rgb": dino_config = FireflyConfigRGB( num_labels=num_classes, image_size=args.image_size, embedding_dim=256, backbone_embed_dim=768, patch_size=16, repo_dir=args.dinov3_repo_dir, model_name="dinov3_vitb16", semantic_loss_ignore_index=255, ) model = FireflyForSemanticSegmentationRGB(dino_config) elif model_type == "thermal": dino_config = FireflyConfigThermal( num_labels=num_classes, image_size=args.image_size, embedding_dim=256, backbone_embed_dim=768, patch_size=16, num_layers=12, rein_token_length=100, feature_layers=[2, 5, 8, 11], repo_dir=args.dinov3_repo_dir, model_name="dinov3_vitb16", semantic_loss_ignore_index=255, ) model = FireflyForSemanticSegmentationThermal(dino_config) model.print_trainable_params() if args.weights_path and os.path.exists(args.weights_path): print(f"Loading weights from {args.weights_path}...") try: if args.weights_path.endswith('.safetensors'): state_dict = load_file(args.weights_path) else: checkpoint = torch.load(args.weights_path, map_location='cpu') if isinstance(checkpoint, dict): if "state_dict" in checkpoint: state_dict = checkpoint["state_dict"] elif "model" in checkpoint: state_dict = checkpoint["model"] elif "student_model" in checkpoint: state_dict = checkpoint["student_model"] else: state_dict = checkpoint else: state_dict = checkpoint new_state_dict = {} for k, v in state_dict.items(): if k.startswith("module."): new_state_dict[k[7:]] = v else: new_state_dict[k] = v msg = model.load_state_dict(new_state_dict, strict=False) print(f"Weights Loaded. Missing keys: {len(msg.missing_keys)}, Unexpected keys: {len(msg.unexpected_keys)}") if len(msg.missing_keys) > 0: print(f"First 5 missing: {msg.missing_keys[:5]}") except Exception as e: print(f"Failed to load weights file: {e}") raise else: print(f"Warning: No valid checkpoint found at {args.weights_path}. Proceeding with initialized weights.") model.eval() model.to(device) torch.backends.cudnn.benchmark = True loader = torch.utils.data.DataLoader( val_set, batch_size=1, num_workers=4, pin_memory=True ) metric = evaluate.load("mean_iou") results, memory, global_gt_counts = generate_metrics_for_dataset( loader, model, device, metric, config_dict, image_processor, visualize=args.visualize ) global_mean_iou = results["mean_iou"] global_mean_acc = results["mean_accuracy"] per_category_iou = results["per_category_iou"] ground_truth_set = global_gt_counts iou = np.array(per_category_iou) iou = np.nan_to_num(iou, nan=0.0) gt_present_mask = (ground_truth_set > 0) total_gt_freq = np.sum(ground_truth_set[gt_present_mask]) if total_gt_freq > 0: fwIoU = np.sum(ground_truth_set[gt_present_mask] * iou[gt_present_mask]) / total_gt_freq else: fwIoU = 0.0 print("=" * 40) print("EVALUATION RESULTS") print("=" * 40) print(f"Mean mIoU: {global_mean_iou:.4f}") print(f"Freq Weighted IoU: {fwIoU:.4f}") print(f"Mean pixel accuracy: {global_mean_acc:.4f}") print("-" * 40) for class_index, class_iou in enumerate(per_category_iou): label = id2label.get(class_index, f"Class {class_index}") print(f"{label}: Mean IoU = {class_iou:.4f}") print("-" * 40) print(f"Mean memory usage: {np.mean(memory):.2f} MB") print(f"Max. GPU memory usage: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB") with open(os.path.join(args.output_dir, "per_class_scores.txt"), "w", encoding="utf-8") as f: f.write(f"Mean mIoU: {global_mean_iou:.4f}\n") f.write(f"Freq Weighted IoU: {fwIoU:.4f}\n") f.write(f"Mean pixel accuracy: {global_mean_acc:.4f}\n") f.write("-" * 40 + "\n") for i, score in enumerate(per_category_iou): label = id2label.get(i, f"Class {i}") f.write(f"{label}: {score:.4f}\n") print(f"\nAll results saved to {args.output_dir}")