[完整案例]编程实现动物头像定位-02

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils import data
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
import torchvision
from torchvision import transforms
import os
from lxml import etree
from matplotlib.patches import Rectangle
import glob
from PIL import Image

img=pil_img.resize((224,224))
width=int(width)
height=int(height)

xmin = (xmin/width)*224
ymin = (ymin/height)*224
xmax = (xmax/width)*224
ymax = (ymax/height)*224

plt.imshow(img)
rect = Rectangle((xmin, ymin), (xmax-xmin), (ymax-ymin), fill=False, color='red')
ax = plt.gca()
ax.axes.add_patch(rect)
plt.show()

images=glob.glob(r'D:\社交应用\QQ\303\第12章\Oxford-IIIT Pets Dataset\dataset\images\*.jpg')
anno=glob.glob(r'D:\社交应用\QQ\303\第12章\Oxford-IIIT Pets Dataset\dataset\annotations\xmls\*.xml')

xml_name=[x.split("\\")[-1].replace('.xml','') for x in anno]
imgs=[x for x in images if x.split('\\')[-1].replace('.jpg','') in xml_name]

def to_labels(path):
    xml=open(r'{}'.format(path)).read()
    sel=etree.HTML(xml)
    width=int(sel.xpath('//size/width/text()')[0])
    height=int(sel.xpath('//size/height/text()')[0])
    xmin=int(sel.xpath('//bndbox/xmin/text()')[0])
    ymin=int(sel.xpath('//bndbox/ymin/text()')[0])
    xmax=int(sel.xpath('//bndbox/xmax/text()')[0])
    ymax=int(sel.xpath('//bndbox/ymax/text()')[0])
    return [xmin/width,ymin/height,xmax/width,ymax/height] #返回比例值

labels=[to_labels(p) for p in anno]

index=np.random.permutation(len(imgs))
imgs=np.array(imgs)[index]
labels=np.array(labels)[index]

labels=labels.astype(np.float32)

i=int(len(imgs)*0.8)
train_imgs=imgs[:i]
train_labels=labels[:i]
test_imgs=imgs[i:]
test_labels=labels[i:]

transform=transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
])

class OXford_dataset(data.Dataset):
    def __init__(self,img_paths,labels_list):
        self.imgs=img_paths
        self.labels=labels_list
        
    def __getitem__(self,index):
        img=self.imgs[index]
        pil_img=Image.open(img)
        img_tensor=transform(pil_img)
        l1,l2,l3,l4=self.labels[index]
        return img_tensor,l1,l2,l3,l4
    
    def __len__(self):
        return len(self.imgs)

train_dataset=OXford_dataset(train_imgs,train_labels)
test_dataset=OXford_dataset(test_imgs,test_labels)

train_dl=data.DataLoader(train_dataset,batch_size=16,shuffle=True)
test_dl=data.DataLoader(test_dataset,batch_size=16,shuffle=False)

plt.figure(figsize=(12,8))
for i,(img,l1,l2,l3,l4) in enumerate(zip(img_batch[:2],out1_b[:2],out2_b[:2],out3_b[:2],out4_b[:2])):
    img=img.permute(1,2,0).numpy()
    plt.subplot(1,2,i+1)
    plt.imshow(img)
    xmin,ymin,xmax,ymax=l1*224,l2*224,l3*224,l4*224
    rect=Rectangle((xmin,ymin),(xmax-xmin),(ymax-ymin),fill=False,color='red')
    ax=plt.gca()
    ax.axes.add_patch(rect)

resnet=torchvision.models.resnet101(pretrained=True)

in_size=resnet.fc.in_features

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.conv_base=nn.Sequential(*list(rensnet.children())[:-1])
        self.fc1=nn.Linear(in_size,1)
        self.fc2=nn.Linear(in_size,1)
        self.fc3=nn.Linear(in_size,1)
        self.fc4=nn.Linear(in_size,1)
    def foward(self,x):
        x=self.conv_base(x)
        x1=self.fc1(x)
        x2=self.fc2(x)
        x3=sell.fc3(x)
        x4=self.fc4(x)
        return x1,x2,x3,x4

model = Net()

if torch.cuda.is_available():
    model.to('cuda')

loss_fn = nn.MSELoss()

from torch.optim import lr_scheduler
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

def fit(epoch, model, trainloader, testloader):
    total = 0
    running_loss = 0
    
    model.train()
    for x, y1, y2, y3, y4 in trainloader:
        if torch.cuda.is_available():
            x, y1, y2, y3, y4 = (x.to('cuda'), 
                                 y1.to('cuda'), y2.to('cuda'),
                                 y3.to('cuda'), y4.to('cuda'))       
        y_pred1, y_pred2, y_pred3, y_pred4 = model(x)
        
        loss1 = loss_fn(y_pred1, y1)
        loss2 = loss_fn(y_pred2, y2)
        loss3 = loss_fn(y_pred3, y3)
        loss4 = loss_fn(y_pred4, y4)
        loss = loss1 + loss2 + loss3 + loss4
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        with torch.no_grad():
            running_loss += loss.item()
    exp_lr_scheduler.step()
    epoch_loss = running_loss / len(trainloader.dataset)
        
        
    test_total = 0
    test_running_loss = 0 
    
    model.eval()
    with torch.no_grad():
        for x, y1, y2, y3, y4 in testloader:
            if torch.cuda.is_available():
                x, y1, y2, y3, y4 = (x.to('cuda'), 
                                     y1.to('cuda'), y2.to('cuda'),
                                     y3.to('cuda'), y4.to('cuda'))
            y_pred1, y_pred2, y_pred3, y_pred4 = model(x)
            loss1 = loss_fn(y_pred1, y1)
            loss2 = loss_fn(y_pred2, y2)
            loss3 = loss_fn(y_pred3, y3)
            loss4 = loss_fn(y_pred4, y4)
            loss = loss1 + loss2 + loss3 + loss4
            test_running_loss += loss.item()
            
    epoch_test_loss = test_running_loss / len(testloader.dataset)
    
        
    print('epoch: ', epoch, 
          'loss： ', round(epoch_loss, 3),
          'test_loss： ', round(epoch_test_loss, 3),
             )
        
    return epoch_loss, epoch_test_loss

epochs = 10
train_loss = []
test_loss = []

for epoch in range(epochs):
    epoch_loss, epoch_test_loss = fit(epoch, model, train_dl, test_dl)
    train_loss.append(epoch_loss)
    test_loss.append(epoch_test_loss)

plt.figure()
plt.plot(range(1, len(train_loss)+1), train_loss, 'r', label='Training loss')
plt.plot(range(1, len(train_loss)+1), test_loss, 'bo', label='Validation loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss Value')
plt.legend()
plt.show()