{"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":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-19T04:50:34.114645Z","iopub.execute_input":"2026-01-19T04:50:34.115160Z","iopub.status.idle":"2026-01-19T04:50:39.570506Z","shell.execute_reply.started":"2026-01-19T04:50:34.115130Z","shell.execute_reply":"2026-01-19T04:50:39.569816Z"}},"outputs":[],"execution_count":null},{"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 = 42\n    model_name = 'resnet50' \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(1): # 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}]}