| 197 | |
| 198 | |
| 199 | def predict_with_args(self, args): |
| 200 | img = cv2.imread(args.img_path) |
| 201 | text_prompts = args.text_prompts.split('.') |
| 202 | text_prompts.append('none') |
| 203 | texts = [text_prompts] |
| 204 | |
| 205 | img_height, img_width = img.shape[:2] |
| 206 | draw_img = img.copy() |
| 207 | |
| 208 | max_size = max(img_height, img_width) |
| 209 | M = np.array([ |
| 210 | [1024.0 / max_size, 0.0, 0.0], |
| 211 | [0.0, 1024.0 / max_size, 0.0]] |
| 212 | ).astype(np.float32) |
| 213 | invM = cv2.invertAffineTransform(M) |
| 214 | |
| 215 | resized_img = cv2.warpAffine(img, M, (1024, 1024), flags=cv2.INTER_LINEAR) |
| 216 | |
| 217 | resized_img = resized_img[:, :, ::-1] |
| 218 | resized_img = np.ascontiguousarray(resized_img) |
| 219 | resized_img = torch.from_numpy(resized_img).permute(2, 0, 1).float() |
| 220 | |
| 221 | resized_img[0, :, :] = (resized_img[0, :, :] - 123.6750) / 58.3950 |
| 222 | resized_img[1, :, :] = (resized_img[1, :, :] - 116.2800) / 57.1200 |
| 223 | resized_img[2, :, :] = (resized_img[2, :, :] - 103.5300) / 57.3750 |
| 224 | |
| 225 | resized_img = resized_img.to(self.device) |
| 226 | resized_img = resized_img.unsqueeze(0) |
| 227 | |
| 228 | with torch.no_grad(): |
| 229 | scores, pred_boxes, pred_masks, pred_classes = self.model(resized_img, texts) |
| 230 | |
| 231 | scoreThres = args.scoreThres |
| 232 | keeps = scores > scoreThres |
| 233 | scores = scores[keeps] |
| 234 | pred_boxes = pred_boxes[keeps] |
| 235 | pred_masks = pred_masks[keeps] |
| 236 | pred_classes = pred_classes[keeps] |
| 237 | |
| 238 | nms_keeps = torchvision.ops.nms(pred_boxes, scores, iou_threshold=args.iouThres) |
| 239 | |
| 240 | scores = scores[nms_keeps].detach().cpu().numpy() |
| 241 | pred_boxes = pred_boxes[nms_keeps].detach().cpu().numpy() |
| 242 | pred_masks = pred_masks[nms_keeps].detach().cpu().numpy() |
| 243 | pred_classes = pred_classes[nms_keeps].detach().cpu().numpy() |
| 244 | |
| 245 | image_pil = Image.fromarray(np.uint8(draw_img[:, :, ::-1])) |
| 246 | draw = ImageDraw.Draw(image_pil) |
| 247 | font = ImageFont.load_default() |
| 248 | |
| 249 | ret = [_COLORS[i] * 255 for i in range(len(_COLORS))] |
| 250 | num_remain = 100 - len(ret) |
| 251 | thing_colors = ret + [random_color(rgb=True, maximum=255).astype(np.int32).tolist() for _ in range(num_remain)] |
| 252 | |
| 253 | info = [] |
| 254 | for index, (score, pred_box, pred_mask, pred_class) in enumerate(zip(scores, pred_boxes, pred_masks, pred_classes)): |
| 255 | pred_box = np.clip(pred_box, 0, 1024) |
| 256 | xmin = pred_box[0] * invM[0, 0] + pred_box[1] * invM[0, 1] + invM[0, 2] |