training/train.py

import argparse
import json
from pathlib import Path

import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
from transformers import RTDetrImageProcessor, RTDetrForObjectDetection

BASE_MODEL = "PekingU/rtdetr_r50vd"

def load_classes(path):
    return [x.strip() for x in Path(path).read_text().splitlines() if x.strip()]

class COCODetectionDataset(Dataset):
    def __init__(self, image_dir, annotation_file, processor):
        self.image_dir = Path(image_dir)
        self.processor = processor
        coco = json.loads(Path(annotation_file).read_text())
        self.images = {x["id"]: x for x in coco["images"]}
        cats = sorted(coco["categories"], key=lambda x: x["id"])
        self.category_id_to_label = {c["id"]: i for i,c in enumerate(cats)}
        anns = {}
        for a in coco["annotations"]:
            if not a.get("iscrowd", 0):
                anns.setdefault(a["image_id"], []).append(a)
        self.records = []
        for image_id, info in self.images.items():
            self.records.append({
                "image_id": image_id, "file_name": info["file_name"],
                "width": info["width"], "height": info["height"],
                "annotations": anns.get(image_id, [])
            })

    def __len__(self): return len(self.records)

    def __getitem__(self, idx):
        r = self.records[idx]
        image = Image.open(self.image_dir / r["file_name"]).convert("RGB")
        anns = []
        for a in r["annotations"]:
            x,y,w,h = a["bbox"]
            if w <= 0 or h <= 0: continue
            anns.append({
                "id": a["id"], "image_id": int(idx),
                "category_id": self.category_id_to_label[a["category_id"]],
                "bbox": [x,y,w,h], "area": float(a.get("area",w*h)),
                "iscrowd": 0
            })
        encoded = self.processor(
            images=image,
            annotations={"image_id": int(idx), "annotations": anns},
            return_tensors="pt"
        )
        encoded["pixel_values"] = encoded["pixel_values"].squeeze(0)
        if "pixel_mask" in encoded:
            encoded["pixel_mask"] = encoded["pixel_mask"].squeeze(0)
        encoded["labels"] = encoded["labels"][0]
        return encoded

def collate_fn(batch):
    out = {"pixel_values": torch.stack([x["pixel_values"] for x in batch]),
           "labels": [x["labels"] for x in batch]}
    if "pixel_mask" in batch[0]:
        out["pixel_mask"] = torch.stack([x["pixel_mask"] for x in batch])
    return out

def move_to_device(obj, device):
    if torch.is_tensor(obj):
        return obj.to(device)
    if isinstance(obj, dict):
        return {k: move_to_device(v, device) for k, v in obj.items()}
    if isinstance(obj, list):
        return [move_to_device(v, device) for v in obj]
    if isinstance(obj, tuple):
        return tuple(move_to_device(v, device) for v in obj)
    return obj

def evaluate(model, loader, device):
    model.eval(); total=0; n=0
    with torch.no_grad():
        for batch in loader:
            batch=move_to_device(batch, device)
            total += float(model(**batch).loss.item()); n += 1
    model.train()
    return total/max(n,1)


def patch_rtdetr_denoising_device():
    """Patch RT-DETR denoising so embedding indices are always on the embedding device.

    Important: the wrapper must be an nn.Module, not a plain Python function.
    Transformers may inspect/use class_embed as a module, and a plain function
    loses the original module parameters/device information.
    """
    try:
        import transformers.models.rt_detr.modeling_rt_detr as rtdetr_mod
        import torch.nn as nn

        original = rtdetr_mod.get_contrastive_denoising_training_group
        if getattr(original, "_icecream_device_patch", False):
            return

        class DeviceSafeEmbedding(nn.Module):
            def __init__(self, embedding):
                super().__init__()
                self.embedding = embedding

            def forward(self, indices):
                if torch.is_tensor(indices):
                    indices = indices.to(self.embedding.weight.device)
                return self.embedding(indices)

        def wrapped(targets, num_classes, num_queries, class_embed,
                    num_denoising_queries=100, label_noise_ratio=0.5,
                    box_noise_scale=1.0, **kwargs):
            # Move every target tensor to the embedding/model device.
            try:
                embed_device = class_embed.weight.device
            except Exception:
                try:
                    embed_device = next(class_embed.parameters()).device
                except Exception:
                    embed_device = None

            if embed_device is not None:
                fixed_targets = []
                for target in targets:
                    if isinstance(target, dict):
                        target = dict(target)
                        for key, value in list(target.items()):
                            if torch.is_tensor(value):
                                target[key] = value.to(embed_device)
                    fixed_targets.append(target)
                targets = fixed_targets
                class_embed = DeviceSafeEmbedding(class_embed)

            return original(
                targets=targets,
                num_classes=num_classes,
                num_queries=num_queries,
                class_embed=class_embed,
                num_denoising_queries=num_denoising_queries,
                label_noise_ratio=label_noise_ratio,
                box_noise_scale=box_noise_scale,
                **kwargs,
            )

        wrapped._icecream_device_patch = True
        rtdetr_mod.get_contrastive_denoising_training_group = wrapped
        print("RT-DETR denoising device patch installed")
    except Exception as exc:
        raise RuntimeError(f"Could not install RT-DETR denoising device patch: {exc}") from exc

