# 2025/2/9 import os from typing import Optional, Any import cv2 import numpy as np import torch from models.config.config_tool import read_yaml from models.line_detect.dataset_LD 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 from models.line_detect.line_net import linenet_resnet50_fpn from torchvision.utils import draw_bounding_boxes from models.wirenet.postprocess import postprocess from torchvision import transforms from collections import OrderedDict device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 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 show_line(img, pred, epoch, writer): im = img.permute(1, 2, 0) writer.add_image("ori", im, epoch, dataformats="HWC") boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), pred[0]["boxes"], colors="yellow", width=1) writer.add_image("boxes", boxed_image.permute(1, 2, 0), epoch, dataformats="HWC") PLTOPTS = {"color": "#33FFFF", "s": 15, "edgecolors": "none", "zorder": 5} H = pred[1]['wires'] 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.8]): 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.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() img2 = transforms.ToTensor()(image_from_plot) writer.add_image("output", img2, epoch) if __name__ == '__main__': cfg = r'./config/wireframe.yaml' cfg = read_yaml(cfg) print(f'cfg:{cfg}') print(cfg['model']['n_dyn_negl']) # net = WirepointPredictor() 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=2, 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=8, 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=2, 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=8, collate_fn=val_collate_fn ) model = linenet_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': # for subdict in losses['loss_wirepoint']['losses']: # for subkey, subvalue in subdict.items(): # 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}") def writer_loss(writer, losses, epoch): try: for key, value in losses.items(): if key == 'loss_wirepoint': for subdict in losses['loss_wirepoint']['losses']: for subkey, subvalue in subdict.items(): writer.add_scalar(f'loss/{subkey}', subvalue.item() if hasattr(subvalue, 'item') else subvalue, epoch) elif isinstance(value, torch.Tensor): writer.add_scalar(f'loss/{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)) # print(losses) loss = _loss(losses) 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)) if batch_idx == 0: show_line(imgs[0], pred, epoch, writer) break