generate_submission_mask2former.py

# %%
CHECKPOINT_PATH = "checkpoints/goose_mask2former_l/best_mIoU_iter_96000.pth"
DATA_DIR = "data"

OUTPUT_DIR = "submission_masks_mask2former"
ZIP_NAME = "submission_mask2former.zip"
LOG_FILE = "evaluation_results_mask2former.log"
VAL_MASK_DIR = "val_masks_m2"

IMG_SCALE  = (2048, 1024)
DEVICE = "cuda:0"
NUM_QUERIES = 200

TEST_DATASET_URLS = {
    'goose_test': 'https://goose-dataset.de/storage/goose_2d_test.zip',
    'gooseEx_test': 'https://goose-dataset.de/storage/gooseEx_2d_test.zip',
}

# %%
import os
import sys
import csv
import subprocess
import shutil
import torch
import numpy as np
import zipfile
import logging
import math
from PIL import Image
from tqdm import tqdm

# Enable TF32 for faster matrix operations on Ampere+
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

os.environ['CN2DCN'] = '1'
sys.path.insert(0, os.path.join(os.getcwd(), 'DCNv4', 'DCNv4_op'))
sys.path.insert(0, os.path.join(os.getcwd(), 'DCNv4', 'segmentation'))

import mmcv_custom
import mmseg_custom
import mmcv
from mmcv import Config
from mmseg.apis import init_segmentor, inference_segmentor
from mmseg.datasets import build_dataset
from mmseg_custom.datasets.dataset_wrappers import ConcatDataset

def clear_directory(dir_path):
    """Remove all files in directory if it exists."""
    if os.path.exists(dir_path):
        for f in os.listdir(dir_path):
            file_path = os.path.join(dir_path, f)
            if os.path.isfile(file_path):
                os.remove(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)

logging.basicConfig(filename=LOG_FILE, level=logging.INFO, format='%(asctime)s - %(message)s')
os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)
clear_directory(OUTPUT_DIR)
print(f"Cleared output directory: {OUTPUT_DIR}")

# %%
zip_path = os.path.join(DATA_DIR, "goose_2d_test.zip")
if not os.path.exists(zip_path):
    subprocess.run(["wget", "-c", TEST_DATASET_URLS['goose_test'], "-O", zip_path], check=True)
if not os.path.exists(os.path.join(DATA_DIR, "goose_2d_test")):
    print("Extracting goose_2d_test.zip...")
    temp = os.path.join(DATA_DIR, "temp_test")
    os.makedirs(temp, exist_ok=True)
    subprocess.run(["unzip", "-n", "-q", zip_path, "-d", temp], check=True)
    os.rename(temp, os.path.join(DATA_DIR, "goose_2d_test"))

zip_path = os.path.join(DATA_DIR, "gooseEx_2d_test.zip")
if not os.path.exists(zip_path):
    subprocess.run(["wget", "-c", TEST_DATASET_URLS['gooseEx_test'], "-O", zip_path], check=True)
if not os.path.exists(os.path.join(DATA_DIR, "gooseEx_2d_test")):
    print("Extracting gooseEx_2d_test.zip...")
    subprocess.run(["unzip", "-n", "-q", zip_path, "-d", DATA_DIR], check=True)

print("All test datasets downloaded and extracted.")

