def transform(img, boxes, labels):
img, boxes = resize(img, boxes, size=(img_size, img_size))
img = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])(img)
return img, boxes, labels
dataset = ListDataset(root='/home/ysdu/hardwareDisk/ysduDir/voc/voc_all_images', \
list_file='torchcv/datasets/voc/voc07_test.txt',
transform=transform)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=2)
box_coder = FPNSSDBoxCoder()
pred_boxes = []
pred_labels = []
pred_scores = []
gt_boxes = []
gt_labels = []
with open('torchcv/datasets/voc/voc07_test_difficult.txt') as f:
gt_difficults = []
for line in f.readlines():
line = line.strip().split()
d = [int(x) for x in line[1:]]
gt_difficults.append(d)
# size_ori = img.size
ow = oh = 512
factor_x = img_ori.size[0]/float(ow)
factor_y = img_ori.size[1]/float(oh)
img = img_ori.copy().resize((ow,oh))
print img_ori.size
print('Predicting..')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5]*3, [0.5]*3)])
x = transform(img)
loc_preds, cls_preds = net(x.unsqueeze(0))
print('Decoding..')
box_coder = FPNSSDBoxCoder()
loc_preds = loc_preds.squeeze().cpu()
cls_preds = F.softmax(cls_preds.squeeze(), dim=1).cpu()
boxes, labels, scores = box_coder.decode(loc_preds, cls_preds, score_thresh=0.2, nms_thresh=0.45)
print img_ori.size
draw = ImageDraw.Draw(img_ori)
# pdb.set_trace()
if not len(boxes) == 0:
boxes[:, 0] *= factor_x
boxes[:, 2] *= factor_x
boxes[:, 1] *= factor_y
boxes[:, 3] *= factor_y
for box in boxes:
draw.rectangle(list(box), outline='red')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
x = transform(img).cuda()
x2 = x.clone()
attention_image = np.zeros((ow, oh, 3), dtype=np.float32)
attention_image[:, :, 1] = 1.0
attention_image = torch.from_numpy(attention_image).permute(
2, 0, 1).unsqueeze(0).cuda()
loc_preds, cls_preds = net(x.unsqueeze(0), attention_image)
print('Decoding..')
box_coder = FPNSSDBoxCoder()
loc_preds = loc_preds.squeeze().cpu()
cls_preds = F.softmax(cls_preds.squeeze(), dim=1).cpu()
print(loc_preds.size())
boxes, labels, scores = box_coder.decode(loc_preds,
cls_preds,
score_thresh=0.6)
for box in boxes:
for j in range(len(box)):
if box[j] < 0:
box[j] = 0
print(boxes)
if len(boxes) == 0:
continue
start_epoch = checkpoint['epoch']
# net = torch.nn.DataParallel(net)
net.eval()
print('Loading image..')
img = Image.open(
'/home/ysdu/hardwareDisk/ysduDir/voc/VOCdevkit_test/VOC2012/JPEGImages/2008_002819.jpg'
)
ow = oh = 512
img = img.resize((ow, oh))
print('Predicting..')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
x = transform(img)
loc_preds, cls_preds = net(x.unsqueeze(0))
print('Decoding..')
box_coder = FPNSSDBoxCoder()
loc_preds = loc_preds.squeeze().cpu()
cls_preds = F.softmax(cls_preds.squeeze(), dim=1).cpu()
boxes, labels, scores = box_coder.decode(loc_preds, cls_preds)
print(labels)
draw = ImageDraw.Draw(img)
for box in boxes:
draw.rectangle(list(box), outline='red')
img.show()