Translate Done

Author: qsyao

cuda_backend/kernel/spatial_deform.cu

@@ -88,3 +88,37 @@ __global__ void flip_3D(float* coords,
         }
     }
 }
+
+__global__ void translate_3D(float* coords,
+                        size_t dim_z,
+                        size_t dim_y, 
+                        size_t dim_x,
+                        float seg_z,
+                        float seg_y,
+                        float seg_x){
+    size_t index = blockIdx.x * blockDim.x + threadIdx.x;
+    size_t total = dim_x * dim_y * dim_z;
+    if(index < total){
+        coords[index] += seg_z;
+        coords[index + total] += seg_y;
+        coords[index + total * 2] += seg_x;
+        __syncthreads();
+    }
+}
+
+__global__ void translate_2D(float* coords, 
+                        size_t dim_y, 
+                        size_t dim_x,
+                        float seg_y,
+                        float seg_x){
+    size_t index = blockIdx.x * blockDim.x + threadIdx.x;
+    size_t total = dim_x * dim_y;
+    if(index < total){
+        coords[index] += seg_y;
+        coords[index + total] += seg_x;
+        __syncthreads();
+    }
+}
+
+
+

cuda_backend/kernel/spatial_deform.cuh

@@ -26,4 +26,19 @@ __global__ void flip_3D(float* coords,
                         int do_y,
                         int do_x);                        
 
+
+__global__ void translate_3D(float* coords,
+                        size_t dim_z,
+                        size_t dim_y, 
+                        size_t dim_x,
+                        float seg_z,
+                        float seg_y,
+                        float seg_x);
+
+__global__ void translate_2D(float* coords, 
+                        size_t dim_y, 
+                        size_t dim_x,
+                        float seg_y,
+                        float seg_x);                      
+
 #endif

cuda_backend/kernel/test.cu

@@ -38,6 +38,10 @@ void cu_flip(Handle* cuda_handle, int do_x, int do_y, int do_z){
     cuda_handle->flip(do_x, do_y, do_z);
 }
 
+void cu_translate(Handle* cuda_handle, float seg_x, float seg_y, float seg_z){
+    cuda_handle->translate(seg_x, seg_y, seg_z);
+}
+
 void endding_flag(Handle* cuda_handle){
     cuda_handle->recenter();
 }

cuda_backend/kernel/utils.cu

@@ -84,6 +84,22 @@ void Handle::flip(int do_x, int do_y, int do_z){
     }
 }
 
+void Handle::translate(float seg_x, float seg_y, float seg_z){
+    if(is_3D){
+        dim3 threads(min(total_size, (long)512), 1, 1);
+        dim3 blocks(total_size/512 + 1, 1, 1);
+        translate_3D<<<blocks, threads, 0, stream>>>(coords, dim_z, dim_y, dim_x,
+                                                seg_z, seg_y, seg_x);
+        checkCudaErrors(cudaStreamSynchronize(stream));
+    }
+    else{
+        dim3 threads(min(total_size, (long)512), 1, 1);
+        dim3 blocks(total_size/512 + 1, 1, 1);
+        translate_2D<<<blocks, threads, 0, stream>>>(coords, dim_y, dim_x, seg_y, seg_x);
+        checkCudaErrors(cudaStreamSynchronize(stream));
+    }    
+}
+
 void Handle::set_3D(size_t z, size_t y, size_t x){
     is_3D = true;
     dim_x = x;

cuda_backend/kernel/utils.cuh

@@ -40,6 +40,8 @@ public:
 
     void flip(int do_x, int do_y, int do_z=0);
 
+    void translate(float seg_x=0, float seg_y=0, float seg_z=0);
+
     ~Handle(){
         checkCudaErrors(cudaFree(img));
         checkCudaErrors(cudaFree(output));

cuda_backend/py_api.py

@@ -29,6 +29,9 @@
 flip = lib.cu_flip
 flip.argtypes = [c_void_p, c_int, c_int, c_int]
 
+translate = lib.cu_translate
+translate.argtypes = [c_void_p, c_float, c_float, c_float]
+
 class Spatial_Deform(object):
     def __init__(self, prob=1.0):
         self.prob = prob
@@ -71,6 +74,21 @@ def defrom(self, handle):
         else:
             return None
 
+class Translate(Spatial_Deform):
+    def __init__(self, seg_x=0.0, seg_y=0.0, seg_z=0.0, prob=1.0):
+        Spatial_Deform.__init__(self, prob)
+        self.label = 'Translate'
+        self.seg_x = seg_x
+        self.seg_y = seg_y
+        self.seg_z = seg_z
+    
+    def defrom(self, handle):
+        if np.random.uniform() < self.prob:
+            translate(handle, self.seg_x, self.seg_y, self.seg_z)
+            return self.label
+        else:
+            return None
+
 class End_Flag(Spatial_Deform):
     def __init__(self, prob=1.0):
         Spatial_Deform.__init__(self, prob)
@@ -106,8 +124,8 @@ def augment(self, img):
         output = np.ones(img.shape).astype(np.float32)
         labels = self.deform_coords()
 
-        # check coords
-        self.get_coords()
+        # # check coords
+        # self.get_coords()
 
         if not self.RGB:
             l_i(self.cuda_handle, output, img, 1)
@@ -126,6 +144,9 @@ def scale(self, sc, prob=1.0):
     def flip(self, do_x=False, do_y=False, do_z=False, prob=1.0):
         self.deform_list.append(Flip(do_x, do_y, do_z, prob))
 
+    def translate(self, seg_x=0.0, seg_y=0.0, seg_z=0.0, prob=1.0):
+        self.deform_list.append(Translate(seg_x, seg_y, seg_z, prob))
+
     def end_flag(self):
         self.deform_list.append(End_Flag())
 

python_augmentation.py

@@ -35,7 +35,8 @@ def test_3D():
 
     cuda_handle = Handle(array_image.shape)
     # cuda_handle.scale(0.5)
-    cuda_handle.flip(do_y=True, do_x=True, do_z=True)
+    # cuda_handle.flip(do_y=True, do_x=True, do_z=True)
+    cuda_handle.translate(100, 100, 20)
     cuda_handle.end_flag()
 
     correct_ret = deform.spatial_augment(array_image)
@@ -68,7 +69,8 @@ def test_2D():
 
     cuda_handle = Handle(array_image.shape, RGB=True)
     # cuda_handle.scale(0.5)
-    cuda_handle.flip(do_y=True)
+    # cuda_handle.flip(do_y=True)
+    cuda_handle.translate(400, 400)
     cuda_handle.end_flag()
 
     if len(array_image.shape) == 2:
@@ -115,5 +117,5 @@ def test_2D():
                                     (end - start) * 1000 / Iters_CPU)) 
 
 if __name__ == "__main__":
-    test_3D()
+    # test_3D()
     test_2D()