:param prob_volume: (b, d, h, w) :param depth_values: (b, d, h, w) :return: (b, h, w)
(prob_volume, depth_values)
| 461 | |
| 462 | |
| 463 | def winner_take_all(prob_volume, depth_values): |
| 464 | """ |
| 465 | :param prob_volume: (b, d, h, w) |
| 466 | :param depth_values: (b, d, h, w) |
| 467 | :return: (b, h, w) |
| 468 | """ |
| 469 | _, idx = torch.max(prob_volume, dim=1, keepdim=True) |
| 470 | depth = torch.gather(depth_values, 1, idx).squeeze(1) |
| 471 | return depth |
| 472 | |
| 473 | |
| 474 |
nothing calls this directly
no outgoing calls
no test coverage detected