| 26 | namespace TEngine { |
| 27 | |
| 28 | bool Crop::InferShape(const std::vector<TEngine::TShape>& ishape, std::vector<TEngine::TShape>& oshape, int layout) |
| 29 | { |
| 30 | const TShape& input_shape = ishape[1]; |
| 31 | int input_h = input_shape.GetH(); |
| 32 | int input_w = input_shape.GetW(); |
| 33 | |
| 34 | int output_h = 0; |
| 35 | int output_w = 0; |
| 36 | // Mxnet |
| 37 | if(param_.flag == 1){ |
| 38 | if(param_.num_args == 2) |
| 39 | { |
| 40 | output_h = input_h; |
| 41 | output_w = input_w; |
| 42 | } |
| 43 | if(param_.num_args == 1) |
| 44 | { |
| 45 | output_h = param_.crop_h; |
| 46 | output_w = param_.crop_h; |
| 47 | } |
| 48 | } |
| 49 | //Caffe |
| 50 | if(param_.flag == 0){ |
| 51 | output_h = input_h; |
| 52 | output_w = input_w; |
| 53 | } |
| 54 | |
| 55 | TShape shape; |
| 56 | if(layout == TENGINE_LAYOUT_NCHW) |
| 57 | { |
| 58 | std::vector<int> dim = {input_shape.GetN(), input_shape.GetC(), output_h, output_w}; |
| 59 | |
| 60 | shape.SetDim(dim); |
| 61 | shape.SetDataLayout(TENGINE_LAYOUT_NCHW); |
| 62 | } |
| 63 | else |
| 64 | { |
| 65 | std::vector<int> dim = {input_shape.GetN(), output_h, output_w, input_shape.GetC()}; |
| 66 | |
| 67 | shape.SetDim(dim); |
| 68 | shape.SetDataLayout(TENGINE_LAYOUT_NHWC); |
| 69 | } |
| 70 | oshape[0] = shape; |
| 71 | |
| 72 | return true; |
| 73 | } |
| 74 | void Crop::SetSchema(void) |
| 75 | { |
| 76 | Input({"input:float"}) |