|
@@ -1,3 +1,4 @@
|
|
|
|
|
+import cv2
|
|
|
import imageio
|
|
import imageio
|
|
|
import numpy as np
|
|
import numpy as np
|
|
|
from skimage.draw import ellipse
|
|
from skimage.draw import ellipse
|
|
@@ -25,13 +26,15 @@ def validate_keypoints(keypoints, image_width, image_height):
|
|
|
raise ValueError(f"Key point ({x}, {y}) is out of bounds for image size ({image_width}, {image_height})")
|
|
raise ValueError(f"Key point ({x}, {y}) is out of bounds for image size ({image_width}, {image_height})")
|
|
|
|
|
|
|
|
|
|
|
|
|
-
|
|
|
|
|
"""
|
|
"""
|
|
|
直接读取xanlabel标注的数据集json格式
|
|
直接读取xanlabel标注的数据集json格式
|
|
|
|
|
|
|
|
"""
|
|
"""
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
class LineDataset(BaseDataset):
|
|
class LineDataset(BaseDataset):
|
|
|
- def __init__(self, dataset_path, data_type, transforms=None,augmentation=False, dataset_type=None,img_type='rgb', target_type='pixel'):
|
|
|
|
|
|
|
+ def __init__(self, dataset_path, data_type, transforms=None, augmentation=False, dataset_type=None, img_type='rgb',
|
|
|
|
|
+ target_type='pixel'):
|
|
|
super().__init__(dataset_path)
|
|
super().__init__(dataset_path)
|
|
|
|
|
|
|
|
self.data_path = dataset_path
|
|
self.data_path = dataset_path
|
|
@@ -43,8 +46,8 @@ class LineDataset(BaseDataset):
|
|
|
self.imgs = os.listdir(self.img_path)
|
|
self.imgs = os.listdir(self.img_path)
|
|
|
self.lbls = os.listdir(self.lbl_path)
|
|
self.lbls = os.listdir(self.lbl_path)
|
|
|
self.target_type = target_type
|
|
self.target_type = target_type
|
|
|
- self.img_type=img_type
|
|
|
|
|
- self.augmentation=augmentation
|
|
|
|
|
|
|
+ self.img_type = img_type
|
|
|
|
|
+ self.augmentation = augmentation
|
|
|
print(f'augmentation:{augmentation}')
|
|
print(f'augmentation:{augmentation}')
|
|
|
# self.default_transform = DefaultTransform()
|
|
# self.default_transform = DefaultTransform()
|
|
|
|
|
|
|
@@ -53,14 +56,13 @@ class LineDataset(BaseDataset):
|
|
|
|
|
|
|
|
if self.data_type == 'tiff':
|
|
if self.data_type == 'tiff':
|
|
|
lbl_path = os.path.join(self.lbl_path, self.imgs[index][:-4] + 'json')
|
|
lbl_path = os.path.join(self.lbl_path, self.imgs[index][:-4] + 'json')
|
|
|
- img = imageio.v3.imread(img_path)[:,:,0]
|
|
|
|
|
|
|
+ img = imageio.v3.imread(img_path)[:, :, 0]
|
|
|
print(f'img shape:{img.shape}')
|
|
print(f'img shape:{img.shape}')
|
|
|
w, h = img.shape[:2]
|
|
w, h = img.shape[:2]
|
|
|
- img=img.reshape(w,h,1)
|
|
|
|
|
|
|
+ img = img.reshape(w, h, 1)
|
|
|
img_3channel = np.zeros((w, h, 3), dtype=img.dtype)
|
|
img_3channel = np.zeros((w, h, 3), dtype=img.dtype)
|
|
|
img_3channel[:, :, 2] = img[:, :, 0]
|
|
img_3channel[:, :, 2] = img[:, :, 0]
|
|
|
|
|
|
|
|
-
|
|
|
|
|
img = torch.from_numpy(img_3channel).permute(2, 1, 0)
|
|
img = torch.from_numpy(img_3channel).permute(2, 1, 0)
|
|
|
else:
|
|
else:
|
|
|
lbl_path = os.path.join(self.lbl_path, self.imgs[index][:-3] + 'json')
|
|
lbl_path = os.path.join(self.lbl_path, self.imgs[index][:-3] + 'json')
|
|
@@ -69,12 +71,10 @@ class LineDataset(BaseDataset):
|
|
|
# wire_labels, target = self.read_target(item=index, lbl_path=lbl_path, shape=(h, w))
|
|
# wire_labels, target = self.read_target(item=index, lbl_path=lbl_path, shape=(h, w))
|
|
|
target = self.read_target(item=index, lbl_path=lbl_path, shape=(h, w))
|
|
target = self.read_target(item=index, lbl_path=lbl_path, shape=(h, w))
|
|
|
|
|
|
|
|
-
|
|
|
|
|
- self.transforms=get_transforms(augmention=self.augmentation)
|
|
|
|
|
|
|
+ self.transforms = get_transforms(augmention=self.augmentation)
|
|
|
|
|
|
|
|
img, target = self.transforms(img, target)
|
|
img, target = self.transforms(img, target)
|
|
|
|
|
|
|
|
-
|
|
|
|
|
return img, target
|
|
return img, target
|
|
|
|
|
|
|
|
def __len__(self):
|
|
def __len__(self):
|
|
@@ -86,20 +86,17 @@ class LineDataset(BaseDataset):
|
|
|
with open(lbl_path, 'r') as file:
|
|
with open(lbl_path, 'r') as file:
|
|
|
lable_all = json.load(file)
|
|
lable_all = json.load(file)
|
|
|
|
|
|
|
|
-
|
|
|
|
|
objs = lable_all["shapes"]
|
|
objs = lable_all["shapes"]
|
|
|
- point_pairs=objs[0]['points']
|
|
|
|
|
-
|
|
|
|
|
|
|
+ point_pairs = objs[0]['points']
|
|
|
|
|
|
|
|
# print(f'point_pairs:{point_pairs}')
|
|
# print(f'point_pairs:{point_pairs}')
|
|
|
target = {}
|
|
target = {}
|
|
|
|
|
|
|
|
target["image_id"] = torch.tensor(item)
|
|
target["image_id"] = torch.tensor(item)
|
|
|
- boxes, lines, points, arc_mask,circle_4points,labels = get_boxes_lines(objs, shape)
|
|
|
|
|
-
|
|
|
|
|
|
|
+ boxes, lines, points, arc_mask, circle_4points, labels, arc_ends, arc_params = get_boxes_lines(objs, shape)
|
|
|
|
|
|
|
|
if points is not None:
|
|
if points is not None:
|
|
|
- target["points"]=points
|
|
|
|
|
|
|
+ target["points"] = points
|
|
|
if lines is not None:
|
|
if lines is not None:
|
|
|
a = torch.full((lines.shape[0],), 2).unsqueeze(1)
|
|
a = torch.full((lines.shape[0],), 2).unsqueeze(1)
|
|
|
lines = torch.cat((lines, a), dim=1)
|
|
lines = torch.cat((lines, a), dim=1)
|
|
@@ -107,36 +104,52 @@ class LineDataset(BaseDataset):
|
|
|
# print(f'lines shape:{ target["lines"].shape}')
|
|
# print(f'lines shape:{ target["lines"].shape}')
|
|
|
|
|
|
|
|
if arc_mask is not None:
|
|
if arc_mask is not None:
|
|
|
- target['arc_mask']=arc_mask
|
|
|
|
|
|
|
+ target['arc_mask'] = arc_mask
|
|
|
# print(f'arc_mask dataset')
|
|
# print(f'arc_mask dataset')
|
|
|
# else:
|
|
# else:
|
|
|
# print(f'not arc_mask dataset')
|
|
# print(f'not arc_mask dataset')
|
|
|
|
|
|
|
|
|
|
+ if arc_ends is not None:
|
|
|
|
|
+ target['mask_ends'] = arc_ends
|
|
|
|
|
+ target['mask_params'] = arc_params
|
|
|
|
|
+
|
|
|
|
|
+ arc_angles = compute_arc_angles(arc_ends, arc_params)
|
|
|
|
|
+ # print(arc_angles)
|
|
|
|
|
+ # print(arc_params)
|
|
|
|
|
+
|
|
|
|
|
+ arc_masks = []
|
|
|
|
|
+ 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))
|
|
|
|
|
+
|
|
|
|
|
+ print(f'arc_masks:{torch.stack(arc_masks, dim=0).shape}')
|
|
|
|
|
+ target['arc_masks'] = torch.stack(arc_masks, dim=0)
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
if circle_4points is not None:
|
|
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['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["labels"]=labels
|
|
|
|
|
|
|
+ target["boxes"] = boxes
|
|
|
|
|
+ target["labels"] = labels
|
|
|
# target["boxes"], lines,target["points"], target["labels"] = get_boxes_lines(objs,shape)
|
|
# target["boxes"], lines,target["points"], target["labels"] = get_boxes_lines(objs,shape)
|
|
|
# print(f'lines:{lines}')
|
|
# print(f'lines:{lines}')
|
|
|
# target["labels"] = torch.ones(len(target["boxes"]), dtype=torch.int64)
|
|
# target["labels"] = torch.ones(len(target["boxes"]), dtype=torch.int64)
|
|
|
# print(f'target points:{target["points"]}')
|
|
# print(f'target points:{target["points"]}')
|
|
|
|
|
|
|
|
-
|
|
|
|
|
-
|
|
|
|
|
# target["lines"] = lines.to(torch.float32).view(-1,2,3)
|
|
# target["lines"] = lines.to(torch.float32).view(-1,2,3)
|
|
|
|
|
|
|
|
# print(f'')
|
|
# print(f'')
|
|
|
|
|
|
|
|
# print(f'lines:{target["lines"].shape}')
|
|
# print(f'lines:{target["lines"].shape}')
|
|
|
- target["img_size"]=shape
|
|
|
|
|
|
|
+ target["img_size"] = shape
|
|
|
|
|
|
|
|
# validate_keypoints(lines, shape[0], shape[1])
|
|
# validate_keypoints(lines, shape[0], shape[1])
|
|
|
return target
|
|
return target
|
|
|
|
|
|
|
|
- def show(self, idx,show_type='all'):
|
|
|
|
|
|
|
+ def show(self, idx, show_type='all'):
|
|
|
image, target = self.__getitem__(idx)
|
|
image, target = self.__getitem__(idx)
|
|
|
|
|
|
|
|
cmap = plt.get_cmap("jet")
|
|
cmap = plt.get_cmap("jet")
|
|
@@ -147,15 +160,26 @@ class LineDataset(BaseDataset):
|
|
|
# img_path = os.path.join(self.img_path, self.imgs[idx])
|
|
# img_path = os.path.join(self.img_path, self.imgs[idx])
|
|
|
# print(f'boxes:{target["boxes"]}')
|
|
# print(f'boxes:{target["boxes"]}')
|
|
|
img = image
|
|
img = image
|
|
|
- if show_type=='all':
|
|
|
|
|
|
|
+
|
|
|
|
|
+ if show_type == 'arc_masks':
|
|
|
boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), target["boxes"],
|
|
boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), target["boxes"],
|
|
|
- colors="yellow", width=1)
|
|
|
|
|
- circle=target['circles']
|
|
|
|
|
- circle_mask=target['circle_masks']
|
|
|
|
|
|
|
+ colors="yellow", width=1)
|
|
|
|
|
+ # arc = target['arc']
|
|
|
|
|
+ arc_mask = target['arc_masks']
|
|
|
|
|
+ # print(f'taget circle:{arc.shape}')
|
|
|
|
|
+ print(f'target circle_masks:{arc_mask.shape}')
|
|
|
|
|
+ plt.imshow(arc_mask.squeeze(0))
|
|
|
|
|
+ plt.show()
|
|
|
|
|
+
|
|
|
|
|
+ if show_type == 'circle_masks':
|
|
|
|
|
+ boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), target["boxes"],
|
|
|
|
|
+ colors="yellow", width=1)
|
|
|
|
|
+ circle = target['circles']
|
|
|
|
|
+ circle_mask = target['circle_masks']
|
|
|
print(f'taget circle:{circle.shape}')
|
|
print(f'taget circle:{circle.shape}')
|
|
|
print(f'target circle_masks:{circle_mask.shape}')
|
|
print(f'target circle_masks:{circle_mask.shape}')
|
|
|
plt.imshow(circle_mask.squeeze(0))
|
|
plt.imshow(circle_mask.squeeze(0))
|
|
|
- keypoint_img=draw_keypoints(boxed_image,circle,colors='red',width=3)
|
|
|
|
|
|
|
+ keypoint_img = draw_keypoints(boxed_image, circle, colors='red', width=3)
|
|
|
# plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
|
|
# plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
|
|
|
plt.show()
|
|
plt.show()
|
|
|
|
|
|
|
@@ -164,23 +188,151 @@ class LineDataset(BaseDataset):
|
|
|
# plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
|
|
# plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
|
|
|
# plt.show()
|
|
# plt.show()
|
|
|
|
|
|
|
|
- if show_type=='points':
|
|
|
|
|
|
|
+ if show_type == 'points':
|
|
|
# print(f'points:{target['points'].shape}')
|
|
# print(f'points:{target['points'].shape}')
|
|
|
- keypoint_img=draw_keypoints((img * 255).to(torch.uint8),target['points'].unsqueeze(1),colors='red',width=3)
|
|
|
|
|
|
|
+ keypoint_img = draw_keypoints((img * 255).to(torch.uint8), target['points'].unsqueeze(1), colors='red',
|
|
|
|
|
+ width=3)
|
|
|
plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
|
|
plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
|
|
|
plt.show()
|
|
plt.show()
|
|
|
|
|
|
|
|
- if show_type=='boxes':
|
|
|
|
|
|
|
+ if show_type == 'boxes':
|
|
|
boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), target["boxes"],
|
|
boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), target["boxes"],
|
|
|
colors="yellow", width=1)
|
|
colors="yellow", width=1)
|
|
|
plt.imshow(boxed_image.permute(1, 2, 0).numpy())
|
|
plt.imshow(boxed_image.permute(1, 2, 0).numpy())
|
|
|
plt.show()
|
|
plt.show()
|
|
|
|
|
|
|
|
|
|
+ def show_img(self, img_path):
|
|
|
|
|
+ pass
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
+def draw_el(all):
|
|
|
|
|
+ # 解析椭圆参数
|
|
|
|
|
+ if isinstance(all, torch.Tensor):
|
|
|
|
|
+ all = all.cpu().numpy()
|
|
|
|
|
+ x, y, a, b, q, q1, q2 = all
|
|
|
|
|
+ theta = np.radians(q)
|
|
|
|
|
+ phi1 = np.radians(q1) # 第一个点的参数角
|
|
|
|
|
+ phi2 = np.radians(q2) # 第二个点的参数角
|
|
|
|
|
|
|
|
- def show_img(self, img_path):
|
|
|
|
|
- pass
|
|
|
|
|
|
|
+ # 生成椭圆上的点
|
|
|
|
|
+ phi = np.linspace(0, 2 * np.pi, 500)
|
|
|
|
|
+ x_ellipse = x + a * np.cos(phi) * np.cos(theta) - b * np.sin(phi) * np.sin(theta)
|
|
|
|
|
+ y_ellipse = y + a * np.cos(phi) * np.sin(theta) + b * np.sin(phi) * np.cos(theta)
|
|
|
|
|
+
|
|
|
|
|
+ # 计算两个指定点的坐标
|
|
|
|
|
+ def param_to_point(phi, xc, yc, a, b, theta):
|
|
|
|
|
+ x = xc + a * np.cos(phi) * np.cos(theta) - b * np.sin(phi) * np.sin(theta)
|
|
|
|
|
+ y = yc + a * np.cos(phi) * np.sin(theta) + b * np.sin(phi) * np.cos(theta)
|
|
|
|
|
+ return x, y
|
|
|
|
|
+
|
|
|
|
|
+ P1 = param_to_point(phi1, x, y, a, b, theta)
|
|
|
|
|
+ P2 = param_to_point(phi2, x, y, a, b, theta)
|
|
|
|
|
+
|
|
|
|
|
+ # 创建画布并显示背景图片(使用传入的background_img,shape为[H, W, C])
|
|
|
|
|
+ plt.figure(figsize=(10, 10))
|
|
|
|
|
+ # plt.imshow(background_img) # 直接显示背景图
|
|
|
|
|
+
|
|
|
|
|
+ # 绘制椭圆及相关元素
|
|
|
|
|
+ plt.plot(x_ellipse, y_ellipse, 'b-', linewidth=2)
|
|
|
|
|
+ plt.plot(x, y, 'ko', markersize=8)
|
|
|
|
|
+ plt.plot(P1[0], P1[1], 'ro', markersize=10)
|
|
|
|
|
+ plt.plot(P2[0], P2[1], 'go', markersize=10)
|
|
|
|
|
+
|
|
|
|
|
+ plt.show()
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def arc_to_mask(arc7, shape, line_width=1):
|
|
|
|
|
+ """
|
|
|
|
|
+ Generate a binary mask of an elliptical arc.
|
|
|
|
|
+
|
|
|
|
|
+ Args:
|
|
|
|
|
+ xc, yc (float): 椭圆中心
|
|
|
|
|
+ a, b (float): 长半轴、短半轴 (a >= b)
|
|
|
|
|
+ theta (float): 椭圆旋转角度(**弧度**,逆时针,相对于 x 轴)
|
|
|
|
|
+ phi1, phi2 (float): 起始和终止参数角(**弧度**,在 [0, 2π) 内)
|
|
|
|
|
+ H, W (int): 输出 mask 的高度和宽度
|
|
|
|
|
+ line_width (int): 弧线宽度(像素)
|
|
|
|
|
+
|
|
|
|
|
+ Returns:
|
|
|
|
|
+ mask (Tensor): [H, W], dtype=torch.uint8, 0/255
|
|
|
|
|
+ """
|
|
|
|
|
+ # 确保 phi1 -> phi2 是正向(可处理跨 2π 的情况)
|
|
|
|
|
+ xc, yc, a, b, theta, phi1, phi2 = arc7
|
|
|
|
|
+ H, W = shape
|
|
|
|
|
+ if phi2 < phi1:
|
|
|
|
|
+ phi2 += 2 * np.pi
|
|
|
|
|
+
|
|
|
|
|
+ # 生成参数角(足够密集,避免断线)
|
|
|
|
|
+ num_points = max(int(200 * abs(phi2 - phi1) / (2 * np.pi)), 10)
|
|
|
|
|
+ phi = np.linspace(phi1, phi2, num_points)
|
|
|
|
|
+
|
|
|
|
|
+ # 椭圆参数方程(先在未旋转坐标系下计算)
|
|
|
|
|
+ x_local = a * np.cos(phi)
|
|
|
|
|
+ y_local = b * np.sin(phi)
|
|
|
|
|
+
|
|
|
|
|
+ # 应用旋转和平移
|
|
|
|
|
+ cos_t = np.cos(theta)
|
|
|
|
|
+ sin_t = np.sin(theta)
|
|
|
|
|
+ x_rot = x_local * cos_t - y_local * sin_t + xc
|
|
|
|
|
+ y_rot = x_local * cos_t + y_local * sin_t + yc
|
|
|
|
|
+
|
|
|
|
|
+ # 转为整数坐标(OpenCV 需要 int32)
|
|
|
|
|
+ points = np.stack([x_rot, y_rot], axis=1).astype(np.int32)
|
|
|
|
|
+
|
|
|
|
|
+ # 创建空白图像
|
|
|
|
|
+ img = np.zeros((H, W), dtype=np.uint8)
|
|
|
|
|
+
|
|
|
|
|
+ # 绘制折线(antialias=False 更适合 mask)
|
|
|
|
|
+ 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
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def compute_arc_angles(gt_mask_ends, gt_mask_params):
|
|
|
|
|
+ """
|
|
|
|
|
+ 给定椭圆上的一个点,计算其对应的参数角 phi(弧度)。
|
|
|
|
|
+
|
|
|
|
|
+ Parameters:
|
|
|
|
|
+ point: tuple or array-like, (x, y)
|
|
|
|
|
+ ellipse_param: tuple or array-like, (xc, yc, a, b, theta)
|
|
|
|
|
+
|
|
|
|
|
+ Returns:
|
|
|
|
|
+ phi: float, in [0, 2*pi)
|
|
|
|
|
+ """
|
|
|
|
|
+ results = []
|
|
|
|
|
+ gt_mask_params_tensor = torch.tensor(gt_mask_params,
|
|
|
|
|
+ dtype=gt_mask_ends.dtype,
|
|
|
|
|
+ device=gt_mask_ends.device)
|
|
|
|
|
+ for ends_img, params_img in zip(gt_mask_ends, gt_mask_params_tensor):
|
|
|
|
|
+ # print(f'params_img:{params_img}')
|
|
|
|
|
+ if torch.norm(params_img) < 1e-6: # L2 norm near zero
|
|
|
|
|
+ results.append(torch.zeros(2, device=params_img.device, dtype=params_img.dtype))
|
|
|
|
|
+ continue
|
|
|
|
|
+ x, y = ends_img
|
|
|
|
|
+ xc, yc, a, b, theta = params_img
|
|
|
|
|
+
|
|
|
|
|
+ # 1. 平移到中心
|
|
|
|
|
+ dx = x - xc
|
|
|
|
|
+ dy = y - yc
|
|
|
|
|
+
|
|
|
|
|
+ # 2. 逆旋转(旋转 -theta)
|
|
|
|
|
+ cos_t = torch.cos(theta)
|
|
|
|
|
+ sin_t = torch.sin(theta)
|
|
|
|
|
+ X = dx * cos_t + dy * sin_t
|
|
|
|
|
+ Y = -dx * sin_t + dy * cos_t
|
|
|
|
|
+
|
|
|
|
|
+ # 3. 归一化到单位圆(除以 a, b)
|
|
|
|
|
+ cos_phi = X / a
|
|
|
|
|
+ sin_phi = Y / b
|
|
|
|
|
+
|
|
|
|
|
+ # 4. 用 atan2 求角度(自动处理象限)
|
|
|
|
|
+ phi = torch.atan2(sin_phi, cos_phi)
|
|
|
|
|
+
|
|
|
|
|
+ # 5. 转换到 [0, 2π)
|
|
|
|
|
+ phi = torch.where(phi < 0, phi + 2 * torch.pi, phi)
|
|
|
|
|
+
|
|
|
|
|
+ results.append(phi)
|
|
|
|
|
+ return results
|
|
|
|
|
|
|
|
|
|
|
|
|
def points_to_ellipse(points):
|
|
def points_to_ellipse(points):
|
|
@@ -194,10 +346,8 @@ def points_to_ellipse(points):
|
|
|
pts = points.numpy()
|
|
pts = points.numpy()
|
|
|
pts = pts.reshape(-1, 2)
|
|
pts = pts.reshape(-1, 2)
|
|
|
|
|
|
|
|
- # 计ç®ä¸å¿ç¹
|
|
|
|
|
center = np.mean(pts, axis=0)
|
|
center = np.mean(pts, axis=0)
|
|
|
|
|
|
|
|
- # ä½¿ç¨æå°äºä¹æ³æåæ¤å
|
|
|
|
|
A = np.hstack(
|
|
A = np.hstack(
|
|
|
[pts[:, 0:1] ** 2, pts[:, 0:1] * pts[:, 1:2], pts[:, 1:2] ** 2, pts[:, :2], np.ones((pts.shape[0], 1))])
|
|
[pts[:, 0:1] ** 2, pts[:, 0:1] * pts[:, 1:2], pts[:, 1:2] ** 2, pts[:, :2], np.ones((pts.shape[0], 1))])
|
|
|
b = np.ones(pts.shape[0])
|
|
b = np.ones(pts.shape[0])
|
|
@@ -211,12 +361,10 @@ def points_to_ellipse(points):
|
|
|
major_axis = np.sqrt(numerator / denominator1)
|
|
major_axis = np.sqrt(numerator / denominator1)
|
|
|
minor_axis = np.sqrt(numerator / denominator2)
|
|
minor_axis = np.sqrt(numerator / denominator2)
|
|
|
|
|
|
|
|
- # ç®åå¤çï¼ç´æ¥ç¨ç¹é´è·ç¦»ä¼°ç®é¿çè½´
|
|
|
|
|
distances = np.linalg.norm(pts - center, axis=1)
|
|
distances = np.linalg.norm(pts - center, axis=1)
|
|
|
long_axis_length = np.max(distances) * 2
|
|
long_axis_length = np.max(distances) * 2
|
|
|
short_axis_length = np.min(distances) * 2
|
|
short_axis_length = np.min(distances) * 2
|
|
|
|
|
|
|
|
- # æ¹åå¯ä»¥éè¿ä¸¤ç¹ä¹é´çè¿çº¿æçæ¥è¿ä¼¼è®¡ç®
|
|
|
|
|
orientation = np.arctan2(pts[1, 1] - pts[0, 1], pts[1, 0] - pts[0, 0])
|
|
orientation = np.arctan2(pts[1, 1] - pts[0, 1], pts[1, 0] - pts[0, 0])
|
|
|
|
|
|
|
|
return center[0], center[1], long_axis_length / 2, short_axis_length / 2, orientation
|
|
return center[0], center[1], long_axis_length / 2, short_axis_length / 2, orientation
|
|
@@ -234,11 +382,12 @@ def generate_ellipse_mask(shape, ellipse_params):
|
|
|
img = np.zeros(shape, dtype=np.uint8)
|
|
img = np.zeros(shape, dtype=np.uint8)
|
|
|
cx, cy, rx, ry = int(cx), int(cy), int(rx), int(ry)
|
|
cx, cy, rx, ry = int(cx), int(cy), int(rx), int(ry)
|
|
|
|
|
|
|
|
- # 注æskimageçellipse彿°ä¸ç´æ¥æ¯ææè½¬ï¼æä»¥è¿éç®åå¤ç
|
|
|
|
|
rr, cc = ellipse(cy, cx, ry, rx, shape)
|
|
rr, cc = ellipse(cy, cx, ry, rx, shape)
|
|
|
img[rr, cc] = 1
|
|
img[rr, cc] = 1
|
|
|
|
|
|
|
|
return img
|
|
return img
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
def sort_points_clockwise(points):
|
|
def sort_points_clockwise(points):
|
|
|
points = np.array(points)
|
|
points = np.array(points)
|
|
|
|
|
|
|
@@ -256,22 +405,26 @@ def sort_points_clockwise(points):
|
|
|
sorted_points = points[sorted_indices]
|
|
sorted_points = points[sorted_indices]
|
|
|
|
|
|
|
|
return sorted_points.tolist()
|
|
return sorted_points.tolist()
|
|
|
-def get_boxes_lines(objs,shape):
|
|
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def get_boxes_lines(objs, shape):
|
|
|
boxes = []
|
|
boxes = []
|
|
|
- labels=[]
|
|
|
|
|
- h,w=shape
|
|
|
|
|
|
|
+ labels = []
|
|
|
|
|
+ h, w = shape
|
|
|
line_point_pairs = []
|
|
line_point_pairs = []
|
|
|
- points=[]
|
|
|
|
|
- line_mask=[]
|
|
|
|
|
- circle_4points=[]
|
|
|
|
|
|
|
+ points = []
|
|
|
|
|
+ arc_mask = []
|
|
|
|
|
+ arc_ends = []
|
|
|
|
|
+ arc_params = []
|
|
|
|
|
+ circle_4points = []
|
|
|
|
|
|
|
|
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]无明确大小
|
|
|
|
|
|
|
|
# print(f"points:{obj['points']}")
|
|
# print(f"points:{obj['points']}")
|
|
|
- label=obj['label']
|
|
|
|
|
- if label =='line' or label=='dseam1':
|
|
|
|
|
- a,b=obj['points'][0],obj['points'][1]
|
|
|
|
|
|
|
+ label = obj['label']
|
|
|
|
|
+ if label == 'line' or label == 'dseam1':
|
|
|
|
|
+ a, b = obj['points'][0], obj['points'][1]
|
|
|
|
|
|
|
|
line_point_pairs.append(a)
|
|
line_point_pairs.append(a)
|
|
|
line_point_pairs.append(b)
|
|
line_point_pairs.append(b)
|
|
@@ -281,90 +434,89 @@ def get_boxes_lines(objs,shape):
|
|
|
ymin = max(0, (min(a[1], b[1]) - 6))
|
|
ymin = max(0, (min(a[1], b[1]) - 6))
|
|
|
ymax = min(h, (max(a[1], b[1]) + 6))
|
|
ymax = min(h, (max(a[1], b[1]) + 6))
|
|
|
|
|
|
|
|
- boxes.append([ xmin,ymin, xmax,ymax])
|
|
|
|
|
|
|
+ boxes.append([xmin, ymin, xmax, ymax])
|
|
|
labels.append(torch.tensor(2))
|
|
labels.append(torch.tensor(2))
|
|
|
|
|
|
|
|
- elif label =='point':
|
|
|
|
|
- p= obj['points'][0]
|
|
|
|
|
- xmin=max(0,p[0]-12)
|
|
|
|
|
- xmax = min(w, p[0] +12)
|
|
|
|
|
- ymin=max(0,p[1]-12)
|
|
|
|
|
- ymax = min(h, p[1] + 12)
|
|
|
|
|
-
|
|
|
|
|
- points.append(p)
|
|
|
|
|
- labels.append(torch.tensor(1))
|
|
|
|
|
- boxes.append([xmin, ymin, xmax, ymax])
|
|
|
|
|
-
|
|
|
|
|
-
|
|
|
|
|
|
|
+ elif label == 'point':
|
|
|
|
|
+ p = obj['points'][0]
|
|
|
|
|
+ xmin = max(0, p[0] - 12)
|
|
|
|
|
+ xmax = min(w, p[0] + 12)
|
|
|
|
|
+ ymin = max(0, p[1] - 12)
|
|
|
|
|
+ ymax = min(h, p[1] + 12)
|
|
|
|
|
|
|
|
|
|
+ points.append(p)
|
|
|
|
|
+ labels.append(torch.tensor(1))
|
|
|
|
|
+ boxes.append([xmin, ymin, xmax, ymax])
|
|
|
|
|
|
|
|
- elif label == 'arc' :
|
|
|
|
|
-
|
|
|
|
|
- line_mask.append(obj['points'])
|
|
|
|
|
-
|
|
|
|
|
- xmin = obj['xmin']
|
|
|
|
|
-
|
|
|
|
|
- xmax = obj['xmax']
|
|
|
|
|
|
|
|
|
|
- ymin = obj['ymin']
|
|
|
|
|
|
|
+ elif label == 'arc':
|
|
|
|
|
+ arc_points = obj['points']
|
|
|
|
|
+ params = obj['params']
|
|
|
|
|
+ ends = obj['ends']
|
|
|
|
|
+ arc_ends.append(ends)
|
|
|
|
|
+ arc_params.append(params)
|
|
|
|
|
|
|
|
- ymax = obj['ymax']
|
|
|
|
|
|
|
+ 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])
|
|
boxes.append([xmin, ymin, xmax, ymax])
|
|
|
-
|
|
|
|
|
labels.append(torch.tensor(3))
|
|
labels.append(torch.tensor(3))
|
|
|
|
|
|
|
|
- elif label == 'circle' :
|
|
|
|
|
|
|
+ elif label == 'circle':
|
|
|
# print(f'len circle_4points: {len(obj['points'])}')
|
|
# print(f'len circle_4points: {len(obj['points'])}')
|
|
|
- points=sort_points_clockwise(obj['points'])
|
|
|
|
|
|
|
+ points = sort_points_clockwise(obj['points'])
|
|
|
circle_4points.append(points)
|
|
circle_4points.append(points)
|
|
|
|
|
|
|
|
xmin = max(obj['xmin'] - 40, 0)
|
|
xmin = max(obj['xmin'] - 40, 0)
|
|
|
-
|
|
|
|
|
xmax = min(obj['xmax'] + 40, w)
|
|
xmax = min(obj['xmax'] + 40, w)
|
|
|
-
|
|
|
|
|
ymin = max(obj['ymin'] - 40, 0)
|
|
ymin = max(obj['ymin'] - 40, 0)
|
|
|
-
|
|
|
|
|
ymax = min(obj['ymax'] + 40, h)
|
|
ymax = min(obj['ymax'] + 40, h)
|
|
|
|
|
|
|
|
boxes.append([xmin, ymin, xmax, ymax])
|
|
boxes.append([xmin, ymin, xmax, ymax])
|
|
|
-
|
|
|
|
|
labels.append(torch.tensor(4))
|
|
labels.append(torch.tensor(4))
|
|
|
|
|
|
|
|
- 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)
|
|
|
|
|
- if len(points)==0:
|
|
|
|
|
- points=None
|
|
|
|
|
|
|
+ labels = torch.tensor(labels)
|
|
|
|
|
+ if len(points) == 0:
|
|
|
|
|
+ points = None
|
|
|
else:
|
|
else:
|
|
|
- points=torch.tensor(points,dtype=torch.float32)
|
|
|
|
|
|
|
+ points = torch.tensor(points, dtype=torch.float32)
|
|
|
print(f'read labels:{labels}')
|
|
print(f'read labels:{labels}')
|
|
|
# print(f'read points:{points}')
|
|
# print(f'read points:{points}')
|
|
|
- if len(line_point_pairs)==0:
|
|
|
|
|
- line_point_pairs=None
|
|
|
|
|
|
|
+ if len(line_point_pairs) == 0:
|
|
|
|
|
+ line_point_pairs = None
|
|
|
else:
|
|
else:
|
|
|
- line_point_pairs=torch.tensor(line_point_pairs)
|
|
|
|
|
|
|
+ line_point_pairs = torch.tensor(line_point_pairs)
|
|
|
# print(f'line_point_pairs:{line_point_pairs.shape},{line_point_pairs.dtype}')
|
|
# 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}')
|
|
# print(f'boxes:{boxes.shape},line_point_pairs:{line_point_pairs.shape}')
|
|
|
|
|
|
|
|
- if len(line_mask)==0:
|
|
|
|
|
- line_mask=None
|
|
|
|
|
|
|
+ if len(arc_mask) == 0:
|
|
|
|
|
+ arc_mask = None
|
|
|
|
|
+ else:
|
|
|
|
|
+ arc_mask = torch.tensor(arc_mask, dtype=torch.float32)
|
|
|
|
|
+ print(f'arc_mask shape :{arc_mask.shape},{arc_mask.dtype}')
|
|
|
|
|
+
|
|
|
|
|
+ if len(arc_ends) == 0:
|
|
|
|
|
+ arc_ends = None
|
|
|
else:
|
|
else:
|
|
|
- line_mask=torch.tensor(line_mask,dtype=torch.float32)
|
|
|
|
|
- print(f'arc_mask shape :{line_mask.shape},{line_mask.dtype}')
|
|
|
|
|
|
|
+ arc_ends = torch.tensor(arc_ends, dtype=torch.float32)
|
|
|
|
|
|
|
|
- if len(circle_4points)==0:
|
|
|
|
|
- circle_4points=None
|
|
|
|
|
|
|
+ if len(circle_4points) == 0:
|
|
|
|
|
+ circle_4points = None
|
|
|
else:
|
|
else:
|
|
|
# for circle_4point in circle_4points:
|
|
# 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_4point len111:{len(circle_4point)}')
|
|
|
|
|
+ circle_4points = torch.tensor(circle_4points, dtype=torch.float32)
|
|
|
# print(f'circle_4points shape:{circle_4points.shape}')
|
|
# print(f'circle_4points shape:{circle_4points.shape}')
|
|
|
|
|
|
|
|
- return boxes,line_point_pairs,points,line_mask,circle_4points, labels
|
|
|
|
|
|
|
+ return boxes, line_point_pairs, points, arc_mask, circle_4points, labels, arc_ends, arc_params
|
|
|
|
|
+
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
if __name__ == '__main__':
|
|
|
- path=r'/data/share/zyh/master_dataset/circle/huyan_eclipse/a_dataset'
|
|
|
|
|
- dataset= LineDataset(dataset_path=path, dataset_type='train',augmentation=False, data_type='jpg')
|
|
|
|
|
- dataset.show(9,show_type='all')
|
|
|
|
|
|
|
+ path = r'\\192.168.50.222/share/lm/1112/a_dataset'
|
|
|
|
|
+ dataset = LineDataset(dataset_path=path, dataset_type='train', augmentation=False, data_type='jpg')
|
|
|
|
|
+ dataset.show(9, show_type='arc_masks')
|