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
from torchvision.utils import draw_bounding_boxes
from torchvision import transforms
from models.wirenet.postprocess import postprocess
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])
# plt.show()
# def _plot_samples(img, i, result, prefix, epoch):
# print(f"prefix:{prefix}")
# def draw_vecl(lines, sline, juncs, junts, fn):
# directory = os.path.dirname(fn)
# if not os.path.exists(directory):
# os.makedirs(directory)
# imshow(img.permute(1, 2, 0))
# 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()
#
# rjuncs = result["juncs"][i].cpu().numpy() * 4
# rjunts = None
# if "junts" in result:
# rjunts = result["junts"][i].cpu().numpy() * 4
#
# vecl_result = result["lines"][i].cpu().numpy() * 4
# score = result["score"][i].cpu().numpy()
#
# draw_vecl(vecl_result, score, rjuncs, rjunts, f"{prefix}_vecl_b.jpg")
#
# img1 = cv2.imread(f"{prefix}_vecl_b.jpg")
# writer.add_image(f'output_epoch_{epoch}', img1, global_step=epoch)
def _plot_samples(img, i, result, prefix, epoch, writer):
# print(f"prefix:{prefix}")
def draw_vecl(lines, sline, juncs, junts, fn):
# 确保目录存在
directory = os.path.dirname(fn)
if not os.path.exists(directory):
os.makedirs(directory)
# 绘制图像
plt.figure()
plt.imshow(img.permute(1, 2, 0).cpu().numpy())
plt.axis('off') # 可选：关闭坐标轴
if len(lines) > 0 and not (lines[0] == 0).all():
for idx, ((a, b), s) in enumerate(zip(lines, sline)):
if idx > 0 and (lines[idx] == lines[0]).all():
break
plt.plot([a[1], b[1]], [a[0], b[0]], c=c(s), linewidth=1)
if not (juncs[0] == 0).all():
for idx, j in enumerate(juncs):
if idx > 0 and (j == juncs[0]).all():
break
plt.scatter(j[1], j[0], c="red", s=20, zorder=100)
if junts is not None and len(junts) > 0 and not (junts[0] == 0).all():
for idx, j in enumerate(junts):
if idx > 0 and (j == junts[0]).all():
break
plt.scatter(j[1], j[0], c="blue", s=20, zorder=100)
# plt.show()
# 将matplotlib图像转换为numpy数组
plt.tight_layout()
fig = plt.gcf()
fig.canvas.draw()
image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8).reshape(
fig.canvas.get_width_height()[::-1] + (3,))
plt.close()
return image_from_plot
# 获取结果数据并转换为numpy数组
rjuncs = result["juncs"][i].cpu().numpy() * 4
rjunts = None
if "junts" in result:
rjunts = result["junts"][i].cpu().numpy() * 4
vecl_result = result["lines"][i].cpu().numpy() * 4
score = result["score"][i].cpu().numpy()
# 调用绘图函数并获取图像
image_path = f"{prefix}_vecl_b.jpg"
image_array = draw_vecl(vecl_result, score, rjuncs, rjunts, image_path)
# 将numpy数组转换为torch tensor，并写入TensorBoard
image_tensor = transforms.ToTensor()(image_array)
writer.add_image(f'output_epoch', image_tensor, global_step=epoch)
writer.add_image(f'ori_epoch', img, global_step=epoch)
def show_line(img, pred, prefix, epoch, write):
fn = f"{prefix}_line.jpg"
directory = os.path.dirname(fn)
if not os.path.exists(directory):
os.makedirs(directory)
print(fn)
PLTOPTS = {"color": "#33FFFF", "s": 15, "edgecolors": "none", "zorder": 5}
H = pred
im = img.permute(1, 2, 0)
lines = H["lines"][0].cpu().numpy() / 128 * im.shape[:2]
scores = H["score"][0].cpu().numpy()
for i in range(1, len(lines)):
if (lines[i] == lines[0]).all():
lines = lines[:i]
scores = scores[:i]
break
# postprocess lines to remove overlapped lines
diag = (im.shape[0] ** 2 + im.shape[1] ** 2) ** 0.5
nlines, nscores = postprocess(lines, scores, diag * 0.01, 0, False)
for i, t in enumerate([0.5]):
plt.gca().set_axis_off()
plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
plt.margins(0, 0)
for (a, b), s in zip(nlines, nscores):
if s < t:
continue
plt.plot([a[1], b[1]], [a[0], b[0]], c=c(s), linewidth=2, zorder=s)
plt.scatter(a[1], a[0], **PLTOPTS)
plt.scatter(b[1], b[0], **PLTOPTS)
plt.gca().xaxis.set_major_locator(plt.NullLocator())
plt.gca().yaxis.set_major_locator(plt.NullLocator())
plt.imshow(im)
plt.savefig(fn, bbox_inches="tight")
plt.show()
plt.close()
img2 = cv2.imread(fn) # 预测图
# img1 = im.resize(img2.shape) # 原图
# writer.add_images(f"{epoch}", torch.tensor([img1, img2]), dataformats='NHWC')
writer.add_image("output", img2, epoch)
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"pred:{pred}")
#
# if batch_idx == 0:
# result = pred[1]['wires'] # pred[0].keys() ['boxes', 'labels', 'scores']
# print(imgs[0].shape) # [3，512，512]
# # imshow(imgs[0].permute(1, 2, 0)) # 改为(512, 512, 3)
# _plot_samples(imgs[0], 0, result, f"{cfg['io']['logdir']}/{epoch}/", epoch, writer)
# show_line(imgs[0], result, f"{cfg['io']['logdir']}/{epoch}", epoch, writer)
# 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}')