mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-07-22 17:52:07 +08:00
firestar5683
75 lines
2.6 KiB
Python
75 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
if __package__ in (None, ""):
|
|
import sys
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
from common import EXTENDED_CLASSIFIER_SPEED_VALUES # type: ignore # noqa: TID251
|
|
else:
|
|
from .common import EXTENDED_CLASSIFIER_SPEED_VALUES
|
|
|
|
|
|
DEFAULT_CLASSES = tuple(str(value) for value in EXTENDED_CLASSIFIER_SPEED_VALUES) + ("reject",)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Expand a YOLO classification head while preserving existing class rows.")
|
|
parser.add_argument("--model", type=Path, required=True, help="Existing Ultralytics classification checkpoint.")
|
|
parser.add_argument("--output", type=Path, required=True, help="Expanded checkpoint path.")
|
|
parser.add_argument("--classes", nargs="+", default=DEFAULT_CLASSES, help="New output classes in dataset index order.")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
from ultralytics import YOLO
|
|
|
|
args = parse_args()
|
|
model = YOLO(str(args.model.expanduser().resolve()))
|
|
old_names = {int(index): str(name) for index, name in model.names.items()}
|
|
new_names = dict(enumerate(args.classes))
|
|
if len(set(new_names.values())) != len(new_names):
|
|
raise ValueError("--classes contains duplicate names")
|
|
if not set(old_names.values()).issubset(new_names.values()):
|
|
missing = sorted(set(old_names.values()) - set(new_names.values()))
|
|
raise ValueError(f"Expanded classes omit existing outputs: {missing}")
|
|
|
|
head = model.model.model[-1]
|
|
old_linear = head.linear
|
|
new_linear = torch.nn.Linear(
|
|
old_linear.in_features,
|
|
len(new_names),
|
|
bias=old_linear.bias is not None,
|
|
device=old_linear.weight.device,
|
|
dtype=old_linear.weight.dtype,
|
|
)
|
|
new_index_by_name = {name: index for index, name in new_names.items()}
|
|
with torch.no_grad():
|
|
for old_index, name in old_names.items():
|
|
new_index = new_index_by_name[name]
|
|
new_linear.weight[new_index].copy_(old_linear.weight[old_index])
|
|
if old_linear.bias is not None:
|
|
new_linear.bias[new_index].copy_(old_linear.bias[old_index])
|
|
|
|
head.linear = new_linear
|
|
head.np = sum(parameter.numel() for parameter in head.parameters())
|
|
model.model.names = new_names
|
|
model.model.yaml["nc"] = len(new_names)
|
|
model.model.args["classes"] = None
|
|
|
|
output = args.output.expanduser().resolve()
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
model.save(str(output))
|
|
print(f"Expanded {len(old_names)} outputs to {len(new_names)}: {output}")
|
|
print(new_names)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|