# %%
def load_classes_from_csv(csv_path):
    classes = []
    palette = []
    with open(csv_path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            classes.append(row['class_name'])
            hex_color = row['hex'].lstrip('#')
            rgb = [int(hex_color[i:i+2], 16) for i in (0, 2, 4)]
            palette.append(rgb)
    return classes, palette

for folder in ['goose_2d_test', 'gooseEx_2d_test', 'goose_2d_train', 'gooseEx_2d_train']:
    csv_try_path = os.path.join(DATA_DIR, folder, 'goose_label_mapping.csv')
    if os.path.exists(csv_try_path):
        CSV_PATH = csv_try_path
        break

CLASSES, PALETTE = load_classes_from_csv(CSV_PATH)
NUM_CLASSES = len(CLASSES)

class_pixels = [
    3332550588, 1964900778, 1911712149, 1876082218, 1122488837, 860703945,
    846706273, 798004899, 665915338, 378936432, 320760570, 318313538,
    238425760, 204121419, 176612293, 175544747, 136367240, 105318316,
    102854343, 74485529, 65863763, 46740199, 45604828, 41963952, 32499742,
    29451667, 21314694, 19221140, 16330303, 16162575, 15428336, 13779710,
    12053148, 11249718, 10525135, 7696453, 7418888, 4633723, 4512670,
    3916332, 3780077, 3525693, 3518320, 3456606, 2460291, 2382179, 1775662,
    1376589, 1228639, 1055547, 1008673, 884133, 768866, 731310, 700429,
    469111, 206935, 42472, 13171, 1852, 1199, 58, 0, 0
]
total_pixels = 16066560000
class_frequencies = [p / total_pixels for p in class_pixels]
class_weights = [1.0 / math.log(f + 1e-6) if f > 0 else 0.0 for f in class_frequencies]
class_weights = [w / sum(class_weights) * NUM_CLASSES for w in class_weights]
CLASS_WEIGHTS_WITH_BG = class_weights + [0.1]

print(f"Loaded {NUM_CLASSES} classes from {CSV_PATH}")

# %%
img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53],
    std=[58.395, 57.12, 57.375],
    to_rgb=True
)

test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=IMG_SCALE,
        img_ratios=[0.75, 1.0, 1.25, 1.5],
        flip=True,
        flip_direction=['horizontal'],
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(type='ResizeToMultiple', size_divisor=32),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]

config = Config(
    dict(
        model=dict(
            type='EncoderDecoderMask2Former',
            pretrained=None,
            backbone=dict(
                type='FlashInternImage',
                core_op='DCNv4',
                channels=160,
                depths=[5, 5, 22, 5],
                groups=[10, 20, 40, 80],
                mlp_ratio=4.,
                drop_path_rate=0.5,
                norm_layer='LN',
                layer_scale=None,
                offset_scale=1.0,
                post_norm=True,
                with_cp=True,
                dw_kernel_size=3,
                use_clip_projector=False,
                level2_post_norm=True,
                level2_post_norm_block_ids=[5, 11, 17],
                center_feature_scale=True,
                remove_center=False,
                out_indices=(0, 1, 2, 3)),
            decode_head=dict(
                type='Mask2FormerHead',
                in_channels=[160, 320, 640, 1280],
                in_index=[0, 1, 2, 3],
                feat_channels=256,
                out_channels=256,
                num_classes=NUM_CLASSES,
                num_queries=NUM_QUERIES,
                pixel_decoder=dict(
                    type='MSDeformAttnPixelDecoder',
                    num_outs=3,
                    norm_cfg=dict(type='GN', num_groups=32),
                    act_cfg=dict(type='ReLU'),
                    encoder=dict(
                        type='DetrTransformerEncoder',
                        num_layers=6,
                        transformerlayers=dict(
                            type='BaseTransformerLayer',
                            attn_cfgs=dict(
                                type='MultiScaleDeformableAttention',
                                embed_dims=256,
                                num_heads=8,
                                num_levels=3,
                                num_points=4,
                                im2col_step=64,
                                dropout=0.0,
                                batch_first=False,
                                norm_cfg=None,
                                init_cfg=None),
                            ffn_cfgs=dict(
                                type='FFN',
                                embed_dims=256,
                                feedforward_channels=2048,
                                num_fcs=2,
                                ffn_drop=0.0,
                                act_cfg=dict(type='ReLU', inplace=True)),
                            operation_order=('self_attn', 'norm', 'ffn', 'norm')),
                        init_cfg=None),
                    positional_encoding=dict(
                        type='SinePositionalEncoding',
                        num_feats=128,
                        normalize=True),
                    init_cfg=None),
                enforce_decoder_input_project=False,
                positional_encoding=dict(
                    type='SinePositionalEncoding',
                    num_feats=128,
                    normalize=True),
                transformer_decoder=dict(
                    type='DetrTransformerDecoder',
                    return_intermediate=True,
                    num_layers=9,
                    transformerlayers=dict(
                        type='DetrTransformerDecoderLayer',
                        attn_cfgs=dict(
                            type='MultiheadAttention',
                            embed_dims=256,
                            num_heads=8,
                            attn_drop=0.0,
                            proj_drop=0.0,
                            dropout_layer=None,
                            batch_first=False),
                        ffn_cfgs=dict(
                            type='FFN',
                            embed_dims=256,
                            feedforward_channels=2048,
                            num_fcs=2,
                            act_cfg=dict(type='ReLU', inplace=True),
                            ffn_drop=0.0,
                            dropout_layer=None,
                            add_identity=True),
                        feedforward_channels=2048,
                        operation_order=('cross_attn', 'norm', 'self_attn', 'norm', 'ffn', 'norm')),
                    init_cfg=None),
                loss_cls=dict(
                    type='CrossEntropyLoss',
                    use_sigmoid=False,
                    loss_weight=2.0,
                    class_weight=CLASS_WEIGHTS_WITH_BG,
                    reduction='mean',
                    avg_non_ignore=True),
                loss_mask=dict(
                    type='CrossEntropyLoss',
                    use_sigmoid=True,
                    reduction='mean',
                    loss_weight=5.0,
                    avg_non_ignore=True),
                loss_dice=dict(
                    type='DiceLoss',
                    use_sigmoid=True,
                    activate=True,
                    reduction='mean',
                    naive_dice=True,
                    eps=1.0,
                    loss_weight=5.0),
                test_cfg=dict(
                    panoptic_on=False,
                    semantic_on=True,
                    instance_on=False,
                    max_per_image=NUM_QUERIES,
                    object_mask_thr=0.0,
                    iou_thr=0.0,
                    filter_low_score=False,
                    mode='whole')),
            train_cfg=dict(),
            test_cfg=dict(mode='whole')),
        data=dict(
            test=dict(
                type='GOOSEDataset',
                data_root='data',
                img_dir='images',
                ann_dir='labels',
                pipeline=test_pipeline)),
        test_pipeline=test_pipeline,
        device=DEVICE,
        gpu_ids=[0]),
    filename='generate_submission_mask2former.py')

