def __getitem__(self, index):
if self.year == 2014:
imgname = self.images[index]
else:
imgname = self.images[index].split('_')[2]
img_path = os.path.join(self.img_folder, imgname)
img = load_image(img_path)
ori_img = img.clone()
img[0].add_(-0.406)
img[1].add_(-0.457)
img[2].add_(-0.480)
imght = img.size(1)
imgwidth = img.size(2)
upLeft = torch.Tensor(
(float(self.annot['xmin'][index]), float(self.annot['ymin'][index])))
bottomRight = torch.Tensor(
(float(self.annot['xmax'][index]), float(self.annot['ymax'][index])))
ht = bottomRight[1] - upLeft[1]
width = bottomRight[0] - upLeft[0]
if width > 100:
scaleRate = 0.2
else:
scaleRate = 0.3
upLeft[0] = max(0, upLeft[0] - width * scaleRate / 2)
upLeft[1] = max(0, upLeft[1] - ht * scaleRate / 2)
bottomRight[0] = max(
min(imgwidth - 1, bottomRight[0] + width * scaleRate / 2), upLeft[0] + 5)
bottomRight[1] = max(
min(imght - 1, bottomRight[1] + ht * scaleRate / 2), upLeft[1] + 5)
inp = cropBox(img, upLeft, bottomRight, opt.inputResH, opt.inputResW)
ori_inp = cropBox(ori_img, upLeft, bottomRight,
opt.inputResH, opt.inputResW)
metaData = (
upLeft,
bottomRight,
ori_inp
)
box = torch.zeros(4)
box[0] = upLeft[0]
box[1] = upLeft[1]
box[2] = bottomRight[0]
box[3] = bottomRight[1]
return inp, box, imgname, metaData
def generateSampleBox(img_path,
bndbox,
part,
nJoints,
imgset,
scale_factor,
dataset,
train=True):
nJoints_coco = 17
nJoints_mpii = 16
img = load_image(img_path)
if train:
img[0].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
img[1].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
img[2].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
ori_img = img.clone()
img[0].add_(-0.406)
img[1].add_(-0.457)
img[2].add_(-0.480)
upLeft = torch.Tensor((int(bndbox[0][0]), int(bndbox[0][1])))
bottomRight = torch.Tensor((int(bndbox[0][2]), int(bndbox[0][3])))
ht = bottomRight[1] - upLeft[1]
width = bottomRight[0] - upLeft[0]
imght = img.shape[1]
imgwidth = img.shape[2]
scaleRate = random.uniform(*scale_factor)
upLeft[0] = max(0, upLeft[0] - width * scaleRate / 2)
upLeft[1] = max(0, upLeft[1] - ht * scaleRate / 2)
bottomRight[0] = min(imgwidth - 1, bottomRight[0] + width * scaleRate / 2)
bottomRight[1] = min(imght - 1, bottomRight[1] + ht * scaleRate / 2)
# Doing Random Sample
if opt.addDPG:
PatchScale = random.uniform(0, 1)
if PatchScale > 0.85:
ratio = ht / width
if (width < ht):
patchWidth = PatchScale * width
patchHt = patchWidth * ratio
else:
patchHt = PatchScale * ht
patchWidth = patchHt / ratio
xmin = upLeft[0] + random.uniform(0, 1) * (width - patchWidth)
ymin = upLeft[1] + random.uniform(0, 1) * (ht - patchHt)
xmax = xmin + patchWidth + 1
ymax = ymin + patchHt + 1
else:
xmin = max(
1,
min(upLeft[0] + np.random.normal(-0.0142, 0.1158) * width,
imgwidth - 3))
ymin = max(
1,
min(upLeft[1] + np.random.normal(0.0043, 0.068) * ht,
imght - 3))
xmax = min(
max(xmin + 2,
bottomRight[0] + np.random.normal(0.0154, 0.1337) * width),
imgwidth - 3)
ymax = min(
max(ymin + 2,
bottomRight[1] + np.random.normal(-0.0013, 0.0711) * ht),
imght - 3)
upLeft[0] = xmin
upLeft[1] = ymin
bottomRight[0] = xmax
bottomRight[1] = ymax
# Counting Joints number
jointNum = 0
if imgset == 'coco':
for i in range(17):
if part[i][0] > 0 and part[i][0] > upLeft[0] and part[i][1] > upLeft[1] \
and part[i][0] < bottomRight[0] and part[i][1] < bottomRight[1]:
jointNum += 1
else:
for i in range(16):
if part[i][0] > 0 and part[i][0] > upLeft[0] and part[i][1] > upLeft[1] \
and part[i][0] < bottomRight[0] and part[i][1] < bottomRight[1]:
jointNum += 1
# Doing Random Crop
if opt.addDPG:
if jointNum > 13 and train:
switch = random.uniform(0, 1)
if switch > 0.96:
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.92:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.88:
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.84:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.80:
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.76:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.72:
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.68:
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
ori_inp = cropBox(ori_img, upLeft, bottomRight, opt.inputResH,
opt.inputResW)
inp = cropBox(img, upLeft, bottomRight, opt.inputResH, opt.inputResW)
if jointNum == 0:
inp = torch.zeros(3, opt.inputResH, opt.inputResW)
out_bigcircle = torch.zeros(nJoints, opt.outputResH, opt.outputResW)
out_smallcircle = torch.zeros(nJoints, opt.outputResH, opt.outputResW)
out = torch.zeros(nJoints, opt.outputResH, opt.outputResW)
setMask = torch.zeros(nJoints, opt.outputResH, opt.outputResW)
# Draw Label
if imgset == 'coco':
for i in range(nJoints_coco):
if part[i][0] > 0 and part[i][0] > upLeft[0] and part[i][1] > upLeft[1] \
and part[i][0] < bottomRight[0] and part[i][1] < bottomRight[1]:
out_bigcircle[i] = drawBigCircle(
out_bigcircle[i],
transformBox(part[i], upLeft, bottomRight, opt.inputResH,
opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss * 2)
out_smallcircle[i] = drawSmallCircle(
out_smallcircle[i],
transformBox(part[i], upLeft, bottomRight, opt.inputResH,
opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
out[i] = drawGaussian(
out[i],
transformBox(part[i], upLeft, bottomRight, opt.inputResH,
opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
setMask[i].add_(1)
elif imgset == 'mpii':
for i in range(nJoints_coco, nJoints_coco + nJoints_mpii):
if part[i - nJoints_coco][0] > 0 and part[i - nJoints_coco][0] > upLeft[0] and part[i - nJoints_coco][1] > upLeft[1] \
and part[i - nJoints_coco][0] < bottomRight[0] and part[i - nJoints_coco][1] < bottomRight[1]:
out_bigcircle[i] = drawBigCircle(
out_bigcircle[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss * 2)
out_smallcircle[i] = drawSmallCircle(
out_smallcircle[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
out[i] = drawGaussian(
out[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
setMask[i].add_(1)
else:
for i in range(nJoints_coco, nJoints_coco + nJoints_mpii):
if part[i - nJoints_coco][0] > 0 and part[i - nJoints_coco][0] > upLeft[0] and part[i - nJoints_coco][1] > upLeft[1] \
and part[i - nJoints_coco][0] < bottomRight[0] and part[i - nJoints_coco][1] < bottomRight[1]:
out_bigcircle[i] = drawBigCircle(
out_bigcircle[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss * 2)
out_smallcircle[i] = drawSmallCircle(
out_smallcircle[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
out[i] = drawGaussian(
out[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
if i != 6 + nJoints_coco and i != 7 + nJoints_coco:
setMask[i].add_(1)
if opt.debug:
preds_hm, preds_img, preds_scores = getPrediction(
out.unsqueeze(0), upLeft.unsqueeze(0), bottomRight.unsqueeze(0),
opt.inputResH, opt.inputResW, opt.outputResH, opt.outputResW)
tmp_preds = preds_hm.mul(opt.inputResH / opt.outputResH)
drawCOCO(ori_inp.unsqueeze(0), tmp_preds, preds_scores)
if train:
# Flip
if random.uniform(0, 1) < 0.5:
inp = flip(inp)
ori_inp = flip(ori_inp)
out_bigcircle = shuffleLR(flip(out_bigcircle), dataset)
out_smallcircle = shuffleLR(flip(out_smallcircle), dataset)
out = shuffleLR(flip(out), dataset)
# Rotate
r = rnd(opt.rotate)
if random.uniform(0, 1) < 0.6:
r = 0
if r != 0:
inp = cv_rotate(inp, r, opt.inputResW, opt.inputResH)
out_bigcircle = cv_rotate(out_bigcircle, r, opt.outputResW,
opt.outputResH)
out_smallcircle = cv_rotate(out_smallcircle, r, opt.outputResW,
opt.outputResH)
out = cv_rotate(out, r, opt.outputResW, opt.outputResH)
return inp, out_bigcircle, out_smallcircle, out, setMask
def generateSampleBox(img_path,
bndbox,
part,
nJoints,
imgset,
scale_factor,
dataset,
train=True,
nJoints_coco=17):
img = load_image(img_path)
if train:
img[0].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
img[1].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
img[2].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
img[0].add_(-0.406)
img[1].add_(-0.457)
img[2].add_(-0.480)
upLeft = torch.Tensor((int(bndbox[0][0]), int(bndbox[0][1])))
bottomRight = torch.Tensor((int(bndbox[0][2]), int(bndbox[0][3])))
ht = bottomRight[1] - upLeft[1]
width = bottomRight[0] - upLeft[0]
imght = img.shape[1]
imgwidth = img.shape[2]
scaleRate = random.uniform(*scale_factor)
upLeft[0] = max(0, upLeft[0] - width * scaleRate / 2)
upLeft[1] = max(0, upLeft[1] - ht * scaleRate / 2)
bottomRight[0] = min(imgwidth - 1, bottomRight[0] + width * scaleRate / 2)
bottomRight[1] = min(imght - 1, bottomRight[1] + ht * scaleRate / 2)
# Doing Random Sample
if opt.addDPG:
PatchScale = random.uniform(0, 1)
if PatchScale > 0.85:
ratio = ht / width
if (width < ht):
patchWidth = PatchScale * width
patchHt = patchWidth * ratio
else:
patchHt = PatchScale * ht
patchWidth = patchHt / ratio
xmin = upLeft[0] + random.uniform(0, 1) * (width - patchWidth)
ymin = upLeft[1] + random.uniform(0, 1) * (ht - patchHt)
xmax = xmin + patchWidth + 1
ymax = ymin + patchHt + 1
else:
xmin = max(
1,
min(upLeft[0] + np.random.normal(-0.0142, 0.1158) * width,
imgwidth - 3))
ymin = max(
1,
min(upLeft[1] + np.random.normal(0.0043, 0.068) * ht,
imght - 3))
xmax = min(
max(xmin + 2,
bottomRight[0] + np.random.normal(0.0154, 0.1337) * width),
imgwidth - 3)
ymax = min(
max(ymin + 2,
bottomRight[1] + np.random.normal(-0.0013, 0.0711) * ht),
imght - 3)
upLeft[0] = xmin
upLeft[1] = ymin
bottomRight[0] = xmax
bottomRight[1] = ymax
# Counting Joints number
jointNum = 0
if imgset == 'coco':
for i in range(17):
if part[i][0] > 0 and part[i][0] > upLeft[0] and part[i][1] > upLeft[1] \
and part[i][0] < bottomRight[0] and part[i][1] < bottomRight[1]:
jointNum += 1
# Doing Random Crop
if opt.addDPG:
if jointNum > 13 and train:
switch = random.uniform(0, 1)
if switch > 0.96:
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.92:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.88:
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.84:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.80:
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.76:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.72:
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.68:
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
inputResH, inputResW = opt.inputResH, opt.inputResW
outputResH, outputResW = opt.outputResH, opt.outputResW
inp = cropBox(img, upLeft, bottomRight, inputResH, inputResW)
if jointNum == 0:
inp = torch.zeros(3, inputResH, inputResW)
out = torch.zeros(nJoints, outputResH, outputResW)
setMask = torch.zeros(nJoints, outputResH, outputResW)
# Draw Label
if imgset == 'coco':
for i in range(nJoints_coco):
if part[i][0] > 0 and part[i][0] > upLeft[0] and part[i][1] > upLeft[1] \
and part[i][0] < bottomRight[0] and part[i][1] < bottomRight[1]:
hm_part = transformBox(part[i], upLeft, bottomRight, inputResH,
inputResW, outputResH, outputResW)
out[i] = drawGaussian(out[i], hm_part, opt.hmGauss)
setMask[i].add_(1)
if train:
# Flip
if random.uniform(0, 1) < 0.5:
inp = flip(inp)
out = shuffleLR(flip(out), dataset)
# Rotate
r = rnd(opt.rotate)
if random.uniform(0, 1) < 0.6:
r = 0
if r != 0:
rotate = tr.Rotate(r)
inp = rotate(inp)
out = rotate(out)
return inp, out, setMask
def generateSampleBox(img_path,
bndbox,
part,
nJoints,
imgset,
scale_factor,
dataset,
train=True,
nJoints_coco=17):
img = load_image(img_path)
if train:
# 将input中的元素限制在[min,max]范围内并返回一个Tensor
img[0].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
img[1].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
img[2].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
img[0].add_(-0.406)
img[1].add_(-0.457)
img[2].add_(-0.480)
# 获取方框左上角和右下角道坐标
upLeft = torch.Tensor((int(bndbox[0][0]), int(bndbox[0][1])))
bottomRight = torch.Tensor((int(bndbox[0][2]), int(bndbox[0][3])))
# 方框道长和宽
ht = bottomRight[1] - upLeft[1]
width = bottomRight[0] - upLeft[0]
# 图片的长和宽
imght = img.shape[1]
imgwidth = img.shape[2]
# 缩放道比例
scaleRate = random.uniform(*scale_factor)
# 横列坐标
upLeft[0] = max(0, upLeft[0] - width * scaleRate / 2)
upLeft[1] = max(0, upLeft[1] - ht * scaleRate / 2)
bottomRight[0] = min(imgwidth - 1, bottomRight[0] + width * scaleRate / 2)
bottomRight[1] = min(imght - 1, bottomRight[1] + ht * scaleRate / 2)
# 做一些变换
if opt.addDPG:
PatchScale = random.uniform(0, 1)
if PatchScale > 0.85:
ratio = ht / width
if (width < ht):
patchWidth = PatchScale * width
patchHt = patchWidth * ratio
else:
patchHt = PatchScale * ht
patchWidth = patchHt / ratio
xmin = upLeft[0] + random.uniform(0, 1) * (width - patchWidth)
ymin = upLeft[1] + random.uniform(0, 1) * (ht - patchHt)
xmax = xmin + patchWidth + 1
ymax = ymin + patchHt + 1
else:
xmin = max(
1,
min(upLeft[0] + np.random.normal(-0.0142, 0.1158) * width,
imgwidth - 3))
ymin = max(
1,
min(upLeft[1] + np.random.normal(0.0043, 0.068) * ht,
imght - 3))
xmax = min(
max(xmin + 2,
bottomRight[0] + np.random.normal(0.0154, 0.1337) * width),
imgwidth - 3)
ymax = min(
max(ymin + 2,
bottomRight[1] + np.random.normal(-0.0013, 0.0711) * ht),
imght - 3)
upLeft[0] = xmin
upLeft[1] = ymin
bottomRight[0] = xmax
bottomRight[1] = ymax
# 计算关键点道个数
jointNum = 0
if imgset == 'coco':
for i in range(17):
# 判断合法性
if part[i][0] > 0 and part[i][0] > upLeft[0] and part[i][1] > upLeft[1] \
and part[i][0] < bottomRight[0] and part[i][1] < bottomRight[1]:
jointNum += 1
# 做随机的裁剪
if opt.addDPG:
if jointNum > 13 and train:
switch = random.uniform(0, 1)
if switch > 0.96:
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.92:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.88:
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.84:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.80:
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.76:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.72:
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.68:
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
inputResH, inputResW = opt.inputResH, opt.inputResW
outputResH, outputResW = opt.outputResH, opt.outputResW
# 裁剪, 将图片裁剪成要输入道图片大小
inp = cropBox(img, upLeft, bottomRight, inputResH, inputResW)
# 如果没有一个关键点，则将输入的图片变成0
if jointNum == 0:
inp = torch.zeros(3, inputResH, inputResW)
out = torch.zeros(nJoints, outputResH, outputResW)
setMask = torch.zeros(nJoints, outputResH, outputResW)
# 画标签的热图
if imgset == 'coco':
for i in range(nJoints_coco):
if part[i][0] > 0 and part[i][0] > upLeft[0] and part[i][1] > upLeft[1] \
and part[i][0] < bottomRight[0] and part[i][1] < bottomRight[1]:
# 将关键点的位置移到相应的位置上
hm_part = transformBox(part[i], upLeft, bottomRight, inputResW,
inputResH, outputResH, outputResW)
# 根据高斯分布画热图
out[i] = drawGaussian(out[i], hm_part, opt.hmGauss)
setMask[i].add_(1)
if train:
# 旋转
if random.uniform(0, 1) < 0.5:
inp = flip(inp)
out = shuffleLR(flip(out), dataset)
# Rotate
r = rnd(opt.rotate)
if random.uniform(0, 1) < 0.6:
inp = cv_rotate(inp, r, opt.inputResW, opt.inputResH)
out = cv_rotate(out, r, opt.outputResW, opt.outputResH)
return inp, out, setMask
