pokouqiege/models/line_detect/line_dataset.py

import imageio
import numpy as np
from skimage.draw import ellipse
from torch.utils.data.dataset import T_co

from libs.vision_libs.utils import draw_keypoints
from models.base.base_dataset import BaseDataset

import json
import os
import PIL
import matplotlib as mpl
from torchvision.utils import draw_bounding_boxes
import torchvision.transforms.v2 as transforms
import torch

import matplotlib.pyplot as plt
from models.base.transforms import get_transforms


def validate_keypoints(keypoints, image_width, image_height):
    for kp in keypoints:
        x, y, v = kp
        if not (0 <= x < image_width and 0 <= y < image_height):
            raise ValueError(f"Key point ({x}, {y}) is out of bounds for image size ({image_width}, {image_height})")



"""
直接读取xanlabel标注的数据集json格式

"""
class LineDataset(BaseDataset):
    def __init__(self, dataset_path, data_type, transforms=None,augmentation=False, dataset_type=None,img_type='rgb', target_type='pixel'):
        super().__init__(dataset_path)

        self.data_path = dataset_path
        self.data_type = data_type
        print(f'data_path:{dataset_path}')
        self.transforms = transforms
        self.img_path = os.path.join(dataset_path, "images/" + dataset_type)
        self.lbl_path = os.path.join(dataset_path, "labels/" + dataset_type)
        self.imgs = os.listdir(self.img_path)
        self.lbls = os.listdir(self.lbl_path)
        self.target_type = target_type
        self.img_type=img_type
        self.augmentation=augmentation
        print(f'augmentation:{augmentation}')
        # self.default_transform = DefaultTransform()

    def __getitem__(self, index) -> T_co:
        img_path = os.path.join(self.img_path, self.imgs[index])

        if self.data_type == 'tiff':
            lbl_path = os.path.join(self.lbl_path, self.imgs[index][:-4] + 'json')
            img = imageio.v3.imread(img_path)[:,:,0]
            print(f'img shape:{img.shape}')
            w, h = img.shape[:2]
            img=img.reshape(w,h,1)
            img_3channel = np.zeros((w, h, 3), dtype=img.dtype)
            img_3channel[:, :, 2] = img[:, :, 0]


            img = torch.from_numpy(img_3channel).permute(2, 1, 0)
        else:
            lbl_path = os.path.join(self.lbl_path, self.imgs[index][:-3] + 'json')
            img = PIL.Image.open(img_path).convert('RGB')
            w, h = img.size
        # 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))


        self.transforms=get_transforms(augmention=self.augmentation)

        img, target = self.transforms(img, target)


        return img, target

    def __len__(self):
        return len(self.imgs)

    def read_target(self, item, lbl_path, shape, extra=None):
        # print(f'shape:{shape}')
        # print(f'lbl_path:{lbl_path}')
        with open(lbl_path, 'r') as file:
            lable_all = json.load(file)


        objs = lable_all["shapes"]
        point_pairs=objs[0]['points']


        # print(f'point_pairs:{point_pairs}')
        target = {}

        target["image_id"] = torch.tensor(item)
        boxes, lines, points, arc_mask,circle_4points,labels = get_boxes_lines(objs, shape)


        if points is not None:
            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)
            # 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')

        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["labels"]=labels
        # target["boxes"], lines,target["points"], target["labels"] = get_boxes_lines(objs,shape)
        # print(f'lines:{lines}')
        # target["labels"] = torch.ones(len(target["boxes"]), dtype=torch.int64)
        # print(f'target points:{target["points"]}')



        # target["lines"] = lines.to(torch.float32).view(-1,2,3)

        # print(f'')

        # print(f'lines:{target["lines"].shape}')
        target["img_size"]=shape

        # validate_keypoints(lines, shape[0], shape[1])
        return target

    def show(self, idx,show_type='all'):
        image, target = self.__getitem__(idx)

        cmap = plt.get_cmap("jet")
        norm = mpl.colors.Normalize(vmin=0.4, vmax=1.0)
        sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
        sm.set_array([])

        # img_path = os.path.join(self.img_path, self.imgs[idx])
        # print(f'boxes:{target["boxes"]}')
        img = image
        if show_type=='all':
            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'target circle_masks:{circle_mask.shape}')
            plt.imshow(circle_mask.squeeze(0))
            keypoint_img=draw_keypoints(boxed_image,circle,colors='red',width=3)
            # plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
            plt.show()

        # if show_type=='lines':
        #     keypoint_img=draw_keypoints((img * 255).to(torch.uint8),target['lines'],colors='red',width=3)
        #     plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
        #     plt.show()

        if show_type=='points':
            # print(f'points:{target['points'].shape}')
            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.show()

        if show_type=='boxes':
            boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), target["boxes"],
                                              colors="yellow", width=1)
            plt.imshow(boxed_image.permute(1, 2, 0).numpy())
            plt.show()




    def show_img(self, img_path):
        pass


