{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":90316,"databundleVersionId":10477888,"sourceType":"competition"}],"dockerImageVersionId":30823,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install evaluate","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:25:24.680770Z","iopub.execute_input":"2025-01-04T08:25:24.681072Z","iopub.status.idle":"2025-01-04T08:25:27.921051Z","shell.execute_reply.started":"2025-01-04T08:25:24.681048Z","shell.execute_reply":"2025-01-04T08:25:27.920171Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport gc\nimport time\nimport copy\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport matplotlib.pyplot as plt\nfrom torch.optim import lr_scheduler\nfrom torchmetrics import AUROC\nfrom PIL import Image\nfrom tqdm import tqdm\nimport evaluate\nfrom torchvision.transforms.functional import to_pil_image\nimport random\n# Albumentations\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom albumentations import Compose, Normalize, RandomCrop, HorizontalFlip, ShiftScaleRotate, HueSaturationValue\n\n# Transformers / HuggingFace\nfrom transformers import (\n    SegformerConfig,\n    SegformerForSemanticSegmentation,\n    SegformerPreTrainedModel,\n)\n\n# Metrics\nfrom sklearn.metrics import f1_score, accuracy_score, recall_score, roc_auc_score, roc_curve\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:25:27.922401Z","iopub.execute_input":"2025-01-04T08:25:27.922638Z","iopub.status.idle":"2025-01-04T08:25:27.930354Z","shell.execute_reply.started":"2025-01-04T08:25:27.922619Z","shell.execute_reply":"2025-01-04T08:25:27.929523Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def crop(imgfile, maskfile, size=224):\n    mask = maskfile.astype(np.float32)\n    if np.sum(mask) / mask.size > 0.9:\n        # if the object mask has more than 50% of the whole mask, return the full image and mask\n        mask = cv2.resize(mask, (size, size), interpolation=cv2.INTER_NEAREST)\n        cropped_img = cv2.resize(imgfile, (size, size), interpolation=cv2.INTER_NEAREST)\n        return cropped_img, mask\n    else:\n        # otherwise, crop the image and mask to the bounding box of the object\n        rows = np.any(mask, axis=1)\n        cols = np.any(mask, axis=0)\n        y1, y2 = np.where(rows)[0][[0, -1]]\n        x1, x2 = np.where(cols)[0][[0, -1]]\n        mask = mask[y1:y2+1, x1:x2+1]\n        mask = cv2.resize(mask, (size, size), interpolation=cv2.INTER_NEAREST)\n        cropped_img = cv2.resize(imgfile[y1:y2+1, x1:x2+1], (size, size), interpolation=cv2.INTER_NEAREST)\n        return cropped_img, mask\n\ndef decode_rle_to_mask(rle, height, width, viz=False):\n    '''\n    rle : run-length as string formated (start value, count)\n    height : height of the mask \n    width : width of the mask\n    returns binary mask\n    '''\n    rle = np.array(rle.split(' ')).reshape(-1, 2)\n    mask = np.zeros((height*width, 1, 3))\n    if viz:\n        color = np.random.rand(3)\n    else:\n        color = [1,1,1]\n    for i in rle:\n        mask[int(i[0]):int(i[0])+int(i[1]), :, :] = color\n\n    return mask.reshape(height, width, 3)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:25:27.932184Z","iopub.execute_input":"2025-01-04T08:25:27.932441Z","iopub.status.idle":"2025-01-04T08:25:27.948661Z","shell.execute_reply.started":"2025-01-04T08:25:27.932419Z","shell.execute_reply":"2025-01-04T08:25:27.947947Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train = pd.read_csv('/kaggle/input/ai-dl-multiclass-segmentation/train.csv')\ndf_test = pd.read_csv('/kaggle/input/ai-dl-multiclass-segmentation/test_class.csv')\ntrainpath = '/kaggle/input/ai-dl-multiclass-segmentation/TrainImages'\ntestpath = '/kaggle/input/ai-dl-multiclass-segmentation/TestImages'\n\ndfcp=df_train.copy()\ndfcp = dfcp.drop_duplicates(subset=['ImageName'], keep='first')\nprint('length of the dataset:', len(dfcp))\n\ntraindata = {\"imgs\": [], \"masks\": []}\ncounter = 0\n\nfor i in dfcp['ImageName']:\n    imgpath=os.path.join(trainpath,i+'.jpg')\n    masks=[]\n    num=0\n    for j in df_train[df_train['ImageName'] == i]['Encoding']:#ImageHeight\n        wi = df_train[df_train['ImageName'] == i]['ImageWidth'].values[0]\n        hi = df_train[df_train['ImageName'] == i]['ImageHeight'].values[0]\n        classes = int(df_train[df_train['ImageName'] == i]['ClassNumber'].values[num])\n        masks.append(decode_rle_to_mask(j, int(hi), int(wi))*classes)\n        num+=1\n    mask = np.sum(masks, axis=0)[:,:,0]\n    try:\n        i,m=crop(cv2.imread(imgpath),mask,size=224)\n        traindata['imgs'].append(np.array(i))\n        traindata['masks'].append(np.array(m))\n        counter += 1\n    except:\n        print(\"Failed to add the image and mask of index\",counter)\n        continue\n\nprint(\"Images added: \",len(traindata['imgs']),\" ,Masks added: \",len(traindata['masks']))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:25:27.949840Z","iopub.execute_input":"2025-01-04T08:25:27.950158Z","iopub.status.idle":"2025-01-04T08:26:12.824448Z","shell.execute_reply.started":"2025-01-04T08:25:27.950125Z","shell.execute_reply":"2025-01-04T08:26:12.823534Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_transforms = Compose([\n                        ToTensorV2()\n                    ])\ntrain_transform = []\nfor i in range(len(traindata['imgs'])):\n    train_transform.append((data_transforms(image=traindata['imgs'][i]),traindata['masks'][i]))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:26:12.825396Z","iopub.execute_input":"2025-01-04T08:26:12.825729Z","iopub.status.idle":"2025-01-04T08:26:12.859626Z","shell.execute_reply.started":"2025-01-04T08:26:12.825695Z","shell.execute_reply":"2025-01-04T08:26:12.858594Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#PyTorch\nALPHA = 0.5\nBETA = 0.5\nGAMMA = 1\nclass CustomSegFormer2(SegformerPreTrainedModel):\n    def __init__(self, config, in_channel, n_classes=5):\n        super().__init__(config)\n        self.segformer = SegformerForSemanticSegmentation.from_pretrained(\"nvidia/segformer-b5-finetuned-ade-640-640\", ignore_mismatched_sizes=True)\n        self.logits_outputs = nn.Conv2d(in_channel, out_channels=n_classes, kernel_size=1, stride=1, padding=0)\n        self.distance_outputs = nn.Conv2d(in_channel, out_channels=1, kernel_size=1, stride=1, padding=0)\n        \n\n    def forward(self, pixel_values):\n        # forward to get raw logits\n        outputs = self.segformer(pixel_values) # what's the number of output channels?\n        logits_outputs = self.logits_outputs(outputs[0])\n        distancemap_outputs = self.distance_outputs(outputs[0])        \n\n        upsampled_logits = nn.functional.interpolate(logits_outputs, size=pixel_values.shape[-2:], mode=\"bilinear\", align_corners=False)\n        upsampled_distancemaps = nn.functional.interpolate(distancemap_outputs , size=pixel_values.shape[-2:], mode=\"bilinear\", align_corners=False)\n\n        \n        return  upsampled_logits, upsampled_distancemaps\nclass AverageMeter:\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, count=1):\n        self.val = val\n        self.sum += val * count\n        self.count += count\n        self.avg = self.sum / self.count\n\nclass FocalTverskyLoss(nn.Module):\n    def __init__(self, weight=None, size_average=True):\n        super(FocalTverskyLoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=1, alpha=ALPHA, beta=BETA, gamma=GAMMA):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        inputs = F.sigmoid(inputs)       \n        print(inputs.shape)\n        print(targets.shape)\n        #flatten label and prediction tensors\n        inputs = inputs.reshape(-1)\n        targets = targets.reshape(-1)\n        print(inputs.shape)\n        print(targets.shape)\n        #True Positives, False Positives & False Negatives\n        TP = (inputs * targets).sum()    \n        FP = ((1-targets) * inputs).sum()\n        FN = (targets * (1-inputs)).sum()\n        \n        Tversky = (TP + smooth) / (TP + alpha*FP + beta*FN + smooth)  \n        FocalTversky = (1 - Tversky)**gamma\n                       \n        return FocalTversky\n\ndef dice_loss(output, target, num_classes = 5):\n    # Compute softmax over the channel dimension\n    output = F.softmax(output, dim=1)\n    \n    # Create a one-hot encoding of the target\n    target_onehot = F.one_hot(target, num_classes=num_classes).permute(0, 3, 1, 2).float()\n\n    # Compute intersection and union between output and target\n    intersection = torch.sum(output * target_onehot, dim=[2, 3])\n    union = torch.sum(output, dim=[2, 3]) + torch.sum(target_onehot, dim=[2, 3])\n    \n    # Compute the Dice coefficient for each class\n    dice = 2 * intersection / union.clamp(min=1e-7)\n    \n    # Compute the average Dice loss across all classes\n    dice_loss = 1 - torch.mean(dice)\n    \n    return dice_loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:26:12.860641Z","iopub.execute_input":"2025-01-04T08:26:12.860919Z","iopub.status.idle":"2025-01-04T08:26:12.881280Z","shell.execute_reply.started":"2025-01-04T08:26:12.860893Z","shell.execute_reply":"2025-01-04T08:26:12.880587Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_model(model,dataloaders, criterion, optimizer, scheduler, name, num_epochs=25, feature_extractor = None, metric = None):\n\n    #Creating a folder to save the model performance.\n    try: os.mkdir('/kaggle/working/modelPerformance')\n    except: print('directory exists')\n    try: os.mkdir(f'/kaggle/working/modelPerformance/{name}')\n    except: print(f\"{name} alreadt exists\")\n    \n    since = time.time()\n    best_model_wts = copy.deepcopy(model.state_dict())\n    best_acc = 0.0  \n    acc = 0.0\n    loss = 0.0\n    loss_avg = AverageMeter()\n    for epoch in range(num_epochs):\n        print('Epoch {}/{}'.format(epoch+1, num_epochs))\n        print('-' * 10)\n        # Each epoch has a training and validation phase\n        for phase in ['train', 'test']:\n            print(phase)\n            if phase == 'train':\n                loss_avg = AverageMeter()\n                model = model.train()  # Set model to training mode\n            else:\n                loss_avg = AverageMeter()\n                model = model.eval()   # Set model to evaluate mode\n            running_loss = 0.0\n            running_corrects = 0\n\n            #epochs\n            epoch=int(len(dataloaders[phase]))\n            for _ in tqdm(range(epoch)):\n                #Loading Data\n                optimizer.zero_grad()\n                inputs, labels = next(iter(dataloaders[phase]))\n                inputs = inputs['image']\n                labels = labels.type(torch.LongTensor)\n                inputs = inputs.to(device)\n                print()\n                labels = labels.to(device)\n                # forward\n                # track history if only in train\n                if phase == 'train':\n                    with torch.set_grad_enabled(phase == 'train'):\n                        outputs, ssn = model(inputs.float())\n                        pred = outputs.softmax(dim=1)\n                        preds=pred.argmax(dim=1)\n                        loss = criterion(outputs, labels)\n                        print(loss.item())\n\n                        loss.backward()\n                        optimizer.step()\n                    loss_avg.update(loss.item())\n\n                if phase == 'test':\n                    outputs, ssn = model(inputs.float())\n                    pred = outputs.softmax(dim=1)\n                    preds=pred.argmax(dim=1)\n\n                    loss = criterion(outputs, labels)\n                    loss_avg.update(loss.item())           \n\n   \n                acc = metric.compute(predictions = preds, references = labels, num_labels = 5, ignore_index = 5, reduce_labels=False)\n                print(\"Acc:\",acc)\n                print(\"Loss:\",loss.item())\n                \n            if phase == 'train':\n                scheduler.step()\n                \n            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, loss_avg.val, acc['overall_accuracy']))\n    \n            \n            # deep copy the model\n            if phase == 'test' and acc['overall_accuracy'] > best_acc:\n                best_acc = acc['overall_accuracy']\n                torch.save(model.state_dict(),'/kaggle/working/modelPerformance/{}/best_model_{:.4f}acc_{}epochs.pth'.format(name,acc['overall_accuracy'],num_epochs))\n\n                train_losses = []\n                valid_losses = []\n\n        print()\n\n\n    time_elapsed = time.time() - since\n    print('Training complete in {:.0f}m {:.0f}s'.format(\n        time_elapsed // 60, time_elapsed % 60))\n    print('Best val Acc: {:4f}'.format(best_acc))\n\n    return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:26:12.882889Z","iopub.execute_input":"2025-01-04T08:26:12.883179Z","iopub.status.idle":"2025-01-04T08:26:12.903410Z","shell.execute_reply.started":"2025-01-04T08:26:12.883157Z","shell.execute_reply":"2025-01-04T08:26:12.902674Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\ndataloaders = {\n    'train': torch.utils.data.DataLoader(train_transform[:1600], batch_size=8, shuffle=True, num_workers=2),\n    'test': torch.utils.data.DataLoader(train_transform[1600:], batch_size=8, shuffle=True, num_workers=2)\n                                            }\nconfiguration = SegformerConfig(hidden_dropout_prob = 0.3, attention_probs_dropout_prob = 0.3, classifier_dropout_prob = 0.5)\nmodel=CustomSegFormer2(configuration,in_channel=150)\nlogits,DM = model(torch.rand(8,3,224,224))\nlogits.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:26:12.904509Z","iopub.execute_input":"2025-01-04T08:26:12.904724Z","iopub.status.idle":"2025-01-04T08:26:17.824200Z","shell.execute_reply.started":"2025-01-04T08:26:12.904704Z","shell.execute_reply":"2025-01-04T08:26:17.823348Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"criterion = nn.CrossEntropyLoss(ignore_index = 5, weight = torch.tensor([1.0, 5.0,5.0,8.0,10.0]).to(device) )\n#criterion = dice_loss\n#criterion = FocalTverskyLoss()\nmean_iou = evaluate.load(\"abidlabs/mean_iou\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:26:17.824995Z","iopub.execute_input":"2025-01-04T08:26:17.825277Z","iopub.status.idle":"2025-01-04T08:26:18.302019Z","shell.execute_reply.started":"2025-01-04T08:26:17.825245Z","shell.execute_reply":"2025-01-04T08:26:18.301393Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#model.load_state_dict(torch.load(\"/kaggle/working/segformers.pth\"))\nmodel_ft = model.to(device)\noptimizer_ft = optim.AdamW(model_ft.parameters(), lr=0.0001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01)\nexp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)\nmodel_ft = train_model(model_ft,dataloaders, criterion, optimizer_ft, exp_lr_scheduler, name=\"Segformer\", num_epochs=2, metric = mean_iou)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:26:18.302764Z","iopub.execute_input":"2025-01-04T08:26:18.303049Z","iopub.status.idle":"2025-01-04T08:47:24.608028Z","shell.execute_reply.started":"2025-01-04T08:26:18.303004Z","shell.execute_reply":"2025-01-04T08:47:24.606857Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.save(model_ft.state_dict(),'/kaggle/working/segformers.pth')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:47:24.609913Z","iopub.execute_input":"2025-01-04T08:47:24.610197Z","iopub.status.idle":"2025-01-04T08:47:25.129697Z","shell.execute_reply.started":"2025-01-04T08:47:24.610175Z","shell.execute_reply":"2025-01-04T08:47:25.128970Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.load_state_dict(torch.load(\"/kaggle/working/segformers.pth\"))\nmodel.to(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:47:25.130851Z","iopub.execute_input":"2025-01-04T08:47:25.131147Z","iopub.status.idle":"2025-01-04T08:47:25.557790Z","shell.execute_reply.started":"2025-01-04T08:47:25.131122Z","shell.execute_reply":"2025-01-04T08:47:25.556964Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"carry, labels = next(iter(dataloaders['test']))\noutput = model(carry['image'].to(device).float())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:49:21.011278Z","iopub.execute_input":"2025-01-04T08:49:21.011668Z","iopub.status.idle":"2025-01-04T08:49:21.390787Z","shell.execute_reply.started":"2025-01-04T08:49:21.011635Z","shell.execute_reply":"2025-01-04T08:49:21.389843Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Extract logits from the first element of the output tuple\nlogits = output[0]\n\n# Apply softmax to get class probabilities\npred = logits.softmax(dim=1)\n\n# Get the predicted classes (argmax over the class dimension)\npreds = pred.argmax(dim=1)\n\nprint(preds.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:50:31.689454Z","iopub.execute_input":"2025-01-04T08:50:31.689754Z","iopub.status.idle":"2025-01-04T08:50:31.694969Z","shell.execute_reply.started":"2025-01-04T08:50:31.689730Z","shell.execute_reply":"2025-01-04T08:50:31.694217Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_images(image, masks, predicted_masks, num_images=1, color_map = None):\n    # Select num_images random indices to plot\n    indices = np.random.choice(range(len(image)), num_images, replace=False)\n\n    # Plot each image and corresponding masks\n    for i, index in enumerate(indices):\n        fig, axs = plt.subplots(1, 3, figsize=(15, 5))\n        #np.transpose(carry['image'][0].numpy(),(1, 2, 0)), cmap='gray')\n        print(image[0].shape)\n        axs[0].imshow(np.transpose(image[index].numpy(), (1,2,0)), cmap='gray')\n        axs[0].set_title(\"Image\")\n        \n        print(masks.shape)\n        axs[1].imshow(masks[index], cmap=color_map)\n        axs[1].set_title(\"Ground Truth\")\n        \n        axs[2].imshow(predicted_masks[index].cpu(), cmap=color_map)\n        axs[2].set_title(\"Prediction\")\n        \n        plt.show()\n\n\n# Example usage\n# images = torch.randn(4, 3, 224, 224)  # 4 images with shape (3, 224, 224)\n# masks = torch.randint(0, 2, (4, 1, 224, 224))  # 4 ground truth masks with shape (1, 224, 224)\n# predicted_masks = torch.randint(0, 2, (4, 1, 224, 224))  # 4 predicted masks with shape (1, 224, 224)\ncolor_map = plt.cm.get_cmap('tab20', 5)\n\nplot_images(carry['image'], labels, preds, 8, color_map)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:57:48.739363Z","iopub.execute_input":"2025-01-04T08:57:48.739687Z","iopub.status.idle":"2025-01-04T08:57:52.665986Z","shell.execute_reply.started":"2025-01-04T08:57:48.739663Z","shell.execute_reply":"2025-01-04T08:57:52.665147Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Load the trained model\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nmodel = CustomSegFormer2(configuration, in_channel=150)\nmodel.load_state_dict(torch.load(\"/kaggle/working/segformers.pth\"))\nmodel.to(device)\nmodel.eval()\n\n# Test data\ntest_csv = '/kaggle/input/ai-dl-multiclass-segmentation/test_class.csv'\n\ntest_images_path = '/kaggle/input/ai-dl-multiclass-segmentation/TestImages'\ntest_data = pd.read_csv(test_csv)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:50:41.507417Z","iopub.execute_input":"2025-01-04T08:50:41.507752Z","iopub.status.idle":"2025-01-04T08:50:43.944133Z","shell.execute_reply.started":"2025-01-04T08:50:41.507726Z","shell.execute_reply":"2025-01-04T08:50:43.943162Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nsubmission_data = []\n# Prepare test loader\ndef prepare_test_data(test_df, images_path):\n    data = []\n    for _, row in test_df.iterrows():\n        img_path = os.path.join(images_path, row['ImageName'] + \".jpg\")\n        if os.path.exists(img_path):\n            image = Image.open(img_path).convert(\"RGB\")\n            data.append((image, row['ClassNumber']))\n    return data\n\ntest_data_prepared = prepare_test_data(test_data, test_images_path)\n\n# Transform function (modify as per your train transforms)\ntransform = Compose([Normalize(), ToTensorV2()])\n\n# Generate predictions\ndef generate_mask(model, image, class_number, transform):\n    input_tensor = transform(image=np.array(image))['image']\n    input_tensor = input_tensor.unsqueeze(0).to(device)\n    with torch.no_grad():\n        logits, _ = model(input_tensor)\n    probs = torch.softmax(logits, dim=1)\n    mask = (probs.argmax(dim=1) == class_number).squeeze().cpu().numpy()\n    return mask\n\ndef encode_mask_to_rle(mask):\n    '''\n    mask: numpy array binary mask \n    1 - mask \n    0 - background\n    Returns encoded run length \n    '''\n    pixels = mask.flatten()\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return ' '.join(str(x) for x in runs)\n\n\n# Visualize first 5 predictions and generate submission data\ndef visualize_first_5_predictions(test_data, test_df, model, transform):\n    # Select the first 5 samples\n    samples = test_data[0:50]\n    for image, class_number in samples:\n        \n        # Extract the image name from the test_df DataFrame using the index\n        # Assuming the test_df has 'ImageName' and 'ClassNumber' columns, we get the first match\n        image_name = test_df.iloc[samples.index((image, class_number))]['ImageName']\n        \n        # Generate the mask for the image\n        mask = generate_mask(model, image, class_number, transform)\n        \n        # Append to submission_data\n        submission_data.append({'ImageName': image_name, 'Encoding': encode_mask_to_rle(mask)})\n        \n        # Print the RLE encoding with the correct image name\n        # print(f\"Image: {image_name}, Encoding: {encode_mask_to_rle(mask)}\")\n        # c+=1\n        # print(c)\n        # Visualize the image and its predicted mask\n        fig, ax = plt.subplots(1, 2, figsize=(12, 6))\n        ax[0].imshow(image)\n        ax[0].set_title(f\"Original Image - Class {class_number}\")\n        ax[1].imshow(mask, cmap=\"gray\")\n        ax[1].set_title(f\"Predicted Mask - Class {class_number}\")\n        plt.show()\n\n# Run visualization for the first 5 images\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:55:49.795916Z","iopub.execute_input":"2025-01-04T08:55:49.796252Z","iopub.status.idle":"2025-01-04T08:55:59.850666Z","shell.execute_reply.started":"2025-01-04T08:55:49.796226Z","shell.execute_reply":"2025-01-04T08:55:59.849978Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"visualize_first_5_predictions(test_data_prepared, test_data, model, transform)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-04T08:55:59.851701Z","iopub.execute_input":"2025-01-04T08:55:59.851907Z","iopub.status.idle":"2025-01-04T08:56:24.962922Z","shell.execute_reply.started":"2025-01-04T08:55:59.851888Z","shell.execute_reply":"2025-01-04T08:56:24.962031Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}