{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"from io import BytesIO\n\nimport lmdb\nfrom PIL import Image\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom torchvision import datasets, transforms, utils\n\n\nclass MultiResolutionDataset(Dataset):\n    def __init__(self, path, transform, resolution=8):\n        self.env = lmdb.open(\n            path,\n            max_readers=32,\n            readonly=True,\n            lock=False,\n            readahead=False,\n            meminit=False,\n        )\n\n        if not self.env:\n            raise IOError('Cannot open lmdb dataset', path)\n\n        with self.env.begin(write=False) as txn:\n            self.length = int(txn.get('length'.encode('utf-8')).decode('utf-8'))\n\n        self.resolution = resolution\n        self.transform = transform\n\n    def __len__(self):\n        return self.length\n\n    def __getitem__(self, index):\n        with self.env.begin(write=False) as txn:\n            key = str(self.resolution) +'-'+ str(index).zfill(5)\n            key = key.encode('utf-8') \n            img_bytes = txn.get(key)\n\n        buffer = BytesIO(img_bytes)\n        img = Image.open(buffer)\n        img = self.transform(img)\n\n        return img\n\n\ntransform = transforms.Compose(\n    [\n        transforms.ToTensor(),\n    ]\n)\n\n\ndataset = MultiResolutionDataset('/kaggle/input/cs470-mvtec/', transform)","metadata":{"execution":{"iopub.status.busy":"2021-10-11T14:01:52.472785Z","iopub.execute_input":"2021-10-11T14:01:52.473127Z","iopub.status.idle":"2021-10-11T14:01:52.492274Z","shell.execute_reply.started":"2021-10-11T14:01:52.473097Z","shell.execute_reply":"2021-10-11T14:01:52.491228Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def sample_data(dataset, batch_size, image_size=256):\n    dataset.resolution = image_size\n    loader = DataLoader(dataset, shuffle=True, batch_size=batch_size, num_workers=8, drop_last=True)\n\n    return loader\n\nloader = sample_data(\n    dataset, 4, 256\n)","metadata":{"execution":{"iopub.status.busy":"2021-10-11T14:01:52.494615Z","iopub.execute_input":"2021-10-11T14:01:52.494976Z","iopub.status.idle":"2021-10-11T14:01:52.504545Z","shell.execute_reply.started":"2021-10-11T14:01:52.49494Z","shell.execute_reply":"2021-10-11T14:01:52.503614Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, datasets, models\n\nimport torch\nimport torch.nn as nn\n\ndef convrelu(in_channels, out_channels, kernel, padding):\n    return nn.Sequential(\n        nn.Conv2d(in_channels, out_channels, kernel, padding=padding),\n        nn.LeakyReLU(0.1)#nn.ReLU(inplace=False),\n    )\n\nclass ResNetUNet(nn.Module):\n\n    def __init__(self, n_class):\n        super().__init__()\n        \n        self.base_model = models.resnet18(pretrained=False)            \n        self.base_layers = list(self.base_model.children())                \n        \n        self.layer0 = nn.Sequential(*self.base_layers[:3])\n        self.layer0_1x1 = convrelu(64, 64, 1, 0)\n        self.layer1 = nn.Sequential(*self.base_layers[3:5])     \n        self.layer1_1x1 = convrelu(64, 64, 1, 0)       \n        self.layer2 = self.base_layers[5]  \n        self.layer2_1x1 = convrelu(128, 128, 1, 0)  \n        self.layer3 = self.base_layers[6]       \n        self.layer3_1x1 = convrelu(256, 256, 1, 0)  \n        self.layer4 = self.base_layers[7] \n        self.layer4_1x1 = convrelu(512, 512, 1, 0) \n        \n        \n        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n        \n        self.conv_up3 = convrelu(256 + 512, 512, 3, 1)\n        self.conv_up2 = convrelu(128 + 512, 256, 3, 1)\n        self.conv_up1 = convrelu(64 + 256, 256, 3, 1)\n        self.conv_up0 = convrelu(64 + 256, 128, 3, 1)\n        \n        self.conv_original_size0 = convrelu(3, 64, 3, 1)\n        self.conv_original_size1 = convrelu(64, 64, 3, 1)\n        self.conv_original_size2 = convrelu(64 + 128, 64, 3, 1)\n        \n        self.conv_last = nn.Conv2d(64, n_class, 1)\n        \n    def forward(self, input):\n        x_original = self.conv_original_size0(input)\n        x_original = self.conv_original_size1(x_original)\n        \n        layer0 = self.layer0(input)            \n        layer1 = self.layer1(layer0)\n        layer2 = self.layer2(layer1)\n        layer3 = self.layer3(layer2)        \n        layer4 = self.layer4(layer3)\n        \n        layer4 = self.layer4_1x1(layer4)\n        x = self.upsample(layer4)\n        layer3 = self.layer3_1x1(layer3)\n        x = torch.cat([x, layer3], dim=1)\n        \n        x = self.conv_up3(x)\n        \n        x = self.upsample(x)\n        layer2 = self.layer2_1x1(layer2)\n        #print(x.shape, layer2.shape)\n        x = torch.cat([x, layer2], dim=1)\n        x = self.conv_up2(x)\n\n        x = self.upsample(x)\n        layer1 = self.layer1_1x1(layer1)\n        x = torch.cat([x, layer1], dim=1)\n        x = self.conv_up1(x)\n\n        x = self.upsample(x)\n        layer0 = self.layer0_1x1(layer0)\n        x = torch.cat([x, layer0], dim=1)\n        x = self.conv_up0(x)\n        \n        x = self.upsample(x)\n        x = torch.cat([x, x_original], dim=1)\n        x = self.conv_original_size2(x)        \n        \n        out = self.conv_last(x)\n        \n        return out","metadata":{"execution":{"iopub.status.busy":"2021-10-11T14:01:52.50631Z","iopub.execute_input":"2021-10-11T14:01:52.506635Z","iopub.status.idle":"2021-10-11T14:01:52.529813Z","shell.execute_reply.started":"2021-10-11T14:01:52.506598Z","shell.execute_reply":"2021-10-11T14:01:52.528659Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.optim as optim\nfrom torch.optim import lr_scheduler\nimport time\nimport copy\n\nimport math\nimport numbers\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\nnum_class = 1\n\nmodel = ResNetUNet(num_class).to(device)\noptimizer = optim.Adam(model.parameters(), lr=1e-4)\n","metadata":{"execution":{"iopub.status.busy":"2021-10-11T14:01:52.531692Z","iopub.execute_input":"2021-10-11T14:01:52.532002Z","iopub.status.idle":"2021-10-11T14:01:52.918812Z","shell.execute_reply.started":"2021-10-11T14:01:52.531963Z","shell.execute_reply":"2021-10-11T14:01:52.917942Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def dice_loss(pred, target, smooth = 1.):\n    pred = pred.contiguous()\n    target = target.contiguous()    \n\n    intersection = (pred * target).sum(dim=2).sum(dim=2)\n    \n    loss = (1 - ((2. * intersection + smooth) / (pred.sum(dim=2).sum(dim=2) + target.sum(dim=2).sum(dim=2) + smooth)))\n    \n    return loss.mean()\n\ndef calc_loss(pred, target, bce_weight=0.5):\n    bce = F.binary_cross_entropy_with_logits(pred, target)\n        \n    pred = torch.sigmoid(pred)\n    dice = dice_loss(pred, target)\n    \n    loss = bce * bce_weight + dice * (1 - bce_weight)    \n    return loss","metadata":{"execution":{"iopub.status.busy":"2021-10-11T14:01:52.92084Z","iopub.execute_input":"2021-10-11T14:01:52.921296Z","iopub.status.idle":"2021-10-11T14:01:52.932217Z","shell.execute_reply.started":"2021-10-11T14:01:52.921257Z","shell.execute_reply":"2021-10-11T14:01:52.931289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"num_epochs=4000\nbest_loss = 1e10\n\nfor epoch in range(num_epochs):\n    print('Epoch {}/{}'.format(epoch, num_epochs - 1))\n    print('-' * 10)\n\n    since = time.time()\n    model.train()\n    \n    for original_img in loader:\n        original_img = original_img.to(device)\n\n        optimizer.zero_grad()\n\n        output = model(original_img) \n        loss = calc_loss(output, torch.mean(original_img, axis=1)[:,None,:,:] )\n\n        loss.backward()#retain_graph=True)\n        optimizer.step()","metadata":{"execution":{"iopub.status.busy":"2021-10-11T14:01:52.933519Z","iopub.execute_input":"2021-10-11T14:01:52.93373Z","iopub.status.idle":"2021-10-11T14:02:08.907826Z","shell.execute_reply.started":"2021-10-11T14:01:52.933704Z","shell.execute_reply":"2021-10-11T14:02:08.906486Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def model_output(input_test_img):\n    #random model\n    #return numpy.random.rand(input_test_img.shape[0],input_test_img.shape[1],3)\n    #naive model\n    mask = model(input_test_img)\n    return mask","metadata":{"execution":{"iopub.status.busy":"2021-10-11T14:02:12.420174Z","iopub.execute_input":"2021-10-11T14:02:12.420526Z","iopub.status.idle":"2021-10-11T14:02:12.426885Z","shell.execute_reply.started":"2021-10-11T14:02:12.420493Z","shell.execute_reply":"2021-10-11T14:02:12.425648Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport numpy as np\n\nnum_file = 0\nroot_dir = \"/kaggle/input/cs470-mvtec/test/capsule/\"\nret_list = []\nfor root, dirs, files in os.walk(root_dir):\n    path = root.split(os.sep)\n    if path[-1] == 'good' or path[-1] == '.ipynb_checkpoints':\n        continue\n    for file in (sorted(files[:])):\n        if 'mask.png' != file[-8:] and '.png' == file[-4:] and path[-1] != '':\n            anomaly = Image.open(root_dir + str(path[-1]) + '/' + file).convert('RGB').resize((256, 256))\n            anomaly_image = np.array(anomaly) / 255.0\n            anomaly_image = torch.from_numpy(anomaly_image.astype(np.float32)).permute(2, 0, 1).unsqueeze(0)#.cuda()\n            \n            ret = model_output(anomaly_image)\n            ret_list.append(ret.data.cpu().numpy())","metadata":{"execution":{"iopub.status.busy":"2021-10-11T14:03:42.344646Z","iopub.execute_input":"2021-10-11T14:03:42.345481Z","iopub.status.idle":"2021-10-11T14:03:53.99004Z","shell.execute_reply.started":"2021-10-11T14:03:42.345442Z","shell.execute_reply":"2021-10-11T14:03:53.988738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ravel_list = np.array(ret_list).ravel()\n\nimport csv\nf = open('answer.csv','w')\nf.write('id,Predicted\\n')\nfor idx, line in enumerate(list(ravel_list)): \n    f.write(str(idx)+',')\n    f.write(str(line)+'\\n')\nf.close()\n\nprint(len(list(ravel_list)) == 7143424)","metadata":{"execution":{"iopub.status.busy":"2021-10-11T14:02:26.787231Z","iopub.execute_input":"2021-10-11T14:02:26.788122Z","iopub.status.idle":"2021-10-11T14:02:26.796579Z","shell.execute_reply.started":"2021-10-11T14:02:26.788085Z","shell.execute_reply":"2021-10-11T14:02:26.795513Z"},"trusted":true},"execution_count":null,"outputs":[]}]}