{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":126777,"databundleVersionId":15314950}],"dockerImageVersionId":31259,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"pip install transformers","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T18:32:05.114269Z","iopub.execute_input":"2026-03-01T18:32:05.114872Z","iopub.status.idle":"2026-03-01T18:32:08.301517Z","shell.execute_reply.started":"2026-03-01T18:32:05.114838Z","shell.execute_reply":"2026-03-01T18:32:08.300749Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pip install ultralytics","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T18:32:08.303309Z","iopub.execute_input":"2026-03-01T18:32:08.303669Z","iopub.status.idle":"2026-03-01T18:32:11.557022Z","shell.execute_reply.started":"2026-03-01T18:32:08.303637Z","shell.execute_reply":"2026-03-01T18:32:11.556263Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport sys\nfrom PIL import Image,ImageFilter\nimport requests\nfrom torch import nn\nimport pandas as pd\nimport random\nimport os\nimport csv\nimport numpy as np\nimport json\nimport timm\nimport math\nimport argparse\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset\nfrom ultralytics.utils import TQDM\nimport torch\nfrom torchvision import transforms\nfrom torch.utils.data import DataLoader\nimport matplotlib.pyplot as plt\nimport requests\nfrom sklearn.metrics.pairwise import cosine_similarity\nfrom sklearn.metrics import accuracy_score\nfrom transformers import (\n    ViTImageProcessor,\n    ViTForImageClassification,\n    AutoConfig,\n    BertForSequenceClassification,\n    AutoTokenizer,\n    EvalPrediction,\n    Trainer,\n    TrainingArguments,\n    default_data_collator,\n    set_seed,\n    get_scheduler\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T18:32:11.558277Z","iopub.execute_input":"2026-03-01T18:32:11.558537Z","iopub.status.idle":"2026-03-01T18:32:11.565220Z","shell.execute_reply.started":"2026-03-01T18:32:11.558506Z","shell.execute_reply":"2026-03-01T18:32:11.564459Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class GeM(nn.Module):\n    def __init__(self, p=3, eps=1e-6):\n        super().__init__()\n        self.p = nn.Parameter(torch.ones(1) * p)\n        self.eps = eps\n    def forward(self, x):\n        return F.avg_pool2d(x.clamp(min=self.eps).pow(self.p), (x.size(-2), x.size(-1))).pow(1.0 / self.p)\n\nclass SubCenterArcFace(nn.Module):\n    def __init__(self, in_features, out_features, s=30.0, m=0.5, k=1):\n        super().__init__()\n        self.s = s\n        self.m = m\n        self.k = k\n        self.weight = nn.Parameter(torch.FloatTensor(out_features * k, in_features))\n        nn.init.xavier_uniform_(self.weight)\n        self.out_features = out_features\n    def forward(self, input, label=None):\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n        cosine = cosine.view(-1, self.out_features, self.k)# B, num_classes, K\n        cosine, _ = cosine.max(dim=2)# B, num_classes\n        if label is None: return cosine\n        phi = cosine - self.m\n        one_hot = torch.zeros_like(cosine)# B, num_classes\n        one_hot.scatter_(1, label.view(-1, 1), 1)\n        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)# B, num_classes\n        return output * self.s# B, num_classes\n\nclass EVAModel(nn.Module):\n    def __init__(\n        self, \n        num_classes,\n        arcface_s,\n        arcface_m,\n        num_subcenters\n    ):\n        super().__init__()\n        # Load local or download if allowed\n        self.backbone = timm.create_model(\"eva02_large_patch14_448.mim_m38m_ft_in22k_in1k\", pretrained=True, num_classes=0)\n        self.backbone.set_grad_checkpointing(True)\n        self.feat_dim = self.backbone.num_features\n        self.gem = GeM(p=3)\n        self.bn = nn.BatchNorm1d(self.feat_dim)\n        self.head = SubCenterArcFace(self.feat_dim, num_classes, s=arcface_s, m=arcface_m, k=num_subcenters)\n\n    def forward(self, x, labels=None):\n        features = self.backbone.forward_features(x)\n        if features.dim() == 3:\n            B, N, C = features.shape# B, 1025, 1024\n            H = W = int(math.sqrt(N))# 32, 32\n            if H * W != N: features = features[:, -H * W :, :]# B, 1024, 1024\n            # [B, N=1024, C=1024]->[B, C=1024, N=1024]->B,1024,32,32\n            features = features.permute(0, 2, 1).reshape(B, C, H, W)\n        emb = self.gem(features).flatten(1)#B,1024,1,1\n        emb = self.bn(emb)#B,1024\n        if labels is not None: \n            logits=self.head(emb, labels)#B, num_classes\n            loss = F.cross_entropy(logits, labels)\n            return {\"loss\": loss, \"logits\": logits}\n        return {\"logits\": F.normalize(emb)}\n\n    def reinit_head(self, new_num_classes):\n        # Create new head with expanded classes\n        old_weight = self.head.weight.data\n        self.head = SubCenterArcFace(self.feat_dim, new_num_classes, s=Config.arcface_s, m=Config.arcface_m, k=Config.num_subcenters)\n        # Copy old weights to preserve knowledge\n        # (This is tricky if we just append, but easiest is to re-init xavier for new ones)\n        # Actually, let's keep it simple: Re-init all. Fine-tuning will adapt quickly.","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T18:32:11.566885Z","iopub.execute_input":"2026-03-01T18:32:11.567107Z","iopub.status.idle":"2026-03-01T18:32:11.586168Z","shell.execute_reply.started":"2026-03-01T18:32:11.567085Z","shell.execute_reply":"2026-03-01T18:32:11.585508Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def seed_everything(seed):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T18:32:11.586951Z","iopub.execute_input":"2026-03-01T18:32:11.587182Z","iopub.status.idle":"2026-03-01T18:32:11.599647Z","shell.execute_reply.started":"2026-03-01T18:32:11.587159Z","shell.execute_reply":"2026-03-01T18:32:11.599083Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TrainDataset(Dataset): \n    def __init__(self,transform=None): \n        self.train_df=pd.read_csv(r'/kaggle/input/jaguar-re-id/train.csv') \n        self.root=r'/kaggle/input/jaguar-re-id/train/train' \n        self.transform=transform\n        self.img_names=self.train_df['filename'].tolist() \n        self.labels=self.train_df['ground_truth'].tolist() \n        self.classes=self.train_df['ground_truth'].unique().tolist() \n    def __len__(self): \n        return len(self.img_names) \n    def __getitem__(self,idx): \n        img_name=self.img_names[idx] \n        img_path=os.path.join(self.root,img_name) \n        img=Image.open(img_path).convert(\"RGB\")\n        if self.transform is not None: \n            img=self.transform(img)\n        label=self.classes.index(self.labels[idx])\n        return {\n            'x': img,\n            'label': torch.tensor(label, dtype=torch.long)\n        }\n\ndef dataset_split(dataset, ratio=0.8): \n    train_size = int(len(dataset) * ratio) \n    val_size = len(dataset) - train_size \n    train_iter, eval_iter = torch.utils.data.random_split(dataset, [train_size, val_size]) \n    return train_iter, eval_iter","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T18:32:11.600535Z","iopub.execute_input":"2026-03-01T18:32:11.600878Z","iopub.status.idle":"2026-03-01T18:32:11.614462Z","shell.execute_reply.started":"2026-03-01T18:32:11.600850Z","shell.execute_reply":"2026-03-01T18:32:11.613865Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SharpenTransform:\n    def __init__(self, p=0.3): \n        self.p = p\n    def __call__(self, img):\n        if random.random() < self.p: \n            return img.filter(ImageFilter.SHARPEN)\n        return img\ntrain_transform = transforms.Compose([\n    transforms.Resize((448, 448)),\n    SharpenTransform(p=0.3),\n    transforms.RandomHorizontalFlip(),\n    transforms.RandomAffine(degrees=15, translate=(0.1, 0.1), scale=(0.9, 1.1)),\n    transforms.ColorJitter(brightness=0.2, contrast=0.2),\n    transforms.ToTensor(),\n    transforms.Normalize([0.481, 0.457, 0.408], [0.268, 0.261, 0.275]),\n    transforms.RandomErasing(p=0.25),\n])\n\ntest_transform = transforms.Compose([\n    transforms.Resize((448, 448)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.481, 0.457, 0.408], [0.268, 0.261, 0.275]),\n])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T18:32:11.615394Z","iopub.execute_input":"2026-03-01T18:32:11.615674Z","iopub.status.idle":"2026-03-01T18:32:11.628596Z","shell.execute_reply.started":"2026-03-01T18:32:11.615650Z","shell.execute_reply":"2026-03-01T18:32:11.628042Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def parse_args():\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--model', type=str, default='/kaggle/working/experiments/checkpoint-240')\n    parser.add_argument('--output_dir', type=str, default='/kaggle/working/experiments')\n    parser.add_argument('--epoch', type=int, default=10)\n    parser.add_argument('--lr', type=float, default=2e-5)\n    parser.add_argument('--bs', type=int, default=4)\n    parser.add_argument('--seed', type=int, default=42, help='random seed')\n    parser.add_argument('--scheduler', type=str, default='cosine', help='consine/linear/cosine_with_restarts/constant')\n    parser.add_argument('--p', type=int, default=3)\n    parser.add_argument('--arcface_s', type=int, default=30)\n    parser.add_argument('--arcface_m', type=float, default=0.5)\n    parser.add_argument('--num_subcenters', type=int, default=1)\n    args ,_= parser.parse_known_args()\n    return args","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T18:32:11.629482Z","iopub.execute_input":"2026-03-01T18:32:11.629854Z","iopub.status.idle":"2026-03-01T18:32:11.642450Z","shell.execute_reply.started":"2026-03-01T18:32:11.629829Z","shell.execute_reply":"2026-03-01T18:32:11.641621Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def main():\n    args = parse_args()\n    print(\"args\",args)\n    print('====Input Arguments====')\n    print(json.dumps(vars(args), indent=2, sort_keys=False))\n    \n    # Set seed before initializing model, for reproduction purpose.\n    set_seed(args.seed)\n    seed_everything(args.seed)\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    num_classes=31\n    model = EVAModel(\n        num_classes,\n        args.arcface_s,\n        args.arcface_m,\n        args.num_subcenters\n    ).to(device)\n    # Load data\n    dataset = TrainDataset(transform=train_transform)\n    train_iter,eval_iter=dataset_split(dataset,ratio=1)\n    \n    training_args = TrainingArguments(\n        output_dir = args.output_dir,\n        overwrite_output_dir=True,\n        do_train=True,\n        do_eval=False,\n        do_predict=True,\n        eval_strategy='no',\n        per_device_train_batch_size=args.bs,\n        per_device_eval_batch_size=args.bs,\n        save_strategy='epoch',\n        save_total_limit=5,\n        save_only_model = True,\n        logging_strategy=\"epoch\",\n        num_train_epochs=args.epoch,\n        learning_rate= args.lr,\n        lr_scheduler_type=args.scheduler,\n        load_best_model_at_end=False,\n        remove_unused_columns=False,\n        report_to=\"none\",\n        fp16=True\n    )\n    # Initialize our Trainer\n    trainer = Trainer(\n        model=model,\n        args=training_args,\n        train_dataset=train_iter,\n        # eval_dataset=eval_iter,\n        data_collator=default_data_collator,\n    )\n    # batch = next(iter(trainer.get_train_dataloader()))\n    # print(batch.keys())\n    # Training\n    if training_args.do_train:\n        print(f'开始训练')\n        # train_result = trainer.train(resume_from_checkpoint=True)\n        train_result = trainer.train()\n        metrics = train_result.metrics\n    \n    #Predicting\n    print(f'开始测试')\n    # Load test data\n    test_df = pd.read_csv('/kaggle/input/jaguar-re-id/test.csv')\n    # Get unique image filenames\n    unique_images = sorted(set(test_df['query_image']) | set(test_df['gallery_image']))\n    n_images = len(unique_images)  # Should be 371\n    # Extract embeddings for all images\n    model.eval()\n    embeddings = []\n    pbar = TQDM(\n            unique_images,\n            desc=\"Predicting\",\n            total=len(unique_images),\n            unit=\"batch\",\n        )\n    for img_file in pbar:\n        img = Image.open(f'/kaggle/input/jaguar-re-id/test/test/{img_file}').convert(\"RGB\")\n        img = test_transform(img).unsqueeze(0).to(device)\n        emb = model(img)[\"logits\"]         \n        emb = emb.squeeze(0)           \n        emb = emb.detach().cpu()       \n        embeddings.append(emb)\n\n    pbar.close()\n    embeddings = torch.stack(embeddings).numpy()\n    # Compute similarity matrix (371 × 371)\n    similarity_matrix = cosine_similarity(embeddings)\n    # Ensure values are in [0, 1]\n    similarity_matrix = np.clip(similarity_matrix, 0, 1)\n    # Create filename to index mapping\n    img_to_idx = {img: idx for idx, img in enumerate(unique_images)}\n    # Extract similarities for each test pair\n    similarities = []\n    for _, row in test_df.iterrows():\n        query_idx = img_to_idx[row['query_image']]\n        gallery_idx = img_to_idx[row['gallery_image']]\n        sim = similarity_matrix[query_idx, gallery_idx]\n        similarities.append(sim)\n    # Create submission\n    submission = pd.DataFrame({\n        'row_id': test_df['row_id'],\n        'similarity': similarities\n    })\n    submission.to_csv('/kaggle/working/submission.csv', index=False)\n    print(\"保存成功\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T18:32:11.643344Z","iopub.execute_input":"2026-03-01T18:32:11.643943Z","iopub.status.idle":"2026-03-01T18:32:11.659258Z","shell.execute_reply.started":"2026-03-01T18:32:11.643918Z","shell.execute_reply":"2026-03-01T18:32:11.658733Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T18:32:11.661269Z","iopub.execute_input":"2026-03-01T18:32:11.661536Z","iopub.status.idle":"2026-03-01T23:17:49.986822Z","shell.execute_reply.started":"2026-03-01T18:32:11.661505Z","shell.execute_reply":"2026-03-01T23:17:49.986068Z"}},"outputs":[],"execution_count":null}]}