train_segment_py.py

import os
import sys
import csv
import subprocess
import numpy as np
import torch
import math

# %%
# 1. Configuration - All parameters at top
EXPERIMENT_NAME = "goose_seg_dcnv4"
CHECKPOINT_DIR = "checkpoints/goose_seg_dcnv4"
DATA_DIR = "data"
CSV_PATH = "data/goose_label_mapping.csv"  # Will be created by prepare_combined_dataset.py

BATCH_SIZE = 2
NUM_WORKERS = 16
MAX_ITERS = 200_000
EVAL_INTERVAL = 1_000
CHECKPOINT_INTERVAL = 1_000

LR = 8e-4
CROP_SIZE = (2048, 1024)
IMG_SCALE = (2048, 1024)

PRETRAINED_BACKBONE = 'https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_l_22kto1k_384.pth'
LOAD_FROM = '/home/ubuntu/goose/checkpoints/goose_seg_dcnv4/best_mIoU_iter_96000.pth'
RESUME_FROM = None

DATASET_URLS = {
    'goose_train': 'https://goose-dataset.de/storage/goose_2d_train.zip',
    'goose_val': 'https://goose-dataset.de/storage/goose_2d_val.zip',
    'gooseEx_train': 'https://goose-dataset.de/storage/gooseEx_2d_train.zip',
    'gooseEx_val': 'https://goose-dataset.de/storage/gooseEx_2d_val.zip',
}

# %%
# 2. Environment setup


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

from mmcv import Config
from mmseg.apis import train_segmentor, init_random_seed, set_random_seed
from mmseg.models import build_segmentor
from mmseg.datasets import build_dataset
from mmseg.utils import get_root_logger, setup_multi_processes, collect_env
from mmseg_custom.datasets.dataset_wrappers import ConcatDataset

os.makedirs(CHECKPOINT_DIR, exist_ok=True)
os.makedirs(DATA_DIR, exist_ok=True)

# Dataset paths (separate goose and gooseEx)
GOOSE_TRAIN_DIR = os.path.join(DATA_DIR, 'goose_2d_train')
GOOSEEX_TRAIN_DIR = os.path.join(DATA_DIR, 'gooseEx_2d_train')
GOOSE_COPYPASTE_TRAIN_DIR = os.path.join(DATA_DIR, 'goose_2d_train_copypaste')
GOOSE_VAL_DIR = os.path.join(DATA_DIR, 'goose_2d_val')
GOOSEEX_VAL_DIR = os.path.join(DATA_DIR, 'gooseEx_2d_val')

# %%
# 3. Verify datasets exist
for dir_path in [GOOSE_TRAIN_DIR, GOOSEEX_TRAIN_DIR, GOOSE_VAL_DIR, GOOSEEX_VAL_DIR]:
    if not os.path.exists(dir_path):
        print("=" * 80)
        print(f"Dataset not found: {dir_path}")
        print("=" * 80)
        sys.exit(1)

if not os.path.exists(GOOSE_COPYPASTE_TRAIN_DIR):
    print(f"Warning: Copy paste sample dataset not found at {GOOSE_COPYPASTE_TRAIN_DIR}")
    sys.exit(1)

print(f"Using datasets:")
print(f"  Train (goose):          {GOOSE_TRAIN_DIR}")
print(f"  Train (gooseEx):        {GOOSEEX_TRAIN_DIR}")
print(f"  Train (copy paste):     {GOOSE_COPYPASTE_TRAIN_DIR}")
print(f"  Val (goose):            {GOOSE_VAL_DIR}")
print(f"  Val (gooseEx):          {GOOSEEX_VAL_DIR}")

# %%
# 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

# CSV should be in data directory from prepare_combined_dataset.py
for folder in ['goose_2d_train', 'gooseEx_2d_train', 'goose_2d_train_copypaste']:
    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

if not os.path.exists(CSV_PATH):
    print(f"Warning: CSV not found at {CSV_PATH}")
    print("Will attempt to use classes from dataset if available.")

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(1.02 + f) if f > 0 else 0.0 for f in class_frequencies]
class_weights = [w / sum(class_weights) * NUM_CLASSES for w in class_weights]

print(f"Loaded {NUM_CLASSES} classes from {CSV_PATH}")
print(f"First 5 classes: {CLASSES[:5]}")

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

train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', reduce_zero_label=False),
    dict(type='Resize', img_scale=IMG_SCALE, ratio_range=(0.5, 2.0)),
    dict(type='RandomCrop', crop_size=CROP_SIZE, cat_max_ratio=0.75),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PhotoMetricDistortion'),
    # dict(type='RGB2Gray'),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='Pad', size=CROP_SIZE, pad_val=0, seg_pad_val=255),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg']),
]

test_pipeline = [
    dict(type='LoadImageFromFile'),
    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'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]

# %%
# 7. Build model configuration
norm_cfg = dict(type='SyncBN', requires_grad=True)

model_cfg = 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),
        init_cfg=dict(type='Pretrained', checkpoint=PRETRAINED_BACKBONE)),
    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=norm_cfg,
        align_corners=False,
        loss_decode=[
            dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1, class_weight=class_weights),
        ]),
    auxiliary_head=dict(
        type='FCNHead',
        in_channels=640,
        in_index=2,
        channels=256,
        num_convs=1,
        concat_input=False,
        dropout_ratio=0.1,
        num_classes=NUM_CLASSES,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)),
    train_cfg=dict(),
    test_cfg=dict(mode='whole'))

