From 3220c74da529db5ec7f7b1f36b62c1ec1af4d27c Mon Sep 17 00:00:00 2001 From: Om Doiphode Date: Sat, 16 Nov 2024 15:18:15 +0530 Subject: [PATCH 1/2] Add MaskRCNN model file --- deepforest/models/MaskRCNN.py | 51 +++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 deepforest/models/MaskRCNN.py diff --git a/deepforest/models/MaskRCNN.py b/deepforest/models/MaskRCNN.py new file mode 100644 index 00000000..a1ac84ae --- /dev/null +++ b/deepforest/models/MaskRCNN.py @@ -0,0 +1,51 @@ +import torch +from torchvision.models.detection import maskrcnn_resnet50_fpn_v2 +from torchvision.models.detection.faster_rcnn import FastRCNNPredictor +from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor +from deepforest.model import Model + + +class Model(Model): + + def __init__(self, config, **kwargs): + super().__init__(config) + + def load_backbone(self): + backbone = maskrcnn_resnet50_fpn_v2(weights='DEFAULT') + + return backbone + + def create_model(self, backbone=None): + model = maskrcnn_resnet50_fpn_v2(weights='DEFAULT') + + # Modify the box predictor for the desired number of classes + in_features_box = model.roi_heads.box_predictor.cls_score.in_features + model.roi_heads.box_predictor = FastRCNNPredictor( + in_features_box, num_classes=self.config["num_classes"]) + + # Modify the mask predictor for the desired number of classes + in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels + dim_reduced = model.roi_heads.mask_predictor.conv5_mask.out_channels + model.roi_heads.mask_predictor = MaskRCNNPredictor( + in_features_mask, dim_reduced, num_classes=self.config["num_classes"]) + + return model + + def check_model(self): + """Ensure that model follows deepforest guidelines, see ##### If fails, + raise ValueError.""" + # This assumes model creation is not expensive + test_model = self.create_model() + test_model.eval() + + # Create a dummy batch of 3 band data. + x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + + predictions = test_model(x) + # Model takes in a batch of images + assert len(predictions) == 2 + + # Returns a list equal to number of images with proper keys per image + model_keys = list(predictions[1].keys()) + model_keys.sort() + assert model_keys == ['boxes', 'labels', 'scores', 'masks'] From 3a45f582e7b8e7cfddb2c25932f5ac668cb1854c Mon Sep 17 00:00:00 2001 From: Om Doiphode Date: Sat, 16 Nov 2024 16:14:50 +0530 Subject: [PATCH 2/2] Added polygon dataset class --- deepforest/dataset.py | 56 ++++++++++++++++++++++++++++++++++- deepforest/models/MaskRCNN.py | 2 +- 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/deepforest/dataset.py b/deepforest/dataset.py index d058fbd7..77204f53 100644 --- a/deepforest/dataset.py +++ b/deepforest/dataset.py @@ -21,11 +21,13 @@ from albumentations.pytorch import ToTensorV2 import torch import typing -from PIL import Image +from PIL import Image, ImageDraw import rasterio as rio from deepforest import preprocess from rasterio.windows import Window from torchvision import transforms +import torchvision +from torchvision.tv_tensors import BoundingBoxes, Mask def get_transform(augment): @@ -248,3 +250,55 @@ def __getitem__(self, idx): image = box return image + + +class PolygonDataset(Dataset): + + def __init__(self, img_keys, annotation_df, class_to_idx, transforms=None): + super(Dataset, self).__init__() + self._img_keys = img_keys + self._annotation_df = annotation_df + self._class_to_idx = class_to_idx + self._transforms = transforms + + def __len__(self): + return len(self._img_keys) + + def __getitem__(self, index): + img_key = self._img_keys[index] + annotation = self._annotation_df.loc[img_key] + image, target = self._load_image_and_target(annotation) + + if self._transforms: + image, target = self._transforms(image, target) + + return image, target + + def create_polygon_mask(self, image_size, vertices): + mask_img = Image.new('L', image_size, 0) + ImageDraw.Draw(mask_img, 'L').polygon(vertices, fill=(255)) + + return mask_img + + def _load_image_and_target(self, annotation): + filepath = annotation['image_path'] + image = Image.open(filepath).convert('RGB') + + labels = [shape['label'] for shape in annotation['shapes']] + labels = torch.Tensor([self._class_to_idx[label] for label in labels]) + labels = labels.to(dtype=torch.int64) + + shape_points = [shape['points'] for shape in annotation['shapes']] + xy_coords = [[tuple(p) for p in points] for points in shape_points] + mask_imgs = [self.create_polygon_mask(image.size, xy) for xy in xy_coords] + masks = Mask( + torch.concat([ + Mask(transforms.PILToTensor()(mask_img), dtype=torch.bool) + for mask_img in mask_imgs + ])) + + bboxes = BoundingBoxes(data=torchvision.ops.masks_to_boxes(masks), + format='xyxy', + canvas_size=image.size[::-1]) + + return image, {'masks': masks, 'boxes': bboxes, 'labels': labels} diff --git a/deepforest/models/MaskRCNN.py b/deepforest/models/MaskRCNN.py index a1ac84ae..12b5dbce 100644 --- a/deepforest/models/MaskRCNN.py +++ b/deepforest/models/MaskRCNN.py @@ -48,4 +48,4 @@ def check_model(self): # Returns a list equal to number of images with proper keys per image model_keys = list(predictions[1].keys()) model_keys.sort() - assert model_keys == ['boxes', 'labels', 'scores', 'masks'] + assert model_keys == ['boxes', 'masks', 'labels', 'scores']