|
@@ -17,6 +17,7 @@ import torch
|
|
|
|
|
|
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.pyplot as plt
|
|
|
from models.base.transforms import get_transforms
|
|
from models.base.transforms import get_transforms
|
|
|
|
|
+from utils.data_process.show_prams import print_params
|
|
|
|
|
|
|
|
|
|
|
|
|
def validate_keypoints(keypoints, image_width, image_height):
|
|
def validate_keypoints(keypoints, image_width, image_height):
|
|
@@ -93,44 +94,52 @@ class LineDataset(BaseDataset):
|
|
|
target = {}
|
|
target = {}
|
|
|
|
|
|
|
|
target["image_id"] = torch.tensor(item)
|
|
target["image_id"] = torch.tensor(item)
|
|
|
- boxes, lines, points, arc_mask, circle_4points, labels, arc_ends, arc_params = get_boxes_lines(objs, shape)
|
|
|
|
|
|
|
+
|
|
|
|
|
+ #boxes, line_point_pairs, points, labels, mask_ends, mask_params
|
|
|
|
|
+ boxes, lines, points, labels, arc_ends, arc_params = get_boxes_lines(objs, shape)
|
|
|
|
|
+
|
|
|
|
|
+ # print_params(arc_ends, arc_params)
|
|
|
|
|
|
|
|
if points is not None:
|
|
if points is not None:
|
|
|
target["points"] = points
|
|
target["points"] = points
|
|
|
|
|
+ # if lines is not None:
|
|
|
|
|
+ # a = torch.full((lines.shape[0],), 2).unsqueeze(1)
|
|
|
|
|
+ # lines = torch.cat((lines, a), dim=1)
|
|
|
|
|
+ # target["lines"] = lines.to(torch.float32).view(-1, 2, 3)
|
|
|
if lines is not None:
|
|
if lines is not None:
|
|
|
- a = torch.full((lines.shape[0],), 2).unsqueeze(1)
|
|
|
|
|
- lines = torch.cat((lines, a), dim=1)
|
|
|
|
|
- target["lines"] = lines.to(torch.float32).view(-1, 2, 3)
|
|
|
|
|
- # print(f'lines shape:{ target["lines"].shape}')
|
|
|
|
|
-
|
|
|
|
|
- if arc_mask is not None:
|
|
|
|
|
- target['arc_mask'] = arc_mask
|
|
|
|
|
- # print(f'arc_mask dataset')
|
|
|
|
|
- # else:
|
|
|
|
|
- # print(f'not arc_mask dataset')
|
|
|
|
|
|
|
+ label_3d = labels.view(-1, 1, 1).expand(-1, 2, -1) # [N] -> [N,2,1]
|
|
|
|
|
+ line1 = torch.cat([lines, label_3d], dim=-1) # [N,2,3]
|
|
|
|
|
+ target["lines"] = line1.to(torch.float32)
|
|
|
|
|
|
|
|
if arc_ends is not None:
|
|
if arc_ends is not None:
|
|
|
target['mask_ends'] = arc_ends
|
|
target['mask_ends'] = arc_ends
|
|
|
|
|
+ if arc_params is not None:
|
|
|
target['mask_params'] = arc_params
|
|
target['mask_params'] = arc_params
|
|
|
|
|
|
|
|
- arc_angles = compute_arc_angles(arc_ends, arc_params)
|
|
|
|
|
- # print(arc_angles)
|
|
|
|
|
- # print(arc_params)
|
|
|
|
|
|
|
|
|
|
|
|
+ arc_angles = compute_arc_angles(arc_ends, arc_params)
|
|
|
|
|
+ # print_params(arc_angles)
|
|
|
arc_masks = []
|
|
arc_masks = []
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
for i in range(len(arc_params)):
|
|
for i in range(len(arc_params)):
|
|
|
- arc7=arc_params[i] + arc_angles[i].tolist()
|
|
|
|
|
- arc_masks.append(arc_to_mask(arc7, shape, line_width=1))
|
|
|
|
|
|
|
+ arc_param_i = arc_params[i].view(-1) # shape (5,)
|
|
|
|
|
+ arc_angle_i = arc_angles[i].view(-1) # shape (2,)
|
|
|
|
|
+ arc7 = torch.cat([arc_param_i, arc_angle_i], dim=0) # shape (7,)
|
|
|
|
|
+
|
|
|
|
|
|
|
|
- print(f'arc_masks:{torch.stack(arc_masks, dim=0).shape}')
|
|
|
|
|
- target['arc_masks'] = torch.stack(arc_masks, dim=0)
|
|
|
|
|
|
|
+ print_params(arc7)
|
|
|
|
|
+ mask = arc_to_mask(arc7, shape, line_width=1)
|
|
|
|
|
|
|
|
|
|
+ arc_masks.append(mask)
|
|
|
|
|
+ # arc7=arc_params[i] + arc_angles[i].tolist()
|
|
|
|
|
+ # arc_masks.append(arc_to_mask(arc7, shape, line_width=1))
|
|
|
|
|
+
|
|
|
|
|
+ # print(f'circle_masks:{torch.stack(arc_masks, dim=0).shape}')
|
|
|
|
|
+ target['circle_masks'] = torch.stack(arc_masks, dim=0)
|
|
|
|
|
|
|
|
|
|
|
|
|
- if circle_4points is not None:
|
|
|
|
|
- target['circles'] = circle_4points
|
|
|
|
|
- circle_masks = generate_ellipse_mask(shape, points_to_ellipse(circle_4points))
|
|
|
|
|
- target['circle_masks'] = torch.tensor(circle_masks, dtype=torch.float32).unsqueeze(0)
|
|
|
|
|
|
|
|
|
|
target["boxes"] = boxes
|
|
target["boxes"] = boxes
|
|
|
target["labels"] = labels
|
|
target["labels"] = labels
|
|
@@ -256,7 +265,11 @@ def arc_to_mask(arc7, shape, line_width=1):
|
|
|
Returns:
|
|
Returns:
|
|
|
mask (Tensor): [H, W], dtype=torch.uint8, 0/255
|
|
mask (Tensor): [H, W], dtype=torch.uint8, 0/255
|
|
|
"""
|
|
"""
|
|
|
|
|
+ # print_params(arc7)
|
|
|
# 确保 phi1 -> phi2 是正向(可处理跨 2π 的情况)
|
|
# 确保 phi1 -> phi2 是正向(可处理跨 2π 的情况)
|
|
|
|
|
+ if torch.all(torch.tensor(arc7) == 0):
|
|
|
|
|
+ return torch.zeros(shape, dtype=torch.uint8)
|
|
|
|
|
+
|
|
|
xc, yc, a, b, theta, phi1, phi2 = arc7
|
|
xc, yc, a, b, theta, phi1, phi2 = arc7
|
|
|
H, W = shape
|
|
H, W = shape
|
|
|
if phi2 < phi1:
|
|
if phi2 < phi1:
|
|
@@ -285,7 +298,7 @@ def arc_to_mask(arc7, shape, line_width=1):
|
|
|
# 绘制折线(antialias=False 更适合 mask)
|
|
# 绘制折线(antialias=False 更适合 mask)
|
|
|
cv2.polylines(img, [points], isClosed=False, color=255, thickness=line_width, lineType=cv2.LINE_AA)
|
|
cv2.polylines(img, [points], isClosed=False, color=255, thickness=line_width, lineType=cv2.LINE_AA)
|
|
|
|
|
|
|
|
- return torch.from_numpy(img).byte() # [H, W], values: 0 or 255
|
|
|
|
|
|
|
+ return torch.from_numpy(img).float() # [H, W], values: 0 or 255
|
|
|
|
|
|
|
|
|
|
|
|
|
def compute_arc_angles(gt_mask_ends, gt_mask_params):
|
|
def compute_arc_angles(gt_mask_ends, gt_mask_params):
|
|
@@ -299,10 +312,13 @@ def compute_arc_angles(gt_mask_ends, gt_mask_params):
|
|
|
Returns:
|
|
Returns:
|
|
|
phi: float, in [0, 2*pi)
|
|
phi: float, in [0, 2*pi)
|
|
|
"""
|
|
"""
|
|
|
|
|
+ # print_params(gt_mask_ends, gt_mask_params)
|
|
|
results = []
|
|
results = []
|
|
|
- gt_mask_params_tensor = torch.tensor(gt_mask_params,
|
|
|
|
|
- dtype=gt_mask_ends.dtype,
|
|
|
|
|
- device=gt_mask_ends.device)
|
|
|
|
|
|
|
+ if not isinstance(gt_mask_params, torch.Tensor):
|
|
|
|
|
+ gt_mask_params_tensor = torch.tensor(gt_mask_params, dtype=gt_mask_ends.dtype, device=gt_mask_ends.device)
|
|
|
|
|
+ else:
|
|
|
|
|
+ gt_mask_params_tensor = gt_mask_params.clone().detach().to(gt_mask_ends)
|
|
|
|
|
+
|
|
|
for ends_img, params_img in zip(gt_mask_ends, gt_mask_params_tensor):
|
|
for ends_img, params_img in zip(gt_mask_ends, gt_mask_params_tensor):
|
|
|
# print(f'params_img:{params_img}')
|
|
# print(f'params_img:{params_img}')
|
|
|
if torch.norm(params_img) < 1e-6: # L2 norm near zero
|
|
if torch.norm(params_img) < 1e-6: # L2 norm near zero
|
|
@@ -411,12 +427,11 @@ def get_boxes_lines(objs, shape):
|
|
|
boxes = []
|
|
boxes = []
|
|
|
labels = []
|
|
labels = []
|
|
|
h, w = shape
|
|
h, w = shape
|
|
|
|
|
+
|
|
|
line_point_pairs = []
|
|
line_point_pairs = []
|
|
|
points = []
|
|
points = []
|
|
|
- arc_mask = []
|
|
|
|
|
- arc_ends = []
|
|
|
|
|
- arc_params = []
|
|
|
|
|
- circle_4points = []
|
|
|
|
|
|
|
+ mask_ends = []
|
|
|
|
|
+ mask_params = []
|
|
|
|
|
|
|
|
for obj in objs:
|
|
for obj in objs:
|
|
|
# plt.plot([a[1], b[1]], [a[0], b[0]], c="red", linewidth=1) # a[1], b[1]无明确大小
|
|
# plt.plot([a[1], b[1]], [a[0], b[0]], c="red", linewidth=1) # a[1], b[1]无明确大小
|
|
@@ -426,8 +441,9 @@ def get_boxes_lines(objs, shape):
|
|
|
if label == 'line' or label == 'dseam1':
|
|
if label == 'line' or label == 'dseam1':
|
|
|
a, b = obj['points'][0], obj['points'][1]
|
|
a, b = obj['points'][0], obj['points'][1]
|
|
|
|
|
|
|
|
- line_point_pairs.append(a)
|
|
|
|
|
- line_point_pairs.append(b)
|
|
|
|
|
|
|
+ # line_point_pairs.append(a)
|
|
|
|
|
+ # line_point_pairs.append(b)
|
|
|
|
|
+ line_point_pairs.append([a, b])
|
|
|
|
|
|
|
|
xmin = max(0, (min(a[0], b[0]) - 6))
|
|
xmin = max(0, (min(a[0], b[0]) - 6))
|
|
|
xmax = min(w, (max(a[0], b[0]) + 6))
|
|
xmax = min(w, (max(a[0], b[0]) + 6))
|
|
@@ -437,6 +453,13 @@ def get_boxes_lines(objs, shape):
|
|
|
boxes.append([xmin, ymin, xmax, ymax])
|
|
boxes.append([xmin, ymin, xmax, ymax])
|
|
|
labels.append(torch.tensor(2))
|
|
labels.append(torch.tensor(2))
|
|
|
|
|
|
|
|
|
|
+ points.append(torch.tensor([0.0]))
|
|
|
|
|
+ mask_ends.append([[0, 0], [0, 0]])
|
|
|
|
|
+ mask_params.append([0, 0, 0, 0, 0])
|
|
|
|
|
+ # circle_4points.append([[0, 0], [0, 0], [0, 0], [0, 0]])
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
elif label == 'point':
|
|
elif label == 'point':
|
|
|
p = obj['points'][0]
|
|
p = obj['points'][0]
|
|
|
xmin = max(0, p[0] - 12)
|
|
xmin = max(0, p[0] - 12)
|
|
@@ -448,72 +471,86 @@ def get_boxes_lines(objs, shape):
|
|
|
labels.append(torch.tensor(1))
|
|
labels.append(torch.tensor(1))
|
|
|
boxes.append([xmin, ymin, xmax, ymax])
|
|
boxes.append([xmin, ymin, xmax, ymax])
|
|
|
|
|
|
|
|
|
|
+ line_point_pairs.append([[0, 0], [0, 0]])
|
|
|
|
|
+ mask_ends.append([[0, 0], [0, 0]])
|
|
|
|
|
+ mask_params.append([0, 0, 0, 0, 0])
|
|
|
|
|
+ # circle_4points.append([[0, 0], [0, 0], [0, 0], [0, 0]])
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+ # elif label == 'arc':
|
|
|
|
|
+ # arc_points = obj['points']
|
|
|
|
|
+ # arc_params = obj['params']
|
|
|
|
|
+ # arc_ends = obj['ends']
|
|
|
|
|
+ # line_mask.append(arc_points)
|
|
|
|
|
+ # mask_ends.append(arc_ends)
|
|
|
|
|
+ # mask_params.append(arc_params)
|
|
|
|
|
+ #
|
|
|
|
|
+ # xs = [p[0] for p in arc_points]
|
|
|
|
|
+ # ys = [p[1] for p in arc_points]
|
|
|
|
|
+ # xmin, xmax = min(xs), max(xs)
|
|
|
|
|
+ # ymin, ymax = min(ys), max(ys)
|
|
|
|
|
+ #
|
|
|
|
|
+ # boxes.append([xmin, ymin, xmax, ymax])
|
|
|
|
|
+ # labels.append(torch.tensor(3))
|
|
|
|
|
+ #
|
|
|
|
|
+ # points.append(torch.tensor([0.0]))
|
|
|
|
|
+ # line_point_pairs.append([[0, 0], [0, 0]])
|
|
|
|
|
+ # circle_4points.append([[0, 0], [0, 0], [0, 0], [0, 0]])
|
|
|
|
|
|
|
|
elif label == 'arc':
|
|
elif label == 'arc':
|
|
|
- arc_points = obj['points']
|
|
|
|
|
- params = obj['params']
|
|
|
|
|
- ends = obj['ends']
|
|
|
|
|
- arc_ends.append(ends)
|
|
|
|
|
- arc_params.append(params)
|
|
|
|
|
|
|
|
|
|
- xs = [p[0] for p in arc_points]
|
|
|
|
|
- ys = [p[1] for p in arc_points]
|
|
|
|
|
- xmin, xmax = min(xs), max(xs)
|
|
|
|
|
- ymin, ymax = min(ys), max(ys)
|
|
|
|
|
-
|
|
|
|
|
- boxes.append([xmin, ymin, xmax, ymax])
|
|
|
|
|
- labels.append(torch.tensor(3))
|
|
|
|
|
-
|
|
|
|
|
- elif label == 'circle':
|
|
|
|
|
- # print(f'len circle_4points: {len(obj['points'])}')
|
|
|
|
|
- points = sort_points_clockwise(obj['points'])
|
|
|
|
|
- circle_4points.append(points)
|
|
|
|
|
-
|
|
|
|
|
- xmin = max(obj['xmin'] - 40, 0)
|
|
|
|
|
- xmax = min(obj['xmax'] + 40, w)
|
|
|
|
|
- ymin = max(obj['ymin'] - 40, 0)
|
|
|
|
|
- ymax = min(obj['ymax'] + 40, h)
|
|
|
|
|
|
|
+ arc_params = obj['params']
|
|
|
|
|
+ arc_ends = obj['ends']
|
|
|
|
|
+ mask_ends.append(arc_ends)
|
|
|
|
|
+ mask_params.append(arc_params)
|
|
|
|
|
+ arc3points = obj['points']
|
|
|
|
|
+ xs = [p[0] for p in arc3points]
|
|
|
|
|
+ ys = [p[1] for p in arc3points]
|
|
|
|
|
+ xmin_raw = min(xs)
|
|
|
|
|
+ xmax_raw = max(xs)
|
|
|
|
|
+ ymin_raw = min(ys)
|
|
|
|
|
+ ymax_raw = max(ys)
|
|
|
|
|
+ xmin = max(xmin_raw - 40, 0)
|
|
|
|
|
+ xmax = min(xmax_raw + 40, w)
|
|
|
|
|
+ ymin = max(ymin_raw - 40, 0)
|
|
|
|
|
+ ymax = min(ymax_raw + 40, h)
|
|
|
|
|
|
|
|
boxes.append([xmin, ymin, xmax, ymax])
|
|
boxes.append([xmin, ymin, xmax, ymax])
|
|
|
labels.append(torch.tensor(4))
|
|
labels.append(torch.tensor(4))
|
|
|
|
|
+ points.append(torch.tensor([0.0]))
|
|
|
|
|
+ line_point_pairs.append([[0, 0], [0, 0]])
|
|
|
|
|
+
|
|
|
|
|
|
|
|
boxes = torch.tensor(boxes, dtype=torch.float32)
|
|
boxes = torch.tensor(boxes, dtype=torch.float32)
|
|
|
print(f'boxes:{boxes.shape}')
|
|
print(f'boxes:{boxes.shape}')
|
|
|
labels = torch.tensor(labels)
|
|
labels = torch.tensor(labels)
|
|
|
- if len(points) == 0:
|
|
|
|
|
- points = None
|
|
|
|
|
- else:
|
|
|
|
|
|
|
+
|
|
|
|
|
+ if points:
|
|
|
points = torch.tensor(points, dtype=torch.float32)
|
|
points = torch.tensor(points, dtype=torch.float32)
|
|
|
- print(f'read labels:{labels}')
|
|
|
|
|
- # print(f'read points:{points}')
|
|
|
|
|
- if len(line_point_pairs) == 0:
|
|
|
|
|
- line_point_pairs = None
|
|
|
|
|
else:
|
|
else:
|
|
|
- line_point_pairs = torch.tensor(line_point_pairs)
|
|
|
|
|
- # print(f'line_point_pairs:{line_point_pairs.shape},{line_point_pairs.dtype}')
|
|
|
|
|
-
|
|
|
|
|
- # print(f'boxes:{boxes.shape},line_point_pairs:{line_point_pairs.shape}')
|
|
|
|
|
|
|
+ points = None
|
|
|
|
|
|
|
|
- if len(arc_mask) == 0:
|
|
|
|
|
- arc_mask = None
|
|
|
|
|
|
|
+ if line_point_pairs:
|
|
|
|
|
+ line_point_pairs = torch.tensor(line_point_pairs, dtype=torch.float32)
|
|
|
else:
|
|
else:
|
|
|
- arc_mask = torch.tensor(arc_mask, dtype=torch.float32)
|
|
|
|
|
- print(f'arc_mask shape :{arc_mask.shape},{arc_mask.dtype}')
|
|
|
|
|
|
|
+ line_point_pairs = None
|
|
|
|
|
|
|
|
- if len(arc_ends) == 0:
|
|
|
|
|
- arc_ends = None
|
|
|
|
|
|
|
+ if mask_ends:
|
|
|
|
|
+ mask_ends = torch.tensor(mask_ends, dtype=torch.float32)
|
|
|
else:
|
|
else:
|
|
|
- arc_ends = torch.tensor(arc_ends, dtype=torch.float32)
|
|
|
|
|
|
|
+ mask_ends = None
|
|
|
|
|
|
|
|
- if len(circle_4points) == 0:
|
|
|
|
|
- circle_4points = None
|
|
|
|
|
|
|
+ if mask_params:
|
|
|
|
|
+ mask_params = torch.tensor(mask_params, dtype=torch.float32)
|
|
|
else:
|
|
else:
|
|
|
- # for circle_4point in circle_4points:
|
|
|
|
|
- # print(f'circle_4point len111:{len(circle_4point)}')
|
|
|
|
|
- circle_4points = torch.tensor(circle_4points, dtype=torch.float32)
|
|
|
|
|
- # print(f'circle_4points shape:{circle_4points.shape}')
|
|
|
|
|
|
|
+ mask_params = None
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
|
|
|
- return boxes, line_point_pairs, points, arc_mask, circle_4points, labels, arc_ends, arc_params
|
|
|
|
|
|
|
+ return boxes, line_point_pairs, points, labels, mask_ends, mask_params
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
if __name__ == '__main__':
|