def __getitem__(self, idx):
item = self.data[idx]
x = torch.from_numpy(item.x).cuda()[None]
y = torch.from_numpy(item.y).cuda().long()
orientation = torch.from_numpy(item.orientation).float()
if self.mode == DataModes.TRAINING: # if training do augmentation
new_orientation = (torch.rand(3) -
0.5) * 2 * self.cfg.augmentation_shift_range
new_orientation = F.normalize(new_orientation, dim=0)
q = orientation + new_orientation
q = F.normalize(q, dim=0)
theta_rotate = transforms.rot_matrix_from_quaternion(q[None])
shift = torch.tensor([
d / (D // 2) for d, D in zip(
2 * (torch.rand(3) - 0.5) *
self.cfg.augmentation_shift_range, y.shape)
])
theta_shift = transforms.shift(shift)
theta = theta_rotate @ theta_shift
x, y = transforms.transform(theta, x, y)
orientation = new_orientation
pose = torch.cat((orientation, shift)).cuda()
else:
pose = torch.zeros(6).cuda()
C, D, H, W = x.shape
center = (D // 2, H // 2, W // 2)
x = crop(x, (C, ) + self.cfg.patch_shape, (0, ) + center)
y = crop(y, self.cfg.patch_shape, center)
if self.cfg.model_name == 'panet':
y = [y, pose]
return x, y
from torch.utils.data import DataLoader
from moduls.modul import Net
from moduls.loss import Loss
from utils.log import Log
# ======================== 开始训练 ========================
if __name__ == "__main__":
print("当前工作路径为:", os.getcwd())
parser = argparse.ArgumentParser("Loss training with Pytorch")
parser.add_argument("--config", help="config file", required=True)
args = parser.parse_args()
assert os.path.exists(args.config), args.config
opt = SourceFileLoader('module.name', args.config).load_module().opt
# ======================== 数据读取 =========================
train_trans, test_trans = transform()
trainset = Rotate(txt_path=opt.read_data.train.file_path, transform=train_trans)
trainloader = DataLoader(trainset, batch_size=opt.read_data.train.batch_size, shuffle=opt.read_data.train.shuffle)
testset = MyDataset(txt_path=opt.read_data.test.file_path, transform=test_trans)
testloader = DataLoader(testset, batch_size=opt.read_data.test.batch_size, shuffle=False)
# ======================== 导入网络 ========================
net = Net(opt).to(opt.device)
if opt.train.is_net_load:
net.load_state_dict(torch.load(opt.train.net_path))
print("模型导入成功！")
criterion = Loss(opt).to(opt.device)
fc = torch.nn.Linear(128, 4, bias=False).to(opt.device)
criterion_rotate = torch.nn.CrossEntropyLoss().to(opt.device)
# ======================== 初始化优化器 =======================
def __getitem__(self, index):
sf = self.scale_factor
rf = self.rot_factor
if self.is_train:
a = self.anno[self.train_list[index]]
else:
a = self.anno[self.valid_list[index]]
img_path = os.path.join(self.img_folder, a['img_paths'])
pts = torch.Tensor(a['joint_self'])
# pts[:, 0:2] -= 1 # Convert pts to zero based
# c = torch.Tensor(a['objpos']) - 1
c = torch.Tensor(a['objpos'])
s = a['scale_provided']
# Adjust center/scale slightly to avoid cropping limbs
if c[0] != -1:
c[1] = c[1] + 15 * s
s = s * 1.25
# For single-person pose estimation with a centered/scaled figure
nparts = pts.size(0)
img = load_image(img_path) # CxHxW
r = 0
if self.is_train:
s = s * torch.randn(1).mul_(sf).add_(1).clamp(1 - sf, 1 + sf)[0]
r = torch.randn(1).mul_(rf).clamp(
-2 * rf, 2 * rf)[0] if random.random() <= 0.6 else 0
# Flip
if random.random() <= 0.5:
img = fliplr(img)
pts = shufflelr(pts, img.size(2), self.DATA_INFO.hflip_indices)
c[0] = img.size(2) - c[0]
# Color
img[0, :, :].mul_(random.uniform(0.8, 1.2)).clamp_(0, 1)
img[1, :, :].mul_(random.uniform(0.8, 1.2)).clamp_(0, 1)
img[2, :, :].mul_(random.uniform(0.8, 1.2)).clamp_(0, 1)
# Prepare image and groundtruth map
inp = crop(img, c, s, self.inp_res, rot=r)
inp = color_normalize(inp, self.DATA_INFO.rgb_mean,
self.DATA_INFO.rgb_stddev)
# Generate ground truth
tpts = pts.clone()
target = torch.zeros(nparts, *self.out_res)
target_weight = tpts[:, 2].clone().view(nparts, 1)
for i in range(nparts):
# if tpts[i, 2] > 0: # This is evil!!
if tpts[i, 1] > 0:
tpts[i, 0:2] = to_torch(
transform(tpts[i, 0:2] + 1, c, s, self.out_res, rot=r))
target[i], vis = draw_labelmap(target[i],
tpts[i] - 1,
self.sigma,
type=self.label_type)
target_weight[i, 0] *= vis
# Meta info
if not isinstance(s, torch.Tensor):
s = torch.Tensor(s)
meta = {
'index': index,
'center': c,
'scale': s,
'pts': pts,
'tpts': tpts,
'target_weight': target_weight
}
return inp, target, meta