# %%
# 7. Build data configuration (separate goose and gooseEx datasets)
goose_train_cfg = dict(
    type='GOOSEDataset',
    data_root=GOOSE_TRAIN_DIR,
    img_dir='images/train',
    ann_dir='labels/train',
    img_suffix='_windshield_vis.png',
    seg_map_suffix='_labelids.png',
    reduce_zero_label=False,
    classes=CLASSES,
    palette=PALETTE,
    pipeline=train_pipeline)

gooseEx_train_cfg = dict(
    type='GOOSEDataset',
    data_root=GOOSEEX_TRAIN_DIR,
    img_dir='images/train',
    ann_dir='labels/train',
    img_suffix='_camera_left.png',
    seg_map_suffix='_labelids.png',
    reduce_zero_label=False,
    classes=CLASSES,
    palette=PALETTE,
    pipeline=train_pipeline)

goose_copypaste_train_cfg = dict(
    type='GOOSEDataset',
    data_root=GOOSE_COPYPASTE_TRAIN_DIR,
    img_dir='images/train',
    ann_dir='labels/train',
    img_suffix='.png',
    seg_map_suffix='_labelids.png',
    reduce_zero_label=False,
    classes=CLASSES,
    palette=PALETTE,
    pipeline=train_pipeline)

goose_val_cfg = dict(
    type='GOOSEDataset',
    data_root=GOOSE_VAL_DIR,
    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=test_pipeline)

gooseEx_val_cfg = dict(
    type='GOOSEDataset',
    data_root=GOOSEEX_VAL_DIR,
    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=test_pipeline)

val_cfg = dict(
    type='ConcatDataset',
    datasets=[goose_val_cfg, gooseEx_val_cfg],
    separate_eval=False)

data_cfg = dict(
    samples_per_gpu=BATCH_SIZE,
    workers_per_gpu=NUM_WORKERS,
    train=goose_train_cfg,
    val=val_cfg,
    test=val_cfg)

# %%
# 8. Build optimizer and scheduler configuration
optimizer_cfg = dict(
    type='AdamW',
    lr=LR,
    betas=(0.9, 0.999),
    weight_decay=0.05,
    constructor='CustomLayerDecayOptimizerConstructor',
    paramwise_cfg=dict(
        num_layers=37,
        layer_decay_rate=0.94,
        depths=[5, 5, 22, 5]))

lr_config_cfg = dict(
    policy='poly',
    warmup='linear',
    warmup_iters=1500,
    warmup_ratio=1e-6,
    power=1.0,
    min_lr=0.0,
    by_epoch=False)

runner_cfg = dict(type='IterBasedRunner', max_iters=MAX_ITERS)
checkpoint_config_cfg = dict(by_epoch=False, interval=CHECKPOINT_INTERVAL, max_keep_ckpts=20)
evaluation_cfg = dict(interval=EVAL_INTERVAL, metric='mIoU', save_best='mIoU')

log_config_cfg = dict(
    interval=50,
    hooks=[dict(type='TextLoggerHook', by_epoch=False)])

# %%
# 9. Assemble full configuration
config = Config(
    dict(
        model=model_cfg,
        data=data_cfg,
        optimizer=optimizer_cfg,
        optimizer_config=dict(),
        lr_config=lr_config_cfg,
        runner=runner_cfg,
        checkpoint_config=checkpoint_config_cfg,
        evaluation=evaluation_cfg,
        log_config=log_config_cfg,
        dist_params=dict(backend='nccl'),
        log_level='WARNING',
        load_from=LOAD_FROM,
        resume_from=RESUME_FROM,
        workflow=[('train', 1)],
        cudnn_benchmark=True,
        find_unused_parameters=False,
        work_dir=CHECKPOINT_DIR,
        gpu_ids=[0],
        device='cuda',
        seed=None),
    filename='train_segment_py.py')

config.dump(os.path.join(CHECKPOINT_DIR, 'config.py'))

print(f"Configuration saved to {CHECKPOINT_DIR}/config.py")
print(f"Experiment: {EXPERIMENT_NAME}")
print(f"Classes: {NUM_CLASSES}")
print(f"Max iterations: {MAX_ITERS}")

# %%
# 10. Initialize random seed
setup_multi_processes(config)
seed = init_random_seed(config.seed, device=config.device)
set_random_seed(seed, deterministic=False)
config.seed = seed

# %%
# 11. Build model
model = build_segmentor(config.model)

model.CLASSES = CLASSES
model.PALETTE = PALETTE
model.init_weights()

print("Model initialized with pretrained backbone")

# %%
# 12. Build datasets
goose_train_ds = build_dataset(goose_train_cfg)
gooseEx_train_ds = build_dataset(gooseEx_train_cfg)

if os.path.exists(GOOSE_COPYPASTE_TRAIN_DIR):
    goose_copypaste_train_ds = build_dataset(goose_copypaste_train_cfg)
    train_dataset = ConcatDataset([goose_train_ds, gooseEx_train_ds, goose_copypaste_train_ds], separate_eval=False)
    print(f"Train: {len(train_dataset)} (goose:{len(goose_train_ds)} + gooseEx:{len(gooseEx_train_ds)} + copy paste sample:{len(goose_copypaste_train_ds)})")
else:
    train_dataset = ConcatDataset([goose_train_ds, gooseEx_train_ds], separate_eval=False)
    print(f"Train: {len(train_dataset)} (goose:{len(goose_train_ds)} + gooseEx:{len(gooseEx_train_ds)})")

datasets = [train_dataset]

print(f"Val: will be built from config (goose+gooseEx)")

# %%
# 13. Start training
print("Starting training...")

train_segmentor(
    model,
    datasets,
    config,
    distributed=False,
    validate=True,
    timestamp=None,
    meta=dict(
        seed=seed,
        exp_name=EXPERIMENT_NAME,
        mmseg_version='0.27.0'))

print("Training finished.")