Spaces:
Running on Zero
Running on Zero
| import argparse | |
| import json | |
| from pathlib import Path | |
| import torch | |
| from PIL import Image | |
| from torch.utils.data import DataLoader, Dataset | |
| from tqdm import tqdm | |
| from transformers import RTDetrForObjectDetection, RTDetrImageProcessor | |
| BASE_MODEL = "PekingU/rtdetr_r50vd" | |
| def classes(path): | |
| return [x.strip() for x in Path(path).read_text().splitlines() if x.strip()] | |
| class COCODataset(Dataset): | |
| def __init__(self, image_dir, ann_file, processor): | |
| self.image_dir = Path(image_dir) | |
| self.processor = processor | |
| coco = json.loads(Path(ann_file).read_text()) | |
| cats = sorted(coco["categories"], key=lambda x: x["id"]) | |
| self.cat_map = {c["id"]: i for i, c in enumerate(cats)} | |
| anns = {} | |
| for a in coco["annotations"]: | |
| anns.setdefault(a["image_id"], []).append(a) | |
| self.records = [] | |
| for im in coco["images"]: | |
| self.records.append({ | |
| "id": im["id"], | |
| "file_name": im["file_name"], | |
| "annotations": anns.get(im["id"], []), | |
| }) | |
| def __len__(self): | |
| return len(self.records) | |
| def __getitem__(self, i): | |
| r = self.records[i] | |
| image = Image.open(self.image_dir / r["file_name"]).convert("RGB") | |
| anns = [] | |
| for a in r["annotations"]: | |
| x, y, w, h = a["bbox"] | |
| anns.append({ | |
| "id": a["id"], | |
| "image_id": r["id"], | |
| "category_id": self.cat_map[a["category_id"]], | |
| "bbox": [x, y, w, h], | |
| "area": float(a["area"]), | |
| "iscrowd": 0, | |
| }) | |
| out = self.processor( | |
| images=image, | |
| annotations={ | |
| "image_id": r["id"], | |
| "annotations": anns, | |
| }, | |
| return_tensors="pt", | |
| ) | |
| out["pixel_values"] = out["pixel_values"].squeeze(0) | |
| if "pixel_mask" in out: | |
| out["pixel_mask"] = out["pixel_mask"].squeeze(0) | |
| out["labels"] = out["labels"][0] | |
| return out | |
| def collate(batch): | |
| result = { | |
| "pixel_values": torch.stack([x["pixel_values"] for x in batch]), | |
| "labels": [x["labels"] for x in batch], | |
| } | |
| if "pixel_mask" in batch[0]: | |
| result["pixel_mask"] = torch.stack([x["pixel_mask"] for x in batch]) | |
| return result | |
| 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", required=True) | |
| p.add_argument("--epochs", type=int, default=30) | |
| p.add_argument("--batch-size", type=int, default=2) | |
| args = p.parse_args() | |
| names = classes(args.classes) | |
| id2label = {i: n for i, n in enumerate(names)} | |
| label2id = {n: i for i, n in enumerate(names)} | |
| processor = RTDetrImageProcessor.from_pretrained(BASE_MODEL) | |
| train = COCODataset( | |
| Path(args.train_dir)/"images", | |
| Path(args.train_dir)/"annotations.json", | |
| processor, | |
| ) | |
| val = COCODataset( | |
| Path(args.val_dir)/"images", | |
| Path(args.val_dir)/"annotations.json", | |
| processor, | |
| ) | |
| model = RTDetrForObjectDetection.from_pretrained( | |
| BASE_MODEL, | |
| num_labels=len(names), | |
| id2label=id2label, | |
| label2id=label2id, | |
| ignore_mismatched_sizes=True, | |
| ) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model.to(device) | |
| train_loader = DataLoader( | |
| train, batch_size=args.batch_size, shuffle=True, | |
| num_workers=0, collate_fn=collate | |
| ) | |
| val_loader = DataLoader( | |
| val, batch_size=args.batch_size, shuffle=False, | |
| num_workers=0, collate_fn=collate | |
| ) | |
| optimizer = torch.optim.AdamW( | |
| model.parameters(), lr=1e-5, weight_decay=1e-4 | |
| ) | |
| outdir = Path(args.output_dir) | |
| outdir.mkdir(parents=True, exist_ok=True) | |
| best = float("inf") | |
| for epoch in range(args.epochs): | |
| model.train() | |
| running = 0.0 | |
| for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}"): | |
| batch = { | |
| k: (v.to(device) if torch.is_tensor(v) else v) | |
| for k, v in batch.items() | |
| } | |
| loss = model(**batch).loss | |
| loss.backward() | |
| optimizer.step() | |
| optimizer.zero_grad(set_to_none=True) | |
| running += float(loss.item()) | |
| model.eval() | |
| val_loss = 0.0 | |
| n = 0 | |
| with torch.no_grad(): | |
| for batch in val_loader: | |
| batch = { | |
| k: (v.to(device) if torch.is_tensor(v) else v) | |
| for k, v in batch.items() | |
| } | |
| val_loss += float(model(**batch).loss.item()) | |
| n += 1 | |
| val_loss /= max(n, 1) | |
| print(f"epoch={epoch+1} validation_loss={val_loss:.4f}") | |
| if val_loss < best: | |
| best = val_loss | |
| model.save_pretrained(outdir) | |
| processor.save_pretrained(outdir) | |
| model.save_pretrained(outdir) | |
| processor.save_pretrained(outdir) | |
| if __name__ == "__main__": | |
| main() | |