lcnn/lcnn/datasets.py

# import glob
# import json
# import math
# import os
# import random
#
# import numpy as np
# import numpy.linalg as LA
# import torch
# from skimage import io
# from torch.utils.data import Dataset
# from torch.utils.data.dataloader import default_collate
#
# from lcnn.config import M
#
# from .dataset_tool import line_boxes_faster, read_masks_from_txt_wire, read_masks_from_pixels_wire
#
#
# class WireframeDataset(Dataset):
#     def __init__(self, rootdir, split):
#         self.rootdir = rootdir
#         filelist = glob.glob(f"{rootdir}/{split}/*_label.npz")
#         filelist.sort()
#
#         # print(f"n{split}:", len(filelist))
#         self.split = split
#         self.filelist = filelist
#
#     def __len__(self):
#         return len(self.filelist)
#
#     def __getitem__(self, idx):
#         iname = self.filelist[idx][:-10].replace("_a0", "").replace("_a1", "") + ".png"
#         image = io.imread(iname).astype(float)[:, :, :3]
#         if "a1" in self.filelist[idx]:
#             image = image[:, ::-1, :]
#         image = (image - M.image.mean) / M.image.stddev
#         image = np.rollaxis(image, 2).copy()
#
#         with np.load(self.filelist[idx]) as npz:
#             target = {
#                 name: torch.from_numpy(npz[name]).float()
#                 for name in ["jmap", "joff", "lmap"]
#             }
#             lpos = np.random.permutation(npz["lpos"])[: M.n_stc_posl]
#             lneg = np.random.permutation(npz["lneg"])[: M.n_stc_negl]
#             npos, nneg = len(lpos), len(lneg)
#             lpre = np.concatenate([lpos, lneg], 0)
#             for i in range(len(lpre)):
#                 if random.random() > 0.5:
#                     lpre[i] = lpre[i, ::-1]
#             ldir = lpre[:, 0, :2] - lpre[:, 1, :2]
#             ldir /= np.clip(LA.norm(ldir, axis=1, keepdims=True), 1e-6, None)
#             feat = [
#                 lpre[:, :, :2].reshape(-1, 4) / 128 * M.use_cood,
#                 ldir * M.use_slop,
#                 lpre[:, :, 2],
#             ]
#             feat = np.concatenate(feat, 1)
#             meta = {
#                 "junc": torch.from_numpy(npz["junc"][:, :2]),
#                 "jtyp": torch.from_numpy(npz["junc"][:, 2]).byte(),
#                 "Lpos": self.adjacency_matrix(len(npz["junc"]), npz["Lpos"]),
#                 "Lneg": self.adjacency_matrix(len(npz["junc"]), npz["Lneg"]),
#                 "lpre": torch.from_numpy(lpre[:, :, :2]),
#                 "lpre_label": torch.cat([torch.ones(npos), torch.zeros(nneg)]),
#                 "lpre_feat": torch.from_numpy(feat),
#             }
#
#         labels = []
#         labels = read_masks_from_pixels_wire(iname, (512, 512))
#         # if self.target_type == 'polygon':
#         #     labels, masks = read_masks_from_txt_wire(iname, (512, 512))
#         # elif self.target_type == 'pixel':
#         #     labels = read_masks_from_pixels_wire(iname, (512, 512))
#
#         target["labels"] = torch.stack(labels)
#         target["boxes"] = line_boxes_faster(meta)
#
#
#         return torch.from_numpy(image).float(), meta, target
#
#     def adjacency_matrix(self, n, link):
#         mat = torch.zeros(n + 1, n + 1, dtype=torch.uint8)
#         link = torch.from_numpy(link)
#         if len(link) > 0:
#             mat[link[:, 0], link[:, 1]] = 1
#             mat[link[:, 1], link[:, 0]] = 1
#         return mat
#
#
# def collate(batch):
#     return (
#         default_collate([b[0] for b in batch]),
#         [b[1] for b in batch],
#         default_collate([b[2] for b in batch]),
#     )

from torch.utils.data.dataset import T_co

from .models.base.base_dataset import BaseDataset

import glob
import json
import math
import os
import random
import cv2
import PIL

import matplotlib.pyplot as plt
import matplotlib as mpl
from torchvision.utils import draw_bounding_boxes

import numpy as np
import numpy.linalg as LA
import torch
from skimage import io
from torch.utils.data import Dataset
from torch.utils.data.dataloader import default_collate

import matplotlib.pyplot as plt
from .dataset_tool import line_boxes_faster, read_masks_from_txt_wire, read_masks_from_pixels_wire, adjacency_matrix



