generate_submission.py

# %%
# 1. Configuration - All parameters at top
CHECKPOINT_PATH = "checkpoints/goose_seg_dcnv4/best_mIoU_iter_96000.pth"
DATA_DIR = "data"

OUTPUT_DIR = "submission_masks"
ZIP_NAME = "submission.zip"
LOG_FILE = "evaluation_results.log"
VAL_MASK_DIR = "val_masks_m1"
VAL_MANIFEST = "val_manifest.csv"

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

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',
}

# %%
# 2. Environment setup
import os
import sys
import csv
import subprocess
import shutil
import torch
import numpy as np
import zipfile
import logging
from PIL import Image
from tqdm import tqdm
from torchvision import transforms

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

# %%
# 3. Download and extract test datasets

# goose_2d_test - special handling (extracts to images/ directly)
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"))

# gooseEx_2d_test - normal extraction
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.")

# %%
# 4. Load class information from CSV
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)

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

# %%
# 5. Load model configuration and checkpoint
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='EncoderDecoder',
            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.4,
                norm_layer='LN',
                layer_scale=1.0,
                offset_scale=2.0,
                post_norm=True,
                with_cp=False,
                dcn_output_bias=True,
                mlp_fc2_bias=True,
                dw_kernel_size=3,
                out_indices=(0, 1, 2, 3)),
            decode_head=dict(
                type='UPerHead',
                in_channels=[160, 320, 640, 1280],
                in_index=[0, 1, 2, 3],
                pool_scales=(1, 2, 3, 6),
                channels=512,
                dropout_ratio=0.1,
                num_classes=NUM_CLASSES,
                norm_cfg=dict(type='SyncBN', requires_grad=True),
                align_corners=False),
            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.py')

# Build model using init_segmentor like image_demo.py
from mmcv.runner import load_checkpoint

# Add numpy scalar to safe globals for checkpoint loading
torch.serialization.add_safe_globals([np.core.multiarray.scalar])

# Build model from config (without checkpoint)
model = init_segmentor(config, checkpoint=None, device=DEVICE)

# Load checkpoint weights manually
checkpoint = load_checkpoint(model, CHECKPOINT_PATH, map_location=DEVICE)

# Set CLASSES and PALETTE
model.CLASSES = CLASSES
model.PALETTE = PALETTE

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

# %%
# 6. File Discovery and Indexing
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")

# %%
# 7. Inference and Mask Generation
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
    
    # Run inference using inference_segmentor like image_demo.py
    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")

# %%
# 8. Packaging Submission
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.")

# %%
# 9. Local Evaluation on Val Split
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)  # true positive
fp = np.zeros(NUM_CLASSES)  # false positive
fn = np.zeros(NUM_CLASSES)  # false negative
total_pixels = 0
correct_pixels = 0

os.makedirs(VAL_MASK_DIR, exist_ok=True)
clear_directory(VAL_MASK_DIR)
val_manifest_rows = []

for idx in tqdm(range(len(val_dataset)), desc="Evaluating"):
    # Get the raw sample info to find the image path
    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)
    
    # Get image info from dataset's img_infos
    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'])
    
    # Run inference on the original image file
    with torch.no_grad():
        result = inference_segmentor(model, img_path)[0]
    if isinstance(result, torch.Tensor):
        pred = result.cpu().numpy()
    else:
        pred = result
    
    
    # Load ground truth annotation
    gt = mmcv.imread(ann_path, flag='unchanged')
    if gt is None:
        continue
    
    # Resize prediction to match GT if needed
    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)
    
    # Save prediction mask for offline ensemble tuning
    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))
    val_manifest_rows.append((key, ds_tag, img_path, ann_path))

    # Calculate per-class metrics
    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()
    
    # Overall accuracy
    total_pixels += gt.size
    correct_pixels += (pred == gt).sum()
    
    eval_count += 1

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

# Precision, Recall, F1 per class
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 accuracy
overall_acc = correct_pixels / total_pixels if total_pixels > 0 else 0

# Log detailed per-class metrics
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")

# Also log to file
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}")

with open(VAL_MANIFEST, 'w', newline='') as mf:
    w = csv.writer(mf)
    w.writerow(['key', 'dataset', 'img_path', 'gt_path'])
    for row in val_manifest_rows:
        w.writerow(row)
print(f"Val manifest saved: {VAL_MANIFEST} ({len(val_manifest_rows)} rows)")
print(f"Val masks saved to: {VAL_MASK_DIR}")