lcnn/models/wirenet/wirepoint_rcnn.py

import os
from typing import Optional, Any

import cv2
import numpy as np
import torch
from tensorboardX import SummaryWriter
from torch import nn
import torch.nn.functional as F
# from torchinfo import summary
from torchvision.io import read_image
from torchvision.models import resnet50, ResNet50_Weights
from torchvision.models.detection import FasterRCNN, MaskRCNN_ResNet50_FPN_V2_Weights
from torchvision.models.detection._utils import overwrite_eps
from torchvision.models.detection.backbone_utils import _resnet_fpn_extractor, _validate_trainable_layers
from torchvision.models.detection.faster_rcnn import TwoMLPHead, FastRCNNPredictor
from torchvision.models.detection.keypoint_rcnn import KeypointRCNNHeads, KeypointRCNNPredictor, \
    KeypointRCNN_ResNet50_FPN_Weights
from torchvision.ops import MultiScaleRoIAlign
from torchvision.ops import misc as misc_nn_ops
# from visdom import Visdom

from models.config import config_tool
from models.config.config_tool import read_yaml
from models.ins.trainer import get_transform
from models.wirenet.head import RoIHeads
from models.wirenet.wirepoint_dataset import WirePointDataset
from tools import utils

from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
import matplotlib as mpl
from skimage import io
import os.path as osp


FEATURE_DIM = 8

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")


def non_maximum_suppression(a):
    ap = F.max_pool2d(a, 3, stride=1, padding=1)
    mask = (a == ap).float().clamp(min=0.0)
    return a * mask


class Bottleneck1D(nn.Module):
    def __init__(self, inplanes, outplanes):
        super(Bottleneck1D, self).__init__()

        planes = outplanes // 2
        self.op = nn.Sequential(
            nn.BatchNorm1d(inplanes),
            nn.ReLU(inplace=True),
            nn.Conv1d(inplanes, planes, kernel_size=1),
            nn.BatchNorm1d(planes),
            nn.ReLU(inplace=True),
            nn.Conv1d(planes, planes, kernel_size=3, padding=1),
            nn.BatchNorm1d(planes),
            nn.ReLU(inplace=True),
            nn.Conv1d(planes, outplanes, kernel_size=1),
        )

    def forward(self, x):
        return x + self.op(x)