def main():
    p=argparse.ArgumentParser()
    p.add_argument("--train-dir",required=True); p.add_argument("--val-dir",required=True)
    p.add_argument("--classes",required=True); p.add_argument("--output-dir",default="model")
    p.add_argument("--epochs",type=int,default=30); p.add_argument("--batch-size",type=int,default=2)
    p.add_argument("--learning-rate",type=float,default=1e-5); p.add_argument("--weight-decay",type=float,default=1e-4)
    p.add_argument("--num-workers",type=int,default=2)
    a=p.parse_args()

    classes=load_classes(a.classes)
    id2label={i:n for i,n in enumerate(classes)}
    label2id={n:i for i,n in enumerate(classes)}

    proc=RTDetrImageProcessor.from_pretrained(BASE_MODEL)
    train=COCODetectionDataset(Path(a.train_dir)/"images",Path(a.train_dir)/"annotations.json",proc)
    val=COCODetectionDataset(Path(a.val_dir)/"images",Path(a.val_dir)/"annotations.json",proc)

    if len(train)==0 or len(val)==0:
        raise ValueError("Training and validation datasets must contain at least one image.")
    if len(train.category_id_to_label)!=len(classes) or len(val.category_id_to_label)!=len(classes):
        raise ValueError("COCO categories do not match classes.txt. Rebuild the dataset after saving the classes.")

    model=RTDetrForObjectDetection.from_pretrained(
        BASE_MODEL,num_labels=len(classes),id2label=id2label,label2id=label2id,
        ignore_mismatched_sizes=True
    )
    # Some Transformers RT-DETR releases enter the denoising path whenever
    # training, regardless of num_denoising. Keep the config disabled where
    # supported, but also install the device-safe embedding patch below.
    for cfg_owner in (model, getattr(model, "model", None)):
        cfg = getattr(cfg_owner, "config", None)
        if cfg is not None:
            for name in ("num_denoising", "num_denoising_queries"):
                if hasattr(cfg, name):
                    setattr(cfg, name, 0)
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    patch_rtdetr_denoising_device()

    tr=DataLoader(train,batch_size=a.batch_size,shuffle=True,num_workers=0,collate_fn=collate_fn)
    va=DataLoader(val,batch_size=a.batch_size,shuffle=False,num_workers=0,collate_fn=collate_fn)
    opt=torch.optim.AdamW(model.parameters(),lr=a.learning_rate,weight_decay=a.weight_decay)

    outdir=Path(a.output_dir); outdir.mkdir(parents=True,exist_ok=True)
    best=float("inf")

    for epoch in range(a.epochs):
        model.train(); running=0
        bar=tqdm(tr,desc=f"epoch {epoch+1}/{a.epochs}")
        for step,batch in enumerate(bar):
            batch=move_to_device(batch, device)
            # RT-DETR's loss matcher uses nested target tensors (boxes/classes).
            # Move every tensor in labels to the same device as the model.
            if "labels" in batch:
                # RT-DETR expects every nested target tensor on the same device as the model.
                for target in batch["labels"]:
                    if isinstance(target, dict):
                        for key, value in list(target.items()):
                            if torch.is_tensor(value):
                                target[key] = value.to(device)
            loss=model(**batch).loss
            loss.backward(); opt.step(); opt.zero_grad(set_to_none=True)
            running += float(loss.item())
            bar.set_postfix(loss=f"{running/(step+1):.4f}")
        vl=evaluate(model,va,device)
        print(f"validation_loss={vl:.4f}")
        if vl<best:
            best=vl
            model.save_pretrained(outdir)
            proc.save_pretrained(outdir)
            (outdir/"classes.json").write_text(json.dumps({"id2label":id2label,"label2id":label2id},indent=2))
    model.save_pretrained(outdir); proc.save_pretrained(outdir)

if __name__=="__main__": main()