torchvision.utils 可视化工具详解:make_grid、检测框、分割掩码、关键点与光流绘制

发布时间:2026/9/21 16:08:07
torchvision.utils 可视化工具详解:make_grid、检测框、分割掩码、关键点与光流绘制 计算机视觉深度学习图像处理数据集【免费下载链接】visionDatasets, Transforms and Models specific to Computer Vision项目地址https://gitcode.com/gh_mirrors/vi/vision点击查看免费下载torchvision.utils是 TorchVision 附带的工具箱模块集中提供计算机视觉开发中最常用的结果可视化能力将小批量图像拼接为网格、在图像上绘制检测框/分割掩码/关键点以及把光流场转为便于人眼观察的 RGB 图像。本文基于当前仓库中 docs/source/utils.rst 的 API 参考draw_bounding_boxes、draw_segmentation_masks、draw_keypoints、flow_to_image、make_grid、save_image 六个公开函数结合 torchvision/utils.py 的完整源码实现、官方可视化示例 与 test/test_utils.py 单元测试逐函数讲解参数含义、底层原理与实战用法。读完本文你将能够独立完成模型输出 → 张量处理 → 图像可视化的完整流水线搭建。一、模块定位与公开 API 总览torchvision.utils是 TorchVision 顶层模块之一在 torchvision/init.py 中被导出因此可以通过import torchvision.utils as utils或from torchvision.utils import draw_bounding_boxes直接使用。按照 docs/source/utils.rst 的说明该模块主要包含各类可视化工具。模块在 torchvision/utils.py 中通过__all__明确声明了全部公开符号公开函数核心用途典型场景make_grid把多个图像拼成一张网格图批量展示样本、训练过程可视化save_image把张量或网格保存为图像文件保存训练/推理结果、断点快照draw_bounding_boxes在图像上绘制含旋转边界框及标签目标检测结果可视化draw_segmentation_masks在图像上叠加半透明分割掩码语义/实例分割结果可视化draw_keypoints绘制关键点及骨架连接线姿态估计、动作识别可视化flow_to_image把光流场编码为彩色图像光流/视频任务结果可视化除_Image_fromarray为内部辅助函数外以上 6 个函数构成了模块的完整公开面。所有绘制类函数都带有torch.no_grad()装饰器且均通过_log_api_usage_once上报调用信息因此在推理/可视化路径中使用不会引入梯度计算开销。需要特别说明的是文档描述这些工具mostly for visualization因此它们的设计目标是与检测、分割、姿态等模型的原始张量输出直接对接而不要求张量在 CUDA 上——多数函数内部会自行将数据搬移到 CPU 进行绘制。二、make_grid把批量图像拼接成网格2.1 签名与参数make_grid是save_image与各类批量展示场景的基石其签名定义于 torchvision/utils.pymake_grid( tensor, # Tensor 或 Tensor 列表 nrow8, # 每行图像数量 padding2, # 图像间间距像素 normalizeFalse, # 是否归一化到 [0, 1] value_rangeNone, # 归一化使用的 (min, max)默认从张量计算 scale_eachFalse, # 是否逐图独立缩放 pad_value0.0, # 填充像素的取值 ) - torch.Tensor参数说明tensor4D 小批量张量形状为(B, C, H, W)也可以传入尺寸完全一致的图像张量列表内部会先torch.stack堆叠见 utils.py#L66-L67。nrow每行展示的图像数最终网格形状约为(ceil(B/nrow), nrow)默认 8。padding图像之间的间距像素数默认 2。normalize若为True会把图像线性拉伸到[0, 1]区间区间由value_range决定未指定时取整张量的全局 min/max。value_range二元组(min, max)仅当normalizeTrue时生效。scale_each为True时对 batch 内每张图独立计算 min/max 归一化而不是使用全 batch 的全局极值。pad_value填充像素的值默认 0黑色。2.2 输入形状的自动适配逻辑源码在 utils.py#L69-L77 中对输入做了多级容错处理2D 张量(H, W)视为单张灰度图先升维3D 张量(C, H, W)视为单张图像若C 1单通道会自动复制 3 份变成 3 通道4D 张量且C 1同样沿通道维复制 3 次保证网格输出为 RGB。也就是说即使你传入单通道的灰度批次最终得到的也是 3 通道网格图便于直接保存或显示。2.3 归一化的实现细节归一化在 utils.py#L79-L98 中实现核心是一个norm_ip函数先clamp_到[low, high]再(img - low) / max(high - low, 1e-5)。其中1e-5的下限防止除零。两个值得注意的细节非原地修改归一化前先tensor.clone()避免原地改写输入。这一点由 test/test_utils.py#L42-L53 的test_make_grid_not_inplace用例专门验证——无论normalize与scale_each如何组合输入张量都必须保持不变。scale_each分支为True时按 batch 维度逐张归一化for t in tensor否则对整体张量统一归一化。2.4 网格拼装与单图快捷返回网格拼装逻辑位于 utils.py#L105-L123按xmaps min(nrow, B)、ymaps ceil(B / xmaps)计算行列数预分配(C, H*ymapspadding, W*xmapspadding)的填充张量再通过narrow copy_把每张图写入对应位置。还有一个易被忽略的快捷行为当 batch 大小为 1 时直接返回squeeze(0)后的单张图utils.py#L102-L103不产生任何拼接开销。2.5 官方示例中的用法在 gallery/others/plot_visualization_utils.py#L37-L53 中官方示例先用torchvision.io.decode_image读取两张狗的图像再直接传给make_gridfrom torchvision.utils import make_grid from torchvision.io import decode_image from pathlib import Path dog1_int decode_image(str(Path(../assets) / dog1.jpg)) dog2_int decode_image(str(Path(../assets) / dog2.jpg)) dog_list [dog1_int, dog2_int] grid make_grid(dog_list) show(grid)三、save_image把张量保存为图像文件save_image定义于 torchvision/utils.py#L256-L281其本质是make_grid的一层封装save_image( tensor, # Tensor 或 Tensor 列表 fp, # 文件名 / 路径 / 二进制文件对象 formatNone, # 图像格式省略时从文件扩展名推断 **kwargs, # 透传给 make_grid 的参数nrow/padding/normalize 等 ) - None实现流程utils.py#L277-L281非常直观调用make_grid(tensor, **kwargs)得到网格张量执行grid.mul(255).add_(0.5).clamp_(0, 255)把[0,1]浮点映射回[0,255]其中0.5用于四舍五入到最近整数permute(1, 2, 0)将CHW转为HWC转 CPU 并截断为uint8再转 numpy用 PIL 的Image.fromarray构造图像并im.save(fp, formatformat)。两个使用要点传入文件对象时务必显式指定format因为从文件对象无法推断扩展名源码 docstring 明确提示如果使用文件对象而非文件名应始终指定该参数。支持的文件目标fp可以是字符串路径、pathlib.Path或任何带save方法的二进制流。test_save_image_file_objecttest/test_utils.py#L91-L100验证了写入BytesIO再读回与直接写文件结果完全一致fp BytesIO() utils.save_image(t, fp, formatpng) img_bytes Image.open(fp)四、draw_bounding_boxes目标检测框可视化4.1 签名与参数draw_bounding_boxes是检测任务最常用的可视化函数签名位于 torchvision/utils.py#L284-L297draw_bounding_boxes( image, # (C, H, W)uint8 或 float boxes, # (N, 4) 或 (N, 8) labelsNone, # 每个框的标签字符串列表 colorsNone, # 颜色字符串/RGB 元组/颜色列表 fillFalse, # 是否填充框内部 width1, # 框线宽 fontNone, # TrueType 字体文件路径 font_sizeNone, # 字号点 label_colorsNone, # 标签文字颜色 label_background_colorsNone, # 标签背景色 fill_labelsFalse, # 是否给标签绘制背景块 ) - torch.Tensor # (C, H, W) uint8或与输入同 dtype4.2 关键参数详解image形状(C, H, W)dtype 必须是uint8值域[0,255]或浮点值域[0,1]支持灰度C1与 RGBC3。源码在 utils.py#L336-L343 做了严格校验不是 Tensor、dtype 不对、不是单图维度不为 3、通道数不是 1 或 3 都会抛异常。boxes两种格式utils.py#L305-L309(N, 4)(xmin, ymin, xmax, ymax)绝对坐标且必须满足0 xmin xmax W、0 ymin ymax H若出现xmin xmax或ymin ymax会抛出提示请用torchvision.ops.box_convert转换的异常utils.py#L344-L347。(N, 8)旋转框的 4 个角点(x1,y1,x2,y2,x3,y3,x4,y4)绝对坐标源码据此选择oriented_rectangle绘制路径utils.py#L397。colors支持 PIL 颜色名red、十六进制串#FF00FF或 RGB 元组(240, 10, 157)也支持颜色列表按框逐一着色不传时由_generate_color_palette自动生成调色板。颜色解析统一走_parse_colorsutils.py#L722-L766若列表长度小于框数会直接报错。fill为True时以颜色 100 的 alpha 值填充框内部utils.py#L398此时输出应保存为 PNG 以保留透明通道。font/font_sizefont为 TrueType 字体文件路径Windows 下会搜索系统fonts/目录macOS 下会搜索/Library/Fonts、/System/Library/Fonts与~/Library/Fonts。不传font时用 PIL 默认字体此时如果传了font_size会触发警告test_draw_boxes_warning用例验证test/test_utils.py#L292-L296。label_colors/label_background_colors/fill_labels默认标签文字颜色与框颜色相同fill_labelsTrue时为黑色fill_labelsTrue时会在文字后绘制背景色块背景默认取框颜色。测试 test/test_utils.py#L137-L193 分别验证了彩色标签文字彩色标签背景标签文字背景双色三种组合并与test/assets/fakedata/下的基准图逐一比对。4.3 灰度图与 dtype 处理灰度图C1会先torch.tile复制成 3 通道再绘制utils.py#L381-L382浮点输入会先转uint8绘制最后再转回原浮点 dtypeutils.py#L384-L386、utils.py#L412-L413。test_draw_boxes_dtypestest/test_utils.py#L218-L232验证了 uint8 与 float 输入输出 dtype 的一致性与像素一致性容差 ±1。4.4 官方示例手工框与模型输出框示例 gallery/others/plot_visualization_utils.py#L56-L101 演示了两种用法。手工指定框from torchvision.utils import draw_bounding_boxes boxes torch.tensor([[50, 50, 100, 200], [210, 150, 350, 430]], dtypetorch.float) colors [blue, yellow] result draw_bounding_boxes(dog1_int, boxes, colorscolors, width5) show(result)以及对接 Faster R-CNN 的真实推理输出先按置信度过滤再绘制from torchvision.models.detection import fasterrcnn_resnet50_fpn, FasterRCNN_ResNet50_FPN_Weights weights FasterRCNN_ResNet50_FPN_Weights.DEFAULT transforms weights.transforms() images [transforms(d) for d in dog_list] model fasterrcnn_resnet50_fpn(weightsweights, progressFalse) model model.eval() outputs model(images) score_threshold .8 dogs_with_boxes [ draw_bounding_boxes(dog_int, boxesoutput[boxes][output[scores] score_threshold], width4) for dog_int, output in zip(dog_list, outputs) ] show(dogs_with_boxes)五、draw_segmentation_masks分割掩码可视化5.1 签名与参数draw_segmentation_masks( image, # (3, H, W)uint8 或 float必须 RGB masks, # (H, W) 或 (num_masks, H, W)dtype 必须为 bool alpha0.8, # 透明度0 全透明1 不透明 colorsNone, # 掩码颜色单个或列表默认自动生成 ) - torch.Tensor5.2 输入约束源码校验utils.py#L444-L459比检测框更严格图像必须是 3 通道 RGBPass an RGB image. Other Image formats are not supported掩码张量必须为booldtype若传入 2D 掩码会先升维为(1, H, W)掩码的空间尺寸必须与图像完全一致。5.3 混合原理与重叠处理绘制算法在 utils.py#L461-L483值得拆解先计算overlapping_masks masks.sum(dim0) 1找出被多个掩码同时覆盖的像素对每张掩码把其覆盖像素直接赋为对应颜色img_to_draw[:, mask] color[:, None]重叠区域强制置 0黑色img_to_draw[:, overlapping_masks] 0最终混合out image * (1 - alpha) img_to_draw * alpha输出保持与原图相同的 dtype。test_draw_segmentation_maskstest/test_utils.py#L308-L360对颜色、透明度、设备做了组合参数化验证并逐一断言非掩码区像素不变、掩码区像素等于颜色与背景的alpha插值、重叠区等于image * (1 - alpha)。5.4 官方示例语义分割与实例分割示例对两类分割模型分别演示gallery/others/plot_visualization_utils.py#L104-L351。语义分割FCN ResNet-50模型输出为(batch, num_classes, H, W)的非归一化分数需先 softmax 归一化为概率再与类别索引比较得到布尔掩码——对于每个像素dog 是否是最可能的类别from torchvision.models.segmentation import fcn_resnet50, FCN_ResNet50_Weights weights FCN_ResNet50_Weights.DEFAULT transforms weights.transforms(resize_sizeNone) model fcn_resnet50(weightsweights, progressFalse) model model.eval() batch torch.stack([transforms(d) for d in dog_list]) output model(batch)[out] normalized_masks torch.nn.functional.softmax(output, dim1) boolean_dog_masks (normalized_masks.argmax(1) sem_class_to_idx[dog]) dogs_with_masks [ draw_segmentation_masks(img, masksmask, alpha0.7) for img, mask in zip(dog_list, boolean_dog_masks) ]实例分割Mask R-CNN输出为每实例的masks概率图布尔化需按 0.5 概率阈值proba_threshold 0.5 score_threshold .75 boolean_masks [ out[masks][out[scores] score_threshold] proba_threshold for out in output ] dogs_with_masks [ draw_segmentation_masks(img, mask.squeeze(1)) for img, mask in zip(dog_list, boolean_masks) ]示例还特别强调想要去除背景掩码只需传入masks[1:]背景类别索引为 0实例分割掩码多出的单例维度需squeeze(1)去除。六、draw_keypoints关键点与骨架连接6.1 签名与参数draw_keypoints( image, # (3, H, W)uint8 或 float必须 RGB keypoints, # (num_instances, K, 2)[x, y] 格式 connectivityNone, # [(i, j), ...] 关键点连接对 colorsNone, # 单一颜色字符串/RGB 元组 radius2, # 关键点圆点半径 width3, # 连接线宽度 visibilityNone, # (num_instances, K) 可见性张量 ) - torch.Tensor6.2 visibility按可见性选择性绘制visibility是姿态任务特有的能力utils.py#L545-L557默认None表示所有关键点可见传入值会被转换为bool也兼容(num_instances, K, 1)形状自动squeeze(-1)不可见的关键点既不画圆点也不参与连接线绘制连接的排除按实例独立计算。绘制逻辑utils.py#L571-L594用 PIL 的draw.ellipse画关键点圆、draw.line画连接线。与draw_bounding_boxes相同浮点输入同样会先转 uint8 绘制再转回。6.3 官方示例骨架连接与可见性实战示例 gallery/others/plot_visualization_utils.py#L354-L522 使用 KeypointRCNN 检测人体关键点并按 COCO 关键点索引定义人类骨架连接索引含义0 鼻、1/2 眼、3/4 耳、5/6 肩、7/8 肘、9/10 腕、11/12 髋、13/14 膝、15/16 踝connect_skeleton [ (0, 1), (0, 2), (1, 3), (2, 4), (0, 5), (0, 6), (5, 7), (6, 8), (7, 9), (8, 10), (5, 11), (6, 12), (11, 13), (12, 14), (13, 15), (14, 16) ] res draw_keypoints(person_int, keypoints, connectivityconnect_skeleton, colorsblue, radius4, width3)对于模型漏检的情形示例展示了如何从模型输出中切分出坐标与可见性KeypointRCNN 输出的第三维就是可见性标志再传入visibility参数使漏检点及其相关连线被自动跳过coordinates, visibility prediction.split([2, 1], dim-1) visibility visibility.bool() res draw_keypoints( person_int, coordinates, visibilityvisibility, connectivityconnect_skeleton, colorsblue, radius4, width3, )示例强调大多数 TorchVision 关键点预测模型都会为每个预测返回可见性可以直接拿来用。对应的测试test_draw_keypoints_visibilitytest/test_utils.py#L484-L528覆盖了bool/float 两种可见性输入与不同连接组合并断言输入张量均不被原地修改。七、flow_to_image光流场的彩色编码7.1 签名与输入约束flow_to_image(flow: torch.Tensor) - torch.Tensor输入形状为(N, 2, H, W)或(2, H, W)dtype 必须为torch.float输出为uint8的 RGB 图形状(N, 3, H, W)或(3, H, W)与输入维度对应。任何形状或 dtype 不符都会抛ValueErrorutils.py#L616-L624对应的test_flow_to_image_errors用五组非法输入做了穷举验证test/test_utils.py#L626-L638。7.2 编码原理颜色轮实现分两步utils.py#L604-L667用整个 batch 的最大位移范数归一化除以max_norm epsilonepsilon取 float 的机器精度防除零调用_normalized_flow_to_image基于方向与大小两个通道编码颜色方向atan2(-fy, -fx) / π映射到颜色轮索引通过k0/k1两近邻做线性插值大小1 - norm * (1 - col)让位移越大颜色越褪白位移为 0 时保持原色。颜色轮_make_colorwheelutils.py#L670-L714按 Baker 等人《A Database and Evaluation Methodology for Optical Flow》(ICCV 2007) 中描述的 55 段方案构造各段长度 RY15、YG6、GC4、CB11、BM13、MR6实现为 55×3 的 RGB 张量。这套可视化代码源自 tomrunia 的开源 OpticalFlow_Visualization 项目源码注释中已标明出处。test_flow_to_imagetest/test_utils.py#L603-L623用一个方向向外辐射的合成光场与test/assets/expected_flow.pt中预生成的基准输出逐元素比对保证单图与 batch 两种模式都精确一致。八、测试与质量保障该模块的可视化正确性由 test/test_utils.py共 642 行系统保障值得在引入新功能时参考其模式基准图比对绘制结果与test/assets/fakedata/下预生成的 PNG 基准图逐像素比对如draw_boxes_vanilla.png、draw_keypoint_vanilla.png、draw_boxes_util.png等。基准图不存在时测试会自动生成之后作为回归基线使用。非原地修改断言几乎所有测试都克隆输入并在调用后assert_equal确保make_grid、draw_bounding_boxes、draw_keypoints等不会污染用户的张量。dtype 往返一致性float 输入经转 uint8 绘制再转回后与 uint8 路径结果在容差 ±1 内一致。错误路径覆盖对每个非法输入非 Tensor、批次图像、通道数错误、标签/颜色数量不匹配、框坐标逆序、可见性形状不符等都断言抛出带特定消息的异常。参数化组合颜色None/名称/十六进制/RGB 元组/列表、透明度alpha ∈ {0, 0.5, 0.7, 1}、设备CPU/CUDA均通过pytest.mark.parametrize交叉覆盖。九、总结与实战建议torchvision.utils六个公开函数覆盖了计算机视觉开发中最常见的四类可视化需求全部以纯张量进、张量出的形态设计天然适配 PyTorch 的模型推理链路批量展示make_gridsave_image一条命令把(B, C, H, W)批次变成网格图片落盘检测draw_bounding_boxes支持普通框与旋转框(N, 8)、自定义颜色/字体/标签样式配合置信度阈值即可画出干净的检测结果分割draw_segmentation_masks直接吃 bool 掩码语义分割argmax 比较与实例分割概率阈值 0.5都可轻松接入姿态draw_keypoints的connectivity与visibility组合可在一行代码内完成骨架绘制并自动跳过不可见点光流flow_to_image内置颜色轮编码输入归一化由函数自动完成。实践中的关键约束可归纳为绘制函数只接受单张3D 图像批次需逐张调用掩码必须是 bool检测框需(xmin, ymin, xmax, ymax)顺序save_image配合文件对象时必须显式传format。沿着 docs/source/utils.rst 的指引进一步阅读 可视化完整示例 与 模块单元测试即可掌握从模型输出到专业可视化结果的全部细节。赞分享计算机视觉深度学习图像处理数据集【免费下载链接】visionDatasets, Transforms and Models specific to Computer Vision项目地址https://gitcode.com/gh_mirrors/vi/vision点击查看免费下载相关推荐U-Net预测结果可视化如何解读分割掩码U Net预测结果可视化如何解读分割掩码 U Net是一种用于图像分割的深度学习框架特别擅长处理医学图像和生物医学图像。通过Keras实现的U Net模型能示例工程深度学习计算机视觉人工智能关键点检测与实例分割实战应用关键点检测与实例分割实战应用 本文详细介绍了三种重要的计算机视觉技术DeepPose回归式关键点检测、HRNet高分辨率特征保持和Mask R CNN实例分割示例工程MMPose关键点检测中的编解码器详解MMPose关键点检测中的编解码器详解 编解码器在关键点检测中的作用 在计算机视觉领域人体姿态估计Human Pose Estimation是一项基础且重计算机视觉人工智能深度学习上一篇faiss IVF 系列索引基准测试实战:bench_all_ivf 的构建、集群运行与精度-速度-压缩权衡分析下一篇Allegro ClearML 实验日志记录全指南从基础到高级应用创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

关于本文作者

来自尧图内容编辑团队

尧图内容编辑团队 内容团队

尧图内容编辑团队

本文由尧图网络内容编辑团队执笔。团队由资深项目经理、前端工程师与设计师组成,所有内容均来自亲手交付的真实项目,先讲清问题、再给出可落地的解法。尧图深耕北京网站建设十年,服务过京华建材集团、智造科技等各行业客户,把一线经验沉淀为可复用的行业观察。

  • 十年建站经验,覆盖建材、制造、服务、文创等
  • 项目经理把关选题与事实准确性
  • 工程师与设计师联合撰写专业细节
  • 统一编辑规范,保证文风与排版一致
  • 每月复盘转化数据,迭代选题方向

延伸阅读

相关资讯与近期热门内容

深度阅读推荐

建站决策前值得细读的三篇

网站改版的5个关键决策
2024-08-12

网站改版的5个关键决策

什么时候该改版、改到什么程度、如何避免流量掉光,京华建材集团改版复盘给出答案。

获取专属建站方案

看完文章,把您的行业与预算告诉我们,免费获取一份量身定制的官网建设方案与报价。

立即免费咨询