class WirepointRCNN(FasterRCNN):
    def __init__(
            self,
            backbone,
            num_classes=None,
            # transform parameters
            min_size=None,
            max_size=1333,
            image_mean=None,
            image_std=None,
            # RPN parameters
            rpn_anchor_generator=None,
            rpn_head=None,
            rpn_pre_nms_top_n_train=2000,
            rpn_pre_nms_top_n_test=1000,
            rpn_post_nms_top_n_train=2000,
            rpn_post_nms_top_n_test=1000,
            rpn_nms_thresh=0.7,
            rpn_fg_iou_thresh=0.7,
            rpn_bg_iou_thresh=0.3,
            rpn_batch_size_per_image=256,
            rpn_positive_fraction=0.5,
            rpn_score_thresh=0.0,
            # Box parameters
            box_roi_pool=None,
            box_head=None,
            box_predictor=None,
            box_score_thresh=0.05,
            box_nms_thresh=0.5,
            box_detections_per_img=100,
            box_fg_iou_thresh=0.5,
            box_bg_iou_thresh=0.5,
            box_batch_size_per_image=512,
            box_positive_fraction=0.25,
            bbox_reg_weights=None,
            # keypoint parameters
            keypoint_roi_pool=None,
            keypoint_head=None,
            keypoint_predictor=None,
            num_keypoints=None,
            wirepoint_roi_pool=None,
            wirepoint_head=None,
            wirepoint_predictor=None,
            **kwargs,
    ):
        if not isinstance(keypoint_roi_pool, (MultiScaleRoIAlign, type(None))):
            raise TypeError(
                "keypoint_roi_pool should be of type MultiScaleRoIAlign or None instead of {type(keypoint_roi_pool)}"
            )
        if min_size is None:
            min_size = (640, 672, 704, 736, 768, 800)

        if num_keypoints is not None:
            if keypoint_predictor is not None:
                raise ValueError("num_keypoints should be None when keypoint_predictor is specified")
        else:
            num_keypoints = 17

        out_channels = backbone.out_channels

        if wirepoint_roi_pool is None:
            wirepoint_roi_pool = MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=128,
                                                    sampling_ratio=2, )

        if wirepoint_head is None:
            keypoint_layers = tuple(512 for _ in range(8))
            # print(f'keypoinyrcnnHeads inchannels:{out_channels},layers{keypoint_layers}')
            wirepoint_head = WirepointHead(out_channels, keypoint_layers)

        if wirepoint_predictor is None:
            keypoint_dim_reduced = 512  # == keypoint_layers[-1]
            wirepoint_predictor = WirepointPredictor()

        super().__init__(
            backbone,
            num_classes,
            # transform parameters
            min_size,
            max_size,
            image_mean,
            image_std,
            # RPN-specific parameters
            rpn_anchor_generator,
            rpn_head,
            rpn_pre_nms_top_n_train,
            rpn_pre_nms_top_n_test,
            rpn_post_nms_top_n_train,
            rpn_post_nms_top_n_test,
            rpn_nms_thresh,
            rpn_fg_iou_thresh,
            rpn_bg_iou_thresh,
            rpn_batch_size_per_image,
            rpn_positive_fraction,
            rpn_score_thresh,
            # Box parameters
            box_roi_pool,
            box_head,
            box_predictor,
            box_score_thresh,
            box_nms_thresh,
            box_detections_per_img,
            box_fg_iou_thresh,
            box_bg_iou_thresh,
            box_batch_size_per_image,
            box_positive_fraction,
            bbox_reg_weights,
            **kwargs,
        )

        if box_roi_pool is None:
            box_roi_pool = MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=14, sampling_ratio=2)

        if box_head is None:
            resolution = box_roi_pool.output_size[0]
            representation_size = 1024
            box_head = TwoMLPHead(out_channels * resolution ** 2, representation_size)

        if box_predictor is None:
            representation_size = 1024
            box_predictor = FastRCNNPredictor(representation_size, num_classes)

        roi_heads = RoIHeads(
            # Box
            box_roi_pool,
            box_head,
            box_predictor,
            box_fg_iou_thresh,
            box_bg_iou_thresh,
            box_batch_size_per_image,
            box_positive_fraction,
            bbox_reg_weights,
            box_score_thresh,
            box_nms_thresh,
            box_detections_per_img,
            # wirepoint_roi_pool=wirepoint_roi_pool,
            # wirepoint_head=wirepoint_head,
            # wirepoint_predictor=wirepoint_predictor,
        )
        self.roi_heads = roi_heads

        self.roi_heads.wirepoint_roi_pool = wirepoint_roi_pool
        self.roi_heads.wirepoint_head = wirepoint_head
        self.roi_heads.wirepoint_predictor = wirepoint_predictor


class WirepointHead(nn.Module):
    def __init__(self, input_channels, num_class):
        super(WirepointHead, self).__init__()
        self.head_size = [[2], [1], [2]]
        m = int(input_channels / 4)
        heads = []
        # print(f'M.head_size:{M.head_size}')
        # for output_channels in sum(M.head_size, []):
        for output_channels in sum(self.head_size, []):
            heads.append(
                nn.Sequential(
                    nn.Conv2d(input_channels, m, kernel_size=3, padding=1),
                    nn.ReLU(inplace=True),
                    nn.Conv2d(m, output_channels, kernel_size=1),
                )
            )
        self.heads = nn.ModuleList(heads)

    def forward(self, x):
        # for idx, head in enumerate(self.heads):
        #     print(f'{idx},multitask head:{head(x).shape},input x:{x.shape}')

        outputs = torch.cat([head(x) for head in self.heads], dim=1)

        features = x
        return outputs, features


