pokouqiege/aaa.py

# import torch
# from torchvision.utils import draw_bounding_boxes
# from torchvision import transforms
# import matplotlib.pyplot as plt
# import numpy as np
#
#
# def c(score):
#     # 根据分数返回颜色的函数，这里仅作示例，您可以根据需要修改
#     return (1, 0, 0) if score > 0.9 else (0, 1, 0)
#
#
# def postprocess(lines, scores, diag_threshold, min_score, remove_overlaps):
#     # 假设的后处理函数，用于过滤线段
#     nlines = []
#     nscores = []
#     for line, score in zip(lines, scores):
#         if score >= min_score:
#             nlines.append(line)
#             nscores.append(score)
#     return np.array(nlines), np.array(nscores)
#
#
# def show_line(img, pred, epoch, writer):
#     im = img.permute(1, 2, 0).cpu().numpy()
#
#     # 绘制边界框
#     boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), pred[0]["boxes"],
#                                       colors="yellow", width=1).permute(1, 2, 0).cpu().numpy()
#
#     H = pred[-1]['wires']
#     lines = H["lines"][0].cpu().numpy() / 128 * im.shape[:2]
#     scores = H["score"][0].cpu().numpy()
#
#     print(f"Lines before deduplication: {len(lines)}")
#
#     # 移除重复的线段
#     for i in range(1, len(lines)):
#         if (lines[i] == lines[0]).all():
#             lines = lines[:i]
#             scores = scores[:i]
#             break
#
#     print(f"Lines after deduplication: {len(lines)}")
#
#     # 后处理线段以移除重叠的线段
#     diag = (im.shape[0] ** 2 + im.shape[1] ** 2) ** 0.5
#     nlines, nscores = postprocess(lines, scores, diag * 0.01, 0, False)
#
#     print(f"Lines after postprocessing: {len(nlines)}")
#
#     # 创建一个新的图像并绘制线段和边界框
#     fig, ax = plt.subplots(figsize=(boxed_image.shape[1] / 100, boxed_image.shape[0] / 100))
#     ax.imshow(boxed_image)
#     ax.set_axis_off()
#     plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
#     plt.margins(0, 0)
#     plt.gca().xaxis.set_major_locator(plt.NullLocator())
#     plt.gca().yaxis.set_major_locator(plt.NullLocator())
#
#     PLTOPTS = {"color": "#33FFFF", "s": 15, "edgecolors": "none", "zorder": 5}
#     for (a, b), s in zip(nlines, nscores):
#         if s < 0.85:  # 调整阈值以筛选显示的线段
#             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.tight_layout()
#     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_with_boxes_and_lines", img2, epoch)
#     print("Image with boxes and lines added to TensorBoard.")


import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter
import random
# 假设我们有一些关键点位置
keypoints = [(0, 0), (70, 80), (90, 30)]

# 创建一个空白的热图
heatmap = np.zeros((100, 100))

# 将关键点位置添加到热图中
for point in keypoints:
    y, x = point
    heatmap[y, x] = random.random()
    # heatmap[y, x] = 1  # 假设置信度为1

print(heatmap)
# 使用高斯滤波平滑热图
heatmap_smooth = gaussian_filter(heatmap, sigma=1)
print(heatmap_smooth)