(args)
| 183 | |
| 184 | |
| 185 | def infer(args): |
| 186 | if not os.path.exists('samples'): |
| 187 | os.mkdir('samples') |
| 188 | faster_rcnn_model, voc, test_dataset = load_model_and_dataset(args) |
| 189 | |
| 190 | # Hard coding the low score threshold for inference on images for now |
| 191 | # Should come from config |
| 192 | faster_rcnn_model.roi_head.low_score_threshold = 0.7 |
| 193 | |
| 194 | for sample_count in tqdm(range(10)): |
| 195 | random_idx = random.randint(0, len(voc)) |
| 196 | im, target, fname = voc[random_idx] |
| 197 | im = im.unsqueeze(0).float().to(device) |
| 198 | |
| 199 | gt_im = cv2.imread(fname) |
| 200 | gt_im_copy = gt_im.copy() |
| 201 | |
| 202 | # Saving images with ground truth boxes |
| 203 | for idx, box in enumerate(target['bboxes']): |
| 204 | x1, y1, x2, y2 = box.detach().cpu().numpy() |
| 205 | x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) |
| 206 | |
| 207 | cv2.rectangle(gt_im, (x1, y1), (x2, y2), thickness=2, color=[0, 255, 0]) |
| 208 | cv2.rectangle(gt_im_copy, (x1, y1), (x2, y2), thickness=2, color=[0, 255, 0]) |
| 209 | text = voc.idx2label[target['labels'][idx].detach().cpu().item()] |
| 210 | text_size, _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_PLAIN, 1, 1) |
| 211 | text_w, text_h = text_size |
| 212 | cv2.rectangle(gt_im_copy , (x1, y1), (x1 + 10+text_w, y1 + 10+text_h), [255, 255, 255], -1) |
| 213 | cv2.putText(gt_im, text=voc.idx2label[target['labels'][idx].detach().cpu().item()], |
| 214 | org=(x1+5, y1+15), |
| 215 | thickness=1, |
| 216 | fontScale=1, |
| 217 | color=[0, 0, 0], |
| 218 | fontFace=cv2.FONT_HERSHEY_PLAIN) |
| 219 | cv2.putText(gt_im_copy, text=text, |
| 220 | org=(x1 + 5, y1 + 15), |
| 221 | thickness=1, |
| 222 | fontScale=1, |
| 223 | color=[0, 0, 0], |
| 224 | fontFace=cv2.FONT_HERSHEY_PLAIN) |
| 225 | cv2.addWeighted(gt_im_copy, 0.7, gt_im, 0.3, 0, gt_im) |
| 226 | cv2.imwrite('samples/output_frcnn_gt_{}.png'.format(sample_count), gt_im) |
| 227 | |
| 228 | # Getting predictions from trained model |
| 229 | rpn_output, frcnn_output = faster_rcnn_model(im, None) |
| 230 | boxes = frcnn_output['boxes'] |
| 231 | labels = frcnn_output['labels'] |
| 232 | scores = frcnn_output['scores'] |
| 233 | im = cv2.imread(fname) |
| 234 | im_copy = im.copy() |
| 235 | |
| 236 | # Saving images with predicted boxes |
| 237 | for idx, box in enumerate(boxes): |
| 238 | x1, y1, x2, y2 = box.detach().cpu().numpy() |
| 239 | x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) |
| 240 | cv2.rectangle(im, (x1, y1), (x2, y2), thickness=2, color=[0, 0, 255]) |
| 241 | cv2.rectangle(im_copy, (x1, y1), (x2, y2), thickness=2, color=[0, 0, 255]) |
| 242 | text = '{} : {:.2f}'.format(voc.idx2label[labels[idx].detach().cpu().item()], |
no test coverage detected