rmbg — fine-tune anime-seg bằng LoRA để tách nền

Fine-tune ISNetDIS (checkpoint skytnt/anime-seg) bằng LoRA trên 1306 cặp ảnh–mask, rồi chạy inference hàng loạt để xuất cutout RGBA.

Trên val split 130 ảnh, IoU trung bình đi từ 0.8449 (checkpoint gốc) lên 0.9523, và số ca thất bại nặng (IoU<0.5) từ 18 xuống 3 — chi tiết trong docs/experiments.md.

Quickstart

git clone <repo> && cd Fine_tune_rm_bg
./scripts/setup_env.sh            # submodule + venv + cài package + chạy test

# đặt dataset vào data/raw/{color_images,mask_images}/ rồi:
make split                        # sinh data/splits/v1.json
make eval                         # kiểm tra checkpoint có sẵn

Nếu dataset nằm ở ổ khác, symlink hoặc set RMBG_DATA_ROOT — xem .env.example.

Các lệnh

Mọi thao tác đi qua một entry point rmbg:

# fine-tune LoRA (cấu hình đầy đủ trong configs/train_lora_r16.yaml)
rmbg train --config configs/train_lora_r16.yaml --steps 6000

# full fine-tune
rmbg train --config configs/train_full.yaml

# fold adapter LoRA vào một checkpoint đầy đủ
rmbg merge --lora artifacts/runs/<run_id>/lora_best.pt \
           --out artifacts/checkpoints/<tên>.ckpt

# đo IoU/MAE/solidity trên mask thô của model
rmbg eval --ckpt artifacts/checkpoints/merged_r16.ckpt --split data/splits/v1.json
rmbg eval --ckpt pretrained --split data/splits/v1.json      # baseline

# inference hàng loạt: cutout RGBA + mask 8-bit (+ ảnh so sánh)
rmbg predict --src data/raw/color_images \
             --dst artifacts/predictions/thu_nghiem \
             --gt data/raw/mask_images --compare

# dựng lại summary từ metrics.csv có sẵn, không chạy model
rmbg report artifacts/predictions/thu_nghiem

# sinh split cố định
rmbg make-split --images data/raw/color_images --val-n 130 --seed 42 \
                --out data/splits/v1.json

make help liệt kê các shortcut tương ứng.

Tên cờ giữ nguyên từ hai script cũ (--img-size, --bs, --acc, --val-n, --no-alpha-gain, ...). Thứ tự ưu tiên cấu hình: mặc định trong code < file YAML (--config) < cờ CLI.

Bố cục

configs/            base.yaml + 3 file kế thừa nó qua `extends`
src/rmbg/           2.7k dòng, 25 file
  paths.py          đường dẫn gốc, đọc từ env
  vendor.py         chỗ DUY NHẤT sửa sys.path
  upstream.py       bề mặt DUY NHẤT tới code bên thứ ba
  config.py         dataclass + nạp YAML (có extends) + override
  metrics.py        nơi DUY NHẤT định nghĩa IoU/MAE/solidity
  postprocess.py    alpha_gain, cutout RGBA, ảnh so sánh
  runtime.py        bộ lọc warning, seed
  provenance.py     git SHA + phiên bản môi trường cho manifest
  model/            lora.py (adapter), loader.py (nạp checkpoint)
  data/             dataset, split, liệt kê file, đọc ảnh
  engine/           trainer, evaluator, merger, predictor, reporter
  cli/main.py       argparse + dispatch (import nặng lười)
tests/              91 test, chạy trên CPU trong ~3s, không cần dataset
docs/               experiments, dataset, architecture, failure_cases
scripts/            setup_env.sh, reproduce_best.sh
third_party/        submodule anime-segmentation @55d8740
data/               [gitignore] raw/, splits/*.json thì commit
artifacts/          [gitignore] checkpoints/, runs/, predictions/

Thư mục data/ và artifacts/ không vào git (dataset 150 MB, checkpoint 200 MB, kết quả inference 1.1 GB). Ngoại lệ duy nhất: data/splits/*.json — file split phải được commit, nếu không thì không so sánh được số liệu giữa các run.

configs/base.yaml giữ phần dùng chung (nguồn dữ liệu, ngưỡng metric, hậu xử lý); ba file còn lại khai extends: base.yaml và chỉ ghi phần khác biệt của chúng. File con ghi đè file cha ở từng khoá lá, không thay cả section.

Đọc số liệu cho đúng

metrics.csv có hai bộ cột:

• iou, mae, pred_fg — đo trên mask thô của model. Dùng bộ này để đánh giá model và so sánh giữa các run.
• iou_post, mae_post, pred_fg_post — đo sau khi kéo giãn alpha (alpha_gain). Dùng bộ này để mô tả chất lượng ảnh xuất ra thực tế.

Ảnh trong cutout/ và mask/ luôn dùng alpha đã hậu xử lý.

Trước refactor chỉ tồn tại bộ _post nhưng lại được đặt tên như thể đo model, nên số IoU của batch test không so được với số IoU in trong lúc train. Chênh lệch là thật: MAE 0.0231 (thô) so với 0.0173 (post) trên cùng 1306 ảnh.

Tái lập

Mỗi lần train tạo artifacts/runs/<run_id>/ gồm:

File	Nội dung
run.json	cấu hình đầy đủ, git SHA của repo và submodule, phiên bản torch/CUDA/GPU, thông tin split, lịch sử val
train_log.csv	loss và lr theo step
val_log.csv	metric val theo step
lora_best.pt / full_best.ckpt	checkpoint tốt nhất, có nhúng manifest

Tái lập checkpoint hiện tại: ./scripts/reproduce_best.sh.

Phiên bản thư viện chính xác của môi trường đã chạy được: requirements.lock.

Phát triển

make check      # ruff check + pytest — chạy trước mỗi commit
make fmt        # ruff format + auto-fix

Quy ước code, cách thêm siêu tham số, cách đổi dataset, cách ghi kết quả thực nghiệm: CONTRIBUTING.md.

91 test chạy trên CPU trong ~3 giây, không tải model và không cần dataset. Chúng bảo vệ những chỗ dễ hỏng nhất: phép fold weight của LoRA (kể cả conv dilated), việc adapter lưu kèm BatchNorm, các ca biên của metric, quy ước RGBA straight, và việc manifest nhúng trong checkpoint phải nạp lại được.

Có sẵn .pre-commit-config.yaml (ruff + chặn file lớn) và .github/workflows/ci.yml (lint + test trên CPU). CI chưa từng chạy thật vì repo còn là git local — theo dõi lần push đầu.

Đọc thêm: docs/architecture.md giải thích các quyết định thiết kế; REFACTOR_GUIDE.md là bản review và kế hoạch refactor đã thực hiện.

Giấy phép

Model và code train gốc: SkyTNT/anime-segmentation (pin ở commit 55d8740), giấy phép Apache-2.0. Checkpoint gốc: skytnt/anime-seg.

Code trong src/, tests/, configs/, scripts/ được phát hành theo Apache License 2.0 (chủ dự án cần điền thông tin copyright vào phần APPENDIX của file LICENSE nếu muốn).