Files
u52f85420 04bf90c09b feat: 导入 IRSTD 研究代码与文档(小目标检测)
- 新增 irstd 核心包、scripts、tests、configs
- 新增研究文档:literature_review / mathematical_analysis / novelty_matrix / experiment_plan / research_decision / research_log / manuscript_outline
- 新增 evidence、results 数据与审计产物
- 新增 pyproject.toml、requirements、.gitignore
2026-11-01 10:10:33 +00:00

21 lines
967 B
Python

import argparse,json
from pathlib import Path
import numpy as np
import torch
from irstd.models import TinyUNet
from irstd.data import read_manifest,read_pair,sha
def main():
p=argparse.ArgumentParser();p.add_argument('--manifest',required=True);p.add_argument('--checkpoint',required=True)
p.add_argument('--out',required=True);p.add_argument('--device',default='cuda');a=p.parse_args()
ck=torch.load(a.checkpoint,map_location='cpu',weights_only=True)
model=TinyUNet().to(a.device);model.load_state_dict(ck['state_dict']);model.eval()
out=Path(a.out);out.mkdir(parents=True,exist_ok=False)
with torch.no_grad():
for r in read_manifest(a.manifest):
x,_=read_pair(r)
prob=model(torch.from_numpy(x)[None,None].to(a.device)).sigmoid()[0,0].cpu().numpy()
np.save(out/(r['id']+'.npy'),prob)
(out/'provenance.json').write_text(json.dumps(dict(checkpoint_sha256=sha(a.checkpoint),manifest_sha256=sha(a.manifest)),indent=2))
if __name__=='__main__':main()