update main.py

main.py

@@ -15,7 +15,7 @@ sys.path.insert(0, str(Path(__file__).parent))
 from train.train import main as train_main, load_config
 from train.eval import evaluate_model
 from train.inference import inference_single_image
-from model.CASWiT import CASWiT
+from model.build_model import build_model
 from dataset.definition_dataset import build_transforms
 
 
@@ -60,13 +60,7 @@ def main():
             print(f"Error: Checkpoint file not found: {args.checkpoint}")
             sys.exit(1)
         
-        model = CASWiT(
-            num_head_xa=cfg.cross_attention_heads,
-            num_classes=cfg.num_classes,
-            model_name=cfg.model_name,
-            mlp_ratio=cfg.fusion_mlp_ratio,
-            drop_path=cfg.fusion_drop_path
-        ).to(device)
+        model = build_model(cfg).to(device)
         
         print(f"Loading checkpoint from: {args.checkpoint}")
         state_dict = torch.load(args.checkpoint, map_location=device)
@@ -81,7 +75,7 @@ def main():
             print(f"   Perfect match! All weights loaded successfully.")
         
         transform = build_transforms()
-        inference_single_image(model, args.image, device, transform, args.output)
+        inference_single_image(model, args.image, device, transform, cfg, args.output)
 
 
 if __name__ == "__main__":