在pytorch中如何实现图像定位

图像定位

import torch
import torch.nn as nn
import torch.optim
import torch.nn.functional as F
from torch.utils import data
import torchvision
import numpy as np
import matplotlib.pyplot as plt
from torchvision import transforms
import os 

from lxml import etree #对xml文件进行解码
from matplotlib.patches import Rectangle #绘制矩形框
import glob

from PIL import Image

images = glob.glob(r'dataset/images/*jpg')
anno = glob.glob(r'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]

从xml中获得矩形框四个顶点的坐标位置

def to_labels(path):
    xml = open(r'{}'.format(path)).read() #防止转译  将path通过{}和format的方法写入
    sel = etree.HTML(xml)#解码
    xmin = int(sel.xpath('//bndbox/xmin/text()')[0])
    xmax = int(sel.xpath('//bndbox/xmax/text()')[0])
    ymin = int(sel.xpath('//bndbox/ymin/text()')[0])
    ymax = int(sel.xpath('//bndbox/ymax/text()')[0])
    width = int(sel.xpath('//size/width/text()')[0])
    height = int(sel.xpath('//size/height/text()')[0])
    return [xmin/width, ymin/height, xmax/width, ymax/height]

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

乱序并区分train和test数据

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