{"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":"none","dataSources":[{"sourceId":126777,"databundleVersionId":15314950,"sourceType":"competition"}],"dockerImageVersionId":31259,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport math\nimport random\nimport warnings\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom tqdm import tqdm\nfrom PIL import Image\nimport timm\n\n# Suppress Pydantic/System warnings\nwarnings.filterwarnings(\"ignore\")\n\n# ==========================================\n# 1. ROBUST CONFIGURATION & PATH FIX\n# ==========================================\nclass Config:\n    seed = 927\n    model_name = 'resnet200d' \n    img_size = 224\n    batch_size = 32\n    embedding_size = 512\n    num_classes = 31\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \n    # Placeholder Paths (Will be auto-corrected below)\n    train_dir = ''\n    test_dir = ''\n    train_csv = ''\n    test_csv = ''\n    submission_file = 'submission.csv'\n\ndef find_dataset_paths():\n    \"\"\"Auto-detects the correct paths for CSVs and Image directories.\"\"\"\n    print(\"Searching for dataset files...\")\n    base_search_paths = ['/kaggle/input', '.']\n    \n    # 1. Find CSVs\n    for root in base_search_paths:\n        for dirpath, _, files in os.walk(root):\n            if 'train.csv' in files:\n                Config.train_csv = os.path.join(dirpath, 'train.csv')\n            if 'test.csv' in files:\n                Config.test_csv = os.path.join(dirpath, 'test.csv')\n                \n    # 2. Find Image Directories (Look for specific sample images)\n    # We look for 'train_0001.png' and 'test_0001.png' to identify the folders\n    sample_train = 'train_0001.png'\n    sample_test = 'test_0001.png'\n    \n    for root in base_search_paths:\n        for dirpath, _, files in os.walk(root):\n            if sample_train in files:\n                Config.train_dir = dirpath\n            if sample_test in files:\n                Config.test_dir = dirpath\n                \n    # 3. Fallbacks / Verification\n    if not Config.train_dir or not Config.train_csv:\n        print(\"WARNING: Dataset not found. Using Dummy Paths (Code will generate random submission).\")\n        # Set dummy paths to prevent crashes before the dummy-data check\n        Config.train_csv = 'dummy_train.csv' \n    else:\n        print(f\"Dataset Found:\\n  Train CSV: {Config.train_csv}\\n  Train Dir: {Config.train_dir}\")\n\n# Set Seeds\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n\nseed_everything(Config.seed)\nfind_dataset_paths() # Run path fixer immediately\n\n# ==========================================\n# 2. DATASET (Robust)\n# ==========================================\nclass JaguarDataset(Dataset):\n    def __init__(self, df, img_dir, transform=None, is_test=False):\n        self.df = df\n        self.img_dir = img_dir\n        self.transform = transform\n        self.is_test = is_test\n        \n        if not self.is_test:\n            # Map labels\n            self.label_map = {name: idx for idx, name in enumerate(df['ground_truth'].unique())}\n            self.df['label'] = self.df['ground_truth'].map(self.label_map)\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_name = row['filename']\n        img_path = os.path.join(self.img_dir, img_name)\n        \n        try:\n            image = Image.open(img_path).convert('RGB')\n        except Exception as e:\n            # Robustness: If a file is missing/corrupt, return a black image instead of crashing\n            # logging only once to avoid spamming\n            if idx == 0: print(f\"Warning: Could not load {img_path}. Returning black image. ({e})\")\n            image = Image.new('RGB', (Config.img_size, Config.img_size), (0, 0, 0))\n        \n        if self.transform:\n            image = self.transform(image)\n            \n        if self.is_test:\n            return image, img_name\n        else:\n            label = torch.tensor(row['label'], dtype=torch.long)\n            return image, label\n\n# ==========================================\n# 3. ARCFACE & MODEL\n# ==========================================\nclass ArcFaceLayer(nn.Module):\n    def __init__(self, in_features, out_features, s=30.0, m=0.50):\n        super(ArcFaceLayer, self).__init__()\n        self.in_features = in_features\n        self.out_features = out_features\n        self.s = s\n        self.m = m\n        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))\n        nn.init.xavier_uniform_(self.weight)\n\n    def forward(self, input, label=None):\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n        if label is None: return cosine\n        phi = cosine - self.m\n        one_hot = torch.zeros(cosine.size(), device=input.device)\n        one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n        output *= self.s\n        return output\n\nclass JaguarReIDModel(nn.Module):\n    def __init__(self, model_name, embedding_size, num_classes, pretrained=True):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0)\n        in_features = self.backbone.num_features\n        self.neck = nn.Sequential(\n            nn.BatchNorm1d(in_features),\n            nn.Linear(in_features, embedding_size),\n            nn.BatchNorm1d(embedding_size),\n        )\n        self.head = ArcFaceLayer(embedding_size, num_classes)\n\n    def forward(self, x, label=None):\n        features = self.backbone(x)\n        embeddings = self.neck(features)\n        if label is not None: return self.head(embeddings, label)\n        return embeddings\n\n# ==========================================\n# 4. PIPELINE\n# ==========================================\ndef run_pipeline():\n    print(f\"Running on device: {Config.device}\")\n    \n    # Transforms\n    transforms_train = transforms.Compose([\n        transforms.Resize((Config.img_size, Config.img_size)),\n        transforms.RandomHorizontalFlip(),\n        transforms.ToTensor(),\n        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n    ])\n    transforms_test = transforms.Compose([\n        transforms.Resize((Config.img_size, Config.img_size)),\n        transforms.ToTensor(),\n        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n    ])\n\n    # Check for Data Presence\n    if not os.path.exists(Config.train_csv):\n        print(\"Data files missing. Generating DUMMY submission for verification.\")\n        # Generate dummy\n        if os.path.exists('sample_submission.csv'):\n            sample = pd.read_csv('sample_submission.csv')\n            pd.DataFrame({'row_id': sample['row_id'], 'similarity': np.random.rand(len(sample))}).to_csv(Config.submission_file, index=False)\n        else:\n            pd.DataFrame({'row_id': range(10), 'similarity': np.random.rand(10)}).to_csv(Config.submission_file, index=False)\n        return\n\n    # Load Data\n    train_df = pd.read_csv(Config.train_csv)\n    test_df_pairs = pd.read_csv(Config.test_csv)\n    \n    train_dataset = JaguarDataset(train_df, Config.train_dir, transform=transforms_train)\n    train_loader = DataLoader(train_dataset, batch_size=Config.batch_size, shuffle=True, num_workers=2)\n\n    # Initialize Model\n    model = JaguarReIDModel(Config.model_name, Config.embedding_size, Config.num_classes).to(Config.device)\n    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n    criterion = nn.CrossEntropyLoss()\n\n    # Training\n    print(\"Starting Training...\")\n    model.train()\n    for epoch in range(5): # Increase epochs for real training\n        loop = tqdm(train_loader, desc=f\"Epoch {epoch+1}\")\n        for images, labels in loop:\n            images, labels = images.to(Config.device), labels.to(Config.device)\n            optimizer.zero_grad()\n            loss = criterion(model(images, labels), labels)\n            loss.backward()\n            optimizer.step()\n            loop.set_postfix(loss=loss.item())\n\n    # Inference\n    print(\"Starting Inference...\")\n    model.eval()\n    \n    unique_images = sorted(list(set(test_df_pairs['query_image']) | set(test_df_pairs['gallery_image'])))\n    unique_df = pd.DataFrame({'filename': unique_images})\n    test_dataset = JaguarDataset(unique_df, Config.test_dir, transform=transforms_test, is_test=True)\n    test_loader = DataLoader(test_dataset, batch_size=Config.batch_size, shuffle=False, num_workers=2)\n    \n    embeddings_map = {}\n    with torch.no_grad():\n        for images, names in tqdm(test_loader, desc=\"Extracting\"):\n            images = images.to(Config.device)\n            emb = F.normalize(model(images), p=2, dim=1).cpu().numpy()\n            for name, e in zip(names, emb):\n                embeddings_map[name] = e\n                \n    # Scoring\n    similarities = []\n    # Optimization: Use list comprehension or vectorization if possible, loop is fine for 130k\n    for _, row in tqdm(test_df_pairs.iterrows(), total=len(test_df_pairs), desc=\"Pairing\"):\n        if row['query_image'] in embeddings_map and row['gallery_image'] in embeddings_map:\n            sim = np.dot(embeddings_map[row['query_image']], embeddings_map[row['gallery_image']])\n            similarities.append((sim + 1) / 2)\n        else:\n            similarities.append(0.5) # Fallback if image load failed\n\n    # Save\n    pd.DataFrame({'row_id': test_df_pairs['row_id'], 'similarity': similarities}).to_csv(Config.submission_file, index=False)\n    print(\"Success! Submission saved.\")\n\nif __name__ == '__main__':\n    run_pipeline()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-19T04:56:09.233415Z","iopub.execute_input":"2026-01-19T04:56:09.233793Z","iopub.status.idle":"2026-01-19T05:07:27.538506Z","shell.execute_reply.started":"2026-01-19T04:56:09.233754Z","shell.execute_reply":"2026-01-19T05:07:27.537163Z"}},"outputs":[],"execution_count":null}]}