| 27 | namespace singa { |
| 28 | |
| 29 | Tensor ImageTransformer::Apply(int flag, Tensor& input) { |
| 30 | CHECK_LE(input.nDim(), 4u); |
| 31 | CHECK_GE(input.nDim(), 2u); |
| 32 | CHECK_EQ(input.data_type(), kFloat32) << "Data type " << input.data_type() |
| 33 | << " is invalid for an raw image"; |
| 34 | srand((unsigned int)time(NULL)); |
| 35 | /// TODO |
| 36 | /// currently only consider one sample each time |
| 37 | |
| 38 | /// resize image using opencv resize |
| 39 | Tensor temp1; |
| 40 | #ifdef USE_OPENCV |
| 41 | temp1 = resize(input, resize_height_, resize_width_, image_dim_order_); |
| 42 | #else |
| 43 | temp1 = input; |
| 44 | #endif |
| 45 | |
| 46 | /// crop |
| 47 | Tensor temp2; |
| 48 | size_t height = 0, width = 0; |
| 49 | if (input.nDim() >= 3u) { |
| 50 | if (image_dim_order_ == "CHW") |
| 51 | height = temp1.shape(input.nDim() - 2), width = temp1.shape(input.nDim() - 1); |
| 52 | else if (image_dim_order_ == "HWC") |
| 53 | height = temp1.shape(input.nDim() - 3), width = temp1.shape(input.nDim() - 2); |
| 54 | else |
| 55 | LOG(FATAL) << "Unknow dimension order for images " << image_dim_order_ |
| 56 | << " Only support 'HWC' and 'CHW'"; |
| 57 | } else /// input is 2D gray image |
| 58 | height = temp1.shape(0), width = temp1.shape(1); |
| 59 | |
| 60 | if (crop_shape_.size() == 2) { |
| 61 | if (flag == kTrain) { |
| 62 | /// random crop |
| 63 | if (crop_shape_[0] > height || crop_shape_[0] > width) |
| 64 | LOG(FATAL) << "Crop size larger than the size of raw image"; |
| 65 | size_t crop_h_offset = rand() % ((height - crop_shape_[0]) / 2), |
| 66 | crop_w_offset = rand() % ((width - crop_shape_[1]) / 2); |
| 67 | temp2 = crop(temp1, crop_shape_[0], crop_shape_[1], |
| 68 | crop_h_offset, crop_w_offset, image_dim_order_); |
| 69 | } else if (flag == kEval) { |
| 70 | /// central crop |
| 71 | size_t crop_h_offset = (height - crop_shape_[0]) / 2, |
| 72 | crop_w_offset = (width - crop_shape_[1]) / 2; |
| 73 | temp2 = crop(temp1, crop_shape_[0], crop_shape_[1], |
| 74 | crop_h_offset, crop_w_offset, image_dim_order_); |
| 75 | } |
| 76 | } else temp2 = temp1; |
| 77 | |
| 78 | /// mirror |
| 79 | Tensor output; |
| 80 | if ((flag == kTrain) && (rand() % 2)) |
| 81 | output = mirror(temp2, true, false, image_dim_order_); |
| 82 | else output = temp2; |
| 83 | return output; |
| 84 | } |
| 85 | |
| 86 | #ifdef USE_OPENCV |