| 80 | |
| 81 | |
| 82 | class Inference: |
| 83 | |
| 84 | def __init__(self, args, device): |
| 85 | cfg = get_cfg() |
| 86 | add_config(cfg) |
| 87 | cfg.merge_from_file(args.config_file) |
| 88 | cfg.freeze() |
| 89 | |
| 90 | ckpt = torch.load(cfg.MODEL.WEIGHTS, map_location=torch.device('cpu')) |
| 91 | self.model = Model(cfg) |
| 92 | out = self.model.load_state_dict(ckpt, strict=False) |
| 93 | print(out) |
| 94 | |
| 95 | self.model.eval() |
| 96 | self.model.to(device) |
| 97 | self.device = device |
| 98 | |
| 99 | def predict(self, img_rgb, text_prompts, scoreThres=0.3, iouThres=0.55): |
| 100 | img = img_rgb |
| 101 | text_prompts = text_prompts.split('.') |
| 102 | text_prompts.append('none') |
| 103 | texts = [text_prompts] |
| 104 | |
| 105 | img_height, img_width = img.shape[:2] |
| 106 | draw_img = img.copy() |
| 107 | |
| 108 | max_size = max(img_height, img_width) |
| 109 | M = np.array([ |
| 110 | [1024.0 / max_size, 0.0, 0.0], |
| 111 | [0.0, 1024.0 / max_size, 0.0]] |
| 112 | ).astype(np.float32) |
| 113 | invM = cv2.invertAffineTransform(M) |
| 114 | |
| 115 | resized_img = cv2.warpAffine(img, M, (1024, 1024), flags=cv2.INTER_LINEAR) |
| 116 | |
| 117 | resized_img = np.ascontiguousarray(resized_img) |
| 118 | resized_img = torch.from_numpy(resized_img).permute(2, 0, 1).float() |
| 119 | |
| 120 | resized_img[0, :, :] = (resized_img[0, :, :] - 123.6750) / 58.3950 |
| 121 | resized_img[1, :, :] = (resized_img[1, :, :] - 116.2800) / 57.1200 |
| 122 | resized_img[2, :, :] = (resized_img[2, :, :] - 103.5300) / 57.3750 |
| 123 | |
| 124 | resized_img = resized_img.to(self.device) |
| 125 | resized_img = resized_img.unsqueeze(0) |
| 126 | |
| 127 | with torch.no_grad(): |
| 128 | scores, pred_boxes, pred_masks, pred_classes = self.model(resized_img, texts) |
| 129 | |
| 130 | keeps = scores > scoreThres |
| 131 | scores = scores[keeps] |
| 132 | pred_boxes = pred_boxes[keeps] |
| 133 | pred_masks = pred_masks[keeps] |
| 134 | pred_classes = pred_classes[keeps] |
| 135 | |
| 136 | nms_keeps = torchvision.ops.nms(pred_boxes, scores, iou_threshold=iouThres) |
| 137 | |
| 138 | scores = scores[nms_keeps].detach().cpu().numpy() |
| 139 | pred_boxes = pred_boxes[nms_keeps].detach().cpu().numpy() |