Method SegDataset

src/SegDataset.cpp:94–115

}

SegDataset::SegDataset(int resize_width, int resize_height, std::vector<std::string> list_images,
                       std::vector<std::string> list_labels, std::vector<std::string> name_list,
					   trainTricks tricks, bool isTrain)
{
	this->tricks = tricks;
	this->name_list = name_list;
	this->resize_width = resize_width;
	this->resize_height = resize_height;
	this->list_images = list_images;
	this->list_labels = list_labels;
	this->isTrain = isTrain;
    for(int i=0; i<name_list.size(); i++){
        name2index.insert(std::pair<std::string, int>(name_list[i], i));
    }
    std::vector<cv::Scalar> color_list = get_color_list();
    if(name_list.size()>color_list.size()){
        std::cout<< "Num of classes exceeds defined color list, please add color to color list in SegDataset.cpp";
    }
    for(int i = 0; i<name_list.size(); i++){
        name2color.insert(std::pair<std::string, cv::Scalar>(name_list[i],color_list[i]));
    }
}

torch::data::Example<> SegDataset::get(size_t index) {
    std::string image_path = list_images.at(index);