class WirepointPredictor(nn.Module):

    def __init__(self):
        super().__init__()
        # self.backbone = backbone
        # self.cfg = read_yaml(cfg)
        self.cfg = read_yaml('wirenet.yaml')
        self.n_pts0 = self.cfg['model']['n_pts0']
        self.n_pts1 = self.cfg['model']['n_pts1']
        self.n_stc_posl = self.cfg['model']['n_stc_posl']
        self.dim_loi = self.cfg['model']['dim_loi']
        self.use_conv = self.cfg['model']['use_conv']
        self.dim_fc = self.cfg['model']['dim_fc']
        self.n_out_line = self.cfg['model']['n_out_line']
        self.n_out_junc = self.cfg['model']['n_out_junc']
        self.loss_weight = self.cfg['model']['loss_weight']
        self.n_dyn_junc = self.cfg['model']['n_dyn_junc']
        self.eval_junc_thres = self.cfg['model']['eval_junc_thres']
        self.n_dyn_posl = self.cfg['model']['n_dyn_posl']
        self.n_dyn_negl = self.cfg['model']['n_dyn_negl']
        self.n_dyn_othr = self.cfg['model']['n_dyn_othr']
        self.use_cood = self.cfg['model']['use_cood']
        self.use_slop = self.cfg['model']['use_slop']
        self.n_stc_negl = self.cfg['model']['n_stc_negl']
        self.head_size = self.cfg['model']['head_size']
        self.num_class = sum(sum(self.head_size, []))
        self.head_off = np.cumsum([sum(h) for h in self.head_size])

        lambda_ = torch.linspace(0, 1, self.n_pts0)[:, None]
        self.register_buffer("lambda_", lambda_)
        self.do_static_sampling = self.n_stc_posl + self.n_stc_negl > 0

        self.fc1 = nn.Conv2d(256, self.dim_loi, 1)
        scale_factor = self.n_pts0 // self.n_pts1
        if self.use_conv:
            self.pooling = nn.Sequential(
                nn.MaxPool1d(scale_factor, scale_factor),
                Bottleneck1D(self.dim_loi, self.dim_loi),
            )
            self.fc2 = nn.Sequential(
                nn.ReLU(inplace=True), nn.Linear(self.dim_loi * self.n_pts1 + FEATURE_DIM, 1)
            )
        else:
            self.pooling = nn.MaxPool1d(scale_factor, scale_factor)
            self.fc2 = nn.Sequential(
                nn.Linear(self.dim_loi * self.n_pts1 + FEATURE_DIM, self.dim_fc),
                nn.ReLU(inplace=True),
                nn.Linear(self.dim_fc, self.dim_fc),
                nn.ReLU(inplace=True),
                nn.Linear(self.dim_fc, 1),
            )
        self.loss = nn.BCEWithLogitsLoss(reduction="none")

    def forward(self, inputs, features, targets=None):

        # outputs, features = input
        # for out in outputs:
        #     print(f'out:{out.shape}')
        # outputs=merge_features(outputs,100)
        batch, channel, row, col = inputs.shape
        # print(f'outputs:{inputs.shape}')
        # print(f'batch:{batch}, channel:{channel}, row:{row}, col:{col}')

        if targets is not None:
            self.training = True
            # print(f'target:{targets}')
            wires_targets = [t["wires"] for t in targets]
            # print(f'wires_target:{wires_targets}')
            # 提取所有 'junc_map', 'junc_offset', 'line_map' 的张量
            junc_maps = [d["junc_map"] for d in wires_targets]
            junc_offsets = [d["junc_offset"] for d in wires_targets]
            line_maps = [d["line_map"] for d in wires_targets]

            junc_map_tensor = torch.stack(junc_maps, dim=0)
            junc_offset_tensor = torch.stack(junc_offsets, dim=0)
            line_map_tensor = torch.stack(line_maps, dim=0)

            wires_meta = {
                "junc_map": junc_map_tensor,
                "junc_offset": junc_offset_tensor,
                # "line_map": line_map_tensor,
            }
        else:
            self.training = False
            t = {
                "junc_coords": torch.zeros(1, 2).to(device),
                "jtyp": torch.zeros(1, dtype=torch.uint8).to(device),
                "line_pos_idx": torch.zeros(2, 2, dtype=torch.uint8).to(device),
                "line_neg_idx": torch.zeros(2, 2, dtype=torch.uint8).to(device),
                "junc_map": torch.zeros([1, 1, 128, 128]).to(device),
                "junc_offset": torch.zeros([1, 1, 2, 128, 128]).to(device),
            }
            wires_targets = [t for b in range(inputs.size(0))]

            wires_meta = {
                "junc_map": torch.zeros([1, 1, 128, 128]).to(device),
                "junc_offset": torch.zeros([1, 1, 2, 128, 128]).to(device),
            }

        T = wires_meta.copy()
        n_jtyp = T["junc_map"].shape[1]
        offset = self.head_off
        result = {}
        for stack, output in enumerate([inputs]):
            output = output.transpose(0, 1).reshape([-1, batch, row, col]).contiguous()
            # print(f"Stack {stack} output shape: {output.shape}")  # 打印每层的输出形状
            jmap = output[0: offset[0]].reshape(n_jtyp, 2, batch, row, col)
            lmap = output[offset[0]: offset[1]].squeeze(0)
            joff = output[offset[1]: offset[2]].reshape(n_jtyp, 2, batch, row, col)

            if stack == 0:
                result["preds"] = {
                    "jmap": jmap.permute(2, 0, 1, 3, 4).softmax(2)[:, :, 1],
                    "lmap": lmap.sigmoid(),
                    "joff": joff.permute(2, 0, 1, 3, 4).sigmoid() - 0.5,
                }

        h = result["preds"]
        # print(f'features shape:{features.shape}')
        x = self.fc1(features)
        n_batch, n_channel, row, col = x.shape
        xs, ys, fs, ps, idx, jcs = [], [], [], [], [0], []

        for i, meta in enumerate(wires_targets):
            p, label, feat, jc = self.sample_lines(
                meta, h["jmap"][i], h["joff"][i],
            )
            # print(f"p.shape:{p.shape},label:{label.shape},feat:{feat.shape},jc:{len(jc)}")
            ys.append(label)
            if self.training and self.do_static_sampling:
                p = torch.cat([p, meta["lpre"]])
                feat = torch.cat([feat, meta["lpre_feat"]])
                ys.append(meta["lpre_label"])
                del jc
            else:
                jcs.append(jc)
                ps.append(p)
            fs.append(feat)

            p = p[:, 0:1, :] * self.lambda_ + p[:, 1:2, :] * (1 - self.lambda_) - 0.5
            p = p.reshape(-1, 2)  # [N_LINE x N_POINT, 2_XY]
            px, py = p[:, 0].contiguous(), p[:, 1].contiguous()
            px0 = px.floor().clamp(min=0, max=127)
            py0 = py.floor().clamp(min=0, max=127)
            px1 = (px0 + 1).clamp(min=0, max=127)
            py1 = (py0 + 1).clamp(min=0, max=127)
            px0l, py0l, px1l, py1l = px0.long(), py0.long(), px1.long(), py1.long()

            # xp: [N_LINE, N_CHANNEL, N_POINT]
            xp = (
                (
                        x[i, :, px0l, py0l] * (px1 - px) * (py1 - py)
                        + x[i, :, px1l, py0l] * (px - px0) * (py1 - py)
                        + x[i, :, px0l, py1l] * (px1 - px) * (py - py0)
                        + x[i, :, px1l, py1l] * (px - px0) * (py - py0)
                )
                    .reshape(n_channel, -1, self.n_pts0)
                    .permute(1, 0, 2)
            )
            xp = self.pooling(xp)
            # print(f'xp.shape:{xp.shape}')
            xs.append(xp)
            idx.append(idx[-1] + xp.shape[0])
            # print(f'idx__:{idx}')

        x, y = torch.cat(xs), torch.cat(ys)
        f = torch.cat(fs)
        x = x.reshape(-1, self.n_pts1 * self.dim_loi)
        x = torch.cat([x, f], 1)
        x = x.to(dtype=torch.float32)
        x = self.fc2(x).flatten()

        # return  x,idx,jcs,n_batch,ps,self.n_out_line,self.n_out_junc
        return x, y, idx, jcs, n_batch, ps, self.n_out_line, self.n_out_junc

        # if mode != "training":
        # self.inference(x, idx, jcs, n_batch, ps)

        # return result

    def sample_lines(self, meta, jmap, joff):
        with torch.no_grad():
            junc = meta["junc_coords"]  # [N, 2]
            jtyp = meta["jtyp"]  # [N]
            Lpos = meta["line_pos_idx"]
            Lneg = meta["line_neg_idx"]

            n_type = jmap.shape[0]
            jmap = non_maximum_suppression(jmap).reshape(n_type, -1)
            joff = joff.reshape(n_type, 2, -1)
            max_K = self.n_dyn_junc // n_type
            N = len(junc)
            # if mode != "training":
            if not self.training:
                K = min(int((jmap > self.eval_junc_thres).float().sum().item()), max_K)
            else:
                K = min(int(N * 2 + 2), max_K)
            if K < 2:
                K = 2
            device = jmap.device

            # index: [N_TYPE, K]
            score, index = torch.topk(jmap, k=K)
            y = (index // 128).float() + torch.gather(joff[:, 0], 1, index) + 0.5
            x = (index % 128).float() + torch.gather(joff[:, 1], 1, index) + 0.5

            # xy: [N_TYPE, K, 2]
            xy = torch.cat([y[..., None], x[..., None]], dim=-1)
            xy_ = xy[..., None, :]
            del x, y, index

            # print(f"xy_.is_cuda: {xy_.is_cuda}")
            # print(f"junc.is_cuda: {junc.is_cuda}")

            # dist: [N_TYPE, K, N]
            dist = torch.sum((xy_ - junc) ** 2, -1)
            cost, match = torch.min(dist, -1)

            # xy: [N_TYPE * K, 2]
            # match: [N_TYPE, K]
            for t in range(n_type):
                match[t, jtyp[match[t]] != t] = N
            match[cost > 1.5 * 1.5] = N
            match = match.flatten()

            _ = torch.arange(n_type * K, device=device)
            u, v = torch.meshgrid(_, _)
            u, v = u.flatten(), v.flatten()
            up, vp = match[u], match[v]
            label = Lpos[up, vp]

            # if mode == "training":
            if self.training:
                c = torch.zeros_like(label, dtype=torch.bool)

                # sample positive lines
                cdx = label.nonzero().flatten()
                if len(cdx) > self.n_dyn_posl:
                    # print("too many positive lines")
                    perm = torch.randperm(len(cdx), device=device)[: self.n_dyn_posl]
                    cdx = cdx[perm]
                c[cdx] = 1

                # sample negative lines
                cdx = Lneg[up, vp].nonzero().flatten()
                if len(cdx) > self.n_dyn_negl:
                    # print("too many negative lines")
                    perm = torch.randperm(len(cdx), device=device)[: self.n_dyn_negl]
                    cdx = cdx[perm]
                c[cdx] = 1

                # sample other (unmatched) lines
                cdx = torch.randint(len(c), (self.n_dyn_othr,), device=device)
                c[cdx] = 1
            else:
                c = (u < v).flatten()

            # sample lines
            u, v, label = u[c], v[c], label[c]
            xy = xy.reshape(n_type * K, 2)
            xyu, xyv = xy[u], xy[v]

            u2v = xyu - xyv
            u2v /= torch.sqrt((u2v ** 2).sum(-1, keepdim=True)).clamp(min=1e-6)
            feat = torch.cat(
                [
                    xyu / 128 * self.use_cood,
                    xyv / 128 * self.use_cood,
                    u2v * self.use_slop,
                    (u[:, None] > K).float(),
                    (v[:, None] > K).float(),
                ],
                1,
            )
            line = torch.cat([xyu[:, None], xyv[:, None]], 1)

            xy = xy.reshape(n_type, K, 2)
            jcs = [xy[i, score[i] > 0.03] for i in range(n_type)]
            return line, label.float(), feat, jcs


def wirepointrcnn_resnet50_fpn(
        *,
        weights: Optional[KeypointRCNN_ResNet50_FPN_Weights] = None,
        progress: bool = True,
        num_classes: Optional[int] = None,
        num_keypoints: Optional[int] = None,
        weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1,
        trainable_backbone_layers: Optional[int] = None,
        **kwargs: Any,
) -> WirepointRCNN:
    weights = KeypointRCNN_ResNet50_FPN_Weights.verify(weights)
    weights_backbone = ResNet50_Weights.verify(weights_backbone)

    is_trained = weights is not None or weights_backbone is not None
    trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3)

    norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d

    backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=norm_layer)
    backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers)
    model = WirepointRCNN(backbone, num_classes=5, **kwargs)

    if weights is not None:
        model.load_state_dict(weights.get_state_dict(progress=progress))
        if weights == KeypointRCNN_ResNet50_FPN_Weights.COCO_V1:
            overwrite_eps(model, 0.0)

    return model


