test_segmentation.py

import os
import glob
import time
import numpy as np
import cv2
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.nn.functional as F

import albumentations as A
from albumentations.pytorch import ToTensorV2

# ═══════════════════════════════════════════════════════════════
# 1. SETUP LOCAL PATHS (Synchronized with code_geass_submission)
# ═══════════════════════════════════════════════════════════════
TEST_IMG_DIR  = './code_geass_submission/Offroad_Segmentation_testImages/Color_Images/'
TEST_MASK_DIR = './code_geass_submission/Offroad_Segmentation_testImages/Segmentation/'

MODEL_PATH    = './code_geass_submission/segmentation_head.pth'
OUTPUT_DIR    = './code_geass_submission/predictions/'

# ═══════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════
# Hardware Detection
if torch.cuda.is_available():
    DEVICE = torch.device('cuda')
    print("✅ NVIDIA GPU Detected! Engaging Full Multi-Scale TTA for Max Accuracy.")
    FAST_CPU_TEST = False  # Full accuracy mode
elif torch.backends.mps.is_available():
    DEVICE = torch.device('mps')
    print("✅ Apple Silicon GPU Detected!")
    FAST_CPU_TEST = False 
else:
    DEVICE = torch.device('cpu')
    print("⚠️ No GPU Detected. Forcing CPU execution in Fast Mode.")
    FAST_CPU_TEST = True   # Skip rotations to save massive time on CPU

VALUE_MAP = {
    0: 0, 100: 1, 200: 2, 300: 3, 500: 4,
    550: 5, 600: 6, 700: 7, 800: 8, 7100: 0, 10000: 9
}

CLASS_NAMES = [
    "Landscape/Bg", "Trees", "Lush Bushes", "Dry Grass", "Dry Bushes",
    "Ground Clutter", "Logs", "Rocks", "Unknown", "Sky"
]

REVERSE_MAP = {
    0: 0, 1: 100, 2: 200, 3: 300, 4: 500,
    5: 550, 6: 600, 7: 700, 8: 800, 9: 10000
}

