Add polynomial function to dummy_segmentation

spinalcordtoolbox/testing/create_test_data.py

@@ -2,14 +2,18 @@
 # Collection of functions to create data for testing
 
 import numpy as np
+import numpy.matlib
 from datetime import datetime
 import itertools
 from skimage.transform import rotate
 
+from random import uniform
+
 import nibabel as nib
 
 from spinalcordtoolbox.image import Image
 from spinalcordtoolbox.resampling import resample_nib
+from sct_image import concat_data
 
 DEBUG = False  # Save img_sub
 
@@ -61,12 +65,11 @@ def dummy_centerline(size_arr=(9, 9, 9), pixdim=(1, 1, 1), subsampling=1, dilate
     :param debug: Bool: Write temp files
     :return:
     """
-    from numpy import poly1d, polyfit
     nx, ny, nz = size_arr
     # define array based on a polynomial function, within X-Z plane, located at y=ny/4, based on the following points:
     x = np.array([round(nx/4.), round(nx/2.), round(3*nx/4.)])
     z = np.array([0, round(nz/2.), nz-1])
-    p = poly1d(polyfit(z, x, deg=3))
+    p = np.poly1d(np.polyfit(z, x, deg=3))
     data = np.zeros((nx, ny, nz))
     arr_ctl = np.array([p(range(nz)).astype(np.int),
                         [round(ny / 4.)] * len(range(nz)),
@@ -109,7 +112,7 @@ def dummy_centerline(size_arr=(9, 9, 9), pixdim=(1, 1, 1), subsampling=1, dilate
 
 def dummy_segmentation(size_arr=(256, 256, 256), pixdim=(1, 1, 1), dtype=np.float64, orientation='LPI',
                        shape='rectangle', angle_RL=0, angle_AP=0, angle_IS=0, radius_RL=5.0, radius_AP=3.0,
-                       zeroslice=[], debug=False):
+                       interleaved=False, factor=1, zeroslice=[], debug=False):
     """Create a dummy Image with a ellipse or ones running from top to bottom in the 3rd dimension, and rotate the image
     to make sure that compute_csa and compute_shape properly estimate the centerline angle.
     :param size_arr: tuple: (nx, ny, nz)
@@ -122,6 +125,7 @@ def dummy_segmentation(size_arr=(256, 256, 256), pixdim=(1, 1, 1), dtype=np.floa
     :param angle_IS: int: angle around IS axis (in deg)
     :param radius_RL: float: 1st radius. With a, b = 50.0, 30.0 (in mm), theoretical CSA of ellipse is 4712.4
     :param radius_AP: float: 2nd radius
+    :param interleaved: bool: use polynomial function to simulate slicewise motion
     :param zeroslice: list int: zero all slices listed in this param
     :param debug: Write temp files for debug
     :return: img: Image object
@@ -132,12 +136,27 @@ def dummy_segmentation(size_arr=(256, 256, 256), pixdim=(1, 1, 1), dtype=np.floa
     nx, ny, nz = [int(size_arr[i] * pixdim[i]) for i in range(3)]
     data = np.random.random((nx, ny, nz)) * 0.
     xx, yy = np.mgrid[:nx, :ny]
-    # loop across slices and add object
-    for iz in range(nz):
-        if shape == 'rectangle':  # theoretical CSA: (a*2+1)(b*2+1)
-            data[:, :, iz] = ((abs(xx - nx / 2) <= radius_RL) & (abs(yy - ny / 2) <= radius_AP)) * 1
-        if shape == 'ellipse':
-            data[:, :, iz] = (((xx - nx / 2) / radius_RL) ** 2 + ((yy - ny / 2) / radius_AP) ** 2 <= 1) * 1
+    if not interleaved:
+        # loop across slices and add object
+        for iz in range(nz):
+            if shape == 'rectangle':  # theoretical CSA: (a*2+1)(b*2+1)
+                data[:, :, iz] = ((abs(xx - nx / 2) <= radius_RL) & (abs(yy - ny / 2) <= radius_AP)) * 1
+            if shape == 'ellipse':
+                data[:, :, iz] = (((xx - nx / 2) / radius_RL) ** 2 + ((yy - ny / 2) / radius_AP) ** 2 <= 1) * 1
+    elif interleaved:
+        # define array based on a polynomial function, within Y-Z plane to simulate slicewise motion in A-P
+        y = np.matlib.repmat([round(nx / 2.) + pixdim[0]*factor, round(nx / 2.) - pixdim[0]*factor], 1, round(nz / 2))
+        if nz % 2 != 0:         # if z-dimension is odd, add one more element to fit size
+            y = numpy.append(y,round(nx / 2.) + pixdim[0]*factor)
+        y = y.reshape(nz)       # reshape to vector (1,R) -> (R,)
+        z = np.arange(0, nz)
+        p = np.poly1d(np.polyfit(z, y, deg=nz))
+        # loop across slices and add object
+        for iz in range(nz):
+            if shape == 'rectangle':  # theoretical CSA: (a*2+1)(b*2+1)
+                data[:, :, iz] = ((abs(xx - nx / 2) <= radius_RL) & (abs(yy - p(iz)) <= radius_AP)) * 1
+            if shape == 'ellipse':
+                data[:, :, iz] = (((xx - nx / 2) / radius_RL) ** 2 + ((yy - p(iz)) / radius_AP) ** 2 <= 1) * 1
 
     # Pad to avoid edge effect during rotation
     data = np.pad(data, padding, 'reflect')
@@ -187,3 +206,52 @@ def dummy_segmentation(size_arr=(256, 256, 256), pixdim=(1, 1, 1), dtype=np.floa
     if debug:
         img.save('tmp_dummy_seg_'+datetime.now().strftime("%Y%m%d%H%M%S%f")+'.nii.gz')
     return img
+
+def dummy_segmentation_4d(vol_num=10, create_bvecs=False, size_arr=(256, 256, 256), pixdim=(1, 1, 1), dtype=np.float64,
+                          orientation='LPI', shape='rectangle', angle_RL=0, angle_AP=0, angle_IS=0, radius_RL=5.0,
+                          radius_AP=3.0, interleaved=False, zeroslice=[], debug=False):
+    """
+    Create a dummy 4D segmentation (dMRI/fMRI) and dummy bvecs file (optional)
+    :param vol_num: int: number of volumes in 4D data
+    :param create_bvecs: bool: create dummy bvecs file (necessary e.g. for sct_dmri_moco)
+    other parameters are same as in dummy_segmentation function
+    :return: Image object
+    """
+
+    img_list = []
+
+    # Loop across individual volumes of 4D data
+    for volume in range(0,vol_num):
+        factor = uniform(0.5, 3.0)          # shift in voxels
+        # set debug=True in line below for saving individual volumes into individual nii files
+        img_list.append(dummy_segmentation(size_arr=size_arr, pixdim=pixdim, dtype=dtype, orientation=orientation,
+                                           shape=shape, angle_RL=angle_RL, angle_AP=angle_AP, angle_IS=angle_IS,
+                                           radius_RL=radius_RL, radius_AP=radius_AP, interleaved=interleaved,
+                                           factor=factor, zeroslice=zeroslice, debug=False))
+
+    # Concatenate individual 3D images into 4D data
+    img_4d = concat_data(img_list, 3)
+    if debug:
+        out_name = datetime.now().strftime("%Y%m%d%H%M%S%f")
+        file_4d_data = 'tmp_dummy_4d_' + out_name + '.nii.gz'
+        img_4d.save(file_4d_data, verbose=0)
+
+    # Create a dummy bvecs file (necessary e.g. for sct_dmri_moco)
+    if create_bvecs:
+        n_b0 = 1                # number of b0
+        n_dwi = vol_num-n_b0    # number of dwi
+        bvecs_dummy = ['', '', '']
+        bvec_b0 = np.array([[0.0, 0.0, 0.0]] * n_b0)
+        bvec_dwi = np.array([[uniform(0,1), uniform(0,1), uniform(0,1)]] * n_dwi)
+        bvec = np.concatenate((bvec_b0,bvec_dwi),axis=0)
+        # Concatenate bvecs
+        for i in (0, 1, 2):
+            bvecs_dummy[i] += ' '.join(str(v) for v in map(lambda n: '%.16f' % n, bvec[:, i]))
+            bvecs_dummy[i] += ' '
+        bvecs_concat = '\n'.join(str(v) for v in bvecs_dummy)  # transform list into lines of strings
+        if debug:
+            new_f = open('tmp_dummy_4d_' + out_name + '.bvec', 'w')
+            new_f.write(bvecs_concat)
+            new_f.close()
+
+    return img_4d