def _loss(losses):
    total_loss = 0
    for i in losses.keys():
        if i != "loss_wirepoint":
            total_loss += losses[i]
        else:
            loss_labels = losses[i]["losses"]
    loss_labels_k = list(loss_labels[0].keys())
    for j, name in enumerate(loss_labels_k):
        loss = loss_labels[0][name].mean()
        total_loss += loss

    return total_loss


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 c(x):
    return sm.to_rgba(x)


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 _plot_samples(self, i, index, result, targets, prefix):
    fn = self.val_loader.dataset.filelist[index][:-10].replace("_a0", "") + ".png"
    img = io.imread(fn)
    imshow(img), plt.savefig(f"{prefix}_img.jpg"), plt.close()

    def draw_vecl(lines, sline, juncs, junts, fn):
        imshow(img)
        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=c(s), linewidth=4)
        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=64, zorder=100)
        if junts is not None and len(junts) > 0 and not (junts[0] == 0).all():
            for i, j in enumerate(junts):
                if i > 0 and (i == junts[0]).all():
                    break
                plt.scatter(j[1], j[0], c="blue", s=64, zorder=100)
        plt.savefig(fn), plt.close()

    junc = targets[i]["junc"].cpu().numpy() * 4
    jtyp = targets[i]["jtyp"].cpu().numpy()
    juncs = junc[jtyp == 0]
    junts = junc[jtyp == 1]
    rjuncs = result["juncs"][i].cpu().numpy() * 4
    rjunts = None
    if "junts" in result:
        rjunts = result["junts"][i].cpu().numpy() * 4

    lpre = targets[i]["lpre"].cpu().numpy() * 4
    vecl_target = targets[i]["lpre_label"].cpu().numpy()
    vecl_result = result["lines"][i].cpu().numpy() * 4
    score = result["score"][i].cpu().numpy()
    lpre = lpre[vecl_target == 1]

    draw_vecl(lpre, np.ones(lpre.shape[0]), juncs, junts, f"{prefix}_vecl_a.jpg")
    draw_vecl(vecl_result, score, rjuncs, rjunts, f"{prefix}_vecl_b.jpg")

    img = cv2.imread(f"{prefix}_vecl_a.jpg")
    img1 = cv2.imread(f"{prefix}_vecl_b.jpg")
    self.writer.add_images(f"{self.epoch}", torch.tensor([img, img1]), dataformats='NHWC')


