Modified factorize to allow dimension selection and permutation

Warning: labels of B columns are copies of first column of M

factorize.m

-function [Ref,C,Value,U,S,V] = factorize(Ref,method,k)
+function [Ref,C,Value,U,S,V] = factorize(Ref,method,k,varargin)
 %COMPLETESET Summary of this function goes here
 %   Detailed explanation goes here
 
@@ -7,12 +7,32 @@ Size = cellfun(@size,Color,'UniformOutput',false);
 [ImgSize,ImgRes] = cellfun(@(x) deal(x(1:3),x(4:end)),Size,'UniformOutput',false);
 RefSize = size(Ref);
 numLab = numel(Label);
+numDims = ndims(Color{1});
+colDims = [];
+rowDims = [];
+
+p = inputParser; p.StructExpand = true; p.KeepUnmatched = true;
+p.addParameter('colDims', colDims, @isnumeric);
+p.addParameter('rowDims', rowDims, @isnumeric);
+
+p.parse(varargin{:});
+colDims = p.Results.colDims;
+rowDims = p.Results.rowDims;
+
+if (isempty(colDims)&&isempty(rowDims)); colDims = 1:min(3,numDims); end
+if isempty(colDims); colDims = setdiff(1:numDims,rowDims); end
+if isempty(rowDims); rowDims = setdiff(1:numDims,colDims); end
 
 for lab = 1:numLab
+    ImgSize{lab} = Size{lab}(colDims);
+    ImgRes {lab} = Size{lab}(rowDims);
+    
     Label{lab} = Label{lab}.*ones(Size{lab});
-    Color{lab} = reshape(Color{lab},prod(ImgSize{lab}),prod(ImgRes{lab}));
+    Label{lab} = permute(Label{lab},[colDims,rowDims]);
     Label{lab} = reshape(Label{lab},prod(ImgSize{lab}),prod(ImgRes{lab}));
-    Color{lab}(~(Label{lab}==lab)) = nan;
+    Color{lab} = permute(Color{lab},[colDims,rowDims]);
+    Color{lab} = reshape(Color{lab},prod(ImgSize{lab}),prod(ImgRes{lab}));
+    Color{lab}(Label{lab}==0) = nan;
 end
 
 switch method
@@ -40,7 +60,11 @@ end
 
 for lab = 1:numLab
     Color{lab} = reshape(Color{lab},[ImgSize{lab},ImgRes{lab}]);
-    Label{lab} = lab*double(any(~isnan(Color{lab}),3));
+    Label{lab} = repmat(Label{lab}(:,1),1,k);
+    ImgRes{lab} = [k,1];
+    Label{lab} = reshape(Label{lab},[ImgSize{lab},ImgRes{lab}]);
+    Color{lab} = ipermute(Color{lab},[colDims,rowDims]);
+    Label{lab} = ipermute(Label{lab},[colDims,rowDims]);
 end
 
 Color = reshape(Color,RefSize);