from mmcv.runner import load_checkpoint

torch.serialization.add_safe_globals([np.core.multiarray.scalar])

model = init_segmentor(config, checkpoint=None, device=DEVICE)

checkpoint = load_checkpoint(model, CHECKPOINT_PATH, map_location=DEVICE)

model.CLASSES = CLASSES
model.PALETTE = PALETTE

print(f"Model loaded from {CHECKPOINT_PATH}")

# %%
print("Indexing test images...")

test_image_map = {}
TEST_IMG_ROOTS = [
    os.path.join(DATA_DIR, 'goose_2d_test', 'images', 'test'),
    os.path.join(DATA_DIR, 'gooseEx_2d_test', 'images', 'test'),
    os.path.join(DATA_DIR, 'gooseEx_2d_test_extra', 'gooseEx_2d_test', 'images', 'test'),
]
print(f"Test image roots: {TEST_IMG_ROOTS}")
for root_path in TEST_IMG_ROOTS:
    if not os.path.exists(root_path):
        continue
    for root, _, files in os.walk(root_path):
        for f in sorted(files):
            if f.endswith((".png", ".jpg")):
                parts = f.split("_")
                timestamp = next((p for p in parts if p.isdigit() and len(p) > 10), None)
                if timestamp:
                    is_vis = "_vis" in f or "camera_left" in f
                    if timestamp not in test_image_map or is_vis:
                        test_image_map[timestamp] = os.path.join(root, f)

print(f"Found {len(test_image_map)} test images")

# %%
os.makedirs(OUTPUT_DIR, exist_ok=True)

generated_count = 0
resolution_stats = {}

def make_mask_filename(img_path):
    base = os.path.basename(img_path)
    for suffix in ["_windshield_vis.png", "_windshield_nir.png", "_realsense.png",
                   "_front.png", "_camera_left.png"]:
        if base.endswith(suffix):
            return base[:-len(suffix)] + "_labelids.png"
    return base.replace(".png", "_labelids.png")


for timestamp, img_path in tqdm(sorted(test_image_map.items()), desc="Generating Submission Masks"):
    filename = make_mask_filename(img_path)

    with Image.open(img_path) as img:
        orig_width, orig_height = img.size
        res_key = f"{orig_width}x{orig_height}"
        resolution_stats[res_key] = resolution_stats.get(res_key, 0) + 1
    
    result = inference_segmentor(model, img_path)[0]
    if isinstance(result, torch.Tensor):
        result = result.cpu().numpy()
    
    
    result_img = Image.fromarray(result.astype(np.uint8)).convert("L")
    result_img = result_img.resize((orig_width, orig_height), Image.Resampling.NEAREST)
    result_img.save(os.path.join(OUTPUT_DIR, filename))
    generated_count += 1

