Method get

lesson6-Segmentation/SegDataset.cpp:113–135

}

torch::data::Example<> SegDataset::get(size_t index) {
    std::string image_path = list_images.at(index);
    std::string label_path = list_labels.at(index);
    cv::Mat image = cv::imread(image_path);
    cv::Mat mask = cv::Mat::zeros(image.rows, image.cols, CV_8UC3);
    draw_mask(label_path,mask);

    //Data augmentation like flip or rotate could be implemented here...
    cv::resize(image, image,cv::Size(resize_width,resize_height));
    cv::resize(mask,mask,cv::Size(resize_width,resize_height));
    torch::Tensor img_tensor = torch::from_blob(image.data, { image.rows, image.cols, 3 }, torch::kByte).permute({ 2, 0, 1 }); // Channels x Height x Width
    torch::Tensor colorful_label_tensor = torch::from_blob(mask.data, { mask.rows, mask.cols, 3 }, torch::kByte);
    torch::Tensor label_tensor = torch::zeros({image.rows, image.cols});

    //encode "colorful" tensor to class_index meaning tensor, [w,h,3]->[w,h], pixel value is the index of a class
    for(int i = 0; i<name_list.size(); i++){
        cv::Scalar color = name2color[name_list[i]];
        torch::Tensor color_tensor = torch::tensor({color.val[0],color.val[1],color.val[2]});
        label_tensor+=torch::all(colorful_label_tensor==color_tensor,-1)*i;
    }
    label_tensor = label_tensor.unsqueeze(0);
    return { img_tensor.clone(), label_tensor.clone() };
}