Extract image features with a list of person bounding boxes. Args: model (nn.Module): The loaded feature extraction model. img_or_path (Union[str, np.ndarray]): Image filename or loaded image. det_results (List(dict)): the item in the dict may contain 'bbox'
(
model,
img_or_path,
det_results,
bbox_thr=None,
format='xywh',
)
| 293 | |
| 294 | |
| 295 | def feature_extract( |
| 296 | model, |
| 297 | img_or_path, |
| 298 | det_results, |
| 299 | bbox_thr=None, |
| 300 | format='xywh', |
| 301 | ): |
| 302 | """Extract image features with a list of person bounding boxes. |
| 303 | |
| 304 | Args: |
| 305 | model (nn.Module): The loaded feature extraction model. |
| 306 | img_or_path (Union[str, np.ndarray]): Image filename or loaded image. |
| 307 | det_results (List(dict)): the item in the dict may contain |
| 308 | 'bbox' and/or 'track_id'. |
| 309 | 'bbox' (4, ) or (5, ): The person bounding box, which contains |
| 310 | 4 box coordinates (and score). |
| 311 | 'track_id' (int): The unique id for each human instance. |
| 312 | bbox_thr (float, optional): Threshold for bounding boxes. |
| 313 | If bbox_thr is None, ignore it. Defaults to None. |
| 314 | format (str, optional): bbox format. Default: 'xywh'. |
| 315 | 'xyxy' means (left, top, right, bottom), |
| 316 | 'xywh' means (left, top, width, height). |
| 317 | |
| 318 | Returns: |
| 319 | list[dict]: The bbox & pose info, |
| 320 | containing the bbox: (left, top, right, bottom, [score]) |
| 321 | and the features. |
| 322 | """ |
| 323 | # only two kinds of bbox format is supported. |
| 324 | assert format in ['xyxy', 'xywh'] |
| 325 | |
| 326 | cfg = model.cfg |
| 327 | device = next(model.parameters()).device |
| 328 | |
| 329 | feature_results = [] |
| 330 | if len(det_results) == 0: |
| 331 | return feature_results |
| 332 | |
| 333 | # Change for-loop preprocess each bbox to preprocess all bboxes at once. |
| 334 | bboxes = np.array([box['bbox'] for box in det_results]) |
| 335 | assert len(bboxes[0]) in [4, 5] |
| 336 | |
| 337 | # Select bboxes by score threshold |
| 338 | if bbox_thr is not None: |
| 339 | assert bboxes.shape[1] == 5 |
| 340 | valid_idx = np.where(bboxes[:, 4] > bbox_thr)[0] |
| 341 | bboxes = bboxes[valid_idx] |
| 342 | det_results = [det_results[i] for i in valid_idx] |
| 343 | |
| 344 | # if bbox_thr remove all bounding box |
| 345 | if len(bboxes) == 0: |
| 346 | return feature_results |
| 347 | |
| 348 | if format == 'xyxy': |
| 349 | bboxes_xyxy = bboxes |
| 350 | bboxes_xywh = xyxy2xywh(bboxes) |
| 351 | else: |
| 352 | # format is already 'xywh' |