def predict(img):
global end_time
print('Predicting..')
w = 1920
h = 1080
x = transform(img)
x = x.unsqueeze(0)
x = Variable(x.cuda(0), volatile=True)
start_time = time.time()
loc_preds, cls_preds = net(x)
print("proid time", time.time() - end_time)
end_time = time.time()
time_cost = end_time - start_time
print("time_cost", time_cost)
print('Decoding..')
encoder = DataEncoder()
n1 = time.time()
# boxes, labels = encoder.decode(loc_preds.data.squeeze().cpu(), cls_preds.data.squeeze().cpu(), (w, h))
ref_table = encoder._get_anchor_boxes(torch.Tensor([1920, 1080]))
cls_preds = cls_preds.flatten()
idx = cls_preds.argmax()
print(idx)
print(cls_preds[idx])
boxes = [ref_table[idx]]
n2 = time.time()
print("useless", n2 - n1)
draw = ImageDraw.Draw(img)
for box in boxes:
draw.rectangle(list(box), outline='red')
return img
image = image.resize((1280, 960))
img = image.copy()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
print(image.size)
image = transform(image)
# forward
loc_preds = traced_script_module(image.unsqueeze(0).cuda())
loc_preds = loc_preds.argmax()
print(loc_preds)
encoder = DataEncoder()
ref_table = encoder._get_anchor_boxes(torch.Tensor([1280, 960]))
boxes = [ref_table[loc_preds]]
box = boxes[0]
print(boxes)
box[0] = (box[0] - box[2] / 2)
box[1] = (box[1] - box[3] / 2)
box[2] = (box[2] + box[0])
box[3] = (box[3] + box[1])
print(ref_table[215999])
draw = ImageDraw.Draw(img)
draw.rectangle(list(box), outline='red')
img.show()