#include<torch/torch.h>
#include<iostream>
#include<vector>
#include<cassert>
#include<stdlib.h>
#include<unordered_map>
#include<fstream>
class double_conv:public torch::nn::Module
{
public:
torch::nn::Conv2d conv1,conv2;
torch::nn::BatchNorm bn1,bn2;
int in_ch,out_ch;
public:
double_conv(int in_ch,int out_ch):in_ch(in_ch),out_ch(out_ch),conv1(torch::nn::Conv2dOptions(in_ch,out_ch,3).padding(1)),bn1(out_ch),
conv2(torch::nn::Conv2dOptions(out_ch,out_ch,3).padding(1)),bn2(out_ch)
{
register_module("conv1",conv1);
register_module("conv2",conv2);
register_module("bn1",bn1);
register_module("bn2",bn2);
}
torch::Tensor forward(torch::Tensor x)
{
x = conv1->forward(x);
x = bn1->forward(x);
x = torch::relu(x);
x = conv2->forward(x);
x = bn2->forward(x);
x = torch::relu(x);
return x;
}
};
class inconv:public torch::nn::Module
{
public:
int in_ch,out_ch;
public:
inconv(int in_ch,int out_ch):in_ch(in_ch),out_ch(out_ch){}
torch::Tensor forward(torch::Tensor x)
{
double_conv dc(in_ch,out_ch);
x = dc.forward(x);
return x;
}
};
class down:public torch::nn::Module
{
public:
int in_ch,out_ch;
public:
down(int in_ch,int out_ch):in_ch(in_ch),out_ch(out_ch){}
torch::Tensor forward(torch::Tensor x)
{
x = torch::max_pool2d(x,2);
double_conv dc(in_ch,out_ch);
x = dc.forward(x);
return x;
}
};
class up:public torch::nn::Module
{
public:
int in_ch,out_ch;
torch::nn::Conv2d upconv;
torch::nn::Conv2d conv1,conv2;
torch::nn::BatchNorm bn1,bn2;
torch::Tensor x;
public:
up(int in_ch,int out_ch):in_ch(in_ch),out_ch(out_ch),upconv(torch::nn::Conv2dOptions(in_ch,out_ch,4).padding(1).stride(2).transposed(new bool(true))),
conv1(torch::nn::Conv2dOptions(out_ch,out_ch,3).padding(1)),bn1(out_ch),conv2(torch::nn::Conv2dOptions(out_ch,out_ch,3).padding(1)),bn2(out_ch)
{
register_module("upconv",upconv);
register_module("conv1",conv2);
register_module("conv2",conv2);
register_module("bn1",bn1);
register_module("bn2",bn2);
}
torch::Tensor forward(torch::Tensor x1,torch::Tensor x2)
{
x = upconv->forward(x1);
x = torch::cat({x,x2},1);
double_conv dc(x.size(1),out_ch);
x = dc.forward(x);
//x = conv1->forward(x);
//x = bn1->forward(x);
//x = torch::relu(x);
//x = conv2->forward(x);
//x = bn2->forward(x);
//x = torch::relu(x);
return x;
}
};
class outconv:public torch::nn::Module
{
public:
int in_ch,out_ch;
torch::nn::Conv2d conv;
public:
outconv(int in_ch,int out_ch):in_ch(in_ch),out_ch(out_ch),conv(torch::nn::Conv2dOptions(in_ch,out_ch,1).padding(0))
{
register_module("conv",conv);
}
torch::Tensor forward(torch::Tensor x)
{
return conv->forward(x);
}
};
class unet:public torch::nn::Module
{
public:
int n_ch,n_class;
inconv *iconv= new inconv(n_ch,64);
down *down1= new down(64,256);
down *down2= new down(256,512);
down *down3= new down(512,512);
down *down4= new down(512,512);
up *up1= new up(512,256);
up *up2= new up(256,128);
up *up3= new up(128,64);
up *up4= new up(64,64);
outconv *oconv= new outconv(64,n_class);
torch::Tensor x1,x2,x3,x4,x5;
public:
unet(int n_ch,int n_class):n_ch(n_ch),n_class(n_class){}
torch::Tensor forward(torch::Tensor x)
{
x1 = iconv->forward(x);
x2 = down1->forward(x1);
x3 = down2->forward(x2);
x4 = down3->forward(x3);
x5 = down4->forward(x4);
x = up1->forward(x5,x4);
x = up2->forward(x,x3);
x = up3->forward(x,x2);
x = up4->forward(x,x1);
x = oconv->forward(x);
return x;
}
};
std::vector<float> Tokenize(const std::string& str,const std::string& delimiters)
{
std::vector<float> tokens;
std::string::size_type lastPos = str.find_first_not_of(delimiters, 0);
std::string::size_type pos = str.find_first_of(delimiters, lastPos);
while (std::string::npos != pos || std::string::npos != lastPos)
{
tokens.push_back(std::atof(str.substr(lastPos, pos - lastPos).c_str()));
lastPos = str.find_first_not_of(delimiters, pos);
pos = str.find_first_of(delimiters, lastPos);
}
return tokens;
}
std::vector<std::vector<float>> readTxt(std::string file)
{
std::ifstream infile;
infile.open(file.data());
assert(infile.is_open());
std::string s;
std::vector<float> vec;
std::vector<std::vector<float>> res;
while(getline(infile,s))
{
std::string tt= static_cast<std::string>(s);
vec = Tokenize(tt, " ");
res.push_back(vec);
}
infile.close();
std::cout<<"gdood"<<std::endl;
return res;
}
torch::Tensor float2TensorLabel()
{
static float tt[2478][3125]={0};
//memset(tt,0,sizeof(tt));
std::vector<std::vector<float>> vec = readTxt("/Users/yanlang/unet/mx-unet/U-Net/LabelData.txt");
int ch = vec.size();
int len = vec[0].size();
for(int i=0;i<ch;i++)
{
for(int j=0;j<len;j++)
{
tt[i][j]=vec[i][j];
}
}
torch::Tensor tmask = torch::CPU(torch::kFloat).tensorFromBlob(tt,{2478,3125});
return tmask;
}
torch::Tensor float2TensorData()
{
static float tt[7][2478*3125] = {0};
std::vector<std::vector<float>> vec = readTxt("/Users/yanlang/unet/mx-unet/U-Net/ImageData.txt");
int ch = vec.size();
int len = vec[0].size();
for(int i=0;i<ch;i++)
{
for(int j=0;j<len;j++)
{
tt[i][j]=vec[i][j];
}
}
torch::Tensor tdata = torch::CPU(torch::kFloat).tensorFromBlob(tt,{7,2478,3125});
return tdata;
}
int imgH=256;
int imgW=256;
torch::Tensor RandData(torch::Tensor data,int hight,int width)
{
//torch::Tensor datat = torch::squeeze(data);
torch::Tensor tmp = torch::zeros({7,imgH,imgW});
for(int i=hight;i<hight+imgH;i++)
{
for(int j=width;j<width+imgW;j++)
{
for(int k=0;k<7;k++)
{
tmp[k][i-hight][j-width]=data[k][i][j];
}
}
}
return tmp;
}
torch::Tensor RandMask(torch::Tensor label, int hight,int width)
{
torch::Tensor tmp = torch::zeros({imgH,imgW});
for(int i=hight;i<hight+imgH;i++)
{
for(int j=width;j<width+imgW;j++)
{
tmp[i-hight][j-width]=label[i][j];
}
}
return tmp;
}
std::vector<torch::Tensor> DataLoader(torch::Tensor data,torch::Tensor label,int batch_size)
{
int imghight = data.size(1);
int imgwidth = data.size(2);
int randhight,randwidth;
torch::Tensor resdata = torch::zeros({batch_size,7,imgH,imgW});
torch::Tensor reslabel = torch::zeros({batch_size,imgH,imgW});
for(int i=0;i<batch_size;i++)
{
randhight = rand()%(imghight-imgH-1);
randwidth = rand()%(imgwidth-imgW-1);
resdata[i] = RandData(data,randhight,randwidth);
reslabel[i] = RandMask(label,randhight,randwidth);
}
return {resdata,reslabel};
}
torch::autograd::Variable Get_predData(torch::autograd::Variable data)
{
//torch::autograd::Variable datat = torch::unsqueeze(data,0);
torch::autograd::Variable tmp = torch::zeros({7,imgH,imgW});
for(int i=500;i<756;i++)
{
for(int j=500;j<756;j++)
{
for(int k=0;k<7;k++)
{
tmp[k][i-500][j-500]=data[k][i][j];
}
}
}
return torch::unsqueeze(tmp,0);
}
void write2Txt(torch::autograd::Variable data)
{
std::ofstream fout("tresult.txt");
for(int i=0;i<data.size(0);i++)
{
for(int j=0;j<data.size(1);j++)
{
fout<<data[i][j]<<std::endl;
}
}
fout.close();
}
void saveModel(std::vector<torch::Tensor> weights,std::vector<std::string> key)
{
std::ofstream fout("unet.txt");
//std::unordered_map<std::string,torch::Tensor> mp;
for(int i=0;i<weights.size();i++)
{
fout<<key[i]<<std::endl;
fout<<weights[i]<<std::endl;
}
fout.close();
}
void trainConvNet(unet model)
{
torch::optim::SGD optimizer(model.parameters(),/*lr=*/0.01);
torch::Tensor pred;
std::cout<<"load data ......"<<std::endl;
torch::autograd::Variable data = torch::autograd::make_variable(float2TensorData());
torch::autograd::Variable label = torch::autograd::make_variable(float2TensorLabel());
std::cout<<"done!!"<<std::endl;
torch::Tensor train_data,train_label;
std::vector<torch::Tensor> vecdata;
for(int epoch=0;epoch<20;epoch++)
{
vecdata = DataLoader(data,label,2);
std::cout<<"vecdata after done!!"<<std::endl;
train_data = vecdata[0];
std::cout<<"train_data after done"<<std::endl;
train_label = vecdata[1];
std::cout<<train_label.size(0)<<std::endl;
std::cout<<"train_label after done"<<std::endl;
pred = model.forward(train_data);
auto loss = torch::nll_loss2d(pred,torch::_cast_Long(train_label));//torch::_cast_Long()
std::cout<<"the loss is"<<loss<<std::endl;
optimizer.zero_grad();
loss.backward();
optimizer.step();
}
std::vector<torch::Tensor> vecValue;
std::vector<std::string> vecKey;
torch::nn::ParameterCursor tt = model.parameters();
for(auto it=tt.begin();it!=tt.end();it++)
{
vecValue.push_back((*it).value);
vecKey.push_back((*it).key);
}
saveModel(vecValue,vecKey);
torch::autograd::Variable predData = Get_predData(data);
torch::autograd::Variable fl = model.forward(predData);
torch::autograd::Variable result = torch::squeeze(fl);
torch::autograd::Variable rt = result.argmax(0);
std::cout<<rt.size(0)<<std::endl;
std::cout<<rt.size(1)<<std::endl;
write2Txt(rt);
}
int main()
{
unet net(7,2);
trainConvNet(net);
return 0;
}
