How to train float32, 2-channel images with detectron2 ?

[Moron9645]

I have several medical images of type float32 (slices from nifty volumes), TIFF format. All of them are grayscale images.
First I tried to convert them into type uint8, PNG format, and train with mask r-cnn, and the result was not satisfying. In this case I want to try to train the float32 images immediately.

However, detectron2 raises error when I tried to train the float32 images.
I got

ValueError: axes don't match array

when the deault data mapper is proceeding this line of code, in detectron2/data/dataset_mapper.py, line 142 :

dataset_dict["image"] = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1)))

It seems that when detectron2 is training 3-channel images, it converts the images from (H,W,C) to (C,H,W) before it transforms them into tensors.

But what for 2-channel images ? I dug in a little and find out that when detectron2 reads 2-channel uint8 images using PIL, it gives the image another default channel, in detectron2/data/detections_utils.py, line 78 :

    # PIL squeezes out the channel dimension for "L", so make it HWC
    if format == "L":
        image = np.expand_dims(image, -1)

Then I took a shortcut by just editing detections_utils.py, line 78, to let detectron2 do the same thing when dealing with float32 images : (of course it's not good to do so...)

    # PIL squeezes out the channel dimension for "L", so make it HWC
    if format == "L" or format == "F":
        image = np.expand_dims(image, -1)

Now I could train the float32 images, but eventually I got almost the same AP score as the uint8 ones...
Could anyone tell me whether I am doing it in the right way ? If not, what and where should I implement ?

The code below is my Train.py :

import os
import json

from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, hooks, launch
from detectron2.evaluation import COCOEvaluator, verify_results
from detectron2.structures import BoxMode
from detectron2.utils import comm


class Trainer(DefaultTrainer):
    """
    We use the "DefaultTrainer" which contains pre-defined default logic for
    standard training workflow. They may not work for you, especially if you
    are working on a new research project. In that case you can write your
    own training loop. You can use "tools/plain_train_net.py" as an example.
    """

    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):

        return COCOEvaluator("liver_val", ("bbox", "segm"), False, output_dir=cfg.OUTPUT_DIR)

def get_liver_dicts(annoFolder):
    # Getting dataset_dicts, omitted
    return dataset_dicts

def main(args):
    for d in ["train", "val"]:
        DatasetCatalog.register("liver_" + d,
                                lambda d=d: get_liver_dicts("/home/xxx/liver/annotations_" + d + "/"))
        MetadataCatalog.get("liver_" + d).set(thing_classes=["tumor"])
        MetadataCatalog.get("liver_" + d).set(thing_colors=[(0, 255, 0)])

    cfg = get_cfg()
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.DATASETS.TRAIN = ("liver_train",)
    cfg.DATASETS.TEST = ("liver_val",)
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
    cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING = "choice"
    # All the images are 512*512
    cfg.INPUT.MAX_SIZE_TRAIN = 1024
    cfg.INPUT.MIN_SIZE_TRAIN = (512, 1024) 
    cfg.INPUT.MAX_SIZE_TEST = 1024
    cfg.INPUT.MIN_SIZE_TEST = 512
    cfg.INPUT.FORMAT = "F"
    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
    cfg.freeze()
    default_setup(cfg, args)

    if args.eval_only:
        model = Trainer.build_model(cfg)
        DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
            cfg.MODEL.WEIGHTS, resume=args.resume
        )
        res = Trainer.test(cfg, model)
        if comm.is_main_process():
            verify_results(cfg, res)
        return res

    trainer = Trainer(cfg)
    trainer.resume_or_load(resume=args.resume)
    return trainer.train()


if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    print("Command Line Args:", args)
    launch(
        main,
        args.num_gpus,
        num_machines=args.num_machines,
        machine_rank=args.machine_rank,
        dist_url=args.dist_url,
        args=(args,),
    )

And I launch training by this command :
python Train.py --config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml --num-gpus 4 SOLVER.IMS_PER_BATCH 8 SOLVER.BASE_LR 0.001 SOLVER.MAX_ITER 14130

Thanks in advance !

[MiXaiLL76]

update you augmentations in DatasetMapper or remove it

[MiXaiLL76]

import detectron2.data.transforms as T


image = np.ones((5,5,1))

aug_input = T.AugInput(image)
transforms = T.RandomFlip(1, horizontal= True, vertical=False)(aug_input)
# transforms = T.RandomRotation(
#                     np.linspace(0, 360, 20).astype(int).tolist(),
#                     sample_style='choice',
#                     expand=False,
#                 )(aug_input)

aug_input.image.shape```

> (5, 5, 1) its OK

[MiXaiLL76]

bat

import numpy as np
import detectron2.data.transforms as T


image = np.ones((5,5,1))

aug_input = T.AugInput(image)
transforms = T.RandomFlip(1, horizontal= True, vertical=False)(aug_input)
transforms = T.RandomRotation(
                    np.linspace(0, 360, 20).astype(int).tolist(),
                    sample_style='choice',
                    expand=False,
                )(aug_input)

aug_input.image.shape

(5, 5) its BAD

[MiXaiLL76]

Append to you code:

from fvcore.transforms.transform import NoOpTransform
import detectron2.data.transforms as T
import numpy as np

class AutoExpandDimsTransform(T.Transform):
    def apply_image(self, img):
        assert len(img.shape) <= 4
        
        if len(img.shape) == 2:
            return np.expand_dims(img, -1)
        else:
            return img
    
    def apply_coords(self, coords):
        return coords

    def inverse(self):
        return NoOpTransform()

    def apply_segmentation(self, segmentation):
        return segmentation

class AutoExpandDims(T.Augmentation):
    def __init__(self):
        super().__init__()
        self._init(locals())

    def get_transform(self, image):
        if len(image.shape) == 2:
            return AutoExpandDimsTransform()
        else:
            return NoOpTransform()