print(f"Generated {generated_count} masks")

# %%
print(f"Packaging {generated_count} masks into {ZIP_NAME}...")
with zipfile.ZipFile(ZIP_NAME, 'w', zipfile.ZIP_DEFLATED) as zipf:
    for f in os.listdir(OUTPUT_DIR):
        if f.endswith(".png"):
            zipf.write(os.path.join(OUTPUT_DIR, f), f)

print(f"Success! All {generated_count} masks packaged.")

# %%
print("\nRunning local evaluation on validation split...")

val_img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53],
    std=[58.395, 57.12, 57.375],
    to_rgb=True
)

val_test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=IMG_SCALE,
        flip=False,
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(type='ResizeToMultiple', size_divisor=32),
            dict(type='RandomFlip', prob=0.0),
            dict(type='Normalize', **val_img_norm_cfg),
            dict(type='ImageToTensor', keys=['img', 'gt_semantic_seg']),
            dict(type='Collect', keys=['img', 'gt_semantic_seg']),
        ])
]

goose_val_cfg = dict(
    type='GOOSEDataset',
    data_root=os.path.join(DATA_DIR, 'goose_2d_val'),
    img_dir='images/val',
    ann_dir='labels/val',
    img_suffix='_windshield_vis.png',
    seg_map_suffix='_labelids.png',
    reduce_zero_label=False,
    classes=CLASSES,
    palette=PALETTE,
    pipeline=val_test_pipeline)

gooseEx_val_cfg = dict(
    type='GOOSEDataset',
    data_root=os.path.join(DATA_DIR, 'gooseEx_2d_val'),
    img_dir='images/val',
    ann_dir='labels/val',
    img_suffix='_camera_left.png',
    seg_map_suffix='_labelids.png',
    reduce_zero_label=False,
    classes=CLASSES,
    palette=PALETTE,
    pipeline=val_test_pipeline)

goose_val_ds = build_dataset(goose_val_cfg)
gooseEx_val_ds = build_dataset(gooseEx_val_cfg)
val_dataset = ConcatDataset([goose_val_ds, gooseEx_val_ds], separate_eval=False)
print(f"Validation samples: {len(val_dataset)}")
print(f"  - goose_2d_val: {len(goose_val_ds)}")
print(f"  - gooseEx_2d_val: {len(gooseEx_val_ds)}")

eval_count = 0
intersection = np.zeros(NUM_CLASSES)
union = np.zeros(NUM_CLASSES)
tp = np.zeros(NUM_CLASSES)
fp = np.zeros(NUM_CLASSES)
fn = np.zeros(NUM_CLASSES)
total_pixels = 0
correct_pixels = 0

os.makedirs(VAL_MASK_DIR, exist_ok=True)
clear_directory(VAL_MASK_DIR)

for idx in tqdm(range(len(val_dataset)), desc="Evaluating"):
    if idx < len(goose_val_ds):
        ds = goose_val_ds
        ds_idx = idx
    else:
        ds = gooseEx_val_ds
        ds_idx = idx - len(goose_val_ds)
    
    img_info = ds.img_infos[ds_idx]
    img_path = os.path.join(ds.img_dir, img_info['filename'])
    ann_path = os.path.join(ds.ann_dir, img_info['ann']['seg_map'])
    
    with torch.no_grad():
        result = inference_segmentor(model, img_path)[0]
    if isinstance(result, torch.Tensor):
        pred = result.cpu().numpy()
    else:
        pred = result
    
    
    gt = mmcv.imread(ann_path, flag='unchanged')
    if gt is None:
        continue
    
    if pred.shape != gt.shape:
        from PIL import Image
        pred_img = Image.fromarray(pred.astype(np.uint8))
        pred_img = pred_img.resize((gt.shape[1], gt.shape[0]), Image.Resampling.NEAREST)
        pred = np.array(pred_img)

    ds_tag = 'goose' if ds is goose_val_ds else 'gooseEx'
    key = f"{ds_tag}__{os.path.basename(img_info['filename']).rsplit('.', 1)[0]}.png"
    Image.fromarray(pred.astype(np.uint8)).convert('L').save(os.path.join(VAL_MASK_DIR, key))

    for cls in range(NUM_CLASSES):
        pred_mask = (pred == cls)
        gt_mask = (gt == cls)
        
        intersection[cls] += np.logical_and(pred_mask, gt_mask).sum()
        union[cls] += np.logical_or(pred_mask, gt_mask).sum()
        
        tp[cls] += np.logical_and(pred_mask, gt_mask).sum()
        fp[cls] += np.logical_and(pred_mask, ~gt_mask).sum()
        fn[cls] += np.logical_and(~pred_mask, gt_mask).sum()
    
    total_pixels += gt.size
    correct_pixels += (pred == gt).sum()
    
    eval_count += 1

