{"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":"none","dataSources":[{"sourceId":76972,"databundleVersionId":8669583,"sourceType":"competition"}],"dockerImageVersionId":30822,"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},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Building Instance Segmentation with MaskRCNN","metadata":{}},{"cell_type":"code","source":"# Downloads helper functions py files\nimport os\n\n# os.system(\"wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/engine.py\")\n# os.system(\"wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/utils.py\")\n# os.system(\"wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/coco_utils.py\")\n# os.system(\"wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/coco_eval.py\")\n# os.system(\"wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/transforms.py\")\n\n#Used only for showing images\n# os.system(\"wget https://raw.githubusercontent.com/pytorch/vision/refs/heads/main/gallery/transforms/helpers.py\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-01T11:46:11.128574Z","iopub.status.idle":"2025-02-01T11:46:11.129016Z","shell.execute_reply":"2025-02-01T11:46:11.128835Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport torch\nimport json\nimport utils\nimport cv2\nimport torchvision\nimport random\n\nimport numpy as np\n\nfrom torchvision import transforms as _transforms\nfrom typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Type, Union\nfrom torchvision.io import read_image\nfrom torchvision.ops.boxes import masks_to_boxes\nfrom torchvision.ops import MultiScaleRoIAlign\nfrom torchvision import tv_tensors\nfrom torchvision.transforms.v2 import functional as F, InterpolationMode, Transform, RandomZoomOut\nfrom torchvision.transforms import v2 as T\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\nfrom torchvision.models.detection.mask_rcnn import MaskRCNNHeads\nfrom torchvision.models.detection.faster_rcnn import TwoMLPHead\nfrom torchvision.models.detection.rpn import AnchorGenerator\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection import MaskRCNN\nfrom torchvision.models.detection.backbone_utils import BackboneWithFPN","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T19:59:36.505649Z","iopub.execute_input":"2026-06-10T19:59:36.505902Z","iopub.status.idle":"2026-06-10T19:59:42.593415Z","shell.execute_reply.started":"2026-06-10T19:59:36.505878Z","shell.execute_reply":"2026-06-10T19:59:42.592558Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(torch.cuda.device_count())\nfor i in range(torch.cuda.device_count()):\n    print(torch.cuda.get_device_name(i))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T20:02:32.484774Z","iopub.execute_input":"2026-06-10T20:02:32.485100Z","iopub.status.idle":"2026-06-10T20:02:32.774600Z","shell.execute_reply.started":"2026-06-10T20:02:32.485071Z","shell.execute_reply":"2026-06-10T20:02:32.773537Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Dataset tranforms functions","metadata":{}},{"cell_type":"code","source":"# Rotating images and masks trasform\nclass Rotate(object):\n    def __call__(self, image, target):        \n        rnd = random.randint(0,3)\n        degree = rnd * 90\n        if (degree != 0):\n            image = F.rotate(image, degree)\n            rotated_masks = F.rotate(target[\"masks\"], degree)\n            target[\"boxes\"] = tv_tensors.BoundingBoxes(masks_to_boxes(rotated_masks), format=\"XYXY\", canvas_size=F.get_size(image))\n            target[\"masks\"] = rotated_masks\n            return image, target\n        else:\n            return image, target\n            \nclass Flip(object):\n    def __call__(self, image, target):\n        # flipped_image = none\n        # flipped_masks = none\n        rnd = random.random()\n        if(rnd >= 0.5):\n            if (rnd >= 0.75):\n                flipped_image = F.hflip(image)\n                flipped_masks = F.hflip(target[\"masks\"])\n            else:\n                flipped_image = F.vflip(image)\n                flipped_masks = F.vflip(target[\"masks\"])\n            target[\"boxes\"] = tv_tensors.BoundingBoxes(masks_to_boxes(flipped_masks), format=\"XYXY\", canvas_size=F.get_size(image))\n            target[\"masks\"] = flipped_masks\n            return flipped_image, target\n        else:\n            return image, target\n        \nclass Zoomout(object):\n    def __call__(self, image, target):\n        zoom_out = RandomZoomOut(side_range=(1.2, 1.4), fill=0.5, p=0.5)\n        out = zoom_out({\"input\": image, \"target\": target[\"masks\"]})                \n        target[\"boxes\"] = tv_tensors.BoundingBoxes(masks_to_boxes(out[\"target\"]), format=\"XYXY\", canvas_size=F.get_size(image))\n        target[\"masks\"] = out[\"target\"]\n        return out[\"input\"], target\n   \n\n# Decreasing image quality transform\nclass DecreaseQuality(object):\n    \n    def __call__(self, image, target):     \n        if(random.random() >= 0.5):\n            width = image.shape[1]\n            height = image.shape[2]\n            decrease_size = T.Resize((int(width/2), int(height/2)))  # Specify desired size\n            increase_size = T.Resize((width, height))\n            resized_image = decrease_size(image)        \n            decreased_quality_img = increase_size(resized_image)\n            return decreased_quality_img, target\n        else:\n            return image, target","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T16:12:26.135114Z","iopub.execute_input":"2026-06-10T16:12:26.135480Z","iopub.status.idle":"2026-06-10T16:12:26.145744Z","shell.execute_reply.started":"2026-06-10T16:12:26.135453Z","shell.execute_reply":"2026-06-10T16:12:26.144804Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Reading files functions","metadata":{}},{"cell_type":"code","source":"# Dataset Helper Functions\ndef read_json_file(file_path):\n    with open(file_path, 'r') as file:\n        data = json.load(file)\n    return data\n    \ndef get_image_name(image_path):\n    splited = image_path.split(\"/\")\n    return \"/\".join(splited[-2:])\n\ndef get_image_annotations(json_data, image_name):\n    image_dic = None\n    annot_list = []\n    for image in json_data['images']:\n        if (image['file_name'] == image_name):\n            image_dic = image\n            break\n    found = False\n    for annot in json_data['annotations']:\n        if (annot['image_id'] == image_dic['id']):\n            annot_list.append(annot)\n            found = True\n        if (found and annot['image_id'] != image_dic['id']):            \n            break\n    \n    # Convert the list of annotations dictionary to a dictionary of annotaions list\n    annot_dic = {key: [d[key][0] if(key == 'segmentation') else d[key] for d in annot_list] for key in annot_list[0].keys()}    \n    annotation = Annotations(annot_dic, image_dic)\n    return annotation\n\ndef create_pairs(input_list):\n  \"\"\"Converts a list to a set of pairs.\n\n  Args:\n    input_list: The input list.\n\n  Returns:\n    A set of pairs.\n  \"\"\"\n  pairs = []\n  for i in range(0, len(input_list) - 1, 2):\n    pairs.append(tuple(input_list[i:i+2]))\n  return pairs\n\ndef create_contour_mask(contour_points, image_shape):\n    \"\"\"\n    Create a binary mask from contour points.\n\n    Args:\n        contour_points: List of tuples [(x,y), (x1,y1), ..., (xn,yn)]\n        image_shape: Tuple of (height, width) for the output mask\n\n    Returns:\n        Binary mask as numpy array where filled contour is 1, rest is 0\n    \"\"\"\n    # Convert list of tuples to numpy array and reshape for cv2    \n    contour = np.array(contour_points).reshape((-1,1,2)).astype(np.int32)\n\n    # Create empty mask\n    mask = np.zeros(image_shape, dtype=np.uint8)\n\n    # Fill the contour\n    cv2.fillPoly(mask, [contour], color=1)\n\n    return torch.tensor(mask)\n\ndef create_masks_from_contours(point_list, image_width, image_height):\n    masks = []\n    image_shape = (image_width, image_height)\n    for seg in point_list:        \n        points_pair = create_pairs(seg)  \n        if (len(points_pair) > 3):\n            masks.append(create_contour_mask(points_pair, image_shape))            \n    if (len(masks) == 0):\n        return None\n    return torch.stack(masks, dim=0)\n\ndef clean_boxes_masks(boxes, masks):        \n    removed_indices = []\n    for i, box in enumerate(boxes):\n        if( box[0] == box[2] or box[1] == box[3]):\n            removed_indices.append(i)\n            \n    if (len(removed_indices) > 0):\n        filtered_numbers = [num for num in range(0, masks.shape[0]) if num not in removed_indices]\n        boxes = boxes[filtered_numbers]\n        masks = masks[filtered_numbers]\n\n    return boxes, masks\n\ndef get_transform(train):\n    transforms = []\n    if train:        \n        # transforms.append(DecreaseQuality())\n        transforms.append(Rotate())\n        transforms.append(Flip())\n        # transforms.append(Zoomout())\n    transforms.append(T.ToDtype(torch.float, scale=True))\n    transforms.append(T.ToPureTensor())\n    return T.Compose(transforms)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T16:12:29.429533Z","iopub.execute_input":"2026-06-10T16:12:29.429863Z","iopub.status.idle":"2026-06-10T16:12:29.441693Z","shell.execute_reply.started":"2026-06-10T16:12:29.429840Z","shell.execute_reply":"2026-06-10T16:12:29.440808Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Data sample class","metadata":{}},{"cell_type":"code","source":"# Annotations Class\nclass Annotations():\n    def __init__(self, annotation_dic, image_dic):\n        self.image_id = image_dic['id']\n        self.image_height = image_dic['height']\n        self.image_width = image_dic['width']\n        self.bbox_list = annotation_dic['bbox']\n        self.segmentation_list = annotation_dic['segmentation']\n        self.area_list = annotation_dic['area']    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T16:12:33.715842Z","iopub.execute_input":"2026-06-10T16:12:33.716206Z","iopub.status.idle":"2026-06-10T16:12:33.721216Z","shell.execute_reply.started":"2026-06-10T16:12:33.716179Z","shell.execute_reply":"2026-06-10T16:12:33.719999Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Dataset class","metadata":{}},{"cell_type":"code","source":"# Create Dataset\nclass BuildingDataset(torch.utils.data.Dataset):\n    def __init__(self, root, annotation_file, transforms):\n        self.root = root\n        self.transforms = transforms\n        # load all image files, sorting them to\n        # ensure that they are aligned\n        self.imgs = list(sorted(os.listdir(os.path.join(root, \"image\"))))\n        self.annot_json = read_json_file(os.path.join(root, annotation_file))\n\n    def __getitem__(self, idx):\n        # load images and masks        \n        img_path = os.path.join(self.root, \"image\", self.imgs[idx])        \n        image_name = get_image_name(img_path)        \n        img = read_image(img_path)                \n        \n        # get image annotations\n        annotations = get_image_annotations(self.annot_json, image_name)                             \n\n        # generate masks from contours          \n        masks = create_masks_from_contours(annotations.segmentation_list, annotations.image_width, annotations.image_height)        \n        \n        if(masks == None):            \n            return None, None\n        \n        # generate boxes from masks\n        boxes = masks_to_boxes(masks)        \n           \n        # remove the boxes that form a line\n        boxes, masks = clean_boxes_masks(boxes, masks)    \n\n        num_objs = len(boxes)\n\n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        \n        # there is only one class\n        labels = torch.ones((num_objs,), dtype=torch.int64)\n\n        image_id = idx        \n        # suppose all instances are not crowd\n        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)\n\n        # Wrap sample and targets into torchvision tv_tensors:\n        img = tv_tensors.Image(img)\n\n        target = {}\n        target[\"boxes\"] = tv_tensors.BoundingBoxes(boxes, format=\"XYXY\", canvas_size=F.get_size(img))        \n        target[\"masks\"] = tv_tensors.Mask(masks)\n        target[\"labels\"] = labels\n        target[\"image_id\"] = image_id\n        target[\"area\"] = area\n        target[\"iscrowd\"] = iscrowd\n\n        if self.transforms is not None:\n            img, target = self.transforms(img, target)\n\n        return img, target\n\n    def __len__(self):\n        return len(self.imgs)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T16:45:09.344516Z","iopub.execute_input":"2026-06-10T16:45:09.344900Z","iopub.status.idle":"2026-06-10T16:45:09.354031Z","shell.execute_reply.started":"2026-06-10T16:45:09.344875Z","shell.execute_reply":"2026-06-10T16:45:09.352937Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Show dataset samples","metadata":{}},{"cell_type":"code","source":"# Show samples from dataset\nimport matplotlib.pyplot as plt\nfrom helpers import plot\nfrom torchvision.tv_tensors import Image\nfrom torchvision.utils import draw_bounding_boxes\n\npath = '/kaggle/input/building-extraction-generalization-2024'\ndataset = BuildingDataset(f'{path}/train', 'train.json', get_transform(train=True))\ndata_loader = torch.utils.data.DataLoader(\n    dataset,\n    batch_size=2,\n    shuffle=False,\n    collate_fn=utils.collate_fn\n)\n\ni = 0\nfor images, targets in iter(data_loader):\n    # rotated_img = F.rotate(images[0], 90)\n    # rotated_mask = F.rotate(targets[0][\"masks\"], 90)\n    # rotated_bbox = masks_to_boxes(rotated_mask)    \n    # print(\"ROTATED BBOX:\",type(rotated_bbox))\n    print(\"BOX numbers:\", targets[0][\"boxes\"].shape)\n    image_with_bbox = draw_bounding_boxes(images[0],  targets[0][\"boxes\"], width=2)\n    print(\"IAMGE ID:\", targets[0][\"image_id\"])    \n    print(images[0].shape, targets[0][\"image_id\"])\n    img = Image(images[0])\n    plot([images[0]])\n    plot([image_with_bbox])\n    plot([targets[0][\"masks\"][0]])\n    i += 1\n    if(i==1):\n        break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T16:12:45.494661Z","iopub.execute_input":"2026-06-10T16:12:45.494971Z","iopub.status.idle":"2026-06-10T16:12:48.354266Z","shell.execute_reply.started":"2026-06-10T16:12:45.494949Z","shell.execute_reply":"2026-06-10T16:12:48.353093Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"'''\n# Create MaskRCNN model efficientnet_v2_m \ndef get_model_instance_segmentation(num_classes, default_weights = True):\n    # load an instance segmentation model pre-trained on COCO    \n    if(default_weights):\n        backbone = torchvision.models.efficientnet_v2_m(weights=\"DEFAULT\")\n        # model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights=\"DEFAULT\")\n    else:\n        backbone = torchvision.models.efficientnet_v2_m()\n        # model = torchvision.models.detection.maskrcnn_resnet50_fpn()\n\n    # Remove the classifier head (we only need the feature extractor)\n    backbone.classifier = torch.nn.Identity()\n    \n    # return_layers = {\"2\": \"0\", \"4\": \"1\"}\n    return_layers = {'2': 'layer1', '3': 'layer2', '4': 'layer3', '6': 'layer4'}\n    # in_channels_list = [48, 160]  # Example channel sizes for the chosen layers.\n    in_channels_list = [48, 80, 160, 304]\n    out_channels = 256  # FPN output channels\n    \n    # Wrap the EfficientNet features with an FPN.\n    # BackboneWithFPN will internally assign an attribute 'out_channels' that MaskRCNN expects.\n    backbone_with_fpn = BackboneWithFPN(\n        backbone=backbone.features,\n        return_layers=return_layers,\n        in_channels_list=in_channels_list,\n        out_channels=out_channels,\n    )\n\n    # dummy_input = torch.randn(1, 3, 224, 224)\n    # features = backbone_with_fpn(dummy_input)\n    # print(\"FPN feature map keys:\", list(features.keys()))\n    \n    model = MaskRCNN(backbone_with_fpn, num_classes=num_classes)\n\n    # Step 4: Modify the ROI Pooling layer to output 14x14 feature maps\n    model.roi_heads.box_roi_pool = MultiScaleRoIAlign(\n        featmap_names=['layer1', 'layer2', 'layer3', 'layer4', 'pool'],  # FPN feature levels\n        output_size=28,  # Change ROI Align output size from 7 to 14\n        sampling_ratio=2\n    )\n    \n    # Update the ROI Pooling layer for the mask head as well\n    model.roi_heads.mask_roi_pool = MultiScaleRoIAlign(\n        featmap_names=['layer1', 'layer2', 'layer3', 'layer4', 'pool'],  # Must match FPN output keys\n        output_size=28,\n        sampling_ratio=2\n    )\n    \n    # sizes=((32,), (64,), (128,), (256), (512)),\n    anchor_generator = AnchorGenerator(\n    sizes = ((8,), (16,), (32,), (64,), (128,)),    \n    aspect_ratios=((0.75, 1.0, 1.5),) * 5\n    )    \n\n    model.rpn.anchor_generator = anchor_generator\n    \n    # Step 5: Update the box head since the feature size has changed\n    in_channels = 256 * 28 * 28  # Adjusted due to new ROI pooling size\n    representation_size = 1024  # Default\n    \n    # Replace the box head to accept the updated feature size\n    model.roi_heads.box_head = TwoMLPHead(in_channels, representation_size)\n    \n    # Replace the box predictor to match the new representation size\n    num_classes = model.roi_heads.box_predictor.cls_score.out_features  # Keep the same number of classes\n    model.roi_heads.box_predictor = FastRCNNPredictor(representation_size, num_classes)\n\n    print(\"ROI:\", model.roi_heads.box_roi_pool.output_size)\n\n    return model\n    '''","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-28T05:42:20.540998Z","iopub.execute_input":"2025-04-28T05:42:20.541309Z","iopub.status.idle":"2025-04-28T05:42:20.548832Z","shell.execute_reply.started":"2025-04-28T05:42:20.541281Z","shell.execute_reply":"2025-04-28T05:42:20.547749Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## CustomRoiHeads\nThis class is created to change mask loss weight","metadata":{}},{"cell_type":"code","source":"from torchvision.models.detection.roi_heads import RoIHeads\nfrom torchvision.models.detection.roi_heads import fastrcnn_loss, maskrcnn_loss\n\n# CustomRoIHead created to change mask loss weight\nclass CustomRoIHeads(RoIHeads):\n    def __init__(self, *args, mask_loss_weight=1.0, **kwargs):\n        super().__init__(*args, **kwargs)\n        self.mask_loss_weight = mask_loss_weight\n\n    def forward(self, features, proposals, image_shapes, targets=None):\n        # Standard forward, mostly same as original torchvision code\n        detections, detector_losses = super().forward(features, proposals, image_shapes, targets)\n\n        # Apply mask loss weighting only if in training and mask is enabled\n        if self.has_mask():\n            if self.training:\n                boxes = [d[\"boxes\"] for d in detections]\n                if any(box.shape[0] > 0 for box in boxes):\n                    mask_features = self.mask_roi_pool(features, boxes, image_shapes)\n                    mask_features = self.mask_head(mask_features)\n                    mask_logits = self.mask_predictor(mask_features)\n    \n                    loss_mask = maskrcnn_loss(mask_logits, detections, self.mask_roi_pool)\n                    detector_losses[\"loss_mask\"] = loss_mask * self.mask_loss_weight\n                else:\n                    detector_losses[\"loss_mask\"] = torch.tensor(0.0, device=features[list(features.keys())[0]].device)\n        return detections, detector_losses","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T16:45:18.811092Z","iopub.execute_input":"2026-06-10T16:45:18.811419Z","iopub.status.idle":"2026-06-10T16:45:18.819301Z","shell.execute_reply.started":"2026-06-10T16:45:18.811393Z","shell.execute_reply":"2026-06-10T16:45:18.817965Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## MaskRCNN model. RESNET 50 Backbone","metadata":{}},{"cell_type":"code","source":"# Create MaskRCNN model RESNET 50\ndef get_model_instance_segmentation(num_classes, default_weights = True):\n    # load an instance segmentation model pre-trained on COCO    \n    if(default_weights):\n        model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights=\"DEFAULT\")\n    else:\n        model = torchvision.models.detection.maskrcnn_resnet50_fpn()\n\n    # changed for better detetion on small buildings \n    anchor_generator = AnchorGenerator(\n    sizes = ((8,), (16,), (32,), (64,), (128,)),    \n    aspect_ratios=((0.75, 1.0, 1.5),) * 5\n    )    \n\n    model.rpn.anchor_generator = anchor_generator\n\n    roi_output_size = 28\n\n    multiscale_roi_align = MultiScaleRoIAlign(\n    featmap_names=['0', '1', '2', '3'],  # These are the default feature map names used with FPN.\n    output_size=roi_output_size,         # Desired output size (if an int is provided, the pooled feature map will be (14, 14))\n    sampling_ratio=2                     # Default sampling ratio (adjust if needed)\n    )\n    model.roi_heads.box_roi_pool = multiscale_roi_align\n    model.roi_heads.mask_roi_pool = multiscale_roi_align\n\n\n    in_channels = 256 * roi_output_size * roi_output_size      # new number of input features\n    representation_size = 1024       # typically remains the same\n\n    # Replace the box head with a new TwoMLPHead\n    model.roi_heads.box_head = TwoMLPHead(in_channels, representation_size)\n\n\n    print(\"CLS IN FEATURES:\", model.roi_heads.box_predictor.cls_score.in_features)\n    # get number of input features for the classifier\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n\n    # replace the pre-trained head with a new one\n    model.roi_heads.box_predictor = FastRCNNPredictor(representation_size, num_classes)\n\n    # Increase mask head conv layers from 4 to 6 and 256 channels to 512\n    mask_head = MaskRCNNHeads(256, (512,512,512,512,512,512), 1)\n    '''Comment this line to use CustomHead'''\n    model.roi_heads.mask_head = mask_head\n    mask_head_out_channels = mask_head[-1][0].out_channels\n\n    print(\"MASK IN CHANNELLS:\", model.roi_heads.mask_predictor.conv5_mask.in_channels)\n    # now get the number of input features for the mask classifier\n    in_features_mask = mask_head_out_channels\n    hidden_layer = mask_head_out_channels\n    # and replace the mask predictor with a new one\n    # model.roi_heads.mask_predictor.conv5_mask = torch.nn.ConvTranspose2d(in_features, 256, kernel_size=2, stride=2)\n    '''Comment this line to use CustomHead'''\n    model.roi_heads.mask_predictor = MaskRCNNPredictor(\n        in_features_mask,\n        hidden_layer,\n        num_classes\n    )\n\n    '''\n    # Replace RoIHeads with custom one\n    model.roi_heads = CustomRoIHeads(\n        # copy args from original roi_heads\n        box_roi_pool = multiscale_roi_align,\n        box_head = TwoMLPHead(in_channels, representation_size),\n        box_predictor = FastRCNNPredictor(representation_size, num_classes),        \n        fg_iou_thresh=0.5,\n        bg_iou_thresh=0.5,\n        batch_size_per_image=512,\n        positive_fraction=0.25,\n        bbox_reg_weights=None,\n        score_thresh=0.05,\n        nms_thresh=0.5,\n        detections_per_img=100,\n        mask_roi_pool = multiscale_roi_align,\n        mask_head = mask_head,\n        mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes),\n        mask_loss_weight = 1.0  # <-- Your custom weight here\n    )\n    '''\n\n    print(\"ROI:\", model.roi_heads.box_roi_pool.output_size)\n\n    return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T16:45:24.449185Z","iopub.execute_input":"2026-06-10T16:45:24.449543Z","iopub.status.idle":"2026-06-10T16:45:24.458240Z","shell.execute_reply.started":"2026-06-10T16:45:24.449515Z","shell.execute_reply":"2026-06-10T16:45:24.457107Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Source code of torchvision.models.detection.mask_rcnn.MaskRCNNHeads:\n\nhttps://docs.pytorch.org/vision/0.12/_modules/torchvision/models/detection/mask_rcnn.html","metadata":{}},{"cell_type":"markdown","source":"## Train functions","metadata":{}},{"cell_type":"code","source":"import math\nimport sys\nimport time\n\nimport torch\nimport torchvision.models.detection.mask_rcnn\nimport utils\nfrom coco_eval import CocoEvaluator\nfrom coco_utils import get_coco_api_from_dataset\n\ndef train_one_epoch_overridden(model, optimizer, data_loader, device, epoch, print_freq, scaler=None):    \n    model.train()\n    metric_logger = utils.MetricLogger(delimiter=\"  \")\n    metric_logger.add_meter(\"lr\", utils.SmoothedValue(window_size=1, fmt=\"{value:.6f}\"))\n    header = f\"Epoch: [{epoch}]\"\n\n    lr_scheduler = None\n    if epoch == 0:\n        warmup_factor = 1.0 / 1000\n        warmup_iters = min(1000, len(data_loader) - 1)\n\n        lr_scheduler = torch.optim.lr_scheduler.LinearLR(\n            optimizer, start_factor=warmup_factor, total_iters=warmup_iters\n        )\n\n    for images, targets in metric_logger.log_every(data_loader, print_freq, header):        \n        images = list(image.to(device) for image in images if image is not None)\n        targets = [{k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in t.items()} for t in targets if t is not None]\n        # with torch.cuda.amp.autocast(enabled=scaler is not None): \n        with torch.amp.autocast('cuda', enabled=scaler is not None):\n            #-----------TODO: Delete-----------------------\n            print(\"model.transform.image_mean:\", model.transform.image_mean)\n            print(\"model.transform.image_std:\", model.transform.image_std)\n            #----------------------------------            \n            loss_dict = model(images, targets)\n            losses = sum(loss for loss in loss_dict.values())\n\n        # reduce losses over all GPUs for logging purposes\n        loss_dict_reduced = utils.reduce_dict(loss_dict)\n        losses_reduced = sum(loss for loss in loss_dict_reduced.values())\n\n        loss_value = losses_reduced.item()\n\n        if not math.isfinite(loss_value):\n            print(f\"Loss is {loss_value}, stopping training\")\n            print(loss_dict_reduced)\n            sys.exit(1)\n\n        optimizer.zero_grad()\n        if scaler is not None:\n            scaler.scale(losses).backward()\n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            losses.backward()\n            optimizer.step()\n\n        if lr_scheduler is not None:\n            lr_scheduler.step()\n\n        metric_logger.update(loss=losses_reduced, **loss_dict_reduced)\n        metric_logger.update(lr=optimizer.param_groups[0][\"lr\"])\n\n    return metric_logger\n\ndef saveModel(model, epoch, optimizer):\n    state = {'epoch': epoch + 1, 'state_dict': model.state_dict(),\n             'optimizer': optimizer.state_dict()}\n    torch.save(state, 'model_weights.pth')    \n\ndef loadModel(model, optimizer, device = None):\n    start_epoch = 0\n    checkpoint = None\n    if os.path.isfile('model_weights.pth'):\n        if (device is None):\n            checkpoint = torch.load('model_weights.pth')\n        else:\n            checkpoint = torch.load('model_weights.pth', map_location=device)\n        start_epoch = checkpoint['epoch']\n        model.load_state_dict(checkpoint['state_dict'])\n        if (optimizer is not None):\n            optimizer.load_state_dict(checkpoint['optimizer'])    \n    \n    return model, optimizer, start_epoch","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T16:51:13.073126Z","iopub.execute_input":"2026-06-10T16:51:13.073545Z","iopub.status.idle":"2026-06-10T16:51:13.088255Z","shell.execute_reply.started":"2026-06-10T16:51:13.073514Z","shell.execute_reply":"2026-06-10T16:51:13.087116Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Train the model","metadata":{}},{"cell_type":"code","source":"# Train and eval model\nfrom engine import train_one_epoch, evaluate\n\ndata_path = '/kaggle/input/building-extraction-generalization-2024'\n\n# train on the GPU or on the CPU, if a GPU is not available\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\n# our dataset has two classes only - background and person\nnum_classes = 2\nstart_epoch = 0\n# use our dataset and defined transformations\ndataset =  BuildingDataset(f'{data_path}/train',  'train.json', get_transform(train=True))\ndataset_val =  BuildingDataset(f'{data_path}/val',  'val.json', get_transform(train=False))\n\n#TODO:delete\n# print(len(dataset))\n# indices = torch.tensor(range(0,len(dataset))).tolist()\n# dataset = torch.utils.data.Subset(dataset, indices[0:200])\n\n# define training and validation data loaders\ndata_loader = torch.utils.data.DataLoader(\n    dataset,\n    batch_size=2,\n    shuffle=True,\n    collate_fn=utils.collate_fn\n)\n\ndata_loader_val = torch.utils.data.DataLoader(\n    dataset_val,\n    batch_size=1,\n    shuffle=False,\n    collate_fn=utils.collate_fn\n)\n\n# get the model using our helper function\nmodel = get_model_instance_segmentation(num_classes)\n# To use a single GPU again, uncomment the next line and comment out\n# the following two lines.\nmodel.to(device)\n# model = torch.nn.DataParallel(model)\n# model = model.cuda()\n\n# Change backbone learning rate\nbackbone_params = []\nother_params = []\nfor name, param in model.named_parameters():\n    if \"backbone\" in name:\n        backbone_params.append(param)\n    else:\n        other_params.append(param)\n\n# params = [p for p in model.parameters() if p.requires_grad]\n# optimizer = torch.optim.SGD(\n#     params,\n#     lr=0.07,\n#     momentum=0.9,\n#     weight_decay=0.0005\n# )\n\n# Define optimizer with different learning rates\noptimizer = torch.optim.SGD([\n    {\"params\": other_params, \"lr\": 0.07, \"weight_decay\":0.0005},\n    {\"params\": backbone_params, \"lr\": 0.0008, \"weight_decay\":0.0005}\n])\n\nprint(\"BACKBONE:\", len(backbone_params))\nprint(\"OTHER:\", len(other_params))\n\nmodel, optimizer, start_epoch = loadModel(model, optimizer)\n\n# for g in optimizer.param_groups:\n#     #backbone params\n#     if(g[\"initial_lr\"] == 0.007):\n#         g[\"lr\"] = 0.0032\n#         g[\"initial_lr\"] = 0.0032\n        \n    \nprint(\"START EPOCH:\", start_epoch)\n\n# move model to the right device\n# Uncomment to use a single GPU again\nmodel.to(device)\n\n# Change trainable layers config\n# for name, param in model.named_parameters():\n#     param.requires_grad = True\n\n\n# and a learning rate scheduler\nlr_scheduler = torch.optim.lr_scheduler.StepLR(\n    optimizer,\n    step_size=3,\n    gamma=0.8\n)\n\n# let's train it just for 2 epochs\nnum_epochs = 5\n\nfor epoch in range(start_epoch, start_epoch + num_epochs):\n    # train for one epoch, printing every 10 iterations\n    train_one_epoch_overridden(model, optimizer, data_loader, device, epoch, print_freq=50)\n    # Saving the model\n    saveModel(model, epoch, optimizer)\n    print(\"Model saved!\")\n    # update the learning rate    \n    lr_scheduler.step()\n    # evaluate on the test dataset    \n    evaluate(model, data_loader_val, device=device)\n\nprint(\"That's it!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T16:51:20.051718Z","iopub.execute_input":"2026-06-10T16:51:20.052044Z","iopub.status.idle":"2026-06-10T16:51:40.865329Z","shell.execute_reply.started":"2026-06-10T16:51:20.052022Z","shell.execute_reply":"2026-06-10T16:51:40.863952Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"'''for idx in range(len(dataset)):\n    try:\n        img, target = dataset[idx]\n\n        if img is None:\n            print(f\"Sample {idx}: image is None\")\n            continue\n\n        if img.ndim != 3:\n            print(f\"Sample {idx}: Invalid shape {img.shape}\")\n            continue\n\n        if img.shape[0] != 3:\n            print(f\"Sample {idx}: {img.shape}\")\n\n    except Exception as e:\n        print(f\"Sample {idx}: ERROR - {e}\")\n'''\nfor batch_idx, (images, targets) in enumerate(data_loader):\n    for img_idx, img in enumerate(images):\n\n        if img is None:\n            print(f\"Batch {batch_idx}, image {img_idx}: None\")\n            continue\n\n        # print(f\"Batch {batch_idx}, image {img_idx}: {img.shape}\")\n\n        if img.ndim != 3 or img.shape[0] != 3:\n            print(f\"PROBLEM: Batch {batch_idx}, image {img_idx}: {img.shape}\")\n            raise RuntimeError(\"Found invalid image\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-10T16:13:46.325948Z","iopub.execute_input":"2026-06-10T16:13:46.326271Z","iopub.status.idle":"2026-06-10T16:18:39.969939Z","shell.execute_reply.started":"2026-06-10T16:13:46.326246Z","shell.execute_reply":"2026-06-10T16:18:39.969093Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# -------------------------- Show Model's Prediction --------------------------------------------","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport pandas as pd\nfrom PIL import Image\nfrom torchvision.utils import draw_bounding_boxes\nfrom torchvision.utils import draw_segmentation_masks","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-07T07:25:55.108957Z","iopub.execute_input":"2026-06-07T07:25:55.109237Z","iopub.status.idle":"2026-06-07T07:25:55.354226Z","shell.execute_reply.started":"2026-06-07T07:25:55.109216Z","shell.execute_reply":"2026-06-07T07:25:55.353207Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\nmodel = get_model_instance_segmentation(num_classes = 2, default_weights = False)\nprint(device)\n# model.load_state_dict(torch.load('model_weights.pth', map_location=device, weights_only=True))\nmodel, _, _ = loadModel(model, None, device)\n# uncomment to get back to single GPU\nmodel.to(device)\n# comment to get back to single GPU\n# model = torch.nn.DataParallel(model)\n# model = model.cuda()\n#---------------------------------\nmodel.eval()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# model.backbone\nmodel.roi_heads.box_roi_pool.output_size","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-06T18:46:36.567465Z","iopub.execute_input":"2026-06-06T18:46:36.567904Z","iopub.status.idle":"2026-06-06T18:46:36.574044Z","shell.execute_reply.started":"2026-06-06T18:46:36.567869Z","shell.execute_reply":"2026-06-06T18:46:36.572856Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pil_image = Image.open(\"/kaggle/input/building-extraction-generalization-2024/test/image/0144.tif\")\n# pil_image = Image.open(\"/kaggle/input/building-extraction-generalization-2024/test/image/0112.tif\")\nimage_tensor = F.to_tensor(pil_image)\nx = [image_tensor.to(device)]\npredictions = model(x)\n\n# print(predictions)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-07T07:26:20.487595Z","iopub.execute_input":"2026-06-07T07:26:20.487906Z","iopub.status.idle":"2026-06-07T07:26:22.420448Z","shell.execute_reply.started":"2026-06-07T07:26:20.487881Z","shell.execute_reply":"2026-06-07T07:26:22.419606Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------ Show Model's prediction ---------------------\nplt.rcParams[\"savefig.bbox\"] = 'tight'\n\ndef show(imgs):\n    if not isinstance(imgs, list):\n        imgs = [imgs]\n    fig, axs = plt.subplots(ncols=len(imgs), squeeze=False)\n    for i, img in enumerate(imgs):\n        img = img.detach()\n        img = F.to_pil_image(img)\n        axs[0, i].imshow(np.asarray(img))\n        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-07T07:26:25.704947Z","iopub.execute_input":"2026-06-07T07:26:25.705237Z","iopub.status.idle":"2026-06-07T07:26:25.710542Z","shell.execute_reply.started":"2026-06-07T07:26:25.705214Z","shell.execute_reply":"2026-06-07T07:26:25.709389Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"score_threshold = 0.5\n# boxes = torch.tensor([[199, 493,  19,   7],[300, 250, 400, 350]], dtype=torch.float)\n# colors = [\"blue\", \"yellow\"]\nindices = torch.where(predictions[0][\"scores\"] > 0.5)\nresult = draw_bounding_boxes(image_tensor, predictions[0]['boxes'][indices], width=3)\n# result = draw_bounding_boxes(image_tensor, boxes, width=5)\nprint(image_tensor.shape)\nshow(result)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-07T07:26:44.935514Z","iopub.execute_input":"2026-06-07T07:26:44.935851Z","iopub.status.idle":"2026-06-07T07:26:45.234617Z","shell.execute_reply.started":"2026-06-07T07:26:44.935824Z","shell.execute_reply":"2026-06-07T07:26:45.233731Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"predictions[0].keys()\nprint(len(predictions[0][\"scores\"]))\nprint(predictions[0][\"scores\"])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-07T07:26:55.767585Z","iopub.execute_input":"2026-06-07T07:26:55.767888Z","iopub.status.idle":"2026-06-07T07:26:55.821362Z","shell.execute_reply.started":"2026-06-07T07:26:55.767864Z","shell.execute_reply":"2026-06-07T07:26:55.820604Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(len(indices[0]))\nmasks = predictions[0][\"masks\"][indices]\nmasks = (masks > 0.5).squeeze(1)\nimage_with_masks = [draw_segmentation_masks(image_tensor, masks=masks, alpha=0.7, colors=\"green\")]\nshow(image_with_masks)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-07T07:27:01.386657Z","iopub.execute_input":"2026-06-07T07:27:01.386971Z","iopub.status.idle":"2026-06-07T07:27:01.611007Z","shell.execute_reply.started":"2026-06-07T07:27:01.386948Z","shell.execute_reply":"2026-06-07T07:27:01.609785Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Create submission file ","metadata":{}},{"cell_type":"code","source":"# Write submission file\ndef convert_mask_to_polygon(mask):\n    # mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)\n    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)\n    coords = []        \n    for obj in contours:        \n        for point in obj:\n            coords.append((int(point[0][0]),int(point[0][1])))\n    return coords","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-07T07:27:24.255262Z","iopub.execute_input":"2026-06-07T07:27:24.255551Z","iopub.status.idle":"2026-06-07T07:27:24.260479Z","shell.execute_reply.started":"2026-06-07T07:27:24.255529Z","shell.execute_reply":"2026-06-07T07:27:24.259395Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Predict masks for list of images and write result in the submission file\ndef create_polygon_from_contours(mask):    \n    threshold = 0.5\n    mask = mask.squeeze(0)\n    mask[mask >= threshold] = 1\n    mask[mask < threshold] = 0\n    mask = mask.to(dtype=torch.uint8)\n    numpy_mask = mask.detach().cpu().numpy()    \n    return convert_mask_to_polygon(numpy_mask)\n\ndef generate_submission(image_list):\n    df = pd.DataFrame(image_list)\n    df.rename(columns={\"imageId\": \"ImageID\", \"objectList\": \"Coordinates\"}, inplace=True)                                                            \n    df.to_csv(\"submission.csv\", index=False, header=True)\n\ntest_data_path = '/kaggle/input/building-extraction-generalization-2024/test'\nimage_file_list = list(sorted(os.listdir(os.path.join(test_data_path, \"image\"))))\nprint(len(image_file_list))\n\nimage_polygon_list = []\nscore_threshold = 0.3\nfor image_id, image_file in enumerate(image_file_list):    \n    pil_image = Image.open(os.path.join(test_data_path, \"image\", image_file))\n    image_tensor = F.to_tensor(pil_image)     \n    x = [image_tensor.to(device)]     \n    predictions = model(x)\n    polygon_list = []\n    for index, mask_pred in enumerate(predictions[0][\"masks\"]):      \n        score = predictions[0][\"scores\"][index].item()   \n        if(score > score_threshold):\n            polygon = create_polygon_from_contours(mask_pred.clone())\n            if (len(polygon) > 3): \n                polygon_list.append(polygon)\n    image_polygon_list.append({'imageId': image_id, 'objectList': str(polygon_list)})\n    print(image_id, len(polygon_list))\n    # if (image_id == 10):\n    #     break\n        \ngenerate_submission(image_polygon_list)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-07T07:27:41.531156Z","iopub.execute_input":"2026-06-07T07:27:41.531506Z","iopub.status.idle":"2026-06-07T07:38:49.161058Z","shell.execute_reply.started":"2026-06-07T07:27:41.531477Z","shell.execute_reply":"2026-06-07T07:38:49.160199Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"!rm model_weights.pth\n# !mv model_weights.pth resnet50-28-28-augmented-data.pth\n# !echo '{\"username\":\"mojtabajahan\",\"key\":\"90d605700d0a1e525e887afc224e4755\"}' > ~/.kaggle/kaggle.json\n# !chmod 600 ~/.kaggle/kaggle.json\n# !kaggle competitions submit -c building-extraction-generalization-2024 -f submission.csv -m \"My notebook submission\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-09T19:00:37.568148Z","iopub.execute_input":"2026-06-09T19:00:37.568482Z","iopub.status.idle":"2026-06-09T19:00:37.884771Z","shell.execute_reply.started":"2026-06-09T19:00:37.568455Z","shell.execute_reply":"2026-06-09T19:00:37.883504Z"}},"outputs":[],"execution_count":null}]}