| 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}") |
|
|