ious = np.divide(intersection, union, out=np.zeros_like(intersection), where=union > 0)
miou = ious.mean()

precision = np.divide(tp, tp + fp, out=np.zeros_like(tp), where=(tp + fp) > 0)
recall = np.divide(tp, tp + fn, out=np.zeros_like(tp), where=(tp + fn) > 0)
f1 = np.divide(2 * precision * recall, precision + recall, out=np.zeros_like(precision), where=(precision + recall) > 0)

overall_acc = correct_pixels / total_pixels if total_pixels > 0 else 0

print("\n" + "="*80)
print("DETAILED VALIDATION METRICS")
print("="*80)
print(f"Overall Accuracy: {overall_acc:.4f}")
print(f"Mean IoU: {miou:.4f}")
print(f"Mean Precision: {precision.mean():.4f}")
print(f"Mean Recall: {recall.mean():.4f}")
print(f"Mean F1: {f1.mean():.4f}")
print("\nPer-Class Metrics:")
print("-"*80)
print(f"{'Class':<40} {'IoU':<10} {'Precision':<10} {'Recall':<10} {'F1':<10} {'TP':<10} {'FP':<10} {'FN':<10}")
print("-"*80)

for idx, class_name in enumerate(CLASSES):
    print(f"{class_name:<40} {ious[idx]:<10.4f} {precision[idx]:<10.4f} {recall[idx]:<10.4f} {f1[idx]:<10.4f} {int(tp[idx]):<10} {int(fp[idx]):<10} {int(fn[idx]):<10}")

print("-"*80)
print(f"Evaluated {eval_count} samples")
print("="*80 + "\n")

detailed_log = f"""
DETAILED VALIDATION METRICS
{'='*80}
Overall Accuracy: {overall_acc:.4f}
Mean IoU: {miou:.4f}
Mean Precision: {precision.mean():.4f}
Mean Recall: {recall.mean():.4f}
Mean F1: {f1.mean():.4f}

Per-Class Metrics:
{'-'*80}
{'Class':<40} {'IoU':<10} {'Precision':<10} {'Recall':<10} {'F1':<10} {'TP':<10} {'FP':<10} {'FN':<10}
{'-'*80}
"""
for idx, class_name in enumerate(CLASSES):
    detailed_log += f"{class_name:<40} {ious[idx]:<10.4f} {precision[idx]:<10.4f} {recall[idx]:<10.4f} {f1[idx]:<10.4f} {int(tp[idx]):<10} {int(fp[idx]):<10} {int(fn[idx]):<10}\n"
detailed_log += f"{'-'*80}\nEvaluated {eval_count} samples\n{'='*80}\n"

logging.info(detailed_log)

summary = f"Evaluation Summary:\n- Generated Masks: {generated_count}\n- Output Resolutions: {resolution_stats}\n- Local mIoU: {miou:.4f}\n- Submission ZIP: {ZIP_NAME}"
print("\n" + summary)
logging.info(summary)

CSV_OUTPUT_PATH = os.path.join(os.path.dirname(CHECKPOINT_PATH), "validation_class_iou.csv")
with open(CSV_OUTPUT_PATH, 'w', newline='') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(['class_name', 'iou', 'precision', 'recall', 'f1'])
    for idx, class_name in enumerate(CLASSES):
        writer.writerow([
            class_name,
            float(ious[idx]) if not np.isnan(ious[idx]) else 0.0,
            float(precision[idx]) if not np.isnan(precision[idx]) else 0.0,
            float(recall[idx]) if not np.isnan(recall[idx]) else 0.0,
            float(f1[idx]) if not np.isnan(f1[idx]) else 0.0
        ])
print(f"Per-class metrics saved to: {CSV_OUTPUT_PATH}")