Neural network with adaptive color channel

Question

If we have a grayscale network ("Input" -> {1, Automatic, Automatic}), it can be applied on both gray and color image by:

ColorCombine[net/@ColorSeparate[img]]

But I want to turn this into an end-to-end network and that is easy to deploy.

So it should be adapted to the number of color channels.

We can find some equivalent layers for normal functions.

• ColorSeparate: no such equivalent form
• Map : NetMapOperator, it's a 4-tensor
• ColorCombine : CatenateLayer[1]

So how do I build an external network to adapt to this?

Answer 1

ClearAll[net];
net[n_Integer] :=
 NetChain[
   {
    (* color separate *)
    ReplicateLayer[1],
    TransposeLayer[],

    (* map *)
    NetMapOperator[{ConvolutionLayer[32, {3, 3}]}],

    (* color combine *)
    FlattenLayer[1, "Input" -> {n, Automatic, Automatic, Automatic}]
    },
   "Input" -> {n, Automatic, Automatic}
   ] // NetInitialize

net[3]

Dimensions@net[3]@RandomReal[1, {3, 128, 128}]

{96, 126, 126}

Dimensions@net[1]@RandomReal[1, {1, 128, 128}]

{32, 126, 126}