eval.py

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