def points_to_ellipse(points):
    """
    根据提供的四个点估计椭圆参数。

    :param points: Tensor of shape (4, 2) 表示椭圆上的四个点
    :return: 返回 (cx, cy, r1, r2, orientation) 其中 cx, cy 是中心坐标，r1, r2 分别是长轴和短轴半径，orientation 是椭圆的方向（弧度）
    """
    # 转换为numpy数组进行计算
    pts = points.numpy()
    pts = pts.reshape(-1, 2)

    # 计算中心点
    center = np.mean(pts, axis=0)

    # 使用最小二乘法拟合椭圆
    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))])
    b = np.ones(pts.shape[0])
    x = np.linalg.lstsq(A, b, rcond=None)[0]

    # 解析解参见 https://en.wikipedia.org/wiki/Ellipse#General_ellipse
    a, b, c, d, f, g = x.ravel()
    numerator = 2 * (a * f * f + c * d * d + g * b * b - 2 * b * d * f - a * c * g)
    denominator1 = (b * b - a * c) * ((c - a) * np.sqrt(1 + 4 * b * b / ((a - c) * (a - c))) - (c + a))
    denominator2 = (b * b - a * c) * ((a - c) * np.sqrt(1 + 4 * b * b / ((a - c) * (a - c))) - (c + a))
    major_axis = np.sqrt(numerator / denominator1)
    minor_axis = np.sqrt(numerator / denominator2)

    # 简化处理：直接用点间距离估算长短轴
    distances = np.linalg.norm(pts - center, axis=1)
    long_axis_length = np.max(distances) * 2
    short_axis_length = np.min(distances) * 2

    # 方向可以通过两点之间的连线斜率来近似计算
    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


def generate_ellipse_mask(shape, ellipse_params):
    """
    在指定形状的图像上生成椭圆mask。

    :param shape: 输出mask的形状 (HxW)
    :param ellipse_params: 椭圆参数 (cx, cy, rx, ry, orientation)
    :return: 椭圆mask
    """
    cx, cy, rx, ry, orientation = ellipse_params
    img = np.zeros(shape, dtype=np.uint8)
    cx, cy, rx, ry = int(cx), int(cy), int(rx), int(ry)

    # 注意skimage的ellipse函数不直接支持旋转，所以这里简化处理
    rr, cc = ellipse(cy, cx, ry, rx, shape)
    img[rr, cc] = 1

    return img
def sort_points_clockwise(points):
    points = np.array(points)

    top_left_idx = np.lexsort((points[:, 0], points[:, 1]))[0]
    reference_point = points[top_left_idx]

    def angle_to_reference(point):
        return np.arctan2(point[1] - reference_point[1], point[0] - reference_point[0])

    angles = np.apply_along_axis(angle_to_reference, 1, points)

    angles[angles < 0] += 2 * np.pi

    sorted_indices = np.argsort(angles)
    sorted_points = points[sorted_indices]

    return sorted_points.tolist()
def get_boxes_lines(objs,shape):
    boxes = []
    labels=[]
    h,w=shape
    line_point_pairs = []
    points=[]
    line_mask=[]
    circle_4points=[]

    for obj in objs:
        # plt.plot([a[1], b[1]], [a[0], b[0]], c="red", linewidth=1)  # a[1], b[1]无明确大小

        # print(f"points:{obj['points']}")
        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(b)

            xmin = max(0, (min(a[0], b[0]) - 6))
            xmax = min(w, (max(a[0], b[0]) + 6))
            ymin = max(0, (min(a[1], b[1]) - 6))
            ymax = min(h, (max(a[1], b[1]) + 6))

            boxes.append([ xmin,ymin,  xmax,ymax])
            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 == 'arc' :

            line_mask.append(obj['points'])

            xmin = obj['xmin']

            xmax = obj['xmax']

            ymin = obj['ymin']

            ymax = obj['ymax']

            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)

            boxes.append([xmin, ymin, xmax, ymax])

            labels.append(torch.tensor(4))

    boxes=torch.tensor(boxes,dtype=torch.float32)
    print(f'boxes:{boxes.shape}')
    labels=torch.tensor(labels)
    if len(points)==0:
        points=None
    else:
        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:
        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}')

    if len(line_mask)==0:
        line_mask=None
    else:
        line_mask=torch.tensor(line_mask,dtype=torch.float32)
        print(f'arc_mask shape :{line_mask.shape},{line_mask.dtype}')

    if len(circle_4points)==0:
        circle_4points=None
    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}')

    return boxes,line_point_pairs,points,line_mask,circle_4points, labels

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')