| 224 | [_generate_random_head_inputs(i) for i in range(len(heads_dicts))], |
| 225 | ) |
| 226 | def test_head(head_dict, input_shape, num_keypoints): |
| 227 | w, h = input_shape |
| 228 | head_dict = copy.deepcopy(head_dict) |
| 229 | |
| 230 | head_type = head_dict["type"] |
| 231 | input_channels = head_dict.pop("input_channels") |
| 232 | output_channels = head_dict.pop("output_channels") |
| 233 | total_stride = head_dict.pop("total_stride") |
| 234 | if head_type == "HeatmapHead": |
| 235 | output_channels = num_keypoints |
| 236 | head_dict["heatmap_config"]["channels"][2] = output_channels |
| 237 | head_dict["locref_config"]["channels"][2] = 2 * output_channels |
| 238 | head_dict["target_generator"]["num_heatmaps"] = output_channels |
| 239 | input_tensor = torch.zeros((1, input_channels, h, w)) |
| 240 | |
| 241 | elif head_type == "TransformerHead": |
| 242 | output_channels = num_keypoints |
| 243 | input_channels = num_keypoints |
| 244 | head_dict["heatmap_dim"] = h * w |
| 245 | head_dict["heatmap_size"] = [h, w] |
| 246 | head_dict["target_generator"]["num_heatmaps"] = output_channels |
| 247 | input_tensor = torch.zeros((1, input_channels, head_dict["dim"] * 3)) |
| 248 | |
| 249 | elif head_type == "DEKRHead": |
| 250 | output_channels = num_keypoints + 1 |
| 251 | head_dict["target_generator"]["num_joints"] = num_keypoints |
| 252 | head_dict["heatmap_config"]["channels"][2] = num_keypoints + 1 |
| 253 | head_dict["offset_config"]["channels"][1] = num_keypoints * head_dict["offset_config"]["num_offset_per_kpt"] |
| 254 | head_dict["offset_config"]["channels"][2] = num_keypoints |
| 255 | input_tensor = torch.zeros((1, input_channels, h, w)) |
| 256 | |
| 257 | if "type" in head_dict["criterion"]: |
| 258 | head_dict["criterion"] = CRITERIONS.build(head_dict["criterion"]) |
| 259 | else: |
| 260 | weights = {} |
| 261 | criterions = {} |
| 262 | for loss_name, criterion_cfg in head_dict["criterion"].items(): |
| 263 | weights[loss_name] = criterion_cfg.get("weight", 1.0) |
| 264 | criterion_cfg = {k: v for k, v in criterion_cfg.items() if k != "weight"} |
| 265 | criterions[loss_name] = CRITERIONS.build(criterion_cfg) |
| 266 | |
| 267 | aggregator_cfg = {"type": "WeightedLossAggregator", "weights": weights} |
| 268 | head_dict["aggregator"] = LOSS_AGGREGATORS.build(aggregator_cfg) |
| 269 | head_dict["criterion"] = criterions |
| 270 | |
| 271 | head_dict["target_generator"] = TARGET_GENERATORS.build(head_dict["target_generator"]) |
| 272 | head_dict["predictor"] = PREDICTORS.build(head_dict["predictor"]) |
| 273 | head = dlc_models.HEADS.build(head_dict) |
| 274 | |
| 275 | output = head(input_tensor)["heatmap"] |
| 276 | _, c_out, h_out, w_out = output.shape |
| 277 | assert (h_out == h * total_stride) and (w_out == w * total_stride) |
| 278 | assert c_out == output_channels |
| 279 | |
| 280 | |
| 281 | def test_msa_hrnet(): |