class  WireframeDataset(BaseDataset):
    def __init__(self, dataset_path, transforms=None, dataset_type=None, target_type='pixel'):
        super().__init__(dataset_path)

        self.data_path = dataset_path
        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.default_transform = DefaultTransform()

    def __getitem__(self, index) -> T_co:
        img_path = os.path.join(self.img_path, self.imgs[index])
        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))
        meta, target, target_b = self.read_target(item=index, lbl_path=lbl_path, shape=(h, w))

        img = self.default_transform(img)

        # print(f'img:{img}')
        return img, meta, target, target_b

    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)

        n_stc_posl = 300
        n_stc_negl = 40
        use_cood = 0
        use_slop = 0

        wire = lable_all["wires"][0]  # 字典
        line_pos_coords = np.random.permutation(wire["line_pos_coords"]["content"])[: n_stc_posl]  # 不足，有多少取多少
        line_neg_coords = np.random.permutation(wire["line_neg_coords"]["content"])[: n_stc_negl]
        npos, nneg = len(line_pos_coords), len(line_neg_coords)
        lpre = np.concatenate([line_pos_coords, line_neg_coords], 0)  # 正负样本坐标合在一起
        for i in range(len(lpre)):
            if random.random() > 0.5:
                lpre[i] = lpre[i, ::-1]
        ldir = lpre[:, 0, :2] - lpre[:, 1, :2]
        ldir /= np.clip(LA.norm(ldir, axis=1, keepdims=True), 1e-6, None)
        feat = [
            lpre[:, :, :2].reshape(-1, 4) / 128 * use_cood,
            ldir * use_slop,
            lpre[:, :, 2],
        ]
        feat = np.concatenate(feat, 1)

        meta = {
            "junc_coords": torch.tensor(wire["junc_coords"]["content"])[:, :2],
            "jtyp": torch.tensor(wire["junc_coords"]["content"])[:, 2].byte(),
            "line_pos_idx": adjacency_matrix(len(wire["junc_coords"]["content"]), wire["line_pos_idx"]["content"]),
            # 真实存在线条的邻接矩阵
            "line_neg_idx": adjacency_matrix(len(wire["junc_coords"]["content"]), wire["line_neg_idx"]["content"]),

            "lpre": torch.tensor(lpre)[:, :, :2],
            "lpre_label": torch.cat([torch.ones(npos), torch.zeros(nneg)]),  # 样本对应标签 1，0
            "lpre_feat": torch.from_numpy(feat),

        }

        target = {
            "junc_map": torch.tensor(wire['junc_map']["content"]),
            "junc_offset": torch.tensor(wire['junc_offset']["content"]),
            "line_map": torch.tensor(wire['line_map']["content"]),
        }

        labels = []
        if self.target_type == 'polygon':
            labels, masks = read_masks_from_txt_wire(lbl_path, shape)
        elif self.target_type == 'pixel':
            labels = read_masks_from_pixels_wire(lbl_path, shape)

        # print(torch.stack(masks).shape)    # [线段数, 512, 512]
        target_b = {}

        # target_b["image_id"] = torch.tensor(item)

        target_b["labels"] = torch.stack(labels)
        target_b["boxes"] = line_boxes_faster(meta)

        return meta, target, target_b


    def show(self, idx):
        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([])

        def imshow(im):
            plt.close()
            plt.tight_layout()
            plt.imshow(im)
            plt.colorbar(sm, fraction=0.046)
            plt.xlim([0, im.shape[0]])
            plt.ylim([im.shape[0], 0])

        def draw_vecl(lines, sline, juncs, junts, fn=None):
            img_path = os.path.join(self.img_path, self.imgs[idx])
            imshow(io.imread(img_path))
            if len(lines) > 0 and not (lines[0] == 0).all():
                for i, ((a, b), s) in enumerate(zip(lines, sline)):
                    if i > 0 and (lines[i] == lines[0]).all():
                        break
                    plt.plot([a[1], b[1]], [a[0], b[0]], c="red", linewidth=1)  # a[1], b[1]无明确大小
            if not (juncs[0] == 0).all():
                for i, j in enumerate(juncs):
                    if i > 0 and (i == juncs[0]).all():
                        break
                    plt.scatter(j[1], j[0], c="red", s=2, zorder=100)  # 原 s=64

            img_path = os.path.join(self.img_path, self.imgs[idx])
            img = PIL.Image.open(img_path).convert('RGB')
            boxed_image = draw_bounding_boxes((self.default_transform(img) * 255).to(torch.uint8), target["boxes"],
                                              colors="yellow", width=1)
            plt.imshow(boxed_image.permute(1, 2, 0).numpy())
            plt.show()

            plt.show()
            if fn != None:
                plt.savefig(fn)

        junc = target['wires']['junc_coords'].cpu().numpy() * 4
        jtyp = target['wires']['jtyp'].cpu().numpy()
        juncs = junc[jtyp == 0]
        junts = junc[jtyp == 1]

        lpre = target['wires']["lpre"].cpu().numpy() * 4
        vecl_target = target['wires']["lpre_label"].cpu().numpy()
        lpre = lpre[vecl_target == 1]

        # draw_vecl(lpre, np.ones(lpre.shape[0]), juncs, junts, save_path)
        draw_vecl(lpre, np.ones(lpre.shape[0]), juncs, junts)

    def show_img(self, img_path):
        pass


def collate(batch):
    return (
        default_collate([b[0] for b in batch]),
        [b[1] for b in batch],
        default_collate([b[2] for b in batch]),
        [b[3] for b in batch],
    )


# if __name__ == '__main__':
#     path = r"D:\python\PycharmProjects\data"
#     dataset = WireframeDataset(dataset_path=path, dataset_type='train')
#     dataset.show(0)