import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

PATH = '/Users/huangyige/Downloads/fig/'def load_img(imgname):#here, only consider pion at 7.7 and 14.5 GeV, px.img = Image.open(imgname).convert('RGB')return imgdef process_img(img):img = img.crop((72,0,360,288))img = img.resize((32,32))return img

img = load_img(PATH+'Pion-7.7GeV-7-P1.png')
img = process_img(img)
plt.imshow(img)

def generate_file(name,num_range):with open('./'+name+'.txt','w') as f:for energy in ['7.7','14.5']:for _ in num_range:imgname = PATH + 'Pion-' + energy + 'GeV-' + str(_+1) + '-P1.png'f.write(imgname+' '+energy+'n')return
generate_file('train',range(0,70))
generate_file('test',range(70,90))

class sets(torch.utils.data.Dataset):def __init__(self,datatxt,transform=None):super(sets,self).__init__()imgs = []with open(datatxt,'r') as f:for line in f:line = line.rstrip('n')words = line.split(' ')imgs.append((words[0],words[1]))self.imgs = imgsself.transform = transformreturndef __getitem__(self,index):imgname,label_o = self.imgs[index]img = load_img(imgname)img = process_img(img)if label_o == '7.7':label = 0else:label = 1if self.transform is not None:img = self.transform(img)return img,labeldef __len__(self):return len(self.imgs)

train_set = sets('./train.txt',transforms.ToTensor())
test_set = sets('./test.txt',transforms.ToTensor())train_loader = torch.utils.data.DataLoader(dataset=train_set,batch_size=1,shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_set,batch_size=1)

class Net(nn.Module):def __init__(self):super(Net,self).__init__()self.conv1 = nn.Conv2d(3,6,5)self.pool = nn.MaxPool2d(2,2)self.conv2 = nn.Conv2d(6,16,5)self.fc1 = nn.Linear(16*5*5,120)self.fc2 = nn.Linear(120,84)#self.fc3 = nn.Linear(84,10)self.fc3 = nn.Linear(84,2)returndef forward(self,x):x = self.pool(F.relu(self.conv1(x)))x = self.pool(F.relu(self.conv2(x)))x = x.view(-1,16*5*5)x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = self.fc3(x)return xnet = Net()

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(),lr=0.001,momentum=0.9)

def accuracy(net,test_loder):correct = 0total = 0with torch.no_grad():for data in test_loader:inputs,labels = dataoutputs = net(inputs)_,pred = torch.max(outputs.data,1)total += labels.size(0)correct += (pred==labels).sum().item()acc = 100.0*correct/totalreturn acc

for epoch in range(2):running_loss = 0.0for i,data in enumerate(train_loader,0):inputs,labels = dataoptimizer.zero_grad()outputs = net(inputs)loss = criterion(outputs,labels)loss.backward()optimizer.step()running_loss += loss.item()acc = accuracy(net,test_loader)print('r[epoch %d >%3.d<] loss:%.3f,acc:%.1f%%'%(epoch+1,i+1,loss,acc),end='')print('')
print('Done!')

torch.save(net.state_dict(),'./model.pth')