From 9802287b5890d9b2cc250adba1b9810657a95c9c Mon Sep 17 00:00:00 2001
From: Joseph Redmon <pjreddie@gmail.com>
Date: Fri, 18 Dec 2015 23:55:58 +0000
Subject: [PATCH] some fixes

---
 src/crop_layer_kernels.cu |   19 ++++++++++++++-----
 1 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/crop_layer_kernels.cu b/src/crop_layer_kernels.cu
index fc7fcbd..8891030 100644
--- a/src/crop_layer_kernels.cu
+++ b/src/crop_layer_kernels.cu
@@ -1,3 +1,7 @@
+#include "cuda_runtime.h"
+#include "curand.h"
+#include "cublas_v2.h"
+
 extern "C" {
 #include "crop_layer.h"
 #include "utils.h"
@@ -91,7 +95,7 @@
     return val;
 }
 
-__global__ void levels_image_kernel(float *image, float *rand, int batch, int w, int h, int train, float saturation, float exposure, float translate, float scale)
+__global__ void levels_image_kernel(float *image, float *rand, int batch, int w, int h, int train, float saturation, float exposure, float translate, float scale, float shift)
 {
     int size = batch * w * h;
     int id = (blockIdx.x + blockIdx.y*gridDim.x) * blockDim.x + threadIdx.x;
@@ -100,6 +104,9 @@
     id /= w;
     int y = id % h;
     id /= h;
+    float rshift = rand[0];
+    float gshift = rand[1];
+    float bshift = rand[2];
     float r0 = rand[8*id + 0];
     float r1 = rand[8*id + 1];
     float r2 = rand[8*id + 2];
@@ -121,10 +128,12 @@
         hsv.y *= saturation;
         hsv.z *= exposure;
         rgb = hsv_to_rgb_kernel(hsv);
+    } else {
+        shift = 0;
     }
-    image[x + w*(y + h*0)] = rgb.x*scale + translate;
-    image[x + w*(y + h*1)] = rgb.y*scale + translate;
-    image[x + w*(y + h*2)] = rgb.z*scale + translate;
+    image[x + w*(y + h*0)] = rgb.x*scale + translate + (rshift - .5)*shift;
+    image[x + w*(y + h*1)] = rgb.y*scale + translate + (gshift - .5)*shift;
+    image[x + w*(y + h*2)] = rgb.z*scale + translate + (bshift - .5)*shift;
 }
 
 __global__ void forward_crop_layer_kernel(float *input, float *rand, int size, int c, int h, int w, int crop_height, int crop_width, int train, int flip, float angle, float *output)
@@ -186,7 +195,7 @@
 
     int size = layer.batch * layer.w * layer.h;
 
-    levels_image_kernel<<<cuda_gridsize(size), BLOCK>>>(state.input, layer.rand_gpu, layer.batch, layer.w, layer.h, state.train, layer.saturation, layer.exposure, translate, scale);
+    levels_image_kernel<<<cuda_gridsize(size), BLOCK>>>(state.input, layer.rand_gpu, layer.batch, layer.w, layer.h, state.train, layer.saturation, layer.exposure, translate, scale, layer.shift);
     check_error(cudaPeekAtLastError());
 
     size = layer.batch*layer.c*layer.crop_width*layer.crop_height;

--
Gitblit v1.10.0