if __name__ == '__main__':
    cfg = 'wirenet.yaml'
    cfg = read_yaml(cfg)
    print(f'cfg:{cfg}')
    print(cfg['model']['n_dyn_negl'])
    # net = WirepointPredictor()

    # if torch.cuda.is_available():
    #     device_name = "cuda"
    #     torch.backends.cudnn.deterministic = True
    #     torch.cuda.manual_seed(0)
    #     print("Let's use", torch.cuda.device_count(), "GPU(s)!")
    # else:
    #     print("CUDA is not available")
    #
    # device = torch.device(device_name)

    dataset_train = WirePointDataset(dataset_path=cfg['io']['datadir'], dataset_type='train')
    train_sampler = torch.utils.data.RandomSampler(dataset_train)
    # test_sampler = torch.utils.data.SequentialSampler(dataset_test)
    train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, batch_size=1, drop_last=True)
    train_collate_fn = utils.collate_fn_wirepoint
    data_loader_train = torch.utils.data.DataLoader(
        dataset_train, batch_sampler=train_batch_sampler, num_workers=0, collate_fn=train_collate_fn
    )

    dataset_val = WirePointDataset(dataset_path=cfg['io']['datadir'], dataset_type='val')
    val_sampler = torch.utils.data.RandomSampler(dataset_val)
    # test_sampler = torch.utils.data.SequentialSampler(dataset_test)
    val_batch_sampler = torch.utils.data.BatchSampler(val_sampler, batch_size=1, drop_last=True)
    val_collate_fn = utils.collate_fn_wirepoint
    data_loader_val = torch.utils.data.DataLoader(
        dataset_val, batch_sampler=val_batch_sampler, num_workers=0, collate_fn=val_collate_fn
    )

    model = wirepointrcnn_resnet50_fpn().to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=cfg['optim']['lr'])
    writer = SummaryWriter(cfg['io']['logdir'])


    def move_to_device(data, device):
        if isinstance(data, (list, tuple)):
            return type(data)(move_to_device(item, device) for item in data)
        elif isinstance(data, dict):
            return {key: move_to_device(value, device) for key, value in data.items()}
        elif isinstance(data, torch.Tensor):
            return data.to(device)
        else:
            return data  # 对于非张量类型的数据不做任何改变


    def writer_loss(writer, losses, epoch):
        # ??????
        try:
            for key, value in losses.items():
                if key == 'loss_wirepoint':
                    # ?? wirepoint ??????
                    for subdict in losses['loss_wirepoint']['losses']:
                        for subkey, subvalue in subdict.items():
                            # ?? .item() ?????
                            writer.add_scalar(f'loss_wirepoint/{subkey}',
                                              subvalue.item() if hasattr(subvalue, 'item') else subvalue,
                                              epoch)
                elif isinstance(value, torch.Tensor):
                    # ????????
                    writer.add_scalar(key, value.item(), epoch)
        except Exception as e:
            print(f"TensorBoard logging error: {e}")


    for epoch in range(cfg['optim']['max_epoch']):
        print(f"epoch:{epoch}")
        model.train()

        for imgs, targets in data_loader_train:
            losses = model(move_to_device(imgs, device), move_to_device(targets, device))
            loss = _loss(losses)
            print(loss)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            writer_loss(writer, losses, epoch)

            model.eval()
            with torch.no_grad():
                for batch_idx, (imgs, targets) in enumerate(data_loader_val):
                    pred = model(move_to_device(imgs, device))
                    print(f"perd:{pred}")

                # if batch_idx == 0:
                #     viz = osp.join(cfg['io']['logdir'], "viz", f"{epoch}")
                #     H = pred["wires"]
                #     _plot_samples(0, 0, H, targets["wires"], f"{viz}/{epoch}")

# imgs, targets = next(iter(data_loader))
#
# model.train()
# pred = model(imgs, targets)
# print(f'pred:{pred}')

# result, losses = model(imgs, targets)
# print(f'result:{result}')
# print(f'pred:{losses}')