{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":88100,"databundleVersionId":10078502,"sourceType":"competition"}],"dockerImageVersionId":30805,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import io\nimport glob\nimport json\nimport math\nimport timm\nimport torch\nimport random\nimport torchvision\n\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nimport albumentations as A\n\nfrom PIL import Image\nfrom tqdm import tqdm\nfrom sklearn.neighbors import KNeighborsClassifier","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:45:51.454644Z","iopub.execute_input":"2025-01-08T00:45:51.455089Z","iopub.status.idle":"2025-01-08T00:46:30.389842Z","shell.execute_reply.started":"2025-01-08T00:45:51.455037Z","shell.execute_reply":"2025-01-08T00:46:30.389062Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"epochs = [50, 50, 50]\nlrs = [1e-4, 1e-4, 1e-4]\ntask = \"strict\"\npclean = 0.4\nbackbones = [\"resnet18\"]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:46:30.391897Z","iopub.execute_input":"2025-01-08T00:46:30.392433Z","iopub.status.idle":"2025-01-08T00:46:30.397546Z","shell.execute_reply.started":"2025-01-08T00:46:30.392402Z","shell.execute_reply":"2025-01-08T00:46:30.396529Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def extract_feature(model, data):\n    dataset = torch.utils.data.DataLoader(TrainDataset(data, \n                                                       [1]*len(data)), \n                                          batch_size=32, \n                                          shuffle=False, \n                                          num_workers=4)\n    features = []\n    pbar = tqdm(dataset)\n    for batch in pbar:\n        images = batch[\"image\"].to(device)\n        with torch.autocast(device_type=\"cuda\" if torch.cuda.is_available() else \"cpu\", dtype=torch.float16, enabled=True):\n            f = model.extract(images)\n        features.append(f.detach().cpu())\n    features = torch.cat(features, dim=0)\n    return features\n\ndef clean(task, pclean, data_positive_test, data_positive_train, data_negative_train, model):\n    if not pclean:\n        return data_positive_test\n\n    model.eval()\n    if task == \"disjoint\":\n        print(\"CLEAN DISJOINT\")\n        test_features = extract_feature(model, data_positive_test)\n        positive_features = extract_feature(model, data_positive_train)\n        negative_features = extract_feature(model, data_negative_train)\n        \n        cos = torch.nn.functional.linear(torch.nn.functional.normalize(test_features.double()), \n                                         torch.nn.functional.normalize(model.fc.weight[1:].detach().cpu().double())).flatten().detach().cpu().numpy()\n        \n        clf = KNeighborsClassifier(n_neighbors=100, metric=\"cosine\")\n        clf.fit(torch.cat([negative_features, positive_features], dim=0),\n                [0]*len(negative_features) + [1]*len(positive_features))\n        probs = clf.predict_proba(test_features)\n\n        score = probs[:, 1] + cos\n        indexes = np.argsort(1-score)[:int((1-pclean)*len(score))]\n        return [data_positive_test[i] for i in indexes]\n\n    elif task == \"strict\":\n        print(\"CLEAN STRICT\")\n        test_features = extract_feature(model, data_positive_test)\n        positive_features = extract_feature(model, data_positive_train)\n        negative_features = extract_feature(model, data_negative_train)\n        \n        cos = torch.nn.functional.linear(torch.nn.functional.normalize(test_features.double()), \n                                         torch.nn.functional.normalize(model.fc.weight[:1].detach().cpu().double())).flatten().detach().cpu().numpy()\n\n        clf = KNeighborsClassifier(n_neighbors=100, metric=\"cosine\")\n        clf.fit(torch.cat([negative_features, positive_features], dim=0),\n                [0]*len(negative_features) + [1]*len(positive_features))\n        probs = clf.predict_proba(test_features)\n        \n        score = probs[:, 0] + cos\n        indexes = np.argsort(score.flatten())[:int((1-pclean)*len(score))]\n        return [data_positive_test[i] for i in indexes]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T01:01:14.949201Z","iopub.execute_input":"2025-01-08T01:01:14.949635Z","iopub.status.idle":"2025-01-08T01:01:14.964863Z","shell.execute_reply.started":"2025-01-08T01:01:14.949598Z","shell.execute_reply":"2025-01-08T01:01:14.963795Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TrainDataset(torch.utils.data.Dataset):\n    def __init__(self, files, labels, img_size=224):\n        self.files = files\n        self.labels = labels\n        self.img_size = img_size\n        self.augmentation = A.Compose([A.LongestMaxSize(max_size=self.img_size, interpolation=1),\n                                       A.PadIfNeeded(min_height=self.img_size, min_width=self.img_size, border_mode=0, value=(0,0,0))])\n        self.transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),\n                                                         torchvision.transforms.Normalize(mean=[0.48141356, 0.47885622, 0.47124268], std=[0.08375391, 0.07815655, 0.0793708])])\n\n\n    def __len__(self):\n        return len(self.files)\n\n    def __getitem__(self, idx):\n        image = np.asarray(Image.open(self.files[idx]))[:, :, :3]\n        image = self.transform(self.augmentation(image=image)['image'])\n        \n        labels = torch.zeros(2)\n        labels[self.labels[idx]] = 1.0\n        return {\"image\": image, \"label\": labels}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:56:44.710467Z","iopub.execute_input":"2025-01-08T00:56:44.710903Z","iopub.status.idle":"2025-01-08T00:56:44.719053Z","shell.execute_reply.started":"2025-01-08T00:56:44.710870Z","shell.execute_reply":"2025-01-08T00:56:44.717974Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TestDataset(torch.utils.data.Dataset):\n    def __init__(self, img_size=224):\n        self.img_size = img_size\n        self.augmentation = A.Compose([A.LongestMaxSize(max_size=self.img_size, interpolation=1),\n                                       A.PadIfNeeded(min_height=self.img_size, min_width=self.img_size, border_mode=0, value=(0,0,0))])\n        self.transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),\n                                                         torchvision.transforms.Normalize(mean=[0.48141356, 0.47885622, 0.47124268], std=[0.08375391, 0.07815655, 0.0793708])])\n\n\n    def __len__(self):\n        return 1500\n\n    def __getitem__(self, idx):\n        image = np.asarray(Image.open(f\"/kaggle/input/concon/unconf/test/images/t0/image_{idx+1}.png\"))[:, :, :3]\n        image = self.transform(self.augmentation(image=image)['image'])\n        return {\"id\": idx+1 ,\"image\": image}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:54:07.711613Z","iopub.execute_input":"2025-01-08T00:54:07.712307Z","iopub.status.idle":"2025-01-08T00:54:07.719640Z","shell.execute_reply.started":"2025-01-08T00:54:07.712271Z","shell.execute_reply":"2025-01-08T00:54:07.718645Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ArcMarginProduct(torch.nn.Module):\n    def __init__(self, in_features, out_features, s=1.0, m=0.1):\n        super(ArcMarginProduct, 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 = torch.nn.Parameter(torch.FloatTensor(out_features, in_features))\n        torch.nn.init.xavier_uniform_(self.weight)\n\n        self.cos_m = math.cos(m)\n        self.sin_m = math.sin(m)\n        self.th = math.cos(math.pi - m)\n        self.mm = math.sin(math.pi - m) * m\n\n    def forward(self, input, label):\n        cos = torch.nn.functional.linear(torch.nn.functional.normalize(input), torch.nn.functional.normalize(self.weight))\n        if label is None:\n            return cos*self.s\n        cos = torch.nn.functional.linear(torch.nn.functional.normalize(input), torch.nn.functional.normalize(self.weight))\n        sin = torch.sqrt((1.0 - torch.pow(cos, 2)).clamp(0, 1))\n        cos_add = cos * self.cos_m - sin * self.sin_m\n        output = (label * cos_add) + ((1.0 - label) * cos)\n        output *= self.s\n        return output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:54:07.864213Z","iopub.execute_input":"2025-01-08T00:54:07.864940Z","iopub.status.idle":"2025-01-08T00:54:07.873906Z","shell.execute_reply.started":"2025-01-08T00:54:07.864899Z","shell.execute_reply":"2025-01-08T00:54:07.872695Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ModelFactory(torch.nn.Module):\n    def __init__(self, image_backbones, n_classes, hashLength=256):\n        super(ModelFactory, self).__init__()\n        self.image_models = [timm.create_model(image_backbone, pretrained=False, num_classes=0).to(device) for image_backbone in image_backbones]\n        if torch.cuda.device_count() > 1:\n            for i in range(len(image_backbones)):\n                self.image_models[i] = torch.nn.DataParallel(self.image_models[i])\n\n        self.hash = torch.nn.Linear(512, hashLength)\n        self.fc = ArcMarginProduct(hashLength, n_classes) # torch.nn.Linear(hashLength, n_classes)\n        self.output_layer = torch.nn.Softmax(dim=1)\n\n    def extract(self, images):\n        image_features = [model(images) for model in self.image_models]\n        features = torch.cat(image_features, dim=1)\n        fhash = self.hash(features)\n        fhash = torch.nn.ReLU()(fhash)\n        return fhash\n\n    def forward(self, images, labels=None):\n        fhash = self.extract(images)\n        logits = self.fc(fhash, labels)\n        return self.output_layer(logits)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:54:08.040812Z","iopub.execute_input":"2025-01-08T00:54:08.041595Z","iopub.status.idle":"2025-01-08T00:54:08.050251Z","shell.execute_reply.started":"2025-01-08T00:54:08.041559Z","shell.execute_reply":"2025-01-08T00:54:08.049216Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Accuracy:\n    def __init__(self):\n        self.correct_count = 0\n        self.total_count = 0\n\n    def update(self, predictions, targets):\n        predicted_labels = torch.argmax(predictions, dim=1)\n        target_labels = torch.argmax(targets, dim=1)\n        self.correct_count += (predicted_labels == target_labels).sum().item()\n        self.total_count += targets.size(0)\n\n    def reset(self):\n        self.correct_count = 0\n        self.total_count = 0\n\n    def compute(self):\n        if self.total_count == 0:\n            return 0.0\n        return self.correct_count / self.total_count","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:47:07.414407Z","iopub.execute_input":"2025-01-08T00:47:07.415305Z","iopub.status.idle":"2025-01-08T00:47:07.421341Z","shell.execute_reply.started":"2025-01-08T00:47:07.415271Z","shell.execute_reply":"2025-01-08T00:47:07.420121Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmetrics = [Accuracy()]\n\nscaler = torch.amp.GradScaler(enabled=True)\nif torch.cuda.device_count() > 1:\n    model = torch.nn.DataParallel(model)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:47:07.575883Z","iopub.execute_input":"2025-01-08T00:47:07.576247Z","iopub.status.idle":"2025-01-08T00:47:07.581695Z","shell.execute_reply.started":"2025-01-08T00:47:07.576216Z","shell.execute_reply":"2025-01-08T00:47:07.580640Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_negative = glob.glob(f\"/kaggle/input/concon/{task}/train/images/t0/0/*\")\ndata_positive = glob.glob(f\"/kaggle/input/concon/{task}/train/images/t0/1/*\")\nfiles = data_negative + data_positive\nlabels = [0]*len(data_negative) + [1]*len(data_positive)\ndata = list(zip(files, labels))\nrandom.shuffle(data)\nfiles, labels = zip(*data)\n\ntrain_dataset = torch.utils.data.DataLoader(TrainDataset(files, labels), batch_size=32, shuffle=True, num_workers=4)\n\ndata_negative = glob.glob(f\"/kaggle/input/concon/{task}/val/images/t0/0/*\")\ndata_positive = glob.glob(f\"/kaggle/input/concon/{task}/val/images/t0/1/*\")\nfiles = data_negative + data_positive\nlabels = [0]*len(data_negative) + [1]*len(data_positive)\ndata = list(zip(files, labels))\nrandom.shuffle(data)\nfiles, labels = zip(*data)\n\nvalid_dataset = torch.utils.data.DataLoader(TrainDataset(files, labels), batch_size=32, shuffle=False, num_workers=4)\n\nmodel = ModelFactory(backbones, 2).to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=lrs[0])\nloss_function = torch.nn.CrossEntropyLoss()\n\nfor epoch in range(epochs[0]):\n    model.train()\n    optimizer.zero_grad()\n    pbar = tqdm(train_dataset)\n    for i, batch in enumerate(pbar):\n        images = batch[\"image\"].to(device)\n        labels = batch[\"label\"].to(device)\n        with torch.autocast(device_type=\"cuda\" if torch.cuda.is_available() else \"cpu\", dtype=torch.float16, enabled=True):\n            output = model(images, labels)\n            loss = loss_function(output, labels)\n\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        optimizer.zero_grad()\n\n        for metric in metrics:\n            metric.update(output, labels)\n        pbar.set_description(f\"Epoch: {epoch}, Loss: {loss.item():.5f}, Accuracy: {metric.compute():.5f}\")\n\n    metric.reset()\n    torch.save(model, \"model.pth\")\n    model.eval()\n    \n    pbar = tqdm(valid_dataset)\n    for batch in pbar:\n        images = batch[\"image\"].to(device)\n        labels = batch[\"label\"].to(device)\n        with torch.autocast(device_type=\"cuda\" if torch.cuda.is_available() else \"cpu\", dtype=torch.float16, enabled=True):\n            output = model(images)\n        for metric in metrics:\n            metric.update(output, labels)\n        pbar.set_description(f\"Valid: {epoch}, Accuracy: {metric.compute():.5f}\")\n    metric.reset()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:47:07.751935Z","iopub.execute_input":"2025-01-08T00:47:07.752285Z","iopub.status.idle":"2025-01-08T00:48:53.601714Z","shell.execute_reply.started":"2025-01-08T00:47:07.752254Z","shell.execute_reply":"2025-01-08T00:48:53.600566Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_negative = glob.glob(f\"/kaggle/input/concon/{task}/train/images/t1/0/*\")\ndata_positive = clean(task, pclean, glob.glob(f\"/kaggle/input/concon/{task}/train/images/t1/1/*\"), glob.glob(f\"/kaggle/input/concon/{task}/train/images/t0/1/*\"), glob.glob(f\"/kaggle/input/concon/{task}/train/images/t0/0/*\"), model)\n\nfiles = data_negative + data_positive\nlabels = [0]*len(data_negative) + [1]*len(data_positive)\ndata = list(zip(files, labels))\nrandom.shuffle(data)\nfiles, labels = zip(*data)\n\ntrain_dataset = torch.utils.data.DataLoader(TrainDataset(files, labels), batch_size=32, shuffle=True, num_workers=4)\n\ndata_negative = glob.glob(f\"/kaggle/input/concon/{task}/val/images/t1/0/*\")\ndata_positive = glob.glob(f\"/kaggle/input/concon/{task}/val/images/t1/1/*\")\nfiles = data_negative + data_positive\nlabels = [0]*len(data_negative) + [1]*len(data_positive)\ndata = list(zip(files, labels))\nrandom.shuffle(data)\nfiles, labels = zip(*data)\n\nvalid_dataset = torch.utils.data.DataLoader(TrainDataset(files, labels), batch_size=32, shuffle=False, num_workers=4)\n\nmodel = ModelFactory(backbones, 2).to(device)\nmodel.load_state_dict(torch.load(\"model.pth\").state_dict())\noptimizer = torch.optim.Adam(model.parameters(), lr=lrs[1])\nloss_function = torch.nn.CrossEntropyLoss()\n\nfor epoch in range(epochs[1]):\n    model.train()\n    optimizer.zero_grad()\n    pbar = tqdm(train_dataset)\n    for i, batch in enumerate(pbar):\n        images = batch[\"image\"].to(device)\n        labels = batch[\"label\"].to(device)\n        with torch.autocast(device_type=\"cuda\" if torch.cuda.is_available() else \"cpu\", dtype=torch.float16, enabled=True):\n            output = model(images, labels)\n            loss = loss_function(output, labels)\n\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        optimizer.zero_grad()\n\n        for metric in metrics:\n            metric.update(output, labels)\n        pbar.set_description(f\"Epoch: {epoch}, Loss: {loss.item():.5f}, Accuracy: {metric.compute():.5f}\")\n\n    metric.reset()\n    torch.save(model, \"model.pth\")\n    model.eval()\n    \n    pbar = tqdm(valid_dataset)\n    for batch in pbar:\n        images = batch[\"image\"].to(device)\n        labels = batch[\"label\"].to(device)\n        with torch.autocast(device_type=\"cuda\" if torch.cuda.is_available() else \"cpu\", dtype=torch.float16, enabled=True):\n            output = model(images)\n        for metric in metrics:\n            metric.update(output, labels)\n        pbar.set_description(f\"Valid: {epoch}, Accuracy: {metric.compute():.5f}\")\n    metric.reset()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:56:51.306921Z","iopub.execute_input":"2025-01-08T00:56:51.307315Z","iopub.status.idle":"2025-01-08T00:57:16.069923Z","shell.execute_reply.started":"2025-01-08T00:56:51.307278Z","shell.execute_reply":"2025-01-08T00:57:16.068240Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_negative = glob.glob(f\"/kaggle/input/concon/{task}/train/images/t2/0/*\")\ndata_positive = clean(task, pclean, glob.glob(f\"/kaggle/input/concon/{task}/train/images/t2/1/*\"), glob.glob(f\"/kaggle/input/concon/{task}/train/images/t1/1/*\"), glob.glob(f\"/kaggle/input/concon/{task}/train/images/t1/0/*\"), model)\n\nfiles = data_negative + data_positive\nlabels = [0]*len(data_negative) + [1]*len(data_positive)\ndata = list(zip(files, labels))\nrandom.shuffle(data)\nfiles, labels = zip(*data)\n\ntrain_dataset = torch.utils.data.DataLoader(TrainDataset(files, labels), batch_size=32, shuffle=True, num_workers=4)\n\ndata_negative = glob.glob(f\"/kaggle/input/concon/{task}/val/images/t2/0/*\")\ndata_positive = glob.glob(f\"/kaggle/input/concon/{task}/val/images/t2/1/*\")\nfiles = data_negative + data_positive\nlabels = [0]*len(data_negative) + [1]*len(data_positive)\ndata = list(zip(files, labels))\nrandom.shuffle(data)\nfiles, labels = zip(*data)\n\nvalid_dataset = torch.utils.data.DataLoader(TrainDataset(files, labels), batch_size=32, shuffle=False, num_workers=4)\n\nmodel = ModelFactory(backbones, 2).to(device)\nmodel.load_state_dict(torch.load(\"model.pth\").state_dict())\noptimizer = torch.optim.Adam(model.parameters(), lr=lrs[2])\nloss_function = torch.nn.CrossEntropyLoss()\n\nfor epoch in range(epochs[2]):\n    model.train()\n    optimizer.zero_grad()\n    pbar = tqdm(train_dataset)\n    for i, batch in enumerate(pbar):\n        images = batch[\"image\"].to(device)\n        labels = batch[\"label\"].to(device)\n        with torch.autocast(device_type=\"cuda\" if torch.cuda.is_available() else \"cpu\", dtype=torch.float16, enabled=True):\n            output = model(images, labels)\n            loss = loss_function(output, labels)\n\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        optimizer.zero_grad()\n\n        for metric in metrics:\n            metric.update(output, labels)\n        pbar.set_description(f\"Epoch: {epoch}, Loss: {loss.item():.5f}, Accuracy: {metric.compute():.5f}\")\n\n    metric.reset()\n    torch.save(model, \"model.pth\")\n    model.eval()\n    \n    pbar = tqdm(valid_dataset)\n    for batch in pbar:\n        images = batch[\"image\"].to(device)\n        labels = batch[\"label\"].to(device)\n        with torch.autocast(device_type=\"cuda\" if torch.cuda.is_available() else \"cpu\", dtype=torch.float16, enabled=True):\n            output = model(images)\n        for metric in metrics:\n            metric.update(output, labels)\n        pbar.set_description(f\"Valid: {epoch}, Accuracy: {metric.compute():.5f}\")\n    metric.reset()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:46:33.230436Z","iopub.status.idle":"2025-01-08T00:46:33.230807Z","shell.execute_reply.started":"2025-01-08T00:46:33.230647Z","shell.execute_reply":"2025-01-08T00:46:33.230666Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_dataset = torch.utils.data.DataLoader(TestDataset(), batch_size=32, shuffle=False, num_workers=0)\npbar = tqdm(test_dataset)\npredictions = []\nfor batch in pbar:\n    images = batch[\"image\"].to(device)\n    with torch.autocast(device_type=\"cuda\" if torch.cuda.is_available() else \"cpu\", dtype=torch.float16, enabled=True):\n        output = model(images).detach().cpu().numpy()\n    for i, prob in enumerate(output):\n        argmax_pred = int(prob.argmax())\n        predictions.append(argmax_pred)\n        \ndf = pd.DataFrame()\ndf[\"ID\"] = [f\"image_{i+1}\" for i in range(1500)]\ndf[\"TARGET\"] = predictions\ndf.to_csv(f\"answer.csv\", index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-08T00:46:33.232613Z","iopub.status.idle":"2025-01-08T00:46:33.232950Z","shell.execute_reply.started":"2025-01-08T00:46:33.232796Z","shell.execute_reply":"2025-01-08T00:46:33.232813Z"}},"outputs":[],"execution_count":null}]}