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 !