| 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() |
| 140 | pred_masks = pred_masks[nms_keeps].detach().cpu().numpy() |
| 141 | pred_classes = pred_classes[nms_keeps].detach().cpu().numpy() |
| 142 | |
| 143 | image_pil = Image.fromarray(np.uint8(draw_img)) |
| 144 | draw = ImageDraw.Draw(image_pil) |
| 145 | font = ImageFont.load_default() |
| 146 | |
| 147 | ret = [_COLORS[i] * 255 for i in range(len(_COLORS))] |
| 148 | num_remain = 1000 - len(ret) |
| 149 | thing_colors = ret + [random_color(rgb=True, maximum=255).astype(np.int32).tolist() for _ in range(num_remain)] |
| 150 | |
| 151 | info = [] |
| 152 | for index, (score, pred_box, pred_mask, pred_class) in enumerate(zip(scores, pred_boxes, pred_masks, pred_classes)): |
| 153 | pred_box = np.clip(pred_box, 0, 1024) |
| 154 | xmin = pred_box[0] * invM[0, 0] + pred_box[1] * invM[0, 1] + invM[0, 2] |
| 155 | ymin = pred_box[0] * invM[1, 0] + pred_box[1] * invM[1, 1] + invM[1, 2] |
| 156 | xmax = pred_box[2] * invM[0, 0] + pred_box[3] * invM[0, 1] + invM[0, 2] |