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()