# ═══════════════════════════════════════════════════════════════
# ARCHITECTURE
# ═══════════════════════════════════════════════════════════════
class PPM(nn.Module):
    def __init__(self, in_ch, out_ch, bins=(1, 2, 3, 6)):
        super().__init__()
        self.features = nn.ModuleList()
        for bin in bins:
            self.features.append(nn.Sequential(
                nn.AdaptiveAvgPool2d(bin),
                nn.Conv2d(in_ch, out_ch, 1, bias=False),
                nn.BatchNorm2d(out_ch),
                nn.ReLU(inplace=True)
            ))
        self.bottleneck = nn.Sequential(
            nn.Conv2d(in_ch + len(bins) * out_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )
    
    def forward(self, x):
        h, w = x.shape[2:]
        out = [x]
        for f in self.features:
            out.append(F.interpolate(f(x), size=(h, w), mode='bilinear', align_corners=False))
        out = torch.cat(out, 1)
        return self.bottleneck(out)

class DINOv2UPerNet(nn.Module):
    def __init__(self, n_classes=10):
        super().__init__()
        self.embed_dim = 768
        self.backbone = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14', pretrained=False)
        
        channels = 256
        self.ppm = PPM(self.embed_dim, channels, bins=(1, 2, 3, 6))
        
        self.fpn_in = nn.ModuleList([
            nn.Sequential(nn.Conv2d(self.embed_dim, channels, 1, bias=False), nn.BatchNorm2d(channels), nn.ReLU(inplace=True))
            for _ in range(3)
        ])
        
        self.fpn_out = nn.ModuleList([
            nn.Sequential(nn.Conv2d(channels, channels, 3, padding=1, bias=False), nn.BatchNorm2d(channels), nn.ReLU(inplace=True))
            for _ in range(3)
        ])
        
        self.fpn_bottleneck = nn.Sequential(
            nn.Conv2d(channels * 4, channels, 3, padding=1, bias=False),
            nn.BatchNorm2d(channels),
            nn.ReLU(inplace=True)
        )
        
        self.final_head = nn.Conv2d(channels, n_classes, 1)
    
    def _extract_features(self, x):
        features = self.backbone.get_intermediate_layers(x, n=4, return_class_token=False)
        B, _, H, W = x.shape
        out = []
        for feat in features:
            feat = feat.reshape(B, H//14, W//14, self.embed_dim).permute(0, 3, 1, 2)
            out.append(feat)
        return out
    
    def _fpn_forward(self, f1, f2, f3, f4):
        p4 = self.ppm(f4)
        
        p3 = self.fpn_in[2](f3) + F.interpolate(p4, size=f3.shape[2:], mode='bilinear', align_corners=False)
        p3 = self.fpn_out[2](p3)
        
        p2 = self.fpn_in[1](f2) + F.interpolate(p3, size=f2.shape[2:], mode='bilinear', align_corners=False)
        p2 = self.fpn_out[1](p2)
        
        p1 = self.fpn_in[0](f1) + F.interpolate(p2, size=f1.shape[2:], mode='bilinear', align_corners=False)
        p1 = self.fpn_out[0](p1)
        
        p4_up = F.interpolate(p4, size=p1.shape[2:], mode='bilinear', align_corners=False)
        p3_up = F.interpolate(p3, size=p1.shape[2:], mode='bilinear', align_corners=False)
        p2_up = F.interpolate(p2, size=p1.shape[2:], mode='bilinear', align_corners=False)
        
        fpn_out = torch.cat([p1, p2_up, p3_up, p4_up], dim=1)
        return self.fpn_bottleneck(fpn_out)
    
    def forward(self, x):
        target_h, target_w = x.shape[2], x.shape[3]
        f1, f2, f3, f4 = self._extract_features(x)
        out = self._fpn_forward(f1, f2, f3, f4)
        out = self.final_head(out)
        return F.interpolate(out, size=(target_h, target_w), mode='bilinear', align_corners=False)

# ═══════════════════════════════════════════════════════════════
# TTA & METRICS
# ═══════════════════════════════════════════════════════════════
def generate_d4_tta(tensor):
    if FAST_CPU_TEST: return [tensor]
    return [
        tensor,
        torch.rot90(tensor, k=1, dims=[2, 3]),
        torch.rot90(tensor, k=2, dims=[2, 3]),
        torch.rot90(tensor, k=3, dims=[2, 3]),
        torch.flip(tensor, dims=[3]),
        torch.rot90(torch.flip(tensor, dims=[3]), k=1, dims=[2, 3]),
        torch.rot90(torch.flip(tensor, dims=[3]), k=2, dims=[2, 3]),
        torch.rot90(torch.flip(tensor, dims=[3]), k=3, dims=[2, 3]),
    ]

def reverse_d4_tta(prob_maps):
    if FAST_CPU_TEST: return prob_maps
    return [
        prob_maps[0],
        torch.rot90(prob_maps[1], k=-1, dims=[2, 3]),
        torch.rot90(prob_maps[2], k=-2, dims=[2, 3]),
        torch.rot90(prob_maps[3], k=-3, dims=[2, 3]),
        torch.flip(prob_maps[4], dims=[3]),
        torch.flip(torch.rot90(prob_maps[5], k=-1, dims=[2, 3]), dims=[3]),
        torch.flip(torch.rot90(prob_maps[6], k=-2, dims=[2, 3]), dims=[3]),
        torch.flip(torch.rot90(prob_maps[7], k=-3, dims=[2, 3]), dims=[3]),
    ]

def remap_mask(mask):
    out = np.zeros_like(mask, dtype=np.int64)
    for raw_id, mapped_id in VALUE_MAP.items():
        out[mask == raw_id] = mapped_id
    return out

def predict_multiscale(model, img_np):
    scales = [1.0] if FAST_CPU_TEST else [0.75, 1.0, 1.25]
    orig_h, orig_w = img_np.shape[:2]
    scale_probs = []
    
    transform = A.Compose([
        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ToTensorV2(),
    ])
    
    for scale in scales:
        new_h = int(orig_h * scale)
        new_w = int(orig_w * scale)
        new_h = (new_h // 14) * 14
        new_w = (new_w // 14) * 14
        
        scaled_img = cv2.resize(img_np, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        tensor = transform(image=scaled_img)['image'].unsqueeze(0).to(DEVICE)
        
        tta_views = generate_d4_tta(tensor)
        prob_maps = []
        
        with torch.no_grad():
            for view in tta_views:
                logits = model(view).float()
                probs = F.softmax(logits, dim=1)
                prob_maps.append(probs)
        
        aligned_probs = reverse_d4_tta(prob_maps)
        scale_prob = torch.stack(aligned_probs, dim=0).mean(dim=0)
        scale_prob = F.interpolate(scale_prob, size=(orig_h, orig_w), mode='bilinear', align_corners=False)
        scale_probs.append(scale_prob)
    
    final_prob = torch.stack(scale_probs, dim=0).mean(dim=0)
    return final_prob

# ═══════════════════════════════════════════════════════════════
# EXECUTION
# ═══════════════════════════════════════════════════════════════
def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    
    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(f"Model not found at {MODEL_PATH}. Check your path!")
    
    print(f"Loading weights from {MODEL_PATH}...")
    model = DINOv2UPerNet(n_classes=10).to(DEVICE)
    ckpt = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=False)
    
    # 🚀 Strip out training wrappers (.block. / _orig_mod) to prevent load crashes
    raw_state_dict = ckpt.get('model_state_dict', ckpt)
    clean_state_dict = {}
    for k, v in raw_state_dict.items():
        clean_k = k.replace('.block.', '.')
        clean_k = clean_k.replace('_orig_mod.', '')
        clean_state_dict[clean_k] = v
    
    model.load_state_dict(clean_state_dict)
    model.eval()
    print("✅ Model Loaded Successfully!")
    
    img_paths = sorted(glob.glob(os.path.join(TEST_IMG_DIR, '*.png')) + glob.glob(os.path.join(TEST_IMG_DIR, '*.jpg')))
    
    if len(img_paths) == 0:
        raise FileNotFoundError(f"No images found in {TEST_IMG_DIR}. Check your paths!")
    
    print(f"Found {len(img_paths)} images to evaluate...")
    
    n_classes = 10
    confusion_matrix = np.zeros((n_classes, n_classes), dtype=np.float64)
    masks_found = 0
    
    for img_path in tqdm(img_paths, desc="Evaluating Images"):
        base_name = os.path.basename(img_path)
        orig_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
        orig_img = cv2.cvtColor(orig_img, cv2.COLOR_BGR2RGB)
        
        mask_path = os.path.join(TEST_MASK_DIR, base_name)
        if not os.path.exists(mask_path):
            mask_path = os.path.join(TEST_MASK_DIR, os.path.splitext(base_name)[0] + '.png')
        
        has_mask = os.path.exists(mask_path)
        
        if has_mask:
            gt_mask_raw = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
            if len(gt_mask_raw.shape) == 3:
                gt_mask_raw = gt_mask_raw[:, :, 0]
            gt_mask = remap_mask(gt_mask_raw)
            masks_found += 1
        
        ensemble_prob = predict_multiscale(model, orig_img)
        pred_idx = ensemble_prob.argmax(dim=1).squeeze(0).cpu().numpy().astype(np.uint8)
        
        if has_mask:
            valid_mask = (gt_mask >= 0) & (gt_mask < n_classes)
            confusion_matrix += np.bincount(
                n_classes * gt_mask[valid_mask].astype(int) + pred_idx[valid_mask],
                minlength=n_classes**2
            ).reshape(n_classes, n_classes)
        
        final_mask = np.zeros_like(pred_idx, dtype=np.uint16)
        for class_idx, raw_id in REVERSE_MAP.items():
            final_mask[pred_idx == class_idx] = raw_id
        
        out_name = os.path.splitext(base_name)[0] + '.png'
        out_path = os.path.join(OUTPUT_DIR, out_name)
        cv2.imwrite(out_path, final_mask)
    
    if masks_found > 0:
        print("\n" + "="*50)
        print(f"FINAL IoU RESULTS (Tested on {masks_found} images)")
        print("="*50)
        
        ious = []
        for c in range(n_classes):
            tp = confusion_matrix[c, c]
            fp = confusion_matrix[:, c].sum() - tp
            fn = confusion_matrix[c, :].sum() - tp
            union = tp + fp + fn
            
            if union == 0:
                iou = float('nan')
            else:
                iou = tp / union
            
            ious.append(iou)
            class_name = CLASS_NAMES[c]
            print(f"Class {c:2d} ({class_name:<15}): {iou:.4f}" if not np.isnan(iou) else f"Class {c:2d} ({class_name:<15}): N/A")
        
        mean_iou = np.nanmean(ious)
        print("-" * 50)
        print(f"Mean IoU (mIoU): {mean_iou:.4f}")
        print("="*50)

if __name__ == '__main__':
    main()