|
@@ -751,16 +751,18 @@ def line_inference(x, line_boxes):
|
|
|
|
|
|
|
|
return lines_probs, lines_scores
|
|
return lines_probs, lines_scores
|
|
|
|
|
|
|
|
-def arc_inference(x, point_boxes):
|
|
|
|
|
|
|
+def arc_inference(x, arc_boxes):
|
|
|
# type: (Tensor, List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]
|
|
# type: (Tensor, List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]
|
|
|
|
|
|
|
|
points_probs = []
|
|
points_probs = []
|
|
|
points_scores = []
|
|
points_scores = []
|
|
|
|
|
|
|
|
- boxes_per_image = [box.size(0) for box in point_boxes]
|
|
|
|
|
|
|
+ print(f'arc_boxes:{len(arc_boxes)}')
|
|
|
|
|
+
|
|
|
|
|
+ boxes_per_image = [box.size(0) for box in arc_boxes]
|
|
|
x2 = x.split(boxes_per_image, dim=0)
|
|
x2 = x.split(boxes_per_image, dim=0)
|
|
|
|
|
|
|
|
- for xx, bb in zip(x2, point_boxes):
|
|
|
|
|
|
|
+ for xx, bb in zip(x2, arc_boxes):
|
|
|
point_prob,point_scores = heatmaps_to_arc(xx, bb)
|
|
point_prob,point_scores = heatmaps_to_arc(xx, bb)
|
|
|
|
|
|
|
|
points_probs.append(point_prob.unsqueeze(1))
|
|
points_probs.append(point_prob.unsqueeze(1))
|
|
@@ -768,7 +770,6 @@ def arc_inference(x, point_boxes):
|
|
|
|
|
|
|
|
return points_probs,points_scores
|
|
return points_probs,points_scores
|
|
|
|
|
|
|
|
-import torch.nn.functional as F
|
|
|
|
|
|
|
|
|
|
import torch.nn.functional as F
|
|
import torch.nn.functional as F
|
|
|
|
|
|