import copy
import detectron2.data as data
import detectron2.data.transforms as T
import numpy as np
import torch
from detectron2.data import detection_utils as utils
import deepdisc.astrodet.astrodet as toolkit
import deepdisc.astrodet.detectron as detectron_addons
[docs]
class DataMapper:
"""Base class that will map data to the format necessary for the model
To implement a data mapper for a new class, the derived class needs to have an
__init__() function that calls super().__init__(*args, **kwargs)
and a custom version of map_data().
"""
def __init__(self, imreader=None, key_mapper=None, augmentations=None):
"""
Parameters
----------
imreader : ImageReader
The class that will load and contrast scale the images.
They can be stored separately from the dataset or with it.
key_mapper : function
The function that takes the data set and returns the key that will be used to load the image.
If the image is stored with the dataset, this is not needed
Default = None
augmentations : detectron2 AugmentationList or a detectron_addons.KRandomAugmentationList
The list of augmentations to apply to the image
Default = None
"""
[docs]
self.augmentations = augmentations
[docs]
def map_data(self, data):
return data
[docs]
class DictMapper(DataMapper):
"""Class that will map COCO dictionary data to the format necessary for the model"""
def __init__(self, *args, **kwargs):
# Pass arguments to the parent function.
super().__init__(*args, **kwargs)
[docs]
def map_data(self, dataset_dict):
"""Map COCO dict data to the correct format
Parameters
----------
dataset_dict: dict
a dictionary of COCO formatted metadata
Returns
-------
reformatted dictionary including image and instances
"""
dataset_dict = copy.deepcopy(dataset_dict)
key = self.km(dataset_dict)
image = self.IR(key)
# Data Augmentation
auginput = T.AugInput(image)
# Transformations to model shapes
if self.augmentations is not None:
augs = self.augmentations(image)
else:
augs = T.AugmentationList([])
transform = augs(auginput)
image = torch.from_numpy(auginput.image.copy().transpose(2, 0, 1))
annos = [
utils.transform_instance_annotations(annotation, [transform], image.shape[1:])
for annotation in dataset_dict.pop("annotations")
]
instances = utils.annotations_to_instances(annos, image.shape[1:])
instances = utils.filter_empty_instances(instances)
return {
# create the format that the model expects
"image": image,
"image_shaped": auginput.image,
"height": image.shape[1],
"width": image.shape[2],
"image_id": dataset_dict["image_id"],
"instances": instances,
}
[docs]
class MagRedshiftDictMapper(DataMapper):
def __init__(self, *args, **kwargs):
# Pass arguments to the parent function.
super().__init__(*args, **kwargs)
[docs]
def map_data(self, dataset_dict):
"""Map COCO dict data to the correct format, add ground truth redhshift
Parameters
----------
dataset_dict: dict
a dictionary of COCO formatted metadata
Returns
-------
reformatted dictionary including image and instances+redshift
"""
dataset_dict = copy.deepcopy(dataset_dict)
key = self.km(dataset_dict)
image = self.IR(key)
# Data Augmentation
auginput = T.AugInput(image)
# Transformations to model shapes
if self.augmentations is not None:
augs = self.augmentations(image)
else:
augs = T.AugmentationList([])
transform = augs(auginput)
image = torch.from_numpy(auginput.image.copy().transpose(2, 0, 1))
annos = [
utils.transform_instance_annotations(annotation, [transform], image.shape[1:])
for annotation in dataset_dict.pop("annotations")
if annotation["redshift"] != 0.0]# and annotation["mag_i"] < 25.3]
instances = utils.annotations_to_instances(annos, image.shape[1:])
instances.gt_magi = torch.tensor([a["mag_i"] for a in annos])
instances.gt_redshift = torch.tensor([a["redshift"] for a in annos])
instances = utils.filter_empty_instances(instances)
return {
# create the format that the model expects
"image": image,
"image_shaped": auginput.image,
"height": image.shape[1],
"width": image.shape[2],
"image_id": dataset_dict["image_id"],
"instances": instances,
"annotations": annos
}
[docs]
class RedshiftDictMapper(DataMapper):
def __init__(self, *args, **kwargs):
# Pass arguments to the parent function.
super().__init__(*args, **kwargs)
[docs]
def map_data(self, dataset_dict):
"""Map COCO dict data to the correct format, add ground truth redhshift
Parameters
----------
dataset_dict: dict
a dictionary of COCO formatted metadata
Returns
-------
reformatted dictionary including image and instances+redshift
"""
dataset_dict = copy.deepcopy(dataset_dict)
key = self.km(dataset_dict)
image = self.IR(key)
# Data Augmentation
auginput = T.AugInput(image)
# Transformations to model shapes
if self.augmentations is not None:
augs = self.augmentations(image)
else:
augs = T.AugmentationList([])
transform = augs(auginput)
image = torch.from_numpy(auginput.image.copy().transpose(2, 0, 1))
annos = [
utils.transform_instance_annotations(annotation, [transform], image.shape[1:])
for annotation in dataset_dict.pop("annotations")
if annotation["redshift"] != 0.0
]
instances = utils.annotations_to_instances(annos, image.shape[1:])
instances.gt_redshift = torch.tensor([a["redshift"] for a in annos])
instances = utils.filter_empty_instances(instances)
return {
# create the format that the model expects
"image": image,
"image_shaped": auginput.image,
"height": image.shape[1],
"width": image.shape[2],
"image_id": dataset_dict["image_id"],
"instances": instances,
"annotations": annos
}
[docs]
class GoldRedshiftDictMapper(DataMapper):
def __init__(self, *args, **kwargs):
# Pass arguments to the parent function.
super().__init__(*args, **kwargs)
[docs]
def map_data(self, dataset_dict):
"""Map COCO dict data to the correct format, add ground truth redhshift
Parameters
----------
dataset_dict: dict
a dictionary of COCO formatted metadata
Returns
-------
reformatted dictionary including image and instances+redshift
"""
dataset_dict = copy.deepcopy(dataset_dict)
key = self.km(dataset_dict)
image = self.IR(key)
# Data Augmentation
auginput = T.AugInput(image)
# Transformations to model shapes
if self.augmentations is not None:
augs = self.augmentations(image)
else:
augs = T.AugmentationList([])
transform = augs(auginput)
image = torch.from_numpy(auginput.image.copy().transpose(2, 0, 1))
annos = [
utils.transform_instance_annotations(annotation, [transform], image.shape[1:])
for annotation in dataset_dict.pop("annotations")
if annotation["redshift"] != 0.0 and annotation["mag_i"] < 25.3
]
instances = utils.annotations_to_instances(annos, image.shape[1:])
instances.gt_redshift = torch.tensor([a["redshift"] for a in annos])
instances = utils.filter_empty_instances(instances)
return {
# create the format that the model expects
"image": image,
"image_shaped": auginput.image,
"height": image.shape[1],
"width": image.shape[2],
"image_id": dataset_dict["image_id"],
"instances": instances,
"annotations": annos
}
class RedshiftEBVDictMapper(DataMapper):
def __init__(self, *args, **kwargs):
# Pass arguments to the parent function.
super().__init__(*args, **kwargs)
def map_data(self, dataset_dict):
"""Map COCO dict data to the correct format, add ground truth redhshift
Parameters
----------
dataset_dict: dict
a dictionary of COCO formatted metadata
Returns
-------
reformatted dictionary including image and instances+redshift
"""
dataset_dict = copy.deepcopy(dataset_dict)
key = self.km(dataset_dict)
image = self.IR(key)
# Data Augmentation
auginput = T.AugInput(image)
# Transformations to model shapes
if self.augmentations is not None:
augs = self.augmentations(image)
else:
augs = T.AugmentationList([])
transform = augs(auginput)
image = torch.from_numpy(auginput.image.copy().transpose(2, 0, 1))
annos = [
utils.transform_instance_annotations(annotation, [transform], image.shape[1:])
for annotation in dataset_dict.pop("annotations")
if annotation["redshift"] != 0.0
]
instances = utils.annotations_to_instances(annos, image.shape[1:])
instances.gt_redshift = torch.tensor([a["redshift"] for a in annos])
instances.gt_ebv = torch.tensor([a["EBV"] for a in annos])
instances = utils.filter_empty_instances(instances)
return {
# create the format that the model expects
"image": image,
"image_shaped": auginput.image,
"height": image.shape[1],
"width": image.shape[2],
"image_id": dataset_dict["image_id"],
"instances": instances,
#"annotations": annos
}
[docs]
class GoldRedshiftDictMapperEval(DataMapper):
def __init__(self, *args, **kwargs):
# Pass arguments to the parent function.
super().__init__(*args, **kwargs)
[docs]
def map_data(self, dataset_dict):
"""Map COCO dict data to the correct format, add ground truth redhshift
Parameters
----------
dataset_dict: dict
a dictionary of COCO formatted metadata
Returns
-------
reformatted dictionary including image and instances+redshift
"""
dataset_dict = copy.deepcopy(dataset_dict)
key = self.km(dataset_dict)
image = self.IR(key)
# Data Augmentation
auginput = T.AugInput(image)
# Transformations to model shapes
if self.augmentations is not None:
augs = self.augmentations(image)
else:
augs = T.AugmentationList([])
transform = augs(auginput)
image = torch.from_numpy(auginput.image.copy().transpose(2, 0, 1))
annotations = [annotation for annotation in dataset_dict["annotations"]
if annotation["redshift"] != 0.0 and annotation["mag_i"] < 25.3]
#annos = [
# utils.transform_instance_annotations(annotation, [transform], image.shape[1:])
# for annotation in dataset_dict.pop("annotations")
# if annotation["redshift"] != 0.0
#]
#instances = utils.annotations_to_instances(annos, image.shape[1:])
#instances.gt_redshift = torch.tensor([a["redshift"] for a in annos])
#instances = utils.filter_empty_instances(instances)
return {
# create the format that the model expects
"image": image,
"image_shaped": auginput.image,
"height": image.shape[1],
"width": image.shape[2],
"image_id": dataset_dict["image_id"],
#"instances": instances,
"annotations": annotations
}
[docs]
class RedshiftDictMapperEval(DataMapper):
def __init__(self, *args, **kwargs):
# Pass arguments to the parent function.
super().__init__(*args, **kwargs)
[docs]
def map_data(self, dataset_dict):
"""Map COCO dict data to the correct format, add ground truth redhshift
Parameters
----------
dataset_dict: dict
a dictionary of COCO formatted metadata
Returns
-------
reformatted dictionary including image and instances+redshift
"""
dataset_dict = copy.deepcopy(dataset_dict)
key = self.km(dataset_dict)
image = self.IR(key)
# Data Augmentation
auginput = T.AugInput(image)
# Transformations to model shapes
if self.augmentations is not None:
augs = self.augmentations(image)
else:
augs = T.AugmentationList([])
transform = augs(auginput)
image = torch.from_numpy(auginput.image.copy().transpose(2, 0, 1))
annotations = [annotation for annotation in dataset_dict["annotations"]
if annotation["redshift"] != 0.0]
#annos = [
# utils.transform_instance_annotations(annotation, [transform], image.shape[1:])
# for annotation in dataset_dict.pop("annotations")
# if annotation["redshift"] != 0.0
#]
#instances = utils.annotations_to_instances(annos, image.shape[1:])
#instances.gt_redshift = torch.tensor([a["redshift"] for a in annos])
#instances = utils.filter_empty_instances(instances)
return {
# create the format that the model expects
"image": image,
"image_shaped": auginput.image,
"height": image.shape[1],
"width": image.shape[2],
"image_id": dataset_dict["image_id"],
#"instances": instances,
"annotations": annotations
}
[docs]
class WCSDictmapper(DataMapper):
def __init__(self, *args, **kwargs):
# Pass arguments to the parent function.
super().__init__(*args, **kwargs)
[docs]
def map_data(self, dataset_dict):
"""Map COCO dict data to the correct format, add ground truth redhshift
Parameters
----------
dataset_dict: dict
a dictionary of COCO formatted metadata
Returns
-------
reformatted dictionary including image and instances+redshift
"""
dataset_dict = copy.deepcopy(dataset_dict)
key = self.km(dataset_dict)
image = self.IR(key)
# Data Augmentation
auginput = T.AugInput(image)
# Transformations to model shapes
if self.augmentations is not None:
augs = self.augmentations(image)
else:
augs = T.AugmentationList([])
transform = augs(auginput)
image = torch.from_numpy(auginput.image.copy().transpose(2, 0, 1))
return {
# create the format that the model expects
"image": image,
"image_shaped": auginput.image,
"height": image.shape[1],
"width": image.shape[2],
#"image_id": dataset_dict["image_id"],
#"instances": instances,
"wcs": dataset_dict['wcs']
}
[docs]
class RedshiftEBVDictMapper(DataMapper):
def __init__(self, *args, **kwargs):
# Pass arguments to the parent function.
super().__init__(*args, **kwargs)
[docs]
def map_data(self, dataset_dict):
"""Map COCO dict data to the correct format, add ground truth redhshift
Parameters
----------
dataset_dict: dict
a dictionary of COCO formatted metadata
Returns
-------
reformatted dictionary including image and instances+redshift
"""
dataset_dict = copy.deepcopy(dataset_dict)
key = self.km(dataset_dict)
image = self.IR(key)
# Data Augmentation
auginput = T.AugInput(image)
# Transformations to model shapes
if self.augmentations is not None:
augs = self.augmentations(image)
else:
augs = T.AugmentationList([])
transform = augs(auginput)
image = torch.from_numpy(auginput.image.copy().transpose(2, 0, 1))
annos = [
utils.transform_instance_annotations(annotation, [transform], image.shape[1:])
for annotation in dataset_dict.pop("annotations")
if annotation["redshift"] != 0.0
]
instances = utils.annotations_to_instances(annos, image.shape[1:])
instances.gt_redshift = torch.tensor([a["redshift"] for a in annos])
instances.gt_ebv = torch.tensor([a["EBV"] for a in annos])
instances = utils.filter_empty_instances(instances)
return {
# create the format that the model expects
"image": image,
"image_shaped": auginput.image,
"height": image.shape[1],
"width": image.shape[2],
"image_id": dataset_dict["image_id"],
"instances": instances,
#"annotations": annos
}
[docs]
def return_train_loader(cfg, mapper):
"""Returns a train loader
Parameters
----------
cfg : LazyConfig
The lazy config, which contains data loader config values
**kwargs for the read_image functionality
Returns
-------
a train loader
"""
loader = data.build_detection_train_loader(cfg, mapper=mapper)
return loader
[docs]
def return_test_loader(cfg, mapper):
"""Returns a test loader
Parameters
----------
cfg : LazyConfig
The lazy config, which contains data loader config values
**kwargs for the read_image functionality
Returns
-------
a test loader
"""
loader = data.build_detection_test_loader(cfg, cfg.DATASETS.TEST, mapper=mapper)
return loader