line_dataset.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660
  1. import cv2
  2. import imageio
  3. import numpy as np
  4. from skimage.draw import ellipse
  5. from torch.utils.data.dataset import T_co
  6. from libs.vision_libs.utils import draw_keypoints
  7. from models.base.base_dataset import BaseDataset
  8. import json
  9. import os
  10. import PIL
  11. import matplotlib as mpl
  12. from torchvision.utils import draw_bounding_boxes
  13. import torchvision.transforms.v2 as transforms
  14. import torch
  15. import matplotlib.pyplot as plt
  16. from models.base.transforms import get_transforms
  17. from utils.data_process.mask.show_mask import save_full_mask
  18. from utils.data_process.show_prams import print_params
  19. def validate_keypoints(keypoints, image_width, image_height):
  20. for kp in keypoints:
  21. x, y, v = kp
  22. if not (0 <= x < image_width and 0 <= y < image_height):
  23. raise ValueError(f"Key point ({x}, {y}) is out of bounds for image size ({image_width}, {image_height})")
  24. """
  25. 直接读取xanlabel标注的数据集json格式
  26. """
  27. class LineDataset(BaseDataset):
  28. def __init__(self, dataset_path, data_type, transforms=None, augmentation=False, dataset_type=None, img_type='rgb',
  29. target_type='pixel'):
  30. super().__init__(dataset_path)
  31. self.data_path = dataset_path
  32. self.data_type = data_type
  33. print(f'data_path:{dataset_path}')
  34. self.transforms = transforms
  35. self.img_path = os.path.join(dataset_path, "images/" + dataset_type)
  36. self.lbl_path = os.path.join(dataset_path, "labels/" + dataset_type)
  37. self.imgs = os.listdir(self.img_path)
  38. self.lbls = os.listdir(self.lbl_path)
  39. self.target_type = target_type
  40. self.img_type = img_type
  41. self.augmentation = augmentation
  42. print(f'augmentation:{augmentation}')
  43. # self.default_transform = DefaultTransform()
  44. def __getitem__(self, index) -> T_co:
  45. img_path = os.path.join(self.img_path, self.imgs[index])
  46. if self.data_type == 'tiff':
  47. lbl_path = os.path.join(self.lbl_path, self.imgs[index][:-4] + 'json')
  48. img = imageio.v3.imread(img_path)[:, :, 0]
  49. print(f'img shape:{img.shape}')
  50. w, h = img.shape[:2]
  51. img = img.reshape(w, h, 1)
  52. img_3channel = np.zeros((w, h, 3), dtype=img.dtype)
  53. img_3channel[:, :, 2] = img[:, :, 0]
  54. img = torch.from_numpy(img_3channel).permute(2, 1, 0)
  55. else:
  56. lbl_path = os.path.join(self.lbl_path, self.imgs[index][:-3] + 'json')
  57. img = PIL.Image.open(img_path).convert('RGB')
  58. w, h = img.size
  59. # wire_labels, target = self.read_target(item=index, lbl_path=lbl_path, shape=(h, w))
  60. print(img_path)
  61. print_params(img)
  62. target = self.read_target(item=index, lbl_path=lbl_path, shape=(h, w),image=img)
  63. self.transforms = get_transforms(augmention=self.augmentation)
  64. img, target = self.transforms(img, target)
  65. return img, target
  66. def __len__(self):
  67. return len(self.imgs)
  68. def read_target(self, item, lbl_path, shape, extra=None,image=None):
  69. # print(f'shape:{shape}')
  70. # print(f'lbl_path:{lbl_path}')
  71. with open(lbl_path, 'r') as file:
  72. lable_all = json.load(file)
  73. objs = lable_all["shapes"]
  74. point_pairs = objs[0]['points']
  75. # print(f'point_pairs:{point_pairs}')
  76. target = {}
  77. target["image_id"] = torch.tensor(item)
  78. boxes, lines, points, labels, arc_ends, arc_params = get_boxes_lines(objs, shape)
  79. if lines is not None:
  80. label_3d = labels.view(-1, 1, 1).expand(-1, 2, -1) # [N] -> [N,2,1]
  81. line1 = torch.cat([lines, label_3d], dim=-1) # [N,2,3]
  82. target["lines"] = line1.to(torch.float32)
  83. if arc_ends is not None:
  84. target['mask_ends'] = arc_ends
  85. if arc_params is not None:
  86. target['mask_params'] = arc_params
  87. # arc_angles = compute_arc_angles(arc_ends, arc_params)
  88. arc_masks = []
  89. for i in range(len(arc_params)):
  90. mask = arc_to_mask_safe(arc_params[i], arc_ends[i], shape=(2000, 2000),debug=False)
  91. arc_masks.append(mask)
  92. target['circle_masks'] = torch.stack(arc_masks, dim=0)
  93. print(f'target[circle_masks]:{target["circle_masks"].shape}')
  94. target["boxes"] = boxes
  95. target["labels"] = labels
  96. target["img_size"] = shape
  97. return target
  98. def show(self, idx, show_type='all'):
  99. image, target = self.__getitem__(idx)
  100. cmap = plt.get_cmap("jet")
  101. norm = mpl.colors.Normalize(vmin=0.4, vmax=1.0)
  102. sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
  103. sm.set_array([])
  104. img = image
  105. # print(f'img:{img.shape}')
  106. if show_type == 'Original_mask':
  107. image = image.permute(1, 2, 0).cpu().numpy() # (3,H,W) -> (H,W,3)
  108. arc_mask = target['circle_masks'][-1]
  109. plt.figure(figsize=(10, 5))
  110. plt.subplot(1, 2, 1)
  111. plt.imshow(image)
  112. plt.title('Original')
  113. # 叠加 mask(用红色轮廓或填充)
  114. plt.subplot(1, 2, 2)
  115. plt.imshow(image)
  116. plt.imshow(arc_mask, cmap='Reds', alpha=0.5) # 半透明红色
  117. plt.title('Image with Mask')
  118. plt.show()
  119. if show_type == 'arc_yuan_point_ellipse':
  120. arc_ends = target['mask_ends']
  121. arc_params = target['mask_params']
  122. fig, ax = plt.subplots()
  123. ax.imshow(img.permute(1, 2, 0))
  124. for params in arc_params:
  125. if torch.all(params == 0):
  126. continue
  127. x, y, a, b, q = params
  128. theta = np.radians(q)
  129. phi = np.linspace(0, 2 * np.pi, 500)
  130. x_ellipse = x + a * np.cos(phi) * np.cos(theta) - b * np.sin(phi) * np.sin(theta)
  131. y_ellipse = y + a * np.cos(phi) * np.sin(theta) + b * np.sin(phi) * np.cos(theta)
  132. plt.plot(x_ellipse, y_ellipse, 'b-', linewidth=2)
  133. for point2 in arc_ends:
  134. if torch.all(point2 == 0):
  135. continue
  136. ends_np = point2.cpu().numpy()
  137. ax.plot(ends_np[:, 0], ends_np[:, 1], 'ro', markersize=6, label='Arc Endpoints')
  138. ax.legend()
  139. plt.axis('image') # 保持比例一致
  140. plt.show()
  141. if show_type == 'circle_masks':
  142. arc_mask = target['circle_masks']
  143. print(f'target circle_masks:{arc_mask.shape}')
  144. plt.imshow(arc_mask[-1])
  145. plt.show()
  146. if show_type == 'circle_masks11':
  147. boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), target["boxes"],
  148. colors="yellow", width=1)
  149. circle = target['circles']
  150. circle_mask = target['circle_masks']
  151. print(f'taget circle:{circle.shape}')
  152. print(f'target circle_masks:{circle_mask.shape}')
  153. plt.imshow(circle_mask.squeeze(0))
  154. keypoint_img = draw_keypoints(boxed_image, circle, colors='red', width=3)
  155. # plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
  156. plt.show()
  157. # if show_type=='lines':
  158. # keypoint_img=draw_keypoints((img * 255).to(torch.uint8),target['lines'],colors='red',width=3)
  159. # plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
  160. # plt.show()
  161. if show_type == 'points':
  162. # print(f'points:{target['points'].shape}')
  163. keypoint_img = draw_keypoints((img * 255).to(torch.uint8), target['points'].unsqueeze(1), colors='red',
  164. width=3)
  165. plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
  166. plt.show()
  167. if show_type == 'boxes':
  168. boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), target["boxes"],
  169. colors="yellow", width=1)
  170. plt.imshow(boxed_image.permute(1, 2, 0).numpy())
  171. plt.show()
  172. def show_img(self, img_path):
  173. pass
  174. import torch
  175. import numpy as np
  176. import cv2
  177. def arc_to_mask_safe(arc_param, arc_end, shape, line_width=5, debug=True, idx=-1):
  178. """
  179. Generate a mask for a small (<180 degree) arc based on arc parameters and endpoints.
  180. Args:
  181. arc_param: torch.Tensor of shape (5,) - [cx, cy, a, b, theta]
  182. arc_end: torch.Tensor of shape (2,2) - [[x1,y1],[x2,y2]]
  183. shape: tuple (H,W) - mask size
  184. line_width: thickness of the arc
  185. debug: bool - if True, print debug info
  186. idx: int or str - index for debugging identification
  187. Returns:
  188. mask: torch.Tensor of shape (H,W)
  189. """
  190. # ------------------ Check for all-zero input ------------------
  191. if torch.all(arc_param == 0) or torch.all(arc_end == 0):
  192. if debug:
  193. print(f"[{idx}] Warning: arc_param or arc_end all zeros. Returning zero mask.")
  194. print(f"[{idx}] arc_param: {arc_param.tolist()}")
  195. print(f"[{idx}] arc_end: {arc_end.tolist()}")
  196. return torch.zeros(shape, dtype=torch.float32)
  197. cx, cy, a, b, theta = arc_param.tolist()
  198. if a <= 0 or b <= 0:
  199. if debug:
  200. print(f"[{idx}] Warning: invalid ellipse axes a={a}, b={b}. Returning zero mask.")
  201. print(f"[{idx}] arc_param: {arc_param.tolist()}")
  202. print(f"[{idx}] arc_end: {arc_end.tolist()}")
  203. return torch.zeros(shape, dtype=torch.float32)
  204. x1, y1 = arc_end[0].tolist()
  205. x2, y2 = arc_end[1].tolist()
  206. cos_t = np.cos(theta)
  207. sin_t = np.sin(theta)
  208. def point_to_angle(x, y):
  209. dx = x - cx
  210. dy = y - cy
  211. x_ = cos_t * dx + sin_t * dy
  212. y_ = -sin_t * dx + cos_t * dy
  213. return np.arctan2(y_ / b, x_ / a)
  214. try:
  215. angle1 = point_to_angle(x1, y1)
  216. angle2 = point_to_angle(x2, y2)
  217. except Exception as e:
  218. if debug:
  219. print(f"[{idx}] Exception in point_to_angle: {e}")
  220. print(f"[{idx}] arc_param: {arc_param.tolist()}, arc_end: {arc_end.tolist()}")
  221. return torch.zeros(shape, dtype=torch.float32)
  222. if np.isnan(angle1) or np.isnan(angle2):
  223. if debug:
  224. print(f"[{idx}] Warning: angle1 or angle2 is NaN. Returning zero mask.")
  225. print(f"[{idx}] arc_param: {arc_param.tolist()}, arc_end: {arc_end.tolist()}")
  226. return torch.zeros(shape, dtype=torch.float32)
  227. # Ensure small arc (<180 degrees)
  228. if angle2 < angle1:
  229. angle2 += 2 * np.pi
  230. if angle2 - angle1 > np.pi:
  231. angle1, angle2 = angle2, angle1 + 2 * np.pi
  232. angles = np.linspace(angle1, angle2, 100)
  233. xs = cx + a * np.cos(angles) * cos_t - b * np.sin(angles) * sin_t
  234. ys = cy + a * np.cos(angles) * sin_t + b * np.sin(angles) * cos_t
  235. xs = np.nan_to_num(xs, nan=0.0).astype(np.int32)
  236. ys = np.nan_to_num(ys, nan=0.0).astype(np.int32)
  237. # ------------------ Debug prints ------------------
  238. if debug:
  239. print(f"[{idx}] arc_param: {arc_param.tolist()}")
  240. print(f"[{idx}] arc_end: {arc_end.tolist()}")
  241. print(f"[{idx}] xs[:5], ys[:5]: {xs[:5]}, {ys[:5]}")
  242. mask = np.zeros(shape, dtype=np.uint8)
  243. pts = np.stack([xs, ys], axis=1)
  244. # Draw the arc with given line_width
  245. for i in range(len(pts) - 1):
  246. cv2.line(mask, tuple(pts[i]), tuple(pts[i + 1]), color=1, thickness=line_width)
  247. # ------------------ Extra check for non-zero mask ------------------
  248. if debug:
  249. mask_sum = mask.sum()
  250. if mask_sum == 0:
  251. print(f"[{idx}] Warning: mask generated is all zeros!")
  252. else:
  253. print(f"[{idx}] mask sum: {mask_sum}")
  254. return torch.tensor(mask, dtype=torch.float32)
  255. def draw_el(all):
  256. # 解析椭圆参数
  257. if isinstance(all, torch.Tensor):
  258. all = all.cpu().numpy()
  259. x, y, a, b, q, q1, q2 = all
  260. theta = np.radians(q)
  261. phi1 = np.radians(q1) # 第一个点的参数角
  262. phi2 = np.radians(q2) # 第二个点的参数角
  263. # 生成椭圆上的点
  264. phi = np.linspace(0, 2 * np.pi, 500)
  265. x_ellipse = x + a * np.cos(phi) * np.cos(theta) - b * np.sin(phi) * np.sin(theta)
  266. y_ellipse = y + a * np.cos(phi) * np.sin(theta) + b * np.sin(phi) * np.cos(theta)
  267. # 计算两个指定点的坐标
  268. def param_to_point(phi, xc, yc, a, b, theta):
  269. x = xc + a * np.cos(phi) * np.cos(theta) - b * np.sin(phi) * np.sin(theta)
  270. y = yc + a * np.cos(phi) * np.sin(theta) + b * np.sin(phi) * np.cos(theta)
  271. return x, y
  272. P1 = param_to_point(phi1, x, y, a, b, theta)
  273. P2 = param_to_point(phi2, x, y, a, b, theta)
  274. # 创建画布并显示背景图片(使用传入的background_img,shape为[H, W, C])
  275. plt.figure(figsize=(10, 10))
  276. # plt.imshow(background_img) # 直接显示背景图
  277. # 绘制椭圆及相关元素
  278. plt.plot(x_ellipse, y_ellipse, 'b-', linewidth=2)
  279. plt.plot(x, y, 'ko', markersize=8)
  280. plt.plot(P1[0], P1[1], 'ro', markersize=10)
  281. plt.plot(P2[0], P2[1], 'go', markersize=10)
  282. plt.show()
  283. def arc_to_mask(arc7, shape, line_width=1):
  284. """
  285. Generate a binary mask of an elliptical arc.
  286. Args:
  287. xc, yc (float): 椭圆中心
  288. a, b (float): 长半轴、短半轴 (a >= b)
  289. theta (float): 椭圆旋转角度(**弧度**,逆时针,相对于 x 轴)
  290. phi1, phi2 (float): 起始和终止参数角(**弧度**,在 [0, 2π) 内)
  291. H, W (int): 输出 mask 的高度和宽度
  292. line_width (int): 弧线宽度(像素)
  293. Returns:
  294. mask (Tensor): [H, W], dtype=torch.uint8, 0/255
  295. """
  296. # print_params(arc7)
  297. # 确保 phi1 -> phi2 是正向(可处理跨 2π 的情况)
  298. if torch.all(arc7 == 0):
  299. return torch.zeros(shape, dtype=torch.uint8)
  300. xc, yc, a, b, theta, phi1, phi2 = arc7
  301. H, W = shape
  302. if phi2 < phi1:
  303. phi2 += 2 * np.pi
  304. # 生成参数角(足够密集,避免断线)
  305. num_points = max(int(200 * abs(phi2 - phi1) / (2 * np.pi)), 10)
  306. phi = np.linspace(phi1, phi2, num_points)
  307. # 椭圆参数方程(先在未旋转坐标系下计算)
  308. x_local = a * np.cos(phi)
  309. y_local = b * np.sin(phi)
  310. # 应用旋转和平移
  311. cos_t = np.cos(theta)
  312. sin_t = np.sin(theta)
  313. x_rot = x_local * cos_t - y_local * sin_t + xc
  314. y_rot = x_local * cos_t + y_local * sin_t + yc
  315. # 转为整数坐标(OpenCV 需要 int32)
  316. points = np.stack([x_rot, y_rot], axis=1).astype(np.int32)
  317. # 创建空白图像
  318. img = np.zeros((H, W), dtype=np.uint8)
  319. # 绘制折线(antialias=False 更适合 mask)
  320. cv2.polylines(img, [points], isClosed=False, color=255, thickness=line_width, lineType=cv2.LINE_AA)
  321. return torch.from_numpy(img).float() # [H, W], values: 0 or 255
  322. def compute_arc_angles(gt_mask_ends, gt_mask_params):
  323. """
  324. 给定椭圆上的一个点,计算其对应的参数角 phi(弧度)。
  325. Parameters:
  326. point: tuple or array-like, (x, y)
  327. ellipse_param: tuple or array-like, (xc, yc, a, b, theta)
  328. Returns:
  329. phi: float, in [0, 2*pi)
  330. """
  331. # print_params(gt_mask_ends, gt_mask_params)
  332. results = []
  333. if not isinstance(gt_mask_params, torch.Tensor):
  334. gt_mask_params_tensor = torch.tensor(gt_mask_params, dtype=gt_mask_ends.dtype, device=gt_mask_ends.device)
  335. else:
  336. gt_mask_params_tensor = gt_mask_params.clone().detach().to(gt_mask_ends)
  337. for ends_img, params_img in zip(gt_mask_ends, gt_mask_params_tensor):
  338. # print(f'params_img:{params_img}')
  339. if torch.norm(params_img) < 1e-6: # L2 norm near zero
  340. results.append(torch.zeros(2, device=params_img.device, dtype=params_img.dtype))
  341. continue
  342. x, y = ends_img
  343. xc, yc, a, b, theta = params_img
  344. # 1. 平移到中心
  345. dx = x - xc
  346. dy = y - yc
  347. # 2. 逆旋转(旋转 -theta)
  348. cos_t = torch.cos(theta)
  349. sin_t = torch.sin(theta)
  350. X = dx * cos_t + dy * sin_t
  351. Y = -dx * sin_t + dy * cos_t
  352. # 3. 归一化到单位圆(除以 a, b)
  353. cos_phi = X / a
  354. sin_phi = Y / b
  355. # 4. 用 atan2 求角度(自动处理象限)
  356. phi = torch.atan2(sin_phi, cos_phi)
  357. # 5. 转换到 [0, 2π)
  358. phi = torch.where(phi < 0, phi + 2 * torch.pi, phi)
  359. results.append(phi)
  360. return results
  361. def points_to_ellipse(points):
  362. """
  363. 根据提供的四个点估计椭圆参数。
  364. :param points: Tensor of shape (4, 2) 表示椭圆上的四个点
  365. :return: 返回 (cx, cy, r1, r2, orientation) 其中 cx, cy 是中心坐标,r1, r2 分别是长轴和短轴半径,orientation 是椭圆的方向(弧度)
  366. """
  367. # 转换为numpy数组进行计算
  368. pts = points.numpy()
  369. pts = pts.reshape(-1, 2)
  370. center = np.mean(pts, axis=0)
  371. A = np.hstack(
  372. [pts[:, 0:1] ** 2, pts[:, 0:1] * pts[:, 1:2], pts[:, 1:2] ** 2, pts[:, :2], np.ones((pts.shape[0], 1))])
  373. b = np.ones(pts.shape[0])
  374. x = np.linalg.lstsq(A, b, rcond=None)[0]
  375. # 解析解参见 https://en.wikipedia.org/wiki/Ellipse#General_ellipse
  376. a, b, c, d, f, g = x.ravel()
  377. numerator = 2 * (a * f * f + c * d * d + g * b * b - 2 * b * d * f - a * c * g)
  378. denominator1 = (b * b - a * c) * ((c - a) * np.sqrt(1 + 4 * b * b / ((a - c) * (a - c))) - (c + a))
  379. denominator2 = (b * b - a * c) * ((a - c) * np.sqrt(1 + 4 * b * b / ((a - c) * (a - c))) - (c + a))
  380. major_axis = np.sqrt(numerator / denominator1)
  381. minor_axis = np.sqrt(numerator / denominator2)
  382. distances = np.linalg.norm(pts - center, axis=1)
  383. long_axis_length = np.max(distances) * 2
  384. short_axis_length = np.min(distances) * 2
  385. orientation = np.arctan2(pts[1, 1] - pts[0, 1], pts[1, 0] - pts[0, 0])
  386. return center[0], center[1], long_axis_length / 2, short_axis_length / 2, orientation
  387. def generate_ellipse_mask(shape, ellipse_params):
  388. """
  389. 在指定形状的图像上生成椭圆mask。
  390. :param shape: 输出mask的形状 (HxW)
  391. :param ellipse_params: 椭圆参数 (cx, cy, rx, ry, orientation)
  392. :return: 椭圆mask
  393. """
  394. cx, cy, rx, ry, orientation = ellipse_params
  395. img = np.zeros(shape, dtype=np.uint8)
  396. cx, cy, rx, ry = int(cx), int(cy), int(rx), int(ry)
  397. rr, cc = ellipse(cy, cx, ry, rx, shape)
  398. img[rr, cc] = 1
  399. return img
  400. def sort_points_clockwise(points):
  401. points = np.array(points)
  402. top_left_idx = np.lexsort((points[:, 0], points[:, 1]))[0]
  403. reference_point = points[top_left_idx]
  404. def angle_to_reference(point):
  405. return np.arctan2(point[1] - reference_point[1], point[0] - reference_point[0])
  406. angles = np.apply_along_axis(angle_to_reference, 1, points)
  407. angles[angles < 0] += 2 * np.pi
  408. sorted_indices = np.argsort(angles)
  409. sorted_points = points[sorted_indices]
  410. return sorted_points.tolist()
  411. def get_boxes_lines(objs, shape):
  412. boxes = []
  413. labels = []
  414. h, w = shape
  415. line_point_pairs = []
  416. points = []
  417. mask_ends = []
  418. mask_params = []
  419. for obj in objs:
  420. label = obj['label']
  421. # if label == 'line' or label == 'dseam1':
  422. # a, b = obj['points'][0], obj['points'][1]
  423. # line_point_pairs.append([a, b])
  424. #
  425. # xmin = max(0, (min(a[0], b[0]) - 6))
  426. # xmax = min(w, (max(a[0], b[0]) + 6))
  427. # ymin = max(0, (min(a[1], b[1]) - 6))
  428. # ymax = min(h, (max(a[1], b[1]) + 6))
  429. #
  430. # boxes.append([xmin, ymin, xmax, ymax])
  431. # labels.append(torch.tensor(2))
  432. #
  433. # points.append(torch.tensor([0.0]))
  434. # mask_ends.append([[0, 0], [0, 0]])
  435. # mask_params.append([0, 0, 0, 0, 0])
  436. # # circle_4points.append([[0, 0], [0, 0], [0, 0], [0, 0]])
  437. #
  438. #
  439. # elif label == 'point':
  440. # p = obj['points'][0]
  441. # xmin = max(0, p[0] - 12)
  442. # xmax = min(w, p[0] + 12)
  443. # ymin = max(0, p[1] - 12)
  444. # ymax = min(h, p[1] + 12)
  445. #
  446. # points.append(p)
  447. # labels.append(torch.tensor(1))
  448. # boxes.append([xmin, ymin, xmax, ymax])
  449. #
  450. # line_point_pairs.append([[0, 0], [0, 0]])
  451. # mask_ends.append([[0, 0], [0, 0]])
  452. # mask_params.append([0, 0, 0, 0, 0])
  453. # # circle_4points.append([[0, 0], [0, 0], [0, 0], [0, 0]])
  454. # elif label == 'arc':
  455. # arc_points = obj['points']
  456. # arc_params = obj['params']
  457. # arc_ends = obj['ends']
  458. # line_mask.append(arc_points)
  459. # mask_ends.append(arc_ends)
  460. # mask_params.append(arc_params)
  461. #
  462. # xs = [p[0] for p in arc_points]
  463. # ys = [p[1] for p in arc_points]
  464. # xmin, xmax = min(xs), max(xs)
  465. # ymin, ymax = min(ys), max(ys)
  466. #
  467. # boxes.append([xmin, ymin, xmax, ymax])
  468. # labels.append(torch.tensor(3))
  469. #
  470. # points.append(torch.tensor([0.0]))
  471. # line_point_pairs.append([[0, 0], [0, 0]])
  472. # circle_4points.append([[0, 0], [0, 0], [0, 0], [0, 0]])
  473. if label == 'arc':
  474. arc_params = obj['params']
  475. arc_ends = obj['ends']
  476. mask_ends.append(arc_ends)
  477. mask_params.append(arc_params)
  478. arc3points = obj['points']
  479. xs = [p[0] for p in arc3points]
  480. ys = [p[1] for p in arc3points]
  481. xmin_raw = min(xs)
  482. xmax_raw = max(xs)
  483. ymin_raw = min(ys)
  484. ymax_raw = max(ys)
  485. xmin = max(xmin_raw - 40, 0)
  486. xmax = min(xmax_raw + 40, w)
  487. ymin = max(ymin_raw - 40, 0)
  488. ymax = min(ymax_raw + 40, h)
  489. boxes.append([xmin, ymin, xmax, ymax])
  490. labels.append(torch.tensor(4))
  491. points.append(torch.tensor([0.0]))
  492. line_point_pairs.append([[0, 0], [0, 0]])
  493. boxes = torch.tensor(boxes, dtype=torch.float32)
  494. print(f'boxes:{boxes.shape}')
  495. labels = torch.tensor(labels)
  496. if points:
  497. points = torch.tensor(points, dtype=torch.float32)
  498. else:
  499. points = None
  500. if line_point_pairs:
  501. line_point_pairs = torch.tensor(line_point_pairs, dtype=torch.float32)
  502. else:
  503. line_point_pairs = None
  504. if mask_ends:
  505. mask_ends = torch.tensor(mask_ends, dtype=torch.float32)
  506. else:
  507. mask_ends = None
  508. if mask_params:
  509. mask_params = torch.tensor(mask_params, dtype=torch.float32)
  510. else:
  511. mask_params = None
  512. return boxes, line_point_pairs, points, labels, mask_ends, mask_params
  513. if __name__ == '__main__':
  514. path = r'/data/share/zyh/master_dataset/dataset_net/pokou_251115_251121/a_dataset'
  515. dataset = LineDataset(dataset_path=path, dataset_type='train', augmentation=False, data_type='jpg')
  516. for i in range(100):
  517. dataset.show(i, show_type='Original_mask')