1 /** 2 * Copyright 2019-2021 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 #ifndef MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_DATASETOPS_SOURCE_IMAGE_FOLDER_OP_H_ 17 #define MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_DATASETOPS_SOURCE_IMAGE_FOLDER_OP_H_ 18 19 #include <deque> 20 #include <memory> 21 #include <queue> 22 #include <string> 23 #include <algorithm> 24 #include <map> 25 #include <set> 26 #include <utility> 27 #include <vector> 28 #include "minddata/dataset/core/tensor.h" 29 30 #include "minddata/dataset/engine/data_schema.h" 31 #include "minddata/dataset/engine/datasetops/parallel_op.h" 32 #include "minddata/dataset/engine/datasetops/source/mappable_leaf_op.h" 33 #include "minddata/dataset/engine/datasetops/source/sampler/sampler.h" 34 #ifndef ENABLE_ANDROID 35 #include "minddata/dataset/kernels/image/image_utils.h" 36 #else 37 #include "minddata/dataset/kernels/image/lite_image_utils.h" 38 #endif 39 #include "minddata/dataset/util/path.h" 40 #include "minddata/dataset/util/queue.h" 41 #include "minddata/dataset/util/services.h" 42 #include "minddata/dataset/util/status.h" 43 #include "minddata/dataset/util/wait_post.h" 44 45 namespace mindspore { 46 namespace dataset { 47 /// Forward declares 48 template <typename T> 49 class Queue; 50 51 using ImageLabelPair = std::shared_ptr<std::pair<std::string, int32_t>>; 52 using FolderImagesPair = std::shared_ptr<std::pair<std::string, std::queue<ImageLabelPair>>>; 53 54 class ImageFolderOp : public MappableLeafOp { 55 public: 56 // Constructor 57 // @param int32_t num_wkrs - Num of workers reading images in parallel 58 // @param std::string - dir directory of ImageNetFolder 59 // @param int32_t queue_size - connector queue size 60 // @param bool recursive - read recursively 61 // @param bool do_decode - decode the images after reading 62 // @param std::set<std::string> &exts - set of file extensions to read, if empty, read everything under the dir 63 // @param std::map<std::string, int32_t> &map- map of folder name and class id 64 // @param std::unique_ptr<dataschema> data_schema - schema of data 65 ImageFolderOp(int32_t num_wkrs, std::string file_dir, int32_t queue_size, bool recursive, bool do_decode, 66 const std::set<std::string> &exts, const std::map<std::string, int32_t> &map, 67 std::unique_ptr<DataSchema> data_schema, std::shared_ptr<SamplerRT> sampler); 68 69 /// Destructor. 70 ~ImageFolderOp() = default; 71 72 /// Initialize ImageFOlderOp related var, calls the function to walk all files 73 /// @param - std::string dir file directory to ImageNetFolder 74 /// @return Status The status code returned 75 Status PrescanMasterEntry(const std::string &dir); 76 77 // Worker thread pulls a number of IOBlock from IOBlock Queue, make a TensorRow and push it to Connector 78 // @param int32_t workerId - id of each worker 79 // @return Status The status code returned 80 Status PrescanWorkerEntry(int32_t worker_id); 81 82 // Method derived from RandomAccess Op, enable Sampler to get all ids for each class 83 // @param (std::map<int64_t, std::vector<int64_t >> * map - key label, val all ids for this class 84 // @return Status The status code returned 85 Status GetClassIds(std::map<int32_t, std::vector<int64_t>> *cls_ids) const override; 86 87 /// A print method typically used for debugging 88 /// @param out 89 /// @param show_all 90 void Print(std::ostream &out, bool show_all) const override; 91 92 /// This function is a hack! It is to return the num_class and num_rows. The result 93 /// returned by this function may not be consistent with what image_folder_op is going to return 94 /// user this at your own risk! 95 static Status CountRowsAndClasses(const std::string &path, const std::set<std::string> &exts, int64_t *num_rows, 96 int64_t *num_classes, std::map<std::string, int32_t> class_index); 97 98 /// Op name getter 99 /// @return Name of the current Op Name()100 std::string Name() const override { return "ImageFolderOp"; } 101 102 // DatasetName name getter 103 // \return DatasetName of the current Op 104 virtual std::string DatasetName(bool upper = false) const { return upper ? "ImageFolder" : "image folder"; } 105 106 //// \brief Base-class override for GetNumClasses 107 //// \param[out] num_classes the number of classes 108 //// \return Status of the function 109 Status GetNumClasses(int64_t *num_classes) override; 110 111 private: 112 // Load a tensor row according to a pair 113 // @param row_id_type row_id - id for this tensor row 114 // @param ImageLabelPair pair - <imagefile,label> 115 // @param TensorRow row - image & label read into this tensor row 116 // @return Status The status code returned 117 Status LoadTensorRow(row_id_type row_id, TensorRow *row) override; 118 119 /// @param std::string & dir - dir to walk all images 120 /// @param int64_t * cnt - number of non folder files under the current dir 121 /// @return 122 Status RecursiveWalkFolder(Path *dir); 123 124 /// start walking of all dirs 125 /// @return 126 Status StartAsyncWalk(); 127 128 // Called first when function is called 129 // @return 130 Status LaunchThreadsAndInitOp() override; 131 132 /// Private function for computing the assignment of the column name map. 133 /// @return - Status 134 Status ComputeColMap() override; 135 136 std::string folder_path_; // directory of image folder 137 bool recursive_; 138 bool decode_; 139 std::set<std::string> extensions_; // extensions allowed 140 std::map<std::string, int32_t> class_index_; 141 std::unique_ptr<DataSchema> data_schema_; 142 int64_t sampler_ind_; 143 uint64_t dirname_offset_; 144 std::vector<ImageLabelPair> image_label_pairs_; 145 std::unique_ptr<Queue<std::string>> folder_name_queue_; 146 std::unique_ptr<Queue<FolderImagesPair>> image_name_queue_; 147 }; 148 } // namespace dataset 149 } // namespace mindspore 150 #endif // MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_DATASETOPS_SOURCE_IMAGE_FOLDER_OP_H_ 151