第1节 · AI 编程规范(Vibe Coding)

🎯 掌握 AI 编程开发规范体系(standards)、pytest 测试与 CI/CD,实战银行营销系统。
📖 笔记

第二天笔记

第二天

1、今日目标

image-20260606204049843


2、昨日回顾

image-20260607085829809

image-20260607090603419

2.1 CI(持续集成)

每当有人(或AI)往代码仓库推送代码, CI就自动启动一个检查员。负责拉取代码, 安装依赖、运行语法/格式检查, 跑单元测试、

计算覆盖率、尝试构建镜像。

CI是否合并代码: 建议人工做一下review, 然后手动合并代码

2.2 CD(持续部署)

代码经过CI-->人工合并代码到主分支, 然后CD自动接手。 它会把代码打包, 推送到服务器,打镜像, 开启服务, 做一下健康检查,确认服务可以正常使用。

2.3 Standards规范下的AI编程流程

AI在本地开发-->PR提交-->CI全绿-->人工Review -->人工合并到main主分支-->CD自动触发-->服务器部署+健康检查

2.4 Standards模板文件

  • README.md: 入口文件, 所有新项目必须先读它
  • 第一次开工prompt是启动项目的标准, 能确保AI自动填写 00-project-context.md、01-requirements.md、PROGRESS.md
  • 00-project-context.md、01-requirements.md、PROGRESS.md文件会更新, 其他文件采用默认内容, 不更新。
请先读取 standards/README.md,并按它的顺序读取 00/01/PROGRESS 与 02~06。
这是一个新项目。我的项目需求是:
<在这里写项目目标、技术栈、希望实现的功能、部署方式>

请先做三件事:
1. 填写 standards/00-project-context.md。
2. 把需求整理进 standards/01-requirements.md,写成用户故事和验收标准。
3. 初始化 standards/PROGRESS.md,列出第一批 TODO。

完成后先停下让我确认,不要直接写代码。
  • 00-project-context.md: 读完README.md以后读取的此文件, 记录项目根本信息(项目名称、目标、核心功能、项目结构、技术选型、不变约束、质量门槛)。这个文件AI会自动填写

  • 01-requirements.md: 唯一的需求来源(Features/Bug/Tech Debt)。这个文件AI会自动填写

  • PROGRESS.md: 作为“存档点”, 记录当前进度、待办、关键决策、已知坑、里程碑, 方便AI或开发者中断后快速恢复。这个文件AI会自动填写, 填写的时机是AI在每次完成一个模块、解决一个故障、或者结束会话时, 主动更新此文件。

  • 当前状态

  • 待办清单

  • 关键决策

  • 已知坑

  • 里程碑

  • 02-coding-standards.md:代码标准, 统一代码风格(命名规范,函数与模块、注释、密钥管理、错误处理)

  • 03-testing-standards.md:测试标准,测试分层、用例用法、好测试的5条铁律、本地测试由谁执行

  • 04-git-workflow.md:统一团队协作方式
  • 分支模型、分支命令
  • 认证区分
  • 提交信息规范
  • PR三道闸(机器闸、人工闸、策略闸)
  • 冲突处理
  • 05-cicd-standards.md:CI/CD自动化的实现标准
  • CD与CD
  • CI门禁
  • CD标准
  • 依赖与镜像构建
  • 远程部署
  • Secrets清单
  • 排错顺序
  • 本地环境准备
  • 06-ai-collab-protocol.md:AI与人类协作的协议
  • 六步交付流程
  • AI开工前读取内容
  • AI的标准开场prompt
  • AI的权限边界
  • 故障反哺
  • 会话结束prompt

3、今日内容

3.1 standards驱动AI实战背景

实战背景:

基于银行数据,通过AI模型能力,过滤筛选高潜客户,进行外呼,实现高转化,降低运营成本。

step1:构建目录以及存放standards

创建项目目录:banksys_lixiaohua

把standard放到我们的项目目录

3.2 standards驱动AI实战开发

实战背景:

基于银行数据,通过AI模型能力,过滤筛选高潜客户,进行外呼,实现高转化,降低运营成本。

image-20260607091009371

step1:构建目录以及存放standards

创建项目目录:banksys_lixiaohua

把standard放到我们的项目目录

banksys_lixiaohua

step2:prompt

有服务器

项目路径:E:\develop\vibe_coding\banksys_lixiaohua
先读 standards/README.md,以及 standards/00-project-context.md、01-requirements.md、PROGRESS.md 和规范 02~06。
这是一个全新项目,需求基于银行营销数据《E:\develop\vibe_coding\banksys_lixiaohua\data》帮我实现一个web应用,但是含有两个功能:
--第一个数据分析交互页面
--第二个基于数据来离线训练出模型之后,构建一个在线预测系统,通过用户输入(点选形式进行输入),得到对应预测是否会认购。
技术栈:Python 3.11、streamlit、pytest、ruff、Docker。
端口:采用8888
仓库名称以及docker容器名称都是: banksys_lixiaohua
仓库是: 开源仓库
目标:跑通完整 CI + CD。
请据此填好 standards/00-project-context.md 与 01-requirements.md(写成带验收标准的用户故事),
并在 PROGRESS.md 写下第一批 TODO。先停下让我确认,不要开始写代码。

image-20260820190454214

没有服务器


项目路径:E:\develop\vibe_coding\banksys_lixiaohua
先读 standards/README.md,以及 standards/00-project-context.md、01-requirements.md、PROGRESS.md 和规范 02~06。
这是一个全新项目,需求基于银行营销数据《E:\develop\vibe_coding\banksys_lixiaohua\data》帮我实现一个web应用,但是含有两个功能:
--第一个数据分析交互页面
--第二个基于数据来离线训练出模型之后,构建一个在线预测系统,通过用户输入(点选形式进行输入),得到对应预测是否会认购。
技术栈:Python 3.11、streamlit、pytest、ruff、Docker。
仓库名称以及docker容器名称都是: banksys_lixiaohua
仓库是: 开源仓库
端口:采用8888
目标:跑通 CI ,不做CD,本地部署,启动后把url地址发给我,我来验证。
请据此填好 standards/00-project-context.md 与 01-requirements.md(写成带验收标准的用户故事),
并在 PROGRESS.md 写下第一批 TODO。先停下让我确认,不要开始写代码。

3.3 openclaw

3.3.1 openclaw是什么

openclaw:一个AI 支持多智能体系统级别智能体平台。

支持多渠道

支持多种模型

支持私有化部署

.....

openclaw应用场景

嵌入到公司业务或者个人学习与工作方方面面。

openclaw 定时做数据分析报告

openclaw做运维

openclaw开发代码

openclaw帮我构建情报系统

....

3.3.2 openclaw安装

3.3.2.1 macOS / Linux / WSL

curl -fsSL https://openclaw.ai/install.sh | bash

3.3.2.2 Windows(PowerShell,以管理员身份运行)

# 安装命令, 注意是powershell, 不是cmd
iwr -useb https://openclaw.ai/install.ps1 | iex

3.3.2.3 Windows (Cmd, 以管理员身份运行)

如果无法powershell安装, 可以通过以下命令在cmd窗口安装

npm install -g openclaw@latest

如果网速不行, 加上镜像源

npm install -g openclaw@latest --registry=https://registry.npmmirror.com

3.3.3 飞书与openclaw配置

3.3.3.1 创建一个飞书机器人(推荐: 选择从智能图模板开始)

获取app id与secret

https://open.feishu.cn/?lang=zh-CN, 进入浏览器页面后, 点击右上角开发者后台

image-20260801225814442

image-20260801230153185

image-20260801230255322

image-20260801230331141

  • 事件配置(默认即可,已经添加了很多事件, 如有需要可以自行添加其他事件)

image-20260801230842512

  • 回调配置(默认即可, 已开启长连接接收回调)

image-20260801231024620

  • 权限配置(默认即可)

image-20260801231413662

  • 版本发布与审核(默认即可, 不用管)

image-20260801231509977

  • 验证: 创建完飞书机器人以后, 你就可以在飞书客户端的-->开发者小助手-->打开应用, 找到你创建的机器人了

image-20260802003554433

image-20260802003645176

3.3.3.2 openclaw配置与飞书的通信策略

在命令窗口(不需要管理员), 输入以下命令

# 这条命令将飞书的私聊(直接消息) 策略设置为 pairing(配对模式)。
# 为什么要这么设:这是一种安全准入机制,防止未经授权的陌生人随意私聊你的机器人,避免资源滥用或信息泄露。适合企业内部核心功能使用
openclaw config set channels.feishu.dmPolicy "pairing"
# 这条命令将飞书的群聊策略设置为 open(开放模式)
# 只要机器人被拉入任何一个飞书群,并且在群里被 @ 提及(或者监听到特定关键词),机器人就会自动响应。
# 群内所有成员都可以通过 @机器人 的方式提问,无需进行任何配对、验证或审批流程。
openclaw config set channels.feishu.groupPolicy "open"

#查整块飞书配置
openclaw config get channels.feishu
#以下用于查看通道状态
openclaw channels status

image-20260801232507524

3.3.3.3 openclaw配置

openclaw onboard 进行交互式配置

第一个配置了模型

第二个配置了飞书

第三个配置搜索

3.3.3.3.1 配置大模型
#通过以下命令,进行一个交互设置:
openclaw onboard --classic
  • 选quickstart

image-20260802000710218

  • 选一个大模型

image-20260802000818790

  • 输入完api key 完成配置

image-20260802000932084

3.3.3.3.2 配置飞书

继续, 配置飞书

image-20260802001010212

  • 输入飞书的app id, app secret, 然后选open

image-20260802001236676

##### 3.3.3.3 配置联网搜索

作用是给你的 AI 助手(Agent)开通“联网搜索”功能。

配置了这个,你的机器人就能上网查实时信息;不配置,它就只能靠之前的训练数据来回答问题。

  • 可以从中选择一个, duckduckGo是免费的

image-20260802005016599

3.3.3.3.4 skills安装

openclaw自动安装一些技能, 可以ctrl + C跳过, 不影响这次飞书的测试. 如果打算安装, 则耐心等待15分钟左右

image-20260802002934280

工具权限的相关命令


# ==================== OpenClaw 权限完整设置 ====================

echo "=== 1. 设置工具基础权限为 full(功能最完整)==="
openclaw config set tools.profile "full"

echo "=== 2. 设置 exec 为白名单模式(最推荐的安全方式)==="
openclaw config set tools.exec.security "allowlist"

echo "=== 3. 让不在白名单的命令必须人工确认(重要安全层)==="
openclaw config set tools.exec.ask "on-miss"

echo "=== 4. 更新 safeBins(保留你常用命令 + 常用安全工具)==="
openclaw config set tools.exec.safeBins '["ls","pwd","cat","echo","git","npm","node","python","python3","pip","pip3","conda","docker","docker-compose","mkdir","rm","cp","mv","touch","grep","find","curl","wget","ps","tar","unzip","zip","make","jq","head","tail","wc","claude"]'

echo "=== 5. 信任完整 conda 目录(解决 conda/python 警告)==="
openclaw config set tools.exec.safeBinTrustedDirs '["/root/miniconda3/bin", "/root/miniconda3/condabin"]'

echo "=== 6. 自动修复 doctor 警告(safeBinProfiles 等)==="
openclaw doctor --fix

echo "=== 7. 重启 gateway 让所有配置生效==="
openclaw gateway restart

echo "=== 设置完成!下面显示最终配置 ==="

echo "=== 查看工具总配置 ==="
openclaw config get tools

echo "=== 查看Exec 执行配置(最重要)==="
openclaw config get tools.exec

echo "=== Approvals 当前状态 ==="
openclaw approvals get

openclaw其他相关命令

#查看状态
openclaw status
#系统自检
openclaw doctor
#通过以下进行解决
openclaw doctor --fix

#也可以考虑用openclaw tui来自己修复自己
openclaw tui

#飞书插件启动
openclaw plugins enable feishu

#飞书插件查看
openclaw plugins list

# 重启 gateway 让所有配置生效===
openclaw gateway restart

工具权限配置

# 例:给 agents.list 第 0 个 agent 单独全量授权
openclaw config set agents.list.0.id "main"
openclaw config set agents.list.0.tools.allow '["read","write","edit","apply_patch","exec","process","web_search","web_fetch","memory_search","memory_get","sessions_list","sessions_history","sessions_send","sessions_spawn","sessions_yield","subagents","session_status","browser","canvas","message","cron","gateway","nodes","agents_list","image","tts"]'
openclaw config set agents.list.0.tools.exec.security "full"
openclaw config set agents.list.0.tools.exec.ask "off"

openclaw config get agents.list.0.tools   # 看某 agent 最终配置


# 以下如果是你不想让其具备相应的能力,则直接设置进行单独收回(则该 agent 不能跑命令也不能改文件)
#openclaw config set agents.list.0.tools.deny #'["exec","process","write","edit","apply_patch"]'

今日总结

主要银行实战ai代码开发:

基于规范standards驱动AI开发:

step1:创建项目、放入数据以及standards

step2:输入prompt

openclaw

如何安装:

第一件事:创建飞书app,给他授权

第二件事:配置app id、app secret到openclaw的agent上

openclaw tui进行测试

技术分享

drawio架构图

下载一个drawio:https://github.com/jgraph/drawio-desktop/releases

自己贴mermaid:调整图形 =》布局=》

image-20260607200358975

指定版本安装

采用以下方式安装指定版本,任何系统window、mac或者linux系统都可以;注意一定要先安装node js,保证有对应的环境:

# 1. 检查 Node.js
node -v

# 2. 卸载旧版
npm uninstall -g openclaw

# 3. 安装指定版本,-g 意思是 --global(全局安装)
npm install -g openclaw@2026.3.23-2

# 4. 自检 + 自动修复配置(最重要一步)
openclaw doctor --fix

# 5. 重启服务
openclaw gateway restart

# 6. 验证
openclaw --version


tips:

#查看所有历史版本
npm view openclaw versions
#查看当前openclaw的版本
npm view openclaw version

📖 笔记

第一天笔记

第1天

1、今日目标

image-20260606200535307

2、今日内容

2.1 关于AI编程

AI编程的本质是什么?

AI编程本质:自然语言到代码的过程。

2.2 关于ai编程三个阶段

三个阶段:

ai coding:辅助编程,以人类为主导,我们基本可以理解架构以及代码

vibe coding:氛围编程,基于自己想要实现效果,来完成对应代码开发,完成最小mvp版本。不关注代码细节。

agentic engineer:智能体工程编程范式,核心是定规则规范,

2.3 AI编程工具类别

  • CLI终端模式开发工具(claude code, copilot)
  • 可视化界面IDE编程工具(cursor, qoder, trae)

2.4 AI编程三件套

  • 大模型API key
  • 开发工具(trae/claude)
  • 语言运行环境

2.5 trae安装

https://www.trae.cn/

选择下载TRAE IDE

image-20260731170219231

2.6 claude安装

2.6.1 下载安装git

https://git-scm.com/install/windows

2.6.2 下载安装nodejs

下载安装nodejs https://nodejs.org/en/download/

检查是否安装成功

image-20260731170604013

2.6.3 安装claude code

#打开终端执行,以下是claudecode安装
npm install -g @anthropic-ai/claude-code
#claude-code版本验证
claude --version

2.6.4 配置大模型

将大模型配置到claude上,来实现真正意义上的代码开发。

配置完以后, 重新打开命令窗口, 才能生效

#1.设置 DeepSeek接口地址
setx ANTHROPIC_BASE_URL "https://api.deepseek.com/anthropic"
#2.设置认证Key(请将 sk-xxxx换为你的真实的 DeepSeekKey)
setx ANTHROPIC_AUTH_TOKEN "sk-6bbfa905bcc34516bb7f81978449b6de"
#3.设置超时时间(10分钟)
setx API_TIMEOUT_MS 1200000

#4.指定模型
setx ANTHROPIC_MODEL "deepseek-v4-pro[1m]"
setx ANTHROPIC_DEFAULT_OPUS_MODEL "deepseek-v4-pro[1m]"
setx ANTHROPIC_DEFAULT_SONNET_MODEL "deepseek-v4-pro[1m]"

setx ANTHROPIC_SMALL_FAST_MODEL "deepseek-v4-flash"
setx ANTHROPIC_DEFAULT_HAIKU_MODEL "deepseek-v4-flash"
setx CLAUDE_CODE_SUBAGENT_MODEL "deepseek-v4-flash"
#5.禁用非必要流量
setX CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC 1

#6.执行代码相关任务时,启用最高投入档位、最大算力 / 最深思考模式
setx CLAUDE_CODE_EFFORT_LEVEL "max"

2.4 AI编程体验

trae测试

基于提供的图片链接,用工具构建一个宠物相册,采用fastapi,然后帮我本地启动,端口号6666
https://qcloud.dpfile.com/pc/Nlz58E1Ql3ZRyl94gMrn3piAYThVPQyJVX9ls7-X1Ehjlfzlnaa5mKt0DviN9Ogk.jpg
https://qcloud.dpfile.com/pc/Cvx06heHNwAKicHSqD1yho_e5eR2HU4AQpv5IyeSfMxdhzWFKOuIXowzVsir5tSo.jpg
https://img2.baidu.com/it/u=1538619194,450600015&fm=253&app=138&f=JPEG?w=800&h=1200
https://img1.baidu.com/it/u=3009263628,4139439784&fm=253&fmt=auto&app=138&f=JPEG?w=800&h=1062

项目存放地址:E:\课程资料\狂野\加餐课\VibeCoding\codes\trae_project
项目名称为:cuteanimalsGemini

这个是我项目的环境解释器D:\soft\Anaconda3\envs\ctm_kg\python
  • claude测试

开发 一个小游戏

你是一个顶级的 3D 游戏开发者,从零开始创建一个可交互的开放世界骑马探索 demo。
核心需求:
玩家骑一匹白马在秋季枫叶森林 + 橙色花田中移动
有木屋村落、烟囱冒烟、路径、树木、草地
支持下雨粒子效果(可切换)
日夜循环或手动切换时间
完整控制:WASD 控制马的速度和转向,鼠标控制视角,滚轮缩放,R 键下马 / 上马
风格参考《荒野大镖客 2》但更风格化一点(不要太写实)

绘制动画

使用Manim可视化二次函数f(x)=x²-4x+3,要求:
1.显示函数表达式
2.绘制函数曲线
3.标注顶点和零点
4.用动画展示函数图像的绘制过程

调度系统

项目描述:
设计一个基于消息队列的生产者-消费者分布式任务调度框架,支持多 Worker 进程、任务持久化和失败重试。

核心功能需求:

使用 Redis 或 RabbitMQ 作为 Broker,实现任务发布(apply_async)和延时任务(countdown)。

Worker 支持多进程并发(multiprocessing.Pool),并实现心跳检测机制,自动剔除宕机节点。

实现任务状态追踪(PENDING/RETRY/SUCCESS/FAILED),结果存储后端支持 Redis 或 SQLite。

支持任务路由(不同队列绑定不同 Worker),以及优先级队列。

提供命令行 CLI 工具:scheduler start --workers=4 --queues=high,low。

约束条件:
不能直接使用 Celery 或 RQ 库,必须基于底层 pika 或 redis-py 自行封装协议。

2.5 AI编程实战

银行数据实战介绍:

实战背景:基于ai实现高潜力客户筛选,进行外呼营销,最大化提升复购率以及降低营销成本。

数据字段信息情况。

ci与cd

ci:持续集成

cd:持续部署

mian=>创建自己一个分支=》开发=》本地测试=》提交PR=》CI=》CD

2.6 三端配置

如果实现cicd,完成以下三端配置:

1、本地开发机:

本地生成公私密钥对

git

github cli

2、github

配置仓库secrets

3、服务器端

生成公私密钥对

2.6.1 本地开发端

2.6.1 本地生成密钥

目的: 将密钥放入github中, 建立本地与github的连接

如果没有密钥:在 CMD 里生成

#查看是否存在
dir %USERPROFILE%\.ssh
#检查用户目录下是否存在 .ssh 文件夹,如果不存在则创建它。
if not exist "%USERPROFILE%\.ssh" mkdir "%USERPROFILE%\.ssh"
#使用 ed25519 算法在指定路径生成一对 SSH 密钥(私钥和公钥),并添加指定注释administrator@your-pc。
ssh-keygen -t ed25519 -C "administrator@your-pc" -f "%USERPROFILE%\.ssh\id_ed25519"

提示 passphrase 时:

  • 直接回车 = 不设密码
  • 或输入密码再确认一次

生成后再查看:

dir %USERPROFILE%\.ssh
type %USERPROFILE%\.ssh\id_ed25519.pub

image-20260606200800059

2.6.2 本地公钥复制到github的SSH key中

image-20260801152802355

2.6.3 配置git

# ① 配置 Git 身份
git config --global user.name  "lixiaohua"
git config --global user.email "332217839@qq.com"
git config --global init.defaultBranch main

# ④ 验证与 GitHub 的连接
ssh -T git@github.com
# 首次问 yes/no → 输入 yes;成功显示 Hi <你的用户名>!

2.6.4 安装github cli

GitHub CLI (gh) 是一个强大的命令行工具,旨在让你在终端中无缝地完成GitHub上的各种工作,无需在浏览器和终端之间频繁切换。其核心功能清单如下:

📁 仓库管理

管理代码仓库是 gh 的基础功能。

  • 查看、创建、克隆和复刻(Fork)仓库
  • 归档、取消归档、重命名、同步和设置默认分支
  • 管理仓库部署密钥和自动链接

💬 议题 (Issues) 与拉取请求 (Pull Requests)

这是 gh 最核心和最常用的功能之一。

  • 完整管理:支持创建、关闭、编辑、查看、列出议题和拉取请求。
  • PR 核心操作:可以在命令行中审查、查看差异(diff)、合并、检视(checkout) 和重新打开拉取请求。
  • 评论与管理:支持为议题和PR添加、编辑评论,以及锁定、转移、指派责任人等。
#winget(Win10/11 自带,最简单,推荐)
winget install --id GitHub.cli
#重开一个 cmd 窗口(让 PATH 生效),然后:
gh --version

# 登录(一次性,用来发 PR / 建仓 / 看 CI)
gh auth login
#查看状态
gh auth status

gh auth login这个命令登录的时候, 按交互依次选:

  • What account do you want to log into? → GitHub.com
  • What is your preferred protocol for Git operations? → HTTPS(或 SSH,你已配好 SSH 可选 SSH)
  • Authenticate Git with your GitHub credentials? → Yes
  • How would you like to authenticate? → Login with a web browser
  • 终端给一个一次性码(如 XXXX-XXXX)→ 回车自动开浏览器 → 粘贴码 → 授权

2.6.5 本地创建远程仓库验证

1)创建项目目录
# 切换到你的工作目录
cd 你的工作目录
mkdir my_test_gh
cd my_test_gh

2)git初始化并提交
git init
echo #mytest_gh> README.md或者在本地新建一个文件, 里边写点内容
git add .
git commit -m "chore: init"

3)创建远程私有仓库, 并push推送
 gh repo create mytest_gh --private --source=. --remote=origin --push

4)查看远程仓库

2.6.2 服务器端

登录的url

https://webterminal-console.jdcloud.com/lavm-remote?regionId=cn-north-1&instanceId=lavm-6gl7eut3f3

image-20260801161737956

# 生成密钥对,-N "" 表示空密码
ssh-keygen -t ed25519 -C "deploy@ai_coding" -f ~/.ssh/deploy_key -N ""
mkdir -p ~/.ssh && chmod 700 ~/.ssh
##加入服务器白名单
cat ~/.ssh/deploy_key.pub >> ~/.ssh/authorized_keys 
chmod 600 ~/.ssh/authorized_keys
#获取私钥用于这个github仓库secrets配置
cat ~/.ssh/deploy_key

2.6.3 github端

  • 在github的仓库中, 点开settings

image-20260801163827845

  • 找到secret and variables-->Actions

image-20260801163920395

github仓库secret 配置

在指定仓库的---->settings--->Actions secrets and variables 中, 配置这三个变量

SSH_PRIVATE_KEY(服务器生成的私钥)

SSH_HOST(服务器的公网IP)

SSH_USER(root)

image-20260801000052831

2.7 standards开发规范

开发规范:提交规范、测试规范、代码规范、进度管理规范等=》给AI看

可读性优先:代码写给人看,其次才是机器执行;拒绝炫技、晦涩写法。

单一职责:函数 / 类 / 组件只做一件事。

最小重复:公共逻辑抽离,禁止复制粘贴代码。

防御性编程:参数校验、空值处理、异常捕获。

可测试:业务逻辑与 UI/IO 解耦,便于单元测试。

兼容性 & 性能前置思考。

约定大于配置:统一项目约定,减少个性化写法。

2.7.1 helloworld跑通验证cicd

第一步:standards存放地址,位于项目目录下

image-20260606200937041

第二步:prompt驱动ai开发

先读 standards/README.md,以及 standards/00-project-context.md、01-requirements.md、PROGRESS.md 和规范 02~06。

这是一个全新项目,需求很简单:用 Python + Flask 写一个最小 Web 服务,
有两个接口:GET / 返回一句欢迎 JSON;GET /health 返回 {"status":"ok"} 供部署后健康检查。端口是8003,通过端口可以打开网页。
技术栈:Python 3.11、Flask、pytest、ruff、Docker。目标:跑通完整 CI + CD。

请据此填好 standards/00-project-context.md 与 01-requirements.md(写成带验收标准的用户故事),
并在 PROGRESS.md 写下第一批 TODO。先停下让我确认,不要开始写代码。
  • 如果有pull requests, AI工具pull requests以后, 需要手动合并, 点开pull requests, 点击下方的merge pull requests按钮即可合并

image-20260801183705767

第三:查看结果

image-20260801204940148

image-20260606200915591

image-20260606200905943

image-20260606200925532

3、今日重点

AI编程:ai编程是通过自然语言来实现代码开发。

ai三个:

ai coding:确定功能模块驱动,人类把控细节,架构清楚。

vibe coding:氛围编程,可控性较差,对细节不清晰

agentic engineer:从架构和逻辑角度,我们人类是可以把控。

ai开发三件套:

大模型:api key

ai编程工具:cli终端、可视化界面工具

项目运行环境

cicd:整个链路整体做了梳理,包括开发流程。

cicd配置:

如何实现从规范化ai编程、测试以及完整ci与cd?
三端配置:

本地开发配置:git、 github cli、ssh、ai开发环境配置

服务端:生成公私秘钥

github:

账户settings:用于git提交代码

仓库settings:用于cd

standards:基于规范进行ai开发

💻 代码

开发规范模板(standards-template)

11 个文件
  • 📁 templates/
    • 📄 ISSUE_TEMPLATE.md 455B
      # Issue 模板
      
      ## 类型
      
      - [ ] Feature 功能
      - [ ] Bug 缺陷
      - [ ] Tech Debt 技术债
      - [ ] Docs 文档
      - [ ] CI/CD 流水线
      
      ## 背景 / 场景
      
      <谁在什么场景下遇到什么问题>
      
      ## 目标
      
      <这次希望达成什么>
      
      ## 验收标准
      
      - [ ] Given <前提>,When <动作>,Then <可验证结果>
      - [ ] <补充标准>
      
      ## 技术备注
      
      <可选:边界、风险、约束、参考资料>
      
      ## 优先级
      
      P0 / P1 / P2
      
    • 📄 PR_TEMPLATE.md 641B
      # Pull Request 模板
      
      ## What
      
      <这个 PR 做了什么>
      
      ## Why
      
      关联 Issue:`closes #<编号>`
      
      ## How to Test
      
      - [ ] 本地格式检查通过
      - [ ] 本地静态检查通过
      - [ ] 本地测试通过
      - [ ] CI 通过
      - [ ] CD / 部署验证通过(如适用)
      
      ## 风险与回滚
      
      <风险点;如果失败如何回滚>
      
      ## 自检清单
      
      - [ ] 没有硬编码密钥/密码
      - [ ] 没有调试 print / 死代码
      - [ ] PR 体积可控
      - [ ] commit message 符合规范
      - [ ] 已同步最新 main,无冲突
      - [ ] 已更新 standards/PROGRESS.md
      
      ---
      
      通过标准:CI 全绿 + Review 通过 + 分支保护通过。
      
  • 📄 00-project-context.md 3.2KB
    # 00 · 项目上下文 〔本项目活记忆 · AI 维护〕
    
    > **作用**:这是项目的“身份档案”。AI 接管项目时先读这里,了解项目目标、技术栈、目录、部署取值。
    > **更新时机**:架构、技术栈、目录结构、端口、部署目录、重要约束变化时更新。
    > **填写方式**:把 `<...>` 替换成真实内容;用不到的行删掉。
    
    ---
    
    ## 1. 项目是什么
    
    - **项目名称**:`<项目名称>`
    - **一句话目标**:`<这个项目要解决什么问题>`
    - **使用者/受益者**:`<谁使用,为谁创造价值>`
    - **核心功能**:
      - `<功能1>`
      - `<功能2>`
    - **输入/数据**(如有):`<数据来源、规模、是否敏感、是否进 Git>`
    
    ## 2. 技术栈
    
    | 层 | 选型 | 理由 |
    |---|---|---|
    | 语言/运行时 | `<如 Python 3.11 / Node 20 / Java 17>` | `<为什么>` |
    | Web/API 框架 | `<如 Flask / FastAPI / Express / Spring Boot>` | `<为什么>` |
    | 测试 | `<如 pytest / jest / junit>` | `<为什么>` |
    | 格式/静态检查 | `<如 ruff / eslint+prettier>` | `<为什么>` |
    | 打包/运行 | `<如 Docker / 直接进程 / systemd>` | `<为什么>` |
    | CI/CD | GitHub Actions | 通用、可视化、适合教学与团队协作 |
    
    ## 3. 目录地图
    
    ```text
    <project>/
    ├── standards/                 # AI 项目记忆与通用规范
    ├── <源码目录>/                 # 如 app.py / src/ / backend/
    ├── <测试目录>/                 # 如 test_app.py / tests/
    ├── requirements.txt            # 生产运行依赖(如 Python 项目)
    ├── requirements-dev.txt        # 本地/CI 检查依赖(如 Python 项目)
    ├── Dockerfile                  # 如需要容器部署
    ├── .github/workflows/
    │   ├── ci.yml
    │   └── cd.yml
    └── README.md
    ```
    
    > 新增目录前先更新本节,避免项目越做越散。
    
    ## 4. 质量门槛
    
    | 类型 | 本项目标准 |
    |---|---|
    | 格式检查 | `<如 ruff format --check .>` |
    | 静态检查 | `<如 ruff check .>` |
    | 单元测试 | `<如 pytest>` |
    | 覆盖率 | `<如 >=80%;没有就写“暂无”>` |
    | 构建 | `<如 docker build 成功>` |
    | 业务/模型指标 | `<如接口成功率、模型指标、业务规则;没有就写“无”>` |
    
    ## 5. 不变约束
    
    - 密钥、密码、私钥、Token **绝不写进代码或文档**,只进 GitHub Secrets / 环境变量。
    - 大文件、数据集、模型产物是否进 Git,必须在本节写清楚;默认不进 Git。
    - `main` 分支受保护,日常开发必须走 feature 分支 + PR。
    - CI 红灯不合并。
    
    ## 6. 部署/CI 占位符取值
    
    > `guides/` 和 workflow 里的通用占位符,在本项目里的真实值只写这里。
    
    | 占位符 | 本项目取值 | 说明 |
    |---|---|---|
    | `<APP>` | `<应用名/镜像名/容器名>` | 如 `<app-name>` |
    | `<DEPLOY_DIR>` | `<服务器部署目录>` | 如 `/opt/<app-name>` |
    | `<PORT>` | `<服务端口>` | 如 `8000` |
    | `<PYVER>` | `<Python/Node/Java 版本>` | 如 `3.11` |
    | `<HEALTHCHECK>` | `<健康检查地址>` | 如 `/health` |
    | `<SSH_USER>` | `<部署用户>` | 如 `root` 或 `deploy` |
    | `<SSH_HOST>` | `<服务器公网 IP 或域名>` | 不写敏感信息以外的密钥 |
    
  • 📄 01-requirements.md 2.5KB
    # 01 · 需求 / 活 PRD 〔本项目活记忆 · AI 维护〕
    
    > **作用**:这是本项目唯一的需求文档。所有新功能、缺陷、技术债都追加到这里,不要另起多个 PRD 文件。
    > **更新时机**:每次有新需求、需求变更、验收标准变化时更新。
    
    ---
    
    ## 1. 需求来源
    
    | 类型 | 来源 | 进入方式 |
    |---|---|---|
    | 功能需求 Feature | 用户 / 老师 / 产品 / 客户 | 写成用户故事 |
    | 缺陷 Bug | 测试 / 线上日志 / 用户反馈 | 写复现步骤和期望结果 |
    | 技术债 Tech Debt | 开发 / Review / CI/CD 故障 | 写影响和修复目标 |
    
    ---
    
    ## 2. Issue 生命周期
    
    | 阶段 | 状态 | 动作 |
    |---|---|---|
    | 提出 | Open | 写清场景、目标、验收标准 |
    | 排期 | Backlog / Todo | 决定优先级和负责人 |
    | 开发 | In Progress | 从 main 开 feature 分支 |
    | 评审 | In Review | 提 PR,等待 CI 和 Review |
    | 合并 | Done | PR 合并 main,自动关闭 Issue |
    | 验收 | Verified | 按验收标准确认 |
    
    **追踪规则**:分支名带 Issue 号,PR 描述写 `closes #<编号>`。
    
    ---
    
    ## 3. 用户故事模板
    
    ```text
    ### US-<编号> <一句话标题> · 状态: Backlog
    作为 <角色>,
    我想要 <能力>,
    以便 <价值>。
    
    验收标准:
    - AC1: Given <前提>,When <动作>,Then <可验证结果>。
    - AC2: <补充标准>
    
    技术备注:
    - <可选:约束、边界、风险>
    ```
    
    ---
    
    ## 4. 需求清单
    
    ### US-1 初始化项目工程化与 CI/CD · 状态: Backlog
    
    作为 **项目开发者**,
    我想要 项目具备基础工程结构、测试、CI 与 CD,
    以便 后续每次开发都能自动检查并自动部署。
    
    验收标准:
    - AC1: 从 `main` 开 feature 分支完成初始化,不直接 push main。
    - AC2: PR 触发 CI,至少包含格式检查、静态检查、单元测试、构建检查。
    - AC3: CI 全绿后合并 main。
    - AC4: 合并 main 自动触发 CD,部署后健康检查通过。
    - AC5: 完成后更新 `standards/PROGRESS.md`。
    
    ### US-2 <你的第一个业务需求> · 状态: Backlog
    
    作为 `<角色>`,
    我想要 `<能力>`,
    以便 `<价值>`。
    
    验收标准:
    - AC1: Given `<前提>`,When `<动作>`,Then `<可验证结果>`。
    
    ---
    
    ## 5. 非功能需求
    
    - **安全**:密钥只进 Secrets,不进 Git。
    - **可维护**:一需求一小 PR,避免大爆炸式提交。
    - **可测试**:核心逻辑必须有单元测试。
    - **可部署**:部署后必须有健康检查或等价验证。
    
  • 📄 02-coding-standards.md 2.1KB
    # 02 · 代码标准
    
    > **目标**:让人和 AI 写出的代码风格一致、可读、可测、可维护。CI 会自动检查其中可机检的部分。
    
    ---
    
    ## 1. 格式与静态检查
    
    - Python 项目推荐:`ruff format .` + `ruff check .`。
    - 其他语言按项目技术栈替换为对应工具,例如 ESLint/Prettier、Checkstyle、gofmt。
    - 行宽默认 100;缩进、import 分组、命名风格由工具强制。
    - 提交前本地必须通过格式与 lint;CI 会再卡一道。
    
    ## 2. 命名
    
    | 对象 | 推荐规则 | 示例 |
    |---|---|---|
    | 变量/函数 | `snake_case` | `load_config` |
    | 类 | `PascalCase` | `UserService` |
    | 常量 | `UPPER_SNAKE` | `DEFAULT_PORT` |
    | 私有实现 | 前缀 `_` | `_parse_row` |
    
    命名必须表达意图。禁止 `process2`、`new_func`、`tmp_final` 这类无法审查的名字。
    
    ## 3. 函数与模块
    
    - 一个函数只做一件事;超过约 40 行优先拆分。
    - 参数过多时使用配置对象/数据类。
    - 业务逻辑与 I/O 分离,核心逻辑优先写成纯函数,方便测试。
    - 新增目录或模块前,先确认能落进 `00-project-context.md` 的目录地图。
    
    ## 4. 注释
    
    - 不写废话注释,例如 `# 调用函数` 配 `call_func()`。
    - 注释只解释“为什么”和“权衡”,不重复“代码做了什么”。
    - 复杂业务规则、部署坑、兼容性原因必须写清楚。
    
    ## 5. 配置与密钥
    
    - 非敏感配置可进 Git,例如 `config.yaml`、`.env.example`。
    - 密钥、密码、Token、私钥绝不进 Git,只放 GitHub Secrets / 环境变量。
    - 代码读取密钥时使用环境变量,并在 README 写清变量名,不写真实值。
    
    ## 6. 错误处理
    
    - 禁止裸 `except: pass`。
    - 外部输入必须校验,错误信息要能帮助定位。
    - 部署脚本必须失败即停;不能让前一步失败后继续执行后续危险动作。
    
    ## 7. 提交前自检
    
    ```text
    [ ] 格式检查通过
    [ ] 静态检查通过
    [ ] 新增逻辑有测试
    [ ] 没有硬编码密钥
    [ ] 没有调试 print / 死代码 / 大段注释代码
    [ ] 目录结构仍符合 00-project-context.md
    ```
    
  • 📄 03-testing-standards.md 2.9KB
    # 03 · 测试标准
    
    > **目标**:定义“什么叫测够了”。开发自测是进入 CI 前的第一道防线。
    
    ---
    
    ## 1. 测试分层
    
    ```text
            /\        E2E 测试:少、慢、贵,覆盖真实用户路径
           /  \
          /----\      集成测试:测试模块/服务边界
         /------\
        /        \    单元测试:多、快、稳定,开发者必写
       /__________\
    ```
    
    - 单元测试优先覆盖核心逻辑。
    - 集成测试覆盖接口、数据库、外部服务边界。
    - E2E 测试只覆盖关键路径,不要用它替代单元测试。
    
    ## 2. 用例写法
    
    推荐 AAA / Given-When-Then:
    
    ```python
    def test_health_returns_ok():
        # Arrange
        client = app.test_client()
    
        # Act
        resp = client.get("/health")
    
        # Assert
        assert resp.status_code == 200
        assert resp.get_json()["status"] == "ok"
    ```
    
    ## 3. 好测试的 5 条铁律
    
    1. 一个用例只测一件事。
    2. 覆盖正常值、边界值、异常值。
    3. 可重复、无外部网络依赖;有随机性必须固定种子。
    4. 测试名能表达场景和预期。
    5. 必须有真实断言,禁止 `assert result is not None` 凑数。
    
    ## 4. 覆盖率
    
    - 默认建议核心代码覆盖率 ≥ 80%。
    - 覆盖率不是目的,关键分支必须被覆盖。
    - 如果项目不适合覆盖率门槛,必须在 `00-project-context.md` 说明原因。
    
    ## 5. 项目特有质量门禁
    
    不同项目可追加不同质量门禁:
    
    | 项目类型 | 可选门禁示例 |
    |---|---|
    | Web/API | 健康检查、接口状态码、契约测试 |
    | 数据/ML | 数据校验、模型指标(AUC/F1/RMSE 等) |
    | 前端 | 组件测试、构建产物检查、可访问性检查 |
    | CLI/脚本 | 命令退出码、输入输出快照 |
    
    > 项目特有门禁写进 `00-project-context.md` 第 4 节,再写入 CI。
    
    ## 6. flaky test 纪律
    
    - 偶发失败必须根治,不能靠“重跑一次”掩盖。
    - 典型原因:时间依赖、随机数、测试顺序依赖、外部网络、共享状态未清理。
    
    ## 7. 本地测试由谁执行
    
    默认规则:
    
    | 环境 | 谁执行本地测试 | 怎么做 |
    |---|---|---|
    | AI 工具有终端权限(Cursor / Claude Code 等) | **AI 执行** | AI 自己运行格式检查、lint、pytest、构建等命令,读取日志,修复问题,把结果写回 `PROGRESS.md` |
    | AI 没有终端权限(纯聊天网页) | **学生执行** | 学生复制 AI 给出的命令到终端运行,把完整日志贴回 AI |
    | 涉及账号/权限/审批/Secrets | **学生确认** | AI 可指导,但不能代替人输入密码、私钥、审批生产发布 |
    
    AI 能执行命令时,不能只告诉学生“你去运行 pytest”。正确做法是:
    
    ```text
    AI 写代码 → AI 运行本地自检 → AI 根据日志修复 → AI 再运行 → 全绿后提交/PR → AI 记录结果
    ```
    
    每次本地自检至少记录:
    
    - 执行的命令。
    - 是否通过。
    - 失败时的关键错误。
    - 修复方式。
    - 最终通过结果。
    
  • 📄 04-git-workflow.md 4.4KB
    # 04 · Git 工作流与 PR/Review 标准
    
    > **目标**:定义代码如何从本地安全进入 `main`,并触发 CI/CD。课堂第一项功能必须完整演示本文件第 2 节。
    
    ---
    
    ## 1. 分支模型:GitHub Flow
    
    ```text
    main 受保护,永远可部署
     ├─ feature/<issue>-<desc>   新功能
     ├─ fix/<issue>-<desc>       普通缺陷
     ├─ docs/<issue>-<desc>      文档
     └─ chore/<issue>-<desc>     杂务/依赖/配置
    ```
    
    - 只有一条长期分支:`main`。
    - 日常开发必须从 `main` 开短分支,不能直接 push `main`。
    - feature 分支寿命建议 ≤ 2-3 天。
    - 分支命名:`<类型>/<issue号>-<英文短描述>`,例如 `feature/1-hello-api`。
    
    ---
    
    ## 2. 课堂必演示:从分支到 CD 的完整命令
    
    > 这段是学生必须亲手跑通的一条链。不要省略 PR,不要直接 push main。
    
    ```bash
    # 1. 确认在 main,并更新主干
    git switch main
    git pull
    
    # 2. 从 main 开功能分支
    git switch -c feature/1-first-feature
    
    # 3. 写代码 + 写测试后,本地自检
    ruff format --check .
    ruff check .
    pytest
    
    # 4. 提交
    git status
    git add .
    git commit -m "feat: add first feature"
    
    # 5. 推送 feature 分支
    git push -u origin feature/1-first-feature
    
    # 6. 创建 PR(AI 做到这里为止)
    gh pr create --base main --head feature/1-first-feature --title "feat: add first feature" --body "closes #1"
    
    # 7. 查看 CI(AI 汇报 CI 状态后停下,等人类审核)
    gh run list --limit 5
    gh run watch
    
    # 8. ★ 合并由人类执行 ★ —— 人工 Review 通过后,由人在网页点 "Merge",或人工运行:
    #    gh pr merge <PR号> --merge      # 默认保留分支,不加 --delete-branch
    #    AI 不得自行合并未获批准的 PR。
    
    # 9. 合并 main 后会触发 CD,查看部署流水线
    gh run list --limit 5
    ```
    
    > 如果课堂暂时没有开启分支保护,也要按这个流程演示 PR。分支保护是“强制执行”,PR 流程是“正确习惯”。
    
    > **合并是人类的动作,不是 AI 的动作。** AI 在第 6/7 步发起 PR、汇报 CI 后必须停下;
    > **只有人工 Review 通过,由人合并**,才触发 CD。AI 绝不替人按下 Merge。
    
    > **默认保留分支**(不加 `--delete-branch`),除非人类明确要求删除已合并分支。
    
    > **PR 必须用你本人 GitHub 账号发起**(`gh auth login` 登录的就是你本人),不要借他人账号代提。
    
    > **⚠️ 建仓后第一件事:配置 Secrets,否则 CD 必失败。**
    > 新仓库创建后(CD 第一次跑之前)立刻去
    > `GitHub → Settings → Secrets and variables → Actions → New repository secret`
    > 配好 `SSH_PRIVATE_KEY` / `SSH_HOST` / `SSH_USER`(详见 `05-cicd-standards.md` 第 5 节)。
    > 没有数据库等部署目标时可只配 SSH 三件套。配置时机:**在合并到 main 触发 CD 之前**。
    
    ---
    
    ## 3. 认证不要混
    
    | 动作 | 认证方式 | 说明 |
    |---|---|---|
    | `git push` / `git clone` | SSH 公钥或 HTTPS 凭证 | 负责搬代码 |
    | `gh pr create` / `gh run list` | GitHub CLI 登录 | 负责调 GitHub API |
    | CD SSH 到服务器 | GitHub Secrets 里的部署私钥 | 由 Actions runner 自动使用 |
    
    如果 SSH 22 端口卡住,可使用 GitHub SSH 443:
    
    ```bash
    ssh -T -p 443 git@ssh.github.com
    git remote set-url origin ssh://git@ssh.github.com:443/<账号>/<仓库>.git
    git push
    ```
    
    ---
    
    ## 4. 提交信息规范
    
    采用 Conventional Commits:
    
    | 前缀 | 含义 |
    |---|---|
    | `feat:` | 新功能 |
    | `fix:` | 修 bug |
    | `docs:` | 文档 |
    | `test:` | 测试 |
    | `refactor:` | 重构 |
    | `ci:` | CI/CD |
    | `build:` | 构建/依赖 |
    | `chore:` | 杂务 |
    
    一个 commit 只做一件事。
    
    ---
    
    ## 5. PR 通过三道闸
    
    | 闸 | 内容 | 不通过怎么办 |
    |---|---|---|
    | 机器闸 | CI 全绿 | 修代码/测试/配置 |
    | 人工闸 | 至少 1 人 Review | 修改并回复评论 |
    | 策略闸 | 无冲突、分支保护通过 | 同步 main / 解冲突 |
    
    PR 描述必须写:
    
    ```text
    ## What
    这个 PR 做了什么
    
    ## Why
    为什么需要
    
    ## Test
    - [ ] 本地检查通过
    - [ ] CI 通过
    
    closes #<issue号>
    ```
    
    ---
    
    ## 6. 冲突处理
    
    ```bash
    git switch main
    git pull
    git switch <你的分支>
    git rebase main
    
    # 如果有冲突:手动修改文件,删掉 <<<<<<< ======= >>>>>>> 标记
    git add <冲突文件>
    git rebase --continue
    
    # rebase 后推送
    git push --force-with-lease
    ```
    
    冲突处理后必须重新跑测试。
    
  • 📄 05-cicd-standards.md 7.1KB
    # 05 · CI/CD 标准
    
    > **目标**:定义最小可复现 CI/CD。CI 管“代码是否可靠”,CD 管“可靠代码如何自动上线”。
    
    ---
    
    ## 1. CI 与 CD
    
    | 名称 | 做什么 | 典型触发 |
    |---|---|---|
    | CI 持续集成 | 格式、lint、测试、覆盖率、构建 | PR / push |
    | CD 持续交付 | 构建可部署产物,部署前可人工审批 | main / tag |
    | CD 持续部署 | 合并 main 后自动部署上线 | main |
    
    课堂简版采用:
    
    ```text
    PR → CI 全绿 → 合并 main → CD 自动部署到服务器 → 健康检查
    ```
    
    ---
    
    ## 2. 最小 CI 门禁
    
    | 关卡 | 推荐命令 | 本地必跑 | CI 必跑 |
    |---|---|:---:|:---:|
    | 格式检查 | `ruff format --check .` 或等价工具 | ✅ | ✅ |
    | 静态检查 | `ruff check .` 或等价工具 | ✅ | ✅ |
    | 单元测试 | `pytest` 或等价工具 | ✅ | ✅ |
    | 覆盖率 | `pytest --cov --cov-fail-under=80` | ✅ | ✅ |
    | 项目特有门禁 | AUC/F1/接口契约/前端构建等 | ✅ | ✅ |
    | 构建 | `docker build ...` 或项目构建命令 | ❌ 不要求 | ✅ |
    
    CI 红灯时禁止合并。不要靠重跑掩盖问题。
    
    > **本地不强制装 Docker**:很多学员/开发机没有 Docker,本地装它只会增加复杂度。
    > 约定 **本地只跑前几项门禁,`docker build` 交给 CI(云端 runner 自带 Docker)与服务器(CD)**。
    > 这样本地零 Docker 负担,镜像构建仍有 CI 兜底。是否要本地构建,**按机器情况决定或先问开发者**。
    
    ---
    
    ## 3. 最小 CD 标准
    
    CD 必须满足:
    
    - 由 GitHub Actions runner 自动触发,不是人手动 SSH。
    - 只引用 GitHub Secrets,不硬编码密钥。
    - 部署脚本可重复执行。
    - 部署后必须健康检查。
    - 失败时 Actions 红灯,不能假成功。
    
    推荐触发:
    
    ```yaml
    on:
      push:
        branches: [ main ]
    ```
    
    ---
    
    ## 4. 依赖与镜像构建标准
    
    真实教学环境常见失败:服务器 `docker build` 时从 PyPI / npm / Docker Hub 下载依赖很慢或超时。
    
    通用标准:
    
    | 标准 | 做法 |
    |---|---|
    | 运行依赖与开发依赖拆分 | 生产镜像只安装运行依赖;测试/lint 工具只在 CI 安装 |
    | 镜像源可配置 | Dockerfile 支持镜像源参数,国内服务器可使用国内源 |
    | 构建失败即停 | 远程部署脚本第一行 `set -e` |
    | 健康检查 | 最后执行 `curl -fsS http://localhost:<PORT>/<HEALTHCHECK>` |
    | 可选预热 | 课前在服务器预拉基础镜像,减少等待 |
    
    Python 项目推荐:
    
    ```text
    requirements.txt      # 生产运行依赖
    requirements-dev.txt  # CI/本地检查依赖
    ```
    
    Dockerfile 推荐:
    
    ```dockerfile
    ARG PIP_INDEX_URL=https://pypi.org/simple
    RUN pip install --no-cache-dir --timeout 120 -i "${PIP_INDEX_URL}" -r requirements.txt
    ```
    
    远程部署脚本推荐(含幂等替换 + 端口自动回退):
    
    ```bash
    set -e
    docker build --build-arg PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple -t <APP>:latest .
    
    # 端口:容器内端口固定;主机端口优先 <PORT>,被占用就在预留区间自动找空闲端口。
    port_in_use() {
      ss -ltnH 2>/dev/null | grep -q ":$1 " && return 0
      docker ps --format "{{.Ports}}" 2>/dev/null | grep -q ":$1->" && return 0
      return 1
    }
    PORT=""
    for p in $(seq <PORT> <PORT_MAX>); do
      if ! port_in_use "$p"; then PORT="$p"; break; fi
    done
    [ -z "$PORT" ] && { echo "预留端口区间已全部占用,部署中止"; exit 1; }
    echo ">> 部署到主机端口 $PORT"
    
    docker rm -f <APP> 2>/dev/null || true   # 一步停删自身旧容器,幂等可重跑
    docker run -d --name <APP> --restart unless-stopped -p ${PORT}:<容器内端口> <APP>:latest
    sleep 3
    curl -fsS "http://localhost:${PORT}/<HEALTHCHECK>"
    echo ">> 部署成功:http://<SSH_HOST>:${PORT}/<HEALTHCHECK>"
    ```
    
    > **端口标准**:容器内端口固定,**主机端口可回退**(预留如 `<PORT>`-`<PORT_MAX>` 一段)。
    > 起容器用 `docker rm -f <APP>` 一步停删自身旧容器保证幂等;**不要**盲目删除占用端口的他人容器。
    > 防火墙/安全组要放行整段预留端口;最终落在哪个端口由 CD 日志打印。
    > 踩坑:`docker run` 报 `port is already allocated`(exit 125)就是主机端口被占用所致。
    
    ---
    
    ## 5. Secrets 标准
    
    常用 SSH 部署需要三个 GitHub Actions Secrets:
    
    | Secret | 含义 |
    |---|---|
    | `SSH_PRIVATE_KEY` | Actions runner 用来登录服务器的私钥全文 |
    | `SSH_HOST` | 服务器公网 IP 或域名 |
    | `SSH_USER` | 服务器登录用户,如 `root` 或 `deploy` |
    
    规则:
    
    - Secret 名大小写必须一致。
    - 私钥必须保留 `BEGIN/END` 和换行。
    - Secrets 必须在 CD run 开始前配置好;配置后旧 run 可能需要 rerun。
    
    ---
    
    ## 6. 本地环境准备(一次性)
    
    环境怎么建,**写进标准、可复现**。优先级:`uv` → `uv` 不行用 `conda` → 再不行用系统 `venv`。
    本课程学员更熟 `conda`,默认用 conda:
    
    ```bash
    # 1) 建并激活环境(示例环境名 envproj,Python 版本按项目 00 文档)
    conda create -y -n envproj python=3.11
    conda activate envproj
    
    # 2) 装依赖(国内用清华源更快)
    pip install -r requirements.txt -r requirements-dev.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
    ```
    
    > **坑:`CondaToSNonInteractiveError`(新版 conda 默认渠道未接受服务条款)**。
    > 解决,对三个渠道各执行一次:
    > ```bash
    > conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
    > conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
    > conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/msys2
    > ```
    
    ---
    
    ## 7. 排错顺序
    
    | 现象 | 优先检查 |
    |---|---|
    | `conda create` 报 ToS 错误 | 执行上面的 `conda tos accept`(main/r/msys2) |
    | Windows 控制台 `UnicodeEncodeError('gbk' ... \u25b6)` | 日志/打印含特殊符号;改用纯 ASCII,或 `set PYTHONIOENCODING=utf-8` / `chcp 65001`。注意 pytest 会捕获输出掩盖此坑,必须真跑一次脚本 |
    | CI 红:`FileNotFoundError` 找不到数据 | 数据被 `.gitignore` 排除,干净 runner 上没有。公开教学数据可入库;敏感数据则在 CI 里下载/造样本 |
    | `docker run` 报 `port is already allocated`(exit 125) | 主机端口被占用;容器内端口固定、主机端口在预留区间自动回退(见第 4 节);`docker rm -f <APP>` 幂等替换自身 |
    | `git push` 卡住/`Recv failure`/`Connection reset` | SSH 22 被拦可切 443;HTTPS 偶发抖动则重试几次 |
    | Actions 读不到变量 | Secret 名是否拼错、是否在正确仓库 |
    | SSH 登录失败 | 私钥/公钥是否匹配、`authorized_keys`、用户权限 |
    | rsync 成功但 docker build 失败 | 依赖下载、镜像源、Dockerfile |
    | 服务器本机 curl 通,外网打不开 | 安全组/防火墙是否放行端口(含回退端口段) |
    | CD 日志二次报错 | 部署脚本是否缺 `set -e` |
    
    ---
    
    ## 8. 标准总结
    
    1. CI 红灯不合并。
    2. CD 必须由 runner 自动触发,不是人手动 SSH。
    3. 密钥进 Secrets,脚本进 Git,环境手动配一次。
    4. 生产镜像只放运行依赖,开发工具留在 CI;**本地不强制装 Docker**。
    5. 每次真实故障必须写回 `PROGRESS.md` 和本标准。
    
  • 📄 06-ai-collab-protocol.md 6.8KB
    # 06 · AI 协作协议
    
    > **目标**:让任意 AI 工具都按同一套上下文、规则和闭环工作,避免换工具/换会话就丢上下文。
    > **核心**:下面第 1 节的「六步交付流程 + 确认门」是 AI 必须严格遵守的主干,**每到一个确认门必须停下,向人类汇报并等待确认,不得自行越过**。
    
    ---
    
    ## 1. 六步交付流程(必须逐步确认)★ 主干
    
    AI 驱动开发分为固定六步。**每一步都有一个「✋ 确认门」**:AI 完成该步动作后,先汇报结果,等人类确认/操作后才进入下一步。绝不一口气从头跑到尾。
    
    ```text
    ① 建仓+配Secrets → ② 开feature分支 → ③ 本地模块化开发(逐模块汇报)
         → ④ 本地CI自检(AI执行) → ⑤ 触发PR(CI在PR复检) → ⑥ 人工审核→合并→CD自动部署
       ✋确认1            ✋确认2          ✋确认3(每个模块)
                        ✋确认4           ✋确认5            ✋确认6(审核+部署结果)
    ```
    
    ### 第①步 · 建仓 + 配 Secrets
    
    - **AI 动作**:用 `gh` 创建 GitHub 仓库;`main` 只放最小引导提交(`.gitignore`/占位 README),**不在 main 上做真实开发**。
    - **✋ 确认门 1**:AI 报出仓库地址后,**必须提示人类去配置 GitHub Secrets**(新仓库默认为空,不配 CD 必失败):
      - `SSH_PRIVATE_KEY` / `SSH_HOST` / `SSH_USER`(见 `05` 第 5 节)。
      - AI 停下,**等人类确认「Secrets 已配置」**再继续(可用 `gh secret list` 核对)。
    
    ### 第②步 · 开 feature 分支
    
    - **AI 动作**:从最新 `main` 切出 `feature/<issue>-<desc>` 分支(类型见 `04`);**严禁直接在 main 上开发**。
    - **✋ 确认门 2**:报出分支名,确认后进入开发。
    
    ### 第③步 · 本地模块化开发 + 进度汇报
    
    - **AI 动作**:按 `00` 目录地图**一个模块一个模块**实现,每个模块同步写测试;每完成一个模块:
      - 更新 `PROGRESS.md`(已完成/下一步)。
      - 向人类**汇报该模块进度**(做了什么、关键决定、是否有坑)。
    - **✋ 确认门 3(每个模块)**:每个模块完成后汇报并确认,再做下一个模块。避免一次性写完全部再回头返工。
    
    ### 第④步 · 本地 CI 自检(由 AI 执行)
    
    - **AI 动作**:到「需要 CI」的节点,**AI 主动确认并亲自执行**本地门禁(具备终端权限时):
      - `ruff format --check .` + `ruff check .`
      - `pytest --cov --cov-fail-under=80`
      - 项目特有门禁(如模型 `--check-auc`)
      - **本地不强制 `docker build`**(交给 CI;见 `05` 第 2 节)。
    - **✋ 确认门 4**:汇报自检结果(覆盖率、指标);**全绿才进入下一步**,红则就地修复后重跑。
    
    ### 第⑤步 · 触发 PR(CI 在 PR 上复检)
    
    - **AI 动作**:`git push` 分支后,**AI 主动 `gh pr create`** 发起 PR(用本人账号);PR 事件触发 CI 在云端再跑一遍同样门禁 + `docker build`。
    - **✋ 确认门 5**:报出 PR 链接与 CI 状态(`gh run watch`);CI 绿才推进到审核。
    
    ### 第⑥步 · 人工审核 → 合并(人类) → CD 自动部署
    
    - **✋ AI 在此硬停**:AI 发起 PR、汇报 CI 绿之后**就停下**,**不得自己合并**。合并是人类的动作。
    - **人类动作**:**PR 必须由人工 Review**;审核通过后**由人类在网页点 Merge 或人工运行 `gh pr merge`**。
    - **分支默认保留**:合并不加 `--delete-branch`,除非人类明确要求删除。
    - **AI 动作(合并发生后)**:CD 自动触发(SSH→同步→构建→运行→健康检查;端口自愈见 `05`);AI 盯流水线。
    - **✋ 确认门 6**:汇报部署结果——**最终落地端口、`/health` 结果、访问地址**;失败则读日志定位根因,修复后回到第②步走一个新分支/PR。
    
    > 一句话:**建仓我做、Secrets 你配;分支我开(默认保留);模块我逐个写并汇报;本地 CI 我跑;PR 我发;
    > 审核与合并都由你做,我绝不替你按 Merge;合并后 CD 自动跑,我盯结果。**
    
    ---
    
    ## 2. AI 开工前必须读
    
    每次新会话,AI 必须按顺序读:
    
    1. `standards/README.md`
    2. `standards/00-project-context.md`
    3. `standards/01-requirements.md`
    4. `standards/PROGRESS.md`
    5. `standards/02`~`06`
    
    然后向人类复述:
    
    ```text
    我理解当前项目是:...
    当前阶段是(对应六步流程第几步):...
    下一步应该做:...
    我将遵守的关键约束是:...
    ```
    
    ---
    
    ## 3. 人给 AI 的标准开场 prompt
    
    ```prompt
    请先读取 standards/README.md,并按顺序读取 00/01/PROGRESS 与 02~06。
    读完后先复述你对项目、当前进度(对应 06 的六步流程第几步)、下一步的理解。
    然后严格按 06 的「六步交付流程 + 确认门」推进:
    - 每到一个确认门必须停下汇报并等我确认,不要一口气跑完。
    - 建仓后提示我配置 Secrets;从 main 开 feature 分支,不直接改 main。
    - 逐模块开发并汇报;本地 CI 由你执行(本地不强制 docker);CI 绿后你主动发 PR。
    - 发完 PR 你就停下,**不要自己合并**;PR 由我审核并由我合并,合并后 CD 自动跑,你盯结果并汇报端口与健康检查。
    - 分支默认保留,不要加 --delete-branch,除非我明确要求删。
    ```
    
    ---
    
    ## 4. AI 写代码边界
    
    AI 可以:
    
    - 新增/修改代码、测试、CI/CD 配置。
    - 在具备终端权限时运行本地测试、lint、构建命令,并根据日志修复。
    - 创建仓库、开分支、发起 PR(用本人账号)。
    - 更新 `00/01/PROGRESS`;根据真实故障更新 `02~06`(须说明原因)。
    
    AI 不可以:
    
    - 硬编码密钥。
    - 绕过 CI 或跳过测试。
    - 明明能执行命令,却只让学生手动跑本地测试。
    - 为了“看起来通过”删除测试或降低门槛。
    - **越过任何确认门**(尤其:未提示配 Secrets 就推进 CD;未经确认大规模重构)。
    - **自行合并任何 PR**——合并是人类的专属动作,AI 只发起 PR 并汇报,绝不替人按 Merge。
    - 合并时默认 `--delete-branch` 删分支——**默认保留分支**,除非人类明确要求删除。
    - 直接把项目特有内容写进通用规范,除非它已沉淀为可复用规则。
    
    ---
    
    ## 5. 故障反哺铁律
    
    CI/CD 一旦真实跑出问题,修完代码还不算结束,必须同步完成三件事:
    
    1. 写回 `PROGRESS.md`:现象、根因、修复、验证结果。
    2. 如是通用规则,写进 `02~06` 对应标准。
    3. 更新学生会照着操作的 README/课件/脚本。
    
    判断标准:
    
    ```text
    换一个学生、换一个 AI 工具,只读 standards + README,是否能避开这次坑并复现?
    ```
    
    如果不能,说明闭环未完成。
    
    ---
    
    ## 6. 会话结束 prompt
    
    ```prompt
    本次会话结束。请更新 standards/PROGRESS.md:
    1. 当前状态(对应六步流程第几步)
    2. 已完成事项
    3. 下一步 TODO
    4. 新增 ADR
    5. 新增 GOTCHAS
    保持时间倒序、简洁、可接力。
    ```
    
  • 📄 PROGRESS.md 1.5KB
    # PROGRESS · <项目名称> 〔本项目活记忆 · 状态机〕
    
    > **作用**:这是项目的“存档点”。任意 AI、任意重启会话,读它即可知道当前做到哪、下一步做什么、踩过什么坑。
    > **更新时机**:每完成一个有意义步骤、每次会话结束前。
    > **格式要求**:时间倒序,最新在上;短、准、可接力。
    
    ---
    
    ## 当前状态 (最后更新: <日期> · by <谁>)
    
    - **阶段**:`<初始化 / 开发中 / 联调 / 已上线>`
    - **上一步完成**:`<刚做完什么>`
    - **下一步 (TODO 第一条)**:`<马上要做什么>`
    - **阻塞项**:`<没有写“无”;有则写清责任人与下一动作>`
    
    ---
    
    ## 待办清单 (TODO,按优先级)
    
    - [ ] 初始化项目上下文:`00-project-context.md`
    - [ ] 确认需求与验收标准:`01-requirements.md`
    - [ ] 从 `main` 开第一条 feature 分支
    - [ ] 实现代码与测试
    - [ ] 本地自检通过
    - [ ] 提交并推送 feature 分支
    - [ ] 创建 PR,等待 CI
    - [ ] CI 全绿后合并 main
    - [ ] 验证 CD 自动部署
    - [ ] 会话结束前更新本文件
    
    ---
    
    ## 关键决策记录 (ADR)
    
    | 日期 | 决策 | 理由 |
    |---|---|---|
    | <日期> | <定了什么> | <为什么这么定> |
    
    ---
    
    ## 已知坑 (GOTCHAS)
    
    - <现象>:<根因>;解决:<怎么处理>;验证:<如何确认已修复>。
    
    ---
    
    ## 里程碑 (DONE)
    
    - [ ] <完成一项就勾选,写一句话说明结果>
    
    > 反臃肿:里程碑超过 15 条时,把更早内容合并成一行摘要,保持本文件可快速阅读。
    
  • 📄 README.md 3.6KB
    # standards-template · 学生项目可复制模板
    
    > **用途**:每个新项目复制一份本目录,改名为项目根目录下的 `standards/`。任意 AI 工具进入项目后,先读 `standards/README.md`,即可按统一规则继续开发、提交、PR、CI/CD 与存档。
    
    ---
    
    ## 1. 学生怎么用
    
    ```bash
    # 在你的项目根目录执行
    cp -r <课件目录>/standards-template ./standards
    ```
    
    Windows 可手动复制文件夹,复制后目录应是:
    
    ```text
    你的项目/
    ├── standards/
    │   ├── README.md
    │   ├── 00-project-context.md
    │   ├── 01-requirements.md
    │   ├── PROGRESS.md
    │   ├── 02-coding-standards.md
    │   ├── 03-testing-standards.md
    │   ├── 04-git-workflow.md
    │   ├── 05-cicd-standards.md
    │   ├── 06-ai-collab-protocol.md
    │   └── templates/
    └── <你的代码>
    ```
    
    ---
    
    ## 2. 两类文件,不要混
    
    | 类型 | 文件 | 是否每个项目要改 | 谁维护 |
    |---|---|---|---|
    | **项目活记忆** | `00-project-context.md`、`01-requirements.md`、`PROGRESS.md` | 必须按项目重填 | AI 写,人审 |
    | **通用规范** | `02`~`06`、`templates/` | 默认不改 | 团队维护 |
    
    > 简单说:`00/01/PROGRESS` 记录“这个项目是什么、要做什么、做到哪了”;`02~06` 规定“怎么写代码、怎么测试、怎么走 Git/PR/CI/CD”。
    
    ---
    
    ## 3. AI 每次会话读取顺序
    
    1. `00-project-context.md` — 项目身份、技术栈、目录地图、部署取值。
    2. `01-requirements.md` — 活 PRD,所有需求与验收标准。
    3. `PROGRESS.md` — 当前状态、下一步、决策、踩坑。
    4. `02`~`06` — 通用工程规范。
    5. `templates/` — Issue / PR 模板。
    
    ---
    
    ## 4. 第一次开工 prompt
    
    ```prompt
    请先读取 standards/README.md,并按它的顺序读取 00/01/PROGRESS 与 02~06。
    这是一个新项目。我的项目需求是:
    <在这里写项目目标、技术栈、希望实现的功能、部署方式>
    
    请先做三件事:
    1. 填写 standards/00-project-context.md。
    2. 把需求整理进 standards/01-requirements.md,写成用户故事和验收标准。
    3. 初始化 standards/PROGRESS.md,列出第一批 TODO。
    
    完成后先停下让我确认,不要直接写代码。
    ```
    
    ---
    
    ## 5. 标准开发闭环 · 固定六步(每步确认)
    
    AI 严格按 `06-ai-collab-protocol.md` 的「六步交付流程 + 确认门」推进,**每到确认门停下汇报、等人确认**:
    
    ```text
    ① 建仓 + 配 Secrets    AI 建仓 → ✋提示人类配 SSH_PRIVATE_KEY/HOST/USER(新仓库默认没有)
    ② 开 feature 分支       从 main 切分支,严禁直接改 main         ✋报分支名
    ③ 本地模块化开发        逐模块写代码+测试,每模块更新 PROGRESS  ✋每个模块汇报进度
    ④ 本地 CI 自检(AI 跑)  ruff + pytest + 覆盖率 + 项目门禁;本地不强制 docker  ✋全绿才继续
    ⑤ 触发 PR              push + AI 主动建 PR,CI 在 PR 上复检    ✋报 PR 链接 + CI 状态
    ⑥ 人工审核 →(人)合并 → CD  ✋AI 发完 PR 即停;Review 与 Merge 都由人做;合并触发 CD → ✋报端口/健康检查
    ```
    
    **课堂第一项功能必须完整演示一次这条链**,不能直接 push main;
    **合并是人类的动作——AI 绝不自行合并 PR;分支默认保留(不自动删)。**
    
    ---
    
    ## 6. 反臃肿纪律
    
    1. 需求只写进 `01-requirements.md`,不要到处新建 PRD。
    2. 进度、决策、坑只写进 `PROGRESS.md`。
    3. 一需求一分支一 PR,PR 尽量小于 400 行。
    4. CI 红灯不合并;flaky test 必须根治。
    5. 新增目录前先更新 `00-project-context.md` 的目录地图。
    
💻 代码

实战项目代码(银行营销系统 / pytest 教学)

58 个文件
  • 📁 banksys/
    • 📁 claude_code/
      • 📁 banksys_szai4/
        • 📁 .claude/
          • ⚙️ settings.local.json 196B
        • 📁 app/
          • 📄 __init__.py 79B
            """Bank marketing analysis and prediction system."""
            
            __version__ = "0.1.0"
            
          • 📄 main.py 1.3KB
            """Bank marketing analysis and prediction system.
            
            Main entry point for Streamlit application.
            """
            
            from pathlib import Path
            
            import streamlit as st
            
            st.set_page_config(
                page_title="Bank Marketing System",
                page_icon="🏦",
                layout="wide",
                initial_sidebar_state="expanded",
            )
            
            
            def main() -> None:
                """Render the home / landing page.
            
                Navigation between pages is handled by Streamlit's native multipage
                mechanism (files under ``app/pages/``), so this entry point only shows
                the landing content and a dataset snapshot.
                """
                st.title("🏦 Bank Marketing Analysis & Prediction")
            
                st.markdown(
                    """
                    ## Welcome
            
                    This system provides:
            
                    - **Data Analysis**: Explore customer demographics and marketing effectiveness
                    - **Prediction**: Predict customer subscription likelihood
            
                    Use the sidebar to navigate between pages.
                    """
                )
            
                data_path = Path("data") / "train.csv"
                if data_path.exists():
                    import pandas as pd
            
                    df = pd.read_csv(data_path)
                    col1, col2 = st.columns(2)
                    col1.metric("Total Records", f"{len(df):,}")
                    col2.metric("Columns", len(df.columns))
            
            
            if __name__ == "__main__":
                main()
            
        • 📁 data/
          • ⚙️ test.csv 901.4KB
        • 📁 standards/
          • 📄 00-project-context.md 5.0KB
            # 00 · 项目上下文 〔本项目活记忆 · AI 维护〕
            
            > **作用**:这是项目的"身份档案"。AI 接管项目时先读这里,了解项目目标、技术栈、目录、部署取值。
            > **更新时机**:架构、技术栈、目录结构、端口、部署目录、重要约束变化时更新。
            > **填写方式**:把 `<...>` 替换成真实内容;用不到的行删掉。
            
            ---
            
            ## 1. 项目是什么
            
            - **项目名称**: `banksys_szai4`
            - **一句话目标**: 基于银行营销数据构建可视化分析与认购预测的 Web 应用
            - **使用者/受益者**: 银行营销团队、数据分析师、业务决策者
            - **核心功能**:
              - 数据分析交互页面:展示客户特征分布、营销效果分析、数据探索
              - 在线预测系统:基于训练模型,通过点选式表单输入客户特征,预测认购意愿
            - **输入/数据**: 银行营销数据集(data/train.csv、data/test.csv),包含 21 个特征(年龄、职业、婚姻、教育、联系方式、经济指标等),目标变量为 subscribe(是否认购定期存款)。**数据进 Git,非敏感教学数据**
            
            ## 2. 技术栈
            
            | 层 | 选型 | 理由 |
            |---|---|---|
            | 语言/运行时 | Python 3.11 | 课程指定版本,生态成熟 |
            | Web/API 框架 | Streamlit 1.x | 快速构建数据应用,适合数据分析与模型演示 |
            | 测试 | pytest | Python 标准测试框架,配合插件强大 |
            | 格式/静态检查 | ruff | 极速 Python linter/formatter,一统格式与 lint |
            | 打包/运行 | Docker | 标准容器化,CI/CD 友好 |
            | CI/CD | GitHub Actions | 通用、可视化、适合教学与团队协作 |
            | 机器学习 | scikit-learn、pandas | 离线训练预测模型、数据处理 |
            
            ## 3. 目录地图
            
            ```text
            banksys_szai4/
            ├── standards/                 # AI 项目记忆与通用规范
            ├── app/                       # 应用主目录
            │   ├── __init__.py
            │   ├── main.py               # Streamlit 主入口
            │   ├── pages/                # 多页面结构
            │   │   ├── __init__.py
            │   │   ├── 01_data_analysis.py    # 数据分析页面
            │   │   └── 02_prediction.py        # 预测系统页面
            │   ├── models/               # 业务逻辑层
            │   │   ├── __init__.py
            │   │   ├── data_loader.py         # 数据加载
            │   │   ├── predictor.py           # 预测逻辑
            │   │   └── visualizer.py         # 可视化逻辑
            │   ├── ml/                   # 机器学习模块
            │   │   ├── __init__.py
            │   │   ├── train.py               # 离线训练脚本
            │   │   └── model/                 # 模型产物目录(.gitignore)
            │   └── utils/                # 工具函数
            │       └── __init__.py
            ├── tests/                     # 测试目录
            │   ├── __init__.py
            │   ├── test_data_loader.py
            │   ├── test_predictor.py
            │   ├── test_visualizer.py
            │   └── conftest.py
            ├── data/                      # 数据目录
            │   ├── train.csv              # 训练数据
            │   └── test.csv               # 测试数据
            ├── requirements.txt            # 生产运行依赖
            ├── requirements-dev.txt        # 本地/CI 检查依赖
            ├── Dockerfile                  # 容器镜像定义
            ├── .github/workflows/
            │   ├── ci.yml                 # 持续集成
            │   └── cd.yml                 # 持续交付
            ├── .gitignore
            ├── README.md
            └── pyproject.toml              # ruff 配置
            ```
            
            > 新增目录前先更新本节,避免项目越做越散。
            
            ## 4. 质量门槛
            
            | 类型 | 本项目标准 |
            |---|---|
            | 格式检查 | `ruff format --check .` |
            | 静态检查 | `ruff check .` |
            | 单元测试 | `pytest` |
            | 覆盖率 | 核心逻辑 ≥80%;UI 页面无覆盖率要求 |
            | 构建 | `docker build` 成功 |
            | 业务/模型指标 | 训练脚本输出模型指标(AUC、准确率等),预测服务响应时间 <1s |
            
            ## 5. 不变约束
            
            - 密钥、密码、私钥、Token **绝不写进代码或文档**,只进 GitHub Secrets / 环境变量。
            - **数据集进 Git**(教学用公开数据),但 **模型产物不进 Git**(ml/model/ 目录加入 .gitignore)。
            - `main` 分支受保护,日常开发必须走 feature 分支 + PR。
            - CI 红灯不合并。
            - 模型训练是离线操作(本地/CI 执行训练脚本),预测服务加载已训练模型。
            
            ## 6. 部署/CI 占位符取值
            
            > `guides/` 和 workflow 里的通用占位符,在本项目里的真实值只写这里。
            
            | 占位符 | 本项目取值 | 说明 |
            |---|---|---|
            | `<APP>` | `banksys` | 应用名/镜像名/容器名 |
            | `<DEPLOY_DIR>` | `/opt/banksys` | 服务器部署目录 |
            | `<PORT>` | `8004` | 服务端口(Streamlit 默认 8501,映射到主机 8004) |
            | `<PORT_MAX>` | `8010` | 端口回退上限 |
            | `<PYVER>` | `3.11` | Python 版本 |
            | `<HEALTHCHECK>` | `/_stcore/health` | Streamlit 健康检查端点 |
            | `<SSH_USER>` | `root` | 部署用户(按实际环境调整) |
            | `<SSH_HOST>` | `<待配置>` | 服务器公网 IP 或域名 |
            
          • 📄 01-requirements.md 7.9KB
            # 01 · 需求 / 活 PRD 〔本项目活记忆 · AI 维护〕
            
            > **作用**:这是本项目唯一的需求文档。所有新功能、缺陷、技术债都追加到这里,不要另起多个 PRD 文件。
            > **更新时机**:每次有新需求、需求变更、验收标准变化时更新。
            
            ---
            
            ## 1. 需求来源
            
            | 类型 | 来源 | 进入方式 |
            |---|---|---|
            | 功能需求 Feature | 课程作业要求 | 写成用户故事 |
            | 缺陷 Bug | 测试 / 线上日志 / 用户反馈 | 写复现步骤和期望结果 |
            | 技术债 Tech Debt | 开发 / Review / CI/CD 故障 | 写影响和修复目标 |
            
            ---
            
            ## 2. Issue 生命周期
            
            | 阶段 | 状态 | 动作 |
            |---|---|---|
            | 提出 | Open | 写清场景、目标、验收标准 |
            | 排期 | Backlog / Todo | 决定优先级和负责人 |
            | 开发 | In Progress | 从 main 开 feature 分支 |
            | 评审 | In Review | 提 PR,等待 CI 和 Review |
            | 合并 | Done | PR 合并 main,自动关闭 Issue |
            | 验收 | Verified | 按验收标准确认 |
            
            **追踪规则**:分支名带 Issue 号,PR 描述写 `closes #<编号>`。
            
            ---
            
            ## 3. 用户故事模板
            
            ```text
            ### US-<编号> <一句话标题> · 状态: Backlog
            作为 <角色>,
            我想要 <能力>,
            以便 <价值>。
            
            验收标准:
            - AC1: Given <前提>,When <动作>,Then <可验证结果>。
            - AC2: <补充标准>
            
            技术备注:
            - <可选:约束、边界、风险>
            ```
            
            ---
            
            ## 4. 需求清单
            
            ### US-1 初始化项目工程化与 CI/CD · 状态: Backlog
            
            作为 **项目开发者**,
            我想要 项目具备基础工程结构、测试、CI 与 CD,
            以便 后续每次开发都能自动检查并自动部署。
            
            验收标准:
            - AC1: 从 `main` 开 feature 分支完成初始化,不直接 push main。
            - AC2: PR 触发 CI,至少包含格式检查(ruff)、静态检查(ruff)、单元测试(pytest)、构建检查(docker build)。
            - AC3: CI 全绿后合并 main。
            - AC4: 合并 main 自动触发 CD,部署后健康检查通过(端口 8004)。
            - AC5: 完成后更新 `standards/PROGRESS.md`。
            
            技术备注:
            - 首个 US 必须完整演示六步交付流程。
            - 建仓后必须提示配置 GitHub Secrets (SSH_PRIVATE_KEY/SSH_HOST/SSH_USER)。
            
            ---
            
            ### US-2 数据加载与预处理模块 · 状态: Backlog
            
            作为 **开发者**,
            我想要 一个可复用的数据加载模块,
            以便 支持数据分析页面和模型训练的数据需求。
            
            验收标准:
            - AC1: Given 数据文件存在,When 调用 `load_train_data()` 和 `load_test_data()`,Then 返回 pandas DataFrame。
            - AC2: Given 原始数据,When 加载数据,Then 列包含 21 个特征(id/age/job/.../subscribe)。
            - AC3: Given 数据加载,When 访问数据,Then 能正确处理中英文列名、缺失值。
            - AC4: Given 数据加载模块,When 编写单元测试,Then 覆盖正常加载、文件不存在、空文件场景。
            - AC5: Given 数据目录,When 部署,Then data/ 目录在镜像内正确挂载。
            
            技术备注:
            - 数据文件是 UTF-8 编码的 CSV。
            - 需要处理可能的缺失值标记为 'unknown'/'nonexistent'。
            
            ---
            
            ### US-3 数据分析交互页面 · 状态: Backlog
            
            作为 **业务分析师**,
            我想要 通过可视化界面探索银行营销数据,
            以便 快速理解客户特征分布和营销效果。
            
            验收标准:
            - AC1: Given 访问应用首页或导航到"数据分析",When 页面加载,Then 显示数据概览(总记录数、认购率)。
            - AC2: Given 数据分析页面,When 选择分析维度,Then 展示对应图表(年龄分布饼图、职业认购率柱状图、教育水平热力图等)。
            - AC3: Given 数据分析页面,When 进行交互操作(选择/筛选),Then 图表实时更新。
            - AC4: Given 页面存在,When 访问非本地环境,Then 页面在 Docker 容器中正常渲染。
            - AC5: Given 页面功能,When 编写测试,Then 核心可视化逻辑有单元测试覆盖。
            
            技术备注:
            - 使用 Streamlit 的 st.metric/st.pyplot/st.plotly_chart 等组件。
            - 图表包括但不限于:年龄分布、职业分布、婚姻状况、教育水平、认购率分析。
            
            ---
            
            ### US-4 模型训练脚本与流程 · 状态: Backlog
            
            作为 **开发者**,
            我想要 一个离线训练脚本,
            以便 从历史数据中学习认购预测模型。
            
            验收标准:
            - AC1: Given 训练数据,When 执行 `python -m app.ml.train`,Then 在 ml/model/ 目录输出模型文件(model.pkl)。
            - AC2: Given 训练过程,When 训练完成,Then 打印关键指标(AUC、准确率、分类报告)到日志。
            - AC3: Given 模型文件,When 模型已存在,Then 训练脚本可选择覆盖或跳过(命令行参数)。
            - AC4: Given 训练脚本,When 在 CI 环境运行,Then 训练可复现(固定随机种子)。
            - AC5: Given 模型产物,When 提交代码,Then ml/model/ 在 .gitignore 中,不进 Git。
            
            技术备注:
            - 使用 scikit-learn(LogisticRegression/RandomForest/XGBoost 任选其一)。
            - 需要处理类别特征编码(OneHot/LabelEncoding)。
            - 固定 random_state 保证可复现性。
            
            ---
            
            ### US-5 预测服务核心逻辑 · 状态: Backlog
            
            作为 **系统**,
            我想要 一个预测服务模块,
            以便 根据输入特征返回认购概率。
            
            验收标准:
            - AC1: Given 模型文件存在,When 调用 `predict(features_dict)`,Then 返回预测结果(是否认购、概率)。
            - AC2: Given 特征输入,When 输入合法特征,Then 正确编码并调用模型。
            - AC3: Given 特征输入,When 输入缺失或非法值,Then 返回友好错误或默认处理。
            - AC4: Given 预测模块,When 编写测试,Then 覆盖正常预测、模型文件缺失、非法输入场景。
            - AC5: Given 预测服务,When 响应请求,Then 单次预测响应时间 <1s。
            
            技术备注:
            - 特征编码必须与训练时一致(建议保存 encoder)。
            - 返回格式:{"subscribe": bool, "probability": float, "confidence": str}。
            
            ---
            
            ### US-6 在线预测页面 · 状态: Backlog
            
            作为 **营销人员**,
            我想要 通过点选式表单输入客户特征,
            以便 快速预测该客户的认购意愿。
            
            验收标准:
            - AC1: Given 访问"预测系统"页面,When 页面加载,Then 显示点选式表单(年龄段、职业、婚姻、教育等选择器)。
            - AC2: Given 表单填写完成,When 点击"预测"按钮,Then 页面显示预测结果(是否认购、概率、置信度)。
            - AC3: Given 预测结果,When 显示结果,Then 提供可视化进度条或仪表盘展示概率。
            - AC4: Given 表单,When 用户操作,Then 表单支持重置、重新预测。
            - AC5: Given 预测页面,When 在生产环境,Then 预测功能在 Docker 容器中正常工作。
            
            技术备注:
            - 使用 Streamlit 的 st.selectbox/st.slider/st.button 组件。
            - 特征选项参考训练数据中的唯一值分布。
            - 结果展示包括:预测标签、概率进度条、建议文案。
            
            ---
            
            ### US-7 健康检查与监控 · 状态: Backlog
            
            作为 **运维/CD 流水线**,
            我想要 一个健康检查端点,
            以便 验证服务是否正常运行。
            
            验收标准:
            - AC1: Given 服务运行,When 访问 `/_stcore/health`,Then 返回 200 状态码。
            - AC2: Given 模型加载,When 服务启动,Then 模型文件成功加载,预测功能可用。
            - AC3: Given 部署脚本,When CD 执行,Then 健康检查失败时 CD 流水线报错。
            
            技术备注:
            - Streamlit 默认提供 `/_stcore/health` 端点。
            - 可选:添加自定义监控指标(请求计数、模型版本)。
            
            ---
            
            ### US-8 测试覆盖与质量门禁 · 状态: Backlog
            
            作为 **CI 流水线**,
            我想要 完整的测试覆盖,
            以便 保证代码质量。
            
            验收标准:
            - AC1: Given 核心业务逻辑,When 运行 `pytest --cov`,Then 覆盖率 ≥80%。
            - AC2: Given CI 触发,When PR 提交,Then CI 执行格式检查、静态检查、单元测试、构建。
            - AC3: Given 任意检查失败,When CI 红灯,Then PR 不能合并。
            - AC4: Given 本地开发,When 提交前,Then 开发者本地执行自检并全绿。
            
            ---
            
            ## 5. 非功能需求
            
            - **安全**:密钥只进 Secrets,不进 Git。
            - **可维护**:一需求一小 PR,避免大爆炸式提交。
            - **可测试**:核心逻辑必须有单元测试。
            - **可部署**:部署后必须有健康检查或等价验证。
            - **性能**:单次预测响应 <1s,页面首屏加载 <3s。
            
          • 📄 02-coding-standards.md 2.1KB
            # 02 · 代码标准
            
            > **目标**:让人和 AI 写出的代码风格一致、可读、可测、可维护。CI 会自动检查其中可机检的部分。
            
            ---
            
            ## 1. 格式与静态检查
            
            - Python 项目推荐:`ruff format .` + `ruff check .`。
            - 其他语言按项目技术栈替换为对应工具,例如 ESLint/Prettier、Checkstyle、gofmt。
            - 行宽默认 100;缩进、import 分组、命名风格由工具强制。
            - 提交前本地必须通过格式与 lint;CI 会再卡一道。
            
            ## 2. 命名
            
            | 对象 | 推荐规则 | 示例 |
            |---|---|---|
            | 变量/函数 | `snake_case` | `load_config` |
            | 类 | `PascalCase` | `UserService` |
            | 常量 | `UPPER_SNAKE` | `DEFAULT_PORT` |
            | 私有实现 | 前缀 `_` | `_parse_row` |
            
            命名必须表达意图。禁止 `process2`、`new_func`、`tmp_final` 这类无法审查的名字。
            
            ## 3. 函数与模块
            
            - 一个函数只做一件事;超过约 40 行优先拆分。
            - 参数过多时使用配置对象/数据类。
            - 业务逻辑与 I/O 分离,核心逻辑优先写成纯函数,方便测试。
            - 新增目录或模块前,先确认能落进 `00-project-context.md` 的目录地图。
            
            ## 4. 注释
            
            - 不写废话注释,例如 `# 调用函数` 配 `call_func()`。
            - 注释只解释“为什么”和“权衡”,不重复“代码做了什么”。
            - 复杂业务规则、部署坑、兼容性原因必须写清楚。
            
            ## 5. 配置与密钥
            
            - 非敏感配置可进 Git,例如 `config.yaml`、`.env.example`。
            - 密钥、密码、Token、私钥绝不进 Git,只放 GitHub Secrets / 环境变量。
            - 代码读取密钥时使用环境变量,并在 README 写清变量名,不写真实值。
            
            ## 6. 错误处理
            
            - 禁止裸 `except: pass`。
            - 外部输入必须校验,错误信息要能帮助定位。
            - 部署脚本必须失败即停;不能让前一步失败后继续执行后续危险动作。
            
            ## 7. 提交前自检
            
            ```text
            [ ] 格式检查通过
            [ ] 静态检查通过
            [ ] 新增逻辑有测试
            [ ] 没有硬编码密钥
            [ ] 没有调试 print / 死代码 / 大段注释代码
            [ ] 目录结构仍符合 00-project-context.md
            ```
            
          • 📄 03-testing-standards.md 2.9KB
            # 03 · 测试标准
            
            > **目标**:定义“什么叫测够了”。开发自测是进入 CI 前的第一道防线。
            
            ---
            
            ## 1. 测试分层
            
            ```text
                    /\        E2E 测试:少、慢、贵,覆盖真实用户路径
                   /  \
                  /----\      集成测试:测试模块/服务边界
                 /------\
                /        \    单元测试:多、快、稳定,开发者必写
               /__________\
            ```
            
            - 单元测试优先覆盖核心逻辑。
            - 集成测试覆盖接口、数据库、外部服务边界。
            - E2E 测试只覆盖关键路径,不要用它替代单元测试。
            
            ## 2. 用例写法
            
            推荐 AAA / Given-When-Then:
            
            ```python
            def test_health_returns_ok():
                # Arrange
                client = app.test_client()
            
                # Act
                resp = client.get("/health")
            
                # Assert
                assert resp.status_code == 200
                assert resp.get_json()["status"] == "ok"
            ```
            
            ## 3. 好测试的 5 条铁律
            
            1. 一个用例只测一件事。
            2. 覆盖正常值、边界值、异常值。
            3. 可重复、无外部网络依赖;有随机性必须固定种子。
            4. 测试名能表达场景和预期。
            5. 必须有真实断言,禁止 `assert result is not None` 凑数。
            
            ## 4. 覆盖率
            
            - 默认建议核心代码覆盖率 ≥ 80%。
            - 覆盖率不是目的,关键分支必须被覆盖。
            - 如果项目不适合覆盖率门槛,必须在 `00-project-context.md` 说明原因。
            
            ## 5. 项目特有质量门禁
            
            不同项目可追加不同质量门禁:
            
            | 项目类型 | 可选门禁示例 |
            |---|---|
            | Web/API | 健康检查、接口状态码、契约测试 |
            | 数据/ML | 数据校验、模型指标(AUC/F1/RMSE 等) |
            | 前端 | 组件测试、构建产物检查、可访问性检查 |
            | CLI/脚本 | 命令退出码、输入输出快照 |
            
            > 项目特有门禁写进 `00-project-context.md` 第 4 节,再写入 CI。
            
            ## 6. flaky test 纪律
            
            - 偶发失败必须根治,不能靠“重跑一次”掩盖。
            - 典型原因:时间依赖、随机数、测试顺序依赖、外部网络、共享状态未清理。
            
            ## 7. 本地测试由谁执行
            
            默认规则:
            
            | 环境 | 谁执行本地测试 | 怎么做 |
            |---|---|---|
            | AI 工具有终端权限(Cursor / Claude Code 等) | **AI 执行** | AI 自己运行格式检查、lint、pytest、构建等命令,读取日志,修复问题,把结果写回 `PROGRESS.md` |
            | AI 没有终端权限(纯聊天网页) | **学生执行** | 学生复制 AI 给出的命令到终端运行,把完整日志贴回 AI |
            | 涉及账号/权限/审批/Secrets | **学生确认** | AI 可指导,但不能代替人输入密码、私钥、审批生产发布 |
            
            AI 能执行命令时,不能只告诉学生“你去运行 pytest”。正确做法是:
            
            ```text
            AI 写代码 → AI 运行本地自检 → AI 根据日志修复 → AI 再运行 → 全绿后提交/PR → AI 记录结果
            ```
            
            每次本地自检至少记录:
            
            - 执行的命令。
            - 是否通过。
            - 失败时的关键错误。
            - 修复方式。
            - 最终通过结果。
            
          • 📄 04-git-workflow.md 4.4KB
            # 04 · Git 工作流与 PR/Review 标准
            
            > **目标**:定义代码如何从本地安全进入 `main`,并触发 CI/CD。课堂第一项功能必须完整演示本文件第 2 节。
            
            ---
            
            ## 1. 分支模型:GitHub Flow
            
            ```text
            main 受保护,永远可部署
             ├─ feature/<issue>-<desc>   新功能
             ├─ fix/<issue>-<desc>       普通缺陷
             ├─ docs/<issue>-<desc>      文档
             └─ chore/<issue>-<desc>     杂务/依赖/配置
            ```
            
            - 只有一条长期分支:`main`。
            - 日常开发必须从 `main` 开短分支,不能直接 push `main`。
            - feature 分支寿命建议 ≤ 2-3 天。
            - 分支命名:`<类型>/<issue号>-<英文短描述>`,例如 `feature/1-hello-api`。
            
            ---
            
            ## 2. 课堂必演示:从分支到 CD 的完整命令
            
            > 这段是学生必须亲手跑通的一条链。不要省略 PR,不要直接 push main。
            
            ```bash
            # 1. 确认在 main,并更新主干
            git switch main
            git pull
            
            # 2. 从 main 开功能分支
            git switch -c feature/1-first-feature
            
            # 3. 写代码 + 写测试后,本地自检
            ruff format --check .
            ruff check .
            pytest
            
            # 4. 提交
            git status
            git add .
            git commit -m "feat: add first feature"
            
            # 5. 推送 feature 分支
            git push -u origin feature/1-first-feature
            
            # 6. 创建 PR(AI 做到这里为止)
            gh pr create --base main --head feature/1-first-feature --title "feat: add first feature" --body "closes #1"
            
            # 7. 查看 CI(AI 汇报 CI 状态后停下,等人类审核)
            gh run list --limit 5
            gh run watch
            
            # 8. ★ 合并由人类执行 ★ —— 人工 Review 通过后,由人在网页点 "Merge",或人工运行:
            #    gh pr merge <PR号> --merge      # 默认保留分支,不加 --delete-branch
            #    AI 不得自行合并未获批准的 PR。
            
            # 9. 合并 main 后会触发 CD,查看部署流水线
            gh run list --limit 5
            ```
            
            > 如果课堂暂时没有开启分支保护,也要按这个流程演示 PR。分支保护是“强制执行”,PR 流程是“正确习惯”。
            
            > **合并是人类的动作,不是 AI 的动作。** AI 在第 6/7 步发起 PR、汇报 CI 后必须停下;
            > **只有人工 Review 通过,由人合并**,才触发 CD。AI 绝不替人按下 Merge。
            
            > **默认保留分支**(不加 `--delete-branch`),除非人类明确要求删除已合并分支。
            
            > **PR 必须用你本人 GitHub 账号发起**(`gh auth login` 登录的就是你本人),不要借他人账号代提。
            
            > **⚠️ 建仓后第一件事:配置 Secrets,否则 CD 必失败。**
            > 新仓库创建后(CD 第一次跑之前)立刻去
            > `GitHub → Settings → Secrets and variables → Actions → New repository secret`
            > 配好 `SSH_PRIVATE_KEY` / `SSH_HOST` / `SSH_USER`(详见 `05-cicd-standards.md` 第 5 节)。
            > 没有数据库等部署目标时可只配 SSH 三件套。配置时机:**在合并到 main 触发 CD 之前**。
            
            ---
            
            ## 3. 认证不要混
            
            | 动作 | 认证方式 | 说明 |
            |---|---|---|
            | `git push` / `git clone` | SSH 公钥或 HTTPS 凭证 | 负责搬代码 |
            | `gh pr create` / `gh run list` | GitHub CLI 登录 | 负责调 GitHub API |
            | CD SSH 到服务器 | GitHub Secrets 里的部署私钥 | 由 Actions runner 自动使用 |
            
            如果 SSH 22 端口卡住,可使用 GitHub SSH 443:
            
            ```bash
            ssh -T -p 443 git@ssh.github.com
            git remote set-url origin ssh://git@ssh.github.com:443/<账号>/<仓库>.git
            git push
            ```
            
            ---
            
            ## 4. 提交信息规范
            
            采用 Conventional Commits:
            
            | 前缀 | 含义 |
            |---|---|
            | `feat:` | 新功能 |
            | `fix:` | 修 bug |
            | `docs:` | 文档 |
            | `test:` | 测试 |
            | `refactor:` | 重构 |
            | `ci:` | CI/CD |
            | `build:` | 构建/依赖 |
            | `chore:` | 杂务 |
            
            一个 commit 只做一件事。
            
            ---
            
            ## 5. PR 通过三道闸
            
            | 闸 | 内容 | 不通过怎么办 |
            |---|---|---|
            | 机器闸 | CI 全绿 | 修代码/测试/配置 |
            | 人工闸 | 至少 1 人 Review | 修改并回复评论 |
            | 策略闸 | 无冲突、分支保护通过 | 同步 main / 解冲突 |
            
            PR 描述必须写:
            
            ```text
            ## What
            这个 PR 做了什么
            
            ## Why
            为什么需要
            
            ## Test
            - [ ] 本地检查通过
            - [ ] CI 通过
            
            closes #<issue号>
            ```
            
            ---
            
            ## 6. 冲突处理
            
            ```bash
            git switch main
            git pull
            git switch <你的分支>
            git rebase main
            
            # 如果有冲突:手动修改文件,删掉 <<<<<<< ======= >>>>>>> 标记
            git add <冲突文件>
            git rebase --continue
            
            # rebase 后推送
            git push --force-with-lease
            ```
            
            冲突处理后必须重新跑测试。
            
          • 📄 05-cicd-standards.md 7.4KB
            # 05 · CI/CD 标准
            
            > **目标**:定义最小可复现 CI/CD。CI 管“代码是否可靠”,CD 管“可靠代码如何自动上线”。
            
            ---
            
            ## 1. CI 与 CD
            
            | 名称 | 做什么 | 典型触发 |
            |---|---|---|
            | CI 持续集成 | 格式、lint、测试、覆盖率、构建 | PR / push |
            | CD 持续交付 | 构建可部署产物,部署前可人工审批 | main / tag |
            | CD 持续部署 | 合并 main 后自动部署上线 | main |
            
            课堂简版采用:
            
            ```text
            PR → CI 全绿 → 合并 main → CD 自动部署到服务器 → 健康检查
            ```
            
            ---
            
            ## 2. 最小 CI 门禁
            
            | 关卡 | 推荐命令 | 本地必跑 | CI 必跑 |
            |---|---|:---:|:---:|
            | 格式检查 | `ruff format --check .` 或等价工具 | ✅ | ✅ |
            | 静态检查 | `ruff check .` 或等价工具 | ✅ | ✅ |
            | 单元测试 | `pytest` 或等价工具 | ✅ | ✅ |
            | 覆盖率 | `pytest --cov --cov-fail-under=80` | ✅ | ✅ |
            | 项目特有门禁 | AUC/F1/接口契约/前端构建等 | ✅ | ✅ |
            | 构建 | `docker build ...` 或项目构建命令 | ❌ 不要求 | ✅ |
            
            CI 红灯时禁止合并。不要靠重跑掩盖问题。
            
            > **本地不强制装 Docker**:很多学员/开发机没有 Docker,本地装它只会增加复杂度。
            > 约定 **本地只跑前几项门禁,`docker build` 交给 CI(云端 runner 自带 Docker)与服务器(CD)**。
            > 这样本地零 Docker 负担,镜像构建仍有 CI 兜底。是否要本地构建,**按机器情况决定或先问开发者**。
            
            ---
            
            ## 3. 最小 CD 标准
            
            CD 必须满足:
            
            - 由 GitHub Actions runner 自动触发,不是人手动 SSH。
            - 只引用 GitHub Secrets,不硬编码密钥。
            - 部署脚本可重复执行。
            - 部署后必须健康检查。
            - 失败时 Actions 红灯,不能假成功。
            
            推荐触发:
            
            ```yaml
            on:
              push:
                branches: [ main ]
            ```
            
            ---
            
            ## 4. 依赖与镜像构建标准
            
            真实教学环境常见失败:服务器 `docker build` 时从 PyPI / npm / Docker Hub 下载依赖很慢或超时。
            
            通用标准:
            
            | 标准 | 做法 |
            |---|---|
            | 运行依赖与开发依赖拆分 | 生产镜像只安装运行依赖;测试/lint 工具只在 CI 安装 |
            | 镜像源可配置 | Dockerfile 支持镜像源参数,国内服务器可使用国内源 |
            | 构建失败即停 | 远程部署脚本第一行 `set -e` |
            | 健康检查 | 最后执行 `curl -fsS http://localhost:<PORT>/<HEALTHCHECK>` |
            | 可选预热 | 课前在服务器预拉基础镜像,减少等待 |
            
            Python 项目推荐:
            
            ```text
            requirements.txt      # 生产运行依赖
            requirements-dev.txt  # CI/本地检查依赖
            ```
            
            Dockerfile 推荐:
            
            ```dockerfile
            ARG PIP_INDEX_URL=https://pypi.org/simple
            RUN pip install --no-cache-dir --timeout 120 -i "${PIP_INDEX_URL}" -r requirements.txt
            ```
            
            远程部署脚本推荐(含幂等替换 + 端口自动回退):
            
            ```bash
            set -e
            docker build --build-arg PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple -t <APP>:latest .
            
            # 端口:容器内端口固定;主机端口优先 <PORT>,被占用就在预留区间自动找空闲端口。
            port_in_use() {
              ss -ltnH 2>/dev/null | grep -q ":$1 " && return 0
              docker ps --format "{{.Ports}}" 2>/dev/null | grep -q ":$1->" && return 0
              return 1
            }
            PORT=""
            for p in $(seq <PORT> <PORT_MAX>); do
              if ! port_in_use "$p"; then PORT="$p"; break; fi
            done
            [ -z "$PORT" ] && { echo "预留端口区间已全部占用,部署中止"; exit 1; }
            echo ">> 部署到主机端口 $PORT"
            
            docker rm -f <APP> 2>/dev/null || true   # 一步停删自身旧容器,幂等可重跑
            docker run -d --name <APP> --restart unless-stopped -p ${PORT}:<容器内端口> <APP>:latest
            sleep 3
            curl -fsS "http://localhost:${PORT}/<HEALTHCHECK>"
            echo ">> 部署成功:http://<SSH_HOST>:${PORT}/<HEALTHCHECK>"
            ```
            
            > **端口标准**:容器内端口固定,**主机端口可回退**(预留如 `<PORT>`-`<PORT_MAX>` 一段)。
            > 起容器用 `docker rm -f <APP>` 一步停删自身旧容器保证幂等;**不要**盲目删除占用端口的他人容器。
            > 防火墙/安全组要放行整段预留端口;最终落在哪个端口由 CD 日志打印。
            > 踩坑:`docker run` 报 `port is already allocated`(exit 125)就是主机端口被占用所致。
            
            ---
            
            ## 5. Secrets 标准
            
            常用 SSH 部署需要三个 GitHub Actions Secrets:
            
            | Secret | 含义 |
            |---|---|
            | `SSH_PRIVATE_KEY` | Actions runner 用来登录服务器的私钥全文 |
            | `SSH_HOST` | 服务器公网 IP 或域名 |
            | `SSH_USER` | 服务器登录用户,如 `root` 或 `deploy` |
            
            规则:
            
            - Secret 名大小写必须一致。
            - 私钥必须保留 `BEGIN/END` 和换行。
            - Secrets 必须在 CD run 开始前配置好;配置后旧 run 可能需要 rerun。
            
            ---
            
            ## 6. 本地环境准备(一次性)
            
            环境怎么建,**写进标准、可复现**。优先级:`uv` → `uv` 不行用 `conda` → 再不行用系统 `venv`。
            本课程学员更熟 `conda`,默认用 conda:
            
            ```bash
            # 1) 建并激活环境(示例环境名 envproj,Python 版本按项目 00 文档)
            conda create -y -n envproj python=3.11
            conda activate envproj
            
            # 2) 装依赖(国内用清华源更快)
            pip install -r requirements.txt -r requirements-dev.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
            ```
            
            > **坑:`CondaToSNonInteractiveError`(新版 conda 默认渠道未接受服务条款)**。
            > 解决,对三个渠道各执行一次:
            > ```bash
            > conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
            > conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
            > conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/msys2
            > ```
            
            ---
            
            ## 7. 排错顺序
            
            | 现象 | 优先检查 |
            |---|---|
            | `conda create` 报 ToS 错误 | 执行上面的 `conda tos accept`(main/r/msys2) |
            | Windows 控制台 `UnicodeEncodeError('gbk' ... \u25b6)` | 日志/打印含特殊符号;改用纯 ASCII,或 `set PYTHONIOENCODING=utf-8` / `chcp 65001`。注意 pytest 会捕获输出掩盖此坑,必须真跑一次脚本 |
            | CI 红:`FileNotFoundError` 找不到数据 | 数据被 `.gitignore` 排除,干净 runner 上没有。公开教学数据可入库;敏感数据则在 CI 里下载/造样本 |
            | `docker run` 报 `port is already allocated`(exit 125) | 主机端口被占用;容器内端口固定、主机端口在预留区间自动回退(见第 4 节);`docker rm -f <APP>` 幂等替换自身 |
            | `git push` 卡住/`Recv failure`/`Connection reset` | SSH 22 被拦可切 443;HTTPS 偶发抖动则重试几次 |
            | Actions 读不到变量 | Secret 名是否拼错、是否在正确仓库 |
            | SSH 登录失败 | 私钥/公钥是否匹配、`authorized_keys`、用户权限 |
            | rsync 成功但 docker build 失败 | 依赖下载、镜像源、Dockerfile |
            | 服务器本机 curl 通,外网打不开 | 安全组/防火墙是否放行端口(含回退端口段) |
            | CD 日志二次报错 | 部署脚本是否缺 `set -e` |
            | CD `Run Command Timeout`,docker build 卡在 `apt-get` | 国内服务器访问 deb.debian.org 极慢;**避免在 Dockerfile 里 apt 装包**,能用基础镜像已有运行时(如 python)替代就别 apt;非装不可则换国内 apt 源 |
            
            ---
            
            ## 8. 标准总结
            
            1. CI 红灯不合并。
            2. CD 必须由 runner 自动触发,不是人手动 SSH。
            3. 密钥进 Secrets,脚本进 Git,环境手动配一次。
            4. 生产镜像只放运行依赖,开发工具留在 CI;**本地不强制装 Docker**。
            5. 每次真实故障必须写回 `PROGRESS.md` 和本标准。
            
          • 📄 06-ai-collab-protocol.md 6.8KB
            # 06 · AI 协作协议
            
            > **目标**:让任意 AI 工具都按同一套上下文、规则和闭环工作,避免换工具/换会话就丢上下文。
            > **核心**:下面第 1 节的「六步交付流程 + 确认门」是 AI 必须严格遵守的主干,**每到一个确认门必须停下,向人类汇报并等待确认,不得自行越过**。
            
            ---
            
            ## 1. 六步交付流程(必须逐步确认)★ 主干
            
            AI 驱动开发分为固定六步。**每一步都有一个「✋ 确认门」**:AI 完成该步动作后,先汇报结果,等人类确认/操作后才进入下一步。绝不一口气从头跑到尾。
            
            ```text
            ① 建仓+配Secrets → ② 开feature分支 → ③ 本地模块化开发(逐模块汇报)
                 → ④ 本地CI自检(AI执行) → ⑤ 触发PR(CI在PR复检) → ⑥ 人工审核→合并→CD自动部署
               ✋确认1            ✋确认2          ✋确认3(每个模块)
                                ✋确认4           ✋确认5            ✋确认6(审核+部署结果)
            ```
            
            ### 第①步 · 建仓 + 配 Secrets
            
            - **AI 动作**:用 `gh` 创建 GitHub 仓库;`main` 只放最小引导提交(`.gitignore`/占位 README),**不在 main 上做真实开发**。
            - **✋ 确认门 1**:AI 报出仓库地址后,**必须提示人类去配置 GitHub Secrets**(新仓库默认为空,不配 CD 必失败):
              - `SSH_PRIVATE_KEY` / `SSH_HOST` / `SSH_USER`(见 `05` 第 5 节)。
              - AI 停下,**等人类确认「Secrets 已配置」**再继续(可用 `gh secret list` 核对)。
            
            ### 第②步 · 开 feature 分支
            
            - **AI 动作**:从最新 `main` 切出 `feature/<issue>-<desc>` 分支(类型见 `04`);**严禁直接在 main 上开发**。
            - **✋ 确认门 2**:报出分支名,确认后进入开发。
            
            ### 第③步 · 本地模块化开发 + 进度汇报
            
            - **AI 动作**:按 `00` 目录地图**一个模块一个模块**实现,每个模块同步写测试;每完成一个模块:
              - 更新 `PROGRESS.md`(已完成/下一步)。
              - 向人类**汇报该模块进度**(做了什么、关键决定、是否有坑)。
            - **✋ 确认门 3(每个模块)**:每个模块完成后汇报并确认,再做下一个模块。避免一次性写完全部再回头返工。
            
            ### 第④步 · 本地 CI 自检(由 AI 执行)
            
            - **AI 动作**:到「需要 CI」的节点,**AI 主动确认并亲自执行**本地门禁(具备终端权限时):
              - `ruff format --check .` + `ruff check .`
              - `pytest --cov --cov-fail-under=80`
              - 项目特有门禁(如模型 `--check-auc`)
              - **本地不强制 `docker build`**(交给 CI;见 `05` 第 2 节)。
            - **✋ 确认门 4**:汇报自检结果(覆盖率、指标);**全绿才进入下一步**,红则就地修复后重跑。
            
            ### 第⑤步 · 触发 PR(CI 在 PR 上复检)
            
            - **AI 动作**:`git push` 分支后,**AI 主动 `gh pr create`** 发起 PR(用本人账号);PR 事件触发 CI 在云端再跑一遍同样门禁 + `docker build`。
            - **✋ 确认门 5**:报出 PR 链接与 CI 状态(`gh run watch`);CI 绿才推进到审核。
            
            ### 第⑥步 · 人工审核 → 合并(人类) → CD 自动部署
            
            - **✋ AI 在此硬停**:AI 发起 PR、汇报 CI 绿之后**就停下**,**不得自己合并**。合并是人类的动作。
            - **人类动作**:**PR 必须由人工 Review**;审核通过后**由人类在网页点 Merge 或人工运行 `gh pr merge`**。
            - **分支默认保留**:合并不加 `--delete-branch`,除非人类明确要求删除。
            - **AI 动作(合并发生后)**:CD 自动触发(SSH→同步→构建→运行→健康检查;端口自愈见 `05`);AI 盯流水线。
            - **✋ 确认门 6**:汇报部署结果——**最终落地端口、`/health` 结果、访问地址**;失败则读日志定位根因,修复后回到第②步走一个新分支/PR。
            
            > 一句话:**建仓我做、Secrets 你配;分支我开(默认保留);模块我逐个写并汇报;本地 CI 我跑;PR 我发;
            > 审核与合并都由你做,我绝不替你按 Merge;合并后 CD 自动跑,我盯结果。**
            
            ---
            
            ## 2. AI 开工前必须读
            
            每次新会话,AI 必须按顺序读:
            
            1. `standards/README.md`
            2. `standards/00-project-context.md`
            3. `standards/01-requirements.md`
            4. `standards/PROGRESS.md`
            5. `standards/02`~`06`
            
            然后向人类复述:
            
            ```text
            我理解当前项目是:...
            当前阶段是(对应六步流程第几步):...
            下一步应该做:...
            我将遵守的关键约束是:...
            ```
            
            ---
            
            ## 3. 人给 AI 的标准开场 prompt
            
            ```prompt
            请先读取 standards/README.md,并按顺序读取 00/01/PROGRESS 与 02~06。
            读完后先复述你对项目、当前进度(对应 06 的六步流程第几步)、下一步的理解。
            然后严格按 06 的「六步交付流程 + 确认门」推进:
            - 每到一个确认门必须停下汇报并等我确认,不要一口气跑完。
            - 建仓后提示我配置 Secrets;从 main 开 feature 分支,不直接改 main。
            - 逐模块开发并汇报;本地 CI 由你执行(本地不强制 docker);CI 绿后你主动发 PR。
            - 发完 PR 你就停下,**不要自己合并**;PR 由我审核并由我合并,合并后 CD 自动跑,你盯结果并汇报端口与健康检查。
            - 分支默认保留,不要加 --delete-branch,除非我明确要求删。
            ```
            
            ---
            
            ## 4. AI 写代码边界
            
            AI 可以:
            
            - 新增/修改代码、测试、CI/CD 配置。
            - 在具备终端权限时运行本地测试、lint、构建命令,并根据日志修复。
            - 创建仓库、开分支、发起 PR(用本人账号)。
            - 更新 `00/01/PROGRESS`;根据真实故障更新 `02~06`(须说明原因)。
            
            AI 不可以:
            
            - 硬编码密钥。
            - 绕过 CI 或跳过测试。
            - 明明能执行命令,却只让学生手动跑本地测试。
            - 为了“看起来通过”删除测试或降低门槛。
            - **越过任何确认门**(尤其:未提示配 Secrets 就推进 CD;未经确认大规模重构)。
            - **自行合并任何 PR**——合并是人类的专属动作,AI 只发起 PR 并汇报,绝不替人按 Merge。
            - 合并时默认 `--delete-branch` 删分支——**默认保留分支**,除非人类明确要求删除。
            - 直接把项目特有内容写进通用规范,除非它已沉淀为可复用规则。
            
            ---
            
            ## 5. 故障反哺铁律
            
            CI/CD 一旦真实跑出问题,修完代码还不算结束,必须同步完成三件事:
            
            1. 写回 `PROGRESS.md`:现象、根因、修复、验证结果。
            2. 如是通用规则,写进 `02~06` 对应标准。
            3. 更新学生会照着操作的 README/课件/脚本。
            
            判断标准:
            
            ```text
            换一个学生、换一个 AI 工具,只读 standards + README,是否能避开这次坑并复现?
            ```
            
            如果不能,说明闭环未完成。
            
            ---
            
            ## 6. 会话结束 prompt
            
            ```prompt
            本次会话结束。请更新 standards/PROGRESS.md:
            1. 当前状态(对应六步流程第几步)
            2. 已完成事项
            3. 下一步 TODO
            4. 新增 ADR
            5. 新增 GOTCHAS
            保持时间倒序、简洁、可接力。
            ```
            
          • 📄 PROGRESS.md 9.4KB
            # PROGRESS · banksys_szai4 〔本项目活记忆 · 状态机〕
            
            > **作用**:这是项目的"存档点"。任意 AI、任意重启会话,读它即可知道当前做到哪、下一步做什么、踩过什么坑。
            > **更新时机**:每完成一个有意义步骤、每次会话结束前。
            > **格式要求**:时间倒序,最新在上;短、准、可接力。
            
            ---
            
            ## 当前状态 (最后更新: 2026-06-07 · by AI)
            
            - **阶段**: `PR #4 已合并;首次 CD 因 apt 超时失败(坑-007),已修 Dockerfile,准备发修复 PR`
            - **开发策略**: 按用户要求 US-3~US-6 全部本地开发并测试完成后,再统一进 CI/CD(一条大分支一个 PR)
            - **上一步完成**: 定位 CD 失败根因(apt-get 装 curl 在国内服务器超时),改为 Python 健康检查移除 apt
            - **下一步**: fix/cd-apt-timeout 分支 → 发 PR → 你合并 → CD 重跑 → 验证端口 8004 健康检查
            - **PR**: #4 已合并;修复 PR 待建
            - **阻塞项**: 无(端口已固定 8004,见坑-005)
            
            ---
            
            ## 待办清单 (TODO,按优先级)
            
            ### 第一批:初始化项目工程化与 CI/CD(US-1)✅ 已完成
            - [x] ① 建仓 + 配 Secrets
            - [x] ② 开 feature 分支
            - [x] ③ 本地模块化开发(8 个模块)
            - [x] ④ 本地 CI 自检
            - [x] ⑤ 触发 PR(PR #1)
            - [x] ⑥ 人工审核 → 合并 → CD 自动部署(端口 8005)
            
            ### 第二批:数据加载与预处理模块(US-2)✅ 已完成
            - [x] 实现 app/models/data_loader.py(load_train_data、load_test_data、辅助函数)
            - [x] 编写 test_data_loader.py(25 个测试,100% 覆盖率)
            - [x] 本地自检 + 提 PR(PR #3)
            - [x] 修复 ruff 版本锁定问题
            - [x] 修复 CD git ownership 问题
            - [x] CD 部署成功(端口 8006)
            
            ### 第三批:数据分析交互页面(US-3)✅ 已完成(本地)
            - [x] 实现数据可视化逻辑(app/models/visualizer.py,10 函数,100% 覆盖)
            - [x] 实现 app/pages/01_data_analysis.py(概览/饼图/类别分析/年龄分析/数值分布,交互式)
            - [x] 精简 main.py 为首页,改用 Streamlit 原生多页导航
            - [x] 编写 test_visualizer.py(24 个测试)+ 真实数据冒烟验证
            - [x] 顺手清理 data_loader.py 死代码、对齐 .gitignore 模型路径
            - [ ] 提 PR(按策略推迟到 US-6 完成后统一进 CI/CD)
            
            ### 第四批:模型训练脚本(US-4)✅ 已完成(本地)
            - [x] 实现 app/ml/train.py(ColumnTransformer+OneHot+RandomForest Pipeline,固定种子)
            - [x] app/ml/model/ 已在 .gitignore(git check-ignore 验证通过)
            - [x] 本地真跑训练:AUC 0.8964 / 准确率 0.8742,产出 model.pkl
            - [x] 编写 test_train.py(10 个测试),CLI 支持 --overwrite/--check-auc
            - [ ] 提 PR(推迟到统一进 CI/CD)
            
            ### 第五批:预测服务核心逻辑(US-5)✅ 已完成(本地)
            - [x] 实现 app/models/predictor.py(load_model 带 lru_cache,predict 返回 subscribe/probability/confidence)
            - [x] 编写 test_predictor.py(16 个测试,100% 覆盖):正常/模型缺失/缺失特征/非法数值/未知类别/响应时间
            - [x] 真实模型端到端验证:高意愿 0.90/high,低意愿 0.35;预热后单次 <1s
            - [ ] 提 PR(推迟到统一进 CI/CD)
            
            ### 第六批:在线预测页面(US-6)✅ 已完成(本地)
            - [x] 实现 app/pages/02_prediction.py(点选式表单,类别选项从数据动态取值)
            - [x] 集成 predictor 模块(结果含标签/概率进度条/置信度/建议文案,含模型缺失友好提示)
            - [x] 本地启动 Streamlit 验证:/_stcore/health=200、data_analysis/prediction 页面均 200
            - [ ] 提 PR(推迟到统一进 CI/CD)
            
            ### 第七批:测试覆盖与质量门禁完善(US-7 & US-8)
            - [x] 核心逻辑覆盖率:data_loader/visualizer/predictor 均 100%(UI 页面按规范不计)
            - [x] 注册 pytest unit marker,消除测试告警
            - [x] 修复 Dockerfile:构建时训练模型 + 装 curl + 健康检查端点 /_stcore/health(坑-006)
            - [x] 推分支 + 发 PR(#4),CI 复检全绿(含 docker build)
            - [ ] 人工审核 → 合并 → CD 部署 → 验证端口 8004 健康检查
            
            ---
            
            ## 关键决策记录 (ADR)
            
            | 日期 | 决策 | 理由 |
            |---|---|---|
            | 2025-06-07 | 选择 Streamlit 作为 Web 框架 | 课程要求;快速构建数据应用;适合数据分析与模型演示场景 |
            | 2025-06-07 | 模型训练离线,预测在线 | 训练是重操作不适合实时请求;预测是轻操作需要快速响应 |
            | 2025-06-07 | 数据集进 Git,模型产物不进 Git | 教学用公开数据,方便复现;模型产物二进制大文件,不应进版本控制 |
            | 2025-06-07 | 端口选择 8004 | 课程指定端口;Streamlit 默认 8501,Docker 映射到主机 8004 |
            
            ---
            
            ## 已知坑 (GOTCHAS)
            
            - **坑-001**: Dockerfile ARG 作用域问题
              - 现象: CI 构建时 `PIP_INDEX_URL` 为空字符串,pip 安装失败
              - 根因: ARG 在 FROM 之前定义,作用域在 FROM 后结束
              - 解决: 在 FROM 后重新声明 ARG,使用 `--index-url` 替代 `-i`
              - 验证: CI Docker 构建步骤通过
            
            - **坑-002**: CD 部署 git 目录检查
              - 现象: 首次部署时 `/opt/banksys` 存在但不是 git 仓库,`git clone || git pull` 失败
              - 根因: 目录存在非空时 `git clone` 失败;非 git 目录时 `git pull` 失败
              - 解决: 先检查是否为 git 目录(`[ -d .git ]`),是则 pull,否则删除后重新克隆
              - 验证: CD 部署成功,端口 8005,健康检查 ok
            
            - **坑-003**: ruff 版本不一致导致 CI 格式检查失败
              - 现象: 本地 `ruff format --check` 通过,CI 失败(1 file would be reformatted)
              - 根因: CI runner 和本地 ruff 版本不同(0.1.0 vs 0.6.4)
              - 解决: 固定 requirements-dev.txt 中 ruff 版本为 `==0.6.4`,CI 中显式安装相同版本
              - 验证: CI 格式检查通过
            
            - **坑-004**: CD git ownership 错误
              - 现象: `fatal: detected dubious ownership in repository at '/opt/banksys'`
              - 根因: Docker 容器操作后目录所有权变化,git 检测到可疑所有权
              - 解决: CD 脚本开头添加 `git config --global --add safe.directory /opt/banksys`
              - 验证: CD 部署成功
            
            - **坑-008**: Streamlit 多页应用容器内 `ModuleNotFoundError: No module named 'app'`(已修复)
              - 现象: 部署后页面报 `from app.models import visualizer` → No module named 'app'(/app/app/pages/01_data_analysis.py)
              - 根因: 容器 `CMD streamlit run app/main.py` 只把主脚本目录 /app/app 加进 sys.path,而 `import app.*` 需要父目录 /app 在路径上;本地用 `python -m streamlit` 会把 cwd 加进路径所以没暴露
              - 隐蔽点: 本地只验了 HTTP 200,但 Streamlit 页面脚本在 websocket 连接时才执行,200 不代表页面渲染成功
              - 解决: Dockerfile 增加 `ENV PYTHONPATH=/app`;本地复现验证:无 PYTHONPATH 报同样错,设置后三模块导入 OK
              - 教训: 验证 Streamlit 多页应用要真正打开每个页面触发脚本执行,不能只看健康检查/HTTP 200
              - 验证: 待 PR #5 合并后 CD 确认
            
            - **坑-007**: Docker 内 apt-get 装 curl 在国内服务器超时导致 CD 失败(已修复)
              - 现象: CD "Deploy to server" 跑 10 分钟后 `Run Command Timeout`;日志 `Fetched 9975 kB in 8min 6s (20.5 kB/s)`
              - 根因: 为容器 HEALTHCHECK 装 curl 加了 `apt-get update/install`,服务器访问 deb.debian.org 极慢(20.5 kB/s),超过 appleboy/ssh-action 默认 10 分钟命令超时。CI 在国外 runner 上 apt 快所以没暴露
              - 解决: Dockerfile 彻底移除 apt-get/curl,容器 HEALTHCHECK 改用镜像自带 Python(urllib)访问 /_stcore/health,完全绕开 apt
              - 教训: 国内服务器构建镜像要避免走 Debian 官方源;能用镜像已有运行时(python)就别额外 apt 装包
              - 验证: 待新 PR 合并后 CD 确认
            
            - **坑-006**: 模型产物不进镜像导致生产预测失效(已修复)
              - 现象: model.pkl 被 .gitignore 排除,CD 克隆代码后镜像内无模型,预测页报"模型未找到"
              - 根因: 模型是离线产物不进 Git,但 Dockerfile 未在构建时训练
              - 解决: Dockerfile 在 COPY app/ 与 data/ 后增加 `RUN python -m app.ml.train --overwrite`,把模型烤进镜像;同时 slim 镜像补装 curl、HEALTHCHECK 端点改为 /_stcore/health
              - 验证: 待 CI docker build 与 CD 部署确认
            
            - **坑-005**: 端口回退问题(已修复)
              - 现象: 每次部署端口不同(8004→8005→8006),端口递增
              - 根因: 端口检查在容器删除之前;多个旧容器占用不同端口
              - 解决: 先停止所有占用 8000-8010 端口的容器,等待释放,固定使用 8004
              - 验证: CD 成功,固定部署在端口 8004
            
            ---
            
            ## 里程碑 (DONE)
            
            - [x] 2025-06-07:填写项目规范文档(00/01/PROGRESS)
            - [x] 2025-06-07:创建 GitHub 仓库(banksys_szai4)
            - [x] 2025-06-07:US-1 完成 - 初始化项目工程化与 CI/CD
              - 完整跑通六步交付流程
              - CI/CD 流水线配置完成
              - CD 成功部署到服务器(端口 8005)
            - [x] 2025-06-07:US-2 完成 - 数据加载与预处理模块
              - 实现 data_loader.py 模块(4 个函数,100% 覆盖率)
              - 编写 25 个测试用例
              - 修复 ruff 版本锁定问题
              - 修复 CD git ownership 问题
              - CD 部署成功(端口 8006)
            
            > 反臃肿:里程碑超过 15 条时,把更早内容合并成一行摘要,保持本文件可快速阅读。
            
          • 📄 README.md 3.6KB
            # standards-template · 学生项目可复制模板
            
            > **用途**:每个新项目复制一份本目录,改名为项目根目录下的 `standards/`。任意 AI 工具进入项目后,先读 `standards/README.md`,即可按统一规则继续开发、提交、PR、CI/CD 与存档。
            
            ---
            
            ## 1. 学生怎么用
            
            ```bash
            # 在你的项目根目录执行
            cp -r <课件目录>/standards-template ./standards
            ```
            
            Windows 可手动复制文件夹,复制后目录应是:
            
            ```text
            你的项目/
            ├── standards/
            │   ├── README.md
            │   ├── 00-project-context.md
            │   ├── 01-requirements.md
            │   ├── PROGRESS.md
            │   ├── 02-coding-standards.md
            │   ├── 03-testing-standards.md
            │   ├── 04-git-workflow.md
            │   ├── 05-cicd-standards.md
            │   ├── 06-ai-collab-protocol.md
            │   └── templates/
            └── <你的代码>
            ```
            
            ---
            
            ## 2. 两类文件,不要混
            
            | 类型 | 文件 | 是否每个项目要改 | 谁维护 |
            |---|---|---|---|
            | **项目活记忆** | `00-project-context.md`、`01-requirements.md`、`PROGRESS.md` | 必须按项目重填 | AI 写,人审 |
            | **通用规范** | `02`~`06`、`templates/` | 默认不改 | 团队维护 |
            
            > 简单说:`00/01/PROGRESS` 记录“这个项目是什么、要做什么、做到哪了”;`02~06` 规定“怎么写代码、怎么测试、怎么走 Git/PR/CI/CD”。
            
            ---
            
            ## 3. AI 每次会话读取顺序
            
            1. `00-project-context.md` — 项目身份、技术栈、目录地图、部署取值。
            2. `01-requirements.md` — 活 PRD,所有需求与验收标准。
            3. `PROGRESS.md` — 当前状态、下一步、决策、踩坑。
            4. `02`~`06` — 通用工程规范。
            5. `templates/` — Issue / PR 模板。
            
            ---
            
            ## 4. 第一次开工 prompt
            
            ```prompt
            请先读取 standards/README.md,并按它的顺序读取 00/01/PROGRESS 与 02~06。
            这是一个新项目。我的项目需求是:
            <在这里写项目目标、技术栈、希望实现的功能、部署方式>
            
            请先做三件事:
            1. 填写 standards/00-project-context.md。
            2. 把需求整理进 standards/01-requirements.md,写成用户故事和验收标准。
            3. 初始化 standards/PROGRESS.md,列出第一批 TODO。
            
            完成后先停下让我确认,不要直接写代码。
            ```
            
            ---
            
            ## 5. 标准开发闭环 · 固定六步(每步确认)
            
            AI 严格按 `06-ai-collab-protocol.md` 的「六步交付流程 + 确认门」推进,**每到确认门停下汇报、等人确认**:
            
            ```text
            ① 建仓 + 配 Secrets    AI 建仓 → ✋提示人类配 SSH_PRIVATE_KEY/HOST/USER(新仓库默认没有)
            ② 开 feature 分支       从 main 切分支,严禁直接改 main         ✋报分支名
            ③ 本地模块化开发        逐模块写代码+测试,每模块更新 PROGRESS  ✋每个模块汇报进度
            ④ 本地 CI 自检(AI 跑)  ruff + pytest + 覆盖率 + 项目门禁;本地不强制 docker  ✋全绿才继续
            ⑤ 触发 PR              push + AI 主动建 PR,CI 在 PR 上复检    ✋报 PR 链接 + CI 状态
            ⑥ 人工审核 →(人)合并 → CD  ✋AI 发完 PR 即停;Review 与 Merge 都由人做;合并触发 CD → ✋报端口/健康检查
            ```
            
            **课堂第一项功能必须完整演示一次这条链**,不能直接 push main;
            **合并是人类的动作——AI 绝不自行合并 PR;分支默认保留(不自动删)。**
            
            ---
            
            ## 6. 反臃肿纪律
            
            1. 需求只写进 `01-requirements.md`,不要到处新建 PRD。
            2. 进度、决策、坑只写进 `PROGRESS.md`。
            3. 一需求一分支一 PR,PR 尽量小于 400 行。
            4. CI 红灯不合并;flaky test 必须根治。
            5. 新增目录前先更新 `00-project-context.md` 的目录地图。
            
        • 📁 tests/
          • 📄 __init__.py 19B
            """Test suite."""
            
          • 📄 conftest.py 526B
            """Pytest configuration and fixtures."""
            
            import sys
            from pathlib import Path
            
            import pytest
            
            # Add app directory to path
            sys.path.insert(0, str(Path(__file__).parent.parent))
            
            
            @pytest.fixture
            def sample_data():
                """Sample bank marketing data for testing."""
                return {
                    "age": 45,
                    "job": "admin.",
                    "marital": "married",
                    "education": "university.degree",
                    "default": "no",
                    "housing": "yes",
                    "loan": "no",
                    "contact": "cellular",
                }
            
          • 📄 test_data_loader.py 7.1KB
            """Tests for data_loader module."""
            
            import pandas as pd
            import pytest
            
            from app.models.data_loader import (
                FEATURE_COLUMNS,
                TRAIN_COLUMNS,
                get_column_statistics,
                get_missing_value_mask,
                load_test_data,
                load_train_data,
            )
            
            
            class TestLoadTrainData:
                """Test suite for load_train_data function."""
            
                def test_load_train_data_returns_dataframe(self):
                    """Test that load_train_data returns a pandas DataFrame."""
                    df = load_train_data()
                    assert isinstance(df, pd.DataFrame)
            
                def test_load_train_data_has_expected_columns(self):
                    """Test that loaded data has all expected columns."""
                    df = load_train_data()
                    assert list(df.columns) == TRAIN_COLUMNS
            
                def test_load_train_data_not_empty(self):
                    """Test that loaded data is not empty."""
                    df = load_train_data()
                    assert len(df) > 0
            
                def test_load_train_data_has_target_column(self):
                    """Test that 'subscribe' column exists."""
                    df = load_train_data()
                    assert "subscribe" in df.columns
            
                def test_load_train_data_from_custom_dir(self, tmp_path):
                    """Test loading from custom data directory."""
                    # Create temporary train.csv
                    train_data = pd.DataFrame({col: ["test"] for col in TRAIN_COLUMNS})
                    train_path = tmp_path / "train.csv"
                    train_data.to_csv(train_path, index=False)
            
                    df = load_train_data(data_dir=tmp_path)
                    assert len(df) == 1
                    assert df.iloc[0]["age"] == "test"
            
            
            class TestLoadTestData:
                """Test suite for load_test_data function."""
            
                def test_load_test_data_returns_dataframe(self):
                    """Test that load_test_data returns a pandas DataFrame."""
                    df = load_test_data()
                    assert isinstance(df, pd.DataFrame)
            
                def test_load_test_data_has_feature_columns(self):
                    """Test that loaded data has feature columns only (no target)."""
                    df = load_test_data()
                    assert list(df.columns) == FEATURE_COLUMNS
                    assert "subscribe" not in df.columns
            
                def test_load_test_data_not_empty(self):
                    """Test that loaded data is not empty."""
                    df = load_test_data()
                    assert len(df) > 0
            
            
            class TestErrorHandling:
                """Test suite for error handling."""
            
                def test_load_train_data_file_not_found(self, tmp_path):
                    """Test that FileNotFoundError is raised when file doesn't exist."""
                    with pytest.raises(FileNotFoundError, match="Data file not found"):
                        load_train_data(data_dir=tmp_path)
            
                def test_load_test_data_file_not_found(self, tmp_path):
                    """Test that FileNotFoundError is raised when file doesn't exist."""
                    with pytest.raises(FileNotFoundError, match="Data file not found"):
                        load_test_data(data_dir=tmp_path)
            
                def test_load_empty_file_raises_error(self, tmp_path):
                    """Test that ValueError is raised for empty file."""
                    empty_path = tmp_path / "train.csv"
                    empty_path.write_text("id,age\n")  # Header only
            
                    with pytest.raises(ValueError, match="Data file is empty"):
                        load_train_data(data_dir=tmp_path)
            
                def test_load_invalid_columns_raises_error(self, tmp_path):
                    """Test that ValueError is raised for invalid columns."""
                    invalid_path = tmp_path / "train.csv"
                    invalid_path.write_text("wrong,columns\n1,2\n")
            
                    with pytest.raises(ValueError, match="Missing expected columns"):
                        load_train_data(data_dir=tmp_path)
            
            
            class TestMissingValueMask:
                """Test suite for get_missing_value_mask function."""
            
                def test_returns_boolean_series(self):
                    """Test that function returns boolean Series."""
                    df = load_train_data()
                    mask = get_missing_value_mask(df, "education")
                    assert mask.dtype == bool
            
                def test_invalid_column_raises_error(self):
                    """Test that ValueError is raised for invalid column."""
                    df = load_train_data()
                    with pytest.raises(ValueError, match="Column not found"):
                        get_missing_value_mask(df, "nonexistent_column")
            
                def test_detects_missing_values(self):
                    """Test that missing values are correctly detected."""
                    df = pd.DataFrame({"col": ["known", "unknown", "nonexistent", ""]})
                    mask = get_missing_value_mask(df, "col")
                    # First value should not be marked as missing
                    assert not mask.iloc[0]
                    # Rest should be marked as missing
                    assert mask.iloc[1:].all()
                    assert mask.sum() == 3
            
            
            class TestColumnStatistics:
                """Test suite for get_column_statistics function."""
            
                @pytest.fixture
                def sample_df(self):
                    """Create sample DataFrame for testing."""
                    return pd.DataFrame(
                        {
                            "age": [25, 30, 35, 40, 45],
                            "job": ["admin.", "blue-collar", "unknown", "admin.", "blue-collar"],
                        }
                    )
            
                def test_returns_statistics_for_numeric_column(self, sample_df):
                    """Test that statistics are returned for numeric columns."""
                    stats = get_column_statistics(sample_df, "age")
                    assert "mean" in stats
                    assert "std" in stats
                    assert "min" in stats
                    assert "max" in stats
                    assert stats["mean"] == 35.0
            
                def test_returns_statistics_for_categorical_column(self, sample_df):
                    """Test that statistics are returned for categorical columns."""
                    stats = get_column_statistics(sample_df, "job")
                    assert "count" in stats
                    assert "unique" in stats
                    assert "value_counts" in stats
                    assert stats["unique"] == 3
            
                def test_invalid_column_raises_error(self):
                    """Test that ValueError is raised for invalid column."""
                    df = load_train_data()
                    with pytest.raises(ValueError, match="Column not found"):
                        get_column_statistics(df, "nonexistent_column")
            
                def test_real_data_age_statistics(self):
                    """Test statistics on real age data."""
                    df = load_train_data()
                    stats = get_column_statistics(df, "age")
                    assert stats["min"] >= 0
                    assert stats["max"] <= 120
                    assert stats["count"] == len(df)
            
            
            class TestIntegration:
                """Integration tests for data loader."""
            
                def test_train_has_more_columns_than_test(self):
                    """Test that train data has target column while test data does not."""
                    train_df = load_train_data()
                    test_df = load_test_data()
                    # Train has subscribe column
                    assert "subscribe" in train_df.columns
                    # Test does not have subscribe column
                    assert "subscribe" not in test_df.columns
                    # Test has all feature columns
                    assert set(test_df.columns) == set(FEATURE_COLUMNS)
            
                def test_data_dimensions_reasonable(self):
                    """Test that data dimensions are reasonable."""
                    train_df = load_train_data()
                    test_df = load_test_data()
                    # Both should have at least 1000 rows
                    assert len(train_df) > 1000
                    assert len(test_df) > 1000
                    # Train has 22 columns, test has 21 columns
                    assert len(train_df.columns) == 22
                    assert len(test_df.columns) == 21
            
          • 📄 test_health.py 828B
            """Basic health tests."""
            
            import pytest
            
            
            def test_python_version():
                """Verify Python version is 3.11+."""
                import sys
            
                major, minor = sys.version_info[:2]
                assert major == 3 and minor >= 11, f"Python 3.11+ required, got {major}.{minor}"
            
            
            def test_app_imports():
                """Verify app modules can be imported."""
                import app
                import app.models
                import app.pages
                import app.utils
            
                assert app.__version__ is not None
            
            
            def test_data_files_exist():
                """Verify data files exist."""
                from pathlib import Path
            
                data_dir = Path(__file__).parent.parent / "data"
                assert (data_dir / "train.csv").exists()
                assert (data_dir / "test.csv").exists()
            
            
            @pytest.mark.unit
            def test_placeholder():
                """Placeholder test to verify pytest runs."""
                assert True
            
          • 📄 test_predictor.py 4.5KB
            """Tests for the predictor module.
            
            A small model is trained on synthetic data and saved to a temp path, so tests
            never depend on the real (gitignored) model artifact.
            """
            
            import time
            
            import pandas as pd
            import pytest
            
            from app.ml import train
            from app.models import predictor
            from app.models.data_loader import MODEL_FEATURES
            
            
            def _make_features(positive: bool = True) -> dict:
                """Build a complete, valid feature dict for one customer."""
                return {
                    "job": "student" if positive else "admin.",
                    "marital": "single" if positive else "married",
                    "education": "university.degree",
                    "default": "no",
                    "housing": "yes",
                    "loan": "no",
                    "contact": "cellular",
                    "month": "may",
                    "day_of_week": "mon",
                    "poutcome": "success" if positive else "failure",
                    "age": 25 if positive else 50,
                    "duration": 800 if positive else 100,
                    "campaign": 1,
                    "pdays": 5 if positive else 999,
                    "previous": 2 if positive else 0,
                    "emp_var_rate": -1.8 if positive else 1.4,
                    "cons_price_index": 92.0,
                    "cons_conf_index": -40.0,
                    "lending_rate3m": 1.0,
                    "nr_employed": 5000.0,
                }
            
            
            @pytest.fixture
            def model_path(tmp_path):
                """Train a tiny model on synthetic data and return its file path as str."""
                rows = [_make_features(positive=(i % 3 == 0)) for i in range(60)]
                df = pd.DataFrame(rows)
                df["subscribe"] = ["yes" if i % 3 == 0 else "no" for i in range(60)]
                pipeline, _ = train.train_model(df)
                out = tmp_path / "model.pkl"
                train.save_model(pipeline, path=out)
                return str(out)
            
            
            class TestLoadModel:
                def test_missing_model_raises(self, tmp_path):
                    missing = str(tmp_path / "nope.pkl")
                    with pytest.raises(predictor.ModelNotFoundError):
                        predictor.load_model(missing)
            
                def test_load_returns_payload(self, model_path):
                    payload = predictor.load_model(model_path)
                    assert "pipeline" in payload
                    assert payload["feature_columns"] == MODEL_FEATURES
            
            
            class TestPredict:
                def test_returns_expected_keys(self, model_path):
                    result = predictor.predict(_make_features(), model_path=model_path)
                    assert set(result.keys()) == {"subscribe", "probability", "confidence"}
            
                def test_types(self, model_path):
                    result = predictor.predict(_make_features(), model_path=model_path)
                    assert isinstance(result["subscribe"], bool)
                    assert isinstance(result["probability"], float)
                    assert isinstance(result["confidence"], str)
            
                def test_probability_in_range(self, model_path):
                    result = predictor.predict(_make_features(), model_path=model_path)
                    assert 0.0 <= result["probability"] <= 1.0
            
                def test_subscribe_consistent_with_probability(self, model_path):
                    result = predictor.predict(_make_features(), model_path=model_path)
                    assert result["subscribe"] == (result["probability"] >= 0.5)
            
                def test_missing_feature_raises(self, model_path):
                    features = _make_features()
                    del features["age"]
                    with pytest.raises(ValueError, match="Missing required features"):
                        predictor.predict(features, model_path=model_path)
            
                def test_invalid_numeric_raises(self, model_path):
                    features = _make_features()
                    features["age"] = "not-a-number"
                    with pytest.raises(ValueError, match="must be numeric"):
                        predictor.predict(features, model_path=model_path)
            
                def test_unknown_category_handled(self, model_path):
                    features = _make_features()
                    features["job"] = "astronaut"  # unseen category
                    result = predictor.predict(features, model_path=model_path)
                    assert 0.0 <= result["probability"] <= 1.0
            
                def test_response_under_one_second(self, model_path):
                    predictor.predict(_make_features(), model_path=model_path)  # warm cache
                    start = time.perf_counter()
                    predictor.predict(_make_features(), model_path=model_path)
                    assert (time.perf_counter() - start) < 1.0
            
            
            class TestConfidenceLabel:
                @pytest.mark.parametrize(
                    "prob,expected",
                    [
                        (0.95, "high"),
                        (0.05, "high"),
                        (0.7, "medium"),
                        (0.3, "medium"),
                        (0.5, "low"),
                        (0.55, "low"),
                    ],
                )
                def test_confidence_labels(self, prob, expected):
                    assert predictor._confidence_label(prob) == expected
            
          • 📄 test_train.py 3.9KB
            """Tests for the model training module.
            
            Uses a small synthetic dataset to keep tests fast and avoid depending on the
            full data file. Verifies pipeline construction, X/y preparation, end-to-end
            training metrics, and model persistence (with feature schema).
            """
            
            import joblib
            import pandas as pd
            import pytest
            from sklearn.pipeline import Pipeline
            
            from app.ml import train
            from app.models.data_loader import CATEGORICAL_FEATURES, MODEL_FEATURES, NUMERIC_FEATURES
            
            
            @pytest.fixture
            def synthetic_df():
                """Synthetic dataset with a learnable signal and both target classes."""
                rows = []
                for i in range(60):
                    positive = i % 3 == 0
                    rows.append(
                        {
                            "job": "student" if positive else "admin.",
                            "marital": "single" if positive else "married",
                            "education": "university.degree",
                            "default": "no",
                            "housing": "yes" if positive else "no",
                            "loan": "no",
                            "contact": "cellular",
                            "month": "may",
                            "day_of_week": "mon",
                            "poutcome": "success" if positive else "failure",
                            "age": 25 if positive else 50,
                            "duration": 800 if positive else 100,
                            "campaign": 1 if positive else 3,
                            "pdays": 5 if positive else 999,
                            "previous": 2 if positive else 0,
                            "emp_var_rate": -1.8 if positive else 1.4,
                            "cons_price_index": 92.0,
                            "cons_conf_index": -40.0,
                            "lending_rate3m": 1.0,
                            "nr_employed": 5000.0,
                            "subscribe": "yes" if positive else "no",
                        }
                    )
                return pd.DataFrame(rows)
            
            
            class TestBuildPipeline:
                def test_returns_pipeline(self):
                    assert isinstance(train.build_pipeline(), Pipeline)
            
                def test_pipeline_has_expected_steps(self):
                    pipeline = train.build_pipeline()
                    assert "preprocess" in pipeline.named_steps
                    assert "model" in pipeline.named_steps
            
            
            class TestPrepareXy:
                def test_x_has_model_features(self, synthetic_df):
                    x, _ = train._prepare_xy(synthetic_df)
                    assert list(x.columns) == MODEL_FEATURES
            
                def test_y_is_binary(self, synthetic_df):
                    _, y = train._prepare_xy(synthetic_df)
                    assert set(y.unique()).issubset({0, 1})
            
                def test_y_positive_count(self, synthetic_df):
                    _, y = train._prepare_xy(synthetic_df)
                    assert int(y.sum()) == 20
            
            
            class TestTrainModel:
                def test_returns_fitted_pipeline_and_metrics(self, synthetic_df):
                    pipeline, metrics = train.train_model(synthetic_df)
                    assert isinstance(pipeline, Pipeline)
                    assert "auc" in metrics and "accuracy" in metrics
            
                def test_metrics_within_range(self, synthetic_df):
                    _, metrics = train.train_model(synthetic_df)
                    assert 0.0 <= metrics["auc"] <= 1.0
                    assert 0.0 <= metrics["accuracy"] <= 1.0
            
                def test_reproducible(self, synthetic_df):
                    _, metrics_a = train.train_model(synthetic_df)
                    _, metrics_b = train.train_model(synthetic_df)
                    assert metrics_a["auc"] == metrics_b["auc"]
            
            
            class TestSaveModel:
                def test_save_creates_file(self, synthetic_df, tmp_path):
                    pipeline, _ = train.train_model(synthetic_df)
                    out = tmp_path / "model.pkl"
                    saved = train.save_model(pipeline, path=out)
                    assert saved.exists()
            
                def test_saved_payload_schema(self, synthetic_df, tmp_path):
                    pipeline, _ = train.train_model(synthetic_df)
                    out = tmp_path / "model.pkl"
                    train.save_model(pipeline, path=out)
                    payload = joblib.load(out)
                    assert payload["categorical_features"] == CATEGORICAL_FEATURES
                    assert payload["numeric_features"] == NUMERIC_FEATURES
                    assert isinstance(payload["pipeline"], Pipeline)
            
          • 📄 test_visualizer.py 6.0KB
            """Tests for the visualizer module.
            
            Computation helpers are tested with a small in-memory DataFrame (no dependency
            on the real data files). Figure builders are smoke-tested to confirm they
            return Plotly figures without raising.
            """
            
            import pandas as pd
            import plotly.graph_objects as go
            import pytest
            
            from app.models import visualizer
            
            
            @pytest.fixture
            def sample_df():
                """Small, hand-built dataset with a known subscribe distribution."""
                return pd.DataFrame(
                    {
                        "age": [23, 30, 41, 52, 67, 38, 29, 45],
                        "job": [
                            "admin.",
                            "admin.",
                            "services",
                            "services",
                            "retired",
                            "admin.",
                            "blue-collar",
                            "services",
                        ],
                        "duration": [100, 200, 300, 400, 500, 150, 250, 350],
                        "subscribe": ["no", "yes", "no", "yes", "yes", "no", "no", "no"],
                    }
                )
            
            
            class TestComputeOverview:
                def test_overview_counts_and_rate(self, sample_df):
                    result = visualizer.compute_overview(sample_df)
                    assert result["total_records"] == 8
                    assert result["subscribed"] == 3
                    assert result["not_subscribed"] == 5
                    assert result["subscribe_rate"] == pytest.approx(3 / 8)
            
                def test_overview_feature_count(self, sample_df):
                    result = visualizer.compute_overview(sample_df)
                    assert result["feature_count"] == 4
            
                def test_overview_missing_target_raises(self):
                    df = pd.DataFrame({"age": [1, 2]})
                    with pytest.raises(ValueError, match="Target column"):
                        visualizer.compute_overview(df)
            
            
            class TestComputeValueCounts:
                def test_value_counts_sums_to_total(self, sample_df):
                    result = visualizer.compute_value_counts(sample_df, "job")
                    assert result["count"].sum() == len(sample_df)
            
                def test_value_counts_columns(self, sample_df):
                    result = visualizer.compute_value_counts(sample_df, "job")
                    assert list(result.columns) == ["job", "count"]
            
                def test_value_counts_admin_count(self, sample_df):
                    result = visualizer.compute_value_counts(sample_df, "job")
                    admin_row = result[result["job"] == "admin."]
                    assert int(admin_row["count"].iloc[0]) == 3
            
                def test_value_counts_missing_column_raises(self, sample_df):
                    with pytest.raises(ValueError, match="Column not found"):
                        visualizer.compute_value_counts(sample_df, "nonexistent")
            
            
            class TestSubscribeRateByCategory:
                def test_rate_values(self, sample_df):
                    result = visualizer.compute_subscribe_rate_by_category(sample_df, "job")
                    retired = result[result["job"] == "retired"]
                    assert retired["rate"].iloc[0] == pytest.approx(1.0)
                    admin = result[result["job"] == "admin."]
                    assert admin["rate"].iloc[0] == pytest.approx(1 / 3)
            
                def test_rate_sorted_descending(self, sample_df):
                    result = visualizer.compute_subscribe_rate_by_category(sample_df, "job")
                    rates = result["rate"].tolist()
                    assert rates == sorted(rates, reverse=True)
            
                def test_totals_match(self, sample_df):
                    result = visualizer.compute_subscribe_rate_by_category(sample_df, "job")
                    assert result["total"].sum() == len(sample_df)
            
                def test_missing_target_raises(self):
                    df = pd.DataFrame({"job": ["a", "b"]})
                    with pytest.raises(ValueError, match="Target column"):
                        visualizer.compute_subscribe_rate_by_category(df, "job")
            
            
            class TestComputeAgeGroups:
                def test_age_group_counts_sum(self, sample_df):
                    result = visualizer.compute_age_groups(sample_df)
                    assert result["count"].sum() == len(sample_df)
            
                def test_age_group_has_rate_column(self, sample_df):
                    result = visualizer.compute_age_groups(sample_df)
                    assert "rate" in result.columns
                    assert (result["rate"] >= 0).all() and (result["rate"] <= 1).all()
            
                def test_age_group_custom_bins(self, sample_df):
                    result = visualizer.compute_age_groups(
                        sample_df, bins=[0, 40, 120], labels=["young", "old"]
                    )
                    assert set(result["age_group"]) == {"young", "old"}
            
                def test_age_group_invalid_labels_raises(self, sample_df):
                    with pytest.raises(ValueError, match="labels length"):
                        visualizer.compute_age_groups(sample_df, bins=[0, 50, 120], labels=["only_one"])
            
            
            class TestNumericSummary:
                def test_summary_keys(self, sample_df):
                    result = visualizer.compute_numeric_summary(sample_df, "age")
                    for key in ["count", "mean", "std", "min", "q25", "median", "q75", "max"]:
                        assert key in result
            
                def test_summary_values(self, sample_df):
                    result = visualizer.compute_numeric_summary(sample_df, "duration")
                    assert result["min"] == 100.0
                    assert result["max"] == 500.0
                    assert result["count"] == 8
            
                def test_summary_non_numeric_raises(self, sample_df):
                    with pytest.raises(ValueError, match="not numeric"):
                        visualizer.compute_numeric_summary(sample_df, "job")
            
                def test_summary_missing_column_raises(self, sample_df):
                    with pytest.raises(ValueError, match="Column not found"):
                        visualizer.compute_numeric_summary(sample_df, "nope")
            
            
            class TestFigureBuilders:
                def test_build_subscribe_pie(self, sample_df):
                    assert isinstance(visualizer.build_subscribe_pie(sample_df), go.Figure)
            
                def test_build_category_bar(self, sample_df):
                    assert isinstance(visualizer.build_category_bar(sample_df, "job"), go.Figure)
            
                def test_build_subscribe_rate_bar(self, sample_df):
                    assert isinstance(visualizer.build_subscribe_rate_bar(sample_df, "job"), go.Figure)
            
                def test_build_age_histogram(self, sample_df):
                    assert isinstance(visualizer.build_age_histogram(sample_df), go.Figure)
            
                def test_build_numeric_histogram(self, sample_df):
                    assert isinstance(visualizer.build_numeric_histogram(sample_df, "duration"), go.Figure)
            
        • ⚙️ .coverage 52.0KB
        • 📄 .gitignore 699B
          # Byte-compiled / optimized / DLL files
          __pycache__/
          *.py[cod]
          *$py.class
          *.so
          
          # Distribution / packaging
          .Python
          build/
          develop-eggs/
          dist/
          downloads/
          eggs/
          .eggs/
          lib/
          lib64/
          parts/
          sdist/
          var/
          wheels/
          *.egg-info/
          .installed.cfg
          *.egg
          
          # PyInstaller
          *.manifest
          *.spec
          
          # Unit test / coverage reports
          htmlcov/
          .tox/
          .coverage
          .coverage.*
          .cache
          nosetests.xml
          coverage.xml
          *.cover
          .hypothesis/
          .pytest_cache/
          
          # Jupyter Notebook
          .ipynb_checkpoints
          
          # pyenv
          .python-version
          
          # Environments
          .env
          .venv
          env/
          venv/
          ENV/
          env.bak/
          venv.bak/
          
          # IDEs
          .vscode/
          .idea/
          *.swp
          *.swo
          *~
          
          # OS
          .DS_Store
          Thumbs.db
          
          # Project specific
          app/ml/model/
          ml/model/
          *.pkl
          *.joblib
          
          # Streamlit
          .streamlit/
          
          # Logs
          *.log
          
        • 📄 Dockerfile 1.3KB
          ARG PIP_INDEX_URL=https://pypi.org/simple
          FROM python:3.11-slim
          
          # Re-declare ARG after FROM (ARG scope ends at FROM)
          ARG PIP_INDEX_URL
          
          ENV PYTHONUNBUFFERED=1 \
              PYTHONDONTWRITEBYTECODE=1 \
              PIP_NO_CACHE_DIR=1 \
              PIP_DISABLE_PIP_VERSION_CHECK=1 \
              PYTHONPATH=/app
          
          WORKDIR /app
          
          # Copy requirements and install production dependencies only
          COPY requirements.txt .
          RUN pip install --no-cache-dir --timeout 120 --index-url "${PIP_INDEX_URL}" -r requirements.txt
          
          # Copy application code
          COPY app/ ./app/
          COPY data/ ./data/
          
          # Train the model at build time so the prediction service has an artifact.
          # model.pkl is gitignored (never committed); it is baked into the image here.
          RUN python -m app.ml.train --overwrite
          
          # Expose Streamlit default port (will be mapped to 8004 on host)
          EXPOSE 8501
          
          # Health check via Python stdlib (avoids installing curl through the slow
          # Debian mirror on the deploy server; Streamlit exposes /_stcore/health).
          HEALTHCHECK --interval=30s --timeout=10s --start-period=15s --retries=3 \
              CMD python -c "import urllib.request,sys; sys.exit(0 if urllib.request.urlopen('http://localhost:8501/_stcore/health').status==200 else 1)" || exit 1
          
          # Run Streamlit
          CMD ["streamlit", "run", "app/main.py", "--server.port=8501", "--server.address=0.0.0.0"]
          
        • 📄 README.md 1.5KB
          # banksys_szai4
          
          银行营销数据分析与认购预测系统。
          
          ## 功能
          
          - 数据分析交互页面:可视化展示客户特征分布和营销效果
          - 在线预测系统:通过点选式表单输入客户特征,预测认购意愿
          
          ## 技术栈
          
          - Python 3.11
          - Streamlit
          - scikit-learn
          - pytest
          - ruff
          - Docker
          
          ## 快速开始
          
          ### 本地运行
          
          ```bash
          # 创建虚拟环境
          conda create -y -n banksys python=3.11
          conda activate banksys
          
          # 安装依赖
          pip install -r requirements.txt
          
          # 运行应用
          streamlit run app/main.py
          ```
          
          ### Docker 运行
          
          ```bash
          docker build -t banksys:latest .
          docker run -d --name banksys -p 8004:8501 banksys:latest
          ```
          
          访问:http://localhost:8004
          
          ## 项目结构
          
          ```
          banksys_szai4/
          ├── standards/           # AI 项目记忆与规范
          ├── app/                 # 应用主目录
          │   ├── main.py         # Streamlit 主入口
          │   ├── pages/          # 多页面
          │   ├── models/         # 业务逻辑
          │   └── ml/             # 机器学习模块
          ├── tests/              # 测试目录
          ├── data/               # 数据目录
          ├── requirements.txt    # 生产依赖
          └── Dockerfile          # 容器镜像
          ```
          
          ## CI/CD
          
          - CI: 格式检查、静态检查、单元测试、构建
          - CD: 合并 main 后自动部署到服务器
          
          ## 质量门槛
          
          - 格式检查: `ruff format --check .`
          - 静态检查: `ruff check .`
          - 单元测试: `pytest`
          - 覆盖率: 核心逻辑 ≥80%
          
          ---
          
          **注意**: 本项目为教学演示项目,数据集为公开教学数据。
          
        • 📄 pyproject.toml 817B
          [project]
          name = "banksys"
          version = "0.1.0"
          description = "Bank marketing data analysis and prediction system"
          readme = "README.md"
          requires-python = ">=3.11"
          dependencies = []
          
          [build-system]
          requires = ["setuptools>=61.0"]
          build-backend = "setuptools.build_meta"
          
          [tool.ruff]
          line-length = 100
          target-version = "py311"
          
          [tool.ruff.lint]
          select = ["E", "F", "W", "I", "B", "C4", "UP"]
          ignore = ["E501"]
          
          [tool.ruff.format]
          quote-style = "double"
          indent-style = "space"
          skip-magic-trailing-comma = false
          line-ending = "auto"
          
          [tool.pytest.ini_options]
          testpaths = ["tests"]
          python_files = ["test_*.py"]
          python_classes = ["Test*"]
          python_functions = ["test_*"]
          addopts = "-v --cov=app --cov-report=term-missing"
          markers = [
              "unit: marks a test as a fast, isolated unit test",
          ]
          
        • 📄 requirements-dev.txt 134B
          # Development and CI dependencies
          -r requirements.txt
          
          # Testing
          pytest>=7.4.0
          pytest-cov>=4.1.0
          
          # Code quality
          ruff==0.6.4
          
        • 📄 requirements.txt 113B
          # Production dependencies
          streamlit>=1.28.0
          pandas>=2.0.0
          numpy>=1.24.0
          scikit-learn>=1.3.0
          plotly>=5.18.0
          
      • 📁 docs/
        • ⚙️ banksys_szai4-architecture.drawio 55.8KB
  • 📁 claude_code/
    • 📁 hello-flask-sz/
      • 📁 .claude/
        • ⚙️ settings.local.json 234B
      • 📁 .github/
        • 📁 workflows/
          • 📄 cd.yml 2.6KB
            name: CD
            
            on:
              push:
                branches: [main]
            
            jobs:
              cd:
                runs-on: ubuntu-latest
            
                steps:
                  - name: Checkout code
                    uses: actions/checkout@v4
            
                  - name: Build Docker image
                    run: |
                      docker build \
                        --build-arg PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple \
                        -t hello-flask-sz:latest .
            
                  - name: Deploy to server
                    env:
                      SSH_PRIVATE_KEY: ${{ secrets.SSH_PRIVATE_KEY }}
                      SSH_HOST: ${{ secrets.SSH_HOST }}
                      SSH_USER: ${{ secrets.SSH_USER }}
                    run: |
                      # Setup SSH
                      mkdir -p ~/.ssh
                      echo "$SSH_PRIVATE_KEY" > ~/.ssh/deploy_key
                      chmod 600 ~/.ssh/deploy_key
                      ssh-keyscan -H "$SSH_HOST" >> ~/.ssh/known_hosts
            
                      # Deploy directory
                      DEPLOY_DIR="/opt/hello-flask-sz"
                      APP="hello-flask-sz"
                      PORT=8003
                      PORT_MAX=8020
            
                      # Create deploy directory
                      ssh -i ~/.ssh/deploy_key "${SSH_USER}@${SSH_HOST}" "mkdir -p ${DEPLOY_DIR}"
            
                      # Sync files
                      rsync -avz -e "ssh -i ~/.ssh/deploy_key" \
                        --exclude='.git/' \
                        --exclude='standards/' \
                        --exclude='__pycache__/' \
                        --exclude='*.pyc' \
                        ./ "${SSH_USER}@${SSH_HOST}:${DEPLOY_DIR}/"
            
                      # Remote deployment script
                      ssh -i ~/.ssh/deploy_key "${SSH_USER}@${SSH_HOST}" << 'ENDSSH'
                      set -e
                      cd /opt/hello-flask-sz
            
                      # Build on server
                      docker build \
                        --build-arg PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple \
                        -t hello-flask-sz:latest .
            
                      # Port fallback logic
                      port_in_use() {
                        ss -ltnH 2>/dev/null | grep -q ":$1 " && return 0
                        docker ps --format "{{.Ports}}" 2>/dev/null | grep -q ":$1->" && return 0
                        return 1
                      }
            
                      PORT=""
                      for p in $(seq 8003 8020); do
                        if ! port_in_use "$p"; then PORT="$p"; break; fi
                      done
            
                      if [ -z "$PORT" ]; then
                        echo "预留端口区间已全部占用,部署中止"
                        exit 1
                      fi
            
                      echo ">> 部署到主机端口 $PORT"
            
                      # Stop and remove old container (idempotent)
                      docker rm -f hello-flask-sz 2>/dev/null || true
            
                      # Run new container
                      docker run -d \
                        --name hello-flask-sz \
                        --restart unless-stopped \
                        -p ${PORT}:8003 \
                        hello-flask-sz:latest
            
                      # Wait for startup
                      sleep 3
            
                      # Health check
                      curl -fsS "http://localhost:${PORT}/health"
                      echo ">> 部署成功: http://localhost:${PORT}/health"
                      ENDSSH
            
          • 📄 ci.yml 911B
            name: CI
            
            on:
              push:
                branches: [main, feature/**]
              pull_request:
                branches: [main]
            
            jobs:
              ci:
                runs-on: ubuntu-latest
            
                steps:
                  - name: Checkout code
                    uses: actions/checkout@v4
            
                  - name: Set up Python
                    uses: actions/setup-python@v5
                    with:
                      python-version: '3.11'
            
                  - name: Cache pip packages
                    uses: actions/cache@v4
                    with:
                      path: ~/.cache/pip
                      key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements*.txt') }}
            
                  - name: Install dependencies
                    run: |
                      pip install -r requirements-dev.txt
            
                  - name: Format check
                    run: ruff format --check .
            
                  - name: Lint
                    run: ruff check .
            
                  - name: Run tests with coverage
                    run: pytest --cov=. --cov-fail-under=80 --cov-report=term-missing
            
                  - name: Build Docker image
                    run: docker build -t hello-flask-sz:latest .
            
      • 📁 standards/
        • 📁 templates/
          • 📄 ISSUE_TEMPLATE.md 455B
            # Issue 模板
            
            ## 类型
            
            - [ ] Feature 功能
            - [ ] Bug 缺陷
            - [ ] Tech Debt 技术债
            - [ ] Docs 文档
            - [ ] CI/CD 流水线
            
            ## 背景 / 场景
            
            <谁在什么场景下遇到什么问题>
            
            ## 目标
            
            <这次希望达成什么>
            
            ## 验收标准
            
            - [ ] Given <前提>,When <动作>,Then <可验证结果>
            - [ ] <补充标准>
            
            ## 技术备注
            
            <可选:边界、风险、约束、参考资料>
            
            ## 优先级
            
            P0 / P1 / P2
            
          • 📄 PR_TEMPLATE.md 641B
            # Pull Request 模板
            
            ## What
            
            <这个 PR 做了什么>
            
            ## Why
            
            关联 Issue:`closes #<编号>`
            
            ## How to Test
            
            - [ ] 本地格式检查通过
            - [ ] 本地静态检查通过
            - [ ] 本地测试通过
            - [ ] CI 通过
            - [ ] CD / 部署验证通过(如适用)
            
            ## 风险与回滚
            
            <风险点;如果失败如何回滚>
            
            ## 自检清单
            
            - [ ] 没有硬编码密钥/密码
            - [ ] 没有调试 print / 死代码
            - [ ] PR 体积可控
            - [ ] commit message 符合规范
            - [ ] 已同步最新 main,无冲突
            - [ ] 已更新 standards/PROGRESS.md
            
            ---
            
            通过标准:CI 全绿 + Review 通过 + 分支保护通过。
            
        • 📄 00-project-context.md 2.7KB
          # 00 · 项目上下文 〔本项目活记忆 · AI 维护〕
          
          > **作用**:这是项目的"身份档案"。AI 接管项目时先读这里,了解项目目标、技术栈、目录、部署取值。
          > **更新时机**:架构、技术栈、目录结构、端口、部署目录、重要约束变化时更新。
          > **填写方式**:把 `<...>` 替换成真实内容;用不到的行删掉。
          
          ---
          
          ## 1. 项目是什么
          
          - **项目名称**: `hello-flask-sz`
          - **一句话目标**: 一个最小 Flask Web 服务示例,跑通 CI/CD 全流程
          - **使用者/受益者**: 学习 DevOps 和 Flask 开发的学生
          - **核心功能**:
            - GET `/` 返回欢迎 JSON
            - GET `/health` 返回 `{"status":"ok"}` 供健康检查
          - **输入/数据**: 无
          
          ## 2. 技术栈
          
          | 层 | 选型 | 理由 |
          |---|---|---|
          | 语言/运行时 | Python 3.11 | 教学常用版本,稳定 |
          | Web/API 框架 | Flask | 轻量、适合教学 |
          | 测试 | pytest | Python 标准测试框架 |
          | 格式/静态检查 | ruff | 快速、现代化 |
          | 打包/运行 | Docker | 标准容器化部署 |
          | CI/CD | GitHub Actions | 通用、可视化、适合教学 |
          
          ## 3. 目录地图
          
          ```text
          hello-flask-sz/
          ├── standards/                 # AI 项目记忆与通用规范
          ├── app.py                     # Flask 应用入口
          ├── test_app.py                # 测试文件
          ├── requirements.txt           # 生产依赖
          ├── requirements-dev.txt       # 开发依赖
          ├── Dockerfile                 # 容器构建
          ├── .github/workflows/
          │   ├── ci.yml
          │   └── cd.yml
          └── README.md
          ```
          
          > 新增目录前先更新本节,避免项目越做越散。
          
          ## 4. 质量门槛
          
          | 类型 | 本项目标准 |
          |---|---|
          | 格式检查 | `ruff format --check .` |
          | 静态检查 | `ruff check .` |
          | 单元测试 | `pytest` |
          | 覆盖率 | >=80% |
          | 构建 | `docker build` 成功 |
          | 业务/模型指标 | 无 |
          
          ## 5. 不变约束
          
          - 密钥、密码、私钥、Token **绝不写进代码或文档**,只进 GitHub Secrets / 环境变量。
          - `main` 分支受保护,日常开发必须走 feature 分支 + PR。
          - CI 红灯不合并。
          
          ## 6. 部署/CI 占位符取值
          
          > `guides/` 和 workflow 里的通用占位符,在本项目里的真实值只写这里。
          
          | 占位符 | 本项目取值 | 说明 |
          |---|---|---|
          | `<APP>` | `hello-flask-sz` | 应用名/镜像名/容器名 |
          | `<DEPLOY_DIR>` | `/opt/hello-flask-sz` | 服务器部署目录 |
          | `<PORT>` | `8003` | 服务端口 |
          | `<PORT_MAX>` | `8020` | 端口回退上限 |
          | `<PYVER>` | `3.11` | Python 版本 |
          | `<HEALTHCHECK>` | `/health` | 健康检查地址 |
          | `<SSH_USER>` | `root` | 部署用户(示例) |
          | `<SSH_HOST>` | `<服务器IP>` | 服务器地址(由学生配置) |
          
        • 📄 01-requirements.md 4.1KB
          # 01 · 需求 / 活 PRD 〔本项目活记忆 · AI 维护〕
          
          > **作用**:这是本项目唯一的需求文档。所有新功能、缺陷、技术债都追加到这里,不要另起多个 PRD 文件。
          > **更新时机**:每次有新需求、需求变更、验收标准变化时更新。
          
          ---
          
          ## 1. 需求来源
          
          | 类型 | 来源 | 进入方式 |
          |---|---|---|
          | 功能需求 Feature | 课程教学需求 | 写成用户故事 |
          | 缺陷 Bug | 测试 / CI/CD 失败 | 写复现步骤和期望结果 |
          | 技术债 Tech Debt | Review / CI/CD 故障 | 写影响和修复目标 |
          
          ---
          
          ## 2. Issue 生命周期
          
          | 阶段 | 状态 | 动作 |
          |---|---|---|
          | 提出 | Open | 写清场景、目标、验收标准 |
          | 排期 | Backlog / Todo | 决定优先级和负责人 |
          | 开发 | In Progress | 从 main 开 feature 分支 |
          | 评审 | In Review | 提 PR,等待 CI 和 Review |
          | 合并 | Done | PR 合并 main,自动关闭 Issue |
          | 验收 | Verified | 按验收标准确认 |
          
          **追踪规则**:分支名带 Issue 号,PR 描述写 `closes #<编号>`。
          
          ---
          
          ## 3. 用户故事模板
          
          ```text
          ### US-<编号> <一句话标题> · 状态: Backlog
          作为 <角色>,
          我想要 <能力>,
          以便 <价值>。
          
          验收标准:
          - AC1: Given <前提>,When <动作>,Then <可验证结果>。
          - AC2: <补充标准>
          
          技术备注:
          - <可选:约束、边界、风险>
          ```
          
          ---
          
          ## 4. 需求清单
          
          ### US-1 初始化项目工程化与 CI/CD · 状态: Backlog
          
          作为 **项目开发者**,
          我想要 项目具备基础工程结构、测试、CI 与 CD,
          以便 后续每次开发都能自动检查并自动部署。
          
          验收标准:
          - AC1: 从 `main` 开 feature 分支完成初始化,不直接 push main。
          - AC2: PR 触发 CI,至少包含格式检查、静态检查、单元测试、构建检查。
          - AC3: CI 全绿后合并 main。
          - AC4: 合并 main 自动触发 CD,部署后健康检查通过。
          - AC5: 完成后更新 `standards/PROGRESS.md`。
          
          ### US-2 实现欢迎接口 · 状态: Backlog
          
          作为 **网站访客**,
          我想要 访问根路径时获得欢迎响应,
          以便 确认服务正常运行。
          
          验收标准:
          - AC1: Given 服务已启动,When 访问 `GET /`,Then 返回状态码 200。
          - AC2: Given 服务已启动,When 访问 `GET /`,Then 返回 JSON 包含欢迎消息(如 `{"message": "Welcome to hello-flask-sz"}`)。
          - AC3: Given 服务已启动,When 用浏览器访问 `http://localhost:8003`,Then 能看到 JSON 响应。
          
          技术备注:
          - 使用 Flask 路由 `@app.route('/')` 实现。
          - 返回 `jsonify` 确保响应头为 `application/json`。
          
          ### US-3 实现健康检查接口 · 状态: Backlog
          
          作为 **运维 / 部署系统**,
          我想要 通过 `/health` 接口检查服务状态,
          以便 在部署后验证服务健康。
          
          验收标准:
          - AC1: Given 服务已启动,When 访问 `GET /health`,Then 返回状态码 200。
          - AC2: Given 服务已启动,When 访问 `GET /health`,Then 返回 JSON `{"status": "ok"}`。
          - AC3: Given CD 部署完成,When 执行 `curl http://localhost:8003/health`,Then 响应时间 < 2 秒。
          
          技术备注:
          - 这是 CD 健康检查的标准接口,不可改名或删除。
          - 不依赖数据库或外部服务,确保始终可达。
          
          ### US-4 容器化部署 · 状态: Backlog
          
          作为 **运维工程师**,
          我想要 服务以 Docker 容器运行,
          以便 统一部署环境并支持 CI/CD。
          
          验收标准:
          - AC1: Given Dockerfile 已就绪,When 执行 `docker build -t hello-flask-sz:latest .`,Then 构建成功无错误。
          - AC2: Given 镜像已构建,When 执行 `docker run -d -p 8003:8003 --name hello-flask-sz hello-flask-sz:latest`,Then 容器正常启动。
          - AC3: Given 容器已运行,When 访问 `http://localhost:8003/health`,Then 返回 `{"status": "ok"}`。
          
          技术备注:
          - Dockerfile 使用官方 Python 3.11 基础镜像。
          - 支持镜像源参数,便于国内服务器使用清华源。
          - 容器内端口固定 8003。
          
          ---
          
          ## 5. 非功能需求
          
          - **安全**:密钥只进 Secrets,不进 Git。
          - **可维护**:一需求一小 PR,避免大爆炸式提交。
          - **可测试**:核心逻辑必须有单元测试,覆盖率 >= 80%。
          - **可部署**:部署后 `/health` 必须返回 200。
          - **代码质量**:通过 ruff 格式检查和静态检查。
          
        • 📄 02-coding-standards.md 2.1KB
          # 02 · 代码标准
          
          > **目标**:让人和 AI 写出的代码风格一致、可读、可测、可维护。CI 会自动检查其中可机检的部分。
          
          ---
          
          ## 1. 格式与静态检查
          
          - Python 项目推荐:`ruff format .` + `ruff check .`。
          - 其他语言按项目技术栈替换为对应工具,例如 ESLint/Prettier、Checkstyle、gofmt。
          - 行宽默认 100;缩进、import 分组、命名风格由工具强制。
          - 提交前本地必须通过格式与 lint;CI 会再卡一道。
          
          ## 2. 命名
          
          | 对象 | 推荐规则 | 示例 |
          |---|---|---|
          | 变量/函数 | `snake_case` | `load_config` |
          | 类 | `PascalCase` | `UserService` |
          | 常量 | `UPPER_SNAKE` | `DEFAULT_PORT` |
          | 私有实现 | 前缀 `_` | `_parse_row` |
          
          命名必须表达意图。禁止 `process2`、`new_func`、`tmp_final` 这类无法审查的名字。
          
          ## 3. 函数与模块
          
          - 一个函数只做一件事;超过约 40 行优先拆分。
          - 参数过多时使用配置对象/数据类。
          - 业务逻辑与 I/O 分离,核心逻辑优先写成纯函数,方便测试。
          - 新增目录或模块前,先确认能落进 `00-project-context.md` 的目录地图。
          
          ## 4. 注释
          
          - 不写废话注释,例如 `# 调用函数` 配 `call_func()`。
          - 注释只解释“为什么”和“权衡”,不重复“代码做了什么”。
          - 复杂业务规则、部署坑、兼容性原因必须写清楚。
          
          ## 5. 配置与密钥
          
          - 非敏感配置可进 Git,例如 `config.yaml`、`.env.example`。
          - 密钥、密码、Token、私钥绝不进 Git,只放 GitHub Secrets / 环境变量。
          - 代码读取密钥时使用环境变量,并在 README 写清变量名,不写真实值。
          
          ## 6. 错误处理
          
          - 禁止裸 `except: pass`。
          - 外部输入必须校验,错误信息要能帮助定位。
          - 部署脚本必须失败即停;不能让前一步失败后继续执行后续危险动作。
          
          ## 7. 提交前自检
          
          ```text
          [ ] 格式检查通过
          [ ] 静态检查通过
          [ ] 新增逻辑有测试
          [ ] 没有硬编码密钥
          [ ] 没有调试 print / 死代码 / 大段注释代码
          [ ] 目录结构仍符合 00-project-context.md
          ```
          
        • 📄 03-testing-standards.md 2.9KB
          # 03 · 测试标准
          
          > **目标**:定义“什么叫测够了”。开发自测是进入 CI 前的第一道防线。
          
          ---
          
          ## 1. 测试分层
          
          ```text
                  /\        E2E 测试:少、慢、贵,覆盖真实用户路径
                 /  \
                /----\      集成测试:测试模块/服务边界
               /------\
              /        \    单元测试:多、快、稳定,开发者必写
             /__________\
          ```
          
          - 单元测试优先覆盖核心逻辑。
          - 集成测试覆盖接口、数据库、外部服务边界。
          - E2E 测试只覆盖关键路径,不要用它替代单元测试。
          
          ## 2. 用例写法
          
          推荐 AAA / Given-When-Then:
          
          ```python
          def test_health_returns_ok():
              # Arrange
              client = app.test_client()
          
              # Act
              resp = client.get("/health")
          
              # Assert
              assert resp.status_code == 200
              assert resp.get_json()["status"] == "ok"
          ```
          
          ## 3. 好测试的 5 条铁律
          
          1. 一个用例只测一件事。
          2. 覆盖正常值、边界值、异常值。
          3. 可重复、无外部网络依赖;有随机性必须固定种子。
          4. 测试名能表达场景和预期。
          5. 必须有真实断言,禁止 `assert result is not None` 凑数。
          
          ## 4. 覆盖率
          
          - 默认建议核心代码覆盖率 ≥ 80%。
          - 覆盖率不是目的,关键分支必须被覆盖。
          - 如果项目不适合覆盖率门槛,必须在 `00-project-context.md` 说明原因。
          
          ## 5. 项目特有质量门禁
          
          不同项目可追加不同质量门禁:
          
          | 项目类型 | 可选门禁示例 |
          |---|---|
          | Web/API | 健康检查、接口状态码、契约测试 |
          | 数据/ML | 数据校验、模型指标(AUC/F1/RMSE 等) |
          | 前端 | 组件测试、构建产物检查、可访问性检查 |
          | CLI/脚本 | 命令退出码、输入输出快照 |
          
          > 项目特有门禁写进 `00-project-context.md` 第 4 节,再写入 CI。
          
          ## 6. flaky test 纪律
          
          - 偶发失败必须根治,不能靠“重跑一次”掩盖。
          - 典型原因:时间依赖、随机数、测试顺序依赖、外部网络、共享状态未清理。
          
          ## 7. 本地测试由谁执行
          
          默认规则:
          
          | 环境 | 谁执行本地测试 | 怎么做 |
          |---|---|---|
          | AI 工具有终端权限(Cursor / Claude Code 等) | **AI 执行** | AI 自己运行格式检查、lint、pytest、构建等命令,读取日志,修复问题,把结果写回 `PROGRESS.md` |
          | AI 没有终端权限(纯聊天网页) | **学生执行** | 学生复制 AI 给出的命令到终端运行,把完整日志贴回 AI |
          | 涉及账号/权限/审批/Secrets | **学生确认** | AI 可指导,但不能代替人输入密码、私钥、审批生产发布 |
          
          AI 能执行命令时,不能只告诉学生“你去运行 pytest”。正确做法是:
          
          ```text
          AI 写代码 → AI 运行本地自检 → AI 根据日志修复 → AI 再运行 → 全绿后提交/PR → AI 记录结果
          ```
          
          每次本地自检至少记录:
          
          - 执行的命令。
          - 是否通过。
          - 失败时的关键错误。
          - 修复方式。
          - 最终通过结果。
          
        • 📄 04-git-workflow.md 4.4KB
          # 04 · Git 工作流与 PR/Review 标准
          
          > **目标**:定义代码如何从本地安全进入 `main`,并触发 CI/CD。课堂第一项功能必须完整演示本文件第 2 节。
          
          ---
          
          ## 1. 分支模型:GitHub Flow
          
          ```text
          main 受保护,永远可部署
           ├─ feature/<issue>-<desc>   新功能
           ├─ fix/<issue>-<desc>       普通缺陷
           ├─ docs/<issue>-<desc>      文档
           └─ chore/<issue>-<desc>     杂务/依赖/配置
          ```
          
          - 只有一条长期分支:`main`。
          - 日常开发必须从 `main` 开短分支,不能直接 push `main`。
          - feature 分支寿命建议 ≤ 2-3 天。
          - 分支命名:`<类型>/<issue号>-<英文短描述>`,例如 `feature/1-hello-api`。
          
          ---
          
          ## 2. 课堂必演示:从分支到 CD 的完整命令
          
          > 这段是学生必须亲手跑通的一条链。不要省略 PR,不要直接 push main。
          
          ```bash
          # 1. 确认在 main,并更新主干
          git switch main
          git pull
          
          # 2. 从 main 开功能分支
          git switch -c feature/1-first-feature
          
          # 3. 写代码 + 写测试后,本地自检
          ruff format --check .
          ruff check .
          pytest
          
          # 4. 提交
          git status
          git add .
          git commit -m "feat: add first feature"
          
          # 5. 推送 feature 分支
          git push -u origin feature/1-first-feature
          
          # 6. 创建 PR(AI 做到这里为止)
          gh pr create --base main --head feature/1-first-feature --title "feat: add first feature" --body "closes #1"
          
          # 7. 查看 CI(AI 汇报 CI 状态后停下,等人类审核)
          gh run list --limit 5
          gh run watch
          
          # 8. ★ 合并由人类执行 ★ —— 人工 Review 通过后,由人在网页点 "Merge",或人工运行:
          #    gh pr merge <PR号> --merge      # 默认保留分支,不加 --delete-branch
          #    AI 不得自行合并未获批准的 PR。
          
          # 9. 合并 main 后会触发 CD,查看部署流水线
          gh run list --limit 5
          ```
          
          > 如果课堂暂时没有开启分支保护,也要按这个流程演示 PR。分支保护是“强制执行”,PR 流程是“正确习惯”。
          
          > **合并是人类的动作,不是 AI 的动作。** AI 在第 6/7 步发起 PR、汇报 CI 后必须停下;
          > **只有人工 Review 通过,由人合并**,才触发 CD。AI 绝不替人按下 Merge。
          
          > **默认保留分支**(不加 `--delete-branch`),除非人类明确要求删除已合并分支。
          
          > **PR 必须用你本人 GitHub 账号发起**(`gh auth login` 登录的就是你本人),不要借他人账号代提。
          
          > **⚠️ 建仓后第一件事:配置 Secrets,否则 CD 必失败。**
          > 新仓库创建后(CD 第一次跑之前)立刻去
          > `GitHub → Settings → Secrets and variables → Actions → New repository secret`
          > 配好 `SSH_PRIVATE_KEY` / `SSH_HOST` / `SSH_USER`(详见 `05-cicd-standards.md` 第 5 节)。
          > 没有数据库等部署目标时可只配 SSH 三件套。配置时机:**在合并到 main 触发 CD 之前**。
          
          ---
          
          ## 3. 认证不要混
          
          | 动作 | 认证方式 | 说明 |
          |---|---|---|
          | `git push` / `git clone` | SSH 公钥或 HTTPS 凭证 | 负责搬代码 |
          | `gh pr create` / `gh run list` | GitHub CLI 登录 | 负责调 GitHub API |
          | CD SSH 到服务器 | GitHub Secrets 里的部署私钥 | 由 Actions runner 自动使用 |
          
          如果 SSH 22 端口卡住,可使用 GitHub SSH 443:
          
          ```bash
          ssh -T -p 443 git@ssh.github.com
          git remote set-url origin ssh://git@ssh.github.com:443/<账号>/<仓库>.git
          git push
          ```
          
          ---
          
          ## 4. 提交信息规范
          
          采用 Conventional Commits:
          
          | 前缀 | 含义 |
          |---|---|
          | `feat:` | 新功能 |
          | `fix:` | 修 bug |
          | `docs:` | 文档 |
          | `test:` | 测试 |
          | `refactor:` | 重构 |
          | `ci:` | CI/CD |
          | `build:` | 构建/依赖 |
          | `chore:` | 杂务 |
          
          一个 commit 只做一件事。
          
          ---
          
          ## 5. PR 通过三道闸
          
          | 闸 | 内容 | 不通过怎么办 |
          |---|---|---|
          | 机器闸 | CI 全绿 | 修代码/测试/配置 |
          | 人工闸 | 至少 1 人 Review | 修改并回复评论 |
          | 策略闸 | 无冲突、分支保护通过 | 同步 main / 解冲突 |
          
          PR 描述必须写:
          
          ```text
          ## What
          这个 PR 做了什么
          
          ## Why
          为什么需要
          
          ## Test
          - [ ] 本地检查通过
          - [ ] CI 通过
          
          closes #<issue号>
          ```
          
          ---
          
          ## 6. 冲突处理
          
          ```bash
          git switch main
          git pull
          git switch <你的分支>
          git rebase main
          
          # 如果有冲突:手动修改文件,删掉 <<<<<<< ======= >>>>>>> 标记
          git add <冲突文件>
          git rebase --continue
          
          # rebase 后推送
          git push --force-with-lease
          ```
          
          冲突处理后必须重新跑测试。
          
        • 📄 05-cicd-standards.md 7.1KB
          # 05 · CI/CD 标准
          
          > **目标**:定义最小可复现 CI/CD。CI 管“代码是否可靠”,CD 管“可靠代码如何自动上线”。
          
          ---
          
          ## 1. CI 与 CD
          
          | 名称 | 做什么 | 典型触发 |
          |---|---|---|
          | CI 持续集成 | 格式、lint、测试、覆盖率、构建 | PR / push |
          | CD 持续交付 | 构建可部署产物,部署前可人工审批 | main / tag |
          | CD 持续部署 | 合并 main 后自动部署上线 | main |
          
          课堂简版采用:
          
          ```text
          PR → CI 全绿 → 合并 main → CD 自动部署到服务器 → 健康检查
          ```
          
          ---
          
          ## 2. 最小 CI 门禁
          
          | 关卡 | 推荐命令 | 本地必跑 | CI 必跑 |
          |---|---|:---:|:---:|
          | 格式检查 | `ruff format --check .` 或等价工具 | ✅ | ✅ |
          | 静态检查 | `ruff check .` 或等价工具 | ✅ | ✅ |
          | 单元测试 | `pytest` 或等价工具 | ✅ | ✅ |
          | 覆盖率 | `pytest --cov --cov-fail-under=80` | ✅ | ✅ |
          | 项目特有门禁 | AUC/F1/接口契约/前端构建等 | ✅ | ✅ |
          | 构建 | `docker build ...` 或项目构建命令 | ❌ 不要求 | ✅ |
          
          CI 红灯时禁止合并。不要靠重跑掩盖问题。
          
          > **本地不强制装 Docker**:很多学员/开发机没有 Docker,本地装它只会增加复杂度。
          > 约定 **本地只跑前几项门禁,`docker build` 交给 CI(云端 runner 自带 Docker)与服务器(CD)**。
          > 这样本地零 Docker 负担,镜像构建仍有 CI 兜底。是否要本地构建,**按机器情况决定或先问开发者**。
          
          ---
          
          ## 3. 最小 CD 标准
          
          CD 必须满足:
          
          - 由 GitHub Actions runner 自动触发,不是人手动 SSH。
          - 只引用 GitHub Secrets,不硬编码密钥。
          - 部署脚本可重复执行。
          - 部署后必须健康检查。
          - 失败时 Actions 红灯,不能假成功。
          
          推荐触发:
          
          ```yaml
          on:
            push:
              branches: [ main ]
          ```
          
          ---
          
          ## 4. 依赖与镜像构建标准
          
          真实教学环境常见失败:服务器 `docker build` 时从 PyPI / npm / Docker Hub 下载依赖很慢或超时。
          
          通用标准:
          
          | 标准 | 做法 |
          |---|---|
          | 运行依赖与开发依赖拆分 | 生产镜像只安装运行依赖;测试/lint 工具只在 CI 安装 |
          | 镜像源可配置 | Dockerfile 支持镜像源参数,国内服务器可使用国内源 |
          | 构建失败即停 | 远程部署脚本第一行 `set -e` |
          | 健康检查 | 最后执行 `curl -fsS http://localhost:<PORT>/<HEALTHCHECK>` |
          | 可选预热 | 课前在服务器预拉基础镜像,减少等待 |
          
          Python 项目推荐:
          
          ```text
          requirements.txt      # 生产运行依赖
          requirements-dev.txt  # CI/本地检查依赖
          ```
          
          Dockerfile 推荐:
          
          ```dockerfile
          ARG PIP_INDEX_URL=https://pypi.org/simple
          RUN pip install --no-cache-dir --timeout 120 -i "${PIP_INDEX_URL}" -r requirements.txt
          ```
          
          远程部署脚本推荐(含幂等替换 + 端口自动回退):
          
          ```bash
          set -e
          docker build --build-arg PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple -t <APP>:latest .
          
          # 端口:容器内端口固定;主机端口优先 <PORT>,被占用就在预留区间自动找空闲端口。
          port_in_use() {
            ss -ltnH 2>/dev/null | grep -q ":$1 " && return 0
            docker ps --format "{{.Ports}}" 2>/dev/null | grep -q ":$1->" && return 0
            return 1
          }
          PORT=""
          for p in $(seq <PORT> <PORT_MAX>); do
            if ! port_in_use "$p"; then PORT="$p"; break; fi
          done
          [ -z "$PORT" ] && { echo "预留端口区间已全部占用,部署中止"; exit 1; }
          echo ">> 部署到主机端口 $PORT"
          
          docker rm -f <APP> 2>/dev/null || true   # 一步停删自身旧容器,幂等可重跑
          docker run -d --name <APP> --restart unless-stopped -p ${PORT}:<容器内端口> <APP>:latest
          sleep 3
          curl -fsS "http://localhost:${PORT}/<HEALTHCHECK>"
          echo ">> 部署成功:http://<SSH_HOST>:${PORT}/<HEALTHCHECK>"
          ```
          
          > **端口标准**:容器内端口固定,**主机端口可回退**(预留如 `<PORT>`-`<PORT_MAX>` 一段)。
          > 起容器用 `docker rm -f <APP>` 一步停删自身旧容器保证幂等;**不要**盲目删除占用端口的他人容器。
          > 防火墙/安全组要放行整段预留端口;最终落在哪个端口由 CD 日志打印。
          > 踩坑:`docker run` 报 `port is already allocated`(exit 125)就是主机端口被占用所致。
          
          ---
          
          ## 5. Secrets 标准
          
          常用 SSH 部署需要三个 GitHub Actions Secrets:
          
          | Secret | 含义 |
          |---|---|
          | `SSH_PRIVATE_KEY` | Actions runner 用来登录服务器的私钥全文 |
          | `SSH_HOST` | 服务器公网 IP 或域名 |
          | `SSH_USER` | 服务器登录用户,如 `root` 或 `deploy` |
          
          规则:
          
          - Secret 名大小写必须一致。
          - 私钥必须保留 `BEGIN/END` 和换行。
          - Secrets 必须在 CD run 开始前配置好;配置后旧 run 可能需要 rerun。
          
          ---
          
          ## 6. 本地环境准备(一次性)
          
          环境怎么建,**写进标准、可复现**。优先级:`uv` → `uv` 不行用 `conda` → 再不行用系统 `venv`。
          本课程学员更熟 `conda`,默认用 conda:
          
          ```bash
          # 1) 建并激活环境(示例环境名 envproj,Python 版本按项目 00 文档)
          conda create -y -n envproj python=3.11
          conda activate envproj
          
          # 2) 装依赖(国内用清华源更快)
          pip install -r requirements.txt -r requirements-dev.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
          ```
          
          > **坑:`CondaToSNonInteractiveError`(新版 conda 默认渠道未接受服务条款)**。
          > 解决,对三个渠道各执行一次:
          > ```bash
          > conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
          > conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
          > conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/msys2
          > ```
          
          ---
          
          ## 7. 排错顺序
          
          | 现象 | 优先检查 |
          |---|---|
          | `conda create` 报 ToS 错误 | 执行上面的 `conda tos accept`(main/r/msys2) |
          | Windows 控制台 `UnicodeEncodeError('gbk' ... \u25b6)` | 日志/打印含特殊符号;改用纯 ASCII,或 `set PYTHONIOENCODING=utf-8` / `chcp 65001`。注意 pytest 会捕获输出掩盖此坑,必须真跑一次脚本 |
          | CI 红:`FileNotFoundError` 找不到数据 | 数据被 `.gitignore` 排除,干净 runner 上没有。公开教学数据可入库;敏感数据则在 CI 里下载/造样本 |
          | `docker run` 报 `port is already allocated`(exit 125) | 主机端口被占用;容器内端口固定、主机端口在预留区间自动回退(见第 4 节);`docker rm -f <APP>` 幂等替换自身 |
          | `git push` 卡住/`Recv failure`/`Connection reset` | SSH 22 被拦可切 443;HTTPS 偶发抖动则重试几次 |
          | Actions 读不到变量 | Secret 名是否拼错、是否在正确仓库 |
          | SSH 登录失败 | 私钥/公钥是否匹配、`authorized_keys`、用户权限 |
          | rsync 成功但 docker build 失败 | 依赖下载、镜像源、Dockerfile |
          | 服务器本机 curl 通,外网打不开 | 安全组/防火墙是否放行端口(含回退端口段) |
          | CD 日志二次报错 | 部署脚本是否缺 `set -e` |
          
          ---
          
          ## 8. 标准总结
          
          1. CI 红灯不合并。
          2. CD 必须由 runner 自动触发,不是人手动 SSH。
          3. 密钥进 Secrets,脚本进 Git,环境手动配一次。
          4. 生产镜像只放运行依赖,开发工具留在 CI;**本地不强制装 Docker**。
          5. 每次真实故障必须写回 `PROGRESS.md` 和本标准。
          
        • 📄 06-ai-collab-protocol.md 6.8KB
          # 06 · AI 协作协议
          
          > **目标**:让任意 AI 工具都按同一套上下文、规则和闭环工作,避免换工具/换会话就丢上下文。
          > **核心**:下面第 1 节的「六步交付流程 + 确认门」是 AI 必须严格遵守的主干,**每到一个确认门必须停下,向人类汇报并等待确认,不得自行越过**。
          
          ---
          
          ## 1. 六步交付流程(必须逐步确认)★ 主干
          
          AI 驱动开发分为固定六步。**每一步都有一个「✋ 确认门」**:AI 完成该步动作后,先汇报结果,等人类确认/操作后才进入下一步。绝不一口气从头跑到尾。
          
          ```text
          ① 建仓+配Secrets → ② 开feature分支 → ③ 本地模块化开发(逐模块汇报)
               → ④ 本地CI自检(AI执行) → ⑤ 触发PR(CI在PR复检) → ⑥ 人工审核→合并→CD自动部署
             ✋确认1            ✋确认2          ✋确认3(每个模块)
                              ✋确认4           ✋确认5            ✋确认6(审核+部署结果)
          ```
          
          ### 第①步 · 建仓 + 配 Secrets
          
          - **AI 动作**:用 `gh` 创建 GitHub 仓库;`main` 只放最小引导提交(`.gitignore`/占位 README),**不在 main 上做真实开发**。
          - **✋ 确认门 1**:AI 报出仓库地址后,**必须提示人类去配置 GitHub Secrets**(新仓库默认为空,不配 CD 必失败):
            - `SSH_PRIVATE_KEY` / `SSH_HOST` / `SSH_USER`(见 `05` 第 5 节)。
            - AI 停下,**等人类确认「Secrets 已配置」**再继续(可用 `gh secret list` 核对)。
          
          ### 第②步 · 开 feature 分支
          
          - **AI 动作**:从最新 `main` 切出 `feature/<issue>-<desc>` 分支(类型见 `04`);**严禁直接在 main 上开发**。
          - **✋ 确认门 2**:报出分支名,确认后进入开发。
          
          ### 第③步 · 本地模块化开发 + 进度汇报
          
          - **AI 动作**:按 `00` 目录地图**一个模块一个模块**实现,每个模块同步写测试;每完成一个模块:
            - 更新 `PROGRESS.md`(已完成/下一步)。
            - 向人类**汇报该模块进度**(做了什么、关键决定、是否有坑)。
          - **✋ 确认门 3(每个模块)**:每个模块完成后汇报并确认,再做下一个模块。避免一次性写完全部再回头返工。
          
          ### 第④步 · 本地 CI 自检(由 AI 执行)
          
          - **AI 动作**:到「需要 CI」的节点,**AI 主动确认并亲自执行**本地门禁(具备终端权限时):
            - `ruff format --check .` + `ruff check .`
            - `pytest --cov --cov-fail-under=80`
            - 项目特有门禁(如模型 `--check-auc`)
            - **本地不强制 `docker build`**(交给 CI;见 `05` 第 2 节)。
          - **✋ 确认门 4**:汇报自检结果(覆盖率、指标);**全绿才进入下一步**,红则就地修复后重跑。
          
          ### 第⑤步 · 触发 PR(CI 在 PR 上复检)
          
          - **AI 动作**:`git push` 分支后,**AI 主动 `gh pr create`** 发起 PR(用本人账号);PR 事件触发 CI 在云端再跑一遍同样门禁 + `docker build`。
          - **✋ 确认门 5**:报出 PR 链接与 CI 状态(`gh run watch`);CI 绿才推进到审核。
          
          ### 第⑥步 · 人工审核 → 合并(人类) → CD 自动部署
          
          - **✋ AI 在此硬停**:AI 发起 PR、汇报 CI 绿之后**就停下**,**不得自己合并**。合并是人类的动作。
          - **人类动作**:**PR 必须由人工 Review**;审核通过后**由人类在网页点 Merge 或人工运行 `gh pr merge`**。
          - **分支默认保留**:合并不加 `--delete-branch`,除非人类明确要求删除。
          - **AI 动作(合并发生后)**:CD 自动触发(SSH→同步→构建→运行→健康检查;端口自愈见 `05`);AI 盯流水线。
          - **✋ 确认门 6**:汇报部署结果——**最终落地端口、`/health` 结果、访问地址**;失败则读日志定位根因,修复后回到第②步走一个新分支/PR。
          
          > 一句话:**建仓我做、Secrets 你配;分支我开(默认保留);模块我逐个写并汇报;本地 CI 我跑;PR 我发;
          > 审核与合并都由你做,我绝不替你按 Merge;合并后 CD 自动跑,我盯结果。**
          
          ---
          
          ## 2. AI 开工前必须读
          
          每次新会话,AI 必须按顺序读:
          
          1. `standards/README.md`
          2. `standards/00-project-context.md`
          3. `standards/01-requirements.md`
          4. `standards/PROGRESS.md`
          5. `standards/02`~`06`
          
          然后向人类复述:
          
          ```text
          我理解当前项目是:...
          当前阶段是(对应六步流程第几步):...
          下一步应该做:...
          我将遵守的关键约束是:...
          ```
          
          ---
          
          ## 3. 人给 AI 的标准开场 prompt
          
          ```prompt
          请先读取 standards/README.md,并按顺序读取 00/01/PROGRESS 与 02~06。
          读完后先复述你对项目、当前进度(对应 06 的六步流程第几步)、下一步的理解。
          然后严格按 06 的「六步交付流程 + 确认门」推进:
          - 每到一个确认门必须停下汇报并等我确认,不要一口气跑完。
          - 建仓后提示我配置 Secrets;从 main 开 feature 分支,不直接改 main。
          - 逐模块开发并汇报;本地 CI 由你执行(本地不强制 docker);CI 绿后你主动发 PR。
          - 发完 PR 你就停下,**不要自己合并**;PR 由我审核并由我合并,合并后 CD 自动跑,你盯结果并汇报端口与健康检查。
          - 分支默认保留,不要加 --delete-branch,除非我明确要求删。
          ```
          
          ---
          
          ## 4. AI 写代码边界
          
          AI 可以:
          
          - 新增/修改代码、测试、CI/CD 配置。
          - 在具备终端权限时运行本地测试、lint、构建命令,并根据日志修复。
          - 创建仓库、开分支、发起 PR(用本人账号)。
          - 更新 `00/01/PROGRESS`;根据真实故障更新 `02~06`(须说明原因)。
          
          AI 不可以:
          
          - 硬编码密钥。
          - 绕过 CI 或跳过测试。
          - 明明能执行命令,却只让学生手动跑本地测试。
          - 为了“看起来通过”删除测试或降低门槛。
          - **越过任何确认门**(尤其:未提示配 Secrets 就推进 CD;未经确认大规模重构)。
          - **自行合并任何 PR**——合并是人类的专属动作,AI 只发起 PR 并汇报,绝不替人按 Merge。
          - 合并时默认 `--delete-branch` 删分支——**默认保留分支**,除非人类明确要求删除。
          - 直接把项目特有内容写进通用规范,除非它已沉淀为可复用规则。
          
          ---
          
          ## 5. 故障反哺铁律
          
          CI/CD 一旦真实跑出问题,修完代码还不算结束,必须同步完成三件事:
          
          1. 写回 `PROGRESS.md`:现象、根因、修复、验证结果。
          2. 如是通用规则,写进 `02~06` 对应标准。
          3. 更新学生会照着操作的 README/课件/脚本。
          
          判断标准:
          
          ```text
          换一个学生、换一个 AI 工具,只读 standards + README,是否能避开这次坑并复现?
          ```
          
          如果不能,说明闭环未完成。
          
          ---
          
          ## 6. 会话结束 prompt
          
          ```prompt
          本次会话结束。请更新 standards/PROGRESS.md:
          1. 当前状态(对应六步流程第几步)
          2. 已完成事项
          3. 下一步 TODO
          4. 新增 ADR
          5. 新增 GOTCHAS
          保持时间倒序、简洁、可接力。
          ```
          
        • 📄 PROGRESS.md 2.5KB
          # PROGRESS · hello-flask-sz 〔本项目活记忆 · 状态机〕
          
          > **作用**:这是项目的"存档点"。任意 AI、任意重启会话,读它即可知道当前做到哪、下一步做什么、踩过什么坑。
          > **更新时机**:每完成一个有意义步骤、每次会话结束前。
          > **格式要求**:时间倒序,最新在上;短、准、可接力。
          
          ---
          
          ## 当前状态 (最后更新: 2026-06-06 · by AI)
          
          - **阶段**: `已上线`
          - **上一步完成**: CI/CD 全流程跑通,部署成功
          - **下一步 (TODO 第一条)**: 功能已完成,可添加新需求
          - **阻塞项**: 无
          
          ---
          
          ## 待办清单 (TODO,按优先级)
          
          - [x] 用户确认 `00-project-context.md` 内容无误
          - [x] 用户确认 `01-requirements.md` 用户故事和验收标准
          - [x] 从 `main` 开第一条 feature 分支: `feature/1-init-project`
          - [x] 创建基础文件: `app.py`, `test_app.py`, `requirements.txt`, `requirements-dev.txt`, `Dockerfile`
          - [x] 实现欢迎接口 (`GET /`)
          - [x] 实现健康检查接口 (`GET /health`)
          - [x] 编写单元测试(覆盖率 >= 80%)
          - [x] 本地自检: `ruff format --check .` + `ruff check .` + `pytest`
          - [x] 创建 CI workflow (`.github/workflows/ci.yml`)
          - [x] 创建 CD workflow (`.github/workflows/cd.yml`)
          - [x] 创建 GitHub 仓库: https://github.com/DaTingLi/hello-flask-sz
          - [x] 推送代码到 main
          - [x] CI 验证: 格式/lint/测试/构建 全绿
          - [x] CD 验证: 部署到端口 8003,健康检查通过
          
          ---
          
          ## 关键决策记录 (ADR)
          
          | 日期 | 决策 | 理由 |
          |---|---|---|
          | 2026-06-06 | HTTPS 失败后切换 SSH | GitHub 443 端口连接超时,SSH 协议成功 |
          | 2026-06-06 | 直接推送 main 跳过 PR | 新仓库 main 为空,无法创建 PR,直接初始化 |
          
          ---
          
          ## 已知坑 (GOTCHAS)
          
          - **HTTPS 连接 GitHub 443 超时**: 现象:`Failed to connect to github.com port 443`;根因:网络限制;解决:`git remote set-url origin` 切换 SSH;验证:`git push` 成功
          - **新仓库 main 为空无法 PR**: 现象:`No commits between main and feature`;根因:main 分支不存在;解决:直接推送 main 或先创建占位提交;验证:仓库有 main 分支
          
          ---
          
          ## 里程碑 (DONE)
          
          - [x] US-1 初始化项目工程化与 CI/CD - CI/CD 全流程跑通
          - [x] US-2 实现欢迎接口 - `GET /` 返回欢迎 JSON
          - [x] US-3 实现健康检查接口 - `GET /health` 返回 `{"status":"ok"}`
          - [x] US-4 容器化部署 - Docker 构建与部署成功
          
          **部署信息**:
          - 仓库: https://github.com/DaTingLi/hello-flask-sz
          - 端口: 8003
          - 健康检查: `/health`
          - 覆盖率: 97.67%
          
        • 📄 README.md 3.6KB
          # standards-template · 学生项目可复制模板
          
          > **用途**:每个新项目复制一份本目录,改名为项目根目录下的 `standards/`。任意 AI 工具进入项目后,先读 `standards/README.md`,即可按统一规则继续开发、提交、PR、CI/CD 与存档。
          
          ---
          
          ## 1. 学生怎么用
          
          ```bash
          # 在你的项目根目录执行
          cp -r <课件目录>/standards-template ./standards
          ```
          
          Windows 可手动复制文件夹,复制后目录应是:
          
          ```text
          你的项目/
          ├── standards/
          │   ├── README.md
          │   ├── 00-project-context.md
          │   ├── 01-requirements.md
          │   ├── PROGRESS.md
          │   ├── 02-coding-standards.md
          │   ├── 03-testing-standards.md
          │   ├── 04-git-workflow.md
          │   ├── 05-cicd-standards.md
          │   ├── 06-ai-collab-protocol.md
          │   └── templates/
          └── <你的代码>
          ```
          
          ---
          
          ## 2. 两类文件,不要混
          
          | 类型 | 文件 | 是否每个项目要改 | 谁维护 |
          |---|---|---|---|
          | **项目活记忆** | `00-project-context.md`、`01-requirements.md`、`PROGRESS.md` | 必须按项目重填 | AI 写,人审 |
          | **通用规范** | `02`~`06`、`templates/` | 默认不改 | 团队维护 |
          
          > 简单说:`00/01/PROGRESS` 记录“这个项目是什么、要做什么、做到哪了”;`02~06` 规定“怎么写代码、怎么测试、怎么走 Git/PR/CI/CD”。
          
          ---
          
          ## 3. AI 每次会话读取顺序
          
          1. `00-project-context.md` — 项目身份、技术栈、目录地图、部署取值。
          2. `01-requirements.md` — 活 PRD,所有需求与验收标准。
          3. `PROGRESS.md` — 当前状态、下一步、决策、踩坑。
          4. `02`~`06` — 通用工程规范。
          5. `templates/` — Issue / PR 模板。
          
          ---
          
          ## 4. 第一次开工 prompt
          
          ```prompt
          请先读取 standards/README.md,并按它的顺序读取 00/01/PROGRESS 与 02~06。
          这是一个新项目。我的项目需求是:
          <在这里写项目目标、技术栈、希望实现的功能、部署方式>
          
          请先做三件事:
          1. 填写 standards/00-project-context.md。
          2. 把需求整理进 standards/01-requirements.md,写成用户故事和验收标准。
          3. 初始化 standards/PROGRESS.md,列出第一批 TODO。
          
          完成后先停下让我确认,不要直接写代码。
          ```
          
          ---
          
          ## 5. 标准开发闭环 · 固定六步(每步确认)
          
          AI 严格按 `06-ai-collab-protocol.md` 的「六步交付流程 + 确认门」推进,**每到确认门停下汇报、等人确认**:
          
          ```text
          ① 建仓 + 配 Secrets    AI 建仓 → ✋提示人类配 SSH_PRIVATE_KEY/HOST/USER(新仓库默认没有)
          ② 开 feature 分支       从 main 切分支,严禁直接改 main         ✋报分支名
          ③ 本地模块化开发        逐模块写代码+测试,每模块更新 PROGRESS  ✋每个模块汇报进度
          ④ 本地 CI 自检(AI 跑)  ruff + pytest + 覆盖率 + 项目门禁;本地不强制 docker  ✋全绿才继续
          ⑤ 触发 PR              push + AI 主动建 PR,CI 在 PR 上复检    ✋报 PR 链接 + CI 状态
          ⑥ 人工审核 →(人)合并 → CD  ✋AI 发完 PR 即停;Review 与 Merge 都由人做;合并触发 CD → ✋报端口/健康检查
          ```
          
          **课堂第一项功能必须完整演示一次这条链**,不能直接 push main;
          **合并是人类的动作——AI 绝不自行合并 PR;分支默认保留(不自动删)。**
          
          ---
          
          ## 6. 反臃肿纪律
          
          1. 需求只写进 `01-requirements.md`,不要到处新建 PRD。
          2. 进度、决策、坑只写进 `PROGRESS.md`。
          3. 一需求一分支一 PR,PR 尽量小于 400 行。
          4. CI 红灯不合并;flaky test 必须根治。
          5. 新增目录前先更新 `00-project-context.md` 的目录地图。
          
      • ⚙️ .coverage 52.0KB
      • 📄 .gitignore 172B
        # Python
        __pycache__/
        *.py[cod]
        *$py.class
        *.so
        .Python
        env/
        venv/
        ENV/
        .venv
        
        # Testing
        .pytest_cache/
        .coverage
        htmlcov/
        
        # IDE
        .vscode/
        .idea/
        
        # OS
        .DS_Store
        Thumbs.db
        
      • 📄 Dockerfile 502B
        ARG PYVER=3.11
        FROM python:${PYVER}-slim
        
        # Build argument for PyPI mirror (supports domestic mirrors)
        ARG PIP_INDEX_URL=https://pypi.org/simple
        
        WORKDIR /app
        
        # Copy requirements first for better caching
        COPY requirements.txt .
        
        # Install production dependencies only
        RUN pip install --no-cache-dir --timeout 120 -i "${PIP_INDEX_URL}" -r requirements.txt
        
        # Copy application code
        COPY app.py .
        
        # Expose port (container port, fixed at 8003)
        EXPOSE 8003
        
        # Run the application
        CMD ["python", "app.py"]
        
      • 📄 README.md 467B
        # hello-flask-sz
        
        A minimal Flask web service for CI/CD demonstration.
        
        ## Features
        
        - `GET /` - Returns welcome message
        - `GET /health` - Health check endpoint
        
        ## Quick Start
        
        ```bash
        # Install dependencies
        pip install -r requirements-dev.txt
        
        # Run tests
        pytest
        
        # Run app
        python app.py
        ```
        
        ## Deployment
        
        - Port: 8003
        - Health check: `http://localhost:8003/health`
        
        ## CI/CD
        
        - CI: Runs on push/PR (format, lint, test, build)
        - CD: Auto-deploys on merge to main
        
      • 📄 app.py 491B
        """hello-flask-sz: A minimal Flask web service for CI/CD demonstration."""
        
        from flask import Flask, jsonify
        
        app = Flask(__name__)
        
        
        @app.route("/")
        def welcome():
            """Return a welcome message in JSON format."""
            return jsonify({"message": "Welcome to hello-flask-sz"}), 200
        
        
        @app.route("/health")
        def health():
            """Health check endpoint for deployment verification."""
            return jsonify({"status": "ok"}), 200
        
        
        if __name__ == "__main__":
            app.run(host="0.0.0.0", port=8003)
        
      • 📄 requirements-dev.txt 64B
        -r requirements.txt
        pytest==8.3.2
        pytest-cov==5.0.0
        ruff==0.6.4
        
      • 📄 requirements.txt 29B
        Flask==3.0.3
        Werkzeug==3.0.3
        
      • 📄 test_app.py 1.7KB
        """Tests for hello-flask-sz Flask application."""
        
        import pytest
        from app import app
        
        
        @pytest.fixture
        def client():
            """Create a test client for the app."""
            app.config["TESTING"] = True
            with app.test_client() as client:
                yield client
        
        
        def test_welcome_returns_json(client):
            """Test GET / returns JSON with welcome message."""
            # Arrange
            expected_message = "Welcome to hello-flask-sz"
        
            # Act
            resp = client.get("/")
        
            # Assert
            assert resp.status_code == 200
            assert resp.content_type == "application/json"
            assert resp.get_json()["message"] == expected_message
        
        
        def test_welcome_returns_200(client):
            """Test GET / returns status code 200."""
            # Act
            resp = client.get("/")
        
            # Assert
            assert resp.status_code == 200
        
        
        def test_health_returns_ok(client):
            """Test GET /health returns status ok."""
            # Act
            resp = client.get("/health")
        
            # Assert
            assert resp.status_code == 200
            assert resp.get_json()["status"] == "ok"
        
        
        def test_health_returns_200(client):
            """Test GET /health returns status code 200."""
            # Act
            resp = client.get("/health")
        
            # Assert
            assert resp.status_code == 200
        
        
        def test_health_json_format(client):
            """Test GET /health returns proper JSON content-type."""
            # Act
            resp = client.get("/health")
        
            # Assert
            assert resp.content_type == "application/json"
            assert isinstance(resp.get_json(), dict)
        
        
        def test_welcome_json_structure(client):
            """Test GET / returns proper JSON structure."""
            # Act
            resp = client.get("/")
            data = resp.get_json()
        
            # Assert
            assert isinstance(data, dict)
            assert "message" in data
            assert isinstance(data["message"], str)
        
  • 📁 pytest_doc/
    • 📁 .claude/
      • ⚙️ settings.local.json 72B
    • 📁 src/
      • 📄 calculator.py 587B
        """一个简单的计算器模块,用于演示 pytest 的基本用法。"""
        
        
        def add(a, b):
            """返回两个数之和。"""
            return a + b
        
        
        def subtract(a, b):
            """返回 a 减 b 的结果。"""
            return a - b
        
        
        def multiply(a, b):
            """返回两个数之积。"""
            return a * b
        
        
        def divide(a, b):
            """返回 a 除以 b 的结果。除数为 0 时抛出 ValueError。"""
            if b == 0:
                raise ValueError("除数不能为 0")
            return a / b
        
        
        def is_even(n):
            """判断一个整数是否为偶数。"""
            return n % 2 == 0
        
    • 📁 tests/
      • 📄 test_calculator.py 1.4KB
        """calculator 模块的 pytest 测试用例。
        
        演示要点:
        1. 基本断言            -> test_add / test_subtract / test_multiply
        2. 异常断言            -> test_divide_by_zero
        3. 参数化(一组数据跑多次) -> test_is_even
        4. fixture(共享前置数据) -> calc_data
        """
        
        import pytest
        from src.calculator import add, divide, is_even, multiply, subtract
        
        
        # ---------- 1. 基本断言 ----------
        def test_add():
            assert add(2, 3) == 5
            assert add(-1, 1) == 0
        
        
        def test_subtract():
            assert subtract(10, 4) == 6
        
        
        def test_multiply():
            assert multiply(3, 4) == 12
        
        
        # ---------- 2. 异常断言 ----------
        def test_divide():
            assert divide(10, 2) == 5
        
        
        def test_divide_by_zero():
            # 期望抛出 ValueError,并校验错误信息
            with pytest.raises(ValueError, match="除数不能为 0"):
                divide(1, 0)
        
        
        # ---------- 3. 参数化 ----------
        @pytest.mark.parametrize(
            "number, expected",
            [
                (2, True),
                (3, False),
                (0, True),
                (-4, True),
                (7, False),
            ],
        )
        def test_is_even(number, expected):
            assert is_even(number) == expected
        
        
        # ---------- 4. fixture ----------
        @pytest.fixture
        def calc_data():
            """提供一组共享的测试数据。"""
            return {"a": 8, "b": 2}
        
        
        def test_with_fixture(calc_data):
            assert add(calc_data["a"], calc_data["b"]) == 10
            assert divide(calc_data["a"], calc_data["b"]) == 4
        
    • 📄 README.md 3.2KB
      # pytest 快速上手示例
      
      一个最小化的 pytest 演示项目,包含一个简单的计算器模块和对应的测试用例。
      
      ## 目录结构
      
      ```
      pytest_doc/
      ├── src/
      │   └── calculator.py        # 被测代码:加减乘除 + 判断偶数
      ├── tests/
      │   └── test_calculator.py   # 测试用例:断言 / 异常 / 参数化 / fixture
      ├── pytest.ini               # pytest 配置
      ├── pyproject.toml           # ruff(lint + format)配置
      ├── requirements.txt         # 依赖
      └── README.md
      ```
      
      ## 环境准备
      
      ```bash
      # 1. (可选)创建并激活虚拟环境
      python -m venv venv
      # Windows PowerShell:
      venv\Scripts\Activate.ps1
      # Linux / macOS:
      source venv/bin/activate
      
      # 2. 安装依赖
      pip install -r requirements.txt
      ```
      
      ## 开发人员如何运行 pytest
      
      在项目根目录 `pytest_doc/` 下执行:
      
      | 目的 | 命令 |
      | --- | --- |
      | 运行全部测试 | `pytest` 或 `python -m pytest` |
      | 详细输出(显示每个用例) | `pytest -v` |
      | 只运行某个文件 | `pytest tests/test_calculator.py` |
      | 只运行某个用例 | `pytest tests/test_calculator.py::test_add` |
      | 按名称关键字筛选 | `pytest -k "divide"` |
      | 遇到第一个失败就停止 | `pytest -x` |
      | 显示 print 输出 | `pytest -s` |
      | 查看测试结果摘要 | `pytest -ra` |
      | 生成覆盖率报告(需 pytest-cov) | `pytest --cov=src` |
      
      > 推荐使用 `python -m pytest`,它会把当前目录加入 `sys.path`,避免导入路径问题。
      
      ## 预期结果
      
      ```
      ============================= test session starts =============================
      collected 11 items
      
      tests/test_calculator.py::test_add PASSED                                [  9%]
      ...
      ============================= 11 passed in 0.08s ==============================
      ```
      
      ## 代码检查(ruff:lint + 格式化)
      
      [ruff](https://docs.astral.sh/ruff/) 是一个用 Rust 写的高性能工具,**一个工具同时替代 flake8(lint)+ isort(导入排序)+ black(格式化)**。配置写在 `pyproject.toml` 里。
      
      在项目根目录 `pytest_doc/` 下执行:
      
      | 目的 | 命令 |
      | --- | --- |
      | 检查代码问题(lint) | `ruff check .` |
      | 自动修复可修复的问题 | `ruff check --fix .` |
      | 检查格式(不修改文件) | `ruff format --check .` |
      | 自动格式化代码 | `ruff format .` |
      
      > 推荐流程:先 `ruff check --fix .` 修 lint,再 `ruff format .` 统一格式,最后 `pytest` 跑测试。
      > 启用的规则集:`E/W`(PEP8) `F`(pyflakes) `I`(isort) `B`(bugbear) `UP`(pyupgrade)。
      
      ### 一键质量检查
      
      ```bash
      ruff check . && ruff format --check . && pytest
      ```
      
      CI 流水线里通常就是这一行:任意一步失败则整体失败。
      
      ## 示例中演示的 pytest 核心知识点
      
      1. **基本断言**:直接用 Python 的 `assert`,pytest 会自动给出详细的失败信息。
      2. **异常断言**:`with pytest.raises(ValueError, match="..."):` 校验函数是否按预期抛异常。
      3. **参数化** `@pytest.mark.parametrize`:一份测试逻辑跑多组输入/输出数据。
      4. **fixture** `@pytest.fixture`:为测试提供可复用的前置数据或资源。
      
    • 📄 pyproject.toml 587B
      [tool.ruff]
      # 目标 Python 版本
      target-version = "py39"
      # 每行最大长度
      line-length = 100
      # 需要检查的目录
      src = ["src", "tests"]
      
      [tool.ruff.lint]
      # 启用的规则集:
      #   E/W = pycodestyle(PEP8 风格)
      #   F   = pyflakes(未使用变量/导入等)
      #   I   = isort(import 排序)
      #   B   = flake8-bugbear(常见 bug 模式)
      #   UP  = pyupgrade(过时写法升级)
      select = ["E", "W", "F", "I", "B", "UP"]
      ignore = []
      
      [tool.ruff.format]
      # 使用双引号,风格与 black 一致
      quote-style = "double"
      indent-style = "space"
      
    • 📄 pytest.ini 462B
      [pytest]
      # 把项目根目录加入导入路径,这样测试里可以写 from src.calculator import ...
      pythonpath = .
      
      # 只在 tests 目录下收集测试
      testpaths = tests
      
      # 识别为测试文件/类/函数的命名规则(pytest 默认值,写出来方便理解)
      python_files = test_*.py
      python_classes = Test*
      python_functions = test_*
      
      # 运行时默认带上的参数:-v 详细输出,-ra 显示测试结果摘要
      addopts = -v -ra
      
    • 📄 pytest_ruff_animation.html 24.4KB
      <!DOCTYPE html>
      <html lang="zh-CN">
      <head>
      <meta charset="UTF-8">
      <meta name="viewport" content="width=device-width, initial-scale=1.0">
      <title>pytest + ruff 原理动画</title>
      <style>
        :root{
          --bg:#0d1117; --panel:#161b22; --border:#30363d;
          --txt:#c9d1d9; --muted:#8b949e;
          --green:#3fb950; --red:#f85149; --blue:#58a6ff;
          --yellow:#d29922; --purple:#bc8cff; --cyan:#39c5cf;
          --mono:'Cascadia Code','Consolas','SFMono-Regular',Menlo,monospace;
        }
        *{box-sizing:border-box;margin:0;padding:0}
        body{
          background:radial-gradient(1200px 600px at 50% -10%,#1b2330,#0d1117 60%);
          color:var(--txt);font-family:'Segoe UI','Microsoft YaHei',sans-serif;
          min-height:100vh;padding:24px;
        }
        .wrap{max-width:1080px;margin:0 auto}
        h1{font-size:24px;text-align:center;margin-bottom:4px}
        h1 .lint{color:var(--yellow)} h1 .test{color:var(--green)}
        .sub{text-align:center;color:var(--muted);margin-bottom:18px;font-size:14px}
      
        /* 顶部进度 / 控制 */
        .controls{display:flex;gap:10px;justify-content:center;align-items:center;
          margin-bottom:18px;flex-wrap:wrap}
        button{background:var(--panel);color:var(--txt);border:1px solid var(--border);
          padding:8px 16px;border-radius:8px;cursor:pointer;font-size:14px;transition:.15s}
        button:hover{border-color:var(--blue);color:#fff}
        button.primary{background:var(--blue);color:#fff;border-color:var(--blue)}
        button:disabled{opacity:.4;cursor:not-allowed}
        .dots{display:flex;gap:8px}
        .dot{width:10px;height:10px;border-radius:50%;background:var(--border);transition:.2s}
        .dot.active{background:var(--blue);transform:scale(1.3)}
        .dot.done{background:var(--green)}
      
        .stage{background:var(--panel);border:1px solid var(--border);border-radius:14px;
          padding:24px;min-height:440px;position:relative;overflow:hidden}
        .scene{display:none;animation:fade .5s ease}
        .scene.active{display:block}
        @keyframes fade{from{opacity:0;transform:translateY(10px)}to{opacity:1;transform:none}}
        .scene h2{font-size:18px;margin-bottom:14px;display:flex;align-items:center;gap:8px}
        .badge{font-size:12px;padding:2px 10px;border-radius:20px;font-weight:600}
        .badge.static{background:rgba(210,153,34,.15);color:var(--yellow);border:1px solid var(--yellow)}
        .badge.dynamic{background:rgba(63,185,80,.15);color:var(--green);border:1px solid var(--green)}
      
        /* 代码块 */
        .code{background:#0d1117;border:1px solid var(--border);border-radius:10px;
          padding:14px 16px;font-family:var(--mono);font-size:13.5px;line-height:1.75;
          white-space:pre;overflow-x:auto;position:relative}
        .code .ln{color:#484f58;user-select:none;display:inline-block;width:22px}
        .kw{color:var(--purple)} .fn{color:var(--blue)} .str{color:#a5d6ff}
        .cm{color:#8b949e} .num{color:#79c0ff}
      
        /* 扫描线 */
        .scanline{position:absolute;left:0;right:0;height:26px;
          background:linear-gradient(90deg,transparent,rgba(88,166,255,.25),transparent);
          border-top:1px solid var(--blue);border-bottom:1px solid var(--blue);
          top:14px;opacity:0;pointer-events:none}
        .scanline.run{animation:scan 2.4s linear forwards;opacity:1}
        @keyframes scan{from{top:14px}to{top:260px}}
      
        .flow{display:flex;align-items:center;justify-content:center;gap:6px;
          flex-wrap:wrap;margin:18px 0}
        .node{background:#0d1117;border:1px solid var(--border);border-radius:10px;
          padding:12px 14px;text-align:center;min-width:104px;opacity:.3;transition:.4s;font-size:13px}
        .node.on{opacity:1;border-color:var(--blue);box-shadow:0 0 14px rgba(88,166,255,.3)}
        .node .ic{font-size:24px;display:block;margin-bottom:4px}
        .node small{color:var(--muted);display:block;font-size:11px;margin-top:2px}
        .arrow{color:var(--muted);font-size:20px;opacity:.3;transition:.4s}
        .arrow.on{opacity:1;color:var(--blue)}
      
        .two{display:grid;grid-template-columns:1fr 1fr;gap:16px}
        @media(max-width:760px){.two{grid-template-columns:1fr}}
        .col{background:#0d1117;border:1px solid var(--border);border-radius:10px;padding:14px}
        .col h3{font-size:14px;margin-bottom:10px}
        .col.lint h3{color:var(--yellow)} .col.test h3{color:var(--green)}
      
        .report{margin-top:14px;font-family:var(--mono);font-size:13px;min-height:24px}
        .report .err{color:var(--red)} .report .ok{color:var(--green)}
        .typed{opacity:0;transform:translateX(-6px);transition:.3s}
        .typed.show{opacity:1;transform:none}
      
        .pill{display:inline-block;background:#0d1117;border:1px solid var(--border);
          border-radius:6px;padding:2px 8px;font-family:var(--mono);font-size:12px;color:var(--cyan)}
        .note{margin-top:14px;color:var(--muted);font-size:13px;line-height:1.7;
          border-left:3px solid var(--blue);padding-left:12px}
        .key{color:var(--yellow);font-weight:600}
      
        /* test 灯 */
        .lights{display:flex;gap:8px;flex-wrap:wrap;margin-top:12px}
        .light{width:30px;height:30px;border-radius:6px;background:var(--border);
          display:flex;align-items:center;justify-content:center;font-size:14px;
          transition:.3s;color:#0d1117;font-weight:700}
        .light.pass{background:var(--green)}
        .light.fail{background:var(--red);color:#fff}
      
        /* 文件树 */
        .tree,.tree ul{list-style:none;margin:0;padding:0}
        .tree ul{margin-left:18px;border-left:1px dashed var(--border);padding-left:6px;
          max-height:400px;overflow:hidden;transition:max-height .3s ease}
        .tree li{margin:2px 0}
        .tree .row{display:inline-flex;align-items:center;gap:6px;padding:4px 8px;border-radius:6px;
          cursor:pointer;font-size:13.5px;transition:.12s;user-select:none}
        .tree .row:hover{background:rgba(88,166,255,.12)}
        .tree .row .ic{font-size:14px}
        .tree .tw{display:inline-block;width:12px;color:var(--muted);transition:transform .2s}
        .tree .dir.open>.row .tw{transform:rotate(90deg)}
        .tree .dir:not(.open)>ul{max-height:0}
        .tree .file.sel>.row{background:rgba(88,166,255,.25);color:#fff;
          box-shadow:inset 2px 0 0 var(--blue)}
        #codeView{display:block;color:var(--txt);font-family:var(--mono);font-size:12.5px;
          line-height:1.7;white-space:pre}
      
        .barwrap{background:#0d1117;border:1px solid var(--border);border-radius:8px;
          height:30px;overflow:hidden;margin-top:8px}
        .bar{height:100%;width:0;background:linear-gradient(90deg,var(--green),var(--cyan));
          transition:width .8s ease;display:flex;align-items:center;justify-content:flex-end;
          padding-right:8px;color:#0d1117;font-weight:700;font-size:12px}
        footer{text-align:center;color:var(--muted);font-size:12px;margin-top:14px}
      </style>
      </head>
      <body>
      <div class="wrap">
        <h1><span class="lint">ruff (lint)</span> 与 <span class="test">pytest</span> 工作原理动画</h1>
        <p class="sub">静态分析「看代码」 vs 动态测试「跑代码」——两者都是外部工具,都不会把自己的逻辑写进你的业务代码</p>
      
        <div class="controls">
          <button id="prev">← 上一步</button>
          <button id="play" class="primary">▶ 自动播放</button>
          <button id="next" class="primary">下一步 →</button>
          <div class="dots" id="dots"></div>
        </div>
      
        <div class="stage">
      
          <!-- 场景 0: 总览 -->
          <section class="scene active" data-scene="0">
            <h2>① 两条完全不同的路线</h2>
            <div class="two">
              <div class="col lint">
                <h3>🔍 ruff / lint(静态分析)</h3>
                <div class="code"><span class="cm"># 把源码当“文本”读入</span>
      <span class="cm"># 解析成 AST(语法树)</span>
      <span class="cm"># 用规则匹配,输出报告</span>
      <span class="key" style="color:var(--yellow)"># ❌ 不运行你的程序</span></div>
                <div class="note">关注:风格、未使用变量、import 排序、潜在 bug 写法……</div>
              </div>
              <div class="col test">
                <h3>🧪 pytest(动态测试)</h3>
                <div class="code"><span class="cm"># 收集 test_* 函数</span>
      <span class="cm"># 真正调用你的函数</span>
      <span class="cm"># 对 assert 求值 → 红/绿</span>
      <span class="key" style="color:var(--green)"># ✅ 真的运行你的程序</span></div>
                <div class="note">关注:逻辑对不对、结果是否符合预期、异常是否抛出……</div>
              </div>
            </div>
            <div class="note">👉 二者都是<span class="key">独立的工具 + 独立的文件</span>,运行时不会混入被测业务代码。</div>
          </section>
      
          <!-- 场景 1: 项目结构(可点击展开) -->
          <section class="scene" data-scene="1">
            <h2>② 项目结构 <span class="badge static">点击文件夹展开 · 点文件看代码</span></h2>
            <div class="two" style="grid-template-columns:0.9fr 1.1fr">
              <div class="col">
                <h3 style="color:var(--cyan)">📦 pytest_doc/ <small style="color:var(--muted);font-weight:400">(点击节点展开/折叠)</small></h3>
                <ul class="tree" id="tree">
                  <li class="dir open" data-name="src">
                    <span class="row"><span class="tw">▸</span><span class="ic">📁</span> src/</span>
                    <ul>
                      <li class="file" data-file="calculator.py"><span class="row"><span class="tw" style="visibility:hidden">▸</span><span class="ic">🐍</span> calculator.py</span></li>
                    </ul>
                  </li>
                  <li class="dir open" data-name="tests">
                    <span class="row"><span class="tw">▸</span><span class="ic">📁</span> tests/</span>
                    <ul>
                      <li class="file" data-file="test_calculator.py"><span class="row"><span class="tw" style="visibility:hidden">▸</span><span class="ic">🧪</span> test_calculator.py</span></li>
                    </ul>
                  </li>
                  <li class="file" data-file="pytest.ini"><span class="row"><span class="tw" style="visibility:hidden">▸</span><span class="ic">⚙️</span> pytest.ini</span></li>
                  <li class="file" data-file="pyproject.toml"><span class="row"><span class="tw" style="visibility:hidden">▸</span><span class="ic">🦀</span> pyproject.toml</span></li>
                  <li class="file" data-file="requirements.txt"><span class="row"><span class="tw" style="visibility:hidden">▸</span><span class="ic">📋</span> requirements.txt</span></li>
                  <li class="file" data-file="README.md"><span class="row"><span class="tw" style="visibility:hidden">▸</span><span class="ic">📖</span> README.md</span></li>
                </ul>
                <div class="note" style="font-size:12px">最后一层(文件)就是代码细节,点击右侧查看 →</div>
              </div>
              <div class="col">
                <h3 id="codeTitle" style="color:var(--blue)">👉 点击左侧任意文件查看代码</h3>
                <div class="code" style="min-height:300px;max-height:340px;overflow:auto"><code id="codeView">// 选择一个文件…</code></div>
              </div>
            </div>
          </section>
      
          <!-- 场景 2: ruff 原理 -->
          <section class="scene" data-scene="2">
            <h2>③ ruff 原理 <span class="badge static">静态分析 · 不执行</span></h2>
            <div class="flow" id="lintflow">
              <div class="node" data-i="0"><span class="ic">📄</span>源码文本<small>calculator.py</small></div>
              <span class="arrow" data-i="0">→</span>
              <div class="node" data-i="1"><span class="ic">🔡</span>词法分析<small>tokens</small></div>
              <span class="arrow" data-i="1">→</span>
              <div class="node" data-i="2"><span class="ic">🌳</span>语法树<small>AST</small></div>
              <span class="arrow" data-i="2">→</span>
              <div class="node" data-i="3"><span class="ic">📐</span>规则匹配<small>E/F/I/B/UP</small></div>
              <span class="arrow" data-i="3">→</span>
              <div class="node" data-i="4"><span class="ic">📋</span>报告<small>不改逻辑</small></div>
            </div>
            <div class="code" style="margin-top:6px">
      <div class="scanline" id="scan1"></div><span class="ln">10</span><span class="kw">import</span> pytest
      <span class="ln">11</span><span class="kw">from</span> src.calculator <span class="kw">import</span> add, subtract, multiply, divide, is_even
      <span class="ln">12</span>
      <span class="ln">13</span><span class="kw">def</span> <span class="fn">test_add</span>():
      <span class="ln">14</span>    <span class="kw">assert</span> <span class="fn">add</span>(<span class="num">2</span>, <span class="num">3</span>) == <span class="num">5</span>
      <span class="ln">15</span>    unused = <span class="num">42</span>   <span class="cm"># 没用到的变量</span></div>
            <div class="report" id="lintreport"></div>
            <div class="note">ruff 全程只是“阅读 + 比对规则”,<span class="key">从不执行</span> <span class="pill">add()</span>。所以它不会把检查逻辑写进代码。</div>
          </section>
      
          <!-- 场景 3: ruff --fix -->
          <section class="scene" data-scene="3">
            <h2>④ <span class="pill">ruff check --fix</span> 会改写「你的源码格式」</h2>
            <div class="two">
              <div class="col">
                <h3 style="color:var(--red)">修复前</h3>
                <div class="code"><span class="kw">from</span> src.calculator <span class="kw">import</span> add, subtract, multiply, divide, is_even</div>
                <div class="report"><span class="err">I001 Import block is un-sorted</span></div>
              </div>
              <div class="col">
                <h3 style="color:var(--green)">修复后(自动排序)</h3>
                <div class="code"><span class="kw">from</span> src.calculator <span class="kw">import</span> <span id="sorted"></span></div>
                <div class="report"><span class="ok">All checks passed! ✅</span></div>
              </div>
            </div>
            <div class="note">
              ⚠️ 注意区别:它修改的是<span class="key">你自己的代码格式</span>(排序/空格/引号),
              而<span class="key">不是</span>“把 lint 规则塞进你的代码”。两回事。
              <br>同理 <span class="pill">ruff format .</span> 只负责统一排版,不改变程序行为。
            </div>
          </section>
      
          <!-- 场景 4: pytest 原理 -->
          <section class="scene" data-scene="4">
            <h2>⑤ pytest 原理 <span class="badge dynamic">动态执行 · 真跑代码</span></h2>
            <div class="flow" id="testflow">
              <div class="node" data-i="0"><span class="ic">🔎</span>收集<small>test_*</small></div>
              <span class="arrow" data-i="0">→</span>
              <div class="node" data-i="1"><span class="ic">⚙️</span>调用函数<small>真的运行</small></div>
              <span class="arrow" data-i="1">→</span>
              <div class="node" data-i="2"><span class="ic">⚖️</span>assert 求值<small>比较结果</small></div>
              <span class="arrow" data-i="2">→</span>
              <div class="node" data-i="3"><span class="ic">🚦</span>红 / 绿<small>pass/fail</small></div>
            </div>
            <div class="code" style="margin-top:6px"><span class="kw">def</span> <span class="fn">test_divide_by_zero</span>():
          <span class="kw">with</span> pytest.<span class="fn">raises</span>(ValueError):
              <span class="fn">divide</span>(<span class="num">1</span>, <span class="num">0</span>)   <span class="cm"># ← pytest 真的调用它,看是否抛异常</span></div>
            <div class="lights" id="lights"></div>
            <div class="report" id="testreport"></div>
            <div class="note">pytest <span class="key">真实执行</span>了 <span class="pill">divide(1,0)</span>,捕获到 <span class="pill">ValueError</span> → 这条用例变绿。</div>
          </section>
      
          <!-- 场景 5: 完整流程 -->
          <section class="scene" data-scene="5">
            <h2>⑥ 完整流程(开发 / CI 质量门禁)</h2>
            <div class="flow" id="ciflow">
              <div class="node" data-i="0"><span class="ic">✍️</span>写代码</div>
              <span class="arrow" data-i="0">→</span>
              <div class="node" data-i="1"><span class="ic">🔧</span>ruff --fix<small>修 lint</small></div>
              <span class="arrow" data-i="1">→</span>
              <div class="node" data-i="2"><span class="ic">🎨</span>ruff format<small>统一排版</small></div>
              <span class="arrow" data-i="2">→</span>
              <div class="node" data-i="3"><span class="ic">🧪</span>pytest<small>跑测试</small></div>
              <span class="arrow" data-i="3">→</span>
              <div class="node" data-i="4"><span class="ic">🚀</span>合并/部署</div>
            </div>
            <div class="code" style="margin-top:6px"><span class="cm"># 一行搞定质量门禁:任意一步失败则整体失败</span>
      ruff check . <span class="kw">&&</span> ruff format --check . <span class="kw">&&</span> pytest</div>
            <div class="barwrap"><div class="bar" id="cibar">0%</div></div>
            <div class="report" id="cireport"></div>
            <div class="note">
              <span class="key">lint 看代码、pytest 跑代码</span>,二者互补:
              前者保证“写得规范、没有低级错误”,后者保证“逻辑真的对”。
              都不会污染你的业务代码。
            </div>
          </section>
      
        </div>
        <footer>用 ← → 或按钮切换步骤 · 可点击「自动播放」· 本文件可离线打开</footer>
      </div>
      
      <script>
      const scenes=[...document.querySelectorAll('.scene')];
      const dotsBox=document.getElementById('dots');
      let cur=0, autoTimer=null;
      
      scenes.forEach((_,i)=>{
        const d=document.createElement('div');d.className='dot'+(i===0?' active':'');
        d.onclick=()=>go(i);dotsBox.appendChild(d);
      });
      const dots=[...document.querySelectorAll('.dot')];
      
      function go(n){
        if(n<0||n>=scenes.length)return;
        scenes[cur].classList.remove('active');
        cur=n;
        scenes[cur].classList.add('active');
        dots.forEach((d,i)=>{
          d.classList.toggle('active',i===cur);
          d.classList.toggle('done',i<cur);
        });
        document.getElementById('prev').disabled=(cur===0);
        document.getElementById('next').disabled=(cur===scenes.length-1);
        runScene(cur);
      }
      
      document.getElementById('next').onclick=()=>{stopAuto();go(cur+1)};
      document.getElementById('prev').onclick=()=>{stopAuto();go(cur-1)};
      document.addEventListener('keydown',e=>{
        if(e.key==='ArrowRight'){stopAuto();go(cur+1)}
        if(e.key==='ArrowLeft'){stopAuto();go(cur-1)}
      });
      
      const playBtn=document.getElementById('play');
      playBtn.onclick=()=>{ autoTimer?stopAuto():startAuto(); };
      function startAuto(){
        playBtn.textContent='⏸ 暂停';
        autoTimer=setInterval(()=>{
          if(cur>=scenes.length-1){go(0);} else {go(cur+1);}
        },3800);
      }
      function stopAuto(){ if(autoTimer){clearInterval(autoTimer);autoTimer=null;playBtn.textContent='▶ 自动播放';} }
      
      function seqOn(nodes,arrows,delay=550){
        nodes.forEach(n=>n.classList.remove('on'));
        arrows.forEach(a=>a.classList.remove('on'));
        let i=0;
        const step=()=>{
          if(i<nodes.length){nodes[i].classList.add('on'); if(arrows[i])arrows[i].classList.add('on'); i++; setTimeout(step,delay);}
        };
        step();
      }
      
      function runScene(n){
        if(n===2){
          const nodes=[...document.querySelectorAll('#lintflow .node')];
          const arrows=[...document.querySelectorAll('#lintflow .arrow')];
          seqOn(nodes,arrows);
          const scan=document.getElementById('scan1');
          scan.classList.remove('run');void scan.offsetWidth;
          setTimeout(()=>scan.classList.add('run'),900);
          const rep=document.getElementById('lintreport');
          rep.innerHTML='';
          setTimeout(()=>{
            rep.innerHTML='<span class="typed err">I001  import 未排序  ·  F841  unused 变量未使用</span>';
            requestAnimationFrame(()=>rep.querySelector('.typed').classList.add('show'));
          },3400);
        }
        if(n===3){
          const s=document.getElementById('sorted');s.textContent='';
          const full='add, divide, is_even, multiply, subtract';
          let i=0;
          (function type(){ if(i<=full.length){s.textContent=full.slice(0,i++);setTimeout(type,40);} })();
        }
        if(n===4){
          const nodes=[...document.querySelectorAll('#testflow .node')];
          const arrows=[...document.querySelectorAll('#testflow .arrow')];
          seqOn(nodes,arrows,600);
          const box=document.getElementById('lights');box.innerHTML='';
          const total=11;
          for(let k=0;k<total;k++){const l=document.createElement('div');l.className='light';box.appendChild(l);}
          const lights=[...box.children];
          const rep=document.getElementById('testreport');rep.innerHTML='';
          setTimeout(()=>{
            lights.forEach((l,k)=>setTimeout(()=>{l.classList.add('pass');l.textContent='✓';},200*k));
            setTimeout(()=>{
              rep.innerHTML='<span class="typed ok">=== 11 passed in 0.04s === ✅</span>';
              requestAnimationFrame(()=>rep.querySelector('.typed').classList.add('show'));
            },200*total+300);
          },1600);
        }
        if(n===5){
          const nodes=[...document.querySelectorAll('#ciflow .node')];
          const arrows=[...document.querySelectorAll('#ciflow .arrow')];
          seqOn(nodes,arrows,500);
          const bar=document.getElementById('cibar');bar.style.width='0';bar.textContent='0%';
          const rep=document.getElementById('cireport');rep.innerHTML='';
          const steps=[['ruff check',33],['ruff format',66],['pytest',100]];
          steps.forEach((st,k)=>setTimeout(()=>{bar.style.width=st[1]+'%';bar.textContent=st[0]+' ✓';},900+k*900));
          setTimeout(()=>{
            rep.innerHTML='<span class="typed ok">质量门禁全部通过 → 可以合并 🚀</span>';
            requestAnimationFrame(()=>rep.querySelector('.typed').classList.add('show'));
          },900+steps.length*900);
        }
      }
      /* ===== 文件树交互 ===== */
      const FILES = {
      "calculator.py":
      `"""一个简单的计算器模块,用于演示 pytest 的基本用法。"""
      
      
      def add(a, b):
          """返回两个数之和。"""
          return a + b
      
      
      def subtract(a, b):
          """返回 a 减 b 的结果。"""
          return a - b
      
      
      def multiply(a, b):
          """返回两个数之积。"""
          return a * b
      
      
      def divide(a, b):
          """除数为 0 时抛出 ValueError。"""
          if b == 0:
              raise ValueError("除数不能为 0")
          return a / b
      
      
      def is_even(n):
          """判断一个整数是否为偶数。"""
          return n % 2 == 0`,
      
      "test_calculator.py":
      `import pytest
      from src.calculator import add, divide, is_even, multiply, subtract
      
      
      # 1) 基本断言
      def test_add():
          assert add(2, 3) == 5
      
      
      # 2) 异常断言
      def test_divide_by_zero():
          with pytest.raises(ValueError, match="除数不能为 0"):
              divide(1, 0)
      
      
      # 3) 参数化:一份逻辑跑多组数据
      @pytest.mark.parametrize("number, expected", [
          (2, True), (3, False), (0, True), (-4, True), (7, False),
      ])
      def test_is_even(number, expected):
          assert is_even(number) == expected
      
      
      # 4) fixture:提供共享数据
      @pytest.fixture
      def calc_data():
          return {"a": 8, "b": 2}
      
      
      def test_with_fixture(calc_data):
          assert add(calc_data["a"], calc_data["b"]) == 10
          assert divide(calc_data["a"], calc_data["b"]) == 4`,
      
      "pytest.ini":
      `[pytest]
      pythonpath = .          ; 把根目录加入导入路径
      testpaths = tests       ; 只在 tests 下收集用例
      python_files = test_*.py
      python_functions = test_*
      addopts = -v -ra        ; 默认参数:详细输出 + 结果摘要`,
      
      "pyproject.toml":
      `[tool.ruff]
      target-version = "py39"
      line-length = 100
      src = ["src", "tests"]
      
      [tool.ruff.lint]
      # E/W=PEP8  F=pyflakes  I=isort  B=bugbear  UP=pyupgrade
      select = ["E", "W", "F", "I", "B", "UP"]
      
      [tool.ruff.format]
      quote-style = "double"
      indent-style = "space"`,
      
      "requirements.txt":
      `pytest>=8.0
      ruff>=0.8`,
      
      "README.md":
      `# pytest + ruff 快速上手示例
      
      运行测试:
          python -m pytest          运行全部
          pytest -v                 详细输出
          pytest -k "divide"        按关键字筛选
          pytest -x                 首个失败即停止
      
      代码检查(ruff):
          ruff check .              lint 检查
          ruff check --fix .        自动修复
          ruff format .             自动格式化
      
      一键质量门禁(CI 常用):
          ruff check . && ruff format --check . && pytest`
      };
      
      const codeView=document.getElementById('codeView');
      const codeTitle=document.getElementById('codeTitle');
      const tree=document.getElementById('tree');
      
      tree.addEventListener('click',e=>{
        const row=e.target.closest('.row'); if(!row)return;
        const li=row.parentElement;
        if(li.classList.contains('dir')){
          li.classList.toggle('open');
          return;
        }
        if(li.classList.contains('file')){
          tree.querySelectorAll('.file.sel').forEach(f=>f.classList.remove('sel'));
          li.classList.add('sel');
          const name=li.dataset.file;
          codeTitle.textContent='📄 '+name;
          codeView.textContent=FILES[name]||'(无内容)';
        }
      });
      
      go(0);
      </script>
      </body>
      </html>
      
    • 📄 requirements.txt 24B
      pytest>=8.0
      ruff>=0.8
      
📎 资料

银行营销数据

  • 📎 bank_marketing_report.xlsx15.2KB
  • 📎 test.csv901.4KB
  • 📎 train.csv2.7MB

ℹ️ 离线版不含附件原文件,需要请到在线网站对应课时下载。

第2节 · 计算机视觉:图像分类

🎯 掌握 CV 基础、AlexNet/VGG/GoogLeNet/ResNet、图像增强与车型分类案例。
📖 笔记

01-计算机视觉基础知识

计算机视觉基础CV笔记

一、今日课程内容

  • 计算机视觉的定义(了解)
  • 计算机视觉常见任务(了解)
  • 计算机视觉的应用场景(了解)

二、计算机视觉介绍(了解)

1、定义

  • 用【摄像机和电脑】及其他相关设备,对【生物视觉】的一种模拟,任务:理解图片和视频

2、主要任务

  • 让【计算机】理解【图片或者视频】中的内容,就像人类和许多其他生物每天所做的那样。

3、任务目标

  • 让计算机【理解】图片中的场景(办公室,客厅,咖啡厅等)
  • 让计算机【识别】场景中包含的物体(宠物,交通工具,人等)
  • 让计算机【定位】物体在图像中的位置(物体的大小,边界等)
  • 让计算机理解物体之间的【关系或行为】(是在对话,比赛或吵架等),以及图像表达的意义(喜庆的,悲伤的等)

三、计算机视觉常见任务(了解)

1、三大经典任务

image-20240423102643136

  • 图像分类(Classification)

  • 将图像结构化为某一【类别】的信息,用事先确定好的【类别】(category)来描述图片。

  • 分类任务关心【整体】,给出的是整张图片的【内容描述】。

  • 目标检测(Detection)

  • 将图像中的某一目标的【类别】信息和【位置】信息获取到

  • 关注【特定】的物体目标的类别信息和位置信息(classification + localization)。

  • 图像分割(Segmentation)

  • 赋予每个像素类别(实例)意义,适用于理解要求较高的场景,如无人驾驶中对道路和非道路的分割。

  • 是对图像的【像素级】描述

四、计算机视觉的应用场景(了解)

1、基本内容

  • 涉及的领域复杂,具有广泛的实际应用范围,从电子商务行业到更经典的各种类型和规模的公司都可以利用其强大的功能
  • 代表企业:
  • “CV四小天鹅” “AI四小龙”(商汤科技、旷视科技、云从科技、依图科技)
  • 人脸识别:旷视科技、依图科技、商汤科技、深醒科技、云从科技等。
  • 视频监控:商汤科技、DeepGlint 格灵深瞳、依图科技、云天励飞、深网视界
  • 辅助驾驶(纵目科技、TuSimple 图森科技、驭势科技、MINIEYE 佑驾创新、中天安驰)

2、发展史

  • CV领域的第一篇专业论文:1963年,Larry Roberts发表,用以对简单几何体进行边缘提取和三维重建。
  • 计算机视觉作为一个科学领域的正式诞生:1966年,麻省理工学院(MIT)发起了一个夏季项目,目标是搭建一个机器视觉系统,完成模式识别(pattern recognition)等工作。
  • 计算机视觉成为了一门独立学科:1982年,学者David Marr发表的著作《Vision》从严谨又长远的角度给出了CV的发展方向和一些基本算法,其中不乏现在为人熟知的“图层”的概念、边缘提取、三维重建等

  • SIFT(尺度不变特征变换):1999年David Lowe提出了尺度不变特征变换(SIFT, Scale-invariant feature transform)目标检测算法,用于匹配不同拍摄方向、纵深、光线等图片中的相同元素。

  • HOG(方向梯度直方图Histogram of Oriented Gradient):2009年,由Felzenszwalb教授在提出基于HOG的deformable parts model,可变形零件模型开发,它是深度学习之前最好的最成功的objectdetection & recognition算法。

  • PASCAL:2006年至2012年间,Everingham等人搭建了一个大型图片数据库,供机器识别和训练,称为PASCAL Visual Object Challenge,该数据库中有20种类别的图片,每种图片数量在一千至一万张不等。

  • ImageNet:2009年,李飞飞教授等在CVPR2009上发表了一篇名为《ImageNet: A Large-Scale Hierarchical Image Database》的论文,发布了ImageNet数据集,这是为了检测计算机视觉能否识别自然万物,回归机器学习,克服过拟合问题。

  • AlexNet:2012 年,Alex Krizhevsky、Ilya Sutskever 和 Geoffrey Hinton 创造了一个“大型的深度卷积神经网络”,也即现在众所周知的 AlexNet,赢得了当年的 ILSVRC。这是史上第一次有模型在 ImageNet 数据集表现如此出色。自那时起,CNN 才成了家喻户晓的名字。

  • 生成对抗网络(GAN):2014年,蒙特利尔大学提出生成对抗网络(GAN):拥有两个相互竞争的神经网络可以使机器学习得更快。一个网络尝试模仿真实数据生成假的数据,而另一个网络则试图将假数据区分出来。随着时间的推移,两个网络都会得到训练,生成对抗网络(GAN)被认为是计算机视觉领域的重大突破。

  • 视频到视频生成(Video-to-Video synthesis):2018年末,英伟达发布的视频到视频生成(Video-to-Video synthesis),它通过精心设计的发生器、鉴别器网络以及时空对抗物镜,合成高分辨率、照片级真实、时间一致的视频,实现了让AI更具物理意识,更强大,并能够推广到新的和看不见的更多场景。

  • BigGAN:2019年,更强大的GAN,BigGAN,是拥有了更聪明的学习技巧的GAN,由它训练生成的图像连它自己都分辨不出真假,因为除非拿显微镜看,否则将无法判断该图像是否有任何问题,因而,它更被誉为史上最强的图像生成器.

📖 笔记

02-图像分类-图像分类简介

计算机视觉基础CV笔记

一、今日课程内容

  • 图像分类的目的(了解)
  • CIFAR数据集(了解)
  • ImageNet数据集(了解)
  • 鲜花分类数据集(了解)

二、图像分类的目的(了解)

1、定义

  • 从给定的【类别集合】中为图像分配对应【标签】的任务
  • 分析一个【输入图像】并返回一个该图像类别的【标签】

2、主要任务

  • 模型给图像分配多个标签,每个标签的【概率值】不同
  • 选取【概率值】最大的类别作为标签

三、常用数据集(了解)

1、CIFAR-10和CIFAR-100

  • CIFAR-10

  • 5w+张训练图像、1w+张测试图像、10个类别、每个类别有6k个图像

  • 图像大小:【32x32】×3

  • CIFAR-100

  • 5w+张训练图像、1w+张测试图像、包含100个类别

  • 图像大小:【32x32】×3

  • 数据加载(torchvision)

  • 导入依赖包:【import torchvision】

  • 调用接口函数:【torchvision.datasets.CIFAR10(root=root,train=True,download=True)】
    • root:【数据集路径】
    • train:【True/False 是否是训练集】
    • download:【True/False,是否下载数据】

2、ImageNet

  • 【ILSVRC】(ImageNet Large-Scale Visual Recognition Challenge)使用的数据集

  • 斯坦福大学李飞飞教授主导

  • 包含超过1400万张全尺寸的有标记图片,大约有【22000】个类别的数据。比赛数据集选用其中【1000】类

  • 举办时间:2010年~2017年

image-20240423115550546

  • ILSVRC历届冠军
  • 2012年:冠军是AlexNet,由于准确率远超传统方法的第二名(top5错误率为15.3%,第二名为26.2%),引起了很大的轰动。自此之后,CNN成为在图像识别分类的核心算法模型,带来了深度学习的大爆发。
  • 2013年:冠军是ZFNet,结构和AlexNet区别不大,分类效果也差不多。
  • 2014年:亚军是VGG网络,网络结构十分简单,因此至今VGG-16仍在广泛使用。
  • 2014年:冠军网络是GoogLeNet ,核心模块是Inception Module。Inception历经了V1、V2、V3、V4等多个版本的发展,不断趋于完善。GoogLeNet取名中L大写是为了向LeNet致敬,而Inception的名字来源于盗梦空间中的"we need to go deeper"梗。
  • 2015年:冠军网络是ResNet。核心是带短连接的残差模块,其中主路径有两层卷积核(Res34),短连接把模块的输入信息直接和经过两次卷积之后的信息融合,相当于加了一个恒等变换。短连接是深度学习又一重要思想,除计算机视觉外,短连接思想也被用到了机器翻译、语音识别/合成领域
  • 2017年:冠军SENet是一个模块,可以和其他的网络架构结合,比如GoogLeNet、ResNet等。

3、鲜花分类数据集

  • 3670 张图像,其中训练图像和测试图像分别为 3306 张和 364张

  • 5 种类型:菊花,蒲公英,玫瑰,向日葵、郁金香

  • 数据集解析:

  • 导入依赖包:from torchvision.datasets import ImageFolder

  • 获取数据:flowers_train = ImageFolder(root, transform=dataset_transform)

    • root:数据路径

    • transform:数据转换处理,例如类型转换、尺寸调整

    python # 先将数据转换为tensor类型,并调整数据的大小为224x224 dataset_transform = torchvision.transforms.Compose([ torchvision.transforms.ToTensor(), torchvision.transforms.Resize((224,224))

  • 获取数据集样本数量

    python len(flowers_train.imgs) len(flowers_test.imgs)

  • 展示数据

    python import matplotlib.pyplot as plt plt.imshow(flowers_train.__getitem__(3000)[0].permute(1,2,0)) plt.show()

📖 笔记

03-图像分类-AlexNet

计算机视觉基础CV笔记

一、今日课程内容

  • AlexNet网络背景及模型地位(了解)
  • AlexNet网络结构(理解)
  • 利用AlexNet完成图像分类(操作+掌握)

二、AlexNet模型背景及模型地位(了解)

1、AlexNet模型背景

(1)发表时间:【2012】

(2)模型名称来源:【论文第一作者的姓名Alex Krizhevsky】

(3)模型荣誉:【赢得了ImageNet 2012图像识别挑战赛】

2、AlexNet模型地位

(1)首次证明:【学习到的特征可以超越手工设计的特征】

(2)时代里程碑:【一举打破计算机视觉研究的方向,计算机视觉进入CNN时代】

三、AlexNet网络结构(理解)

1、网络结构

image-20240503155039843

2、网络特点

(1)网络层数:【8】

(2)卷积层和池化层参数:【卷积核:11,5,3 池化:3X3 ,步长2】

(3)激活函数:【ReLU】

(4)模型复杂度调节:【Dropout】

(5)缓解过拟合操作:【引入了大量的图像增强,如翻转、裁剪和颜色变化,扩大数据集】

3、网络层参数计算

(1)计算公式

image-20240503155253184

(2)AlexNet网络层计算

image-20240309105724554

4、Pytorch构建AlexNet网络流程

(1)网络构建基本方法

  • 一个继承:【继承自类nn.Module】

  • 两个方法:【__init__和forward方法】

(2)网络构建代码架构

# 1 导入工具包
import torch
from torch import nn

# 2 定义网络类
class AlexNet(nn.Module):
    # 3 构建网络层
    def __init__(self,in_dim,n_class):
        super().__init__()
        # 卷积部分
        self.conv = nn.Sequential(【代码块】)    
        # 全连接部分
        self.fc = nn.Sequential(【代码块】)
    # 4 构建前向传播
    def forward(self, x):
        【代码块】
        return out

(3)完整代码

# AlexNet 模型构建
# 1 导入工具包
import torch
from torch import nn

# 2 定义网络类
class AlexNet(nn.Module):
    # 3 构建网络层
    def __init__(self,in_dim,n_class):
        super().__init__()
        # 卷积部分
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=in_dim,out_channels=96,kernel_size=11,stride=4,padding=0),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3,stride=2),
            nn.Conv2d(96,256,5,1,2),
            nn.ReLU(True),
            nn.MaxPool2d(3,2),
            nn.Conv2d(256,384,3,1,1),
            nn.ReLU(True),
            nn.Conv2d(384, 384, 3, 1, 1),
            nn.ReLU(True),
            nn.Conv2d(384, 256, 3, 1, 1),
            nn.ReLU(True),
            nn.MaxPool2d(3,2)
        )
        # 全连接部分
        self.fc = nn.Sequential(
            nn.Linear(9216,4096),
            nn.ReLU(True),
            nn.Dropout(0.5),
            nn.Linear(4096,4096),
            nn.ReLU(True),
            nn.Dropout(0.5),
            nn.Linear(4096,n_class)

        )
    # 4 构建前向传播
    def forward(self, x):
        # 输入数据送入卷积部分进行处理
        x = self.conv(x)
        #将特征图抻成一维向量
        x = x.view(x.size(0), -1)  
        # 全连接层获取输出结果
        output = self.fc(x)
        return output

(4)查看网络架构(代码)

调用模块介绍
from torchsummary import summary
# 【参数解读】
summary(model, input_size, batch_size, device)
代码实操
import torch
from torchsummary import summary

# 模型实例化
net = Alexnet(3,1000)
# 查看网络的形状
summary(model=net,input_size=(3,227,227),batch_size=1,device='cpu')

四、利用AlexNet完成图像分类(操作+掌握)

1、基本流程

【数据读取】—>【模型实例化和参数设置】—>【定义模型评估方法】—>【模型训练】

2、读取数据

(1)代码架构

# 数据获取
# 1 指定数据集路径
# 2 定义数据处理方法(resize)
# 3 ImageFolder读取数据
# 4 DataLoader获取batch数据

(2)完整代码

from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
# 指定批次大小
batch_size = 2
# 指定数据集路径
flower_train_path = './dataset/flower_datas/train/'
flower_test_path = './dataset/flower_datas/val/'
# 先将数据转换为tensor类型,并调整数据的大小为227x227
dataset_transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Resize((227, 227))])
# 获取训练集数据和测试集数据
flower_train = ImageFolder(flower_train_path, transform=dataset_transform)
flower_test = ImageFolder(flower_test_path, transform=dataset_transform)
# 获取数据的迭代
train_loader = DataLoader(dataset=flower_train,
                          batch_size=batch_size,
                          shuffle=True)
test_loader = DataLoader(dataset=flower_ttest
                         batch_size=batch_size,
                         shuffle=False)

3、模型实例化和参数设置

(1)代码架构

# 模型实例化
# 1 模型实例化

# 参数设置
# 1 学习率lr设置
# 2 轮次设置
# 3 损失函数设置
# 4 优化器设置

(2)完整代码

# 模型实例化:输入数据3通道,进行5类的分类处理
model = Alexnet(3, 5)
# 模型训练的参数设置
# 学习率
learning_rate = 1e-3
# 训练轮数
num_epochs = 10
# 优化算法Adam = RMSProp + Momentum
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# 交叉熵损失函数
loss_fn = torch.nn.CrossEntropyLoss()

4、定义模型评估方法

(1)代码架构

# 评估方法:准确率
# 1 定义模型准确率评估方法
def eval_acc(test_data,model):
    # 2 设置参数初始值(预测正确个数corr,预测总个数total)
    # 3 遍历数据,进行模型预测
    with torch.no_grad():
        model.eval()  # 将模型切换到评估模式
        【模型预测和准确率计算】    
    return corr/total

(2)完整代码

# 计算模型预测精度:测试集数据,模型
def evaluate_accuracy(data_iter, model):
    total = 0
    correct = 0
    # 不进行梯度计算
    with torch.no_grad():
        # 模型是验证模式
        model.eval()
        # 获取每一个batch的数据,进行预测
        for images, labels in data_iter:
            outputs = model(images)
            # 获取预测结果
            _, predicts = torch.max(outputs.data, dim=1)
            # 预测的次数
            total += labels.size(0)
            # 预测正确的个数
            correct += (predicts == labels).cpu().sum()
    # 获取准确率
    return correct / total

5、模型训练

(1)代码架构

# 模型训练
# 1 定义train方法
def train(data_loader = flower_train_loader,model = model,optim = opti,loss_fn = loss):
    # 2 遍历每个轮次
    for epoch in range(num_epoch):
        # 3 定义参数初始值
        # 4 遍历每个batch
        for i,(image,targtes) in enumerate(data_loader):
            ## 将模型切换到训练模式
            model.train()  
            # (1)模型前向传播计算
            # (2)损失计算
            # (3)优化器梯度归零
            # (4)反向传播
            # (5)参数更新

            ## 计算准确率,loss的变化
        # 5 评估模型效果

(2)完整代码

# 定义模型训练过程:指定数据集,优化器,损失函数和轮次
def train(data_loader=train_loader,
          optimizer=optimizer,
          loss_fn=loss_fn,
          epochs=num_epochs):
    # 遍历每一个轮次进行训练
    for epoch in range(epochs):
        print('current epoch = {}'.format(epoch))
        # 每一个轮次的损失,预测个数和预测正确个数的初始化
        train_accuracy_total = 0
        train_correct = 0
        # 损失和的初始化
        train_loss_sum = 0
        # 迭代次数
        iter = 0
        for i, (images, labels) in enumerate(data_loader):
            # 模型定义为训练模式
            model.train()
            # 对数据进行预测
            outputs = model(images)
            # 计算模型的损失
            loss = loss_fn(outputs, labels)  
            # 在做反向传播前先清除网络状态
            optimizer.zero_grad() 
            # 损失值进行反向传播
            loss.backward()  
            # 参数迭代更新
            optimizer.step()
            # 求损失的和
            train_loss_sum += loss.item()  
            # 输出模型预测结果
            _, predicts = torch.max(outputs.data, dim=1)  
            # 获取训练集预测正确的个数
            train_accuracy_total += labels.size(0)
            train_correct += (predicts == labels).cpu().sum().item()
            iter += 1
        # 测试集预测的准确率
        test_acc = evaluate_accuracy(test_loader, model)
        print(f'epoch:{epoch},   loss:{train_loss_sum / iter:.4f},   train accuracy:{train_correct / train_accuracy_total:.3f},  test accuracy:{test_acc:.3f}')
    print('------------finish training-------------')


train()    # 99.99%   95% +

(3)运行结果

image-20240503160227720

📖 笔记

04-图像分类-VGG

计算机视觉基础CV笔记

一、今日课程内容

  • VGG网络结构的特点(了解)
  • 能够利用VGG完成图像分类(实操+掌握)

二、VGG的网络架构(了解)

1、基本内容

  • 模型归属:牛津大学计算机视觉组(Visual Geometry Group)和 Google DeepMind公司的研究员
  • 荣誉:ILSVRC2014比赛分类项目的第二名
  • 特点:【3X3】的小卷积核 + 【2X2】的池化核,能够取得较好的识别精度
  • 常用模型:
  • VGG-16(【13】层卷积+2层全连接隐藏层+1层全连接输出层)
  • VGG-19(【16】层卷积+2层全连接隐藏层+1层全连接输出层)

2、模型调用及查看

  • 模型调用:

python import torch import torchvision.models as models # 模型实例化 model = models.vgg16()

  • 模型架构查看

python from torchsummary import summary # 通过net.summay()查看网络的形状 summary(model=model,input_size=(3,224,224),batch_size=1,device='cpu')

三、鲜花种类识别案例(实操+掌握)

1、读取数据

  • 图像Resize:224 X 224

  • 读取数据

```python
# 1.导入依赖包
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

# 2.定义数据转换方式
dataset_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Resize((224, 224))])

# 3.获取训练集数据和测试集数据,并进行数据转换
flower_train = ImageFolder(flower_train_path, transform=dataset_transform)

# 4.加载数据
train_loader = DataLoader(dataset=flower_train,
batch_size=batch_size,
shuffle=True)
```

  • 展示数据

python # 1.导入依赖包 import matplotlib.pyplot as plt # 2.定义图像和坐标1行2列 fig, axes = plt.subplots(1, 2) # 3.展示图像 axes.imshow(imgs.permute(1, 2, 0)) plt.show()

2、模型实例化和参数设置

(1)代码架构

# 1.模型实例化:输入数据3通道,进行5类的分类处理
# 2.模型训练的参数设置
# 2.1 学习率
# 2.2 训练轮数
# 2.3 优化算法Adam = RMSProp + Momentum
# 2.4 交叉熵损失函数

(2)完整代码

# 1.模型实例化:输入数据3通道,进行5类的分类处理
model = models.vgg16(num_classes=5)
# 2.模型训练的参数设置
# 2.1 学习率
learning_rate = 1e-3
# 2.2 训练轮数
num_epochs = 10
# 2.3 优化算法Adam = RMSProp + Momentum
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# 2.4 交叉熵损失函数
loss_fn = torch.nn.CrossEntropyLoss()

3、模型评估

(1)代码架构

# 1.定义模型评估函数
def evaluate_accuracy(data_iter, model):
    # 2.不进行梯度计算
        # 3.模型设置验证模式
        # 4.获取每一个batch的数据,进行预测
            # 4.1 获取预测结果
            # 4.2 预测的次数
            # 4.3 预测正确的个数
    # 5.获取准确率

(2)完整代码

# 1.定义模型评估函数
def evaluate_accuracy(data_iter, model):
    total = 0
    correct = 0
    # 2.不进行梯度计算
    with torch.no_grad():
        # 3.模型设置验证模式
        model.eval()
        # 4.获取每一个batch的数据,进行预测
        for images, labels in data_iter:
            outputs = model(images)
            # 4.1 获取预测结果
            _, predicts = torch.max(outputs.data, dim=1)
            # 4.2 预测的次数
            total += labels.size(0)
            # 4.3 预测正确的个数
            correct += (predicts == labels).cpu().sum()
            break
    # 5.获取准确率
    return correct / total

4、模型训练

(1)代码架构

# 1.定义模型训练过程:指定数据集,优化器,损失函数和轮次
def train(data_loader, optimizer, loss_fn, epochs):
    # 2.遍历每一个轮次进行训练
        # 3.每一个轮次的损失,预测个数和预测正确个数的初始化
        # 3.1 预测个数和预测正确个数的初始化
        # 3.2 损失值的和
        # 3.3 迭代次数
        # 4.遍历数据进行模型训练
            # 4.1 模型定义为训练模式
            # 4.2 对数据进行预测
            # 4.3 计算模型的损失
            # 4.4 在做反向传播前先清除网络状态
            # 4.5 损失值进行反向传播
            # 4.6 参数迭代更新
            # 4.7 求损失的和
            # 4.8 输出模型预测结果
            # 4.9 获取训练集预测正确的个数
        # 5.测试集预测的准确率

(2)完整代码

# 1.定义模型训练过程:指定数据集,优化器,损失函数和轮次
def train(data_loader=train_loader,
          optimizer=optimizer,
          loss_fn=loss_fn,
          epochs=num_epochs):
    # 2.遍历每一个轮次进行训练
    for epoch in range(epochs):
        print('current epoch = {}'.format(epoch))
        # 3.每一个轮次的损失,预测个数和预测正确个数的初始化
        # 3.1 预测个数和预测正确个数的初始化
        train_accuracy_total = 0
        train_correct = 0
        # 3.2 损失值的和
        train_loss_sum = 0
        # 3.3 迭代次数
        iter = 0
        # 4.遍历数据进行模型训练
        for i, (images, labels) in enumerate(data_loader):
            # 4.1 模型定义为训练模式
            model.train()
            # 4.2 对数据进行预测
            outputs = model(images)
            # 4.3 计算模型的损失
            loss = loss_fn(outputs, labels)  
            # 4.4 在做反向传播前先清除网络状态
            optimizer.zero_grad() 
            # 4.5 损失值进行反向传播
            loss.backward()  
            # 4.6 参数迭代更新
            optimizer.step()
            # 4.7 求损失的和
            train_loss_sum += loss.item()  
            # 4.8 输出模型预测结果
            _, predicts = torch.max(outputs.data, dim=1)  
            # 4.9 获取训练集预测正确的个数
            train_accuracy_total += labels.size(0)
            train_correct += (predicts == labels).cpu().sum().item()
            iter += 1
        # 5.测试集预测的准确率
        test_acc = evaluate_accuracy(test_loader, model)
        print('epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
            .format(epoch, train_loss_sum / (iter + 0.01),
                    train_correct / train_accuracy_total, test_acc))
📖 笔记

05-图像分类-GoogLeNet

计算机视觉基础CV笔记

一、今日课程内容

  • GoogLeNet网络结构的特点(了解)
  • 能够利用GoogLeNet完成图像分类(实操+掌握)

二、GoogLeNet网络结构的特点(了解)

1、基本内容

  • 模型归属:Google
  • 荣誉:在ImageNet2014分类比赛上的Top-5错误率降低到了6.7%
  • 特点:22层网络,在加深度的同时做了结构上的创新,引入叫做Inception的结构来代替之前的卷积加激活的经典组件。
  • 常用模块:
  • InceptionV1
  • InceptionV2
  • InceptionV3

2、Inception 块

image-20240423162241336

  • 4条并行线路
  • 前3条线路使用窗口大小分别是1×1、3×3和5×5的卷积层来抽取不同空间尺寸下的信息
  • 中间2个线路会对输入先做【1X1】卷积来减少输入通道数,以降低模型复杂度
  • 第4条线路则使用3×3最大池化层,后接1×1卷积层来改变通道数

  • 1 X 1卷积

  • 实现跨通道的交互和信息整合
  • 卷积核通道数的降维和升维,减少网络参数

2、GoogLeNet模型

  • 由Inception模块构成

  • 整个网络架构我们分为五个模块,每个模块之间使用步幅为2的3×3最大池化层来减小输出高宽。

  • B1模块:使用一个64通道的7×7卷积层
  • B2模块:使用2个卷积层:首先是64通道的1×1卷积层,然后是将通道增大3倍的3×3卷积层。
  • B3模块:串联2个完整的Inception块, 分别是Inception3a和Inception3b
  • B4模块:串联了5个Inception块,分别是Inception4a,Inception4b,Inception4c,Inception4d和Inception4e,并且在4b和4e模块添加了辅助输出端,用于模型训练
  • B5模块:串联了2个Inception块,分别是Inception5a,Inception5b,后面紧跟输出层,该模块使用全局平均池化层(GAP)来将每个通道的高和宽变成1。最后输出变成二维数组后接输出个数为标签类别数的全连接层。

  • 模型调用

```python
# 1.导入相关的工具包
import torch
import torchvision.models as models
# 2.模型实例化
model = models.googlenet()

# 3.B1模块
model.conv1
model.maxpool1
# 4.B2模块
model.conv2
model.conv3
model.maxpool2
# 5.B3模块(串联2个完整的Inception块)
model.inception3a
model.inception3b
# 6.B4模块(串联了5个Inception块)
model.inception4a
......
model.inception4e
model.aux1 # 4b和4e模块添加辅助输出端
# 7.B5模块(串联了2个Inception块,后面紧跟输出层)
model.inception5a
model.inception5b
```

三、鲜花种类识别案例(实操+掌握)

1、读取数据

  • 图像Resize:224 X 224

  • 读取数据

```python
# 1.导入依赖包
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

# 2.设置初始参数
# 2.1 指定批次大小
batch_size = 2
# 2.2 指定数据集路径
flower_train_path = './dataset/flower_datas/train/'
flower_test_path = './dataset/flower_datas/val/'

# 3.定义转换类型(先将数据转换为tensor类型,并调整数据的大小为224x224)
dataset_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Resize((224, 224))])

# 4.获取训练集数据和测试集数据
flower_train = ImageFolder(flower_train_path, transform=dataset_transform)
flower_test = ImageFolder(flower_test_path, transform=dataset_transform)

# 5.获取数据的迭代
train_loader = DataLoader(dataset=flower_train,
batch_size=batch_size,
shuffle=True)
test_loader = DataLoader(dataset=flower_test,
batch_size=batch_size,
shuffle=False)
```

  • 展示数据

python # 1.导入依赖包 import matplotlib.pyplot as plt # 2.定义图像和坐标1行2列 fig, axes = plt.subplots(1, 2) # 3.展示图像 axes.imshow(imgs.permute(1, 2, 0)) plt.show()

2、模型实例化和参数设置

(1)代码架构

# 1.模型实例化:输入数据3通道,进行5类的分类处理
# 2.模型训练的参数设置
# 2.1 学习率
# 2.2 训练轮数
# 2.3 优化算法Adam = RMSProp + Momentum
# 2.4 交叉熵损失函数

(2)完整代码

# 1.模型实例化:输入数据3通道,进行5类的分类处理
model = models.googlenet(num_classes=5)
# 2.模型训练的参数设置
# 2.1 学习率
learning_rate = 1e-3
# 2.2 训练轮数
num_epochs = 10
# 2.3 优化算法Adam = RMSProp + Momentum
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# 2.4 交叉熵损失函数
loss_fn = torch.nn.CrossEntropyLoss()

3、模型评估

(1)代码架构

# 1.定义模型评估函数
def evaluate_accuracy(data_iter, model):
    # 2.不进行梯度计算
        # 3.模型设置验证模式
        # 4.获取每一个batch的数据,进行预测
            # 4.1 获取预测结果
            # 4.2 预测的次数
            # 4.3 预测正确的个数
    # 5.获取准确率

(2)完整代码

# 1.定义模型评估函数
def evaluate_accuracy(data_iter, model):
    total = 0
    correct = 0
    # 2.不进行梯度计算
    with torch.no_grad():
        # 3.模型设置验证模式
        model.eval()
        # 4.获取每一个batch的数据,进行预测
        for images, labels in data_iter:
            outputs = model(images)
            # 4.1 获取预测结果
            _, predicts = torch.max(outputs.data, dim=1)
            # 4.2 预测的次数
            total += labels.size(0)
            # 4.3 预测正确的个数
            correct += (predicts == labels).cpu().sum()
    # 5.获取准确率
    return correct / total

4、模型训练

(1)代码架构

# 1.定义模型训练过程:指定数据集,优化器,损失函数和轮次
def train(data_loader, optimizer, loss_fn, epochs):
    # 2.遍历每一个轮次进行训练
        # 3.每一个轮次的损失,预测个数和预测正确个数的初始化
        # 3.1 预测个数和预测正确个数的初始化
        # 3.2 损失值的和
        # 3.3 迭代次数
        # 4.遍历数据进行模型训练
            # 4.1 模型定义为训练模式
            # 4.2 对数据进行预测
            # 4.3 各个输出端和模型的损失
            # 4.4 在做反向传播前先清除网络状态
            # 4.5 损失值进行反向传播
            # 4.6 参数迭代更新
            # 4.7 求损失的和
            # 4.8 输出模型预测结果
            # 4.9 获取训练集预测正确的个数
        # 5.测试集预测的准确率

(2)完整代码

# 1.定义模型训练过程:指定数据集,优化器,损失函数和轮次
def train(data_loader=train_loader,
          optimizer=optimizer,
          loss_fn=loss_fn,
          epochs=num_epochs):
    # 2.遍历每一个轮次进行训练
    for epoch in range(epochs):
        print('current epoch = {}'.format(epoch))
        # 3.每一个轮次的损失,预测个数和预测正确个数的初始化
        # 3.1 预测个数和预测正确个数的初始化
        train_accuracy_total = 0
        train_correct = 0
        # 3.2 损失值的和
        train_loss_sum = 0
        # 3.3 迭代次数
        iter = 0
        # 4.遍历数据进行模型训练
        for i, (images, labels) in enumerate(data_loader):
            # 4.1 模型定义为训练模式
            model.train()
            # 4.2 对数据进行预测
            outputs,aux2,aux1 = model(images)
            # 4.3 各个输出端和模型的损失
            loss0 = loss_fn(outputs, labels)
            loss1 = loss_fn(aux1,labels)
            loss2 = loss_fn(aux2,labels)
            loss = loss0 + 0.3*loss2+0.2*loss1
            # 4.4 在做反向传播前先清除网络状态
            optimizer.zero_grad() 
            # 4.5 损失值进行反向传播
            loss.backward()  
            # 4.6 参数迭代更新
            optimizer.step()
            # 4.7 求损失的和
            train_loss_sum += loss.item()  
            # 4.8 输出模型预测结果
            _, predicts = torch.max(outputs.data, dim=1)  
            # 4.9 获取训练集预测正确的个数
            train_accuracy_total += labels.size(0)
            train_correct += (predicts == labels).cpu().sum().item()
            iter += 1
        # 5.测试集预测的准确率
        test_acc = evaluate_accuracy(test_loader, model)
        print('epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
            .format(epoch, train_loss_sum / (iter + 0.01),
                    train_correct / train_accuracy_total, test_acc))

四、拓展知识

1、InceptionV2

  • 在InceptionV2中将大卷积核拆分为小卷积核
  • 将V1中的5×5的卷积用两个3×3的卷积替代
  • 增加网络的深度,减少了参数

2、InceptionV3

  • 将n×n卷积分割为1×n和n×1两个卷积
  • 参数量和计算量都比原来降低
📖 笔记

06-图像分类-ResNet

计算机视觉基础CV笔记

一、今日课程内容

  • ResNet网络结构的特点(了解)
  • 能够利用ResNet完成图像分类(实操+掌握)

二、ResNet模型及残差块(了解)

1、基本内容

  • 模型归属:何恺明等人
  • 荣誉:ILSVRC2015比赛分类项目的第一名
  • 特点:提出残差结构,使用shortcut使得模型网络可以更深

2、残差块

  • 直接从输入连接到输出的结构也称为shortcut
  • 整个结构就是残差块,ResNet的基础模块。
  • 2种模块
  • 基础模块:2个有相同输出通道数的3×3卷积层。每个卷积层后接BN层和ReLU激活函数,然后将输入直接加在最后的ReLU激活函数前。
    • ResNet18和ResNet34
  • 瓶颈模块(bottleneck):引入1×1卷积层来调整输入的通道数,达到减少模型参数量的作用。
    • ResNet50,ResNet101和ResNet152

3、ResNet模型

  • 网络结构
  • 第一个模块使用输出通道数为64、步幅为2的7×7卷积层,后接步幅为2的3×3的最大池化层。
  • 之后的每个模块在第一个残差块里将上一个模块的通道数翻倍,并将高和宽减半。
  • 最后,与GoogLeNet一样,加入全局平均池化层(GAP)后接上全连接层输出。
  • Max Pooling(固定窗口,固定步长,取最大值)
  • Avg Pooling(固定窗口,固定步长,取平均值)
  • GAP(不固定窗口,不固定步长,动态调整窗口和步长)

4、模型调用及查看

  • 模型调用:

python # 导入相关的工具包 import torch import torchvision.models as models # 模型实例化 model = models.resnet18()

  • 模型架构查看

python model.parameters

  • 模型网络层参看

python model.layer1

三、鲜花种类识别案例(实操+掌握)

1、读取数据

  • 图像Resize:224 X 224

  • 读取数据

```python
# 1.导入依赖包
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

# 2.定义数据转换方式
dataset_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Resize((224, 224))])

# 3.获取训练集数据和测试集数据,并进行数据转换
flower_train = ImageFolder(flower_train_path, transform=dataset_transform)

# 4.加载数据
train_loader = DataLoader(dataset=flower_train,
batch_size=batch_size,
shuffle=True)
```

  • 展示数据

python # 1.导入依赖包 import matplotlib.pyplot as plt # 2.定义图像和坐标1行2列 fig, axes = plt.subplots(1, 2) # 3.展示图像 axes.imshow(imgs.permute(1, 2, 0)) plt.show()

2、模型实例化和参数设置

(1)代码架构

# 1.模型实例化:输入数据3通道,进行5类的分类处理
# 2.模型训练的参数设置
# 2.1 学习率
# 2.2 训练轮数
# 2.3 优化算法Adam = RMSProp + Momentum
# 2.4 交叉熵损失函数

(2)完整代码

# 1.模型实例化:输入数据3通道,进行5类的分类处理
model = models.resnet18(num_classes=5)
# 2.模型训练的参数设置
# 2.1 学习率
learning_rate = 1e-3
# 2.2 训练轮数
num_epochs = 10
# 2.3 优化算法Adam = RMSProp + Momentum
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# 2.4 交叉熵损失函数
loss_fn = torch.nn.CrossEntropyLoss()

3、模型评估

(1)代码架构

# 1.定义模型评估函数
def evaluate_accuracy(data_iter, model):
    # 2.不进行梯度计算
        # 3.模型设置验证模式
        # 4.获取每一个batch的数据,进行预测
            # 4.1 获取预测结果
            # 4.2 预测的次数
            # 4.3 预测正确的个数
    # 5.获取准确率

(2)完整代码

# 1.定义模型评估函数
def evaluate_accuracy(data_iter, model):
    total = 0
    correct = 0
    # 2.不进行梯度计算
    with torch.no_grad():
        # 3.模型设置验证模式
        model.eval()
        # 4.获取每一个batch的数据,进行预测
        for images, labels in data_iter:
            outputs = model(images)
            # 4.1 获取预测结果
            _, predicts = torch.max(outputs.data, dim=1)
            # 4.2 预测的次数
            total += labels.size(0)
            # 4.3 预测正确的个数
            correct += (predicts == labels).cpu().sum()
    # 5.获取准确率
    return correct / total

4、模型训练

(1)代码架构

# 1.定义模型训练过程:指定数据集,优化器,损失函数和轮次
def train(data_loader, optimizer, loss_fn, epochs):
    # 2.遍历每一个轮次进行训练
        # 3.每一个轮次的损失,预测个数和预测正确个数的初始化
        # 3.1 预测个数和预测正确个数的初始化
        # 3.2 损失值的和
        # 3.3 迭代次数
        # 4.遍历数据进行模型训练
            # 4.1 模型定义为训练模式
            # 4.2 对数据进行预测
            # 4.3 计算模型的损失
            # 4.4 在做反向传播前先清除网络状态
            # 4.5 损失值进行反向传播
            # 4.6 参数迭代更新
            # 4.7 求损失的和
            # 4.8 输出模型预测结果
            # 4.9 获取训练集预测正确的个数
        # 5.测试集预测的准确率

(2)完整代码

# 1.定义模型训练过程:指定数据集,优化器,损失函数和轮次
def train(data_loader=train_loader,
          optimizer=optimizer,
          loss_fn=loss_fn,
          epochs=num_epochs):
    # 2.遍历每一个轮次进行训练
    for epoch in range(epochs):
        print('current epoch = {}'.format(epoch))
        # 3.每一个轮次的损失,预测个数和预测正确个数的初始化
        # 3.1 预测个数和预测正确个数的初始化
        train_accuracy_total = 0
        train_correct = 0
        # 3.2 损失值的和
        train_loss_sum = 0
        # 3.3 迭代次数
        iter = 0
        # 4.遍历数据进行模型训练
        for i, (images, labels) in enumerate(data_loader):
            # 4.1 模型定义为训练模式
            model.train()
            # 4.2 对数据进行预测
            outputs = model(images)
            # 4.3 计算模型的损失
            loss = loss_fn(outputs, labels)  
            # 4.4 在做反向传播前先清除网络状态
            optimizer.zero_grad() 
            # 4.5 损失值进行反向传播
            loss.backward()  
            # 4.6 参数迭代更新
            optimizer.step()
            # 4.7 求损失的和
            train_loss_sum += loss.item()  
            # 4.8 输出模型预测结果
            _, predicts = torch.max(outputs.data, dim=1)  
            # 4.9 获取训练集预测正确的个数
            train_accuracy_total += labels.size(0)
            train_correct += (predicts == labels).cpu().sum().item()
            iter += 1
        # 5.测试集预测的准确率
        test_acc = evaluate_accuracy(test_loader, model)
        print('epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
            .format(epoch, train_loss_sum / (iter + 0.01),
                    train_correct / train_accuracy_total, test_acc))
📖 笔记

07-图像分类-图像增强

计算机视觉基础CV笔记

一、今日课程内容

  • 图像增强的常用方法(了解)
  • 能够利用torchvision来完成图像增强(实操 + 掌握)

二、常用的图像增强方法(了解)

1、定义

  • 通过剪切、旋转/反射/翻转变换、缩放变换、平移变换、尺度变换、对比度变换、噪声扰动、颜色变换等一种或多种组合数据增强变换的方式来【transforms.Compose】

2、主要任务

  • 通过对训练图像做一系列随机改变,来产生相似但又不同的训练样本,从而扩大训练数据集的规模,而且随机改变训练样本可以降低模型对某些属性的依赖,从而提高模型的【泛化性能】。

3、常见方式

  • 几何变换类
  • 对图像进行几何变换操作
  • 翻转,旋转,裁剪,变形,缩放等
  • 颜色变换类
  • 通过模糊、颜色变换、擦除、填充等方式对图像进行处理

4、API实现

torchvision.transforms

三、torchvision.transforms进行图像增强(实操+掌握)

1、基本实现方式

  • 主要针对 PIL.Image 对象和 torch.Tensor 对象

  • 代码实现

python # 1.导入相应的工具包:图像读取,图像增强和绘制图像的 from PIL import Image # pip install pillow from torchvision import transforms import matplotlib.pyplot as plt # 2.读取图像并进行显示 dog = Image.open('./dog.jpg') plt.imshow(dog) plt.show()

2、几何变换

  • 1 左右翻转(水平翻转)

  • API

    python transforms.RandomHorizontalFlip(p=0.5)

  • 案例

    ```python

    水平翻转:以0.5的概率进行水平翻转

    transforms_HF =transforms.RandomHorizontalFlip(p=0.5)
    dog_HF = transforms_HF(dog)
    plt.imshow(dog_HF)
    plt.show()
    ```

  • 2 上下翻转(垂直翻转)

  • API

    python transforms.RandomVerticalFlip(p=0.5)

  • 案例

    ```python

    垂直上下翻转:以0.5的概率进行翻转

    transforms_VF=transforms.RandomVerticalFlip(p=0.5)
    dog_VF = transforms_VF(dog)
    plt.imshow(dog_VF)
    ```

  • 3 随机裁剪

  • API

    python transforms.RandomCrop(size=(H, W))

  • 案例

    ```python

    随机裁剪:宽高要小于原始图像的宽高

    transforms_RC=transforms.RandomCrop(size=(600, 600))
    dog_RC = transforms_RC(dog)
    plt.imshow(dog_RC)
    ```

  • 4 尺寸调整

  • API

    python transforms.Resize(size=(H, W))

  • 案例

    ```python

    尺寸调整

    transforms_RS=transforms.Resize(size=(224, 224))
    dog_RS = transforms_RS(dog)
    plt.imshow(dog_RS)
    ```

3、颜色变换

  • 从4个方面改变图像的颜色:亮度、对比度、饱和度和色调

  • API:

python transforms.ColorJitter(brightness=0, contrast=0, saturation=0, hue=0)

  • brightness:亮度,取值范围为 [0.0, 1.0]
    • 当取值为0.5时,表示将图像的亮度随机变化为原图亮度的50%(即1−0.5)∼150%(即1+0.5)。
  • contrast:对比度,取值范围为 [0.0, 1.0]
    • 当取值为0.5时,表示将图像的对比度随机变化为原图对比度的50%(即1−0.5)∼150%(即1+0.5)。
  • saturation:饱和度,取值范围为 [0.0, 1.0]
    • 当取值为0.5时,表示将图像的饱和度随机变化为原图饱和度的50%(即1−0.5)∼150%(即1+0.5)。
  • hue:色调,取值范围为 [-0.5, 0.5]

    • 当取值为0.2时,表示将图像的色调就在[-0.2,0.2]之间取值。
  • 案例

python # 颜色调整 transforms_color =transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.2) dog_color = transforms_color(dog) plt.imshow(dog_color) plt.show()

4、增强方法的组合

  • 多种增强方式组合在一起,例如:几何增强+ 颜色增强

  • API

python transforms.Compose([transforms])

  • 案例

python # 将几何变换和颜色变换组合在一起 transform = transforms.Compose([transforms_RS,transforms_color]) dog_ = transform(dog) plt.imshow(dog_) plt.show()

📖 笔记

08-图像分类-模型微调

计算机视觉基础CV笔记

一、今日课程内容

  • 微调的原理(了解)
  • 能够利用微调模型来完成图像的分类任务(实操 + 掌握)

二、微调(了解)

1、定义

  • 给定大规模数据集训练好预训练模型,基于该预训练模型在指定小数据集上进行训练,调整模型。
  • 相对于从头开始训练,微调会省去大量计算资源和计算时间,提高了计算效率,甚至提高准确率。

2、主要步骤

(1)预训练源模型:在源数据集(如ImageNet数据集)上预训练一个神经网络模型,即源模型。

(2)创建目标模型:创建一个新的神经网络模型,即目标模型。

(3)修改目标模型输出层:为目标模型添加一个输出大小为目标数据集类别个数的输出层,并随机初始化该层的模型参数。

(4)在目标数据集(如鲜花数据集)上训练目标模型

三、鲜花种类识别(实操 + 掌握)

1、基本内容

  • 对在ImageNet数据集上训练好的VGG模型进行微调

  • 依赖包

python import torch import torch.nn as nn from torchvision import transforms from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder import torchvision.models as models

2、获取数据集

  • Resize到224 X224

  • 代码实现

```python
# 1.初始化参数设置
# 1.1 指定批次大小
batch_size = 8
# 1.2 指定数据集路径
flower_train_path = './dataset/flower_datas/train/'
flower_test_path = './dataset/flower_datas/val/'

# 2.定义数据转换方式,先将数据转换为tensor类型,并调整数据的大小为224x224
dataset_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Resize((224,224))
])

# 3.获取训练集数据和测试集数据
flower_train = ImageFolder(flower_train_path, transform=dataset_transform)
flower_test = ImageFolder(flower_test_path,transform=dataset_transform)

# 4.获取数据的迭代
train_loader = DataLoader(dataset = flower_train, batch_size = batch_size, shuffle = True)
test_loader = DataLoader(dataset = flower_test, batch_size = batch_size, shuffle = False)
```

3、构建微调模型

  • 使用在ImageNet数据集上预训练的Vgg-16作为源模型

  • torchvision小于等于0.12版本我们指定pretrain=True即可

  • torchvision大于0.12版本的指定weights来自动下载并加载预训练的模型参数

  • 模型构建

python class VGG16_finetune(nn.Module): def __init__(self, num_class=5): super(VGG16_finetune, self).__init__() # 获取预训练模型 # 0.12 vgg16_net = models.vgg16(pretrained=True) # >0.12 # weights= models.VGG16_Weights.DEFAULT # vgg16_net = model.vgg16(weights = weights) self.num_class = num_class # 获取预训练模型的特征和GAP的结果 self.features = vgg16_net.features self.avgpool = vgg16_net.avgpool # 构建当前任务的输出端:FC self.classifier = nn.Sequential( nn.Linear(25088, 512), nn.ReLU(True), nn.Dropout(), nn.Linear(512, 128), nn.ReLU(True), nn.Dropout(), nn.Linear(128, self.num_class), ) # 前向传播过程 def forward(self, x): x = self.features(x) x = self.avgpool(x) x = torch.flatten(x, 1) x = self.classifier(x) return x

  • 实例化并冻结卷积层部分的参数

python # 模型实例化 model = VGG16_finetune(num_class = 5) # 冻结特征提取部分的参数,不进行梯度计算 for param in model.features.parameters(): param.requires_grad = False # 学习率 learning_rate = 1e-3 # 训练轮数 num_epochs = 10 # 优化算法Adam = RMSProp + Momentum,在这里我们只更新后面的分类输出层的参数。卷积层的参数使用预训练模型的结果 optimizer = torch.optim.Adam(model.classifier.parameters(), lr=learning_rate) # 交叉熵损失函数 loss_fn = torch.nn.CrossEntropyLoss()

4、模型训练

  • 定义评估函数
# 1.计算模型预测精度:测试集数据,模型
def evaluate_accuracy(data_iter, model):
    total = 0
    correct = 0
    # 2.不进行梯度计算
    with torch.no_grad():
        # 3.模型是验证模式
        model.eval()
        # 4.获取每一个batch的数据,进行预测
        for images, labels in data_iter:
            outputs = model(images)
            # 4.1 获取预测结果
            _, predicts = torch.max(outputs.data, dim=1)
            # 4.2 预测的次数
            total += labels.size(0)
            # 4.3 预测正确的个数
            correct += (predicts == labels).cpu().sum()
    # 5.获取准确率
    return correct / total
  • 模型训练
# 1.定义模型训练过程:指定数据集,优化器,损失函数和轮次
def train(data_loader=train_loader,
          optimizer=optimizer,
          loss_fn=loss_fn,
          epochs=num_epochs):
    # 2.遍历每一个轮次进行训练
    for epoch in range(epochs):
        print('current epoch = {}'.format(epoch))
        # 2.1 每一个轮次的损失,预测个数和预测正确个数的初始化
        train_accuracy_total = 0
        train_correct = 0
        # 2.2 损失值的和
        train_loss_sum = 0
        # 2.3 迭代次数
        iter = 0
        # 3.遍历数据集进行训练
        for i, (images, labels) in enumerate(data_loader):
            # 3.1 模型定义为训练模式
            model.train()
            # 3.2 对数据进行预测
            outputs = model(images)
            # 3.3 计算模型的损失
            loss = loss_fn(outputs, labels)  
            # 3.4 在做反向传播前先清除网络状态
            optimizer.zero_grad() 
            # 3.5 损失值进行反向传播
            loss.backward()  
            # 3.6 参数迭代更新
            optimizer.step()
            # 3.7 求损失的和
            train_loss_sum += loss.item()  
            # 3.8 输出模型预测结果
            _, predicts = torch.max(outputs.data, dim=1)  
            # 3.9获取训练集预测正确的个数
            train_accuracy_total += labels.size(0)
            train_correct += (predicts == labels).cpu().sum().item()
            iter+=1
        # 4.测试集预测的准确率
        test_acc = evaluate_accuracy(test_loader, model)
        print(
            'epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
            .format(epoch, train_loss_sum / (iter+0.01),
                    train_correct / train_accuracy_total, test_acc))
📖 笔记

09-图像分类-汽车车型分类案例

计算机视觉基础CV笔记

一、今日课程内容

  • 数据增强的使用方法(了解)
  • 能够实现模型微调的网络构建(实操 + 掌握)
  • 能够实现模型训练(实操 + 掌握)
  • 能够使用训练好的模型进行预测(实操 + 掌握)

二、数据集介绍(了解)

1、基本内容

  • 斯坦福汽车数据集
  • 196 种汽车类型的图像数据集
  • 共有 16,185 张图像,其中训练图像和测试图像分别为 8,144 张和 8,041 张
  • 分类基于汽车品牌、车型和年份

三、数据增强及获取(实操+掌握)

1、基本内容

  • 几何增强 + 颜色增强

2、代码架构

# 图像增强方法
# 1.水平翻转
# 2.随机旋转
# 3.尺寸调整
# 4.颜色调整

# 5.训练集增强方法的组合:首先转换为tensor,然后进行几何和颜色增强

# 6.验证集增强方法

# 7.获取并加载数据

# 8.展示数据

3、完整代码

# 图像增强方法
# 1.水平翻转
transforms_HF =transforms.RandomHorizontalFlip(p=0.5)
# 2.随机旋转
transforms_RR=transforms.RandomRotation(45)
# 3.尺寸调整
transforms_RS=transforms.Resize(size=(256,256))
# 4.颜色调整
transforms_color =transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.2)

# 5.训练集增强方法的组合:首先转换为tensor,然后进行几何和颜色增强
train_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms_HF,
    transforms_RR,
    transforms_RS,
    transforms_color
])

# 6.验证集增强方法
test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms_RS,
])

# 7.获取并加载数据
cars_train = ImageFolder(cars_train_path, transform=train_transform)
cars_test = ImageFolder(cars_test_path,transform=test_transform)

batch_size = 2
train_loader = DataLoader(dataset = cars_train, batch_size = batch_size, shuffle = True)
test_loader = DataLoader(dataset = cars_test, batch_size = batch_size, shuffle = False)

# 8.展示数据
import matplotlib.pyplot as plt
# 8.1 遍历每个迭代的数据,将其结果展示出来
for b, (imgs, targets) in enumerate(train_loader):
    # 8.2 获取第一个batch的图像
    if b == 1:
        # 8.3 将其进行展示
        fig, axes = plt.subplots(1, 2)
        # 8.4 遍历batch中的每个图像
        for i in range(batch_size):
            # 8.5 图像显示出来
            axes[i].imshow(imgs[i].permute(1, 2, 0))
            # 8.6 设置图像标题
            axes[i].set_title(targets[i].item())
        plt.show()
    elif b > 0:
        break

四、模型构建(实操+掌握)

1、基本内容

  • 使用imageNet的预训练模型RestNet34进行微调

2、代码架构

# 1.模型加载:进行模型微调,使用imageNet的预训练模型
# 2.将所有的参数层进行冻结
# 3.获取到fc层的输入
# 4.定义一个新的FC层,作为模型的输出层,进行车辆的分类

3、完整代码

# 1.模型加载:进行模型微调,使用imageNet的预训练模型
model = models.resnet34(pretrained=True)
# 2.将所有的参数层进行冻结
for param in model.parameters():
    param.requires_grad = False
# 3.获取到fc层的输入
num_fc_ftr = model.fc.in_features
# 4.定义一个新的FC层,作为模型的输出层,进行车辆的分类
model.fc = nn.Linear(num_fc_ftr, 196)

五、模型训练(实操+掌握)

1、基本内容

  • 定义训练过程中的评估方法和模型训练方法

2、代码架构

# 1.初始参数设置
# 1.1 学习率
# 1.2 训练轮数
# 1.3 优化算法Adam = RMSProp + Momentum (梯度、lr两方面优化下降更快更稳)
# 1.4 交叉熵损失函数

# 2.计算模型预测精度:测试集数据,模型
def evaluate_accuracy(data_iter, model):
    # 2.1 不进行梯度计算
        # 2.2 模型是验证模式
        # 2.3 获取每一个batch的数据,进行预测
        for images, labels in data_iter:
            # 2.4 获取预测结果
            # 2.5 预测的次数
            # 2.6 预测正确的个数
    # 2.7 获取准确率
    return correct / total

# 3.定义模型训练过程:指定数据集,优化器,损失函数和轮次
def train(data_loader=train_loader,
          optimizer=optimizer,
          loss_fn=loss_fn,
          epochs=num_epochs):
    # 3.1 遍历每一个轮次进行训练
    for epoch in range(epochs):
        # 3.2 每一个轮次的损失,预测个数和预测正确个数的初始化
        # 3.3 损失值的和
        # 3.4 迭代次数
        # 3.5 遍历数据进行模型训练
        for i, (images, labels) in enumerate(data_loader):
            # 3.5.1 模型定义为训练模式
            # 3.5.2 对数据进行预测
            # 3.5.3 计算模型的损失
            # 3.5.4 在做反向传播前先清除网络状态
            # 3.5.5 损失值进行反向传播
            # 3.5.6 参数迭代更新
            # 3.5.7 求损失的和
            # 3.5.8 输出模型预测结果
            # 3.5.9 获取训练集预测正确的个数
        # 3.6 测试集预测的准确率

3、完整代码

# 1.初始参数设置
# 1.1 学习率
learning_rate = 1e-3
# 1.2 训练轮数
num_epochs = 10
# 1.3 优化算法Adam = RMSProp + Momentum (梯度、lr两方面优化下降更快更稳)
optimizer = torch.optim.Adam(model.fc.parameters(), lr=learning_rate)
# 1.4 交叉熵损失函数
loss_fn = torch.nn.CrossEntropyLoss()
# 2.计算模型预测精度:测试集数据,模型
def evaluate_accuracy(data_iter, model):
    total = 0
    correct = 0
    # 2.1 不进行梯度计算
    with torch.no_grad():
        # 2.2 模型是验证模式
        model.eval()
        # 2.3 获取每一个batch的数据,进行预测
        for images, labels in data_iter:
            outputs = model(images)
            # 2.4 获取预测结果
            _, predicts = torch.max(outputs.data, dim=1)
            # 2.5 预测的次数
            total += labels.size(0)
            # 2.6 预测正确的个数
            correct += (predicts == labels).cpu().sum()
    # 2.7 获取准确率
    return correct / total

# 3.定义模型训练过程:指定数据集,优化器,损失函数和轮次
def train(data_loader=train_loader,
          optimizer=optimizer,
          loss_fn=loss_fn,
          epochs=num_epochs):
    # 3.1 遍历每一个轮次进行训练
    for epoch in range(epochs):
        print('current epoch = {}'.format(epoch))
        # 3.2 每一个轮次的损失,预测个数和预测正确个数的初始化
        train_accuracy_total = 0
        train_correct = 0
        # 3.3 损失值的和
        train_loss_sum = 0
        # 3.4 迭代次数
        iter = 0
        # 3.5 遍历数据进行模型训练
        for i, (images, labels) in enumerate(data_loader):
            # 3.5.1 模型定义为训练模式
            model.train()
            # 3.5.2 对数据进行预测
            outputs = model(images)
            # 3.5.3 计算模型的损失
            loss = loss_fn(outputs, labels)  
            # 3.5.4 在做反向传播前先清除网络状态
            optimizer.zero_grad() 
            # 3.5.5 损失值进行反向传播
            loss.backward()  
            # 3.5.6 参数迭代更新
            optimizer.step()
            # 3.5.7 求损失的和
            train_loss_sum += loss.item()  
            # 3.5.8 输出模型预测结果
            _, predicts = torch.max(outputs.data, dim=1)  
            # 3.5.9 获取训练集预测正确的个数
            train_accuracy_total += labels.size(0)
            train_correct += (predicts == labels).cpu().sum().item()
            iter+=1
        # 3.6 测试集预测的准确率
        test_acc = evaluate_accuracy(test_loader, model)
        torch.save(model.state_dict(), "./weights/model-{}.pth".format(epoch))
        print(
            'epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
            .format(epoch, train_loss_sum / (iter+0.01),
                    train_correct / train_accuracy_total, test_acc))

四、模型预测(实操+掌握)

1、基本内容

  • 加载训练好的模型,对图片数据进行预测

2、代码架构

# 1.导入工具包
# 2.加载训练好的模型参数
# 3.指定要预测的图像
# 4.对图像进行预测
with torch.no_grad():
    # 4.1 指定模型为验证模式
    # 4.2 获取图像的路径
    # 4.3 遍历每一幅图像进行预测
    for img_path in image_paths:
        # 4.3.1 打开图片
        # 4.3.2 增加batch维度
        # 4.3.3 将图像送入网络中进行预测
        # 4.3.4 获取预测结果:类别索引
        # 4.3.5 获取预测结果:类别分数
        # 4.3.6 获取对应的类别名称
        # 4.3.7 预测结果的展示

3、完整代码

# 1.导入工具包
import glob
import os
from PIL import Image
# 2.加载训练好的模型参数
model.load_state_dict(torch.load("./weights/model.pth",map_location=torch.device('cpu')))
# 3.指定要预测的图像
file_path= "./car_image/"
# 4.对图像进行预测
with torch.no_grad():
    # 4.1 指定模型为验证模式
    model.eval()
    # 4.2 获取图像的路径
    image_paths = glob.glob(os.path.join(file_path, '*.jpg'))
    # 4.3 遍历每一幅图像进行预测
    for img_path in image_paths:
        # 4.3.1 打开图片
        img = Image.open(img_path)
        # 4.3.2 增加batch维度
        batch = test_transform(img).unsqueeze(0)
        # 4.3.3 将图像送入网络中进行预测
        outputs = model(batch).squeeze(0).softmax(0)
        # 4.3.4 获取预测结果:类别索引
        class_id = outputs.argmax().item()
        # 4.3.5 获取预测结果:类别分数
        score = outputs[class_id].item()
        # 4.3.6 获取对应的类别名称
        category_name = cars_train.classes[class_id]
        # 4.3.7 预测结果的展示
        plt.imshow(img)
        plt.title(category_name)
        plt.show()
💻 代码

图像分类代码

18 个文件
  • 📁 car_image/
    • ⚙️ AUDi R8 2012.jpg 294.2KB
    • ⚙️ BMW M3 2012.jpg 129.5KB
    • ⚙️ BMW X6 2012.jpg 67.3KB
    • ⚙️ Tesla Model S 2012.jpg 79.2KB
    • ⚙️ Toyota Corolla 2012.jpg 192.7KB
  • 📄 01-dataset.py 1.0KB
    # 1.导入工具包
    import torchvision
    from torchvision.datasets import ImageFolder
    import matplotlib.pyplot as plt
    
    # 2.加载数据(公开数据集)
    # train_data = torchvision.datasets.MNIST(root='./dataset', train=True, download=True)
    # test_data = torchvision.datasets.MNIST(root='./dataset', train=False, download=True)
    # print(f'类别索引:{train_data.class_to_idx}')
    # print(len(train_data))
    
    # 3.加载数据集(自定义的数据集)
    from torchvision.datasets import ImageFolder
    from torchvision.transforms import ToTensor, Resize
    import matplotlib.pyplot as plt
    
    train_dataset = ImageFolder(root='dataset/flower_datas/train',
                                transform=torchvision.transforms.Compose(
                                    [ToTensor(), Resize((224, 224))]))
    
    
    print(f'训练集的数量:{len(train_dataset.imgs)}')
    
    # 绘图展示
    # plt.show()
    print(type(train_dataset.__getitem__(2000)[0]))
    print(train_dataset.__getitem__(2000)[0].shape)
    plt.imshow(train_dataset.__getitem__(2000)[0].permute(1, 2, 0))
    plt.show()
    
  • 📄 02-AlexNet.py 5.0KB
    # AlexNet模型构建
    # 1.导入工具包
    import torch.nn
    from torch import nn
    from torchsummary import summary
    from torchvision import transforms
    from torchvision.datasets import ImageFolder
    from torch.utils.data import DataLoader
    import matplotlib.pyplot as plt
    
    device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
    
    
    # 2.定义网络类
    class AlexNet(nn.Module):
        # init:构建层架构
        def __init__(self, in_dim, n_class):
            super().__init__()
            # 卷积层
            self.conv = nn.Sequential(
                nn.Conv2d(in_channels=in_dim, out_channels=96, kernel_size=11, stride=4, padding=0),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(3, 2),
                nn.Conv2d(96, 256, 5, 1, 2),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(3, 2),
                nn.Conv2d(256, 384, 3, 1, 1),
                nn.ReLU(inplace=True),
                nn.Conv2d(384, 384, 3, 1, 1),
                nn.ReLU(inplace=True),
                nn.Conv2d(384, 256, 3, 1, 1),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(3, 2)
            )
            # 全连接层
            self.fc = nn.Sequential(
                nn.Linear(9216, 4096),
                nn.ReLU(inplace=True),
                nn.Dropout(0.5),
                nn.Linear(4096, 4096),
                nn.ReLU(inplace=True),
                nn.Dropout(0.5),
                nn.Linear(4096, n_class)
            )
    
        # forward:完成前向传播
        def forward(self, x):
            x = self.conv(x)
            x = x.view(x.size(0), -1)
            out = self.fc(x)
            return out
    
    
    #
    # net = AlexNet(3, 1000)
    # print(summary(net, input_size=(3, 227, 227), batch_size=1))
    
    # --------------------------------鲜花分类案例------------------------------------------ #
    
    # 1.数据获取
    # 1.1 指定数据集路径
    flower_train_path = './dataset/flower_datas/train/'
    flower_test_path = './dataset/flower_datas/val/'
    # 1.2 定义数据处理方法(resize)
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize((227, 227))
    ])
    # 1.3 ImageFolder读取数据
    flower_train_data = ImageFolder(flower_train_path, transform=transform)
    flower_test_data = ImageFolder(flower_test_path, transform=transform)
    # 1.4 DataLoader获取batch数据
    
    batch_size = 1024  # config.py && 参数解析器
    
    flower_train_loader = DataLoader(dataset=flower_train_data, batch_size=batch_size, shuffle=True)
    flower_test_loader = DataLoader(dataset=flower_test_data, batch_size=batch_size, shuffle=False)
    # 1.5 数据的展示
    # for b, (images, targtes) in enumerate(flower_train_loader):
    #     if b == 10:
    #         fig, axes = plt.subplots(1, 2)
    #         for i in range(batch_size):
    #             axes[i].imshow(images[i].permute(1, 2, 0))
    #             axes[i].set_title(targtes[i].item())
    #         plt.show()
    #     elif b > 10:
    #         break
    
    # 2.模型实例化和参数设置
    # 2.1 模型实例化
    model = AlexNet(3, 5).to(device)
    # 2.2 参数设置
    # 2.2.1 lr
    lr = 0.001
    # 2.2.2 轮次
    num_epoch = 20
    # 2.2.3 损失
    loss = torch.nn.CrossEntropyLoss()
    # 2.2.4 优化器
    opti = torch.optim.Adam(model.parameters(), lr=lr)
    
    
    # 3.评估方法:准确率
    def eval_acc(test_data, model):
        # 3.1 设置参数
        corr = 0
        total = 0
        # 3.2 遍历数据,进行模型预测
        with torch.no_grad():
            model.eval()
            for imags, label in test_data:
                imags, label = imags.to(device), label.to(device)
                out_put = model(imags).to(device)
                _, pred = torch.max(out_put.data, dim=1)
                total += label.size(0)
                corr += (pred == label).cpu().sum()
        return corr / total
    
    
    # 4.模型训练
    # 4.1 定义训练方法train
    def train(data_loader=flower_train_loader, model=model, optim=opti, loss_fn=loss):
        # 4.1.1 遍历每个轮次
        for epoch in range(num_epoch):
            print(f'current epoch:{epoch}')
            # 4.1.2 定义一些参数
            train_data_num = 0
            train_data_corr = 0
            loss_sum = 0
            iter = 0
            # 4.1.3 遍历每个batch
            for i, (image, targtes) in enumerate(data_loader):
                image,targtes = image.to(device), targtes.to(device)
                # 4.1.4 将模型设置为训练模式
                model.train()
                # 4.1.5 前向传播
                output = model(image).to(device)
                # 4.1.6 损失计算
                loss = loss_fn(output, targtes)
                # 4.1.7 梯度归零
                optim.zero_grad()
                # 4.1.8 反向传播
                loss.backward()
                # 4.1.9 参数更新
                optim.step()
                # 4.1.10 计算准确率,loss的变化
                train_data_num += targtes.size(0)
                _, pred = torch.max(output.data, dim=1)
                train_data_corr += (pred == targtes).cpu().sum()
                loss_sum += loss.item()
                iter += 1
            corr = train_data_corr / train_data_num  # 训练集准确率
            # 4.1.11 评估模型效果
            test_corr = eval_acc(flower_test_loader, model)
            print(
                f'epoch:{epoch + 1},   loss:{loss_sum / (iter + 0.01):.4f},   train accuracy:{corr:.3f},  test accuracy:{test_corr:.3f}')
    
    
    if __name__ == '__main__':
        train()
    
  • 📄 03-VGGNet.py 5.1KB
    # 1.导入工具包
    import torch
    import torchvision.models as models
    from torchvision import models
    from torchsummary import summary
    import matplotlib.pyplot as plt
    from torchvision import transforms
    from torch.utils.data import DataLoader
    from torchvision.datasets import ImageFolder
    
    # 2. 实例化模型
    # model = models.vgg16()
    # 观察模型结构
    # print(summary(model=model, input_size=(3, 224, 224), batch_size=1))
    # print(model.parameters)
    
    # -----------------------------------鲜花识别案例-------------------------------------- #
    
    # 1.数据获取
    # 1.1 指定批次大小
    batch_size = 2
    # 1.2 指定数据集路径
    flower_train_path = './dataset/flower_datas/train/'
    flower_test_path = './dataset/flower_datas/val/'
    
    # 1.3 先将数据转换为tensor类型,并调整数据的大小为224x225
    dataset_transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Resize((224, 224))])
    
    # 1.4 获取训练集数据和测试集数据
    flower_train = ImageFolder(flower_train_path, transform=dataset_transform)
    flower_test = ImageFolder(flower_test_path, transform=dataset_transform)
    
    # 1.5 获取数据的迭代
    train_loader = DataLoader(dataset=flower_train,
                              batch_size=batch_size,
                              shuffle=True)
    test_loader = DataLoader(dataset=flower_test,
                             batch_size=batch_size,
                             shuffle=False)
    
    # 1.6 遍历每个迭代的数据,将其结果展示出来
    for b, (imgs, targets) in enumerate(train_loader):
        # 获取第一个batch的图像
        if b == 1:
            # 将其进行展示
            fig, axes = plt.subplots(1, 2)
            # 遍历batch中的每个图像
            for i in range(batch_size):
                # 图像显示出来
                axes[i].imshow(imgs[i].permute(1, 2, 0))
                # 设置图像标题
                axes[i].set_title(targets[i].item())
            # plt.show()
        elif b > 0:
            break
    
    # 2.模型实例化和参数设置
    # 2.1 模型实例化
    model = models.vgg16(num_classes=5)
    print(model)
    # print(model.parameters)
    # 2.2 模型训练的参数设置
    # 1 学习率
    learning_rate = 1e-3
    # 2 训练轮数
    num_epochs = 20
    # 3 优化算法Adam = RMSProp + Momentum
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    # 4 交叉熵损失函数
    loss_fn = torch.nn.CrossEntropyLoss()
    
    
    # 3.模型评估方法
    # 计算模型预测精度:测试集数据,模型
    def evaluate_accuracy(data_iter, model):
        total = 0
        correct = 0
        # 3.1 不进行梯度计算
        with torch.no_grad():
            # 3.2 模型是验证模式
            model.eval()
            # 3.3 获取每一个batch的数据,进行预测
            for images, labels in data_iter:
                outputs = model(images)
                # 3.4 获取预测结果
                _, predicts = torch.max(outputs.data, dim=1)
                # 3.5 预测的次数
                total += labels.size(0)
                # 3.6 预测正确的个数
                correct += (predicts == labels).cpu().sum()
        # 3.2 获取准确率
        return correct / total
    
    
    # 4.定义模型训练过程:指定数据集,优化器,损失函数和轮次
    # 4.1 定义训练方法train
    def train(data_loader=train_loader,
              optimizer=optimizer,
              loss_fn=loss_fn,
              epochs=num_epochs):
        # 4.1.1 遍历每一个轮次进行训练
        for epoch in range(epochs):
            print('current epoch = {}'.format(epoch))
            # 4.1.2 每一个轮次的损失,预测个数和预测正确个数的初始化
            train_accuracy_total = 0
            train_correct = 0
            # 4.1.3 损失值的和
            train_loss_sum = 0
            # 4.1.4 迭代次数
            iter = 0
            for i, (images, labels) in enumerate(data_loader):
                # 4.1.5 模型定义为训练模式
                model.train()
                # 4.1.6 对数据进行预测
                outputs = model(images)
                # 4.1.7 计算模型的损失
                loss = loss_fn(outputs, labels)
                # 4.1.8 在做反向传播前先清除网络状态
                optimizer.zero_grad()
                # 4.1.9 损失值进行反向传播
                loss.backward()
                # 4.1.10 参数迭代更新
                optimizer.step()
                # 4.1.11 求损失的和
                train_loss_sum += loss.item()
                # 4.1.12 输出模型预测结果
                _, predicts = torch.max(outputs.data, dim=1)
                # 4.1.13 获取训练集预测正确的个数
                train_accuracy_total += labels.size(0)
                train_correct += (predicts == labels).cpu().sum().item()
                iter += 1
            # 4.1.14 测试集预测的准确率
            test_acc = evaluate_accuracy(test_loader, model)
            print(
                'epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
                .format(epoch, train_loss_sum / (iter + 0.01),
                        train_correct / train_accuracy_total, test_acc))
            print(f'轮次:{epoch + 1},损失:{train_loss_sum / (iter + 0.01):.4f}, '
                  f'训练集准确率:{train_correct / train_accuracy_total:.3f},测试集准确率:{test_acc:.3f}')
        print('------------finish training-------------')
    
    
    # # train()
    # if __name__ == '__main__':
    #     train()
    
  • 📄 04-googleNet.py 4.7KB
    # 1.导入工具包
    import torch
    import torchvision.models as models
    from torchvision import transforms
    from torch.utils.data import DataLoader
    from torchvision.datasets import ImageFolder
    import matplotlib.pyplot as plt
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 2.模型实例化
    model = models.googlenet()
    # 3.模型展示
    # print(model.parameters)
    # print(model.inception5a)
    # print(model.aux1)
    
    # print(model)
    
    
    # -------------------------------鲜花案例的实现--------------------------------------------- #
    
    # 1.获取数据
    # 1.1 指定批次大小
    batch_size = 256
    # 1.2 指定数据集路径
    flower_train_path = './dataset/flower_datas/train/'
    flower_test_path = './dataset/flower_datas/val/'
    # 1.3 先将数据转换为tensor类型,并调整数据的大小为224x224
    dataset_transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Resize((224, 224))])
    # 1.4 获取训练集数据和测试集数据
    flower_train = ImageFolder(flower_train_path, transform=dataset_transform)
    flower_test = ImageFolder(flower_test_path, transform=dataset_transform)
    # 1.5 获取数据的迭代
    train_loader = DataLoader(dataset=flower_train,
                              batch_size=batch_size,
                              shuffle=True)
    test_loader = DataLoader(dataset=flower_test,
                             batch_size=batch_size,
                             shuffle=False)
    
    # # 1.6 图像展示,遍历每个迭代的数据,将其结果展示出来
    # for b, (imgs, targets) in enumerate(train_loader):
    #     # 获取第一个batch的图像
    #     if b == 20:
    #         # 将其进行展示
    #         fig, axes = plt.subplots(1, 2)
    #         # 遍历batch中的每个图像
    #         for i in range(batch_size):
    #             # 图像显示出来
    #             axes[i].imshow(imgs[i].permute(1, 2, 0))
    #             # 设置图像标题
    #             axes[i].set_title(targets[i].item())
    #         # plt.show()
    #     elif b > 20:
    #         break
    
    # 2.模型实例化及参数设置
    # 2.1 模型实例化
    model = models.googlenet(num_classes=5).to(device)
    
    # print(model.parameters)
    # 2.2 参数的设置
    # 2.2.1 学习率
    learning_rate = 1e-3
    # 2.2.2 训练轮数
    num_epochs = 3
    # 2.2.3 优化算法Adam = RMSProp + Momentum
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    # 2.2.4 交叉熵损失函数
    loss_fn = torch.nn.CrossEntropyLoss()
    
    
    # 3.定义模型评估方法,计算模型预测精度:测试集数据,模型
    def evaluate_accuracy(data_iter, model):
        total = 0
        correct = 0
        # 3.1 不进行梯度计算
        with torch.no_grad():
            # 3.2 模型是验证模式
            model.eval()
            # 3.3 获取每一个batch的数据,进行预测
            for images, labels in data_iter:
                images,labels = images.to(device), labels.to(device)
                # 3.4 google进行模型预测时只返回最终的结果
                outputs = model(images).to(device)
                # 3.5 获取预测结果
                _, predicts = torch.max(outputs.data, dim=1)
                # 3.6 预测的次数
                total += labels.size(0)
                # 3.7 预测正确的个数
                correct += (predicts == labels).cpu().sum()
        # 3.8 获取准确率
        return correct / total
    
    
    # 4.模型训练
    def train(data_loader=train_loader, optim=optimizer, loss_fn=loss_fn, epochs=num_epochs):
        # 4.1 遍历每个轮次
        for i in range(epochs):
            # 4.2 准确率和loss初试参数设置
            total_image = 0
            correct_image = 0
            loss_sum = 0
            iter = 0
            # 4.3 遍历每个bacth
            for b, (images, labels) in enumerate(data_loader):
                images, labels = images.to(device), labels.to(device)
                model.train()
                # 4.4 模型预测
                outputs, aux2, aux1 = model(images)
                # 4.5 损失计算
                loss0 = loss_fn(outputs, labels)
                loss1 = loss_fn(aux1, labels)
                loss2 = loss_fn(aux2, labels)
    
                loss = loss0 + 0.3 * loss2 + 0.2 * loss1
    
                loss_sum += loss.item()
                # 4.6 梯度清零
                optimizer.zero_grad()
                # 4.7 反向传播
                loss.backward()
                # 4.8 参数更新
                optimizer.step()
                # 4.9 统计数据
                total_image += labels.size(0)
                _, pred = torch.max(outputs.data, dim=1)
                correct_image += (pred == labels).cpu().sum().item()
                iter += 1
            # 4.10 测试集准确率
            test_acc = evaluate_accuracy(test_loader, model)
            print('epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
                  .format(i, loss_sum / (iter + 0.01),
                          correct_image / total_image, test_acc))
    
    if __name__ == '__main__':
        train()
    
  • 📄 05-resnet.py 5.1KB
    # 1.导入工具包
    import torch
    import torchvision.models as models
    # 2.实例化模型
    # model = models.resnet18()
    model = models.resnet50()
    # print(model)
    
    # 3.查看模型结构
    # print(model.layer4)
    
    # --------------------------------鲜花识别案例---------------------------------------- #
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 1.读取数据
    from torchvision import transforms
    from torch.utils.data import DataLoader
    from torchvision.datasets import ImageFolder
    # 1.1 指定批次大小
    batch_size = 256
    # 1.2 指定数据集路径
    flower_train_path = './dataset/flower_datas/train/'
    flower_test_path = './dataset/flower_datas/val/'
    # 1.3 先将数据转换为tensor类型,并调整数据的大小为224x224
    dataset_transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Resize((224,224))])
    # 1.4 获取训练集数据和测试集数据
    flower_train = ImageFolder(flower_train_path, transform=dataset_transform)
    flower_test = ImageFolder(flower_train_path, transform=dataset_transform)
    # 1.5 获取数据的迭代
    train_loader = DataLoader(dataset=flower_train,
                              batch_size=batch_size,
                              shuffle=True)
    test_loader = DataLoader(dataset=flower_test,
                             batch_size=batch_size,
                             shuffle=False)
    
    # import matplotlib.pyplot as plt
    # 1.6 遍历每个迭代的数据,将其结果展示出来
    # for b, (imgs, targets) in enumerate(train_loader):
    #     # 获取第一个batch的图像
    #     if b == 10:
    #         # 将其进行展示
    #         fig, axes = plt.subplots(1, 2)
    #         # 遍历batch中的每个图像
    #         for i in range(batch_size):
    #             # 图像显示出来
    #             axes[i].imshow(imgs[i].permute(1, 2, 0))
    #             # 设置图像标题
    #             axes[i].set_title(targets[i].item())
    #         plt.show()
    #     elif b > 10:
    #         break
    # 模型实例化:输入数据3通道,进行5类的分类处理
    
    # 2.模型实例化和参数设置
    # 2.1 模型实例化:输入数据3通道,进行5类的分类处理
    model = models.resnet18(num_classes=5).to(device)
    # 2.2 模型训练的参数设置
    # 2.2.1 学习率
    learning_rate = 1e-3
    # 2.2.2 训练轮数
    num_epochs = 10
    # 2.2.3 优化算法Adam = RMSProp + Momentum
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    # 2.2.4 交叉熵损失函数
    loss_fn = torch.nn.CrossEntropyLoss()
    
    
    # 3.定义评估函数
    def evaluate_accuracy(data_iter, model):
        total = 0
        correct = 0
        # 3.1 不进行梯度计算
        with torch.no_grad():
            # 3.2 模型是验证模式
            model.eval()
            # 3.3 获取每一个batch的数据,进行预测
            for images, labels in data_iter:
                images,labels = images.to(device),labels.to(device)
                outputs = model(images).to(device)
                # 3.4 获取预测结果
                _, predicts = torch.max(outputs.data, dim=1)
                # 3.5 预测的次数
                total += labels.size(0)
                # 3.6 预测正确的个数
                correct += (predicts == labels).cpu().sum()
        # 3.7 获取准确率
        return correct / total
    
    
    # 4.定义模型训练过程:指定数据集,优化器,损失函数和轮次
    def train(data_loader=train_loader,
              optimizer=optimizer,
              loss_fn=loss_fn,
              epochs=num_epochs):
        # 4.1 遍历每一个轮次进行训练
        for epoch in range(epochs):
            print('current epoch = {}'.format(epoch))
            # 4.2 定义初试参数
            # 每一个轮次的损失,预测个数和预测正确个数的初始化
            train_accuracy_total = 0
            train_correct = 0
            # 损失值的和
            train_loss_sum = 0
            # 迭代次数
            iter = 0
            # 4.3 遍历每一个batch的数据进行训练
            for i, (images, labels) in enumerate(data_loader):
                images,labels = images.to(device),labels.to(device)
                # 4.4 模型定义为训练模式
                model.train()
                # 4.5 对数据进行预测
                outputs = model(images).to(device)
                # 4.6 计算模型的损失
                loss = loss_fn(outputs, labels)
                # 4.7 在做反向传播前先清除网络状态
                optimizer.zero_grad()
                # 4.8 损失值进行反向传播
                loss.backward()
                # 4.9 参数迭代更新
                optimizer.step()
                # 4.10 求损失的和
                train_loss_sum += loss.item()
                # 4.11 输出模型预测结果
                _, predicts = torch.max(outputs.data, dim=1)
                # 4.12 获取训练集预测正确的个数
                train_accuracy_total += labels.size(0)
                train_correct += (predicts == labels).cpu().sum().item()
                iter+=1
            # 4.13 测试集预测的准确率
            test_acc = evaluate_accuracy(test_loader, model)
            print(
                'epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
                .format(epoch, train_loss_sum / (iter+0.01),
                        train_correct / train_accuracy_total, test_acc))
        print('------------finish training-------------')
    
    if __name__ == '__main__':
        train()
    
  • 📄 06-imageAgu.py 1.2KB
    # 1.导入工具包
    from PIL import Image
    from torchvision import transforms
    import matplotlib.pyplot as plt
    
    # 2.读取图片
    dog = Image.open('./dog.jpg')
    # 3. 图像展示
    plt.imshow(dog)
    plt.show()
    
    # -----------------------------图像增强(几何+颜色)---------------------------------------- #
    
    # 1.几何增强
    # 1.1水平翻转
    transforms_RH = transforms.RandomHorizontalFlip(p=0.8)
    dog_RH = transforms_RH(dog)
    plt.imshow(dog_RH)
    plt.show()
    
    # 1.2垂直翻转
    transforms_RV = transforms.RandomVerticalFlip(p=0.1)
    dog_RV = transforms_RV(dog)
    plt.imshow(dog_RV)
    plt.show()
    #
    # 1.3 裁剪
    transforms_Crop = transforms.RandomCrop(size=(100, 100))
    dog_crop = transforms_Crop(dog)
    plt.imshow(dog_crop)
    plt.show()
    
    # 1.4 缩放
    transforms_Resize = transforms.Resize(size=(200, 100))
    dog_resize = transforms_Resize(dog)
    plt.imshow(dog_resize)
    plt.show()
    
    
    # 2.颜色增强
    transforms_color = transforms.ColorJitter(brightness=0.6, contrast=0.6, saturation=0.6, hue=0.05)
    dog_color = transforms_color(dog)
    plt.imshow(dog_color)
    # plt.show()
    
    # 3.增强方式的组合
    transform = transforms.Compose([
        transforms_Resize,
        transforms_RH,
        transforms_color
    ])
    
    dog_ = transform(dog)
    plt.imshow(dog_)
    # plt.show()
    
  • 📄 07-finetune.py 6.1KB
    # 1.导入工具包
    import torch
    import torch.nn as nn
    from torchvision import transforms
    from torchvision.datasets import ImageFolder
    from torchvision import models
    from torch.utils.data import DataLoader
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    
    # 2.数据获取
    # 2.1 指定批次大小
    batch_size = 256
    # 2.2 指定数据集路径
    flower_train_path = './dataset/flower_datas/train/'
    flower_test_path = './dataset/flower_datas/val/'
    # 2.3 先将数据转换为tensor类型,并调整数据的大小为224x224
    dataset_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize((224, 224))
    ])
    # 2.4 获取训练集数据和测试集数据
    flower_train = ImageFolder(flower_train_path, transform=dataset_transform)
    flower_test = ImageFolder(flower_test_path, transform=dataset_transform)
    # 2.5 获取数据的迭代
    train_loader = DataLoader(dataset=flower_train, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(dataset=flower_test, batch_size=batch_size, shuffle=False)
    
    
    # 3.构建模型:预训练模型
    # class Vgg16_finetune(nn.Module):
    #     def __init__(self, num_class=5):
    #         super(Vgg16_finetune, self).__init__()
    #         # # torchvision <0.13
    #         # vgg16_pre = models.vgg16(pretrain=True)
    #         # torchvision >0.13
    #         vgg16_pre = models.vgg16(weights=models.VGG16_Weights.DEFAULT)
    #         self.features = vgg16_pre.features
    #         self.avgpool = vgg16_pre.avgpool
    #         self.classfier = nn.Sequential(
    #             nn.Linear(512 * 7 * 7, 512),
    #             nn.ReLU(True),
    #             nn.Dropout(0.2),
    #             nn.Linear(512, 128),
    #             nn.ReLU(True),
    #             nn.Dropout(0.2),
    #             nn.Linear(128, num_class),
    #         )
    #
    #     def forward(self, x):
    #         x = self.features(x)
    #         x = self.avgpool(x)
    #         x = torch.flatten(x, 1)
    #         output = self.classfier(x)
    #         return output
    
    
    # 微调模型构建  1个继承 2个方法
    class Vgg16_finetune(nn.Module):
        def __init__(self, num_classes=5):
            super().__init__()
            # 1.获取预训练模型
            ## torchvison <0.12
            # vgg16_net = models.vgg16(pretrained=True)
            ## torchvision >0.12
            vgg16_net = models.vgg16(weights=models.VGG16_Weights.DEFAULT)
    
            # 2.获取特征提取层(包含池化层)
            ## 1.特征提取层
            self.features = vgg16_net.features
            ## 2.全局平均池化层
            self.avgpool = vgg16_net.avgpool
            ## 3.重写classifier层
            self.classifier = nn.Sequential(
                nn.Linear(25088, 256),
                nn.ReLU(True),
                nn.Dropout(0.5),
                nn.Linear(256, 128),
                nn.ReLU(True),
                nn.Dropout(0.5),
                nn.Linear(128, num_classes)
            )
    
        def forward(self, x):
            x = self.features(x)
            x = self.avgpool(x)
            x = torch.flatten(x, 1)
            output = self.classifier(x)
            return output
    
    
    # 4.模型实例化和参数设置
    # 4.1 模型实例化
    model = Vgg16_finetune(num_classes=5).to(device)
    # 4.2 参数设置
    # 4.2.1 冻结参数
    for param in model.features.parameters():
        param.requires_grad = False
    # 4.2.2 学习率/优化器/损失函数/轮次
    lr = 0.001
    optimizer = torch.optim.Adam(model.classifier.parameters(), lr=lr)
    loss_fn = torch.nn.CrossEntropyLoss()
    num_epochs = 10
    
    
    # 5.模型训练
    # 5.1 定义评估函数
    def evaluate_accuracy(data_iter, model):
        total = 0
        correct = 0
        # 5.1.1 不进行梯度计算
        with torch.no_grad():
            # 5.1.2 模型是验证模式
            model.eval()
            # 5.1.3 获取每一个batch的数据,进行预测
            for images, labels in data_iter:
                images,labels = images.to(device), labels.to(device)
                outputs = model(images).to(device)
                # 5.1.4 获取预测结果
                _, predicts = torch.max(outputs.data, dim=1)
                # 5.1.5 预测的次数
                total += labels.size(0)
                # 5.1.6 预测正确的个数
                correct += (predicts == labels).cpu().sum()
        # 5.1.7 获取准确率
        return correct / total
    
    
    # 5.2 定义模型训练过程:指定数据集,优化器,损失函数和轮次
    def train(data_loader=train_loader,
              optimizer=optimizer,
              loss_fn=loss_fn,
              epochs=num_epochs):
        # 5.2.1 遍历每一个轮次进行训练
        for epoch in range(epochs):
            print('current epoch = {}'.format(epoch))
            # 5.2.2 每一个轮次的损失,预测个数和预测正确个数的初始化
            train_accuracy_total = 0
            train_correct = 0
            # 损失值的和
            train_loss_sum = 0
            # 迭代次数
            iter = 0
            # 5.2.3 遍历每一个批次的数据进行训练
            for i, (images, labels) in enumerate(data_loader):
                images, labels = images.to(device), labels.to(device)
                # 5.2.4 模型定义为训练模式
                model.train()
                # 5.2.5 对数据进行预测
                outputs = model(images).to(device)
                # 5.2.6 计算模型的损失
                loss = loss_fn(outputs, labels)
                # 5.2.7 在做反向传播前先清除网络状态
                optimizer.zero_grad()
                # 5.2.8 损失值进行反向传播
                loss.backward()
                # 5.2.9 参数迭代更新
                optimizer.step()
                # 5.2.10 求损失的和
                train_loss_sum += loss.item()
                # 5.2.11 输出模型预测结果
                _, predicts = torch.max(outputs.data, dim=1)
                # 5.2.12 获取训练集预测正确的个数
                train_accuracy_total += labels.size(0)
                train_correct += (predicts == labels).cpu().sum().item()
                iter += 1
            # 5.2.13 测试集预测的准确率
            test_acc = evaluate_accuracy(test_loader, model)
            print(
                'epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
                .format(epoch, train_loss_sum / (iter + 0.01),
                        train_correct / train_accuracy_total, test_acc))
        print('------------finish training-------------')
    
    
    if __name__ == '__main__':
        train()
    
  • 📄 08-carcls.py 6.4KB
    # 1.导入工具包
    import torch
    import torch.nn as nn
    from torchvision.datasets import ImageFolder
    from torchvision import transforms
    from torch.utils.data import DataLoader
    from torchvision import models
    
    # 2.读取数据
    # 2.1 指定数据集位置
    car_train_path = './dataset/Stanford_Cars/train/'
    car_test_path = './dataset/Stanford_Cars/test/'
    
    # 2.2 设置增强方式
    transform_HF = transforms.RandomHorizontalFlip(p=0.5)
    transform_RR = transforms.RandomRotation(45)
    transform_RS = transforms.Resize(size=(256, 256))
    transform_color = transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.05)
    train_transform = transforms.Compose([
        transforms.ToTensor(),
        transform_RS,
        transform_RR,
        transform_HF,
        transform_color
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transform_RS
    ])
    
    # 2.3 读取数据
    car_train = ImageFolder(car_train_path, transform=train_transform)
    car_test = ImageFolder(car_test_path, transform=test_transform)
    # 2.4 获取batch数据
    batch_size = 16
    train_loader = DataLoader(dataset=car_train, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(dataset=car_test, batch_size=batch_size, shuffle=False)
    
    # 2.5 遍历每个迭代的数据,将其结果展示出来
    # import matplotlib.pyplot as plt
    # for b, (imgs, targets) in enumerate(train_loader):
    #     # 获取第一个batch的图像
    #     if b == 1:
    #         # 将其进行展示
    #         fig, axes = plt.subplots(1, 2)
    #         # 遍历batch中的每个图像
    #         for i in range(batch_size):
    #             # 图像显示出来
    #             axes[i].imshow(imgs[i].permute(1, 2, 0))
    #             # 设置图像标题
    #             axes[i].set_title(targets[i].item())
    #         plt.show()
    #     elif b > 0:
    #         break
    # print(car_train.class_to_idx)
    
    # 模型构建
    
    # 3.模型实例化和参数设置
    # 3.1 模型实例化
    ## torcvison <0.13
    # model = models.resnet34(pretrained=True)
    ## torcvison >0.13
    # model = models.resnet34(weights=models.ResNet34_Weights.DEFAULT)
    # model = models.resnet50(pretrained=True)
    model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
    # print(model.parameters)
    # 3.2 参数设置
    # 3.2.1 冻结参数
    for param in model.parameters():
        param.requires_grad = False
    # 3.2.2 获取模型中fc层的输入特征数
    # 3.2.3 构建新的fc层
    model.fc = nn.Linear(model.fc.in_features, 196)
    
    # print(model.parameters)
    
    # 4.模型训练
    # 4.1 训练参数设置
    # 4.1.1 学习率
    learning_rate = 1e-3
    # 4.1.2 训练轮数
    num_epochs = 100
    # 4.1.3 优化算法Adam = RMSProp + Momentum (梯度、lr两方面优化下降更快更稳)
    optimizer = torch.optim.Adam(model.fc.parameters(), lr=learning_rate)
    # 4.1.4 交叉熵损失函数
    loss_fn = torch.nn.CrossEntropyLoss()
    
    
    # 4.2 定义评估方法
    def evaluate_accuracy(data_iter, model):
        total = 0
        correct = 0
        # 4.2.1 不进行梯度计算
        with torch.no_grad():
            # 4.2.2 模型是验证模式
            model.eval()
            # 4.2.3 获取每一个batch的数据,进行预测
            for images, labels in data_iter:
                outputs = model(images)
                # 4.2.4 获取预测结果
                _, predicts = torch.max(outputs.data, dim=1)
                # 4.2.5 预测的次数
                total += labels.size(0)
                # 4.2.6 预测正确的个数
                correct += (predicts == labels).cpu().sum()
        # 4.2.7 获取准确率
        return correct / total
    
    
    # 4.3 定义模型训练过程:指定数据集,优化器,损失函数和轮次
    def train(data_loader=train_loader,
              optimizer=optimizer,
              loss_fn=loss_fn,
              epochs=num_epochs):
        # 4.3.1 遍历每一个轮次进行训练
        for epoch in range(epochs):
            print('current epoch = {}'.format(epoch))
            # 4.3.2 每一个轮次的损失,迭代次数,预测个数和预测正确个数的初始化
            train_accuracy_total = 0
            train_correct = 0
            #  损失值的和
            train_loss_sum = 0
            #  迭代次数
            iter = 0
            for i, (images, labels) in enumerate(data_loader):
                # 4.3.3 模型定义为训练模式
                model.train()
                # 4.3.4 对数据进行预测
                outputs = model(images)
                # 4.3.5 计算模型的损失
                loss = loss_fn(outputs, labels)
                # 4.3.6 在做反向传播前先清除网络状态
                optimizer.zero_grad()
                # 4.3.7 损失值进行反向传播
                loss.backward()
                # 4.3.8 参数迭代更新
                optimizer.step()
                # 4.3.9 求损失的和
                train_loss_sum += loss.item()
                # 4.3.10 输出模型预测结果
                _, predicts = torch.max(outputs.data, dim=1)
                # 4.3.11 获取训练集预测正确的个数
                train_accuracy_total += labels.size(0)
                train_correct += (predicts == labels).cpu().sum().item()
                iter += 1
            # 4.3.12 测试集预测的准确率
            test_acc = evaluate_accuracy(test_loader, model)
            # 保存模型
            torch.save(model.state_dict(), "./weights/model-{}.pth".format(epoch))
            print(
                'epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
                .format(epoch + 1, train_loss_sum / (iter + 0.01),
                        train_correct / train_accuracy_total, test_acc))
        print('------------finish training-------------')
    
    
    # if __name__ == '__main__':
    #     train()
    
    
    # 5.模型预测
    import glob
    import os
    from PIL import Image
    import matplotlib.pyplot as plt
    
    # 5.1 加载模型
    model.load_state_dict(torch.load('./weights/model-149.pth', map_location=torch.device('cpu')),strict=False)
    # 5.2 待预测数据
    file_path = './car_image/'
    # 5.3 禁用梯度计算
    with torch.no_grad():
        # 5.4 设置为验证模式
        model.eval()
        # 5.5 遍历每一个图片
        image_paths = glob.glob(os.path.join(file_path, '*.jpg'))  # 批量获取图像数据
        for image_path in image_paths:
            # 5.6 读取图片
            img = Image.open(image_path)
            batch = test_transform(img).unsqueeze(0)  # (B,C,H,W)
            # 5.7 网络预测
            outputs = model(batch).squeeze(0).softmax(0)  # {0:0.78}
            class_id = outputs.argmax().item()
            score = outputs[class_id].item()  # 置信度
            print(score)
            labels_name = car_train.classes[class_id]  # {0:audi,1:Benz}
            # 5.8 显示图片
            plt.imshow(img)
            plt.title(labels_name)
            plt.show()
    
  • 📄 AlexNet_Flower.py 4.2KB
    # 1.导入依赖包
    import torch
    import torchvision
    from torch import nn
    from torchvision.datasets import ImageFolder
    from torchvision.transforms import Resize, ToTensor, Compose
    from torch.utils.data import DataLoader
    
    batch_size = 2
    num_epochs = 10
    
    # 2.获取数据  227*227  Resize ToTensor
    train_data = ImageFolder(root='dataset/flower_datas/train', transform=Compose([ToTensor(),
                                                                                   Resize((227, 227))]))
    test_data = ImageFolder(root='dataset/flower_datas/val', transform=Compose([ToTensor(),
                                                                                Resize((227, 227))]))
    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=True)
    
    # 查看数据集
    import matplotlib.pyplot as plt
    
    # 遍历每个迭代的数据,将其结果展示出来
    for b, (imgs, targets) in enumerate(train_loader):
        # 获取第一个batch的图像
        if b == 1:
            # 将其进行展示
            fig, axes = plt.subplots(1, 2)
            # 遍历batch中的每个图像
            for i in range(batch_size):
                # 图像显示出来
                axes[i].imshow(imgs[i].permute(1, 2, 0))
                # 设置图像标题
                axes[i].set_title(targets[i].item())
            plt.show()
        elif b > 0:
            break
    
    
    # 3.构建模型
    class AlexNet_Model(nn.Module):
        # 2.2 定义构造方法__init__
        def __init__(self, num_classes):
            super().__init__()
            self.conv = nn.Sequential(
                nn.Conv2d(3, 96, 11, 4, 0),
                nn.ReLU(inplace=True),  # inplace=True 内存复用
                nn.MaxPool2d(3, 2),
    
                nn.Conv2d(96, 256, 5, 1, 2),
                nn.ReLU(inplace=True),  # inplace=True 内存复用
                nn.MaxPool2d(3, 2),
    
                nn.Conv2d(256, 384, 3, 1, 1),
                nn.ReLU(inplace=True),  # inplace=True 内存复用
                nn.Conv2d(384, 384, 3, 1, 1),
                nn.ReLU(inplace=True),  # inplace=True 内存复用
                nn.Conv2d(384, 256, 3, 1, 1),
                nn.ReLU(inplace=True),  # inplace=True 内存复用
                nn.MaxPool2d(3, 2),
            )
            self.fc = nn.Sequential(
                nn.Linear(9216, 4096),
                nn.ReLU(inplace=True),
                nn.Dropout(p=0.5),
    
                nn.Linear(4096, 4096),
                nn.ReLU(inplace=True),
                nn.Dropout(p=0.5),
    
                nn.Linear(4096, num_classes)
            )
    
        # 2.3 定义前向传播方法forward
        def forward(self, x):
            x = self.conv(x)
            x = x.view(x.size(0), -1)
            x = self.fc(x)
            return x
    
    
    model = AlexNet_Model(num_classes=5)
    
    # 4.模型训练(225原则)
    def train():
        # 225----2(初始化参数) 损失+优化器
        loss_fn = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
        # 225----2(遍历-轮次-数据)
        for i in range(num_epochs):
            train_accuracy_total = 0
            train_correct = 0
            epoch_loss = 0
            iter = 0
            for b, (images, labels) in enumerate(train_loader):
                # 225----5(前,损失,梯度,反向,参数)
                # 1.前向传播
                output = model(images)
                # 2.损失计算
                loss = loss_fn(output, labels)
                # 3.梯度归零
                optimizer.zero_grad()
                # 4.反向传播  # w.grad b.grad
                loss.backward()
                # 5.参数更新
                optimizer.step()
    
                # 获取输出类别 pred==label
                _, pred = torch.max(output.data, dim=1)
    
                # 累积训练准确率
                train_accuracy_total += labels.size(0)
                train_correct += (pred == labels).cpu().sum().item()
                train_acc = train_correct / train_accuracy_total
    
                # 统计loss  每一个轮次平均loss
                epoch_loss += loss.item()
                iter += 1
                avg_loss = epoch_loss / iter
    
            print(f'epoch:{i+1},   loss:{avg_loss:.4f},   train accuracy:{train_acc:.3f}')
    
    
    if __name__ == '__main__':
        train()
    
    # 5.模型测试
    
    # 6.模型部署
    
  • 📄 Alexnet_demo.py 1.7KB
    # 构建一个AlexNet模型
    
    # 1.导入依赖包
    import torch.nn as nn
    
    
    # 2.  1个继承  2个方法
    # 2.1 定义类,继承nn.Module
    class AlexNet_Model(nn.Module):
        # 2.2 定义构造方法__init__
        def __init__(self, num_classes):
            super().__init__()
            self.conv = nn.Sequential(
                nn.Conv2d(3, 96, 11, 4, 0),
                nn.ReLU(inplace=True),  # inplace=True 内存复用
                nn.MaxPool2d(3, 2),
    
                nn.Conv2d(96, 256, 5, 1, 2),
                nn.ReLU(inplace=True),  # inplace=True 内存复用
                nn.MaxPool2d(3, 2),
    
                nn.Conv2d(256, 384, 3, 1, 1),
                nn.ReLU(inplace=True),  # inplace=True 内存复用
                nn.Conv2d(384, 384, 3, 1, 1),
                nn.ReLU(inplace=True),  # inplace=True 内存复用
                nn.Conv2d(384, 256, 3, 1, 1),
                nn.ReLU(inplace=True),  # inplace=True 内存复用
                nn.MaxPool2d(3, 2),
            )
            self.fc = nn.Sequential(
                nn.Linear(9216, 4096),
                nn.ReLU(inplace=True),
                nn.Dropout(p=0.5),
    
                nn.Linear(4096, 4096),
                nn.ReLU(inplace=True),
                nn.Dropout(p=0.5),
    
                nn.Linear(4096, num_classes)
            )
    
        # 2.3 定义前向传播方法forward
        def forward(self, x):
            x = self.conv(x)
            x = x.view(x.size(0), -1)
            x = self.fc(x)
            return x
    
    
    if __name__ == '__main__':
        # 查看模型结构 torchsummary--summary
        from torchsummary import summary
    
        # 1.实例化模型
        model = AlexNet_Model(num_classes=1000)
    
        # 2.查看模型结构
        summary(model, input_size=(3, 224, 224), device='cpu')   # 如果输入224*224    第一个线性层输入6400
    
  • 📄 IA_img_Aug.py 1.2KB
    # 1.导入依赖包
    from PIL import Image
    import matplotlib.pyplot as plt
    from torchvision.transforms import RandomHorizontalFlip, RandomVerticalFlip, RandomCrop, Resize, RandomRotation
    from torchvision.transforms import ColorJitter,Compose
    
    # 2.读取图像 PIL--pillow----Image.open() [C,H,W]
    img = Image.open('dog.jpg')
    plt.imshow(img)
    plt.show()
    
    # 3.几何增强
    # 3.1 随机水平翻转
    HF = RandomHorizontalFlip(p=0.8)
    hf_img = HF(img)
    plt.imshow(hf_img)
    # plt.show()
    # 3.2 随机垂直翻转
    VF = RandomVerticalFlip(p=0.8)
    VF_img = VF(img)
    plt.imshow(VF_img)
    # plt.show()
    # 3.3 随机裁剪
    Crop = RandomCrop((224, 224))
    crop_img = Crop(img)
    plt.imshow(crop_img)
    # plt.show()
    # 3.4 随机旋转
    RR = RandomRotation(45)
    RR_img = RR(img)
    plt.imshow(RR_img)
    # plt.show()
    # 3.5 尺寸调整
    Resize = Resize((224, 224))
    Resize_img = Resize(img)
    plt.imshow(Resize_img)
    # plt.show()
    
    # 4.颜色增强
    CJ = ColorJitter(brightness=0.5, contrast=0.5,saturation=0.5, hue=0.5)
    CJ_img = CJ(img)
    plt.imshow(CJ_img)
    # plt.show()
    
    
    # 5.组合增强(几何+颜色)
    compose = Compose([VF, RR, Resize, HF, CJ])
    compose_img = compose(img)
    plt.imshow(compose_img)
    plt.show()
    
  • 📄 dataset_demo.py 1.2KB
    # 数据集导入(公开数据集)
    # 1.导入依赖包
    import torchvision
    
    # 2.获取数据集
    # 2.1 训练集
    train_data = torchvision.datasets.CIFAR10(root='./dataset', train=True, download=True)
    test_data = torchvision.datasets.CIFAR10(root='./dataset', train=False, download=True)
    # 2.2 查看数据量
    # print(f'训练集数据量:{len(train_data)}')
    # print(f'测试集数据量:{len(test_data)}')
    # print(f'类别:{train_data.class_to_idx}')
    
    # 数据集导入(自定义数据集)
    # 1.导入依赖包
    from torchvision.datasets import ImageFolder
    from torchvision.transforms import ToTensor, Resize, Compose  # 组合
    
    # 2.读取数据集(数据路径 +数据增强)
    train_data = ImageFolder(root='dataset/flower_datas/train', transform=Compose([ToTensor(),
                                                                                   Resize((224, 224))]))  # H W
    test_data = ImageFolder(root='dataset/flower_datas/val', transform=Compose([ToTensor(),
                                                                                   Resize((224, 224))]))
    # 查看数据情况
    print(f'训练集数据量:{len(train_data)}')
    print(f'测试集数据量:{len(test_data)}')
    print(f'类别:{train_data.class_to_idx}')
    
  • ⚙️ dog.jpg 143.6KB

第3节 · 计算机视觉:目标检测

🎯 掌握 IoU / NMS、Faster R-CNN 与 YOLO 检测流程。
📖 笔记

01-目标检测-目标检测概述

计算机视觉基础CV笔记

一、今日课程内容

  • 目标检测的任务(了解)
  • 目标检测的常用数据集(了解)
  • 目标检测算法的评价指标(了解)
  • 非极大值NMS算法的应用(掌握)
  • 常用的目标检测算法分类(了解)

二、目标检测(了解)

1、定义

  • 找出图像中所有感兴趣的目标,并确定它们的【类别】和【位置】。

2、主要内容

  • 类别:取决于当前任务(数据集)需要检测的物体类别

  • 位置:极坐标表示和中心点坐标表示

  • 1、极坐标表示:(xmin, ymin, xmax, ymax)

    • xmin , ymin: x,y坐标的最小值
    • xmin , ymin: x,y坐标的最大值
  • 2、中心点坐标:(x_center , y_center, w , h)

  • x_center , y_center: 目标检测框的中心点坐标

  • w , h: 目标检测框的宽、高

三、常用的开源数据集(了解)

1、PASCAL VOC数据集

  • 目标检测领域的经典基准数据集
  • 包含约10,000张带有边界框的图片用于训练和验证
  • 共20个类别
  • 1.人: 人
  • 2.动物: 鸟,猫,牛,狗,马,羊
  • 3.交通工具: 飞机,自行车,船,公共汽车,汽车,摩托车,火车
  • 4.室内: 瓶子,椅子,餐桌,盆栽,沙发,电视/显示器
  • 常用的是VOC2007和VOC2012两个版本
  • 下载地址:https://pjreddie.com/projects/pascal-voc-dataset-mirror/
  • 目录结构
  • JPEGImages存放图片文件
  • Annotations下存放的是xml文件,描述了图片信息
    • xmin , ymin构成了boundingbox的左上角
    • xmax , ymax构成了boundingbox的右下角
  • ImageSets包含以下4个文件夹:
    • Action下存放的是人的动作(例如running、jumping等等)
    • Layout下存放的是具有人体部位的数据(人的head、hand、feet等等)
    • Segmentation下存放的是可用于分割的数据。
    • Main下存放的是图像物体识别的数据,总共分为20类,这是进行目标检测的重点。

2、MS COCO数据集

  • 全称是Microsoft Common Objects in Context
  • 微软于2014年出资标注的Microsoft COCO数据集
  • 一个大型的、丰富的物体检测,分割和字幕数据集
  • 以场景理解为目标,主要从复杂的日常场景中截取,图像中的目标通过精确的分割进行位置的标定
  • 图像包括91类目标,328,000影像和2,500,000个label
  • 80 类,有超过33 万张图片,其中20 万张有标注,整个数据集中个体的数目超过150 万个。
  • 标签文件:
  • segmentation+bounding box (精度均为小数点后两位)

四、常用的评价指标(了解)

1、IOU

  • IoU(intersection over union,交并比)是目标检测算法中用来评价2个矩形框之间相似度的指标

  • IoU = 两个矩形框相交的面积 / 两个矩形框相并的面积

  • 代码实现

python import numpy as np # 1.定义方法计算IOU def Iou(box1, box2): # 2.使用极坐标形式表示:直接获取两个bbox的坐标 # 2.1 获取两个框的左上角坐标和右下角坐标 xmin1, ymin1, xmax1, ymax1 = box1 xmin2, ymin2, xmax2, ymax2 = box2 # 3.获取矩形框交集对应的左上角和右下角的坐标(intersection) xx1 = np.max([xmin1, xmin2]) yy1 = np.max([ymin1, ymin2]) xx2 = np.min([xmax1, xmax2]) yy2 = np.min([ymax1, ymax2]) # 4.计算两个矩形框面积 area1 = (xmax1-xmin1) * (ymax1-ymin1) area2 = (xmax2-xmin2) * (ymax2-ymin2) # 5.计算交集面积 inter_area = (np.max([0, xx2-xx1])) * (np.max([0, yy2-yy1])) # 6.计算交并比 iou = inter_area / (area1+area2-inter_area) return iou

2、mAP

  • 评估模型的物体分类和定位性能的指标
  • mAP是多个分类任务的AP的平均值,而AP(average precision)是PR曲线下的面积
  • TP、FP、FN、TN
  • TP(True Positive):IOU大于阈值的检测框数量
  • FP(False Positive):IOU小于等于阈值的检测框数量
  • FN(False Negative ):没有检测到的真实框的数量
  • TN(True Negative):真负例,在 mAP 评价指标中不会使用到

  • 查准率、查全率

  • 查准率(Precision): TP/(TP + FP)
  • 查全率(Recall): TP/(TP + FN)

  • AP值计算方法

  • (1)VOC2010以前
    • 选取当Recall >= 0, 0.1, 0.2, ..., 1共11个点时的Precision最大值,然后AP就是这11个Precision的平均值
  • (2)VOC2010以后
    • 针对每一个不同的Recall值(包括0和1),选取其大于等于这些Recall值时的Precision最大值

五、NMS(非极大值抑制)

1、基本内容

  • 非极大值抑制(Non-Maximum Suppression,NMS),顾名思义就是抑制不是极大值的元素

  • 选取那些邻域里分数最高,并且抑制那些分数低的窗口

  • 目的就是去除冗余的检测框,保留最好的一个。

  • 使用流程

  • 首先是检测出一系列的检测框

  • 将检测框按照类别进行分类
  • 对同一类别的检测框应用NMS获取最终的检测结果

  • 代码实现

```python
import numpy as np
def nms(bboxes, confidence_score, threshold):
"""非极大抑制过程
:param bboxes: 同类别候选框坐标
:param confidence: 同类别候选框分数
:param threshold: iou阈值
:return:
"""
# 1、强转数组,并取出每个框左上角坐标和右下角坐标
bboxes = np.array(bboxes)
score = np.array(confidence_score)

  # 取出n个的极坐标点(左上角坐标和右下角坐标)
  x1 = bboxes[:, 0]
  y1 = bboxes[:, 1]
  x2 = bboxes[:, 2]
  y2 = bboxes[:, 3]
  # 计算每个框的面积
  areas = (x2 - x1) * (y2 - y1)
  # 2、对候选框进行NMS筛选
  # 2.1 返回的框坐标和分数
  picked_boxes = []
  picked_score = []
  # 2.2 对置信度进行排序, 获取排序后的下标序号, argsort默认从小到大排序
  order = np.argsort(score)
  while order.size > 0:
      # 2.3 将当前置信度最大的框加入返回值列表中
      index = order[-1]
      # 2.4 保留该类剩余box中得分最高的一个
      picked_boxes.append(bboxes[index])
      picked_score.append(confidence_score[index])
      # 2.5 计算交并比
      # 2.5.1 获取当前置信度最大的候选框与其他任意候选框的相交面积
      x11 = np.maximum(x1[index], x1[order[:-1]])
      y11 = np.maximum(y1[index], y1[order[:-1]])
      x22 = np.minimum(x2[index], x2[order[:-1]])
      y22 = np.minimum(y2[index], y2[order[:-1]])
      # 2.5.2 计算交集的面积,不重叠时面积为0
      w = np.maximum(0.0, x22 - x11)
      h = np.maximum(0.0, y22 - y11)
      intersection = w * h
      # 2.5.3 利用相交的面积和两个框自身的面积计算框的交并比
      ratio = intersection / (areas[index] + areas[order[:-1]] - intersection)
      # 2.5.4 保留IoU小于阈值的box
      keep_boxes_indics = np.where(ratio < threshold)
      # 2.5.5 保留剩余的框
      order = order[keep_boxes_indics]

  # 2.6 返回NMS后的框及分类结果   
  return picked_boxes, picked_score

```

六、目标检测方法分类

1、基本内容

  • two-stage(两阶段)
  • 先由算法生成一系列作为样本的候选框,再通过卷积神经网络进行样本分类
  • 特点:网络的准确度高、速度相对较慢
  • 代表算法:RCNN系列:R-CNN到Faster R-CNN网络
  • one-stage(单阶段)
  • 直接通过主干网络给出目标的类别和位置信息,没有使用候选区域的筛选网路
  • 特点:算法速度快,但是精度相对Two-stage目标检测网络降低了很多。
  • 代表算法: YOLO系列:YOLOv1、YOLOv2、YOLOv3、 SSD等
📖 笔记

02-目标检测-R-CNN网络基础

计算机视觉基础CV笔记

一、今日课程内容

  • Overfeat模型的移动窗口方法(了解)
  • RCNN目标检测的思想(掌握)
  • fastRCNN目标检测的思想(掌握)
  • 多任务损失(掌握)

二、Overfeat模型(了解)

1、定义

  • Overfeat方法:使用滑动窗口和神经网络来检测目标的方法
  • 实现方法:滑动窗口使用固定宽度和高度的矩形区域,在图像上“滑动”,并将扫描结果送入到神经网络中进行分类和回归。
  • 缺点:类似一种暴力穷举的方式,会消耗大量的计算力,并且由于窗口大小问题可能会造成效果不准确。

三、RCNN模型(掌握)

1、基本内容

  • 2014年提出,使用候选区域方法(region proposal method)创建目标检测的区域来完成目标检测的任务
  • R-CNN是以深度神经网络为基础的目标检测的模型 ,以R-CNN为基点,后续的Fast R-CNN、Faster R-CNN模型都延续了这种目标检测思路。

2、算法流程

  1. 候选区域生成:使用选择性搜索(Selective Search)的方法找出图片中可能存在目标的侯选区域
    - 候选区域搜索依据:使用语义分割的方法,将颜色、边界、纹理等信息作为合并条件,采用多尺度的综合方法
  2. CNN网络提取特征:选取预训练卷积神经网网络(AlexNet或VGG)用于进行特征提取。
    - 采用预训练模型(AlexNet或VGG)在生成的候选区域上进行特征提取
    - 特征提取流程:
    • (1)候选区域需进行裁剪或变形为固定的尺寸,再送入CNN网络中进行特征提取(全连接层输入数据的尺寸是固定的)
    • (2)全连接层由1000修改为N+1(预训练模型在ImageNet数据集上获得)
    • (3)利用微调后的CNN网络,提取每一个候选区域的特征,获取一个4096维的特征,一幅图像就是2000x4096维特征存储到磁盘中。
  3. 目标分类:训练支持向量机(SVM)来辨别目标物体和背景,对每个类别,都要训练一个二元SVM。
    - 对于N个类别的检测任务,需要训练N(目标类别数目)个SVM分类器
    - 对候选区域的特征向量(4096维)进行二分类,判断其是某一类别的目标,还是背景来完成目标分类。
  4. 目标定位:训练一个线性回归模型,为每个辨识到的物体生成更精确的边界框。
    - 训练一个线性回归模型在给定的候选区域的结果上去预测一个新的检测窗口,能够获得更精确的位置
  5. 模型预测
    - 使用选择性搜索的方法从一张图片中提取2000个候选区域
    - 将每个区域送入CNN网络中进行特征提取,然后送入到SVM中进行分类
    - 使用候选框回归器,计算出每个候选区域的位置。
    - 候选区域较多,有2000个,需要剔除掉部分检测结果。 针对每个类,通过计算IOU,采取非最大值抑制NMS的方法,保留比较好的检测结果。

3、算法总结

1、训练阶段多,训练耗时: 微调CNN网络+训练SVM+训练边框回归器。

2、预测速度慢: 使用GPU, VGG16模型处理一张图像需要47s。

3、占用磁盘空间大:5000张图像产生几百G的特征文件。

4、数据的形状变化:候选区域要经过缩放来固定大小,无法保证目标的不变形

四、Fast RCNN模型(了解)

1、基本内容

  • 提出时间:2015年
  • 改进点:
  • 1、提高训练和预测的速度
    • 将整张图输入到CNN中提取特征,将候选区域映射到特征图上,这样就避免了对图像区域进行重复处理,提高效率减少时间
  • 2、不需要额外的空间保存CNN网络提取的特征向量
    • 将类别判断和边框回归统一使用CNN实现,不需要在额外的空间存储特征
  • 3、不在直接对候选区域进行缩放
    • 使用ROIpooling的方法进行尺寸的调整。

2、算法流程

  • 1、候选区域生成:使用选择性搜索(Selective Search)的方法找出图片中可能存在目标的侯选区域,只需要候选区域的位置信息

  • 2、CNN网络特征提取:将整张图像输入到CNN网络中,得到整副图的特征图,并将上一步获取的候选区域位置从原图映射到该特征图上

  • 3、ROIPooling: 对于每个特征图上候选框,RoI pooling层从特征图中提取固定长度的特征向量每个特征向量被送入一系列全连接(fc)层中。

  • 候选区域从原图映射到特征图中后,进行ROIpooling的计算
  • ROI Pooling层使用最大池化将输入的特征图中的任意区域(候选区域对应的区域)内的特征转化为固定的𝐻×𝑊的特征图,其中𝐻和𝑊是超参数
  • 对于任意输入的ℎ×𝑤的候选区域,将其分割为𝐻×𝑊的子网格,每个子网格的大小为:(h/H) x (w/W),取每个子网格中的最大值,送入后续网络中进行处理

  • 4、目标检测:分两部分完成,一个输出各类别加上1个背景类别的Softmax概率估计,另一个为各类别的每一个类别输出四个实数值,来确定目标的位置信息。

  • 原网络的最后一个全连接层替换为两个同级层:K+1个类别的SoftMax分类层和边框的回归层。

3、模型训练

  • R-CNN中的特征提取和检测部分是分开进行的
  • Fast R-CNN提出一个高效的训练方法:多任务训练
  • (1)一部分输出在K+1个类别上的离散概率分布(每个候选区域)
  • (2)另一部分输出对于由K个类别中的每一个检测框回归偏移

image-20240424174300084

4、模型预测

(1)输入图像

(2)图像被送入到卷积网络进行特征提取,将通过选择性搜索获取的候选区域映射到特征图中

(3)在特征图上Rol中应用RoIPooling,获取尺寸相同的特征向量

(4)将这些区域传递到全连接的网络中进行分类和回归,得到目标检测的结果。

5、模型总结

Fast R-CNN是对R-CNN模型的一种改进:

  1. CNN网络不再对每个候选区域进行特征提取,而是直接对整张图像进行出路,这样减少了很多重复计算。
  2. 用ROI pooling进行特征的尺寸变换,来满足FC全连接层对输入数据尺度的要求。
  3. 将目标的回归和分类统一在一个网络中,使用FC+softmax进行目标分类,使用FC Layer进行目标框的回归。

缺点:候选区域提取方法耗时较长,而且和目标检测网络是分离的,并不是端到端的

后续优化:2016年,提出了Faster-RCNN模型

📖 笔记

03-目标检测-Faster-RCNN网络

计算机视觉基础CV笔记

一、今日课程内容

  • FasterRCNN目标检测的思想(了解)
  • anchor(锚框)的思想(了解)
  • RPN网络是如何进行候选区域的生成的(了解)
  • ROIPooling的使用方法(掌握)
  • FasterRCNN的训练方法(掌握)

二、网络工作流程(了解)

1、基本流程

  • 1、特征提取:将整个图像缩放至固定的大小输入到CNN网络中进行特征提取,得到特征图。
  • 2、候选区域提取:输入特征图,使用区域生成网络RPN,产生一些列的候选区域
  • 3、ROIPooling: 与Fast RCNN网络中一样,使用最大池化固定候选区域的尺寸,送入后续网络中进行处理
  • 4、目标分类和回归:与Fast RCNN网络中一样,使用两个同级层:K+1个类别的SoftMax分类层和边框的回归层,来完成目标的分类和回归。

2、数据加载及转换

(1)代码架构

# 1.导入依赖包
# 1.1 pytorh相关的工具包
# 1.2 读取图像的工具包
# 1.3 绘图工具包

# 2.获取图片

# 3.将图片格式转换为tensor形式,大小转换为800x800的大小
# 3.1 查看图像大小
# 3.2 展示图像

#### (2)完整代码

# 1.导入依赖包
# 1.1 pytorh相关的工具包
import torch
import torchvision
import torchvision.transforms as transforms
from torchvision.io import read_image
# 1.2 读取图像的工具包
from PIL import Image
# 1.3 绘图工具包
import matplotlib.pyplot as plt

# 2.获取图片
img = Image.open('img3.jpg')   # [C,H,W]      opencv     # BGR
plt.imshow(img)
plt.show()

# 3.将图片格式转换为tensor形式,大小转换为800x800的大小
transform = transforms.Compose([transforms.ToTensor(),transforms.Resize((800,800))])
img = transform(img)
# 3.1 查看图像大小
img.shape
# 3.2 展示图像
plt.imshow(img.permute(1,2,0))   # [H,W,C]

3、模型加载

(1)代码架构

# 1.实例化模型
# 2.将模型设置为eval模式     

(2)完整代码

backbone   + Neck   +  head

# 1.实例化模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# 2.将模型设置为eval模式
model.eval()

4、模型预测过程

(1)代码架构

# 1.直接进行预测

# 2.获取类别名称,框及对应的置信度

# 3.将框绘制在图像上
# 3.1 图像读取
# 3.2 图像展示
# 3.3 将框绘制在图像上

# 4.过滤分数较低的预测
# 4.1 因pred_score是从大大小进行排列的,只要获取最后一个索引即可

# 5.将过滤后的结构绘制在图像上
# 5.1 绘制图像
# 5.2 将框绘制在图像上

(2)完整代码

# 1.直接进行预测
pred = model([img])

# 2.获取类别名称,框及对应的置信度
pred_class = [coco_names[i] for i in list(pred[0]['labels'].numpy())] 
pred_boxes = [[(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]['boxes'].detach().numpy())]  
pred_score = list(pred[0]['scores'].detach().numpy())

# 3.将框绘制在图像上
# 3.1 图像读取
img = Image.open('img3.jpg').resize((800,800))
# 3.2 图像展示
fig = plt.imshow(img)
# 3.3 将框绘制在图像上
for index, boxes in enumerate(pred_boxes):
    fig.axes.add_patch(plt.Rectangle(
    xy=(boxes[0]), width=boxes[1][0]-boxes[0][0], height=boxes[1][1]-boxes[0][1],
    fill=False, edgecolor="blue", linewidth=1))
plt.show()

# 4.过滤分数较低的预测
threshold=0.5
# 4.1 因pred_score是从大大小进行排列的,只要获取最后一个索引即可
pred_t = [pred_score.index(x) for x in pred_score if x > threshold][-1]
pred_boxes = pred_boxes[:pred_t + 1]
pred_class = pred_class[:pred_t + 1]

# 5.将过滤后的结构绘制在图像上
img = Image.open('img3.jpg').resize((800,800))
# 5.1 绘制图像
fig = plt.imshow(img)
# 5.2 将框绘制在图像上
for index, boxes in enumerate(pred_boxes):
    fig.axes.add_patch(plt.Rectangle(
    xy=(boxes[0]), width=boxes[1][0]-boxes[0][0], height=boxes[1][1]-boxes[0][1],
    fill=False, edgecolor="g", linewidth=1))
plt.show()

三、模型结构详解(了解)

1、基本结构

  • Backbone:Backbone由CNN卷积神经网络构成,常用的是VGG和resnet, 用来提取图像中的特征,获取图像的特征图。该特征图被共享用于后续RPN层生成候选区域和ROIPooling层中。

  • 一般为VGG,ResNet等网络构成

  • 主要进行特征提取,将最后的全连接层舍弃,得到特征图送入后续网络中进行处理。

  • Resnet进行特征提取,FPN结构作用是当前层的特征图会融合未来层的特征进行上采样,并加以利用

  • 完整流程:Resnet进行特征提取,FPN网络进行特征融合获取多个特征图后,输入到RPN网络中的特征图是[p2,p3,p4,p5,p6] ,而作为后续目标检测网络FastRCNN的输入则是 [p2,p3,p4,p5] 。

  • 代码实现

    ```python

    1.使用backbone获取特征图

    features_body = model.backbone.body(images.tensors)
    C2,C3,C4,C5 = features_body.values()

    2.FPN网络融合:C2,C3,C4,C5是resnet提取的特征结果

    features_fpn = model.backbone.fpn(features_body)
    P2,P3,P4,P5,P6 = features_fpn.values()
    ```

  • RPN网络:RPN网络用于生成候选区域,用于后续的目标检测。

  • 特点:极大提升检测速度

  • 主要流程:

    • 1、生成一系列的固定参考框anchors,覆盖图像的任意位置,然后送入后续网络中进行分类和回归
    • 2、分类分支:通过softmax分类判断anchor中是否包含目标
    • 3、回归分支:计算目标框对于anchors的偏移量,以获得精确的候选区域
    • 4、最后的Proposal层则负责综合含有目标的anchors和对应bbox回归偏移量获取候选区域,同时剔除太小和超出边界的候选区域。
  • anchor:表示固定的参考框,预设一组不同尺度不同长宽比的固定参考框,覆盖几乎所有位置, 每个参考框负责检测与其交并比大于阈值 (训练预设值,常用0.5或0.7) 的目标

  • FastRCNN中anchor设置:尺度为32,64,128,长宽比为1:1,1:2,2:1的一组anchors

  • 代码实现

    ```python

    1.生成anchors,送入图片信息及相应的特征图

    anchors=model.rpn.anchor_generator(images,[P2,P3,P4,P5,P6])

    2.将生成的anchor绘制在图像上

    img = Image.open('img3.jpg').resize((800,800))

    2.1 绘制图像

    fig = plt.imshow(img)

    2.2 将框绘制在图像上

    for index, boxes in enumerate(anchors[0]):
    if index <1000:
    fig.axes.add_patch(plt.Rectangle(
    xy=(boxes[0],boxes[1]), width=boxes[2]-boxes[0], height=boxes[3]-boxes[1],
    fill=False, edgecolor="red", linewidth=1))
    plt.show()
    ```

  • RPN分类和回归

    ```python

    RPN网络预测分类,返回:logits送入softmax之前的分数

    cls_logits = model.rpn.head.cls_logits(P2)

    RPN网络预测目标框

    box_pred = model.rpn.head.bbox_pred(P2)
    ```

  • Proposal层

    • 负责综合RPN网络对anchors分类和回归的结果,利用回归的结果对包含目标的anchors进行修正,计算出候选区域,送入后续RoI Pooling层中。
  • Roi Pooling: 该部分收集图像的特征图和RPN网络提取的候选区域位置,综合信息后获取固定尺寸的特征,送入后续全连接层判定目标类别和确定目标位置。

  • 负责收集RPN网络生成的候选区域,并将其映射到特征图中并固定维度,送入后续网络中进行分类和回归。

  • 使用最大池化将任何有效的RoI区域内的特征转换成具有pool_H×pool_W的固定空间范围的小的特征图,其中pool_H和pool_W是超参数,比如设置为7x7, 它们独立于任何特定的RoI

  • 代码实现

    ```python

    ROI Pooling层实现:

    pool_region_list = model.roi_heads.box_roi_pool(features_fpn, proposal[0],images.image_sizes)
    ```

  • 目标分类与回归: 该部分利用ROIpooling输出特征向量计算候选区域的类别,并通过回归获得检测框最终的精确位置。

  • 利用获得的候选区域的特征图,通过全连接层与softmax计算每个候选区域具体属于的类别(如人,车,电视等),输出概率值;

  • 同时再次利用回归方法获得每个候选区域的位置偏移量,用于回归更加精确的目标检测框。

  • 代码实现:

    ```python

    1.模型最终输出的结果

    detection = model.roi_heads(features, proposal[0],images.image_sizes)

    2.检测结果的绘制

    detection_boxes = detection[0][0]['boxes'].detach().numpy()
    img = Image.open('img3.jpg').resize((800,800))

    2.1 显示图像

    fig = plt.imshow(img)

    2.2 将检测结果绘制在图像上

    for index, boxes in enumerate(detection_boxes):
    fig.axes.add_patch(plt.Rectangle(
    xy=(boxes[0],boxes[1]), width=boxes[2]-boxes[0], height=boxes[3]-boxes[1],
    fill=False, edgecolor="purple", linewidth=1))
    plt.show()
    ```

四、FasterRCNN的训练(了解)

1、基本内容

  • (1)数据集加载

  • (2)模型加载

  • (3)RPN网络的训练

  • 1.正负样本标记

    • 与真实框ground truth(GT)交并比IOU大于0.7的anchor是正样本,即anchor中包含目标,目标值设为1
    • 与真实框ground truth(GT)交并比IOU小于0.3的anchor是负样本,即anchor中不包含目标,目标值设为-1
    • 其他的anchor舍弃,不参与网络的训练,目标值设为0
  • 2.RPN网络的损失函数

    image-20240424185154480

  • (4)FastRCNN网络的训练

  • 训练流程

    • 首先将与真实框ground truth(GT)交并比IOU大于0.5的候选区域设为正样本,类别的目标值是GT的类别1
    • 将与真实框ground truth(GT)交并比IOU小于0.5的候选区域设为负样本,类别的目标值是0
  • FastRCNN的损失函数

    • FastRCNN的输出由两部分组成:一部分是softmax层进行分类,输出类别有K个类别加上”背景”类,另一部分是回归bounding box regressor
    • 一部分输出在K+1个类别上的离散概率分布(每个候选区域)
    • 另一部分输出对于由K个类别中的每一个检测框回归偏移
  • (5)FasterRCNN的训练

  • 训练流程

    • 使用Dataloader加载训练数据
    • 指定优化器:在这里我们使用加动量的SGD方法
    • 设置轮次epoch,遍历每个轮次获取batch数据送入网络中进行预测,计算损失函数,使用反向传播更新参数,完成模型训练

    • loss_dict = model(images,targets)

    • fast RCNN loss
    • RPN loss

课堂笔记

image-20240903110414146

image-20240903124642866

image-20240903161627315

📖 笔记

04-目标检测-麦穗检测案例-FasterRCNN

计算机视觉基础CV笔记

一、今日课程内容

  • 麦穗检测的数据集的构成(了解)
  • 项目解决的全流程(了解)
  • 能够利用fasterRCNN网络模型进行训练和预测(实操 + 掌握)

二、麦穗数据集介绍与数据分析(了解)

1、基本内容

  • 来自Global Wheat Head detection (GWHD) dataset,是第一个从现场光学图像进行小麦头检测的大型数据集

  • 包括多个地区的不同品种,涵盖了由欧洲(法国,英国,瑞士),亚洲(中国,日本),大洋洲(澳大利亚)和北美(加拿大)的标注图像

  • 图像的通用格式为1024 × 1024 px,分辨率为每像素0.1~0.3 mm

2、数据读取

# 数据获取与分析:指定训练集和验证集的路径
train_data_dir = 'wheatData/train'
test_data_dir = 'wheatData/test'
# 读取CSV文件
train_data_label = pd.read_csv('wheatData/train_data.csv')
print(train_data_label.head())
  • 主要内容包括:
  • image_id:图像文件的名称,指当前标注框所在的图像
  • width和height: 图像的宽高,在当前数据集中,所有图片的大小均为1024x1024
  • x和y: 表示当前标注框的左上角坐标
  • w和h: 表示标注框的宽和高。

3、数据分析

  • 数据量的分析
# 训练数据大小(这里指的是目标框)
print("目标框个数:{}".format(train_data_label.shape))
  • 训练集图片的数量
# 有标注信息图片的数量
print("图片数量:{}".format(train_data_label['image_id'].nunique()))
# 所有图片的数量
print("所有图片数量:{}".format(len(os.listdir(train_data_dir))))
  • 图片中目标框数量分布
counts = train_data_label['image_id'].value_counts()
print("单张图片中标注框的最大值{}和最小值{}".format(max(counts), min(counts)))
sns.displot(counts, kde=True, color="g")
plt.xlabel('boxes')
plt.ylabel('images')
plt.show()
  • 标注框位置左上角坐标的分布
# 标注框左上角坐标的分布
sns.histplot(data=train_data_label, x='x', y='y', bins=50, pmax=0.9)
plt.xlabel('x')
plt.ylabel('y')
plt.show()
  • 标注框位置中心点坐标的分布
# 获取数据的中心点坐标
train_data_label['cx'] = train_data_label['x'] + train_data_label['w'] / 2
train_data_label['cy'] = train_data_label['y'] + train_data_label['h'] / 2
# 绘制分布图像
sns.histplot(data=train_data_label, x='cx', y='cy', bins=50, pmax=0.9)
# 设置轴描述
plt.xlabel('cx')
plt.ylabel('cy')
plt.show()
  • 标注框的宽高的分布
# 标注框宽高的分布
sns.histplot(data=train_data_label, x='w', y='h', bins=50, pmax=0.9)
# 轴描述
plt.xlabel('w')
plt.ylabel('h')
plt.show()
  • 标注框面积的分布
# 标注框面积的分布:用来设置anchor的尺度
aeras = train_data_label['w'] * train_data_label['h']
print("目标框最小面积为{}最大面积为{}".format(min(aeras), max(aeras)))
sns.histplot(aeras, bins=50, kde=False)
plt.show()
  • 数据展示
# 图片及标注框预览
# 设置坐标轴的个数
num_rows, num_cols = 1, 2
# 获取标注信息
ids = train_data_label['image_id'].unique()[100:100 + num_rows * num_cols]
# 读取图片
imgs = [plt.imread(f'{train_data_dir}/{n}.jpg') for n in ids]
# 图片显示
axes = show_images(imgs, num_rows, num_cols)
# 显示标注框
for ax, id in zip(axes, ids):
    datas = train_data_label[train_data_label['image_id'] == id]
    bboxes = [(d['x'], d['y'], d['x'] + d['w'], d['y'] + d['h']) for _, d in datas.iterrows()]
    show_bboxes(ax, bboxes, labels=None, colors=['r'])
    plt.show()

三、数据加载DataSet的构建(了解)

1、基本内容

  • 继承 torch.utils.data.Dataset 抽象类,实现__init__ __len__ __getitem__

  • __init__方法进行参数的初始化,包括

  • __len__方法获取数量集中图片的数量

  • __getitem__ 获取数据:

  • image: 图像数据

  • target: 字典数据,包括模型训练时的目标值信息,介绍如下:

    • boxes : 表示标注框的左上角坐标和右下角坐标 [x0, y0, x1, y1]。
    • label : 标志标注框的类别,在当前任务中只有小麦一个类别,所以所有标注框的类别标签值设为1

四、模型训练(掌握)

1、基本内容

# 1.创建模型,使用torchvision中的fasterRCNN模型,使用预训练模型; 训练过程中不修改backbone部分的参数
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True, trainable_backbone_layers=False)
print("预训练模型的输出端{}".format(model.roi_heads.box_predictor))

# 2.当前任务检测一类数据(小麦),fasterRcnn中需要设置为N+1(背景),设置为2
num_classes = 2
# 2.1 获取输出端的特征向量的维度
in_features = model.roi_heads.box_predictor.cls_score.in_features
# 2.2 设置输出端:
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
print("当前任务的输出端{}".format(model.roi_heads.box_predictor))

2、训练参数设置

  • 设置设备信息,优化器,轮次,批量大小等参数
  • 代码实现
# 训练设备的设置,有GPU的使用GPU,否则使用CPU进行训练
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
# 将模型写入到设备中
model.to(device)
# 获取所有要进行训练的参数,设置优化器
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
# 设置训练的轮次
num_epochs = 10
# 批次大小
batch_size = 4

3、数据加载

  • 需要按照批次的大小来获取数据送入网络中进行训练,创建DataLoader来获取数据

  • 代码实现

```python
# 获取数据和标签信息
train_data_dir = 'wheatData/train'
train_data_label = pd.read_csv('wheatData/train_data.csv')
# 实例化数据获取类
train_dataset = Wheat(train_data_label, train_data_dir,transforms=Wheat.get_transform())

def collate_fn(batch):
return tuple(zip(*batch))

# 获取batch的数据,送入网络中进行训练
train_data_loader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
collate_fn=collate_fn
)
```

4、模型训练

  • 基本流程

  • 遍历每个轮次

  • 遍历每个batch的图像
  • 模型前向传播进行预测
  • 损失计算
  • 反向传播更新梯度
  • 模型保存

  • 代码实现

```python
# 迭代次数
itr = 1
# 存放训练损失
total_train_loss = []
# 损失值
losses_value = 0
# 遍历每个轮次进行训练
for epoch in range(num_epochs):
# 计时
start_time = time.time()
# 训练模式
model.train()
train_loss = []
# 进度条
pbar = tqdm(train_data_loader, desc='let\'s train')
# 遍历数据,获取图像,目标值和图像id
for images, targets, image_ids in pbar:
# 将图像写入设备中
images = list(image.to(device) for image in images)
# 将目标值写入设备中:key:value
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
# 将数据送入网络中进行训练,获取损失值
loss_dict = model(images, targets)
print(loss_dict)
# 将损失求和
losses = sum(loss for loss in loss_dict.values())
# 获取loss值
losses_value = losses.item()
train_loss.append(losses_value)
# 进行反向传播,更新参数
optimizer.zero_grad()
losses.backward()
optimizer.step()
# 日志信息
pbar.set_description(f"Epoch: {epoch + 1}, Batch: {itr}, Loss: {losses_value}")
# 迭代次数增1
itr += 1
# 获取当前轮次的损失
epoch_train_loss = np.mean(train_loss)
# 轮次损失写入到列表中
total_train_loss.append(epoch_train_loss)

# 模型保存
torch.save(model.state_dict(), 'fasterrcnn_resnet50_fpn.pth')
```

五、模型测试(掌握)

1、基本内容

  • 加载要预测的图片数据,加载训练好的模型,并利用训练好的模型进行预测

2、数据和模型加载

  • 基本流程

  • 1 数据的加载

    • 设置要预测的图片的路径
    • 实例化测试数据集的读取类
    • 创建dataloader获取送入网络中的数据
  • 2 模型参数加载

    • 指定训练好的模型参数的路径
    • 加载模型结构(与训练阶段是一样的),并加载模型参数
    • 将模型设置为预测模式,并写入到设备中
  • 代码实现

```python
# 1.数据加载
# 设置测试数据集路径
test_data_dir = 'wheatData/test'
# 加载数据
test_data = pd.read_csv('wheatData/submission.csv')
print(test_data.tail())
# 实例化数据读取类
test_data = Wheat(test_data, test_data_dir, phase="test", transforms=Wheat.get_test_transform())
def collate_fn(batch):
return tuple(zip(*batch))
# 设置batch_size
batchsize = 2
# 加载测试集数据
test_data_loader = DataLoader(
# 测试数据
test_data,
# 批次大小
batch_size=batchsize,
# 测试集不需要打乱
shuffle=False,
# 不满batch的数据依然要进行预测
drop_last=False,
collate_fn=collate_fn
)

# 2.模型参数加载
# 指明模型参数的路径
weight_dir = 'weights/fasterrcnn_resnet50_fpn.pth'
# 2.加载训练好的模型
# 设置设备信息
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
# 构建fasterRCNN模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)
# 类别个数:N+1
num_classes = 2
# 获取输出层输入特征向量的维度
in_features = model.roi_heads.box_predictor.cls_score.in_features
# 构建fasterRCNN的输出端
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
print("当前任务的输出端{}".format(model.roi_heads.box_predictor))
# 加载训练好的参数
model.load_state_dict(torch.load(weight_dir, map_location=device))
# 模型预测模式
model.eval()
# 将模型写入到设备中
model.to(device)
```

### 3、模型预测

  • 基本流程

    • 设置预测的置信度阈值
    • 遍历每个batch的图片进行预测
    • 遍历预测结果,使用置信度阈值进行筛选
  • 代码实现

    ```python

    类别概率的阈值

    score_threshold = 0.7

    存放预测结果

    image_outputs = []

    遍历每个批次的测试数据

    for images, image_ids in test_data_loader:
    # 获取每一幅图片并写入到设备中
    images = list(image.to(device) for image in images)
    # 模型预测返回结果
    outputs = model(images)
    # 遍历图像和预测结果
    for image_id, output in zip(image_ids, outputs):
    # 框
    boxes = output['boxes'].data.cpu().numpy()
    # 类别概率值
    scores = output['scores'].data.cpu().numpy()
    # 使用预测进行筛选
    mask = scores >= score_threshold
    # 获取最终的检测框和类别概率值
    boxes = boxes[mask].astype(np.int32)
    scores = scores[mask]
    # 添加到列表中进行保存
    image_outputs.append((image_id, boxes, scores))
    ```

#### 课堂笔记

image-20240904090902773

image-20240904101156283

📖 笔记

05-目标检测-YOLO系列

计算机视觉基础CV笔记

一、今日课程内容

  • yolo网络架构,及其输入输出(了解)
  • yolo模型的训练样本构建的方法(了解)
  • yolo模型的损失函数(了解)
  • yoloV2模型的改进方法
  • yoloV3模型的网络结构及网络输出
  • yoloV3模型正负样本设计和损失函数的计算

二、yolo算法(了解)

1、基本内容

  • (You Only Look Once),采用一个单独的CNN模型实现end-to-end的目标检测
  • 核心思想:利用整张图作为网络的输入,直接在输出层回归 bounding box(边界框) 的位置及其所属的类别
  • 创造性的将候选区和目标分类合二为一,看一眼图片就能知道有哪些对象以及它们的位置。
  • 采用预定义预测区域的方法来完成目标检测,具体而言是将原始图像划分为 7x7=49 个网格(grid),每个网格允许预测出2个边框(bounding box,包含某个对象的矩形框),总共 49x2=98 个bounding box。我们将其理解为98个预测区,很粗略的覆盖了图片的整个区域,就在这98个预测区中进行目标检测。

2、损失函数

image-20240424191852831

3、模型预测

  • 将图片resize成448x448的大小,送入到yolo网络中
  • 输出一个 7x7x30 的张量(tensor)来表示图片中所有网格包含的对象(概率)以及该对象可能的2个位置(bounding box)和可信程度(置信度)。
  • 再采用NMS(Non-maximal suppression,非极大值抑制)算法选出最有可能是目标的结果。

4、总结

优点

  • 速度非常快,处理速度可以达到45fps,其快速版本(网络较小)甚至可以达到155fps。
  • 训练和预测可以端到端的进行,非常简便。

缺点

  • 准确率会打折扣
  • 对于小目标和靠的很近的目标检测效果并不好

三、YOLOv2(了解)

1、基本内容

  • 相对v1版本,在继续保持处理速度的基础上,从预测更准确(Better),速度更快(Faster),识别对象更多(Stronger)这三个方面进行了改进
  • 识别更多对象也就是扩展到能够检测9000种不同对象,称之为YOLO9000
  • 优化点:
  • batch normalization
  • 使用高分辨率图像微调分类模型
  • 采用Anchor Boxes
  • 聚类提取anchor尺度
  • 边框位置的预测
  • 细粒度特征融合
  • 多尺度训练

四、YOLOv3(了解)

1、基本内容

  • 以V1,V2为基础进行的改进,主要有:利用多尺度特征进行目标检测;先验框更丰富;调整了网络结构;对象分类使用logistic代替了softmax,更适用于多标签分类任务。

  • 相比之前的算法,尤其是针对小目标,精度有显著提升。

  • 网络结构
  • image-20240904110835525
  • 整个v3结构里面,没有池化层和全连接层,网络的下采样是通过设置卷积的步长为2来达到的,每当通过这个卷积层之后图像的尺寸就会减小到一半。
  • 基本组件
    • 1、CBL:Yolov3网络结构中的最小组件,由Conv+Bn+Leaky_relu激活函数三者组成。
    • 2、Res unit:借鉴Resnet网络中的残差结构,让网络可以构建的更深。
    • 3、ResX:由一个CBL和X个残差组件构成,是Yolov3中的大组件。每个Res模块前面的CBL都起到下采样的作用,因此经过5次Res模块后,得到的特征图是416->208->104->52->26->13大小。
  • 网络输入:在V3版本中默认是416x416。
  • Backbone:采用了Darknet-53的网络结构(含有53个卷积层),借鉴了残差网络ResNet的做法,在层之间设置了shortcut,来解决深层网络梯度的问题
  • neck部分:引入了FPN网络进行特征融合,并获取多种尺度的特征图,进而完成多尺度检测
  • 输出部分:采用K-means聚类得到先验框的尺寸,为每种尺度设定3种先验框,总共聚类出9种尺寸的先验框。
  • 网络训练
  • 正负样本
    • 正样本:首先计算目标中心点落在哪个grid上,然后计算这个grid对应的3个先验框(anchor)和目标真实位置的IOU值,取IOU值最大的先验框和目标匹配。那么该anchor 就负责预测这个目标,那这个anchor就作为正样本,将其置信度设为1,其他的目标值根据标注信息设置。
    • 负样本:所有不是正样本的anchor都是负样本,将其置信度设为0,参与损失计算,其它的值不参与损失计算,默认为0
  • 损失函数
    • box的损失
    • 置信度的损失
    • 分类的损失
  • 基本流程

    • 根据训练集的标注信息获取每幅图像对应的目标值,即设置正负样本(三个尺度的目标值)
    • 将图像送入网络中进行预测,计算预测结果与目标值之间的差异(损失函数)
    • 反向传播更新参数,直到损失收敛。
  • 模型预测

  • 图片resize成416x416的大小,送入到yoloV3网络中,输出3个尺度的预测结果,每个结果中包含边框坐标(4个数值),边框置信度(1个数值),对象类别的概率(对于COCO数据集,有80种对象)。
  • 筛选出置信度较高的检测结果(>0.5),在采用NMS(Non-maximal suppression,非极大值抑制)算法选出最有可能是目标的结果,就完成了目标的检测
📖 笔记

06-目标检测-YOLOv4算法详解

计算机视觉基础CV笔记

一、今日课程内容

  • yoloV4网络架构(了解)
  • 目标检测中常见的数据增强方法(掌握)
  • yoloV4的损失函数的计算方法(掌握)

二、算法简介(了解)

1、基本内容

  • YOLO之父在2020年初宣布退出CV界,YOLOv4 的作者并不是YOLO系列 的原作者
  • 得到了Joseph Redmon的官方认可,被认为是当时最强的实时对象检测模型之一
  • 作者总结了目标检测的各种技巧,然后经过筛选,排列组合,挨个实验(ablation study)哪些方法有效,提升了yolo在COCO数据集上的检测效果。
  • 在COCO数据集上,yoloV4的检测速度和检测精度都有较大的提升。

三、YOLOv4的网络结构(了解)

1、基本内容

  • 整体架构与YOLOv3相似
  • image-20240904112240115
  • image-20240904113220546

  • 基本组件:

  • CBM:Yolov4网络结构中的最小组件,由Conv+Bn+Mish激活函数三者组成。
  • CBL:由Conv+Bn+Leaky_relu激活函数三者组成。
  • Res unit:借鉴Resnet网络中的残差结构,让网络可以构建的更深。
  • CSPX:由三个卷积层和X个Res unint模块Concate组成。
  • SPP:采用1×1,5×5,9×9,13×13的最大池化的方式,进行多尺度融合。

  • 其他基础操作:

  • Concat:张量拼接,通道维度会扩充。
  • Add:张量相加,不会扩充维度。
  • Backbone中卷积层的数量: 每个CSPX中包含3+2×X个卷积层,因此整个主干网络Backbone中一共包含2+(3+2×1)+2+(3+2×2)+2+(3+2×8)+2+(3+2×8)+2+(3+2×4)+1=72。

  • 基本架构

  • 输入(input)
  • backbone
  • neck
  • image-20240904113230550
  • 输出(output)

四、模型训练(了解)

1、基本内容

  • 正负样本的设置方式,损失函数的构成及训练过程与YOLOv3一致
  • 唯一不同是在计算回归值得损失时,V4中采用了IOU的损失

2、IOU损失

  • GIOU损失
  • image-20240904114134679
  • DIOU损失
  • image-20240904114140307
  • CIOU损失
  • image-20240904114147488
  • .........

五、模型预测

  • 将图片resize成608x608【32倍数】的大小,送入到yoloV4网络中,输出3个尺度的预测结果,每个结果中包含边框坐标(4个数值),边框置信度(1个数值),对象类别的概率(对于COCO数据集,有80种对象)。
  • 筛选出置信度较高的检测结果(>0.5),在采用DIOU-NMS算法选出最有可能是目标的结果,就完成了目标的检测。
📖 笔记

07-目标检测-YOLOv5算法详解及案例

计算机视觉基础CV笔记

一、今日课程内容

  • YOLOv5网络架构(了解)
  • YOLOv5中使用的策略(掌握)
  • 能够使用YOLOv5框架进行目标检测(掌握)

二、YOLOv5模型介绍(了解)

1、基本内容

  • 基于pytorch实现
  • 包含多种类型的网络结构,按模型大小从小到大是yoloV5n,yoloV5s,yoloV5m,yoloV5l,yoloV5x

2、网络架构

  • 基本组件:
  • CBL:由Conv+Bn+Leaky_relu激活函数三者组成。
  • Res unit:借鉴Resnet网络中的残差结构,让网络可以构建的更深。
  • CSP1_X:由三个卷积层和X个Res unint模块Concate组成,模型实现使用C3来完成,将shortcut设为True即可
  • CSP2_X:这个在V4中没有使用,由三个卷积层和2X个Res unint模块Concate组成,模型实现使用C3来完成,将shortcut设为False即可
  • SPPF:V4中使用的SPP进行融合,在这里采用1x1卷积和3个串联的5x5的最大池化的方式,进行多尺度融合。

  • 其他基础操作:

  • Concat:张量拼接,通道维度会扩充。
  • Add:张量相加,不会扩充维度。

  • 基础结构

  • 输入(input)
  • backbone
  • neck
  • 输出(output)

3、模型训练

  • anchor的自动调整
  • 正负样本的设计

4、模型预测

  • 将图片resize后,送入到yoloV5网络中,输出3个尺度的预测结果,每个结果中包含边框坐标(4个数值),边框置信度(1个数值),对象类别的概率(对于COCO数据集,有80种对象)。
  • 筛选出置信度较高的检测结果(>0.5),在采用NMS算法选出最有可能是目标的结果,就完成了目标的检测。

三、案例实现(掌握)

1、数据构建

  • 数据标注:labelImg
  • 数据加载:coco128.yaml

2、模型选择(构建)

  • yolov5-master/models文件夹中yaml文件

3、模型训练

4、模型预测

image-20240904175220766

💻 代码

目标检测代码(IoU/NMS/FasterRCNN/YOLO)

154 个文件
  • 📁 yolov5-master/
    • 📁 .github/
      • 📁 ISSUE_TEMPLATE/
        • 📄 bug-report.yml 2.9KB
          name: 🐛 Bug Report
          # title: " "
          description: Problems with YOLOv5
          labels: [bug, triage]
          body:
            - type: markdown
              attributes:
                value: |
                  Thank you for submitting a YOLOv5 🐛 Bug Report!
          
            - type: checkboxes
              attributes:
                label: Search before asking
                description: >
                  Please search the [issues](https://github.com/ultralytics/yolov5/issues) to see if a similar bug report already exists.
                options:
                  - label: >
                      I have searched the YOLOv5 [issues](https://github.com/ultralytics/yolov5/issues) and found no similar bug report.
                    required: true
          
            - type: dropdown
              attributes:
                label: YOLOv5 Component
                description: |
                  Please select the part of YOLOv5 where you found the bug.
                multiple: true
                options:
                  - "Training"
                  - "Validation"
                  - "Detection"
                  - "Export"
                  - "PyTorch Hub"
                  - "Multi-GPU"
                  - "Evolution"
                  - "Integrations"
                  - "Other"
              validations:
                required: false
          
            - type: textarea
              attributes:
                label: Bug
                description: Provide console output with error messages and/or screenshots of the bug.
                placeholder: |
                  💡 ProTip! Include as much information as possible (screenshots, logs, tracebacks etc.) to receive the most helpful response.
              validations:
                required: true
          
            - type: textarea
              attributes:
                label: Environment
                description: Please specify the software and hardware you used to produce the bug.
                placeholder: |
                  - YOLO: YOLOv5 🚀 v6.0-67-g60e42e1 torch 1.9.0+cu111 CUDA:0 (A100-SXM4-40GB, 40536MiB)
                  - OS: Ubuntu 20.04
                  - Python: 3.9.0
              validations:
                required: false
          
            - type: textarea
              attributes:
                label: Minimal Reproducible Example
                description: >
                  When asking a question, people will be better able to provide help if you provide code that they can easily understand and use to **reproduce** the problem.
                  This is referred to by community members as creating a [minimal reproducible example](https://docs.ultralytics.com/help/minimum_reproducible_example/).
                placeholder: |
                  ```
                  # Code to reproduce your issue here
                  ```
              validations:
                required: false
          
            - type: textarea
              attributes:
                label: Additional
                description: Anything else you would like to share?
          
            - type: checkboxes
              attributes:
                label: Are you willing to submit a PR?
                description: >
                  (Optional) We encourage you to submit a [Pull Request](https://github.com/ultralytics/yolov5/pulls) (PR) to help improve YOLOv5 for everyone, especially if you have a good understanding of how to implement a fix or feature.
                  See the YOLOv5 [Contributing Guide](https://docs.ultralytics.com/help/contributing) to get started.
                options:
                  - label: Yes I'd like to help by submitting a PR!
          
        • 📄 config.yml 360B
          blank_issues_enabled: true
          contact_links:
            - name: 📄 Docs
              url: https://docs.ultralytics.com/yolov5
              about: View Ultralytics YOLOv5 Docs
            - name: 💬 Forum
              url: https://community.ultralytics.com/
              about: Ask on Ultralytics Community Forum
            - name: 🎧 Discord
              url: https://ultralytics.com/discord
              about: Ask on Ultralytics Discord
          
        • 📄 feature-request.yml 1.7KB
          name: 🚀 Feature Request
          description: Suggest a YOLOv5 idea
          # title: " "
          labels: [enhancement]
          body:
            - type: markdown
              attributes:
                value: |
                  Thank you for submitting a YOLOv5 🚀 Feature Request!
          
            - type: checkboxes
              attributes:
                label: Search before asking
                description: >
                  Please search the [issues](https://github.com/ultralytics/yolov5/issues) to see if a similar feature request already exists.
                options:
                  - label: >
                      I have searched the YOLOv5 [issues](https://github.com/ultralytics/yolov5/issues) and found no similar feature requests.
                    required: true
          
            - type: textarea
              attributes:
                label: Description
                description: A short description of your feature.
                placeholder: |
                  What new feature would you like to see in YOLOv5?
              validations:
                required: true
          
            - type: textarea
              attributes:
                label: Use case
                description: |
                  Describe the use case of your feature request. It will help us understand and prioritize the feature request.
                placeholder: |
                  How would this feature be used, and who would use it?
          
            - type: textarea
              attributes:
                label: Additional
                description: Anything else you would like to share?
          
            - type: checkboxes
              attributes:
                label: Are you willing to submit a PR?
                description: >
                  (Optional) We encourage you to submit a [Pull Request](https://github.com/ultralytics/yolov5/pulls) (PR) to help improve YOLOv5 for everyone, especially if you have a good understanding of how to implement a fix or feature.
                  See the YOLOv5 [Contributing Guide](https://docs.ultralytics.com/help/contributing) to get started.
                options:
                  - label: Yes I'd like to help by submitting a PR!
          
        • 📄 question.yml 1.1KB
          name: ❓ Question
          description: Ask a YOLOv5 question
          # title: " "
          labels: [question]
          body:
            - type: markdown
              attributes:
                value: |
                  Thank you for asking a YOLOv5 ❓ Question!
          
            - type: checkboxes
              attributes:
                label: Search before asking
                description: >
                  Please search the [issues](https://github.com/ultralytics/yolov5/issues) and [discussions](https://github.com/ultralytics/yolov5/discussions) to see if a similar question already exists.
                options:
                  - label: >
                      I have searched the YOLOv5 [issues](https://github.com/ultralytics/yolov5/issues) and [discussions](https://github.com/ultralytics/yolov5/discussions) and found no similar questions.
                    required: true
          
            - type: textarea
              attributes:
                label: Question
                description: What is your question?
                placeholder: |
                  💡 ProTip! Include as much information as possible (screenshots, logs, tracebacks etc.) to receive the most helpful response.
              validations:
                required: true
          
            - type: textarea
              attributes:
                label: Additional
                description: Anything else you would like to share?
          
      • 📁 workflows/
        • 📄 ci-testing.yml 6.8KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # YOLOv5 Continuous Integration (CI) GitHub Actions tests
          
          name: YOLOv5 CI
          
          on:
            push:
              branches: [master]
            pull_request:
              branches: [master]
            schedule:
              - cron: "0 0 * * *" # runs at 00:00 UTC every day
          
          jobs:
            Benchmarks:
              runs-on: ${{ matrix.os }}
              strategy:
                fail-fast: false
                matrix:
                  os: [ubuntu-latest]
                  python-version: ["3.11"] # requires python<=3.11
                  model: [yolov5n]
              steps:
                - uses: actions/checkout@v4
                - uses: actions/setup-python@v5
                  with:
                    python-version: ${{ matrix.python-version }}
                    cache: "pip" # cache pip dependencies
                - name: Install requirements
                  run: |
                    python -m pip install --upgrade pip wheel
                    pip install -r requirements.txt coremltools openvino-dev "tensorflow-cpu<2.15.1" --extra-index-url https://download.pytorch.org/whl/cpu
                    yolo checks
                    pip list
                - name: Benchmark DetectionModel
                  run: |
                    python benchmarks.py --data coco128.yaml --weights ${{ matrix.model }}.pt --img 320 --hard-fail 0.29
                - name: Benchmark SegmentationModel
                  run: |
                    python benchmarks.py --data coco128-seg.yaml --weights ${{ matrix.model }}-seg.pt --img 320 --hard-fail 0.22
                - name: Test predictions
                  run: |
                    python export.py --weights ${{ matrix.model }}-cls.pt --include onnx --img 224
                    python detect.py --weights ${{ matrix.model }}.onnx --img 320
                    python segment/predict.py --weights ${{ matrix.model }}-seg.onnx --img 320
                    python classify/predict.py --weights ${{ matrix.model }}-cls.onnx --img 224
          
            Tests:
              timeout-minutes: 60
              runs-on: ${{ matrix.os }}
              strategy:
                fail-fast: false
                matrix:
                  os: [ubuntu-latest, windows-latest, macos-14] # macos-latest bug https://github.com/ultralytics/yolov5/pull/9049
                  python-version: ["3.11"]
                  model: [yolov5n]
                  include:
                    - os: ubuntu-latest
                      python-version: "3.8" # torch 1.8.0 requires python >=3.6, <=3.8
                      model: yolov5n
                      torch: "1.8.0" # min torch version CI https://pypi.org/project/torchvision/
              steps:
                - uses: actions/checkout@v4
                - uses: actions/setup-python@v5
                  with:
                    python-version: ${{ matrix.python-version }}
                    cache: "pip" # caching pip dependencies
                - name: Install requirements
                  run: |
                    python -m pip install --upgrade pip wheel
                    torch=""
                    if [ "${{ matrix.torch }}" == "1.8.0" ]; then
                      torch="torch==1.8.0 torchvision==0.9.0"
                    fi
                    pip install -r requirements.txt $torch --extra-index-url https://download.pytorch.org/whl/cpu
                  shell: bash # for Windows compatibility
                - name: Check environment
                  run: |
                    yolo checks
                    pip list
                - name: Test detection
                  shell: bash # for Windows compatibility
                  run: |
                    # export PYTHONPATH="$PWD"  # to run '$ python *.py' files in subdirectories
                    m=${{ matrix.model }}  # official weights
                    b=runs/train/exp/weights/best  # best.pt checkpoint
                    python train.py --imgsz 64 --batch 32 --weights $m.pt --cfg $m.yaml --epochs 1 --device cpu  # train
                    for d in cpu; do  # devices
                      for w in $m $b; do  # weights
                        python val.py --imgsz 64 --batch 32 --weights $w.pt --device $d  # val
                        python detect.py --imgsz 64 --weights $w.pt --device $d  # detect
                      done
                    done
                    python hubconf.py --model $m  # hub
                    # python models/tf.py --weights $m.pt  # build TF model
                    python models/yolo.py --cfg $m.yaml  # build PyTorch model
                    python export.py --weights $m.pt --img 64 --include torchscript  # export
                    python - <<EOF
                    import torch
                    im = torch.zeros([1, 3, 64, 64])
                    for path in '$m', '$b':
                        model = torch.hub.load('.', 'custom', path=path, source='local')
                        print(model('data/images/bus.jpg'))
                        model(im)  # warmup, build grids for trace
                        torch.jit.trace(model, [im])
                    EOF
                - name: Test segmentation
                  shell: bash # for Windows compatibility
                  run: |
                    m=${{ matrix.model }}-seg  # official weights
                    b=runs/train-seg/exp/weights/best  # best.pt checkpoint
                    python segment/train.py --imgsz 64 --batch 32 --weights $m.pt --cfg $m.yaml --epochs 1 --device cpu  # train
                    python segment/train.py --imgsz 64 --batch 32 --weights '' --cfg $m.yaml --epochs 1 --device cpu  # train
                    for d in cpu; do  # devices
                      for w in $m $b; do  # weights
                        python segment/val.py --imgsz 64 --batch 32 --weights $w.pt --device $d  # val
                        python segment/predict.py --imgsz 64 --weights $w.pt --device $d  # predict
                        python export.py --weights $w.pt --img 64 --include torchscript --device $d  # export
                      done
                    done
                - name: Test classification
                  shell: bash # for Windows compatibility
                  run: |
                    m=${{ matrix.model }}-cls.pt  # official weights
                    b=runs/train-cls/exp/weights/best.pt  # best.pt checkpoint
                    python classify/train.py --imgsz 32 --model $m --data mnist160 --epochs 1  # train
                    python classify/val.py --imgsz 32 --weights $b --data ../datasets/mnist160  # val
                    python classify/predict.py --imgsz 32 --weights $b --source ../datasets/mnist160/test/7/60.png  # predict
                    python classify/predict.py --imgsz 32 --weights $m --source data/images/bus.jpg  # predict
                    python export.py --weights $b --img 64 --include torchscript  # export
                    python - <<EOF
                    import torch
                    for path in '$m', '$b':
                        model = torch.hub.load('.', 'custom', path=path, source='local')
                    EOF
          
            Summary:
              runs-on: ubuntu-latest
              needs: [Benchmarks, Tests] # Add job names that you want to check for failure
              if: always() # This ensures the job runs even if previous jobs fail
              steps:
                - name: Check for failure and notify
                  if: (needs.Benchmarks.result == 'failure' || needs.Tests.result == 'failure' || needs.Benchmarks.result == 'cancelled' || needs.Tests.result == 'cancelled') && github.repository == 'ultralytics/yolov5' && (github.event_name == 'schedule' || github.event_name == 'push')
                  uses: slackapi/slack-github-action@v1.26.0
                  with:
                    payload: |
                      {"text": "<!channel> GitHub Actions error for ${{ github.workflow }} ❌\n\n\n*Repository:* https://github.com/${{ github.repository }}\n*Action:* https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}\n*Author:* ${{ github.actor }}\n*Event:* ${{ github.event_name }}\n"}
                  env:
                    SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL_YOLO }}
          
        • 📄 cla.yml 1.6KB
          # Ultralytics YOLO 🚀, AGPL-3.0 license
          # Ultralytics Contributor License Agreement (CLA) action https://docs.ultralytics.com/help/CLA
          # This workflow automatically requests Pull Requests (PR) authors to sign the Ultralytics CLA before PRs can be merged
          
          name: CLA Assistant
          on:
            issue_comment:
              types:
                - created
            pull_request_target:
              types:
                - reopened
                - opened
                - synchronize
          
          jobs:
            CLA:
              if: github.repository == 'ultralytics/yolov5'
              runs-on: ubuntu-latest
              steps:
                - name: CLA Assistant
                  if: (github.event.comment.body == 'recheck' || github.event.comment.body == 'I have read the CLA Document and I sign the CLA') || github.event_name == 'pull_request_target'
                  uses: contributor-assistant/github-action@v2.4.0
                  env:
                    GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
                    # must be repository secret token
                    PERSONAL_ACCESS_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }}
                  with:
                    path-to-signatures: "signatures/version1/cla.json"
                    path-to-document: "https://docs.ultralytics.com/help/CLA" # CLA document
                    # branch should not be protected
                    branch: "main"
                    allowlist: dependabot[bot],github-actions,[pre-commit*,pre-commit*,bot*
          
                    remote-organization-name: ultralytics
                    remote-repository-name: cla
                    custom-pr-sign-comment: "I have read the CLA Document and I sign the CLA"
                    custom-allsigned-prcomment: All Contributors have signed the CLA. ✅
                    #custom-notsigned-prcomment: 'pull request comment with Introductory message to ask new contributors to sign'
          
        • 📄 codeql-analysis.yml 2.0KB
          # This action runs GitHub's industry-leading static analysis engine, CodeQL, against a repository's source code to find security vulnerabilities.
          # https://github.com/github/codeql-action
          
          name: "CodeQL"
          
          on:
            schedule:
              - cron: "0 0 1 * *" # Runs at 00:00 UTC on the 1st of every month
            workflow_dispatch:
          
          jobs:
            analyze:
              name: Analyze
              runs-on: ubuntu-latest
          
              strategy:
                fail-fast: false
                matrix:
                  language: ["python"]
                  # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python' ]
                  # Learn more:
                  # https://docs.github.com/en/free-pro-team@latest/github/finding-security-vulnerabilities-and-errors-in-your-code/configuring-code-scanning#changing-the-languages-that-are-analyzed
          
              steps:
                - name: Checkout repository
                  uses: actions/checkout@v4
          
                # Initializes the CodeQL tools for scanning.
                - name: Initialize CodeQL
                  uses: github/codeql-action/init@v3
                  with:
                    languages: ${{ matrix.language }}
                    # If you wish to specify custom queries, you can do so here or in a config file.
                    # By default, queries listed here will override any specified in a config file.
                    # Prefix the list here with "+" to use these queries and those in the config file.
                    # queries: ./path/to/local/query, your-org/your-repo/queries@main
          
                # Autobuild attempts to build any compiled languages  (C/C++, C#, or Java).
                # If this step fails, then you should remove it and run the build manually (see below)
                - name: Autobuild
                  uses: github/codeql-action/autobuild@v3
          
                # ℹ️ Command-line programs to run using the OS shell.
                # 📚 https://git.io/JvXDl
          
                # ✏️ If the Autobuild fails above, remove it and uncomment the following three lines
                #    and modify them (or add more) to build your code if your project
                #    uses a compiled language
          
                #- run: |
                #   make bootstrap
                #   make release
          
                - name: Perform CodeQL Analysis
                  uses: github/codeql-action/analyze@v3
          
        • 📄 docker.yml 1.7KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Builds ultralytics/yolov5:latest images on DockerHub https://hub.docker.com/r/ultralytics/yolov5
          
          name: Publish Docker Images
          
          on:
            push:
              branches: [master]
            workflow_dispatch:
          
          jobs:
            docker:
              if: github.repository == 'ultralytics/yolov5'
              name: Push Docker image to Docker Hub
              runs-on: ubuntu-latest
              steps:
                - name: Checkout repo
                  uses: actions/checkout@v4
                  with:
                    fetch-depth: 0 # copy full .git directory to access full git history in Docker images
          
                - name: Set up QEMU
                  uses: docker/setup-qemu-action@v3
          
                - name: Set up Docker Buildx
                  uses: docker/setup-buildx-action@v3
          
                - name: Login to Docker Hub
                  uses: docker/login-action@v3
                  with:
                    username: ${{ secrets.DOCKERHUB_USERNAME }}
                    password: ${{ secrets.DOCKERHUB_TOKEN }}
          
                - name: Build and push arm64 image
                  uses: docker/build-push-action@v5
                  continue-on-error: true
                  with:
                    context: .
                    platforms: linux/arm64
                    file: utils/docker/Dockerfile-arm64
                    push: true
                    tags: ultralytics/yolov5:latest-arm64
          
                - name: Build and push CPU image
                  uses: docker/build-push-action@v5
                  continue-on-error: true
                  with:
                    context: .
                    file: utils/docker/Dockerfile-cpu
                    push: true
                    tags: ultralytics/yolov5:latest-cpu
          
                - name: Build and push GPU image
                  uses: docker/build-push-action@v5
                  continue-on-error: true
                  with:
                    context: .
                    file: utils/docker/Dockerfile
                    push: true
                    tags: ultralytics/yolov5:latest
          
        • 📄 format.yml 1.0KB
          # Ultralytics 🚀 - AGPL-3.0 license
          # Ultralytics Actions https://github.com/ultralytics/actions
          # This workflow automatically formats code and documentation in PRs to official Ultralytics standards
          
          name: Ultralytics Actions
          
          on:
            push:
              branches: [main, master]
            pull_request_target:
              branches: [main, master]
          
          jobs:
            format:
              runs-on: ubuntu-latest
              steps:
                - name: Run Ultralytics Formatting
                  uses: ultralytics/actions@main
                  with:
                    token: ${{ secrets.GITHUB_TOKEN }} # automatically generated, do not modify
                    python: true # format Python code and docstrings
                    markdown: true # format Markdown and YAML
                    spelling: true # check spelling
                    links: false # check broken links
                    summary: true # print PR summary with GPT4 (requires 'openai_api_key' or 'openai_azure_api_key' and 'openai_azure_endpoint')
                    openai_azure_api_key: ${{ secrets.OPENAI_AZURE_API_KEY }}
                    openai_azure_endpoint: ${{ secrets.OPENAI_AZURE_ENDPOINT }}
          
        • 📄 greetings.yml 5.4KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          name: Greetings
          
          on:
            pull_request_target:
              types: [opened]
            issues:
              types: [opened]
          
          jobs:
            greeting:
              runs-on: ubuntu-latest
              steps:
                - uses: actions/first-interaction@v1
                  with:
                    repo-token: ${{ secrets.GITHUB_TOKEN }}
                    pr-message: |
                      👋 Hello @${{ github.actor }}, thank you for submitting a YOLOv5 🚀 PR! To allow your work to be integrated as seamlessly as possible, we advise you to:
          
                      - ✅ Verify your PR is **up-to-date** with `ultralytics/yolov5` `master` branch. If your PR is behind you can update your code by clicking the 'Update branch' button or by running `git pull` and `git merge master` locally.
                      - ✅ Verify all YOLOv5 Continuous Integration (CI) **checks are passing**.
                      - ✅ Reduce changes to the absolute **minimum** required for your bug fix or feature addition. _"It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is."_  — Bruce Lee
          
                    issue-message: |
                      👋 Hello @${{ github.actor }}, thank you for your interest in YOLOv5 🚀! Please visit our ⭐️ [Tutorials](https://docs.ultralytics.com/yolov5/) to get started, where you can find quickstart guides for simple tasks like [Custom Data Training](https://docs.ultralytics.com/yolov5/tutorials/train_custom_data/) all the way to advanced concepts like [Hyperparameter Evolution](https://docs.ultralytics.com/yolov5/tutorials/hyperparameter_evolution/).
          
                      If this is a 🐛 Bug Report, please provide a **minimum reproducible example** to help us debug it.
          
                      If this is a custom training ❓ Question, please provide as much information as possible, including dataset image examples and training logs, and verify you are following our [Tips for Best Training Results](https://docs.ultralytics.com/yolov5/tutorials/tips_for_best_training_results/).
          
                      ## Requirements
          
                      [**Python>=3.8.0**](https://www.python.org/) with all [requirements.txt](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) installed including [**PyTorch>=1.8**](https://pytorch.org/get-started/locally/). To get started:
                      ```bash
                      git clone https://github.com/ultralytics/yolov5  # clone
                      cd yolov5
                      pip install -r requirements.txt  # install
                      ```
          
                      ## Environments
          
                      YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including [CUDA](https://developer.nvidia.com/cuda)/[CUDNN](https://developer.nvidia.com/cudnn), [Python](https://www.python.org/) and [PyTorch](https://pytorch.org/) preinstalled):
          
                      - **Notebooks** with free GPU: <a href="https://bit.ly/yolov5-paperspace-notebook"><img src="https://assets.paperspace.io/img/gradient-badge.svg" alt="Run on Gradient"></a> <a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> <a href="https://www.kaggle.com/ultralytics/yolov5"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
                      - **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://docs.ultralytics.com/yolov5/environments/google_cloud_quickstart_tutorial/)
                      - **Amazon** Deep Learning AMI. See [AWS Quickstart Guide](https://docs.ultralytics.com/yolov5/environments/aws_quickstart_tutorial/)
                      - **Docker Image**. See [Docker Quickstart Guide](https://docs.ultralytics.com/yolov5/environments/docker_image_quickstart_tutorial/) <a href="https://hub.docker.com/r/ultralytics/yolov5"><img src="https://img.shields.io/docker/pulls/ultralytics/yolov5?logo=docker" alt="Docker Pulls"></a>
          
                      ## Status
          
                      <a href="https://github.com/ultralytics/yolov5/actions/workflows/ci-testing.yml"><img src="https://github.com/ultralytics/yolov5/actions/workflows/ci-testing.yml/badge.svg" alt="YOLOv5 CI"></a>
          
                      If this badge is green, all [YOLOv5 GitHub Actions](https://github.com/ultralytics/yolov5/actions) Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 [training](https://github.com/ultralytics/yolov5/blob/master/train.py), [validation](https://github.com/ultralytics/yolov5/blob/master/val.py), [inference](https://github.com/ultralytics/yolov5/blob/master/detect.py), [export](https://github.com/ultralytics/yolov5/blob/master/export.py) and [benchmarks](https://github.com/ultralytics/yolov5/blob/master/benchmarks.py) on macOS, Windows, and Ubuntu every 24 hours and on every commit.
          
                      ## Introducing YOLOv8 🚀
          
                      We're excited to announce the launch of our latest state-of-the-art (SOTA) object detection model for 2023 - [YOLOv8](https://github.com/ultralytics/ultralytics) 🚀!
          
                      Designed to be fast, accurate, and easy to use, YOLOv8 is an ideal choice for a wide range of object detection, image segmentation and image classification tasks. With YOLOv8, you'll be able to quickly and accurately detect objects in real-time, streamline your workflows, and achieve new levels of accuracy in your projects.
          
                      Check out our [YOLOv8 Docs](https://docs.ultralytics.com/) for details and get started with:
                      ```bash
                      pip install ultralytics
                      ```
          
        • 📄 links.yml 2.4KB
          # Ultralytics YOLO 🚀, AGPL-3.0 license
          # Continuous Integration (CI) GitHub Actions tests broken link checker using https://github.com/lycheeverse/lychee
          # Ignores the following status codes to reduce false positives:
          #   - 403(OpenVINO, 'forbidden')
          #   - 429(Instagram, 'too many requests')
          #   - 500(Zenodo, 'cached')
          #   - 502(Zenodo, 'bad gateway')
          #   - 999(LinkedIn, 'unknown status code')
          
          name: Check Broken links
          
          on:
            workflow_dispatch:
            schedule:
              - cron: "0 0 * * *" # runs at 00:00 UTC every day
          
          jobs:
            Links:
              runs-on: ubuntu-latest
              steps:
                - uses: actions/checkout@v4
          
                - name: Download and install lychee
                  run: |
                    LYCHEE_URL=$(curl -s https://api.github.com/repos/lycheeverse/lychee/releases/latest | grep "browser_download_url" | grep "x86_64-unknown-linux-gnu.tar.gz" | cut -d '"' -f 4)
                    curl -L $LYCHEE_URL -o lychee.tar.gz
                    tar xzf lychee.tar.gz
                    sudo mv lychee /usr/local/bin
          
                - name: Test Markdown and HTML links with retry
                  uses: nick-invision/retry@v3
                  with:
                    timeout_minutes: 5
                    retry_wait_seconds: 60
                    max_attempts: 3
                    command: |
                      lychee \
                      --scheme 'https' \
                      --timeout 60 \
                      --insecure \
                      --accept 403,429,500,502,999 \
                      --exclude-all-private \
                      --exclude 'https?://(www\.)?(linkedin\.com|twitter\.com|instagram\.com|kaggle\.com|fonts\.gstatic\.com|url\.com)' \
                      --exclude-path '**/ci.yaml' \
                      --github-token ${{ secrets.GITHUB_TOKEN }} \
                      './**/*.md' \
                      './**/*.html'
          
                - name: Test Markdown, HTML, YAML, Python and Notebook links with retry
                  if: github.event_name == 'workflow_dispatch'
                  uses: nick-invision/retry@v3
                  with:
                    timeout_minutes: 5
                    retry_wait_seconds: 60
                    max_attempts: 3
                    command: |
                      lychee \
                      --scheme 'https' \
                      --timeout 60 \
                      --insecure \
                      --accept 429,999 \
                      --exclude-all-private \
                      --exclude 'https?://(www\.)?(linkedin\.com|twitter\.com|instagram\.com|kaggle\.com|fonts\.gstatic\.com|url\.com)' \
                      --exclude-path '**/ci.yaml' \
                      --github-token ${{ secrets.GITHUB_TOKEN }} \
                      './**/*.md' \
                      './**/*.html' \
                      './**/*.yml' \
                      './**/*.yaml' \
                      './**/*.py' \
                      './**/*.ipynb'
          
        • 📄 merge-main-into-prs.yml 1.7KB
          # Ultralytics YOLO 🚀, AGPL-3.0 license
          # Automatically merges repository 'main' branch into all open PRs to keep them up-to-date
          # Action runs on updates to main branch so when one PR merges to main all others update
          
          name: Merge main into PRs
          
          on:
            workflow_dispatch:
            push:
              branches:
                - main
                - master
            
          jobs:
            Merge:
              if: github.repository == 'ultralytics/yolov5'
              runs-on: ubuntu-latest
              steps:
              - name: Checkout repository
                uses: actions/checkout@v4
                with:
                  fetch-depth: 0
              - uses: actions/setup-python@v5
                with:
                  python-version: "3.11"
                  cache: "pip" # caching pip dependencies
              - name: Install requirements
                run: |
                  pip install pygithub
              - name: Merge main into PRs
                shell: python
                run: |
                  from github import Github
                  import os
                  
                  # Authenticate with the GitHub Token
                  g = Github(os.getenv('GITHUB_TOKEN'))
                  
                  # Get the repository dynamically
                  repo = g.get_repo(os.getenv('GITHUB_REPOSITORY'))
                  
                  # List all open pull requests
                  open_pulls = repo.get_pulls(state='open', sort='created')
                  
                  for pr in open_pulls:
                      # Compare PR head with main to see if it's behind
                      try:
                          # Merge main into the PR branch
                          success = pr.update_branch()
                          assert success, "Branch update failed"
                          print(f"Merged 'master' into PR #{pr.number} ({pr.head.ref}) successfully.")
                      except Exception as e:
                          print(f"Could not merge 'master' into PR #{pr.number} ({pr.head.ref}): {e}")
                env:
                  GITHUB_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }}
                  GITHUB_REPOSITORY: ${{ github.repository }}
          
        • 📄 stale.yml 2.3KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          name: Close stale issues
          on:
            schedule:
              - cron: "0 0 * * *" # Runs at 00:00 UTC every day
          
          jobs:
            stale:
              runs-on: ubuntu-latest
              steps:
                - uses: actions/stale@v9
                  with:
                    repo-token: ${{ secrets.GITHUB_TOKEN }}
          
                    stale-issue-message: |
                      👋 Hello there! We wanted to give you a friendly reminder that this issue has not had any recent activity and may be closed soon, but don't worry - you can always reopen it if needed. If you still have any questions or concerns, please feel free to let us know how we can help.
          
                      For additional resources and information, please see the links below:
          
                      - **Docs**: https://docs.ultralytics.com
                      - **HUB**: https://hub.ultralytics.com
                      - **Community**: https://community.ultralytics.com
          
                      Feel free to inform us of any other **issues** you discover or **feature requests** that come to mind in the future. Pull Requests (PRs) are also always welcomed!
          
                      Thank you for your contributions to YOLO 🚀 and Vision AI ⭐
          
                    stale-pr-message: |
                      👋 Hello there! We wanted to let you know that we've decided to close this pull request due to inactivity. We appreciate the effort you put into contributing to our project, but unfortunately, not all contributions are suitable or aligned with our product roadmap.
          
                      We hope you understand our decision, and please don't let it discourage you from contributing to open source projects in the future. We value all of our community members and their contributions, and we encourage you to keep exploring new projects and ways to get involved.
          
                      For additional resources and information, please see the links below:
          
                      - **Docs**: https://docs.ultralytics.com
                      - **HUB**: https://hub.ultralytics.com
                      - **Community**: https://community.ultralytics.com
          
                      Thank you for your contributions to YOLO 🚀 and Vision AI ⭐
          
                    days-before-issue-stale: 30
                    days-before-issue-close: 10
                    days-before-pr-stale: 90
                    days-before-pr-close: 30
                    exempt-issue-labels: "documentation,tutorial,TODO"
                    operations-per-run: 300 # The maximum number of operations per run, used to control rate limiting.
          
      • 📄 dependabot.yml 647B
        # Ultralytics YOLO 🚀, AGPL-3.0 license
        # Dependabot for package version updates
        # https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates
        
        version: 2
        updates:
          - package-ecosystem: pip
            directory: "/"
            schedule:
              interval: weekly
              time: "04:00"
            open-pull-requests-limit: 10
            reviewers:
              - glenn-jocher
            labels:
              - dependencies
        
          - package-ecosystem: github-actions
            directory: "/.github/workflows"
            schedule:
              interval: weekly
              time: "04:00"
            open-pull-requests-limit: 5
            reviewers:
              - glenn-jocher
            labels:
              - dependencies
        
    • 📁 classify/
      • 📄 predict.py 11.7KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """
        Run YOLOv5 classification inference on images, videos, directories, globs, YouTube, webcam, streams, etc.
        
        Usage - sources:
            $ python classify/predict.py --weights yolov5s-cls.pt --source 0                               # webcam
                                                                           img.jpg                         # image
                                                                           vid.mp4                         # video
                                                                           screen                          # screenshot
                                                                           path/                           # directory
                                                                           list.txt                        # list of images
                                                                           list.streams                    # list of streams
                                                                           'path/*.jpg'                    # glob
                                                                           'https://youtu.be/LNwODJXcvt4'  # YouTube
                                                                           'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
        
        Usage - formats:
            $ python classify/predict.py --weights yolov5s-cls.pt                 # PyTorch
                                                   yolov5s-cls.torchscript        # TorchScript
                                                   yolov5s-cls.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                                   yolov5s-cls_openvino_model     # OpenVINO
                                                   yolov5s-cls.engine             # TensorRT
                                                   yolov5s-cls.mlmodel            # CoreML (macOS-only)
                                                   yolov5s-cls_saved_model        # TensorFlow SavedModel
                                                   yolov5s-cls.pb                 # TensorFlow GraphDef
                                                   yolov5s-cls.tflite             # TensorFlow Lite
                                                   yolov5s-cls_edgetpu.tflite     # TensorFlow Edge TPU
                                                   yolov5s-cls_paddle_model       # PaddlePaddle
        """
        
        import argparse
        import os
        import platform
        import sys
        from pathlib import Path
        
        import torch
        import torch.nn.functional as F
        
        FILE = Path(__file__).resolve()
        ROOT = FILE.parents[1]  # YOLOv5 root directory
        if str(ROOT) not in sys.path:
            sys.path.append(str(ROOT))  # add ROOT to PATH
        ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
        
        from ultralytics.utils.plotting import Annotator
        
        from models.common import DetectMultiBackend
        from utils.augmentations import classify_transforms
        from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
        from utils.general import (
            LOGGER,
            Profile,
            check_file,
            check_img_size,
            check_imshow,
            check_requirements,
            colorstr,
            cv2,
            increment_path,
            print_args,
            strip_optimizer,
        )
        from utils.torch_utils import select_device, smart_inference_mode
        
        
        @smart_inference_mode()
        def run(
            weights=ROOT / "yolov5s-cls.pt",  # model.pt path(s)
            source=ROOT / "data/images",  # file/dir/URL/glob/screen/0(webcam)
            data=ROOT / "data/coco128.yaml",  # dataset.yaml path
            imgsz=(224, 224),  # inference size (height, width)
            device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
            view_img=False,  # show results
            save_txt=False,  # save results to *.txt
            nosave=False,  # do not save images/videos
            augment=False,  # augmented inference
            visualize=False,  # visualize features
            update=False,  # update all models
            project=ROOT / "runs/predict-cls",  # save results to project/name
            name="exp",  # save results to project/name
            exist_ok=False,  # existing project/name ok, do not increment
            half=False,  # use FP16 half-precision inference
            dnn=False,  # use OpenCV DNN for ONNX inference
            vid_stride=1,  # video frame-rate stride
        ):
            source = str(source)
            save_img = not nosave and not source.endswith(".txt")  # save inference images
            is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
            is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
            webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
            screenshot = source.lower().startswith("screen")
            if is_url and is_file:
                source = check_file(source)  # download
        
            # Directories
            save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
            (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
        
            # Load model
            device = select_device(device)
            model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
            stride, names, pt = model.stride, model.names, model.pt
            imgsz = check_img_size(imgsz, s=stride)  # check image size
        
            # Dataloader
            bs = 1  # batch_size
            if webcam:
                view_img = check_imshow(warn=True)
                dataset = LoadStreams(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride)
                bs = len(dataset)
            elif screenshot:
                dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
            else:
                dataset = LoadImages(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride)
            vid_path, vid_writer = [None] * bs, [None] * bs
        
            # Run inference
            model.warmup(imgsz=(1 if pt else bs, 3, *imgsz))  # warmup
            seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))
            for path, im, im0s, vid_cap, s in dataset:
                with dt[0]:
                    im = torch.Tensor(im).to(model.device)
                    im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
                    if len(im.shape) == 3:
                        im = im[None]  # expand for batch dim
        
                # Inference
                with dt[1]:
                    results = model(im)
        
                # Post-process
                with dt[2]:
                    pred = F.softmax(results, dim=1)  # probabilities
        
                # Process predictions
                for i, prob in enumerate(pred):  # per image
                    seen += 1
                    if webcam:  # batch_size >= 1
                        p, im0, frame = path[i], im0s[i].copy(), dataset.count
                        s += f"{i}: "
                    else:
                        p, im0, frame = path, im0s.copy(), getattr(dataset, "frame", 0)
        
                    p = Path(p)  # to Path
                    save_path = str(save_dir / p.name)  # im.jpg
                    txt_path = str(save_dir / "labels" / p.stem) + ("" if dataset.mode == "image" else f"_{frame}")  # im.txt
        
                    s += "%gx%g " % im.shape[2:]  # print string
                    annotator = Annotator(im0, example=str(names), pil=True)
        
                    # Print results
                    top5i = prob.argsort(0, descending=True)[:5].tolist()  # top 5 indices
                    s += f"{', '.join(f'{names[j]} {prob[j]:.2f}' for j in top5i)}, "
        
                    # Write results
                    text = "\n".join(f"{prob[j]:.2f} {names[j]}" for j in top5i)
                    if save_img or view_img:  # Add bbox to image
                        annotator.text([32, 32], text, txt_color=(255, 255, 255))
                    if save_txt:  # Write to file
                        with open(f"{txt_path}.txt", "a") as f:
                            f.write(text + "\n")
        
                    # Stream results
                    im0 = annotator.result()
                    if view_img:
                        if platform.system() == "Linux" and p not in windows:
                            windows.append(p)
                            cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
                            cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
                        cv2.imshow(str(p), im0)
                        cv2.waitKey(1)  # 1 millisecond
        
                    # Save results (image with detections)
                    if save_img:
                        if dataset.mode == "image":
                            cv2.imwrite(save_path, im0)
                        else:  # 'video' or 'stream'
                            if vid_path[i] != save_path:  # new video
                                vid_path[i] = save_path
                                if isinstance(vid_writer[i], cv2.VideoWriter):
                                    vid_writer[i].release()  # release previous video writer
                                if vid_cap:  # video
                                    fps = vid_cap.get(cv2.CAP_PROP_FPS)
                                    w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                                    h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                                else:  # stream
                                    fps, w, h = 30, im0.shape[1], im0.shape[0]
                                save_path = str(Path(save_path).with_suffix(".mp4"))  # force *.mp4 suffix on results videos
                                vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
                            vid_writer[i].write(im0)
        
                # Print time (inference-only)
                LOGGER.info(f"{s}{dt[1].dt * 1E3:.1f}ms")
        
            # Print results
            t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
            LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}" % t)
            if save_txt or save_img:
                s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
                LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
            if update:
                strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
        
        
        def parse_opt():
            """Parses command line arguments for YOLOv5 inference settings including model, source, device, and image size."""
            parser = argparse.ArgumentParser()
            parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s-cls.pt", help="model path(s)")
            parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
            parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="(optional) dataset.yaml path")
            parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[224], help="inference size h,w")
            parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
            parser.add_argument("--view-img", action="store_true", help="show results")
            parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
            parser.add_argument("--nosave", action="store_true", help="do not save images/videos")
            parser.add_argument("--augment", action="store_true", help="augmented inference")
            parser.add_argument("--visualize", action="store_true", help="visualize features")
            parser.add_argument("--update", action="store_true", help="update all models")
            parser.add_argument("--project", default=ROOT / "runs/predict-cls", help="save results to project/name")
            parser.add_argument("--name", default="exp", help="save results to project/name")
            parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
            parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
            parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
            parser.add_argument("--vid-stride", type=int, default=1, help="video frame-rate stride")
            opt = parser.parse_args()
            opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
            print_args(vars(opt))
            return opt
        
        
        def main(opt):
            """Executes YOLOv5 model inference with options for ONNX DNN and video frame-rate stride adjustments."""
            check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
            run(**vars(opt))
        
        
        if __name__ == "__main__":
            opt = parse_opt()
            main(opt)
        
      • 📄 train.py 16.0KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """
        Train a YOLOv5 classifier model on a classification dataset.
        
        Usage - Single-GPU training:
            $ python classify/train.py --model yolov5s-cls.pt --data imagenette160 --epochs 5 --img 224
        
        Usage - Multi-GPU DDP training:
            $ python -m torch.distributed.run --nproc_per_node 4 --master_port 2022 classify/train.py --model yolov5s-cls.pt --data imagenet --epochs 5 --img 224 --device 0,1,2,3
        
        Datasets:           --data mnist, fashion-mnist, cifar10, cifar100, imagenette, imagewoof, imagenet, or 'path/to/data'
        YOLOv5-cls models:  --model yolov5n-cls.pt, yolov5s-cls.pt, yolov5m-cls.pt, yolov5l-cls.pt, yolov5x-cls.pt
        Torchvision models: --model resnet50, efficientnet_b0, etc. See https://pytorch.org/vision/stable/models.html
        """
        
        import argparse
        import os
        import subprocess
        import sys
        import time
        from copy import deepcopy
        from datetime import datetime
        from pathlib import Path
        
        import torch
        import torch.distributed as dist
        import torch.hub as hub
        import torch.optim.lr_scheduler as lr_scheduler
        import torchvision
        from torch.cuda import amp
        from tqdm import tqdm
        
        FILE = Path(__file__).resolve()
        ROOT = FILE.parents[1]  # YOLOv5 root directory
        if str(ROOT) not in sys.path:
            sys.path.append(str(ROOT))  # add ROOT to PATH
        ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
        
        from classify import val as validate
        from models.experimental import attempt_load
        from models.yolo import ClassificationModel, DetectionModel
        from utils.dataloaders import create_classification_dataloader
        from utils.general import (
            DATASETS_DIR,
            LOGGER,
            TQDM_BAR_FORMAT,
            WorkingDirectory,
            check_git_info,
            check_git_status,
            check_requirements,
            colorstr,
            download,
            increment_path,
            init_seeds,
            print_args,
            yaml_save,
        )
        from utils.loggers import GenericLogger
        from utils.plots import imshow_cls
        from utils.torch_utils import (
            ModelEMA,
            de_parallel,
            model_info,
            reshape_classifier_output,
            select_device,
            smart_DDP,
            smart_optimizer,
            smartCrossEntropyLoss,
            torch_distributed_zero_first,
        )
        
        LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))  # https://pytorch.org/docs/stable/elastic/run.html
        RANK = int(os.getenv("RANK", -1))
        WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
        GIT_INFO = check_git_info()
        
        
        def train(opt, device):
            """Trains a YOLOv5 model, managing datasets, model optimization, logging, and saving checkpoints."""
            init_seeds(opt.seed + 1 + RANK, deterministic=True)
            save_dir, data, bs, epochs, nw, imgsz, pretrained = (
                opt.save_dir,
                Path(opt.data),
                opt.batch_size,
                opt.epochs,
                min(os.cpu_count() - 1, opt.workers),
                opt.imgsz,
                str(opt.pretrained).lower() == "true",
            )
            cuda = device.type != "cpu"
        
            # Directories
            wdir = save_dir / "weights"
            wdir.mkdir(parents=True, exist_ok=True)  # make dir
            last, best = wdir / "last.pt", wdir / "best.pt"
        
            # Save run settings
            yaml_save(save_dir / "opt.yaml", vars(opt))
        
            # Logger
            logger = GenericLogger(opt=opt, console_logger=LOGGER) if RANK in {-1, 0} else None
        
            # Download Dataset
            with torch_distributed_zero_first(LOCAL_RANK), WorkingDirectory(ROOT):
                data_dir = data if data.is_dir() else (DATASETS_DIR / data)
                if not data_dir.is_dir():
                    LOGGER.info(f"\nDataset not found ⚠️, missing path {data_dir}, attempting download...")
                    t = time.time()
                    if str(data) == "imagenet":
                        subprocess.run(["bash", str(ROOT / "data/scripts/get_imagenet.sh")], shell=True, check=True)
                    else:
                        url = f"https://github.com/ultralytics/yolov5/releases/download/v1.0/{data}.zip"
                        download(url, dir=data_dir.parent)
                    s = f"Dataset download success ✅ ({time.time() - t:.1f}s), saved to {colorstr('bold', data_dir)}\n"
                    LOGGER.info(s)
        
            # Dataloaders
            nc = len([x for x in (data_dir / "train").glob("*") if x.is_dir()])  # number of classes
            trainloader = create_classification_dataloader(
                path=data_dir / "train",
                imgsz=imgsz,
                batch_size=bs // WORLD_SIZE,
                augment=True,
                cache=opt.cache,
                rank=LOCAL_RANK,
                workers=nw,
            )
        
            test_dir = data_dir / "test" if (data_dir / "test").exists() else data_dir / "val"  # data/test or data/val
            if RANK in {-1, 0}:
                testloader = create_classification_dataloader(
                    path=test_dir,
                    imgsz=imgsz,
                    batch_size=bs // WORLD_SIZE * 2,
                    augment=False,
                    cache=opt.cache,
                    rank=-1,
                    workers=nw,
                )
        
            # Model
            with torch_distributed_zero_first(LOCAL_RANK), WorkingDirectory(ROOT):
                if Path(opt.model).is_file() or opt.model.endswith(".pt"):
                    model = attempt_load(opt.model, device="cpu", fuse=False)
                elif opt.model in torchvision.models.__dict__:  # TorchVision models i.e. resnet50, efficientnet_b0
                    model = torchvision.models.__dict__[opt.model](weights="IMAGENET1K_V1" if pretrained else None)
                else:
                    m = hub.list("ultralytics/yolov5")  # + hub.list('pytorch/vision')  # models
                    raise ModuleNotFoundError(f"--model {opt.model} not found. Available models are: \n" + "\n".join(m))
                if isinstance(model, DetectionModel):
                    LOGGER.warning("WARNING ⚠️ pass YOLOv5 classifier model with '-cls' suffix, i.e. '--model yolov5s-cls.pt'")
                    model = ClassificationModel(model=model, nc=nc, cutoff=opt.cutoff or 10)  # convert to classification model
                reshape_classifier_output(model, nc)  # update class count
            for m in model.modules():
                if not pretrained and hasattr(m, "reset_parameters"):
                    m.reset_parameters()
                if isinstance(m, torch.nn.Dropout) and opt.dropout is not None:
                    m.p = opt.dropout  # set dropout
            for p in model.parameters():
                p.requires_grad = True  # for training
            model = model.to(device)
        
            # Info
            if RANK in {-1, 0}:
                model.names = trainloader.dataset.classes  # attach class names
                model.transforms = testloader.dataset.torch_transforms  # attach inference transforms
                model_info(model)
                if opt.verbose:
                    LOGGER.info(model)
                images, labels = next(iter(trainloader))
                file = imshow_cls(images[:25], labels[:25], names=model.names, f=save_dir / "train_images.jpg")
                logger.log_images(file, name="Train Examples")
                logger.log_graph(model, imgsz)  # log model
        
            # Optimizer
            optimizer = smart_optimizer(model, opt.optimizer, opt.lr0, momentum=0.9, decay=opt.decay)
        
            # Scheduler
            lrf = 0.01  # final lr (fraction of lr0)
            # lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - lrf) + lrf  # cosine
            lf = lambda x: (1 - x / epochs) * (1 - lrf) + lrf  # linear
            scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
            # scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=lr0, total_steps=epochs, pct_start=0.1,
            #                                    final_div_factor=1 / 25 / lrf)
        
            # EMA
            ema = ModelEMA(model) if RANK in {-1, 0} else None
        
            # DDP mode
            if cuda and RANK != -1:
                model = smart_DDP(model)
        
            # Train
            t0 = time.time()
            criterion = smartCrossEntropyLoss(label_smoothing=opt.label_smoothing)  # loss function
            best_fitness = 0.0
            scaler = amp.GradScaler(enabled=cuda)
            val = test_dir.stem  # 'val' or 'test'
            LOGGER.info(
                f'Image sizes {imgsz} train, {imgsz} test\n'
                f'Using {nw * WORLD_SIZE} dataloader workers\n'
                f"Logging results to {colorstr('bold', save_dir)}\n"
                f'Starting {opt.model} training on {data} dataset with {nc} classes for {epochs} epochs...\n\n'
                f"{'Epoch':>10}{'GPU_mem':>10}{'train_loss':>12}{f'{val}_loss':>12}{'top1_acc':>12}{'top5_acc':>12}"
            )
            for epoch in range(epochs):  # loop over the dataset multiple times
                tloss, vloss, fitness = 0.0, 0.0, 0.0  # train loss, val loss, fitness
                model.train()
                if RANK != -1:
                    trainloader.sampler.set_epoch(epoch)
                pbar = enumerate(trainloader)
                if RANK in {-1, 0}:
                    pbar = tqdm(enumerate(trainloader), total=len(trainloader), bar_format=TQDM_BAR_FORMAT)
                for i, (images, labels) in pbar:  # progress bar
                    images, labels = images.to(device, non_blocking=True), labels.to(device)
        
                    # Forward
                    with amp.autocast(enabled=cuda):  # stability issues when enabled
                        loss = criterion(model(images), labels)
        
                    # Backward
                    scaler.scale(loss).backward()
        
                    # Optimize
                    scaler.unscale_(optimizer)  # unscale gradients
                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)  # clip gradients
                    scaler.step(optimizer)
                    scaler.update()
                    optimizer.zero_grad()
                    if ema:
                        ema.update(model)
        
                    if RANK in {-1, 0}:
                        # Print
                        tloss = (tloss * i + loss.item()) / (i + 1)  # update mean losses
                        mem = "%.3gG" % (torch.cuda.memory_reserved() / 1e9 if torch.cuda.is_available() else 0)  # (GB)
                        pbar.desc = f"{f'{epoch + 1}/{epochs}':>10}{mem:>10}{tloss:>12.3g}" + " " * 36
        
                        # Test
                        if i == len(pbar) - 1:  # last batch
                            top1, top5, vloss = validate.run(
                                model=ema.ema, dataloader=testloader, criterion=criterion, pbar=pbar
                            )  # test accuracy, loss
                            fitness = top1  # define fitness as top1 accuracy
        
                # Scheduler
                scheduler.step()
        
                # Log metrics
                if RANK in {-1, 0}:
                    # Best fitness
                    if fitness > best_fitness:
                        best_fitness = fitness
        
                    # Log
                    metrics = {
                        "train/loss": tloss,
                        f"{val}/loss": vloss,
                        "metrics/accuracy_top1": top1,
                        "metrics/accuracy_top5": top5,
                        "lr/0": optimizer.param_groups[0]["lr"],
                    }  # learning rate
                    logger.log_metrics(metrics, epoch)
        
                    # Save model
                    final_epoch = epoch + 1 == epochs
                    if (not opt.nosave) or final_epoch:
                        ckpt = {
                            "epoch": epoch,
                            "best_fitness": best_fitness,
                            "model": deepcopy(ema.ema).half(),  # deepcopy(de_parallel(model)).half(),
                            "ema": None,  # deepcopy(ema.ema).half(),
                            "updates": ema.updates,
                            "optimizer": None,  # optimizer.state_dict(),
                            "opt": vars(opt),
                            "git": GIT_INFO,  # {remote, branch, commit} if a git repo
                            "date": datetime.now().isoformat(),
                        }
        
                        # Save last, best and delete
                        torch.save(ckpt, last)
                        if best_fitness == fitness:
                            torch.save(ckpt, best)
                        del ckpt
        
            # Train complete
            if RANK in {-1, 0} and final_epoch:
                LOGGER.info(
                    f'\nTraining complete ({(time.time() - t0) / 3600:.3f} hours)'
                    f"\nResults saved to {colorstr('bold', save_dir)}"
                    f'\nPredict:         python classify/predict.py --weights {best} --source im.jpg'
                    f'\nValidate:        python classify/val.py --weights {best} --data {data_dir}'
                    f'\nExport:          python export.py --weights {best} --include onnx'
                    f"\nPyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', '{best}')"
                    f'\nVisualize:       https://netron.app\n'
                )
        
                # Plot examples
                images, labels = (x[:25] for x in next(iter(testloader)))  # first 25 images and labels
                pred = torch.max(ema.ema(images.to(device)), 1)[1]
                file = imshow_cls(images, labels, pred, de_parallel(model).names, verbose=False, f=save_dir / "test_images.jpg")
        
                # Log results
                meta = {"epochs": epochs, "top1_acc": best_fitness, "date": datetime.now().isoformat()}
                logger.log_images(file, name="Test Examples (true-predicted)", epoch=epoch)
                logger.log_model(best, epochs, metadata=meta)
        
        
        def parse_opt(known=False):
            """Parses command line arguments for YOLOv5 training including model path, dataset, epochs, and more, returning
            parsed arguments.
            """
            parser = argparse.ArgumentParser()
            parser.add_argument("--model", type=str, default="yolov5s-cls.pt", help="initial weights path")
            parser.add_argument("--data", type=str, default="imagenette160", help="cifar10, cifar100, mnist, imagenet, ...")
            parser.add_argument("--epochs", type=int, default=10, help="total training epochs")
            parser.add_argument("--batch-size", type=int, default=64, help="total batch size for all GPUs")
            parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=224, help="train, val image size (pixels)")
            parser.add_argument("--nosave", action="store_true", help="only save final checkpoint")
            parser.add_argument("--cache", type=str, nargs="?", const="ram", help='--cache images in "ram" (default) or "disk"')
            parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
            parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
            parser.add_argument("--project", default=ROOT / "runs/train-cls", help="save to project/name")
            parser.add_argument("--name", default="exp", help="save to project/name")
            parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
            parser.add_argument("--pretrained", nargs="?", const=True, default=True, help="start from i.e. --pretrained False")
            parser.add_argument("--optimizer", choices=["SGD", "Adam", "AdamW", "RMSProp"], default="Adam", help="optimizer")
            parser.add_argument("--lr0", type=float, default=0.001, help="initial learning rate")
            parser.add_argument("--decay", type=float, default=5e-5, help="weight decay")
            parser.add_argument("--label-smoothing", type=float, default=0.1, help="Label smoothing epsilon")
            parser.add_argument("--cutoff", type=int, default=None, help="Model layer cutoff index for Classify() head")
            parser.add_argument("--dropout", type=float, default=None, help="Dropout (fraction)")
            parser.add_argument("--verbose", action="store_true", help="Verbose mode")
            parser.add_argument("--seed", type=int, default=0, help="Global training seed")
            parser.add_argument("--local_rank", type=int, default=-1, help="Automatic DDP Multi-GPU argument, do not modify")
            return parser.parse_known_args()[0] if known else parser.parse_args()
        
        
        def main(opt):
            """Executes YOLOv5 training with given options, handling device setup and DDP mode; includes pre-training checks."""
            if RANK in {-1, 0}:
                print_args(vars(opt))
                check_git_status()
                check_requirements(ROOT / "requirements.txt")
        
            # DDP mode
            device = select_device(opt.device, batch_size=opt.batch_size)
            if LOCAL_RANK != -1:
                assert opt.batch_size != -1, "AutoBatch is coming soon for classification, please pass a valid --batch-size"
                assert opt.batch_size % WORLD_SIZE == 0, f"--batch-size {opt.batch_size} must be multiple of WORLD_SIZE"
                assert torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command"
                torch.cuda.set_device(LOCAL_RANK)
                device = torch.device("cuda", LOCAL_RANK)
                dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
        
            # Parameters
            opt.save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)  # increment run
        
            # Train
            train(opt, device)
        
        
        def run(**kwargs):
            """
            Executes YOLOv5 model training or inference with specified parameters, returning updated options.
        
            Example: from yolov5 import classify; classify.train.run(data=mnist, imgsz=320, model='yolov5m')
            """
            opt = parse_opt(True)
            for k, v in kwargs.items():
                setattr(opt, k, v)
            main(opt)
            return opt
        
        
        if __name__ == "__main__":
            opt = parse_opt()
            main(opt)
        
      • 📄 tutorial.ipynb 101.2KB

        Run on Gradient Open In Colab Open In Kaggle
        This YOLOv5 🚀 notebook by Ultralytics presents simple train, validate and predict examples to help start your AI adventure.
        See GitHub for community support or contact us for professional support.

        Setup

        Clone GitHub repository, install dependencies and check PyTorch and GPU.

        !git clone https://github.com/ultralytics/yolov5  # clone
        %cd yolov5
        %pip install -qr requirements.txt  # install
        
        import torch
        import utils
        display = utils.notebook_init()  # checks
        
        YOLOv5 🚀 v7.0-3-g61ebf5e Python-3.7.15 torch-1.12.1+cu113 CUDA:0 (Tesla T4, 15110MiB)
        Setup complete ✅ (2 CPUs, 12.7 GB RAM, 22.6/78.2 GB disk)
        

        1. Predict

        classify/predict.py runs YOLOv5 Classification inference on a variety of sources, downloading models automatically from the latest YOLOv5 release, and saving results to runs/predict-cls. Example inference sources are:

        python classify/predict.py --source 0  # webcam
                                      img.jpg  # image 
                                      vid.mp4  # video
                                      screen  # screenshot
                                      path/  # directory
                                      'path/*.jpg'  # glob
                                      'https://youtu.be/LNwODJXcvt4'  # YouTube
                                      'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
        
        !python classify/predict.py --weights yolov5s-cls.pt --img 224 --source data/images
        # display.Image(filename='runs/predict-cls/exp/zidane.jpg', width=600)
        
        classify/predict: weights=['yolov5s-cls.pt'], source=data/images, data=data/coco128.yaml, imgsz=[224, 224], device=, view_img=False, save_txt=False, nosave=False, augment=False, visualize=False, update=False, project=runs/predict-cls, name=exp, exist_ok=False, half=False, dnn=False, vid_stride=1
        YOLOv5 🚀 v7.0-3-g61ebf5e Python-3.7.15 torch-1.12.1+cu113 CUDA:0 (Tesla T4, 15110MiB)
        
        Downloading https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s-cls.pt to yolov5s-cls.pt...
        100% 10.5M/10.5M [00:00<00:00, 12.3MB/s]
        
        Fusing layers... 
        Model summary: 117 layers, 5447688 parameters, 0 gradients, 11.4 GFLOPs
        image 1/2 /content/yolov5/data/images/bus.jpg: 224x224 minibus 0.39, police van 0.24, amphibious vehicle 0.05, recreational vehicle 0.04, trolleybus 0.03, 3.9ms
        image 2/2 /content/yolov5/data/images/zidane.jpg: 224x224 suit 0.38, bow tie 0.19, bridegroom 0.18, rugby ball 0.04, stage 0.02, 4.6ms
        Speed: 0.3ms pre-process, 4.3ms inference, 1.5ms NMS per image at shape (1, 3, 224, 224)
        Results saved to runs/predict-cls/exp
        

                

        2. Validate

        Validate a model's accuracy on the Imagenet dataset's val or test splits. Models are downloaded automatically from the latest YOLOv5 release. To show results by class use the --verbose flag.

        # Download Imagenet val (6.3G, 50000 images)
        !bash data/scripts/get_imagenet.sh --val
        
        --2022-11-22 19:53:40--  https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar
        Resolving image-net.org (image-net.org)... 171.64.68.16
        Connecting to image-net.org (image-net.org)|171.64.68.16|:443... connected.
        HTTP request sent, awaiting response... 200 OK
        Length: 6744924160 (6.3G) [application/x-tar]
        Saving to: ‘ILSVRC2012_img_val.tar’
        
        ILSVRC2012_img_val. 100%[===================>]   6.28G  16.1MB/s    in 10m 52s 
        
        2022-11-22 20:04:32 (9.87 MB/s) - ‘ILSVRC2012_img_val.tar’ saved [6744924160/6744924160]
        
        
        # Validate YOLOv5s on Imagenet val
        !python classify/val.py --weights yolov5s-cls.pt --data ../datasets/imagenet --img 224 --half
        
        classify/val: data=../datasets/imagenet, weights=['yolov5s-cls.pt'], batch_size=128, imgsz=224, device=, workers=8, verbose=True, project=runs/val-cls, name=exp, exist_ok=False, half=True, dnn=False
        YOLOv5 🚀 v7.0-3-g61ebf5e Python-3.7.15 torch-1.12.1+cu113 CUDA:0 (Tesla T4, 15110MiB)
        
        Fusing layers... 
        Model summary: 117 layers, 5447688 parameters, 0 gradients, 11.4 GFLOPs
        validating: 100% 391/391 [04:57<00:00,  1.31it/s]
                           Class      Images    top1_acc    top5_acc
                             all       50000       0.715       0.902
                           tench          50        0.94        0.98
                        goldfish          50        0.88        0.92
               great white shark          50        0.78        0.96
                     tiger shark          50        0.68        0.96
                hammerhead shark          50        0.82        0.92
                    electric ray          50        0.76         0.9
                        stingray          50         0.7         0.9
                            cock          50        0.78        0.92
                             hen          50        0.84        0.96
                         ostrich          50        0.98           1
                       brambling          50         0.9        0.96
                       goldfinch          50        0.92        0.98
                     house finch          50        0.88        0.96
                           junco          50        0.94        0.98
                  indigo bunting          50        0.86        0.88
                  American robin          50         0.9        0.96
                          bulbul          50        0.84        0.96
                             jay          50         0.9        0.96
                          magpie          50        0.84        0.96
                       chickadee          50         0.9           1
                 American dipper          50        0.82        0.92
                            kite          50        0.76        0.94
                      bald eagle          50        0.92           1
                         vulture          50        0.96           1
                  great grey owl          50        0.94        0.98
                 fire salamander          50        0.96        0.98
                     smooth newt          50        0.58        0.94
                            newt          50        0.74         0.9
              spotted salamander          50        0.86        0.94
                         axolotl          50        0.86        0.96
               American bullfrog          50        0.78        0.92
                       tree frog          50        0.84        0.96
                     tailed frog          50        0.48         0.8
           loggerhead sea turtle          50        0.68        0.94
          leatherback sea turtle          50         0.5         0.8
                      mud turtle          50        0.64        0.84
                        terrapin          50        0.52        0.98
                      box turtle          50        0.84        0.98
                    banded gecko          50         0.7        0.88
                    green iguana          50        0.76        0.94
                  Carolina anole          50        0.58        0.96
        desert grassland whiptail lizard          50        0.82        0.94
                           agama          50        0.74        0.92
           frilled-necked lizard          50        0.84        0.86
                alligator lizard          50        0.58        0.78
                    Gila monster          50        0.72         0.8
           European green lizard          50        0.42         0.9
                       chameleon          50        0.76        0.84
                   Komodo dragon          50        0.86        0.96
                  Nile crocodile          50         0.7        0.84
              American alligator          50        0.76        0.96
                     triceratops          50         0.9        0.94
                      worm snake          50        0.76        0.88
               ring-necked snake          50         0.8        0.92
         eastern hog-nosed snake          50        0.58        0.88
              smooth green snake          50         0.6        0.94
                       kingsnake          50        0.82         0.9
                    garter snake          50        0.88        0.94
                     water snake          50         0.7        0.94
                      vine snake          50        0.66        0.76
                     night snake          50        0.34        0.82
                 boa constrictor          50         0.8        0.96
             African rock python          50        0.48        0.76
                    Indian cobra          50        0.82        0.94
                     green mamba          50        0.54        0.86
                       sea snake          50        0.62         0.9
            Saharan horned viper          50        0.56        0.86
        eastern diamondback rattlesnake          50         0.6        0.86
                      sidewinder          50        0.28        0.86
                       trilobite          50        0.98        0.98
                      harvestman          50        0.86        0.94
                        scorpion          50        0.86        0.94
            yellow garden spider          50        0.92        0.96
                     barn spider          50        0.38        0.98
          European garden spider          50        0.62        0.98
            southern black widow          50        0.88        0.94
                       tarantula          50        0.94           1
                     wolf spider          50        0.82        0.92
                            tick          50        0.74        0.84
                       centipede          50        0.68        0.82
                    black grouse          50        0.88        0.98
                       ptarmigan          50        0.78        0.94
                   ruffed grouse          50        0.88           1
                  prairie grouse          50        0.92           1
                         peacock          50        0.88         0.9
                           quail          50         0.9        0.94
                       partridge          50        0.74        0.96
                     grey parrot          50         0.9        0.96
                           macaw          50        0.88        0.98
        sulphur-crested cockatoo          50        0.86        0.92
                        lorikeet          50        0.96           1
                          coucal          50        0.82        0.88
                       bee eater          50        0.96        0.98
                        hornbill          50         0.9        0.96
                     hummingbird          50        0.88        0.96
                         jacamar          50        0.92        0.94
                          toucan          50        0.84        0.94
                            duck          50        0.76        0.94
          red-breasted merganser          50        0.86        0.96
                           goose          50        0.74        0.96
                      black swan          50        0.94        0.98
                          tusker          50        0.54        0.92
                         echidna          50        0.98           1
                        platypus          50        0.72        0.84
                         wallaby          50        0.78        0.88
                           koala          50        0.84        0.92
                          wombat          50        0.78        0.84
                       jellyfish          50        0.88        0.96
                     sea anemone          50        0.72         0.9
                     brain coral          50        0.88        0.96
                        flatworm          50         0.8        0.98
                        nematode          50        0.86         0.9
                           conch          50        0.74        0.88
                           snail          50        0.78        0.88
                            slug          50        0.74        0.82
                        sea slug          50        0.88        0.98
                          chiton          50        0.88        0.98
              chambered nautilus          50        0.88        0.92
                  Dungeness crab          50        0.78        0.94
                       rock crab          50        0.68        0.86
                    fiddler crab          50        0.64        0.86
                   red king crab          50        0.76        0.96
                American lobster          50        0.78        0.96
                   spiny lobster          50        0.74        0.88
                        crayfish          50        0.56        0.86
                     hermit crab          50        0.78        0.96
                          isopod          50        0.66        0.78
                     white stork          50        0.88        0.96
                     black stork          50        0.84        0.98
                       spoonbill          50        0.96           1
                        flamingo          50        0.94           1
               little blue heron          50        0.92        0.98
                     great egret          50         0.9        0.96
                         bittern          50        0.86        0.94
                    crane (bird)          50        0.62         0.9
                         limpkin          50        0.98           1
                common gallinule          50        0.92        0.96
                   American coot          50         0.9        0.98
                         bustard          50        0.92        0.96
                 ruddy turnstone          50        0.94           1
                          dunlin          50        0.86        0.94
                 common redshank          50         0.9        0.96
                       dowitcher          50        0.84        0.96
                   oystercatcher          50        0.86        0.94
                         pelican          50        0.92        0.96
                    king penguin          50        0.88        0.96
                       albatross          50         0.9           1
                      grey whale          50        0.84        0.92
                    killer whale          50        0.92           1
                          dugong          50        0.84        0.96
                        sea lion          50        0.82        0.92
                       Chihuahua          50        0.66        0.84
                   Japanese Chin          50        0.72        0.98
                         Maltese          50        0.76        0.94
                       Pekingese          50        0.84        0.94
                        Shih Tzu          50        0.74        0.96
            King Charles Spaniel          50        0.88        0.98
                        Papillon          50        0.86        0.94
                     toy terrier          50        0.48        0.94
             Rhodesian Ridgeback          50        0.76        0.98
                    Afghan Hound          50        0.84           1
                    Basset Hound          50         0.8        0.92
                          Beagle          50        0.82        0.96
                      Bloodhound          50        0.48        0.72
              Bluetick Coonhound          50        0.86        0.94
         Black and Tan Coonhound          50        0.54         0.8
        Treeing Walker Coonhound          50        0.66        0.98
                English foxhound          50        0.32        0.84
               Redbone Coonhound          50        0.62        0.94
                          borzoi          50        0.92           1
                 Irish Wolfhound          50        0.48        0.88
               Italian Greyhound          50        0.76        0.98
                         Whippet          50        0.74        0.92
                    Ibizan Hound          50         0.6        0.86
              Norwegian Elkhound          50        0.88        0.98
                      Otterhound          50        0.62         0.9
                          Saluki          50        0.72        0.92
              Scottish Deerhound          50        0.86        0.98
                      Weimaraner          50        0.88        0.94
        Staffordshire Bull Terrier          50        0.66        0.98
        American Staffordshire Terrier          50        0.64        0.92
              Bedlington Terrier          50         0.9        0.92
                  Border Terrier          50        0.86        0.92
              Kerry Blue Terrier          50        0.78        0.98
                   Irish Terrier          50         0.7        0.96
                 Norfolk Terrier          50        0.68         0.9
                 Norwich Terrier          50        0.72           1
               Yorkshire Terrier          50        0.66         0.9
                Wire Fox Terrier          50        0.64        0.98
                Lakeland Terrier          50        0.74        0.92
                Sealyham Terrier          50        0.76         0.9
                Airedale Terrier          50        0.82        0.92
                   Cairn Terrier          50        0.76         0.9
              Australian Terrier          50        0.48        0.84
          Dandie Dinmont Terrier          50        0.82        0.92
                  Boston Terrier          50        0.92           1
             Miniature Schnauzer          50        0.68         0.9
                 Giant Schnauzer          50        0.72        0.98
              Standard Schnauzer          50        0.74           1
                Scottish Terrier          50        0.76        0.96
                 Tibetan Terrier          50        0.48           1
        Australian Silky Terrier          50        0.66        0.96
        Soft-coated Wheaten Terrier          50        0.74        0.96
        West Highland White Terrier          50        0.88        0.96
                      Lhasa Apso          50        0.68        0.96
           Flat-Coated Retriever          50        0.72        0.94
          Curly-coated Retriever          50        0.82        0.94
                Golden Retriever          50        0.86        0.94
              Labrador Retriever          50        0.82        0.94
        Chesapeake Bay Retriever          50        0.76        0.96
        German Shorthaired Pointer          50         0.8        0.96
                          Vizsla          50        0.68        0.96
                  English Setter          50         0.7           1
                    Irish Setter          50         0.8         0.9
                   Gordon Setter          50        0.84        0.92
                        Brittany          50        0.84        0.96
                 Clumber Spaniel          50        0.92        0.96
        English Springer Spaniel          50        0.88           1
          Welsh Springer Spaniel          50        0.92           1
                 Cocker Spaniels          50         0.7        0.94
                  Sussex Spaniel          50        0.72        0.92
             Irish Water Spaniel          50        0.88        0.98
                          Kuvasz          50        0.66         0.9
                      Schipperke          50         0.9        0.98
                     Groenendael          50         0.8        0.94
                        Malinois          50        0.86        0.98
                          Briard          50        0.52         0.8
               Australian Kelpie          50         0.6        0.88
                        Komondor          50        0.88        0.94
            Old English Sheepdog          50        0.94        0.98
               Shetland Sheepdog          50        0.74         0.9
                          collie          50         0.6        0.96
                   Border Collie          50        0.74        0.96
            Bouvier des Flandres          50        0.78        0.94
                      Rottweiler          50        0.88        0.96
             German Shepherd Dog          50         0.8        0.98
                       Dobermann          50        0.68        0.96
              Miniature Pinscher          50        0.76        0.88
        Greater Swiss Mountain Dog          50        0.68        0.94
            Bernese Mountain Dog          50        0.96           1
          Appenzeller Sennenhund          50        0.22           1
          Entlebucher Sennenhund          50        0.64        0.98
                           Boxer          50         0.7        0.92
                     Bullmastiff          50        0.78        0.98
                 Tibetan Mastiff          50        0.88        0.96
                  French Bulldog          50        0.84        0.94
                      Great Dane          50        0.54         0.9
                     St. Bernard          50        0.92           1
                           husky          50        0.46        0.98
                Alaskan Malamute          50        0.76        0.96
                  Siberian Husky          50        0.46        0.98
                       Dalmatian          50        0.94        0.98
                   Affenpinscher          50        0.78         0.9
                         Basenji          50        0.92        0.94
                             pug          50        0.94        0.98
                      Leonberger          50           1           1
                    Newfoundland          50        0.78        0.96
           Pyrenean Mountain Dog          50        0.78        0.96
                         Samoyed          50        0.96           1
                      Pomeranian          50        0.98           1
                       Chow Chow          50         0.9        0.96
                        Keeshond          50        0.88        0.94
              Griffon Bruxellois          50        0.84        0.98
            Pembroke Welsh Corgi          50        0.82        0.94
            Cardigan Welsh Corgi          50        0.66        0.98
                      Toy Poodle          50        0.52        0.88
                Miniature Poodle          50        0.52        0.92
                 Standard Poodle          50         0.8           1
            Mexican hairless dog          50        0.88        0.98
                       grey wolf          50        0.82        0.92
             Alaskan tundra wolf          50        0.78        0.98
                        red wolf          50        0.48         0.9
                          coyote          50        0.64        0.86
                           dingo          50        0.76        0.88
                           dhole          50         0.9        0.98
                African wild dog          50        0.98           1
                           hyena          50        0.88        0.96
                         red fox          50        0.54        0.92
                         kit fox          50        0.72        0.98
                      Arctic fox          50        0.94           1
                        grey fox          50         0.7        0.94
                       tabby cat          50        0.54        0.92
                       tiger cat          50        0.22        0.94
                     Persian cat          50         0.9        0.98
                     Siamese cat          50        0.96           1
                    Egyptian Mau          50        0.54         0.8
                          cougar          50         0.9           1
                            lynx          50        0.72        0.88
                         leopard          50        0.78        0.98
                    snow leopard          50         0.9        0.98
                          jaguar          50         0.7        0.94
                            lion          50         0.9        0.98
                           tiger          50        0.92        0.98
                         cheetah          50        0.94        0.98
                      brown bear          50        0.94        0.98
             American black bear          50         0.8           1
                      polar bear          50        0.84        0.96
                      sloth bear          50        0.72        0.92
                        mongoose          50         0.7        0.92
                         meerkat          50        0.82        0.92
                    tiger beetle          50        0.92        0.94
                         ladybug          50        0.86        0.94
                   ground beetle          50        0.64        0.94
                 longhorn beetle          50        0.62        0.88
                     leaf beetle          50        0.64        0.98
                     dung beetle          50        0.86        0.98
               rhinoceros beetle          50        0.86        0.94
                          weevil          50         0.9           1
                             fly          50        0.78        0.94
                             bee          50        0.68        0.94
                             ant          50        0.68        0.78
                     grasshopper          50         0.5        0.92
                         cricket          50        0.64        0.92
                    stick insect          50        0.64        0.92
                       cockroach          50        0.72         0.8
                          mantis          50        0.64        0.86
                          cicada          50         0.9        0.96
                      leafhopper          50        0.88        0.94
                

        3. Train

        Close the active learning loop by sampling images from your inference conditions with the roboflow pip package


        Train a YOLOv5s Classification model on the Imagenette dataset with --data imagenet, starting from pretrained --pretrained yolov5s-cls.pt.

        • Pretrained Models are downloaded
          automatically from the latest YOLOv5 release
        • Training Results are saved to runs/train-cls/ with incrementing run directories, i.e. runs/train-cls/exp2, runs/train-cls/exp3 etc.


        A Mosaic Dataloader is used for training which combines 4 images into 1 mosaic.

        Train on Custom Data with Roboflow 🌟 NEW

        Roboflow enables you to easily organize, label, and prepare a high quality dataset with your own custom data. Roboflow also makes it easy to establish an active learning pipeline, collaborate with your team on dataset improvement, and integrate directly into your model building workflow with the roboflow pip package.

        Label images lightning fast (including with model-assisted labeling)

        #@title Select YOLOv5 🚀 logger {run: 'auto'}
        logger = 'Comet' #@param ['Comet', 'ClearML', 'TensorBoard']
        
        if logger == 'Comet':
          %pip install -q comet_ml
          import comet_ml; comet_ml.init()
        elif logger == 'ClearML':
          %pip install -q clearml
          import clearml; clearml.browser_login()
        elif logger == 'TensorBoard':
          %load_ext tensorboard
          %tensorboard --logdir runs/train
        
        # Train YOLOv5s Classification on Imagenette160 for 3 epochs
        !python classify/train.py --model yolov5s-cls.pt --data imagenette160 --epochs 5 --img 224 --cache
        
        classify/train: model=yolov5s-cls.pt, data=imagenette160, epochs=5, batch_size=64, imgsz=224, nosave=False, cache=ram, device=, workers=8, project=runs/train-cls, name=exp, exist_ok=False, pretrained=True, optimizer=Adam, lr0=0.001, decay=5e-05, label_smoothing=0.1, cutoff=None, dropout=None, verbose=False, seed=0, local_rank=-1
        github: up to date with https://github.com/ultralytics/yolov5 ✅
        YOLOv5 🚀 v7.0-3-g61ebf5e Python-3.7.15 torch-1.12.1+cu113 CUDA:0 (Tesla T4, 15110MiB)
        
        TensorBoard: Start with 'tensorboard --logdir runs/train-cls', view at http://localhost:6006/
        
        Dataset not found ⚠️, missing path /content/datasets/imagenette160, attempting download...
        Downloading https://github.com/ultralytics/yolov5/releases/download/v1.0/imagenette160.zip to /content/datasets/imagenette160.zip...
        100% 103M/103M [00:00<00:00, 347MB/s] 
        Unzipping /content/datasets/imagenette160.zip...
        Dataset download success ✅ (3.3s), saved to /content/datasets/imagenette160
        
        albumentations: RandomResizedCrop(p=1.0, height=224, width=224, scale=(0.08, 1.0), ratio=(0.75, 1.3333333333333333), interpolation=1), HorizontalFlip(p=0.5), ColorJitter(p=0.5, brightness=[0.6, 1.4], contrast=[0.6, 1.4], saturation=[0.6, 1.4], hue=[0, 0]), Normalize(p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0), ToTensorV2(always_apply=True, p=1.0, transpose_mask=False)
        Model summary: 149 layers, 4185290 parameters, 4185290 gradients, 10.5 GFLOPs
        optimizer: Adam(lr=0.001) with parameter groups 32 weight(decay=0.0), 33 weight(decay=5e-05), 33 bias
        Image sizes 224 train, 224 test
        Using 1 dataloader workers
        Logging results to runs/train-cls/exp
        Starting yolov5s-cls.pt training on imagenette160 dataset with 10 classes for 5 epochs...
        
             Epoch   GPU_mem  train_loss    val_loss    top1_acc    top5_acc
               1/5     1.47G        1.05       0.974       0.828       0.975: 100% 148/148 [00:38<00:00,  3.82it/s]
               2/5     1.73G       0.895       0.766       0.911       0.994: 100% 148/148 [00:36<00:00,  4.03it/s]
               3/5     1.73G        0.82       0.704       0.934       0.996: 100% 148/148 [00:35<00:00,  4.20it/s]
               4/5     1.73G       0.766       0.664       0.951       0.998: 100% 148/148 [00:36<00:00,  4.05it/s]
               5/5     1.73G       0.724       0.634       0.959       0.997: 100% 148/148 [00:37<00:00,  3.94it/s]
        
        Training complete (0.052 hours)
        Results saved to runs/train-cls/exp
        Predict:         python classify/predict.py --weights runs/train-cls/exp/weights/best.pt --source im.jpg
        Validate:        python classify/val.py --weights runs/train-cls/exp/weights/best.pt --data /content/datasets/imagenette160
        Export:          python export.py --weights runs/train-cls/exp/weights/best.pt --include onnx
        PyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', 'runs/train-cls/exp/weights/best.pt')
        Visualize:       https://netron.app
        
        

        4. Visualize

        Comet Logging and Visualization 🌟 NEW

        Comet is now fully integrated with YOLOv5. Track and visualize model metrics in real time, save your hyperparameters, datasets, and model checkpoints, and visualize your model predictions with Comet Custom Panels! Comet makes sure you never lose track of your work and makes it easy to share results and collaborate across teams of all sizes!

        Getting started is easy:

        pip install comet_ml  # 1. install
        export COMET_API_KEY=<Your API Key>  # 2. paste API key
        python train.py --img 640 --epochs 3 --data coco128.yaml --weights yolov5s.pt  # 3. train
        

        To learn more about all of the supported Comet features for this integration, check out the Comet Tutorial. If you'd like to learn more about Comet, head over to our documentation. Get started by trying out the Comet Colab Notebook:
        Open In Colab


        Comet Dashboard

        ClearML Logging and Automation 🌟 NEW

        ClearML is completely integrated into YOLOv5 to track your experimentation, manage dataset versions and even remotely execute training runs. To enable ClearML (check cells above):

        You'll get all the great expected features from an experiment manager: live updates, model upload, experiment comparison etc. but ClearML also tracks uncommitted changes and installed packages for example. Thanks to that ClearML Tasks (which is what we call experiments) are also reproducible on different machines! With only 1 extra line, we can schedule a YOLOv5 training task on a queue to be executed by any number of ClearML Agents (workers).

        You can use ClearML Data to version your dataset and then pass it to YOLOv5 simply using its unique ID. This will help you keep track of your data without adding extra hassle. Explore the ClearML Tutorial for details!


        ClearML Experiment Management UI

        Local Logging

        Training results are automatically logged with Tensorboard and CSV loggers to runs/train, with a new experiment directory created for each new training as runs/train/exp2, runs/train/exp3, etc.

        This directory contains train and val statistics, mosaics, labels, predictions and augmentated mosaics, as well as metrics and charts including precision-recall (PR) curves and confusion matrices.

        Local logging results

        Environments

        YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

        Status

        YOLOv5 CI

        If this badge is green, all YOLOv5 GitHub Actions Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training (train.py), testing (val.py), inference (detect.py) and export (export.py) on macOS, Windows, and Ubuntu every 24 hours and on every commit.

        Appendix

        Additional content below.

        # YOLOv5 PyTorch HUB Inference (DetectionModels only)
        import torch
        
        model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True, trust_repo=True)  # or yolov5n - yolov5x6 or custom
        im = 'https://ultralytics.com/images/zidane.jpg'  # file, Path, PIL.Image, OpenCV, nparray, list
        results = model(im)  # inference
        results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
        
      • 📄 val.py 7.9KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """
        Validate a trained YOLOv5 classification model on a classification dataset.
        
        Usage:
            $ bash data/scripts/get_imagenet.sh --val  # download ImageNet val split (6.3G, 50000 images)
            $ python classify/val.py --weights yolov5m-cls.pt --data ../datasets/imagenet --img 224  # validate ImageNet
        
        Usage - formats:
            $ python classify/val.py --weights yolov5s-cls.pt                 # PyTorch
                                               yolov5s-cls.torchscript        # TorchScript
                                               yolov5s-cls.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                               yolov5s-cls_openvino_model     # OpenVINO
                                               yolov5s-cls.engine             # TensorRT
                                               yolov5s-cls.mlmodel            # CoreML (macOS-only)
                                               yolov5s-cls_saved_model        # TensorFlow SavedModel
                                               yolov5s-cls.pb                 # TensorFlow GraphDef
                                               yolov5s-cls.tflite             # TensorFlow Lite
                                               yolov5s-cls_edgetpu.tflite     # TensorFlow Edge TPU
                                               yolov5s-cls_paddle_model       # PaddlePaddle
        """
        
        import argparse
        import os
        import sys
        from pathlib import Path
        
        import torch
        from tqdm import tqdm
        
        FILE = Path(__file__).resolve()
        ROOT = FILE.parents[1]  # YOLOv5 root directory
        if str(ROOT) not in sys.path:
            sys.path.append(str(ROOT))  # add ROOT to PATH
        ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
        
        from models.common import DetectMultiBackend
        from utils.dataloaders import create_classification_dataloader
        from utils.general import (
            LOGGER,
            TQDM_BAR_FORMAT,
            Profile,
            check_img_size,
            check_requirements,
            colorstr,
            increment_path,
            print_args,
        )
        from utils.torch_utils import select_device, smart_inference_mode
        
        
        @smart_inference_mode()
        def run(
            data=ROOT / "../datasets/mnist",  # dataset dir
            weights=ROOT / "yolov5s-cls.pt",  # model.pt path(s)
            batch_size=128,  # batch size
            imgsz=224,  # inference size (pixels)
            device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
            workers=8,  # max dataloader workers (per RANK in DDP mode)
            verbose=False,  # verbose output
            project=ROOT / "runs/val-cls",  # save to project/name
            name="exp",  # save to project/name
            exist_ok=False,  # existing project/name ok, do not increment
            half=False,  # use FP16 half-precision inference
            dnn=False,  # use OpenCV DNN for ONNX inference
            model=None,
            dataloader=None,
            criterion=None,
            pbar=None,
        ):
            # Initialize/load model and set device
            training = model is not None
            if training:  # called by train.py
                device, pt, jit, engine = next(model.parameters()).device, True, False, False  # get model device, PyTorch model
                half &= device.type != "cpu"  # half precision only supported on CUDA
                model.half() if half else model.float()
            else:  # called directly
                device = select_device(device, batch_size=batch_size)
        
                # Directories
                save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
                save_dir.mkdir(parents=True, exist_ok=True)  # make dir
        
                # Load model
                model = DetectMultiBackend(weights, device=device, dnn=dnn, fp16=half)
                stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
                imgsz = check_img_size(imgsz, s=stride)  # check image size
                half = model.fp16  # FP16 supported on limited backends with CUDA
                if engine:
                    batch_size = model.batch_size
                else:
                    device = model.device
                    if not (pt or jit):
                        batch_size = 1  # export.py models default to batch-size 1
                        LOGGER.info(f"Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models")
        
                # Dataloader
                data = Path(data)
                test_dir = data / "test" if (data / "test").exists() else data / "val"  # data/test or data/val
                dataloader = create_classification_dataloader(
                    path=test_dir, imgsz=imgsz, batch_size=batch_size, augment=False, rank=-1, workers=workers
                )
        
            model.eval()
            pred, targets, loss, dt = [], [], 0, (Profile(device=device), Profile(device=device), Profile(device=device))
            n = len(dataloader)  # number of batches
            action = "validating" if dataloader.dataset.root.stem == "val" else "testing"
            desc = f"{pbar.desc[:-36]}{action:>36}" if pbar else f"{action}"
            bar = tqdm(dataloader, desc, n, not training, bar_format=TQDM_BAR_FORMAT, position=0)
            with torch.cuda.amp.autocast(enabled=device.type != "cpu"):
                for images, labels in bar:
                    with dt[0]:
                        images, labels = images.to(device, non_blocking=True), labels.to(device)
        
                    with dt[1]:
                        y = model(images)
        
                    with dt[2]:
                        pred.append(y.argsort(1, descending=True)[:, :5])
                        targets.append(labels)
                        if criterion:
                            loss += criterion(y, labels)
        
            loss /= n
            pred, targets = torch.cat(pred), torch.cat(targets)
            correct = (targets[:, None] == pred).float()
            acc = torch.stack((correct[:, 0], correct.max(1).values), dim=1)  # (top1, top5) accuracy
            top1, top5 = acc.mean(0).tolist()
        
            if pbar:
                pbar.desc = f"{pbar.desc[:-36]}{loss:>12.3g}{top1:>12.3g}{top5:>12.3g}"
            if verbose:  # all classes
                LOGGER.info(f"{'Class':>24}{'Images':>12}{'top1_acc':>12}{'top5_acc':>12}")
                LOGGER.info(f"{'all':>24}{targets.shape[0]:>12}{top1:>12.3g}{top5:>12.3g}")
                for i, c in model.names.items():
                    acc_i = acc[targets == i]
                    top1i, top5i = acc_i.mean(0).tolist()
                    LOGGER.info(f"{c:>24}{acc_i.shape[0]:>12}{top1i:>12.3g}{top5i:>12.3g}")
        
                # Print results
                t = tuple(x.t / len(dataloader.dataset.samples) * 1e3 for x in dt)  # speeds per image
                shape = (1, 3, imgsz, imgsz)
                LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms post-process per image at shape {shape}" % t)
                LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
        
            return top1, top5, loss
        
        
        def parse_opt():
            """Parses and returns command line arguments for YOLOv5 model evaluation and inference settings."""
            parser = argparse.ArgumentParser()
            parser.add_argument("--data", type=str, default=ROOT / "../datasets/mnist", help="dataset path")
            parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s-cls.pt", help="model.pt path(s)")
            parser.add_argument("--batch-size", type=int, default=128, help="batch size")
            parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=224, help="inference size (pixels)")
            parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
            parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
            parser.add_argument("--verbose", nargs="?", const=True, default=True, help="verbose output")
            parser.add_argument("--project", default=ROOT / "runs/val-cls", help="save to project/name")
            parser.add_argument("--name", default="exp", help="save to project/name")
            parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
            parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
            parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
            opt = parser.parse_args()
            print_args(vars(opt))
            return opt
        
        
        def main(opt):
            """Executes the YOLOv5 model prediction workflow, handling argument parsing and requirement checks."""
            check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
            run(**vars(opt))
        
        
        if __name__ == "__main__":
            opt = parse_opt()
            main(opt)
        
    • 📁 data/
      • 📁 hyps/
        • 📄 hyp.Objects365.yaml 674B
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Hyperparameters for Objects365 training
          # python train.py --weights yolov5m.pt --data Objects365.yaml --evolve
          # See Hyperparameter Evolution tutorial for details https://github.com/ultralytics/yolov5#tutorials
          
          lr0: 0.00258
          lrf: 0.17
          momentum: 0.779
          weight_decay: 0.00058
          warmup_epochs: 1.33
          warmup_momentum: 0.86
          warmup_bias_lr: 0.0711
          box: 0.0539
          cls: 0.299
          cls_pw: 0.825
          obj: 0.632
          obj_pw: 1.0
          iou_t: 0.2
          anchor_t: 3.44
          anchors: 3.2
          fl_gamma: 0.0
          hsv_h: 0.0188
          hsv_s: 0.704
          hsv_v: 0.36
          degrees: 0.0
          translate: 0.0902
          scale: 0.491
          shear: 0.0
          perspective: 0.0
          flipud: 0.0
          fliplr: 0.5
          mosaic: 1.0
          mixup: 0.0
          copy_paste: 0.0
          
        • 📄 hyp.VOC.yaml 1.1KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Hyperparameters for VOC training
          # python train.py --batch 128 --weights yolov5m6.pt --data VOC.yaml --epochs 50 --img 512 --hyp hyp.scratch-med.yaml --evolve
          # See Hyperparameter Evolution tutorial for details https://github.com/ultralytics/yolov5#tutorials
          
          # YOLOv5 Hyperparameter Evolution Results
          # Best generation: 467
          # Last generation: 996
          #    metrics/precision,       metrics/recall,      metrics/mAP_0.5, metrics/mAP_0.5:0.95,         val/box_loss,         val/obj_loss,         val/cls_loss
          #              0.87729,              0.85125,              0.91286,              0.72664,            0.0076739,            0.0042529,            0.0013865
          
          lr0: 0.00334
          lrf: 0.15135
          momentum: 0.74832
          weight_decay: 0.00025
          warmup_epochs: 3.3835
          warmup_momentum: 0.59462
          warmup_bias_lr: 0.18657
          box: 0.02
          cls: 0.21638
          cls_pw: 0.5
          obj: 0.51728
          obj_pw: 0.67198
          iou_t: 0.2
          anchor_t: 3.3744
          fl_gamma: 0.0
          hsv_h: 0.01041
          hsv_s: 0.54703
          hsv_v: 0.27739
          degrees: 0.0
          translate: 0.04591
          scale: 0.75544
          shear: 0.0
          perspective: 0.0
          flipud: 0.0
          fliplr: 0.5
          mosaic: 0.85834
          mixup: 0.04266
          copy_paste: 0.0
          anchors: 3.412
          
        • 📄 hyp.no-augmentation.yaml 1.6KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Hyperparameters when using Albumentations frameworks
          # python train.py --hyp hyp.no-augmentation.yaml
          # See https://github.com/ultralytics/yolov5/pull/3882 for YOLOv5 + Albumentations Usage examples
          
          lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
          lrf: 0.1 # final OneCycleLR learning rate (lr0 * lrf)
          momentum: 0.937 # SGD momentum/Adam beta1
          weight_decay: 0.0005 # optimizer weight decay 5e-4
          warmup_epochs: 3.0 # warmup epochs (fractions ok)
          warmup_momentum: 0.8 # warmup initial momentum
          warmup_bias_lr: 0.1 # warmup initial bias lr
          box: 0.05 # box loss gain
          cls: 0.3 # cls loss gain
          cls_pw: 1.0 # cls BCELoss positive_weight
          obj: 0.7 # obj loss gain (scale with pixels)
          obj_pw: 1.0 # obj BCELoss positive_weight
          iou_t: 0.20 # IoU training threshold
          anchor_t: 4.0 # anchor-multiple threshold
          # anchors: 3  # anchors per output layer (0 to ignore)
          # this parameters are all zero since we want to use albumentation framework
          fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
          hsv_h: 0 # image HSV-Hue augmentation (fraction)
          hsv_s: 0 # image HSV-Saturation augmentation (fraction)
          hsv_v: 0 # image HSV-Value augmentation (fraction)
          degrees: 0.0 # image rotation (+/- deg)
          translate: 0 # image translation (+/- fraction)
          scale: 0 # image scale (+/- gain)
          shear: 0 # image shear (+/- deg)
          perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
          flipud: 0.0 # image flip up-down (probability)
          fliplr: 0.0 # image flip left-right (probability)
          mosaic: 0.0 # image mosaic (probability)
          mixup: 0.0 # image mixup (probability)
          copy_paste: 0.0 # segment copy-paste (probability)
          
        • 📄 hyp.scratch-high.yaml 1.6KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Hyperparameters for high-augmentation COCO training from scratch
          # python train.py --batch 32 --cfg yolov5m6.yaml --weights '' --data coco.yaml --img 1280 --epochs 300
          # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
          
          lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
          lrf: 0.1 # final OneCycleLR learning rate (lr0 * lrf)
          momentum: 0.937 # SGD momentum/Adam beta1
          weight_decay: 0.0005 # optimizer weight decay 5e-4
          warmup_epochs: 3.0 # warmup epochs (fractions ok)
          warmup_momentum: 0.8 # warmup initial momentum
          warmup_bias_lr: 0.1 # warmup initial bias lr
          box: 0.05 # box loss gain
          cls: 0.3 # cls loss gain
          cls_pw: 1.0 # cls BCELoss positive_weight
          obj: 0.7 # obj loss gain (scale with pixels)
          obj_pw: 1.0 # obj BCELoss positive_weight
          iou_t: 0.20 # IoU training threshold
          anchor_t: 4.0 # anchor-multiple threshold
          # anchors: 3  # anchors per output layer (0 to ignore)
          fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
          hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
          hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
          hsv_v: 0.4 # image HSV-Value augmentation (fraction)
          degrees: 0.0 # image rotation (+/- deg)
          translate: 0.1 # image translation (+/- fraction)
          scale: 0.9 # image scale (+/- gain)
          shear: 0.0 # image shear (+/- deg)
          perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
          flipud: 0.0 # image flip up-down (probability)
          fliplr: 0.5 # image flip left-right (probability)
          mosaic: 1.0 # image mosaic (probability)
          mixup: 0.1 # image mixup (probability)
          copy_paste: 0.1 # segment copy-paste (probability)
          
        • 📄 hyp.scratch-low.yaml 1.6KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Hyperparameters for low-augmentation COCO training from scratch
          # python train.py --batch 64 --cfg yolov5n6.yaml --weights '' --data coco.yaml --img 640 --epochs 300 --linear
          # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
          
          lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
          lrf: 0.01 # final OneCycleLR learning rate (lr0 * lrf)
          momentum: 0.937 # SGD momentum/Adam beta1
          weight_decay: 0.0005 # optimizer weight decay 5e-4
          warmup_epochs: 3.0 # warmup epochs (fractions ok)
          warmup_momentum: 0.8 # warmup initial momentum
          warmup_bias_lr: 0.1 # warmup initial bias lr
          box: 0.05 # box loss gain
          cls: 0.5 # cls loss gain
          cls_pw: 1.0 # cls BCELoss positive_weight
          obj: 1.0 # obj loss gain (scale with pixels)
          obj_pw: 1.0 # obj BCELoss positive_weight
          iou_t: 0.20 # IoU training threshold
          anchor_t: 4.0 # anchor-multiple threshold
          # anchors: 3  # anchors per output layer (0 to ignore)
          fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
          hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
          hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
          hsv_v: 0.4 # image HSV-Value augmentation (fraction)
          degrees: 0.0 # image rotation (+/- deg)
          translate: 0.1 # image translation (+/- fraction)
          scale: 0.5 # image scale (+/- gain)
          shear: 0.0 # image shear (+/- deg)
          perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
          flipud: 0.0 # image flip up-down (probability)
          fliplr: 0.5 # image flip left-right (probability)
          mosaic: 1.0 # image mosaic (probability)
          mixup: 0.0 # image mixup (probability)
          copy_paste: 0.0 # segment copy-paste (probability)
          
        • 📄 hyp.scratch-med.yaml 1.6KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Hyperparameters for medium-augmentation COCO training from scratch
          # python train.py --batch 32 --cfg yolov5m6.yaml --weights '' --data coco.yaml --img 1280 --epochs 300
          # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
          
          lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
          lrf: 0.1 # final OneCycleLR learning rate (lr0 * lrf)
          momentum: 0.937 # SGD momentum/Adam beta1
          weight_decay: 0.0005 # optimizer weight decay 5e-4
          warmup_epochs: 3.0 # warmup epochs (fractions ok)
          warmup_momentum: 0.8 # warmup initial momentum
          warmup_bias_lr: 0.1 # warmup initial bias lr
          box: 0.05 # box loss gain
          cls: 0.3 # cls loss gain
          cls_pw: 1.0 # cls BCELoss positive_weight
          obj: 0.7 # obj loss gain (scale with pixels)
          obj_pw: 1.0 # obj BCELoss positive_weight
          iou_t: 0.20 # IoU training threshold
          anchor_t: 4.0 # anchor-multiple threshold
          # anchors: 3  # anchors per output layer (0 to ignore)
          fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
          hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
          hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
          hsv_v: 0.4 # image HSV-Value augmentation (fraction)
          degrees: 0.0 # image rotation (+/- deg)
          translate: 0.1 # image translation (+/- fraction)
          scale: 0.9 # image scale (+/- gain)
          shear: 0.0 # image shear (+/- deg)
          perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
          flipud: 0.0 # image flip up-down (probability)
          fliplr: 0.5 # image flip left-right (probability)
          mosaic: 1.0 # image mosaic (probability)
          mixup: 0.1 # image mixup (probability)
          copy_paste: 0.0 # segment copy-paste (probability)
          
      • 📁 images/
        • ⚙️ bus.jpg 476.0KB
        • ⚙️ zidane.jpg 165.0KB
      • 📁 scripts/
        • 📄 download_weights.sh 641B
          #!/bin/bash
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Download latest models from https://github.com/ultralytics/yolov5/releases
          # Example usage: bash data/scripts/download_weights.sh
          # parent
          # └── yolov5
          #     ├── yolov5s.pt  ← downloads here
          #     ├── yolov5m.pt
          #     └── ...
          
          python - <<EOF
          from utils.downloads import attempt_download
          
          p5 = list('nsmlx')  # P5 models
          p6 = [f'{x}6' for x in p5]  # P6 models
          cls = [f'{x}-cls' for x in p5]  # classification models
          seg = [f'{x}-seg' for x in p5]  # classification models
          
          for x in p5 + p6 + cls + seg:
              attempt_download(f'weights/yolov5{x}.pt')
          
          EOF
          
        • 📄 get_coco.sh 1.5KB
          #!/bin/bash
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Download COCO 2017 dataset http://cocodataset.org
          # Example usage: bash data/scripts/get_coco.sh
          # parent
          # ├── yolov5
          # └── datasets
          #     └── coco  ← downloads here
          
          # Arguments (optional) Usage: bash data/scripts/get_coco.sh --train --val --test --segments
          if [ "$#" -gt 0 ]; then
            for opt in "$@"; do
              case "${opt}" in
              --train) train=true ;;
              --val) val=true ;;
              --test) test=true ;;
              --segments) segments=true ;;
              esac
            done
          else
            train=true
            val=true
            test=false
            segments=false
          fi
          
          # Download/unzip labels
          d='../datasets' # unzip directory
          url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
          if [ "$segments" == "true" ]; then
            f='coco2017labels-segments.zip' # 168 MB
          else
            f='coco2017labels.zip' # 46 MB
          fi
          echo 'Downloading' $url$f ' ...'
          curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
          
          # Download/unzip images
          d='../datasets/coco/images' # unzip directory
          url=http://images.cocodataset.org/zips/
          if [ "$train" == "true" ]; then
            f='train2017.zip' # 19G, 118k images
            echo 'Downloading' $url$f '...'
            curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
          fi
          if [ "$val" == "true" ]; then
            f='val2017.zip' # 1G, 5k images
            echo 'Downloading' $url$f '...'
            curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
          fi
          if [ "$test" == "true" ]; then
            f='test2017.zip' # 7G, 41k images (optional)
            echo 'Downloading' $url$f '...'
            curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
          fi
          wait # finish background tasks
          
        • 📄 get_coco128.sh 619B
          #!/bin/bash
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Download COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017)
          # Example usage: bash data/scripts/get_coco128.sh
          # parent
          # ├── yolov5
          # └── datasets
          #     └── coco128  ← downloads here
          
          # Download/unzip images and labels
          d='../datasets' # unzip directory
          url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
          f='coco128.zip' # or 'coco128-segments.zip', 68 MB
          echo 'Downloading' $url$f ' ...'
          curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
          
          wait # finish background tasks
          
        • 📄 get_imagenet.sh 1.6KB
          #!/bin/bash
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Download ILSVRC2012 ImageNet dataset https://image-net.org
          # Example usage: bash data/scripts/get_imagenet.sh
          # parent
          # ├── yolov5
          # └── datasets
          #     └── imagenet  ← downloads here
          
          # Arguments (optional) Usage: bash data/scripts/get_imagenet.sh --train --val
          if [ "$#" -gt 0 ]; then
            for opt in "$@"; do
              case "${opt}" in
              --train) train=true ;;
              --val) val=true ;;
              esac
            done
          else
            train=true
            val=true
          fi
          
          # Make dir
          d='../datasets/imagenet' # unzip directory
          mkdir -p $d && cd $d
          
          # Download/unzip train
          if [ "$train" == "true" ]; then
            wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_train.tar # download 138G, 1281167 images
            mkdir train && mv ILSVRC2012_img_train.tar train/ && cd train
            tar -xf ILSVRC2012_img_train.tar && rm -f ILSVRC2012_img_train.tar
            find . -name "*.tar" | while read NAME; do
              mkdir -p "${NAME%.tar}"
              tar -xf "${NAME}" -C "${NAME%.tar}"
              rm -f "${NAME}"
            done
            cd ..
          fi
          
          # Download/unzip val
          if [ "$val" == "true" ]; then
            wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar # download 6.3G, 50000 images
            mkdir val && mv ILSVRC2012_img_val.tar val/ && cd val && tar -xf ILSVRC2012_img_val.tar
            wget -qO- https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh | bash # move into subdirs
          fi
          
          # Delete corrupted image (optional: PNG under JPEG name that may cause dataloaders to fail)
          # rm train/n04266014/n04266014_10835.JPEG
          
          # TFRecords (optional)
          # wget https://raw.githubusercontent.com/tensorflow/models/master/research/slim/datasets/imagenet_lsvrc_2015_synsets.txt
          
        • 📄 get_imagenet10.sh 734B
          #!/bin/bash
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Download ILSVRC2012 ImageNet dataset https://image-net.org
          # Example usage: bash data/scripts/get_imagenet.sh
          # parent
          # ├── yolov5
          # └── datasets
          #     └── imagenet  ← downloads here
          
          # Arguments (optional) Usage: bash data/scripts/get_imagenet.sh --train --val
          if [ "$#" -gt 0 ]; then
            for opt in "$@"; do
              case "${opt}" in
              --train) train=true ;;
              --val) val=true ;;
              esac
            done
          else
            train=true
            val=true
          fi
          
          # Make dir
          d='../datasets/imagenet10' # unzip directory
          mkdir -p $d && cd $d
          
          # Download/unzip train
          wget https://github.com/ultralytics/yolov5/releases/download/v1.0/imagenet10.zip
          unzip imagenet10.zip && rm imagenet10.zip
          
        • 📄 get_imagenet100.sh 738B
          #!/bin/bash
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Download ILSVRC2012 ImageNet dataset https://image-net.org
          # Example usage: bash data/scripts/get_imagenet.sh
          # parent
          # ├── yolov5
          # └── datasets
          #     └── imagenet  ← downloads here
          
          # Arguments (optional) Usage: bash data/scripts/get_imagenet.sh --train --val
          if [ "$#" -gt 0 ]; then
            for opt in "$@"; do
              case "${opt}" in
              --train) train=true ;;
              --val) val=true ;;
              esac
            done
          else
            train=true
            val=true
          fi
          
          # Make dir
          d='../datasets/imagenet100' # unzip directory
          mkdir -p $d && cd $d
          
          # Download/unzip train
          wget https://github.com/ultralytics/yolov5/releases/download/v1.0/imagenet100.zip
          unzip imagenet100.zip && rm imagenet100.zip
          
        • 📄 get_imagenet1000.sh 742B
          #!/bin/bash
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Download ILSVRC2012 ImageNet dataset https://image-net.org
          # Example usage: bash data/scripts/get_imagenet.sh
          # parent
          # ├── yolov5
          # └── datasets
          #     └── imagenet  ← downloads here
          
          # Arguments (optional) Usage: bash data/scripts/get_imagenet.sh --train --val
          if [ "$#" -gt 0 ]; then
            for opt in "$@"; do
              case "${opt}" in
              --train) train=true ;;
              --val) val=true ;;
              esac
            done
          else
            train=true
            val=true
          fi
          
          # Make dir
          d='../datasets/imagenet1000' # unzip directory
          mkdir -p $d && cd $d
          
          # Download/unzip train
          wget https://github.com/ultralytics/yolov5/releases/download/v1.0/imagenet1000.zip
          unzip imagenet1000.zip && rm imagenet1000.zip
          
      • 📄 AI_coco128.yaml 1.8KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
        # Example usage: python train.py --data coco128.yaml
        # parent
        # ├── yolov5
        # └── datasets
        #     └── coco128  ← downloads here (7 MB)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/coco128 # dataset root dir
        train: images/train2017 # train images (relative to 'path') 128 images
        val: images/train2017 # val images (relative to 'path') 128 images
        test: # test images (optional)
        
        # Classes
        names:
          0: person
          1: bicycle
          2: car
          3: motorcycle
          4: airplane
          5: bus
          6: train
          7: truck
          8: boat
          9: traffic light
          10: fire hydrant
          11: stop sign
          12: parking meter
          13: bench
          14: bird
          15: cat
          16: dog
          17: horse
          18: sheep
          19: cow
          20: elephant
          21: bear
          22: zebra
          23: giraffe
          24: backpack
          25: umbrella
          26: handbag
          27: tie
          28: suitcase
          29: frisbee
          30: skis
          31: snowboard
          32: sports ball
          33: kite
          34: baseball bat
          35: baseball glove
          36: skateboard
          37: surfboard
          38: tennis racket
          39: bottle
          40: wine glass
          41: cup
          42: fork
          43: knife
          44: spoon
          45: bowl
          46: banana
          47: apple
          48: sandwich
          49: orange
          50: broccoli
          51: carrot
          52: hot dog
          53: pizza
          54: donut
          55: cake
          56: chair
          57: couch
          58: potted plant
          59: bed
          60: dining table
          61: toilet
          62: tv
          63: laptop
          64: mouse
          65: remote
          66: keyboard
          67: cell phone
          68: microwave
          69: oven
          70: toaster
          71: sink
          72: refrigerator
          73: book
          74: clock
          75: vase
          76: scissors
          77: teddy bear
          78: hair drier
          79: toothbrush
        
        # Download script/URL (optional)
        download: https://ultralytics.com/assets/coco128.zip
        
      • 📄 Argoverse.yaml 2.7KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/ by Argo AI
        # Example usage: python train.py --data Argoverse.yaml
        # parent
        # ├── yolov5
        # └── datasets
        #     └── Argoverse  ← downloads here (31.3 GB)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/Argoverse # dataset root dir
        train: Argoverse-1.1/images/train/ # train images (relative to 'path') 39384 images
        val: Argoverse-1.1/images/val/ # val images (relative to 'path') 15062 images
        test: Argoverse-1.1/images/test/ # test images (optional) https://eval.ai/web/challenges/challenge-page/800/overview
        
        # Classes
        names:
          0: person
          1: bicycle
          2: car
          3: motorcycle
          4: bus
          5: truck
          6: traffic_light
          7: stop_sign
        
        # Download script/URL (optional) ---------------------------------------------------------------------------------------
        download: |
          import json
        
          from tqdm import tqdm
          from utils.general import download, Path
        
        
          def argoverse2yolo(set):
              labels = {}
              a = json.load(open(set, "rb"))
              for annot in tqdm(a['annotations'], desc=f"Converting {set} to YOLOv5 format..."):
                  img_id = annot['image_id']
                  img_name = a['images'][img_id]['name']
                  img_label_name = f'{img_name[:-3]}txt'
        
                  cls = annot['category_id']  # instance class id
                  x_center, y_center, width, height = annot['bbox']
                  x_center = (x_center + width / 2) / 1920.0  # offset and scale
                  y_center = (y_center + height / 2) / 1200.0  # offset and scale
                  width /= 1920.0  # scale
                  height /= 1200.0  # scale
        
                  img_dir = set.parents[2] / 'Argoverse-1.1' / 'labels' / a['seq_dirs'][a['images'][annot['image_id']]['sid']]
                  if not img_dir.exists():
                      img_dir.mkdir(parents=True, exist_ok=True)
        
                  k = str(img_dir / img_label_name)
                  if k not in labels:
                      labels[k] = []
                  labels[k].append(f"{cls} {x_center} {y_center} {width} {height}\n")
        
              for k in labels:
                  with open(k, "w") as f:
                      f.writelines(labels[k])
        
        
          # Download
          dir = Path(yaml['path'])  # dataset root dir
          urls = ['https://argoverse-hd.s3.us-east-2.amazonaws.com/Argoverse-HD-Full.zip']
          download(urls, dir=dir, delete=False)
        
          # Convert
          annotations_dir = 'Argoverse-HD/annotations/'
          (dir / 'Argoverse-1.1' / 'tracking').rename(dir / 'Argoverse-1.1' / 'images')  # rename 'tracking' to 'images'
          for d in "train.json", "val.json":
              argoverse2yolo(dir / annotations_dir / d)  # convert VisDrone annotations to YOLO labels
        
      • 📄 GlobalWheat2020.yaml 1.8KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # Global Wheat 2020 dataset http://www.global-wheat.com/ by University of Saskatchewan
        # Example usage: python train.py --data GlobalWheat2020.yaml
        # parent
        # ├── yolov5
        # └── datasets
        #     └── GlobalWheat2020  ← downloads here (7.0 GB)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/GlobalWheat2020 # dataset root dir
        train: # train images (relative to 'path') 3422 images
          - images/arvalis_1
          - images/arvalis_2
          - images/arvalis_3
          - images/ethz_1
          - images/rres_1
          - images/inrae_1
          - images/usask_1
        val: # val images (relative to 'path') 748 images (WARNING: train set contains ethz_1)
          - images/ethz_1
        test: # test images (optional) 1276 images
          - images/utokyo_1
          - images/utokyo_2
          - images/nau_1
          - images/uq_1
        
        # Classes
        names:
          0: wheat_head
        
        # Download script/URL (optional) ---------------------------------------------------------------------------------------
        download: |
          from utils.general import download, Path
        
        
          # Download
          dir = Path(yaml['path'])  # dataset root dir
          urls = ['https://zenodo.org/record/4298502/files/global-wheat-codalab-official.zip',
                  'https://github.com/ultralytics/yolov5/releases/download/v1.0/GlobalWheat2020_labels.zip']
          download(urls, dir=dir)
        
          # Make Directories
          for p in 'annotations', 'images', 'labels':
              (dir / p).mkdir(parents=True, exist_ok=True)
        
          # Move
          for p in 'arvalis_1', 'arvalis_2', 'arvalis_3', 'ethz_1', 'rres_1', 'inrae_1', 'usask_1', \
                   'utokyo_1', 'utokyo_2', 'nau_1', 'uq_1':
              (dir / p).rename(dir / 'images' / p)  # move to /images
              f = (dir / p).with_suffix('.json')  # json file
              if f.exists():
                  f.rename((dir / 'annotations' / p).with_suffix('.json'))  # move to /annotations
        
      • 📄 ImageNet.yaml 18.4KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # ImageNet-1k dataset https://www.image-net.org/index.php by Stanford University
        # Simplified class names from https://github.com/anishathalye/imagenet-simple-labels
        # Example usage: python classify/train.py --data imagenet
        # parent
        # ├── yolov5
        # └── datasets
        #     └── imagenet  ← downloads here (144 GB)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/imagenet # dataset root dir
        train: train # train images (relative to 'path') 1281167 images
        val: val # val images (relative to 'path') 50000 images
        test: # test images (optional)
        
        # Classes
        names:
          0: tench
          1: goldfish
          2: great white shark
          3: tiger shark
          4: hammerhead shark
          5: electric ray
          6: stingray
          7: cock
          8: hen
          9: ostrich
          10: brambling
          11: goldfinch
          12: house finch
          13: junco
          14: indigo bunting
          15: American robin
          16: bulbul
          17: jay
          18: magpie
          19: chickadee
          20: American dipper
          21: kite
          22: bald eagle
          23: vulture
          24: great grey owl
          25: fire salamander
          26: smooth newt
          27: newt
          28: spotted salamander
          29: axolotl
          30: American bullfrog
          31: tree frog
          32: tailed frog
          33: loggerhead sea turtle
          34: leatherback sea turtle
          35: mud turtle
          36: terrapin
          37: box turtle
          38: banded gecko
          39: green iguana
          40: Carolina anole
          41: desert grassland whiptail lizard
          42: agama
          43: frilled-necked lizard
          44: alligator lizard
          45: Gila monster
          46: European green lizard
          47: chameleon
          48: Komodo dragon
          49: Nile crocodile
          50: American alligator
          51: triceratops
          52: worm snake
          53: ring-necked snake
          54: eastern hog-nosed snake
          55: smooth green snake
          56: kingsnake
          57: garter snake
          58: water snake
          59: vine snake
          60: night snake
          61: boa constrictor
          62: African rock python
          63: Indian cobra
          64: green mamba
          65: sea snake
          66: Saharan horned viper
          67: eastern diamondback rattlesnake
          68: sidewinder
          69: trilobite
          70: harvestman
          71: scorpion
          72: yellow garden spider
          73: barn spider
          74: European garden spider
          75: southern black widow
          76: tarantula
          77: wolf spider
          78: tick
          79: centipede
          80: black grouse
          81: ptarmigan
          82: ruffed grouse
          83: prairie grouse
          84: peacock
          85: quail
          86: partridge
          87: grey parrot
          88: macaw
          89: sulphur-crested cockatoo
          90: lorikeet
          91: coucal
          92: bee eater
          93: hornbill
          94: hummingbird
          95: jacamar
          96: toucan
          97: duck
          98: red-breasted merganser
          99: goose
          100: black swan
          101: tusker
          102: echidna
          103: platypus
          104: wallaby
          105: koala
          106: wombat
          107: jellyfish
          108: sea anemone
          109: brain coral
          110: flatworm
          111: nematode
          112: conch
          113: snail
          114: slug
          115: sea slug
          116: chiton
          117: chambered nautilus
          118: Dungeness crab
          119: rock crab
          120: fiddler crab
          121: red king crab
          122: American lobster
          123: spiny lobster
          124: crayfish
          125: hermit crab
          126: isopod
          127: white stork
          128: black stork
          129: spoonbill
          130: flamingo
          131: little blue heron
          132: great egret
          133: bittern
          134: crane (bird)
          135: limpkin
          136: common gallinule
          137: American coot
          138: bustard
          139: ruddy turnstone
          140: dunlin
          141: common redshank
          142: dowitcher
          143: oystercatcher
          144: pelican
          145: king penguin
          146: albatross
          147: grey whale
          148: killer whale
          149: dugong
          150: sea lion
          151: Chihuahua
          152: Japanese Chin
          153: Maltese
          154: Pekingese
          155: Shih Tzu
          156: King Charles Spaniel
          157: Papillon
          158: toy terrier
          159: Rhodesian Ridgeback
          160: Afghan Hound
          161: Basset Hound
          162: Beagle
          163: Bloodhound
          164: Bluetick Coonhound
          165: Black and Tan Coonhound
          166: Treeing Walker Coonhound
          167: English foxhound
          168: Redbone Coonhound
          169: borzoi
          170: Irish Wolfhound
          171: Italian Greyhound
          172: Whippet
          173: Ibizan Hound
          174: Norwegian Elkhound
          175: Otterhound
          176: Saluki
          177: Scottish Deerhound
          178: Weimaraner
          179: Staffordshire Bull Terrier
          180: American Staffordshire Terrier
          181: Bedlington Terrier
          182: Border Terrier
          183: Kerry Blue Terrier
          184: Irish Terrier
          185: Norfolk Terrier
          186: Norwich Terrier
          187: Yorkshire Terrier
          188: Wire Fox Terrier
          189: Lakeland Terrier
          190: Sealyham Terrier
          191: Airedale Terrier
          192: Cairn Terrier
          193: Australian Terrier
          194: Dandie Dinmont Terrier
          195: Boston Terrier
          196: Miniature Schnauzer
          197: Giant Schnauzer
          198: Standard Schnauzer
          199: Scottish Terrier
          200: Tibetan Terrier
          201: Australian Silky Terrier
          202: Soft-coated Wheaten Terrier
          203: West Highland White Terrier
          204: Lhasa Apso
          205: Flat-Coated Retriever
          206: Curly-coated Retriever
          207: Golden Retriever
          208: Labrador Retriever
          209: Chesapeake Bay Retriever
          210: German Shorthaired Pointer
          211: Vizsla
          212: English Setter
          213: Irish Setter
          214: Gordon Setter
          215: Brittany
          216: Clumber Spaniel
          217: English Springer Spaniel
          218: Welsh Springer Spaniel
          219: Cocker Spaniels
          220: Sussex Spaniel
          221: Irish Water Spaniel
          222: Kuvasz
          223: Schipperke
          224: Groenendael
          225: Malinois
          226: Briard
          227: Australian Kelpie
          228: Komondor
          229: Old English Sheepdog
          230: Shetland Sheepdog
          231: collie
          232: Border Collie
          233: Bouvier des Flandres
          234: Rottweiler
          235: German Shepherd Dog
          236: Dobermann
          237: Miniature Pinscher
          238: Greater Swiss Mountain Dog
          239: Bernese Mountain Dog
          240: Appenzeller Sennenhund
          241: Entlebucher Sennenhund
          242: Boxer
          243: Bullmastiff
          244: Tibetan Mastiff
          245: French Bulldog
          246: Great Dane
          247: St. Bernard
          248: husky
          249: Alaskan Malamute
          250: Siberian Husky
          251: Dalmatian
          252: Affenpinscher
          253: Basenji
          254: pug
          255: Leonberger
          256: Newfoundland
          257: Pyrenean Mountain Dog
          258: Samoyed
          259: Pomeranian
          260: Chow Chow
          261: Keeshond
          262: Griffon Bruxellois
          263: Pembroke Welsh Corgi
          264: Cardigan Welsh Corgi
          265: Toy Poodle
          266: Miniature Poodle
          267: Standard Poodle
          268: Mexican hairless dog
          269: grey wolf
          270: Alaskan tundra wolf
          271: red wolf
          272: coyote
          273: dingo
          274: dhole
          275: African wild dog
          276: hyena
          277: red fox
          278: kit fox
          279: Arctic fox
          280: grey fox
          281: tabby cat
          282: tiger cat
          283: Persian cat
          284: Siamese cat
          285: Egyptian Mau
          286: cougar
          287: lynx
          288: leopard
          289: snow leopard
          290: jaguar
          291: lion
          292: tiger
          293: cheetah
          294: brown bear
          295: American black bear
          296: polar bear
          297: sloth bear
          298: mongoose
          299: meerkat
          300: tiger beetle
          301: ladybug
          302: ground beetle
          303: longhorn beetle
          304: leaf beetle
          305: dung beetle
          306: rhinoceros beetle
          307: weevil
          308: fly
          309: bee
          310: ant
          311: grasshopper
          312: cricket
          313: stick insect
          314: cockroach
          315: mantis
          316: cicada
          317: leafhopper
          318: lacewing
          319: dragonfly
          320: damselfly
          321: red admiral
          322: ringlet
          323: monarch butterfly
          324: small white
          325: sulphur butterfly
          326: gossamer-winged butterfly
          327: starfish
          328: sea urchin
          329: sea cucumber
          330: cottontail rabbit
          331: hare
          332: Angora rabbit
          333: hamster
          334: porcupine
          335: fox squirrel
          336: marmot
          337: beaver
          338: guinea pig
          339: common sorrel
          340: zebra
          341: pig
          342: wild boar
          343: warthog
          344: hippopotamus
          345: ox
          346: water buffalo
          347: bison
          348: ram
          349: bighorn sheep
          350: Alpine ibex
          351: hartebeest
          352: impala
          353: gazelle
          354: dromedary
          355: llama
          356: weasel
          357: mink
          358: European polecat
          359: black-footed ferret
          360: otter
          361: skunk
          362: badger
          363: armadillo
          364: three-toed sloth
          365: orangutan
          366: gorilla
          367: chimpanzee
          368: gibbon
          369: siamang
          370: guenon
          371: patas monkey
          372: baboon
          373: macaque
          374: langur
          375: black-and-white colobus
          376: proboscis monkey
          377: marmoset
          378: white-headed capuchin
          379: howler monkey
          380: titi
          381: Geoffroy's spider monkey
          382: common squirrel monkey
          383: ring-tailed lemur
          384: indri
          385: Asian elephant
          386: African bush elephant
          387: red panda
          388: giant panda
          389: snoek
          390: eel
          391: coho salmon
          392: rock beauty
          393: clownfish
          394: sturgeon
          395: garfish
          396: lionfish
          397: pufferfish
          398: abacus
          399: abaya
          400: academic gown
          401: accordion
          402: acoustic guitar
          403: aircraft carrier
          404: airliner
          405: airship
          406: altar
          407: ambulance
          408: amphibious vehicle
          409: analog clock
          410: apiary
          411: apron
          412: waste container
          413: assault rifle
          414: backpack
          415: bakery
          416: balance beam
          417: balloon
          418: ballpoint pen
          419: Band-Aid
          420: banjo
          421: baluster
          422: barbell
          423: barber chair
          424: barbershop
          425: barn
          426: barometer
          427: barrel
          428: wheelbarrow
          429: baseball
          430: basketball
          431: bassinet
          432: bassoon
          433: swimming cap
          434: bath towel
          435: bathtub
          436: station wagon
          437: lighthouse
          438: beaker
          439: military cap
          440: beer bottle
          441: beer glass
          442: bell-cot
          443: bib
          444: tandem bicycle
          445: bikini
          446: ring binder
          447: binoculars
          448: birdhouse
          449: boathouse
          450: bobsleigh
          451: bolo tie
          452: poke bonnet
          453: bookcase
          454: bookstore
          455: bottle cap
          456: bow
          457: bow tie
          458: brass
          459: bra
          460: breakwater
          461: breastplate
          462: broom
          463: bucket
          464: buckle
          465: bulletproof vest
          466: high-speed train
          467: butcher shop
          468: taxicab
          469: cauldron
          470: candle
          471: cannon
          472: canoe
          473: can opener
          474: cardigan
          475: car mirror
          476: carousel
          477: tool kit
          478: carton
          479: car wheel
          480: automated teller machine
          481: cassette
          482: cassette player
          483: castle
          484: catamaran
          485: CD player
          486: cello
          487: mobile phone
          488: chain
          489: chain-link fence
          490: chain mail
          491: chainsaw
          492: chest
          493: chiffonier
          494: chime
          495: china cabinet
          496: Christmas stocking
          497: church
          498: movie theater
          499: cleaver
          500: cliff dwelling
          501: cloak
          502: clogs
          503: cocktail shaker
          504: coffee mug
          505: coffeemaker
          506: coil
          507: combination lock
          508: computer keyboard
          509: confectionery store
          510: container ship
          511: convertible
          512: corkscrew
          513: cornet
          514: cowboy boot
          515: cowboy hat
          516: cradle
          517: crane (machine)
          518: crash helmet
          519: crate
          520: infant bed
          521: Crock Pot
          522: croquet ball
          523: crutch
          524: cuirass
          525: dam
          526: desk
          527: desktop computer
          528: rotary dial telephone
          529: diaper
          530: digital clock
          531: digital watch
          532: dining table
          533: dishcloth
          534: dishwasher
          535: disc brake
          536: dock
          537: dog sled
          538: dome
          539: doormat
          540: drilling rig
          541: drum
          542: drumstick
          543: dumbbell
          544: Dutch oven
          545: electric fan
          546: electric guitar
          547: electric locomotive
          548: entertainment center
          549: envelope
          550: espresso machine
          551: face powder
          552: feather boa
          553: filing cabinet
          554: fireboat
          555: fire engine
          556: fire screen sheet
          557: flagpole
          558: flute
          559: folding chair
          560: football helmet
          561: forklift
          562: fountain
          563: fountain pen
          564: four-poster bed
          565: freight car
          566: French horn
          567: frying pan
          568: fur coat
          569: garbage truck
          570: gas mask
          571: gas pump
          572: goblet
          573: go-kart
          574: golf ball
          575: golf cart
          576: gondola
          577: gong
          578: gown
          579: grand piano
          580: greenhouse
          581: grille
          582: grocery store
          583: guillotine
          584: barrette
          585: hair spray
          586: half-track
          587: hammer
          588: hamper
          589: hair dryer
          590: hand-held computer
          591: handkerchief
          592: hard disk drive
          593: harmonica
          594: harp
          595: harvester
          596: hatchet
          597: holster
          598: home theater
          599: honeycomb
          600: hook
          601: hoop skirt
          602: horizontal bar
          603: horse-drawn vehicle
          604: hourglass
          605: iPod
          606: clothes iron
          607: jack-o'-lantern
          608: jeans
          609: jeep
          610: T-shirt
          611: jigsaw puzzle
          612: pulled rickshaw
          613: joystick
          614: kimono
          615: knee pad
          616: knot
          617: lab coat
          618: ladle
          619: lampshade
          620: laptop computer
          621: lawn mower
          622: lens cap
          623: paper knife
          624: library
          625: lifeboat
          626: lighter
          627: limousine
          628: ocean liner
          629: lipstick
          630: slip-on shoe
          631: lotion
          632: speaker
          633: loupe
          634: sawmill
          635: magnetic compass
          636: mail bag
          637: mailbox
          638: tights
          639: tank suit
          640: manhole cover
          641: maraca
          642: marimba
          643: mask
          644: match
          645: maypole
          646: maze
          647: measuring cup
          648: medicine chest
          649: megalith
          650: microphone
          651: microwave oven
          652: military uniform
          653: milk can
          654: minibus
          655: miniskirt
          656: minivan
          657: missile
          658: mitten
          659: mixing bowl
          660: mobile home
          661: Model T
          662: modem
          663: monastery
          664: monitor
          665: moped
          666: mortar
          667: square academic cap
          668: mosque
          669: mosquito net
          670: scooter
          671: mountain bike
          672: tent
          673: computer mouse
          674: mousetrap
          675: moving van
          676: muzzle
          677: nail
          678: neck brace
          679: necklace
          680: nipple
          681: notebook computer
          682: obelisk
          683: oboe
          684: ocarina
          685: odometer
          686: oil filter
          687: organ
          688: oscilloscope
          689: overskirt
          690: bullock cart
          691: oxygen mask
          692: packet
          693: paddle
          694: paddle wheel
          695: padlock
          696: paintbrush
          697: pajamas
          698: palace
          699: pan flute
          700: paper towel
          701: parachute
          702: parallel bars
          703: park bench
          704: parking meter
          705: passenger car
          706: patio
          707: payphone
          708: pedestal
          709: pencil case
          710: pencil sharpener
          711: perfume
          712: Petri dish
          713: photocopier
          714: plectrum
          715: Pickelhaube
          716: picket fence
          717: pickup truck
          718: pier
          719: piggy bank
          720: pill bottle
          721: pillow
          722: ping-pong ball
          723: pinwheel
          724: pirate ship
          725: pitcher
          726: hand plane
          727: planetarium
          728: plastic bag
          729: plate rack
          730: plow
          731: plunger
          732: Polaroid camera
          733: pole
          734: police van
          735: poncho
          736: billiard table
          737: soda bottle
          738: pot
          739: potter's wheel
          740: power drill
          741: prayer rug
          742: printer
          743: prison
          744: projectile
          745: projector
          746: hockey puck
          747: punching bag
          748: purse
          749: quill
          750: quilt
          751: race car
          752: racket
          753: radiator
          754: radio
          755: radio telescope
          756: rain barrel
          757: recreational vehicle
          758: reel
          759: reflex camera
          760: refrigerator
          761: remote control
          762: restaurant
          763: revolver
          764: rifle
          765: rocking chair
          766: rotisserie
          767: eraser
          768: rugby ball
          769: ruler
          770: running shoe
          771: safe
          772: safety pin
          773: salt shaker
          774: sandal
          775: sarong
          776: saxophone
          777: scabbard
          778: weighing scale
          779: school bus
          780: schooner
          781: scoreboard
          782: CRT screen
          783: screw
          784: screwdriver
          785: seat belt
          786: sewing machine
          787: shield
          788: shoe store
          789: shoji
          790: shopping basket
          791: shopping cart
          792: shovel
          793: shower cap
          794: shower curtain
          795: ski
          796: ski mask
          797: sleeping bag
          798: slide rule
          799: sliding door
          800: slot machine
          801: snorkel
          802: snowmobile
          803: snowplow
          804: soap dispenser
          805: soccer ball
          806: sock
          807: solar thermal collector
          808: sombrero
          809: soup bowl
          810: space bar
          811: space heater
          812: space shuttle
          813: spatula
          814: motorboat
          815: spider web
          816: spindle
          817: sports car
          818: spotlight
          819: stage
          820: steam locomotive
          821: through arch bridge
          822: steel drum
          823: stethoscope
          824: scarf
          825: stone wall
          826: stopwatch
          827: stove
          828: strainer
          829: tram
          830: stretcher
          831: couch
          832: stupa
          833: submarine
          834: suit
          835: sundial
          836: sunglass
          837: sunglasses
          838: sunscreen
          839: suspension bridge
          840: mop
          841: sweatshirt
          842: swimsuit
          843: swing
          844: switch
          845: syringe
          846: table lamp
          847: tank
          848: tape player
          849: teapot
          850: teddy bear
          851: television
          852: tennis ball
          853: thatched roof
          854: front curtain
          855: thimble
          856: threshing machine
          857: throne
          858: tile roof
          859: toaster
          860: tobacco shop
          861: toilet seat
          862: torch
          863: totem pole
          864: tow truck
          865: toy store
          866: tractor
          867: semi-trailer truck
          868: tray
          869: trench coat
          870: tricycle
          871: trimaran
          872: tripod
          873: triumphal arch
          874: trolleybus
          875: trombone
          876: tub
          877: turnstile
          878: typewriter keyboard
          879: umbrella
          880: unicycle
          881: upright piano
          882: vacuum cleaner
          883: vase
          884: vault
          885: velvet
          886: vending machine
          887: vestment
          888: viaduct
          889: violin
          890: volleyball
          891: waffle iron
          892: wall clock
          893: wallet
          894: wardrobe
          895: military aircraft
          896: sink
          897: washing machine
          898: water bottle
          899: water jug
          900: water tower
          901: whiskey jug
          902: whistle
          903: wig
          904: window screen
          905: window shade
          906: Windsor tie
          907: wine bottle
          908: wing
          909: wok
          910: wooden spoon
          911: wool
          912: split-rail fence
          913: shipwreck
          914: yawl
          915: yurt
          916: website
          917: comic book
          918: crossword
          919: traffic sign
          920: traffic light
          921: dust jacket
          922: menu
          923: plate
          924: guacamole
          925: consomme
          926: hot pot
          927: trifle
          928: ice cream
          929: ice pop
          930: baguette
          931: bagel
          932: pretzel
          933: cheeseburger
          934: hot dog
          935: mashed potato
          936: cabbage
          937: broccoli
          938: cauliflower
          939: zucchini
          940: spaghetti squash
          941: acorn squash
          942: butternut squash
          943: cucumber
          944: artichoke
          945: bell pepper
          946: cardoon
          947: mushroom
          948: Granny Smith
          949: strawberry
          950: orange
          951: lemon
          952: fig
          953: pineapple
          954: banana
          955: jackfruit
          956: custard apple
          957: pomegranate
          958: hay
          959: carbonara
          960: chocolate syrup
          961: dough
          962: meatloaf
          963: pizza
          964: pot pie
          965: burrito
          966: red wine
          967: espresso
          968: cup
          969: eggnog
          970: alp
          971: bubble
          972: cliff
          973: coral reef
          974: geyser
          975: lakeshore
          976: promontory
          977: shoal
          978: seashore
          979: valley
          980: volcano
          981: baseball player
          982: bridegroom
          983: scuba diver
          984: rapeseed
          985: daisy
          986: yellow lady's slipper
          987: corn
          988: acorn
          989: rose hip
          990: horse chestnut seed
          991: coral fungus
          992: agaric
          993: gyromitra
          994: stinkhorn mushroom
          995: earth star
          996: hen-of-the-woods
          997: bolete
          998: ear
          999: toilet paper
        
        # Download script/URL (optional)
        download: data/scripts/get_imagenet.sh
        
      • 📄 ImageNet10.yaml 932B
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # ImageNet-1k dataset https://www.image-net.org/index.php by Stanford University
        # Simplified class names from https://github.com/anishathalye/imagenet-simple-labels
        # Example usage: python classify/train.py --data imagenet
        # parent
        # ├── yolov5
        # └── datasets
        #     └── imagenet10  ← downloads here
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/imagenet10 # dataset root dir
        train: train # train images (relative to 'path') 1281167 images
        val: val # val images (relative to 'path') 50000 images
        test: # test images (optional)
        
        # Classes
        names:
          0: tench
          1: goldfish
          2: great white shark
          3: tiger shark
          4: hammerhead shark
          5: electric ray
          6: stingray
          7: cock
          8: hen
          9: ostrich
        
        # Download script/URL (optional)
        download: data/scripts/get_imagenet10.sh
        
      • 📄 ImageNet100.yaml 2.6KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # ImageNet-1k dataset https://www.image-net.org/index.php by Stanford University
        # Simplified class names from https://github.com/anishathalye/imagenet-simple-labels
        # Example usage: python classify/train.py --data imagenet
        # parent
        # ├── yolov5
        # └── datasets
        #     └── imagenet100  ← downloads here
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/imagenet100 # dataset root dir
        train: train # train images (relative to 'path') 1281167 images
        val: val # val images (relative to 'path') 50000 images
        test: # test images (optional)
        
        # Classes
        names:
          0: tench
          1: goldfish
          2: great white shark
          3: tiger shark
          4: hammerhead shark
          5: electric ray
          6: stingray
          7: cock
          8: hen
          9: ostrich
          10: brambling
          11: goldfinch
          12: house finch
          13: junco
          14: indigo bunting
          15: American robin
          16: bulbul
          17: jay
          18: magpie
          19: chickadee
          20: American dipper
          21: kite
          22: bald eagle
          23: vulture
          24: great grey owl
          25: fire salamander
          26: smooth newt
          27: newt
          28: spotted salamander
          29: axolotl
          30: American bullfrog
          31: tree frog
          32: tailed frog
          33: loggerhead sea turtle
          34: leatherback sea turtle
          35: mud turtle
          36: terrapin
          37: box turtle
          38: banded gecko
          39: green iguana
          40: Carolina anole
          41: desert grassland whiptail lizard
          42: agama
          43: frilled-necked lizard
          44: alligator lizard
          45: Gila monster
          46: European green lizard
          47: chameleon
          48: Komodo dragon
          49: Nile crocodile
          50: American alligator
          51: triceratops
          52: worm snake
          53: ring-necked snake
          54: eastern hog-nosed snake
          55: smooth green snake
          56: kingsnake
          57: garter snake
          58: water snake
          59: vine snake
          60: night snake
          61: boa constrictor
          62: African rock python
          63: Indian cobra
          64: green mamba
          65: sea snake
          66: Saharan horned viper
          67: eastern diamondback rattlesnake
          68: sidewinder
          69: trilobite
          70: harvestman
          71: scorpion
          72: yellow garden spider
          73: barn spider
          74: European garden spider
          75: southern black widow
          76: tarantula
          77: wolf spider
          78: tick
          79: centipede
          80: black grouse
          81: ptarmigan
          82: ruffed grouse
          83: prairie grouse
          84: peacock
          85: quail
          86: partridge
          87: grey parrot
          88: macaw
          89: sulphur-crested cockatoo
          90: lorikeet
          91: coucal
          92: bee eater
          93: hornbill
          94: hummingbird
          95: jacamar
          96: toucan
          97: duck
          98: red-breasted merganser
          99: goose
        # Download script/URL (optional)
        download: data/scripts/get_imagenet100.sh
        
      • 📄 ImageNet1000.yaml 18.4KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # ImageNet-1k dataset https://www.image-net.org/index.php by Stanford University
        # Simplified class names from https://github.com/anishathalye/imagenet-simple-labels
        # Example usage: python classify/train.py --data imagenet
        # parent
        # ├── yolov5
        # └── datasets
        #     └── imagenet100  ← downloads here
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/imagenet1000 # dataset root dir
        train: train # train images (relative to 'path') 1281167 images
        val: val # val images (relative to 'path') 50000 images
        test: # test images (optional)
        
        # Classes
        names:
          0: tench
          1: goldfish
          2: great white shark
          3: tiger shark
          4: hammerhead shark
          5: electric ray
          6: stingray
          7: cock
          8: hen
          9: ostrich
          10: brambling
          11: goldfinch
          12: house finch
          13: junco
          14: indigo bunting
          15: American robin
          16: bulbul
          17: jay
          18: magpie
          19: chickadee
          20: American dipper
          21: kite
          22: bald eagle
          23: vulture
          24: great grey owl
          25: fire salamander
          26: smooth newt
          27: newt
          28: spotted salamander
          29: axolotl
          30: American bullfrog
          31: tree frog
          32: tailed frog
          33: loggerhead sea turtle
          34: leatherback sea turtle
          35: mud turtle
          36: terrapin
          37: box turtle
          38: banded gecko
          39: green iguana
          40: Carolina anole
          41: desert grassland whiptail lizard
          42: agama
          43: frilled-necked lizard
          44: alligator lizard
          45: Gila monster
          46: European green lizard
          47: chameleon
          48: Komodo dragon
          49: Nile crocodile
          50: American alligator
          51: triceratops
          52: worm snake
          53: ring-necked snake
          54: eastern hog-nosed snake
          55: smooth green snake
          56: kingsnake
          57: garter snake
          58: water snake
          59: vine snake
          60: night snake
          61: boa constrictor
          62: African rock python
          63: Indian cobra
          64: green mamba
          65: sea snake
          66: Saharan horned viper
          67: eastern diamondback rattlesnake
          68: sidewinder
          69: trilobite
          70: harvestman
          71: scorpion
          72: yellow garden spider
          73: barn spider
          74: European garden spider
          75: southern black widow
          76: tarantula
          77: wolf spider
          78: tick
          79: centipede
          80: black grouse
          81: ptarmigan
          82: ruffed grouse
          83: prairie grouse
          84: peacock
          85: quail
          86: partridge
          87: grey parrot
          88: macaw
          89: sulphur-crested cockatoo
          90: lorikeet
          91: coucal
          92: bee eater
          93: hornbill
          94: hummingbird
          95: jacamar
          96: toucan
          97: duck
          98: red-breasted merganser
          99: goose
          100: black swan
          101: tusker
          102: echidna
          103: platypus
          104: wallaby
          105: koala
          106: wombat
          107: jellyfish
          108: sea anemone
          109: brain coral
          110: flatworm
          111: nematode
          112: conch
          113: snail
          114: slug
          115: sea slug
          116: chiton
          117: chambered nautilus
          118: Dungeness crab
          119: rock crab
          120: fiddler crab
          121: red king crab
          122: American lobster
          123: spiny lobster
          124: crayfish
          125: hermit crab
          126: isopod
          127: white stork
          128: black stork
          129: spoonbill
          130: flamingo
          131: little blue heron
          132: great egret
          133: bittern
          134: crane (bird)
          135: limpkin
          136: common gallinule
          137: American coot
          138: bustard
          139: ruddy turnstone
          140: dunlin
          141: common redshank
          142: dowitcher
          143: oystercatcher
          144: pelican
          145: king penguin
          146: albatross
          147: grey whale
          148: killer whale
          149: dugong
          150: sea lion
          151: Chihuahua
          152: Japanese Chin
          153: Maltese
          154: Pekingese
          155: Shih Tzu
          156: King Charles Spaniel
          157: Papillon
          158: toy terrier
          159: Rhodesian Ridgeback
          160: Afghan Hound
          161: Basset Hound
          162: Beagle
          163: Bloodhound
          164: Bluetick Coonhound
          165: Black and Tan Coonhound
          166: Treeing Walker Coonhound
          167: English foxhound
          168: Redbone Coonhound
          169: borzoi
          170: Irish Wolfhound
          171: Italian Greyhound
          172: Whippet
          173: Ibizan Hound
          174: Norwegian Elkhound
          175: Otterhound
          176: Saluki
          177: Scottish Deerhound
          178: Weimaraner
          179: Staffordshire Bull Terrier
          180: American Staffordshire Terrier
          181: Bedlington Terrier
          182: Border Terrier
          183: Kerry Blue Terrier
          184: Irish Terrier
          185: Norfolk Terrier
          186: Norwich Terrier
          187: Yorkshire Terrier
          188: Wire Fox Terrier
          189: Lakeland Terrier
          190: Sealyham Terrier
          191: Airedale Terrier
          192: Cairn Terrier
          193: Australian Terrier
          194: Dandie Dinmont Terrier
          195: Boston Terrier
          196: Miniature Schnauzer
          197: Giant Schnauzer
          198: Standard Schnauzer
          199: Scottish Terrier
          200: Tibetan Terrier
          201: Australian Silky Terrier
          202: Soft-coated Wheaten Terrier
          203: West Highland White Terrier
          204: Lhasa Apso
          205: Flat-Coated Retriever
          206: Curly-coated Retriever
          207: Golden Retriever
          208: Labrador Retriever
          209: Chesapeake Bay Retriever
          210: German Shorthaired Pointer
          211: Vizsla
          212: English Setter
          213: Irish Setter
          214: Gordon Setter
          215: Brittany
          216: Clumber Spaniel
          217: English Springer Spaniel
          218: Welsh Springer Spaniel
          219: Cocker Spaniels
          220: Sussex Spaniel
          221: Irish Water Spaniel
          222: Kuvasz
          223: Schipperke
          224: Groenendael
          225: Malinois
          226: Briard
          227: Australian Kelpie
          228: Komondor
          229: Old English Sheepdog
          230: Shetland Sheepdog
          231: collie
          232: Border Collie
          233: Bouvier des Flandres
          234: Rottweiler
          235: German Shepherd Dog
          236: Dobermann
          237: Miniature Pinscher
          238: Greater Swiss Mountain Dog
          239: Bernese Mountain Dog
          240: Appenzeller Sennenhund
          241: Entlebucher Sennenhund
          242: Boxer
          243: Bullmastiff
          244: Tibetan Mastiff
          245: French Bulldog
          246: Great Dane
          247: St. Bernard
          248: husky
          249: Alaskan Malamute
          250: Siberian Husky
          251: Dalmatian
          252: Affenpinscher
          253: Basenji
          254: pug
          255: Leonberger
          256: Newfoundland
          257: Pyrenean Mountain Dog
          258: Samoyed
          259: Pomeranian
          260: Chow Chow
          261: Keeshond
          262: Griffon Bruxellois
          263: Pembroke Welsh Corgi
          264: Cardigan Welsh Corgi
          265: Toy Poodle
          266: Miniature Poodle
          267: Standard Poodle
          268: Mexican hairless dog
          269: grey wolf
          270: Alaskan tundra wolf
          271: red wolf
          272: coyote
          273: dingo
          274: dhole
          275: African wild dog
          276: hyena
          277: red fox
          278: kit fox
          279: Arctic fox
          280: grey fox
          281: tabby cat
          282: tiger cat
          283: Persian cat
          284: Siamese cat
          285: Egyptian Mau
          286: cougar
          287: lynx
          288: leopard
          289: snow leopard
          290: jaguar
          291: lion
          292: tiger
          293: cheetah
          294: brown bear
          295: American black bear
          296: polar bear
          297: sloth bear
          298: mongoose
          299: meerkat
          300: tiger beetle
          301: ladybug
          302: ground beetle
          303: longhorn beetle
          304: leaf beetle
          305: dung beetle
          306: rhinoceros beetle
          307: weevil
          308: fly
          309: bee
          310: ant
          311: grasshopper
          312: cricket
          313: stick insect
          314: cockroach
          315: mantis
          316: cicada
          317: leafhopper
          318: lacewing
          319: dragonfly
          320: damselfly
          321: red admiral
          322: ringlet
          323: monarch butterfly
          324: small white
          325: sulphur butterfly
          326: gossamer-winged butterfly
          327: starfish
          328: sea urchin
          329: sea cucumber
          330: cottontail rabbit
          331: hare
          332: Angora rabbit
          333: hamster
          334: porcupine
          335: fox squirrel
          336: marmot
          337: beaver
          338: guinea pig
          339: common sorrel
          340: zebra
          341: pig
          342: wild boar
          343: warthog
          344: hippopotamus
          345: ox
          346: water buffalo
          347: bison
          348: ram
          349: bighorn sheep
          350: Alpine ibex
          351: hartebeest
          352: impala
          353: gazelle
          354: dromedary
          355: llama
          356: weasel
          357: mink
          358: European polecat
          359: black-footed ferret
          360: otter
          361: skunk
          362: badger
          363: armadillo
          364: three-toed sloth
          365: orangutan
          366: gorilla
          367: chimpanzee
          368: gibbon
          369: siamang
          370: guenon
          371: patas monkey
          372: baboon
          373: macaque
          374: langur
          375: black-and-white colobus
          376: proboscis monkey
          377: marmoset
          378: white-headed capuchin
          379: howler monkey
          380: titi
          381: Geoffroy's spider monkey
          382: common squirrel monkey
          383: ring-tailed lemur
          384: indri
          385: Asian elephant
          386: African bush elephant
          387: red panda
          388: giant panda
          389: snoek
          390: eel
          391: coho salmon
          392: rock beauty
          393: clownfish
          394: sturgeon
          395: garfish
          396: lionfish
          397: pufferfish
          398: abacus
          399: abaya
          400: academic gown
          401: accordion
          402: acoustic guitar
          403: aircraft carrier
          404: airliner
          405: airship
          406: altar
          407: ambulance
          408: amphibious vehicle
          409: analog clock
          410: apiary
          411: apron
          412: waste container
          413: assault rifle
          414: backpack
          415: bakery
          416: balance beam
          417: balloon
          418: ballpoint pen
          419: Band-Aid
          420: banjo
          421: baluster
          422: barbell
          423: barber chair
          424: barbershop
          425: barn
          426: barometer
          427: barrel
          428: wheelbarrow
          429: baseball
          430: basketball
          431: bassinet
          432: bassoon
          433: swimming cap
          434: bath towel
          435: bathtub
          436: station wagon
          437: lighthouse
          438: beaker
          439: military cap
          440: beer bottle
          441: beer glass
          442: bell-cot
          443: bib
          444: tandem bicycle
          445: bikini
          446: ring binder
          447: binoculars
          448: birdhouse
          449: boathouse
          450: bobsleigh
          451: bolo tie
          452: poke bonnet
          453: bookcase
          454: bookstore
          455: bottle cap
          456: bow
          457: bow tie
          458: brass
          459: bra
          460: breakwater
          461: breastplate
          462: broom
          463: bucket
          464: buckle
          465: bulletproof vest
          466: high-speed train
          467: butcher shop
          468: taxicab
          469: cauldron
          470: candle
          471: cannon
          472: canoe
          473: can opener
          474: cardigan
          475: car mirror
          476: carousel
          477: tool kit
          478: carton
          479: car wheel
          480: automated teller machine
          481: cassette
          482: cassette player
          483: castle
          484: catamaran
          485: CD player
          486: cello
          487: mobile phone
          488: chain
          489: chain-link fence
          490: chain mail
          491: chainsaw
          492: chest
          493: chiffonier
          494: chime
          495: china cabinet
          496: Christmas stocking
          497: church
          498: movie theater
          499: cleaver
          500: cliff dwelling
          501: cloak
          502: clogs
          503: cocktail shaker
          504: coffee mug
          505: coffeemaker
          506: coil
          507: combination lock
          508: computer keyboard
          509: confectionery store
          510: container ship
          511: convertible
          512: corkscrew
          513: cornet
          514: cowboy boot
          515: cowboy hat
          516: cradle
          517: crane (machine)
          518: crash helmet
          519: crate
          520: infant bed
          521: Crock Pot
          522: croquet ball
          523: crutch
          524: cuirass
          525: dam
          526: desk
          527: desktop computer
          528: rotary dial telephone
          529: diaper
          530: digital clock
          531: digital watch
          532: dining table
          533: dishcloth
          534: dishwasher
          535: disc brake
          536: dock
          537: dog sled
          538: dome
          539: doormat
          540: drilling rig
          541: drum
          542: drumstick
          543: dumbbell
          544: Dutch oven
          545: electric fan
          546: electric guitar
          547: electric locomotive
          548: entertainment center
          549: envelope
          550: espresso machine
          551: face powder
          552: feather boa
          553: filing cabinet
          554: fireboat
          555: fire engine
          556: fire screen sheet
          557: flagpole
          558: flute
          559: folding chair
          560: football helmet
          561: forklift
          562: fountain
          563: fountain pen
          564: four-poster bed
          565: freight car
          566: French horn
          567: frying pan
          568: fur coat
          569: garbage truck
          570: gas mask
          571: gas pump
          572: goblet
          573: go-kart
          574: golf ball
          575: golf cart
          576: gondola
          577: gong
          578: gown
          579: grand piano
          580: greenhouse
          581: grille
          582: grocery store
          583: guillotine
          584: barrette
          585: hair spray
          586: half-track
          587: hammer
          588: hamper
          589: hair dryer
          590: hand-held computer
          591: handkerchief
          592: hard disk drive
          593: harmonica
          594: harp
          595: harvester
          596: hatchet
          597: holster
          598: home theater
          599: honeycomb
          600: hook
          601: hoop skirt
          602: horizontal bar
          603: horse-drawn vehicle
          604: hourglass
          605: iPod
          606: clothes iron
          607: jack-o'-lantern
          608: jeans
          609: jeep
          610: T-shirt
          611: jigsaw puzzle
          612: pulled rickshaw
          613: joystick
          614: kimono
          615: knee pad
          616: knot
          617: lab coat
          618: ladle
          619: lampshade
          620: laptop computer
          621: lawn mower
          622: lens cap
          623: paper knife
          624: library
          625: lifeboat
          626: lighter
          627: limousine
          628: ocean liner
          629: lipstick
          630: slip-on shoe
          631: lotion
          632: speaker
          633: loupe
          634: sawmill
          635: magnetic compass
          636: mail bag
          637: mailbox
          638: tights
          639: tank suit
          640: manhole cover
          641: maraca
          642: marimba
          643: mask
          644: match
          645: maypole
          646: maze
          647: measuring cup
          648: medicine chest
          649: megalith
          650: microphone
          651: microwave oven
          652: military uniform
          653: milk can
          654: minibus
          655: miniskirt
          656: minivan
          657: missile
          658: mitten
          659: mixing bowl
          660: mobile home
          661: Model T
          662: modem
          663: monastery
          664: monitor
          665: moped
          666: mortar
          667: square academic cap
          668: mosque
          669: mosquito net
          670: scooter
          671: mountain bike
          672: tent
          673: computer mouse
          674: mousetrap
          675: moving van
          676: muzzle
          677: nail
          678: neck brace
          679: necklace
          680: nipple
          681: notebook computer
          682: obelisk
          683: oboe
          684: ocarina
          685: odometer
          686: oil filter
          687: organ
          688: oscilloscope
          689: overskirt
          690: bullock cart
          691: oxygen mask
          692: packet
          693: paddle
          694: paddle wheel
          695: padlock
          696: paintbrush
          697: pajamas
          698: palace
          699: pan flute
          700: paper towel
          701: parachute
          702: parallel bars
          703: park bench
          704: parking meter
          705: passenger car
          706: patio
          707: payphone
          708: pedestal
          709: pencil case
          710: pencil sharpener
          711: perfume
          712: Petri dish
          713: photocopier
          714: plectrum
          715: Pickelhaube
          716: picket fence
          717: pickup truck
          718: pier
          719: piggy bank
          720: pill bottle
          721: pillow
          722: ping-pong ball
          723: pinwheel
          724: pirate ship
          725: pitcher
          726: hand plane
          727: planetarium
          728: plastic bag
          729: plate rack
          730: plow
          731: plunger
          732: Polaroid camera
          733: pole
          734: police van
          735: poncho
          736: billiard table
          737: soda bottle
          738: pot
          739: potter's wheel
          740: power drill
          741: prayer rug
          742: printer
          743: prison
          744: projectile
          745: projector
          746: hockey puck
          747: punching bag
          748: purse
          749: quill
          750: quilt
          751: race car
          752: racket
          753: radiator
          754: radio
          755: radio telescope
          756: rain barrel
          757: recreational vehicle
          758: reel
          759: reflex camera
          760: refrigerator
          761: remote control
          762: restaurant
          763: revolver
          764: rifle
          765: rocking chair
          766: rotisserie
          767: eraser
          768: rugby ball
          769: ruler
          770: running shoe
          771: safe
          772: safety pin
          773: salt shaker
          774: sandal
          775: sarong
          776: saxophone
          777: scabbard
          778: weighing scale
          779: school bus
          780: schooner
          781: scoreboard
          782: CRT screen
          783: screw
          784: screwdriver
          785: seat belt
          786: sewing machine
          787: shield
          788: shoe store
          789: shoji
          790: shopping basket
          791: shopping cart
          792: shovel
          793: shower cap
          794: shower curtain
          795: ski
          796: ski mask
          797: sleeping bag
          798: slide rule
          799: sliding door
          800: slot machine
          801: snorkel
          802: snowmobile
          803: snowplow
          804: soap dispenser
          805: soccer ball
          806: sock
          807: solar thermal collector
          808: sombrero
          809: soup bowl
          810: space bar
          811: space heater
          812: space shuttle
          813: spatula
          814: motorboat
          815: spider web
          816: spindle
          817: sports car
          818: spotlight
          819: stage
          820: steam locomotive
          821: through arch bridge
          822: steel drum
          823: stethoscope
          824: scarf
          825: stone wall
          826: stopwatch
          827: stove
          828: strainer
          829: tram
          830: stretcher
          831: couch
          832: stupa
          833: submarine
          834: suit
          835: sundial
          836: sunglass
          837: sunglasses
          838: sunscreen
          839: suspension bridge
          840: mop
          841: sweatshirt
          842: swimsuit
          843: swing
          844: switch
          845: syringe
          846: table lamp
          847: tank
          848: tape player
          849: teapot
          850: teddy bear
          851: television
          852: tennis ball
          853: thatched roof
          854: front curtain
          855: thimble
          856: threshing machine
          857: throne
          858: tile roof
          859: toaster
          860: tobacco shop
          861: toilet seat
          862: torch
          863: totem pole
          864: tow truck
          865: toy store
          866: tractor
          867: semi-trailer truck
          868: tray
          869: trench coat
          870: tricycle
          871: trimaran
          872: tripod
          873: triumphal arch
          874: trolleybus
          875: trombone
          876: tub
          877: turnstile
          878: typewriter keyboard
          879: umbrella
          880: unicycle
          881: upright piano
          882: vacuum cleaner
          883: vase
          884: vault
          885: velvet
          886: vending machine
          887: vestment
          888: viaduct
          889: violin
          890: volleyball
          891: waffle iron
          892: wall clock
          893: wallet
          894: wardrobe
          895: military aircraft
          896: sink
          897: washing machine
          898: water bottle
          899: water jug
          900: water tower
          901: whiskey jug
          902: whistle
          903: wig
          904: window screen
          905: window shade
          906: Windsor tie
          907: wine bottle
          908: wing
          909: wok
          910: wooden spoon
          911: wool
          912: split-rail fence
          913: shipwreck
          914: yawl
          915: yurt
          916: website
          917: comic book
          918: crossword
          919: traffic sign
          920: traffic light
          921: dust jacket
          922: menu
          923: plate
          924: guacamole
          925: consomme
          926: hot pot
          927: trifle
          928: ice cream
          929: ice pop
          930: baguette
          931: bagel
          932: pretzel
          933: cheeseburger
          934: hot dog
          935: mashed potato
          936: cabbage
          937: broccoli
          938: cauliflower
          939: zucchini
          940: spaghetti squash
          941: acorn squash
          942: butternut squash
          943: cucumber
          944: artichoke
          945: bell pepper
          946: cardoon
          947: mushroom
          948: Granny Smith
          949: strawberry
          950: orange
          951: lemon
          952: fig
          953: pineapple
          954: banana
          955: jackfruit
          956: custard apple
          957: pomegranate
          958: hay
          959: carbonara
          960: chocolate syrup
          961: dough
          962: meatloaf
          963: pizza
          964: pot pie
          965: burrito
          966: red wine
          967: espresso
          968: cup
          969: eggnog
          970: alp
          971: bubble
          972: cliff
          973: coral reef
          974: geyser
          975: lakeshore
          976: promontory
          977: shoal
          978: seashore
          979: valley
          980: volcano
          981: baseball player
          982: bridegroom
          983: scuba diver
          984: rapeseed
          985: daisy
          986: yellow lady's slipper
          987: corn
          988: acorn
          989: rose hip
          990: horse chestnut seed
          991: coral fungus
          992: agaric
          993: gyromitra
          994: stinkhorn mushroom
          995: earth star
          996: hen-of-the-woods
          997: bolete
          998: ear
          999: toilet paper
        
        # Download script/URL (optional)
        download: data/scripts/get_imagenet1000.sh
        
      • 📄 Objects365.yaml 9.0KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # Objects365 dataset https://www.objects365.org/ by Megvii
        # Example usage: python train.py --data Objects365.yaml
        # parent
        # ├── yolov5
        # └── datasets
        #     └── Objects365  ← downloads here (712 GB = 367G data + 345G zips)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/Objects365 # dataset root dir
        train: images/train # train images (relative to 'path') 1742289 images
        val: images/val # val images (relative to 'path') 80000 images
        test: # test images (optional)
        
        # Classes
        names:
          0: Person
          1: Sneakers
          2: Chair
          3: Other Shoes
          4: Hat
          5: Car
          6: Lamp
          7: Glasses
          8: Bottle
          9: Desk
          10: Cup
          11: Street Lights
          12: Cabinet/shelf
          13: Handbag/Satchel
          14: Bracelet
          15: Plate
          16: Picture/Frame
          17: Helmet
          18: Book
          19: Gloves
          20: Storage box
          21: Boat
          22: Leather Shoes
          23: Flower
          24: Bench
          25: Potted Plant
          26: Bowl/Basin
          27: Flag
          28: Pillow
          29: Boots
          30: Vase
          31: Microphone
          32: Necklace
          33: Ring
          34: SUV
          35: Wine Glass
          36: Belt
          37: Monitor/TV
          38: Backpack
          39: Umbrella
          40: Traffic Light
          41: Speaker
          42: Watch
          43: Tie
          44: Trash bin Can
          45: Slippers
          46: Bicycle
          47: Stool
          48: Barrel/bucket
          49: Van
          50: Couch
          51: Sandals
          52: Basket
          53: Drum
          54: Pen/Pencil
          55: Bus
          56: Wild Bird
          57: High Heels
          58: Motorcycle
          59: Guitar
          60: Carpet
          61: Cell Phone
          62: Bread
          63: Camera
          64: Canned
          65: Truck
          66: Traffic cone
          67: Cymbal
          68: Lifesaver
          69: Towel
          70: Stuffed Toy
          71: Candle
          72: Sailboat
          73: Laptop
          74: Awning
          75: Bed
          76: Faucet
          77: Tent
          78: Horse
          79: Mirror
          80: Power outlet
          81: Sink
          82: Apple
          83: Air Conditioner
          84: Knife
          85: Hockey Stick
          86: Paddle
          87: Pickup Truck
          88: Fork
          89: Traffic Sign
          90: Balloon
          91: Tripod
          92: Dog
          93: Spoon
          94: Clock
          95: Pot
          96: Cow
          97: Cake
          98: Dinning Table
          99: Sheep
          100: Hanger
          101: Blackboard/Whiteboard
          102: Napkin
          103: Other Fish
          104: Orange/Tangerine
          105: Toiletry
          106: Keyboard
          107: Tomato
          108: Lantern
          109: Machinery Vehicle
          110: Fan
          111: Green Vegetables
          112: Banana
          113: Baseball Glove
          114: Airplane
          115: Mouse
          116: Train
          117: Pumpkin
          118: Soccer
          119: Skiboard
          120: Luggage
          121: Nightstand
          122: Tea pot
          123: Telephone
          124: Trolley
          125: Head Phone
          126: Sports Car
          127: Stop Sign
          128: Dessert
          129: Scooter
          130: Stroller
          131: Crane
          132: Remote
          133: Refrigerator
          134: Oven
          135: Lemon
          136: Duck
          137: Baseball Bat
          138: Surveillance Camera
          139: Cat
          140: Jug
          141: Broccoli
          142: Piano
          143: Pizza
          144: Elephant
          145: Skateboard
          146: Surfboard
          147: Gun
          148: Skating and Skiing shoes
          149: Gas stove
          150: Donut
          151: Bow Tie
          152: Carrot
          153: Toilet
          154: Kite
          155: Strawberry
          156: Other Balls
          157: Shovel
          158: Pepper
          159: Computer Box
          160: Toilet Paper
          161: Cleaning Products
          162: Chopsticks
          163: Microwave
          164: Pigeon
          165: Baseball
          166: Cutting/chopping Board
          167: Coffee Table
          168: Side Table
          169: Scissors
          170: Marker
          171: Pie
          172: Ladder
          173: Snowboard
          174: Cookies
          175: Radiator
          176: Fire Hydrant
          177: Basketball
          178: Zebra
          179: Grape
          180: Giraffe
          181: Potato
          182: Sausage
          183: Tricycle
          184: Violin
          185: Egg
          186: Fire Extinguisher
          187: Candy
          188: Fire Truck
          189: Billiards
          190: Converter
          191: Bathtub
          192: Wheelchair
          193: Golf Club
          194: Briefcase
          195: Cucumber
          196: Cigar/Cigarette
          197: Paint Brush
          198: Pear
          199: Heavy Truck
          200: Hamburger
          201: Extractor
          202: Extension Cord
          203: Tong
          204: Tennis Racket
          205: Folder
          206: American Football
          207: earphone
          208: Mask
          209: Kettle
          210: Tennis
          211: Ship
          212: Swing
          213: Coffee Machine
          214: Slide
          215: Carriage
          216: Onion
          217: Green beans
          218: Projector
          219: Frisbee
          220: Washing Machine/Drying Machine
          221: Chicken
          222: Printer
          223: Watermelon
          224: Saxophone
          225: Tissue
          226: Toothbrush
          227: Ice cream
          228: Hot-air balloon
          229: Cello
          230: French Fries
          231: Scale
          232: Trophy
          233: Cabbage
          234: Hot dog
          235: Blender
          236: Peach
          237: Rice
          238: Wallet/Purse
          239: Volleyball
          240: Deer
          241: Goose
          242: Tape
          243: Tablet
          244: Cosmetics
          245: Trumpet
          246: Pineapple
          247: Golf Ball
          248: Ambulance
          249: Parking meter
          250: Mango
          251: Key
          252: Hurdle
          253: Fishing Rod
          254: Medal
          255: Flute
          256: Brush
          257: Penguin
          258: Megaphone
          259: Corn
          260: Lettuce
          261: Garlic
          262: Swan
          263: Helicopter
          264: Green Onion
          265: Sandwich
          266: Nuts
          267: Speed Limit Sign
          268: Induction Cooker
          269: Broom
          270: Trombone
          271: Plum
          272: Rickshaw
          273: Goldfish
          274: Kiwi fruit
          275: Router/modem
          276: Poker Card
          277: Toaster
          278: Shrimp
          279: Sushi
          280: Cheese
          281: Notepaper
          282: Cherry
          283: Pliers
          284: CD
          285: Pasta
          286: Hammer
          287: Cue
          288: Avocado
          289: Hamimelon
          290: Flask
          291: Mushroom
          292: Screwdriver
          293: Soap
          294: Recorder
          295: Bear
          296: Eggplant
          297: Board Eraser
          298: Coconut
          299: Tape Measure/Ruler
          300: Pig
          301: Showerhead
          302: Globe
          303: Chips
          304: Steak
          305: Crosswalk Sign
          306: Stapler
          307: Camel
          308: Formula 1
          309: Pomegranate
          310: Dishwasher
          311: Crab
          312: Hoverboard
          313: Meat ball
          314: Rice Cooker
          315: Tuba
          316: Calculator
          317: Papaya
          318: Antelope
          319: Parrot
          320: Seal
          321: Butterfly
          322: Dumbbell
          323: Donkey
          324: Lion
          325: Urinal
          326: Dolphin
          327: Electric Drill
          328: Hair Dryer
          329: Egg tart
          330: Jellyfish
          331: Treadmill
          332: Lighter
          333: Grapefruit
          334: Game board
          335: Mop
          336: Radish
          337: Baozi
          338: Target
          339: French
          340: Spring Rolls
          341: Monkey
          342: Rabbit
          343: Pencil Case
          344: Yak
          345: Red Cabbage
          346: Binoculars
          347: Asparagus
          348: Barbell
          349: Scallop
          350: Noddles
          351: Comb
          352: Dumpling
          353: Oyster
          354: Table Tennis paddle
          355: Cosmetics Brush/Eyeliner Pencil
          356: Chainsaw
          357: Eraser
          358: Lobster
          359: Durian
          360: Okra
          361: Lipstick
          362: Cosmetics Mirror
          363: Curling
          364: Table Tennis
        
        # Download script/URL (optional) ---------------------------------------------------------------------------------------
        download: |
          from tqdm import tqdm
        
          from utils.general import Path, check_requirements, download, np, xyxy2xywhn
        
          check_requirements('pycocotools>=2.0')
          from pycocotools.coco import COCO
        
          # Make Directories
          dir = Path(yaml['path'])  # dataset root dir
          for p in 'images', 'labels':
              (dir / p).mkdir(parents=True, exist_ok=True)
              for q in 'train', 'val':
                  (dir / p / q).mkdir(parents=True, exist_ok=True)
        
          # Train, Val Splits
          for split, patches in [('train', 50 + 1), ('val', 43 + 1)]:
              print(f"Processing {split} in {patches} patches ...")
              images, labels = dir / 'images' / split, dir / 'labels' / split
        
              # Download
              url = f"https://dorc.ks3-cn-beijing.ksyun.com/data-set/2020Objects365%E6%95%B0%E6%8D%AE%E9%9B%86/{split}/"
              if split == 'train':
                  download([f'{url}zhiyuan_objv2_{split}.tar.gz'], dir=dir, delete=False)  # annotations json
                  download([f'{url}patch{i}.tar.gz' for i in range(patches)], dir=images, curl=True, delete=False, threads=8)
              elif split == 'val':
                  download([f'{url}zhiyuan_objv2_{split}.json'], dir=dir, delete=False)  # annotations json
                  download([f'{url}images/v1/patch{i}.tar.gz' for i in range(15 + 1)], dir=images, curl=True, delete=False, threads=8)
                  download([f'{url}images/v2/patch{i}.tar.gz' for i in range(16, patches)], dir=images, curl=True, delete=False, threads=8)
        
              # Move
              for f in tqdm(images.rglob('*.jpg'), desc=f'Moving {split} images'):
                  f.rename(images / f.name)  # move to /images/{split}
        
              # Labels
              coco = COCO(dir / f'zhiyuan_objv2_{split}.json')
              names = [x["name"] for x in coco.loadCats(coco.getCatIds())]
              for cid, cat in enumerate(names):
                  catIds = coco.getCatIds(catNms=[cat])
                  imgIds = coco.getImgIds(catIds=catIds)
                  for im in tqdm(coco.loadImgs(imgIds), desc=f'Class {cid + 1}/{len(names)} {cat}'):
                      width, height = im["width"], im["height"]
                      path = Path(im["file_name"])  # image filename
                      try:
                          with open(labels / path.with_suffix('.txt').name, 'a') as file:
                              annIds = coco.getAnnIds(imgIds=im["id"], catIds=catIds, iscrowd=False)
                              for a in coco.loadAnns(annIds):
                                  x, y, w, h = a['bbox']  # bounding box in xywh (xy top-left corner)
                                  xyxy = np.array([x, y, x + w, y + h])[None]  # pixels(1,4)
                                  x, y, w, h = xyxy2xywhn(xyxy, w=width, h=height, clip=True)[0]  # normalized and clipped
                                  file.write(f"{cid} {x:.5f} {y:.5f} {w:.5f} {h:.5f}\n")
                      except Exception as e:
                          print(e)
        
      • 📄 SKU-110K.yaml 2.3KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # SKU-110K retail items dataset https://github.com/eg4000/SKU110K_CVPR19 by Trax Retail
        # Example usage: python train.py --data SKU-110K.yaml
        # parent
        # ├── yolov5
        # └── datasets
        #     └── SKU-110K  ← downloads here (13.6 GB)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/SKU-110K # dataset root dir
        train: train.txt # train images (relative to 'path')  8219 images
        val: val.txt # val images (relative to 'path')  588 images
        test: test.txt # test images (optional)  2936 images
        
        # Classes
        names:
          0: object
        
        # Download script/URL (optional) ---------------------------------------------------------------------------------------
        download: |
          import shutil
          from tqdm import tqdm
          from utils.general import np, pd, Path, download, xyxy2xywh
        
        
          # Download
          dir = Path(yaml['path'])  # dataset root dir
          parent = Path(dir.parent)  # download dir
          urls = ['http://trax-geometry.s3.amazonaws.com/cvpr_challenge/SKU110K_fixed.tar.gz']
          download(urls, dir=parent, delete=False)
        
          # Rename directories
          if dir.exists():
              shutil.rmtree(dir)
          (parent / 'SKU110K_fixed').rename(dir)  # rename dir
          (dir / 'labels').mkdir(parents=True, exist_ok=True)  # create labels dir
        
          # Convert labels
          names = 'image', 'x1', 'y1', 'x2', 'y2', 'class', 'image_width', 'image_height'  # column names
          for d in 'annotations_train.csv', 'annotations_val.csv', 'annotations_test.csv':
              x = pd.read_csv(dir / 'annotations' / d, names=names).values  # annotations
              images, unique_images = x[:, 0], np.unique(x[:, 0])
              with open((dir / d).with_suffix('.txt').__str__().replace('annotations_', ''), 'w') as f:
                  f.writelines(f'./images/{s}\n' for s in unique_images)
              for im in tqdm(unique_images, desc=f'Converting {dir / d}'):
                  cls = 0  # single-class dataset
                  with open((dir / 'labels' / im).with_suffix('.txt'), 'a') as f:
                      for r in x[images == im]:
                          w, h = r[6], r[7]  # image width, height
                          xywh = xyxy2xywh(np.array([[r[1] / w, r[2] / h, r[3] / w, r[4] / h]]))[0]  # instance
                          f.write(f"{cls} {xywh[0]:.5f} {xywh[1]:.5f} {xywh[2]:.5f} {xywh[3]:.5f}\n")  # write label
        
      • 📄 VOC.yaml 3.4KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC by University of Oxford
        # Example usage: python train.py --data VOC.yaml
        # parent
        # ├── yolov5
        # └── datasets
        #     └── VOC  ← downloads here (2.8 GB)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/VOC
        train: # train images (relative to 'path')  16551 images
          - images/train2012
          - images/train2007
          - images/val2012
          - images/val2007
        val: # val images (relative to 'path')  4952 images
          - images/test2007
        test: # test images (optional)
          - images/test2007
        
        # Classes
        names:
          0: aeroplane
          1: bicycle
          2: bird
          3: boat
          4: bottle
          5: bus
          6: car
          7: cat
          8: chair
          9: cow
          10: diningtable
          11: dog
          12: horse
          13: motorbike
          14: person
          15: pottedplant
          16: sheep
          17: sofa
          18: train
          19: tvmonitor
        
        # Download script/URL (optional) ---------------------------------------------------------------------------------------
        download: |
          import xml.etree.ElementTree as ET
        
          from tqdm import tqdm
          from utils.general import download, Path
        
        
          def convert_label(path, lb_path, year, image_id):
              def convert_box(size, box):
                  dw, dh = 1. / size[0], 1. / size[1]
                  x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]
                  return x * dw, y * dh, w * dw, h * dh
        
              in_file = open(path / f'VOC{year}/Annotations/{image_id}.xml')
              out_file = open(lb_path, 'w')
              tree = ET.parse(in_file)
              root = tree.getroot()
              size = root.find('size')
              w = int(size.find('width').text)
              h = int(size.find('height').text)
        
              names = list(yaml['names'].values())  # names list
              for obj in root.iter('object'):
                  cls = obj.find('name').text
                  if cls in names and int(obj.find('difficult').text) != 1:
                      xmlbox = obj.find('bndbox')
                      bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])
                      cls_id = names.index(cls)  # class id
                      out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n')
        
        
          # Download
          dir = Path(yaml['path'])  # dataset root dir
          url = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/'
          urls = [f'{url}VOCtrainval_06-Nov-2007.zip',  # 446MB, 5012 images
                  f'{url}VOCtest_06-Nov-2007.zip',  # 438MB, 4953 images
                  f'{url}VOCtrainval_11-May-2012.zip']  # 1.95GB, 17126 images
          download(urls, dir=dir / 'images', delete=False, curl=True, threads=3)
        
          # Convert
          path = dir / 'images/VOCdevkit'
          for year, image_set in ('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test'):
              imgs_path = dir / 'images' / f'{image_set}{year}'
              lbs_path = dir / 'labels' / f'{image_set}{year}'
              imgs_path.mkdir(exist_ok=True, parents=True)
              lbs_path.mkdir(exist_ok=True, parents=True)
        
              with open(path / f'VOC{year}/ImageSets/Main/{image_set}.txt') as f:
                  image_ids = f.read().strip().split()
              for id in tqdm(image_ids, desc=f'{image_set}{year}'):
                  f = path / f'VOC{year}/JPEGImages/{id}.jpg'  # old img path
                  lb_path = (lbs_path / f.name).with_suffix('.txt')  # new label path
                  f.rename(imgs_path / f.name)  # move image
                  convert_label(path, lb_path, year, id)  # convert labels to YOLO format
        
      • 📄 VisDrone.yaml 2.9KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset by Tianjin University
        # Example usage: python train.py --data VisDrone.yaml
        # parent
        # ├── yolov5
        # └── datasets
        #     └── VisDrone  ← downloads here (2.3 GB)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/VisDrone # dataset root dir
        train: VisDrone2019-DET-train/images # train images (relative to 'path')  6471 images
        val: VisDrone2019-DET-val/images # val images (relative to 'path')  548 images
        test: VisDrone2019-DET-test-dev/images # test images (optional)  1610 images
        
        # Classes
        names:
          0: pedestrian
          1: people
          2: bicycle
          3: car
          4: van
          5: truck
          6: tricycle
          7: awning-tricycle
          8: bus
          9: motor
        
        # Download script/URL (optional) ---------------------------------------------------------------------------------------
        download: |
          from utils.general import download, os, Path
        
          def visdrone2yolo(dir):
              from PIL import Image
              from tqdm import tqdm
        
              def convert_box(size, box):
                  # Convert VisDrone box to YOLO xywh box
                  dw = 1. / size[0]
                  dh = 1. / size[1]
                  return (box[0] + box[2] / 2) * dw, (box[1] + box[3] / 2) * dh, box[2] * dw, box[3] * dh
        
              (dir / 'labels').mkdir(parents=True, exist_ok=True)  # make labels directory
              pbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}')
              for f in pbar:
                  img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).size
                  lines = []
                  with open(f, 'r') as file:  # read annotation.txt
                      for row in [x.split(',') for x in file.read().strip().splitlines()]:
                          if row[4] == '0':  # VisDrone 'ignored regions' class 0
                              continue
                          cls = int(row[5]) - 1
                          box = convert_box(img_size, tuple(map(int, row[:4])))
                          lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n")
                          with open(str(f).replace(os.sep + 'annotations' + os.sep, os.sep + 'labels' + os.sep), 'w') as fl:
                              fl.writelines(lines)  # write label.txt
        
        
          # Download
          dir = Path(yaml['path'])  # dataset root dir
          urls = ['https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-train.zip',
                  'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-val.zip',
                  'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-dev.zip',
                  'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-challenge.zip']
          download(urls, dir=dir, curl=True, threads=4)
        
          # Convert
          for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':
              visdrone2yolo(dir / d)  # convert VisDrone annotations to YOLO labels
        
      • 📄 coco.yaml 2.4KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # COCO 2017 dataset http://cocodataset.org by Microsoft
        # Example usage: python train.py --data coco.yaml
        # parent
        # ├── yolov5
        # └── datasets
        #     └── coco  ← downloads here (20.1 GB)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/coco # dataset root dir
        train: train2017.txt # train images (relative to 'path') 118287 images
        val: val2017.txt # val images (relative to 'path') 5000 images
        test: test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
        
        # Classes
        names:
          0: person
          1: bicycle
          2: car
          3: motorcycle
          4: airplane
          5: bus
          6: train
          7: truck
          8: boat
          9: traffic light
          10: fire hydrant
          11: stop sign
          12: parking meter
          13: bench
          14: bird
          15: cat
          16: dog
          17: horse
          18: sheep
          19: cow
          20: elephant
          21: bear
          22: zebra
          23: giraffe
          24: backpack
          25: umbrella
          26: handbag
          27: tie
          28: suitcase
          29: frisbee
          30: skis
          31: snowboard
          32: sports ball
          33: kite
          34: baseball bat
          35: baseball glove
          36: skateboard
          37: surfboard
          38: tennis racket
          39: bottle
          40: wine glass
          41: cup
          42: fork
          43: knife
          44: spoon
          45: bowl
          46: banana
          47: apple
          48: sandwich
          49: orange
          50: broccoli
          51: carrot
          52: hot dog
          53: pizza
          54: donut
          55: cake
          56: chair
          57: couch
          58: potted plant
          59: bed
          60: dining table
          61: toilet
          62: tv
          63: laptop
          64: mouse
          65: remote
          66: keyboard
          67: cell phone
          68: microwave
          69: oven
          70: toaster
          71: sink
          72: refrigerator
          73: book
          74: clock
          75: vase
          76: scissors
          77: teddy bear
          78: hair drier
          79: toothbrush
        
        # Download script/URL (optional)
        download: |
          from utils.general import download, Path
        
        
          # Download labels
          segments = False  # segment or box labels
          dir = Path(yaml['path'])  # dataset root dir
          url = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/'
          urls = [url + ('coco2017labels-segments.zip' if segments else 'coco2017labels.zip')]  # labels
          download(urls, dir=dir.parent)
        
          # Download data
          urls = ['http://images.cocodataset.org/zips/train2017.zip',  # 19G, 118k images
                  'http://images.cocodataset.org/zips/val2017.zip',  # 1G, 5k images
                  'http://images.cocodataset.org/zips/test2017.zip']  # 7G, 41k images (optional)
          download(urls, dir=dir / 'images', threads=3)
        
      • 📄 coco128-seg.yaml 1.8KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # COCO128-seg dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
        # Example usage: python train.py --data coco128.yaml
        # parent
        # ├── yolov5
        # └── datasets
        #     └── coco128-seg  ← downloads here (7 MB)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/coco128-seg # dataset root dir
        train: images/train2017 # train images (relative to 'path') 128 images
        val: images/train2017 # val images (relative to 'path') 128 images
        test: # test images (optional)
        
        # Classes
        names:
          0: person
          1: bicycle
          2: car
          3: motorcycle
          4: airplane
          5: bus
          6: train
          7: truck
          8: boat
          9: traffic light
          10: fire hydrant
          11: stop sign
          12: parking meter
          13: bench
          14: bird
          15: cat
          16: dog
          17: horse
          18: sheep
          19: cow
          20: elephant
          21: bear
          22: zebra
          23: giraffe
          24: backpack
          25: umbrella
          26: handbag
          27: tie
          28: suitcase
          29: frisbee
          30: skis
          31: snowboard
          32: sports ball
          33: kite
          34: baseball bat
          35: baseball glove
          36: skateboard
          37: surfboard
          38: tennis racket
          39: bottle
          40: wine glass
          41: cup
          42: fork
          43: knife
          44: spoon
          45: bowl
          46: banana
          47: apple
          48: sandwich
          49: orange
          50: broccoli
          51: carrot
          52: hot dog
          53: pizza
          54: donut
          55: cake
          56: chair
          57: couch
          58: potted plant
          59: bed
          60: dining table
          61: toilet
          62: tv
          63: laptop
          64: mouse
          65: remote
          66: keyboard
          67: cell phone
          68: microwave
          69: oven
          70: toaster
          71: sink
          72: refrigerator
          73: book
          74: clock
          75: vase
          76: scissors
          77: teddy bear
          78: hair drier
          79: toothbrush
        
        # Download script/URL (optional)
        download: https://ultralytics.com/assets/coco128-seg.zip
        
      • 📄 coco128.yaml 1.8KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
        # Example usage: python train.py --data coco128.yaml
        # parent
        # ├── yolov5
        # └── datasets
        #     └── coco128  ← downloads here (7 MB)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/coco128 # dataset root dir
        train: images/train2017 # train images (relative to 'path') 128 images
        val: images/train2017 # val images (relative to 'path') 128 images
        test: # test images (optional)
        
        # Classes
        names:
          0: person
          1: bicycle
          2: car
          3: motorcycle
          4: airplane
          5: bus
          6: train
          7: truck
          8: boat
          9: traffic light
          10: fire hydrant
          11: stop sign
          12: parking meter
          13: bench
          14: bird
          15: cat
          16: dog
          17: horse
          18: sheep
          19: cow
          20: elephant
          21: bear
          22: zebra
          23: giraffe
          24: backpack
          25: umbrella
          26: handbag
          27: tie
          28: suitcase
          29: frisbee
          30: skis
          31: snowboard
          32: sports ball
          33: kite
          34: baseball bat
          35: baseball glove
          36: skateboard
          37: surfboard
          38: tennis racket
          39: bottle
          40: wine glass
          41: cup
          42: fork
          43: knife
          44: spoon
          45: bowl
          46: banana
          47: apple
          48: sandwich
          49: orange
          50: broccoli
          51: carrot
          52: hot dog
          53: pizza
          54: donut
          55: cake
          56: chair
          57: couch
          58: potted plant
          59: bed
          60: dining table
          61: toilet
          62: tv
          63: laptop
          64: mouse
          65: remote
          66: keyboard
          67: cell phone
          68: microwave
          69: oven
          70: toaster
          71: sink
          72: refrigerator
          73: book
          74: clock
          75: vase
          76: scissors
          77: teddy bear
          78: hair drier
          79: toothbrush
        
        # Download script/URL (optional)
        download: https://ultralytics.com/assets/coco128.zip
        
      • 📄 xView.yaml 5.0KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # DIUx xView 2018 Challenge https://challenge.xviewdataset.org by U.S. National Geospatial-Intelligence Agency (NGA)
        # --------  DOWNLOAD DATA MANUALLY and jar xf val_images.zip to 'datasets/xView' before running train command!  --------
        # Example usage: python train.py --data xView.yaml
        # parent
        # ├── yolov5
        # └── datasets
        #     └── xView  ← downloads here (20.7 GB)
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/xView # dataset root dir
        train: images/autosplit_train.txt # train images (relative to 'path') 90% of 847 train images
        val: images/autosplit_val.txt # train images (relative to 'path') 10% of 847 train images
        
        # Classes
        names:
          0: Fixed-wing Aircraft
          1: Small Aircraft
          2: Cargo Plane
          3: Helicopter
          4: Passenger Vehicle
          5: Small Car
          6: Bus
          7: Pickup Truck
          8: Utility Truck
          9: Truck
          10: Cargo Truck
          11: Truck w/Box
          12: Truck Tractor
          13: Trailer
          14: Truck w/Flatbed
          15: Truck w/Liquid
          16: Crane Truck
          17: Railway Vehicle
          18: Passenger Car
          19: Cargo Car
          20: Flat Car
          21: Tank car
          22: Locomotive
          23: Maritime Vessel
          24: Motorboat
          25: Sailboat
          26: Tugboat
          27: Barge
          28: Fishing Vessel
          29: Ferry
          30: Yacht
          31: Container Ship
          32: Oil Tanker
          33: Engineering Vehicle
          34: Tower crane
          35: Container Crane
          36: Reach Stacker
          37: Straddle Carrier
          38: Mobile Crane
          39: Dump Truck
          40: Haul Truck
          41: Scraper/Tractor
          42: Front loader/Bulldozer
          43: Excavator
          44: Cement Mixer
          45: Ground Grader
          46: Hut/Tent
          47: Shed
          48: Building
          49: Aircraft Hangar
          50: Damaged Building
          51: Facility
          52: Construction Site
          53: Vehicle Lot
          54: Helipad
          55: Storage Tank
          56: Shipping container lot
          57: Shipping Container
          58: Pylon
          59: Tower
        
        # Download script/URL (optional) ---------------------------------------------------------------------------------------
        download: |
          import json
          import os
          from pathlib import Path
        
          import numpy as np
          from PIL import Image
          from tqdm import tqdm
        
          from utils.dataloaders import autosplit
          from utils.general import download, xyxy2xywhn
        
        
          def convert_labels(fname=Path('xView/xView_train.geojson')):
              # Convert xView geoJSON labels to YOLO format
              path = fname.parent
              with open(fname) as f:
                  print(f'Loading {fname}...')
                  data = json.load(f)
        
              # Make dirs
              labels = Path(path / 'labels' / 'train')
              os.system(f'rm -rf {labels}')
              labels.mkdir(parents=True, exist_ok=True)
        
              # xView classes 11-94 to 0-59
              xview_class2index = [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 0, 1, 2, -1, 3, -1, 4, 5, 6, 7, 8, -1, 9, 10, 11,
                                   12, 13, 14, 15, -1, -1, 16, 17, 18, 19, 20, 21, 22, -1, 23, 24, 25, -1, 26, 27, -1, 28, -1,
                                   29, 30, 31, 32, 33, 34, 35, 36, 37, -1, 38, 39, 40, 41, 42, 43, 44, 45, -1, -1, -1, -1, 46,
                                   47, 48, 49, -1, 50, 51, -1, 52, -1, -1, -1, 53, 54, -1, 55, -1, -1, 56, -1, 57, -1, 58, 59]
        
              shapes = {}
              for feature in tqdm(data['features'], desc=f'Converting {fname}'):
                  p = feature['properties']
                  if p['bounds_imcoords']:
                      id = p['image_id']
                      file = path / 'train_images' / id
                      if file.exists():  # 1395.tif missing
                          try:
                              box = np.array([int(num) for num in p['bounds_imcoords'].split(",")])
                              assert box.shape[0] == 4, f'incorrect box shape {box.shape[0]}'
                              cls = p['type_id']
                              cls = xview_class2index[int(cls)]  # xView class to 0-60
                              assert 59 >= cls >= 0, f'incorrect class index {cls}'
        
                              # Write YOLO label
                              if id not in shapes:
                                  shapes[id] = Image.open(file).size
                              box = xyxy2xywhn(box[None].astype(np.float), w=shapes[id][0], h=shapes[id][1], clip=True)
                              with open((labels / id).with_suffix('.txt'), 'a') as f:
                                  f.write(f"{cls} {' '.join(f'{x:.6f}' for x in box[0])}\n")  # write label.txt
                          except Exception as e:
                              print(f'WARNING: skipping one label for {file}: {e}')
        
        
          # Download manually from https://challenge.xviewdataset.org
          dir = Path(yaml['path'])  # dataset root dir
          # urls = ['https://d307kc0mrhucc3.cloudfront.net/train_labels.zip',  # train labels
          #         'https://d307kc0mrhucc3.cloudfront.net/train_images.zip',  # 15G, 847 train images
          #         'https://d307kc0mrhucc3.cloudfront.net/val_images.zip']  # 5G, 282 val images (no labels)
          # download(urls, dir=dir, delete=False)
        
          # Convert labels
          convert_labels(dir / 'xView_train.geojson')
        
          # Move images
          images = Path(dir / 'images')
          images.mkdir(parents=True, exist_ok=True)
          Path(dir / 'train_images').rename(dir / 'images' / 'train')
          Path(dir / 'val_images').rename(dir / 'images' / 'val')
        
          # Split
          autosplit(dir / 'images' / 'train')
        
      • 📄 zcy880.yaml 14.6KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        # ImageNet-1k dataset https://www.image-net.org/index.php by Stanford University
        # Simplified class names from https://github.com/anishathalye/imagenet-simple-labels
        # Example usage: python classify/train.py --data imagenet
        # parent
        # ├── yolov5
        # └── datasets
        #     └── imagenet10  ← downloads here
        
        # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
        path: ../datasets/zcy_dataset # dataset root dir
        train: train # train images (relative to 'path') 1281167 images
        val: val # val images (relative to 'path') 50000 images
        test: # test images (optional)
        
        # Classes
        names:
          0: 阿胶
          1: 阿魏
          2: 矮地茶
          3: 艾叶
          4: 安息香
          5: 桉叶油
          6: 八角枫
          7: 八角茴香
          8: 八角莲
          9: 八楞木
          10: 巴豆
          11: 巴戟天
          12: 菝葜
          13: 白背叶根
          14: 白扁豆
          15: 白矾
          16: 白附子
          17: 白果
          18: 白鹤灵芝
          19: 白花丹
          20: 白花蛇舌草
          21: 白及
          22: 白芥子
          23: 白兰花
          24: 白蔹
          25: 白马骨
          26: 白茅根
          27: 白前
          28: 白屈菜
          29: 白芍
          30: 白石花
          31: 白石英
          32: 白首乌
          33: 白术
          34: 白苏子
          35: 白头翁
          36: 白薇
          37: 白鲜皮
          38: 白药子
          39: 白英
          40: 白芷
          41: 百部
          42: 百合
          43: 柏子仁
          44: 败酱草
          45: 斑蝥
          46: 板蓝根
          47: 半边莲
          48: 半枫荷
          49: 半夏
          50: 半枝莲
          51: 薄荷
          52: 暴马丁香
          53: 北豆根
          54: 北刘寄奴
          55: 北沙参
          56: 荜茇
          57: 荜澄茄
          58: 蓖麻子
          59: 扁蓄
          60: 鳖甲
          61: 槟榔
          62: 冰片
          63: 卜芥
          64: 补骨脂
          65: 布渣叶
          66: 蚕沙
          67: 苍耳子
          68: 苍术
          69: 藏菖蒲
          70: 草豆蔻
          71: 草果
          72: 草乌
          73: 草乌叶
          74: 侧柏叶
          75: 柴胡
          76: 蝉花
          77: 蝉蜕
          78: 蟾酥
          79: 常山
          80: 车前草
          81: 车前子
          82: 沉香
          83: 陈皮
          84: 澄茄子
          85: 赤芍
          86: 赤石脂
          87: 赤小豆
          88: 茺蔚子
          89: 虫白蜡
          90: 虫牙药
          91: 臭灵丹草
          92: 臭牡丹
          93: 臭梧桐
          94: 楮实子
          95: 川贝母
          96: 川楝子
          97: 川木通
          98: 川木香
          99: 川牛膝
          100: 川射干
          101: 川乌
          102: 川芎
          103: 川续断
          104: 穿破石
          105: 穿山甲
          106: 穿山龙
          107: 穿心莲
          108: 垂盆草
          109: 春不见
          110: 椿皮
          111: 磁石
          112: 刺猬皮
          113: 刺五加
          114: 刺苋菜
          115: 葱白
          116: 葱子
          117: 楤木
          118: 翠云草
          119: 大驳骨
          120: 大豆黄卷
          121: 大飞扬草
          122: 大风艾
          123: 大风子
          124: 大腹皮
          125: 大黄
          126: 大戟
          127: 大蓟
          128: 大青盐
          129: 大青叶
          130: 大蒜
          131: 大血藤
          132: 大叶桉叶
          133: 大叶紫珠
          134: 大枣
          135: 玳瑁
          136: 丹参
          137: 单叶蔓荆子
          138: 胆矾
          139: 胆木
          140: 胆南星
          141: 淡豆豉
          142: 淡竹叶
          143: 当归
          144: 当药
          145: 党参
          146: 刀豆
          147: 倒扣草
          148: 稻芽
          149: 地不容
          150: 地耳草
          151: 地枫皮
          152: 地肤子
          153: 地骨皮
          154: 地黄
          155: 地锦草
          156: 地龙
          157: 地稔根
          158: 地桃花
          159: 地榆
          160: 灯心草
          161: 灯盏细辛
          162: 颠茄草
          163: 丁公藤
          164: 丁香
          165: 丁香罗勒油
          166: 东风橘
          167: 冬虫夏草
          168: 冬瓜皮
          169: 冬葵果
          170: 冬葵子
          171: 冬凌草
          172: 豆豉姜
          173: 豆蔻
          174: 独活
          175: 独脚金
          176: 独一味
          177: 杜仲
          178: 杜仲叶
          179: 断血流
          180: 莪术
          181: 峨参
          182: 鹅不食草
          183: 儿茶
          184: 番石榴叶
          185: 番泻叶
          186: 翻白草
          187: 方解石
          188: 防风
          189: 防风草
          190: 防己
          191: 飞龙掌血
          192: 飞天蠄蟧
          193: 榧子
          194: 粉萆薢
          195: 粉葛
          196: 枫香脂
          197: 蜂房
          198: 凤凰衣
          199: 凤尾草
          200: 凤眼草
          201: 凤眼果
          202: 伏龙肝
          203: 扶芳藤
          204: 扶桑花
          205: 佛甲草
          206: 佛手
          207: 茯苓
          208: 茯神
          209: 浮萍
          210: 浮小麦
          211: 附子
          212: 覆盆子
          213: 甘草
          214: 甘松
          215: 甘肃丹参
          216: 甘遂
          217: 干姜
          218: 干漆
          219: 岗梅根
          220: 岗稔根
          221: 杠板归
          222: 高良姜
          223: 藁本
          224: 葛根
          225: 功劳木
          226: 钩藤
          227: 狗肝菜
          228: 狗脊
          229: 狗肾
          230: 枸骨叶
          231: 枸杞子
          232: 谷精草
          233: 谷芽
          234: 骨碎补
          235: 瓜蒌
          236: 瓜子金
          237: 关黄柏
          238: 贯叶金丝桃
          239: 贯众
          240: 广东海桐皮
          241: 广东合欢花
          242: 广东刘寄奴
          243: 广东络石藤
          244: 广东王不留行
          245: 广藿香
          246: 广金钱草
          247: 广昆布
          248: 广山药
          249: 广升麻
          250: 广枣
          251: 龟甲
          252: 鬼箭羽
          253: 鬼针草
          254: 桂花
          255: 桂皮
          256: 桂枝
          257: 蛤蚧
          258: 海参
          259: 海风藤
          260: 海浮石
          261: 海蛤壳
          262: 海金沙
          263: 海龙
          264: 海马
          265: 海螵蛸
          266: 海桐皮
          267: 海芋
          268: 海藻
          269: 蔊菜
          270: 汉桃叶
          271: 诃子
          272: 合欢花
          273: 合欢皮
          274: 何首乌
          275: 荷叶
          276: 核桃仁
          277: 鹤草芽
          278: 鹤虱
          279: 黑豆
          280: 黑芝麻
          281: 黑种草子
          282: 红参
          283: 红大戟
          284: 红豆蔻
          285: 红粉
          286: 红花
          287: 红花龙胆
          288: 红景天
          289: 红芪
          290: 红曲
          291: 红丝线
          292: 红天葵
          293: 红药子
          294: 红鱼眼
          295: 猴耳环
          296: 厚朴
          297: 厚朴花
          298: 胡黄连
          299: 胡椒
          300: 胡芦巴
          301: 胡荽
          302: 胡颓子叶
          303: 葫芦
          304: 葫芦茶
          305: 湖北贝母
          306: 槲寄生
          307: 虎耳草
          308: 虎杖
          309: 琥珀
          310: 花椒
          311: 花蕊石
          312: 华山参
          313: 滑石
          314: 化橘红
          315: 槐花
          316: 槐角
          317: 槐枝
          318: 黄柏
          319: 黄花倒水莲
          320: 黄荆子
          321: 黄精
          322: 黄连
          323: 黄皮核
          324: 黄芪
          325: 黄芩
          326: 黄蜀葵
          327: 黄藤
          328: 黄药子
          329: 回回蒜
          330: 火麻仁
          331: 火炭母
          332: 藿香
          333: 鸡蛋花
          334: 鸡骨草
          335: 鸡骨香
          336: 鸡冠花
          337: 鸡内金
          338: 鸡矢藤
          339: 鸡血藤
          340: 鸡眼草
          341: 鸡爪大黄
          342: 积雪草
          343: 吉祥草
          344: 急性子
          345: 蒺藜
          346: 荠菜
          347: 荠苨
          348: 寄生黄
          349: 檵木叶
          350: 假芫茜
          351: 剑花
          352: 姜黄
          353: 僵蚕
          354: 降香
          355: 椒目
          356: 绞股蓝
          357: 芥子
          358: 金沸草
          359: 金果榄
          360: 金莲花
          361: 金礞石
          362: 金钱白花蛇
          363: 金钱草
          364: 金荞麦
          365: 金丝草
          366: 金铁锁
          367: 金银花
          368: 金樱子
          369: 筋骨草
          370: 锦灯笼
          371: 荆芥
          372: 景天三七
          373: 九层塔
          374: 九节菖蒲
          375: 九里香
          376: 九龙藤
          377: 九香虫
          378: 韭菜子
          379: 救必应
          380: 菊花
          381: 菊苣
          382: 橘核
          383: 卷柏
          384: 决明子
          385: 榼藤子
          386: 苦参
          387: 苦地胆
          388: 苦地丁
          389: 苦丁茶
          390: 苦瓜
          391: 苦楝皮
          392: 苦木
          393: 苦石莲
          394: 苦杏仁
          395: 苦玄参
          396: 宽筋藤
          397: 款冬花
          398: 昆布
          399: 辣椒
          400: 辣蓼
          401: 莱菔子
          402: 蓝布正
          403: 狼毒
          404: 榔榆皮
          405: 莨菪叶
          406: 老鹳草
          407: 了哥王
          408: 雷公藤
          409: 雷丸
          410: 藜
          411: 藜芦
          412: 荔枝核
          413: 连钱草
          414: 连翘
          415: 莲房
          416: 莲子
          417: 莲子心
          418: 凉粉草
          419: 两面针
          420: 两头尖
          421: 蓼大青叶
          422: 灵香草
          423: 灵芝
          424: 凌霄花
          425: 羚羊角
          426: 刘寄奴
          427: 硫黄
          428: 柳枝
          429: 六棱菊
          430: 六神曲
          431: 六月寒
          432: 龙齿
          433: 龙胆
          434: 龙骨
          435: 龙葵
          436: 龙脷叶
          437: 龙涎香
          438: 龙眼肉
          439: 漏芦
          440: 芦根
          441: 芦荟
          442: 炉甘石
          443: 陆英
          444: 鹿茸
          445: 鹿衔草
          446: 鹿药
          447: 路路通
          448: 罗布麻叶
          449: 罗汉果
          450: 罗勒
          451: 萝芙木
          452: 萝藦
          453: 络石藤
          454: 落地生根
          455: 绿豆
          456: 葎草
          457: 麻黄
          458: 麻黄根
          459: 麻油
          460: 马宝
          461: 马鞭草
          462: 马勃
          463: 马齿苋
          464: 马兜铃
          465: 马兰
          466: 马蔺子
          467: 马钱子
          468: 马桑根
          469: 马蹄金
          470: 麦冬
          471: 麦芽
          472: 满山红
          473: 蔓荆子
          474: 芒硝
          475: 猫爪草
          476: 毛草龙
          477: 毛诃子
          478: 没药
          479: 玫瑰花
          480: 梅花
          481: 虻虫
          482: 猕猴桃
          483: 密蒙花
          484: 绵萆薢
          485: 绵马贯众
          486: 明党参
          487: 茉莉花
          488: 墨旱莲
          489: 母丁香
          490: 牡丹皮
          491: 牡蒿
          492: 牡荆叶
          493: 牡蛎
          494: 木鳖子
          495: 木耳
          496: 木防己
          497: 木芙蓉叶
          498: 木瓜
          499: 木蝴蝶
          500: 木槿花
          501: 木槿皮
          502: 木棉花
          503: 木通
          504: 木香
          505: 木贼
          506: 南板蓝根
          507: 南瓜子
          508: 南鹤虱
          509: 南沙参
          510: 南山楂
          511: 南天竹叶
          512: 南五味子
          513: 硇砂
          514: 闹羊花
          515: 牛蒡子
          516: 牛大力
          517: 牛耳枫
          518: 牛黄
          519: 牛膝
          520: 牛至
          521: 扭肚藤
          522: 糯稻根
          523: 女贞子
          524: 藕节
          525: 爬山虎
          526: 盘龙参
          527: 盘龙七
          528: 胖大海
          529: 佩兰
          530: 硼砂
          531: 枇杷叶
          532: 啤酒花
          533: 辟汗草
          534: 片姜黄
          535: 平贝母
          536: 蒲公英
          537: 蒲黄
          538: 七叶莲
          539: 七叶一枝花
          540: 千斤拔
          541: 千金子
          542: 千里光
          543: 千年健
          544: 千日红
          545: 牵牛子
          546: 前胡
          547: 芡实
          548: 茜草
          549: 羌活
          550: 秦艽
          551: 秦皮
          552: 青黛
          553: 青风藤
          554: 青果
          555: 青蒿
          556: 青礞石
          557: 青木香
          558: 青牛胆
          559: 青皮
          560: 青天葵
          561: 青葙子
          562: 青羊参
          563: 青叶胆
          564: 苘麻子
          565: 瞿麦
          566: 全蝎
          567: 拳参
          568: 人参
          569: 人参叶
          570: 忍冬藤
          571: 榕树叶
          572: 肉苁蓉
          573: 肉豆蔻
          574: 肉桂
          575: 乳香
          576: 蕤仁
          577: 三白草
          578: 三分三
          579: 三棵针
          580: 三棱
          581: 三七
          582: 三丫苦
          583: 桑白皮
          584: 桑寄生
          585: 桑螵蛸
          586: 桑椹
          587: 桑叶
          588: 桑枝
          589: 沙棘
          590: 沙苑子
          591: 砂仁
          592: 山慈菇
          593: 山大颜
          594: 山豆根
          595: 山胡椒
          596: 山菅兰
          597: 山姜
          598: 山麦冬
          599: 山木通
          600: 山柰
          601: 山药
          602: 山银花
          603: 山萸肉
          604: 山楂
          605: 山楂叶
          606: 山芝麻
          607: 山茱萸
          608: 商陆
          609: 蛇床子
          610: 蛇莓
          611: 蛇蜕
          612: 射干
          613: 麝香
          614: 伸筋草
          615: 肾茶
          616: 升麻
          617: 生姜
          618: 生姜皮
          619: 蓍草
          620: 十萼茄
          621: 石蝉草
          622: 石菖蒲
          623: 石吊兰
          624: 石膏
          625: 石斛
          626: 石决明
          627: 石榴皮
          628: 石南藤
          629: 石楠叶
          630: 石上柏
          631: 石蒜
          632: 石韦
          633: 使君子
          634: 柿蒂
          635: 首乌藤
          636: 熟地黄
          637: 鼠李
          638: 鼠曲草
          639: 鼠尾草
          640: 双参
          641: 水半夏
          642: 水朝阳
          643: 水飞蓟
          644: 水河剑
          645: 水红花子
          646: 水茴香
          647: 水牛角
          648: 水翁花
          649: 水蛭
          650: 丝瓜络
          651: 四季青
          652: 四块瓦
          653: 四叶参
          654: 松花粉
          655: 松香
          656: 苏合香
          657: 苏木
          658: 酸枣仁
          659: 娑罗子
          660: 锁阳
          661: 太白参
          662: 太白花
          663: 太白米
          664: 太子参
          665: 檀香
          666: 桃仁
          667: 桃枝
          668: 藤黄
          669: 天冬
          670: 天胡荽
          671: 天花粉
          672: 天葵子
          673: 天麻
          674: 天南星
          675: 天山雪莲
          676: 天仙藤
          677: 天仙子
          678: 天竺黄
          679: 甜地丁
          680: 甜瓜蒂
          681: 甜瓜子
          682: 铁包金
          683: 铁皮石斛
          684: 铁苋
          685: 葶苈子
          686: 通草
          687: 透骨草
          688: 透骨香
          689: 土贝母
          690: 土鳖虫
          691: 土大黄
          692: 土党参
          693: 土茯苓
          694: 土荆芥
          695: 土荆皮
          696: 土木香
          697: 土牛膝
          698: 土人参
          699: 菟丝子
          700: 瓦楞子
          701: 瓦松
          702: 王不留行
          703: 威灵仙
          704: 委陵菜
          705: 问荆
          706: 乌桕子
          707: 乌蔹莓
          708: 乌梅
          709: 乌梢蛇
          710: 乌药
          711: 巫山淫羊藿
          712: 无根藤
          713: 无花果
          714: 无患子
          715: 吴茱萸
          716: 梧桐子
          717: 蜈蚣
          718: 五倍子
          719: 五加皮
          720: 五灵脂
          721: 五味子
          722: 五指毛
          723: 西番莲
          724: 西瓜皮
          725: 西河柳
          726: 西红花
          727: 西青果
          728: 西洋参
          729: 菥蓂
          730: 溪黄草
          731: 豨莶草
          732: 喜树果
          733: 细辛
          734: 霞草
          735: 夏枯草
          736: 夏天无
          737: 仙鹤草
          738: 仙茅
          739: 仙人掌
          740: 相思子
          741: 香附
          742: 香加皮
          743: 香薷
          744: 香橼
          745: 小驳骨
          746: 小飞扬草
          747: 小茴香
          748: 小蓟
          749: 小通草
          750: 小叶莲
          751: 缬草
          752: 薤白
          753: 辛夷
          754: 雄黄
          755: 熊胆
          756: 徐长卿
          757: 续断
          758: 萱草根
          759: 玄参
          760: 玄明粉
          761: 旋覆花
          762: 雪上一枝蒿
          763: 血风藤
          764: 血竭
          765: 血余炭
          766: 寻骨风
          767: 荨麻
          768: 鸦胆子
          769: 鸭脚木皮
          770: 鸭跖草
          771: 亚麻子
          772: 延胡索
          773: 芫花
          774: 芫荽子
          775: 岩参
          776: 岩黄连
          777: 岩椒草
          778: 盐蛇干
          779: 燕窝
          780: 羊耳菊
          781: 羊蹄
          782: 羊蹄草
          783: 阳起石
          784: 洋金花
          785: 野菊花
          786: 野马追
          787: 叶下珠
          788: 叶象花
          789: 一点红
          790: 一点血
          791: 一支箭
          792: 一枝黄花
          793: 伊贝母
          794: 益母草
          795: 益智
          796: 薏苡仁
          797: 翼首草
          798: 阴香皮
          799: 茵陈
          800: 银柴胡
          801: 银耳
          802: 银杏叶
          803: 淫羊藿
          804: 罂粟壳
          805: 樱桃核
          806: 鹰不泊
          807: 迎春花
          808: 油松节
          809: 余甘子
          810: 鱼胆草
          811: 鱼腥草
          812: 禹余粮
          813: 禹州漏芦
          814: 玉米须
          815: 玉叶金花
          816: 玉竹
          817: 郁金
          818: 郁李仁
          819: 预知子
          820: 远志
          821: 月季花
          822: 粤丝瓜络
          823: 芸薹子
          824: 皂角
          825: 皂角刺
          826: 泽兰
          827: 泽漆
          828: 泽泻
          829: 樟木
          830: 樟脑
          831: 长春花
          832: 掌叶半夏
          833: 赭石
          834: 浙贝母
          835: 珍珠
          836: 珍珠母
          837: 知母
          838: 栀子
          839: 蜘蛛香
          840: 枳椇子
          841: 枳壳
          842: 枳实
          843: 炙甘草
          844: 炙黄芪
          845: 钟乳石
          846: 肿节风
          847: 重楼
          848: 朱砂
          849: 朱砂根
          850: 珠子参
          851: 猪胆粉
          852: 猪苓
          853: 猪牙皂
          854: 竹节参
          855: 竹沥
          856: 竹茹
          857: 竹叶椒
          858: 苎麻根
          859: 梓实
          860: 紫贝齿
          861: 紫草
          862: 紫草茸
          863: 紫河车
          864: 紫花地丁
          865: 紫花前胡
          866: 紫梢花
          867: 紫石英
          868: 紫苏梗
          869: 紫苏叶
          870: 紫苏子
          871: 紫玉簪
          872: 紫苑
          873: 紫珠叶
          874: 自然铜
          875: 棕榈
          876: 走马箭
          877: 走马胎
          878: 醉鱼草
          879: 酢浆草
        
        
        # Download script/URL (optional)
        download: data/scripts/get_imagenet10.sh
        
    • 📁 models/
      • 📁 hub/
        • 📄 anchors.yaml 3.3KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Default anchors for COCO data
          
          # P5 -------------------------------------------------------------------------------------------------------------------
          # P5-640:
          anchors_p5_640:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # P6 -------------------------------------------------------------------------------------------------------------------
          # P6-640:  thr=0.25: 0.9964 BPR, 5.54 anchors past thr, n=12, img_size=640, metric_all=0.281/0.716-mean/best, past_thr=0.469-mean: 9,11,  21,19,  17,41,  43,32,  39,70,  86,64,  65,131,  134,130,  120,265,  282,180,  247,354,  512,387
          anchors_p6_640:
            - [9, 11, 21, 19, 17, 41] # P3/8
            - [43, 32, 39, 70, 86, 64] # P4/16
            - [65, 131, 134, 130, 120, 265] # P5/32
            - [282, 180, 247, 354, 512, 387] # P6/64
          
          # P6-1280:  thr=0.25: 0.9950 BPR, 5.55 anchors past thr, n=12, img_size=1280, metric_all=0.281/0.714-mean/best, past_thr=0.468-mean: 19,27,  44,40,  38,94,  96,68,  86,152,  180,137,  140,301,  303,264,  238,542,  436,615,  739,380,  925,792
          anchors_p6_1280:
            - [19, 27, 44, 40, 38, 94] # P3/8
            - [96, 68, 86, 152, 180, 137] # P4/16
            - [140, 301, 303, 264, 238, 542] # P5/32
            - [436, 615, 739, 380, 925, 792] # P6/64
          
          # P6-1920:  thr=0.25: 0.9950 BPR, 5.55 anchors past thr, n=12, img_size=1920, metric_all=0.281/0.714-mean/best, past_thr=0.468-mean: 28,41,  67,59,  57,141,  144,103,  129,227,  270,205,  209,452,  455,396,  358,812,  653,922,  1109,570,  1387,1187
          anchors_p6_1920:
            - [28, 41, 67, 59, 57, 141] # P3/8
            - [144, 103, 129, 227, 270, 205] # P4/16
            - [209, 452, 455, 396, 358, 812] # P5/32
            - [653, 922, 1109, 570, 1387, 1187] # P6/64
          
          # P7 -------------------------------------------------------------------------------------------------------------------
          # P7-640:  thr=0.25: 0.9962 BPR, 6.76 anchors past thr, n=15, img_size=640, metric_all=0.275/0.733-mean/best, past_thr=0.466-mean: 11,11,  13,30,  29,20,  30,46,  61,38,  39,92,  78,80,  146,66,  79,163,  149,150,  321,143,  157,303,  257,402,  359,290,  524,372
          anchors_p7_640:
            - [11, 11, 13, 30, 29, 20] # P3/8
            - [30, 46, 61, 38, 39, 92] # P4/16
            - [78, 80, 146, 66, 79, 163] # P5/32
            - [149, 150, 321, 143, 157, 303] # P6/64
            - [257, 402, 359, 290, 524, 372] # P7/128
          
          # P7-1280:  thr=0.25: 0.9968 BPR, 6.71 anchors past thr, n=15, img_size=1280, metric_all=0.273/0.732-mean/best, past_thr=0.463-mean: 19,22,  54,36,  32,77,  70,83,  138,71,  75,173,  165,159,  148,334,  375,151,  334,317,  251,626,  499,474,  750,326,  534,814,  1079,818
          anchors_p7_1280:
            - [19, 22, 54, 36, 32, 77] # P3/8
            - [70, 83, 138, 71, 75, 173] # P4/16
            - [165, 159, 148, 334, 375, 151] # P5/32
            - [334, 317, 251, 626, 499, 474] # P6/64
            - [750, 326, 534, 814, 1079, 818] # P7/128
          
          # P7-1920:  thr=0.25: 0.9968 BPR, 6.71 anchors past thr, n=15, img_size=1920, metric_all=0.273/0.732-mean/best, past_thr=0.463-mean: 29,34,  81,55,  47,115,  105,124,  207,107,  113,259,  247,238,  222,500,  563,227,  501,476,  376,939,  749,711,  1126,489,  801,1222,  1618,1227
          anchors_p7_1920:
            - [29, 34, 81, 55, 47, 115] # P3/8
            - [105, 124, 207, 107, 113, 259] # P4/16
            - [247, 238, 222, 500, 563, 227] # P5/32
            - [501, 476, 376, 939, 749, 711] # P6/64
            - [1126, 489, 801, 1222, 1618, 1227] # P7/128
          
        • 📄 yolov3-spp.yaml 1.6KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.0 # model depth multiple
          width_multiple: 1.0 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # darknet53 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [32, 3, 1]], # 0
              [-1, 1, Conv, [64, 3, 2]], # 1-P1/2
              [-1, 1, Bottleneck, [64]],
              [-1, 1, Conv, [128, 3, 2]], # 3-P2/4
              [-1, 2, Bottleneck, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 5-P3/8
              [-1, 8, Bottleneck, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 7-P4/16
              [-1, 8, Bottleneck, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32
              [-1, 4, Bottleneck, [1024]], # 10
            ]
          
          # YOLOv3-SPP head
          head: [
              [-1, 1, Bottleneck, [1024, False]],
              [-1, 1, SPP, [512, [5, 9, 13]]],
              [-1, 1, Conv, [1024, 3, 1]],
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large)
          
              [-2, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 8], 1, Concat, [1]], # cat backbone P4
              [-1, 1, Bottleneck, [512, False]],
              [-1, 1, Bottleneck, [512, False]],
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium)
          
              [-2, 1, Conv, [128, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P3
              [-1, 1, Bottleneck, [256, False]],
              [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small)
          
              [[27, 22, 15], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov3-tiny.yaml 1.2KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.0 # model depth multiple
          width_multiple: 1.0 # layer channel multiple
          anchors:
            - [10, 14, 23, 27, 37, 58] # P4/16
            - [81, 82, 135, 169, 344, 319] # P5/32
          
          # YOLOv3-tiny backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [16, 3, 1]], # 0
              [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 1-P1/2
              [-1, 1, Conv, [32, 3, 1]],
              [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 3-P2/4
              [-1, 1, Conv, [64, 3, 1]],
              [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 5-P3/8
              [-1, 1, Conv, [128, 3, 1]],
              [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 7-P4/16
              [-1, 1, Conv, [256, 3, 1]],
              [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 9-P5/32
              [-1, 1, Conv, [512, 3, 1]],
              [-1, 1, nn.ZeroPad2d, [[0, 1, 0, 1]]], # 11
              [-1, 1, nn.MaxPool2d, [2, 1, 0]], # 12
            ]
          
          # YOLOv3-tiny head
          head: [
              [-1, 1, Conv, [1024, 3, 1]],
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, Conv, [512, 3, 1]], # 15 (P5/32-large)
          
              [-2, 1, Conv, [128, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 8], 1, Concat, [1]], # cat backbone P4
              [-1, 1, Conv, [256, 3, 1]], # 19 (P4/16-medium)
          
              [[19, 15], 1, Detect, [nc, anchors]], # Detect(P4, P5)
            ]
          
        • 📄 yolov3.yaml 1.5KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.0 # model depth multiple
          width_multiple: 1.0 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # darknet53 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [32, 3, 1]], # 0
              [-1, 1, Conv, [64, 3, 2]], # 1-P1/2
              [-1, 1, Bottleneck, [64]],
              [-1, 1, Conv, [128, 3, 2]], # 3-P2/4
              [-1, 2, Bottleneck, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 5-P3/8
              [-1, 8, Bottleneck, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 7-P4/16
              [-1, 8, Bottleneck, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32
              [-1, 4, Bottleneck, [1024]], # 10
            ]
          
          # YOLOv3 head
          head: [
              [-1, 1, Bottleneck, [1024, False]],
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, Conv, [1024, 3, 1]],
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large)
          
              [-2, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 8], 1, Concat, [1]], # cat backbone P4
              [-1, 1, Bottleneck, [512, False]],
              [-1, 1, Bottleneck, [512, False]],
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium)
          
              [-2, 1, Conv, [128, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P3
              [-1, 1, Bottleneck, [256, False]],
              [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small)
          
              [[27, 22, 15], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov5-bifpn.yaml 1.4KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.0 # model depth multiple
          width_multiple: 1.0 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 BiFPN head
          head: [
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 13
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 14, 6], 1, Concat, [1]], # cat P4 <--- BiFPN change
              [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 10], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
          
              [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov5-fpn.yaml 1.2KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.0 # model depth multiple
          width_multiple: 1.0 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 FPN head
          head: [
              [-1, 3, C3, [1024, False]], # 10 (P5/32-large)
          
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 3, C3, [512, False]], # 14 (P4/16-medium)
          
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 3, C3, [256, False]], # 18 (P3/8-small)
          
              [[18, 14, 10], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov5-p2.yaml 1.7KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.0 # model depth multiple
          width_multiple: 1.0 # layer channel multiple
          anchors: 3 # AutoAnchor evolves 3 anchors per P output layer
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 head with (P2, P3, P4, P5) outputs
          head: [
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 13
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, Conv, [128, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 2], 1, Concat, [1]], # cat backbone P2
              [-1, 1, C3, [128, False]], # 21 (P2/4-xsmall)
          
              [-1, 1, Conv, [128, 3, 2]],
              [[-1, 18], 1, Concat, [1]], # cat head P3
              [-1, 3, C3, [256, False]], # 24 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 27 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 10], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [1024, False]], # 30 (P5/32-large)
          
              [[21, 24, 27, 30], 1, Detect, [nc, anchors]], # Detect(P2, P3, P4, P5)
            ]
          
        • 📄 yolov5-p34.yaml 1.2KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 0.33 # model depth multiple
          width_multiple: 0.50 # layer channel multiple
          anchors: 3 # AutoAnchor evolves 3 anchors per P output layer
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 head with (P3, P4) outputs
          head: [
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 13
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
          
              [[17, 20], 1, Detect, [nc, anchors]], # Detect(P3, P4)
            ]
          
        • 📄 yolov5-p6.yaml 1.7KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.0 # model depth multiple
          width_multiple: 1.0 # layer channel multiple
          anchors: 3 # AutoAnchor evolves 3 anchors per P output layer
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [768, 3, 2]], # 7-P5/32
              [-1, 3, C3, [768]],
              [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 11
            ]
          
          # YOLOv5 v6.0 head with (P3, P4, P5, P6) outputs
          head: [
              [-1, 1, Conv, [768, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 8], 1, Concat, [1]], # cat backbone P5
              [-1, 3, C3, [768, False]], # 15
          
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 19
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 23 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 20], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 26 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 16], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [768, False]], # 29 (P5/32-large)
          
              [-1, 1, Conv, [768, 3, 2]],
              [[-1, 12], 1, Concat, [1]], # cat head P6
              [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge)
          
              [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6)
            ]
          
        • 📄 yolov5-p7.yaml 2.1KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.0 # model depth multiple
          width_multiple: 1.0 # layer channel multiple
          anchors: 3 # AutoAnchor evolves 3 anchors per P output layer
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [768, 3, 2]], # 7-P5/32
              [-1, 3, C3, [768]],
              [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64
              [-1, 3, C3, [1024]],
              [-1, 1, Conv, [1280, 3, 2]], # 11-P7/128
              [-1, 3, C3, [1280]],
              [-1, 1, SPPF, [1280, 5]], # 13
            ]
          
          # YOLOv5 v6.0 head with (P3, P4, P5, P6, P7) outputs
          head: [
              [-1, 1, Conv, [1024, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 10], 1, Concat, [1]], # cat backbone P6
              [-1, 3, C3, [1024, False]], # 17
          
              [-1, 1, Conv, [768, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 8], 1, Concat, [1]], # cat backbone P5
              [-1, 3, C3, [768, False]], # 21
          
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 25
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 29 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 26], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 32 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 22], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [768, False]], # 35 (P5/32-large)
          
              [-1, 1, Conv, [768, 3, 2]],
              [[-1, 18], 1, Concat, [1]], # cat head P6
              [-1, 3, C3, [1024, False]], # 38 (P6/64-xlarge)
          
              [-1, 1, Conv, [1024, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P7
              [-1, 3, C3, [1280, False]], # 41 (P7/128-xxlarge)
          
              [[29, 32, 35, 38, 41], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6, P7)
            ]
          
        • 📄 yolov5-panet.yaml 1.4KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.0 # model depth multiple
          width_multiple: 1.0 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 PANet head
          head: [
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 13
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 10], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
          
              [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov5l6.yaml 1.8KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.0 # model depth multiple
          width_multiple: 1.0 # layer channel multiple
          anchors:
            - [19, 27, 44, 40, 38, 94] # P3/8
            - [96, 68, 86, 152, 180, 137] # P4/16
            - [140, 301, 303, 264, 238, 542] # P5/32
            - [436, 615, 739, 380, 925, 792] # P6/64
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [768, 3, 2]], # 7-P5/32
              [-1, 3, C3, [768]],
              [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 11
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [768, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 8], 1, Concat, [1]], # cat backbone P5
              [-1, 3, C3, [768, False]], # 15
          
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 19
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 23 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 20], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 26 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 16], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [768, False]], # 29 (P5/32-large)
          
              [-1, 1, Conv, [768, 3, 2]],
              [[-1, 12], 1, Concat, [1]], # cat head P6
              [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge)
          
              [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6)
            ]
          
        • 📄 yolov5m6.yaml 1.8KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 0.67 # model depth multiple
          width_multiple: 0.75 # layer channel multiple
          anchors:
            - [19, 27, 44, 40, 38, 94] # P3/8
            - [96, 68, 86, 152, 180, 137] # P4/16
            - [140, 301, 303, 264, 238, 542] # P5/32
            - [436, 615, 739, 380, 925, 792] # P6/64
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [768, 3, 2]], # 7-P5/32
              [-1, 3, C3, [768]],
              [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 11
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [768, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 8], 1, Concat, [1]], # cat backbone P5
              [-1, 3, C3, [768, False]], # 15
          
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 19
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 23 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 20], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 26 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 16], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [768, False]], # 29 (P5/32-large)
          
              [-1, 1, Conv, [768, 3, 2]],
              [[-1, 12], 1, Concat, [1]], # cat head P6
              [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge)
          
              [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6)
            ]
          
        • 📄 yolov5n6.yaml 1.8KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 0.33 # model depth multiple
          width_multiple: 0.25 # layer channel multiple
          anchors:
            - [19, 27, 44, 40, 38, 94] # P3/8
            - [96, 68, 86, 152, 180, 137] # P4/16
            - [140, 301, 303, 264, 238, 542] # P5/32
            - [436, 615, 739, 380, 925, 792] # P6/64
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [768, 3, 2]], # 7-P5/32
              [-1, 3, C3, [768]],
              [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 11
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [768, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 8], 1, Concat, [1]], # cat backbone P5
              [-1, 3, C3, [768, False]], # 15
          
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 19
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 23 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 20], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 26 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 16], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [768, False]], # 29 (P5/32-large)
          
              [-1, 1, Conv, [768, 3, 2]],
              [[-1, 12], 1, Concat, [1]], # cat head P6
              [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge)
          
              [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6)
            ]
          
        • 📄 yolov5s-LeakyReLU.yaml 1.5KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          activation: nn.LeakyReLU(0.1) # <----- Conv() activation used throughout entire YOLOv5 model
          depth_multiple: 0.33 # model depth multiple
          width_multiple: 0.50 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 13
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 10], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
          
              [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov5s-ghost.yaml 1.5KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 0.33 # model depth multiple
          width_multiple: 0.50 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, GhostConv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3Ghost, [128]],
              [-1, 1, GhostConv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3Ghost, [256]],
              [-1, 1, GhostConv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3Ghost, [512]],
              [-1, 1, GhostConv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3Ghost, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, GhostConv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3Ghost, [512, False]], # 13
          
              [-1, 1, GhostConv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3Ghost, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, GhostConv, [256, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P4
              [-1, 3, C3Ghost, [512, False]], # 20 (P4/16-medium)
          
              [-1, 1, GhostConv, [512, 3, 2]],
              [[-1, 10], 1, Concat, [1]], # cat head P5
              [-1, 3, C3Ghost, [1024, False]], # 23 (P5/32-large)
          
              [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov5s-transformer.yaml 1.4KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 0.33 # model depth multiple
          width_multiple: 0.50 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3TR, [1024]], # 9 <--- C3TR() Transformer module
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 13
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 10], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
          
              [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov5s6.yaml 1.8KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 0.33 # model depth multiple
          width_multiple: 0.50 # layer channel multiple
          anchors:
            - [19, 27, 44, 40, 38, 94] # P3/8
            - [96, 68, 86, 152, 180, 137] # P4/16
            - [140, 301, 303, 264, 238, 542] # P5/32
            - [436, 615, 739, 380, 925, 792] # P6/64
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [768, 3, 2]], # 7-P5/32
              [-1, 3, C3, [768]],
              [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 11
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [768, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 8], 1, Concat, [1]], # cat backbone P5
              [-1, 3, C3, [768, False]], # 15
          
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 19
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 23 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 20], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 26 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 16], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [768, False]], # 29 (P5/32-large)
          
              [-1, 1, Conv, [768, 3, 2]],
              [[-1, 12], 1, Concat, [1]], # cat head P6
              [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge)
          
              [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6)
            ]
          
        • 📄 yolov5x6.yaml 1.8KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.33 # model depth multiple
          width_multiple: 1.25 # layer channel multiple
          anchors:
            - [19, 27, 44, 40, 38, 94] # P3/8
            - [96, 68, 86, 152, 180, 137] # P4/16
            - [140, 301, 303, 264, 238, 542] # P5/32
            - [436, 615, 739, 380, 925, 792] # P6/64
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [768, 3, 2]], # 7-P5/32
              [-1, 3, C3, [768]],
              [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 11
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [768, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 8], 1, Concat, [1]], # cat backbone P5
              [-1, 3, C3, [768, False]], # 15
          
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 19
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 23 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 20], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 26 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 16], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [768, False]], # 29 (P5/32-large)
          
              [-1, 1, Conv, [768, 3, 2]],
              [[-1, 12], 1, Concat, [1]], # cat head P6
              [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge)
          
              [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6)
            ]
          
      • 📁 segment/
        • 📄 yolov5l-seg.yaml 1.4KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.0 # model depth multiple
          width_multiple: 1.0 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 13
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 10], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
          
              [[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov5m-seg.yaml 1.4KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 0.67 # model depth multiple
          width_multiple: 0.75 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 13
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 10], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
          
              [[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov5n-seg.yaml 1.4KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 0.33 # model depth multiple
          width_multiple: 0.25 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 13
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 10], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
          
              [[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov5s-seg.yaml 1.4KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 0.33 # model depth multiple
          width_multiple: 0.5 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 13
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 10], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
          
              [[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5)
            ]
          
        • 📄 yolov5x-seg.yaml 1.4KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          
          # Parameters
          nc: 80 # number of classes
          depth_multiple: 1.33 # model depth multiple
          width_multiple: 1.25 # layer channel multiple
          anchors:
            - [10, 13, 16, 30, 33, 23] # P3/8
            - [30, 61, 62, 45, 59, 119] # P4/16
            - [116, 90, 156, 198, 373, 326] # P5/32
          
          # YOLOv5 v6.0 backbone
          backbone:
            # [from, number, module, args]
            [
              [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
              [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
              [-1, 3, C3, [128]],
              [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
              [-1, 6, C3, [256]],
              [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
              [-1, 9, C3, [512]],
              [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
              [-1, 3, C3, [1024]],
              [-1, 1, SPPF, [1024, 5]], # 9
            ]
          
          # YOLOv5 v6.0 head
          head: [
              [-1, 1, Conv, [512, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 6], 1, Concat, [1]], # cat backbone P4
              [-1, 3, C3, [512, False]], # 13
          
              [-1, 1, Conv, [256, 1, 1]],
              [-1, 1, nn.Upsample, [None, 2, "nearest"]],
              [[-1, 4], 1, Concat, [1]], # cat backbone P3
              [-1, 3, C3, [256, False]], # 17 (P3/8-small)
          
              [-1, 1, Conv, [256, 3, 2]],
              [[-1, 14], 1, Concat, [1]], # cat head P4
              [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
          
              [-1, 1, Conv, [512, 3, 2]],
              [[-1, 10], 1, Concat, [1]], # cat head P5
              [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
          
              [[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5)
            ]
          
      • 📄 AI_yolov5s.yaml 1.4KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        
        # Parameters
        nc: 80 # number of classes
        depth_multiple: 0.33 # model depth multiple
        width_multiple: 0.50 # layer channel multiple
        anchors:
          - [10, 13, 16, 30, 33, 23] # P3/8
          - [30, 61, 62, 45, 59, 119] # P4/16
          - [116, 90, 156, 198, 373, 326] # P5/32
        
        # YOLOv5 v6.0 backbone
        backbone:
          # [from, number, module, args]
          [
            [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
            [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
            [-1, 3, C3, [128]],  # 2
            [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
            [-1, 6, C3, [256]], # 4
            [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
            [-1, 9, C3, [512]], # 6
            [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
            [-1, 3, C3, [1024]],
            [-1, 1, SPPF, [1024, 5]], # 9
          ]
        
        # YOLOv5 v6.0 head
        head: [
            [-1, 1, Conv, [512, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 6], 1, Concat, [1]], # cat backbone P4
            [-1, 3, C3, [512, False]], # 13
        
            [-1, 1, Conv, [256, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 4], 1, Concat, [1]], # cat backbone P3
            [-1, 3, C3, [256, False]], # 17 (P3/8-small)
        
            [-1, 1, Conv, [256, 3, 2]],
            [[-1, 14], 1, Concat, [1]], # cat head P4
            [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
        
            [-1, 1, Conv, [512, 3, 2]],
            [[-1, 10], 1, Concat, [1]], # cat head P5
            [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
        
            [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
          ]
        
      • 📄 __init__.py 0B
        
        
      • 📄 common.py 49.4KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Common modules."""
        
        import ast
        import contextlib
        import json
        import math
        import platform
        import warnings
        import zipfile
        from collections import OrderedDict, namedtuple
        from copy import copy
        from pathlib import Path
        from urllib.parse import urlparse
        
        import cv2
        import numpy as np
        import pandas as pd
        import requests
        import torch
        import torch.nn as nn
        from PIL import Image
        from torch.cuda import amp
        
        # Import 'ultralytics' package or install if missing
        try:
            import ultralytics
        
            assert hasattr(ultralytics, "__version__")  # verify package is not directory
        except (ImportError, AssertionError):
            import os
        
            os.system("pip install -U ultralytics")
            import ultralytics
        
        from ultralytics.utils.plotting import Annotator, colors, save_one_box
        
        from utils import TryExcept
        from utils.dataloaders import exif_transpose, letterbox
        from utils.general import (
            LOGGER,
            ROOT,
            Profile,
            check_requirements,
            check_suffix,
            check_version,
            colorstr,
            increment_path,
            is_jupyter,
            make_divisible,
            non_max_suppression,
            scale_boxes,
            xywh2xyxy,
            xyxy2xywh,
            yaml_load,
        )
        from utils.torch_utils import copy_attr, smart_inference_mode
        
        
        def autopad(k, p=None, d=1):
            """
            Pads kernel to 'same' output shape, adjusting for optional dilation; returns padding size.
        
            `k`: kernel, `p`: padding, `d`: dilation.
            """
            if d > 1:
                k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
            if p is None:
                p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
            return p
        
        
        class Conv(nn.Module):
            # Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)
            default_act = nn.SiLU()  # default activation
        
            def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
                """Initializes a standard convolution layer with optional batch normalization and activation."""
                super().__init__()
                self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
                self.bn = nn.BatchNorm2d(c2)
                self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
        
            def forward(self, x):
                """Applies a convolution followed by batch normalization and an activation function to the input tensor `x`."""
                return self.act(self.bn(self.conv(x)))
        
            def forward_fuse(self, x):
                """Applies a fused convolution and activation function to the input tensor `x`."""
                return self.act(self.conv(x))
        
        
        class DWConv(Conv):
            # Depth-wise convolution
            def __init__(self, c1, c2, k=1, s=1, d=1, act=True):
                """Initializes a depth-wise convolution layer with optional activation; args: input channels (c1), output
                channels (c2), kernel size (k), stride (s), dilation (d), and activation flag (act).
                """
                super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), d=d, act=act)
        
        
        class DWConvTranspose2d(nn.ConvTranspose2d):
            # Depth-wise transpose convolution
            def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0):
                """Initializes a depth-wise transpose convolutional layer for YOLOv5; args: input channels (c1), output channels
                (c2), kernel size (k), stride (s), input padding (p1), output padding (p2).
                """
                super().__init__(c1, c2, k, s, p1, p2, groups=math.gcd(c1, c2))
        
        
        class TransformerLayer(nn.Module):
            # Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
            def __init__(self, c, num_heads):
                """
                Initializes a transformer layer, sans LayerNorm for performance, with multihead attention and linear layers.
        
                See  as described in https://arxiv.org/abs/2010.11929.
                """
                super().__init__()
                self.q = nn.Linear(c, c, bias=False)
                self.k = nn.Linear(c, c, bias=False)
                self.v = nn.Linear(c, c, bias=False)
                self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
                self.fc1 = nn.Linear(c, c, bias=False)
                self.fc2 = nn.Linear(c, c, bias=False)
        
            def forward(self, x):
                """Performs forward pass using MultiheadAttention and two linear transformations with residual connections."""
                x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
                x = self.fc2(self.fc1(x)) + x
                return x
        
        
        class TransformerBlock(nn.Module):
            # Vision Transformer https://arxiv.org/abs/2010.11929
            def __init__(self, c1, c2, num_heads, num_layers):
                """Initializes a Transformer block for vision tasks, adapting dimensions if necessary and stacking specified
                layers.
                """
                super().__init__()
                self.conv = None
                if c1 != c2:
                    self.conv = Conv(c1, c2)
                self.linear = nn.Linear(c2, c2)  # learnable position embedding
                self.tr = nn.Sequential(*(TransformerLayer(c2, num_heads) for _ in range(num_layers)))
                self.c2 = c2
        
            def forward(self, x):
                """Processes input through an optional convolution, followed by Transformer layers and position embeddings for
                object detection.
                """
                if self.conv is not None:
                    x = self.conv(x)
                b, _, w, h = x.shape
                p = x.flatten(2).permute(2, 0, 1)
                return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, w, h)
        
        
        class Bottleneck(nn.Module):
            # Standard bottleneck
            def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):
                """Initializes a standard bottleneck layer with optional shortcut and group convolution, supporting channel
                expansion.
                """
                super().__init__()
                c_ = int(c2 * e)  # hidden channels
                self.cv1 = Conv(c1, c_, 1, 1)
                self.cv2 = Conv(c_, c2, 3, 1, g=g)
                self.add = shortcut and c1 == c2
        
            def forward(self, x):
                """Processes input through two convolutions, optionally adds shortcut if channel dimensions match; input is a
                tensor.
                """
                return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
        
        
        class BottleneckCSP(nn.Module):
            # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
            def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
                """Initializes CSP bottleneck with optional shortcuts; args: ch_in, ch_out, number of repeats, shortcut bool,
                groups, expansion.
                """
                super().__init__()
                c_ = int(c2 * e)  # hidden channels
                self.cv1 = Conv(c1, c_, 1, 1)
                self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
                self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
                self.cv4 = Conv(2 * c_, c2, 1, 1)
                self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
                self.act = nn.SiLU()
                self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
        
            def forward(self, x):
                """Performs forward pass by applying layers, activation, and concatenation on input x, returning feature-
                enhanced output.
                """
                y1 = self.cv3(self.m(self.cv1(x)))
                y2 = self.cv2(x)
                return self.cv4(self.act(self.bn(torch.cat((y1, y2), 1))))
        
        
        class CrossConv(nn.Module):
            # Cross Convolution Downsample
            def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
                """
                Initializes CrossConv with downsampling, expanding, and optionally shortcutting; `c1` input, `c2` output
                channels.
        
                Inputs are ch_in, ch_out, kernel, stride, groups, expansion, shortcut.
                """
                super().__init__()
                c_ = int(c2 * e)  # hidden channels
                self.cv1 = Conv(c1, c_, (1, k), (1, s))
                self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
                self.add = shortcut and c1 == c2
        
            def forward(self, x):
                """Performs feature sampling, expanding, and applies shortcut if channels match; expects `x` input tensor."""
                return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
        
        
        class C3(nn.Module):
            # CSP Bottleneck with 3 convolutions
            def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
                """Initializes C3 module with options for channel count, bottleneck repetition, shortcut usage, group
                convolutions, and expansion.
                """
                super().__init__()
                c_ = int(c2 * e)  # hidden channels
                self.cv1 = Conv(c1, c_, 1, 1)
                self.cv2 = Conv(c1, c_, 1, 1)
                self.cv3 = Conv(2 * c_, c2, 1)  # optional act=FReLU(c2)
                self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
        
            def forward(self, x):
                """Performs forward propagation using concatenated outputs from two convolutions and a Bottleneck sequence."""
                return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
        
        
        class C3x(C3):
            # C3 module with cross-convolutions
            def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
                """Initializes C3x module with cross-convolutions, extending C3 with customizable channel dimensions, groups,
                and expansion.
                """
                super().__init__(c1, c2, n, shortcut, g, e)
                c_ = int(c2 * e)
                self.m = nn.Sequential(*(CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)))
        
        
        class C3TR(C3):
            # C3 module with TransformerBlock()
            def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
                """Initializes C3 module with TransformerBlock for enhanced feature extraction, accepts channel sizes, shortcut
                config, group, and expansion.
                """
                super().__init__(c1, c2, n, shortcut, g, e)
                c_ = int(c2 * e)
                self.m = TransformerBlock(c_, c_, 4, n)
        
        
        class C3SPP(C3):
            # C3 module with SPP()
            def __init__(self, c1, c2, k=(5, 9, 13), n=1, shortcut=True, g=1, e=0.5):
                """Initializes a C3 module with SPP layer for advanced spatial feature extraction, given channel sizes, kernel
                sizes, shortcut, group, and expansion ratio.
                """
                super().__init__(c1, c2, n, shortcut, g, e)
                c_ = int(c2 * e)
                self.m = SPP(c_, c_, k)
        
        
        class C3Ghost(C3):
            # C3 module with GhostBottleneck()
            def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
                """Initializes YOLOv5's C3 module with Ghost Bottlenecks for efficient feature extraction."""
                super().__init__(c1, c2, n, shortcut, g, e)
                c_ = int(c2 * e)  # hidden channels
                self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))
        
        
        class SPP(nn.Module):
            # Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729
            def __init__(self, c1, c2, k=(5, 9, 13)):
                """Initializes SPP layer with Spatial Pyramid Pooling, ref: https://arxiv.org/abs/1406.4729, args: c1 (input channels), c2 (output channels), k (kernel sizes)."""
                super().__init__()
                c_ = c1 // 2  # hidden channels
                self.cv1 = Conv(c1, c_, 1, 1)
                self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
                self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
        
            def forward(self, x):
                """Applies convolution and max pooling layers to the input tensor `x`, concatenates results, and returns output
                tensor.
                """
                x = self.cv1(x)
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")  # suppress torch 1.9.0 max_pool2d() warning
                    return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
        
        
        class SPPF(nn.Module):
            # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
            def __init__(self, c1, c2, k=5):
                """
                Initializes YOLOv5 SPPF layer with given channels and kernel size for YOLOv5 model, combining convolution and
                max pooling.
        
                Equivalent to SPP(k=(5, 9, 13)).
                """
                super().__init__()
                c_ = c1 // 2  # hidden channels
                self.cv1 = Conv(c1, c_, 1, 1)
                self.cv2 = Conv(c_ * 4, c2, 1, 1)
                self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
        
            def forward(self, x):
                """Processes input through a series of convolutions and max pooling operations for feature extraction."""
                x = self.cv1(x)
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")  # suppress torch 1.9.0 max_pool2d() warning
                    y1 = self.m(x)
                    y2 = self.m(y1)
                    return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
        
        
        class Focus(nn.Module):
            # Focus wh information into c-space
            def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
                """Initializes Focus module to concentrate width-height info into channel space with configurable convolution
                parameters.
                """
                super().__init__()
                self.conv = Conv(c1 * 4, c2, k, s, p, g, act=act)
                # self.contract = Contract(gain=2)
        
            def forward(self, x):
                """Processes input through Focus mechanism, reshaping (b,c,w,h) to (b,4c,w/2,h/2) then applies convolution."""
                return self.conv(torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1))
                # return self.conv(self.contract(x))
        
        
        class GhostConv(nn.Module):
            # Ghost Convolution https://github.com/huawei-noah/ghostnet
            def __init__(self, c1, c2, k=1, s=1, g=1, act=True):
                """Initializes GhostConv with in/out channels, kernel size, stride, groups, and activation; halves out channels
                for efficiency.
                """
                super().__init__()
                c_ = c2 // 2  # hidden channels
                self.cv1 = Conv(c1, c_, k, s, None, g, act=act)
                self.cv2 = Conv(c_, c_, 5, 1, None, c_, act=act)
        
            def forward(self, x):
                """Performs forward pass, concatenating outputs of two convolutions on input `x`: shape (B,C,H,W)."""
                y = self.cv1(x)
                return torch.cat((y, self.cv2(y)), 1)
        
        
        class GhostBottleneck(nn.Module):
            # Ghost Bottleneck https://github.com/huawei-noah/ghostnet
            def __init__(self, c1, c2, k=3, s=1):
                """Initializes GhostBottleneck with ch_in `c1`, ch_out `c2`, kernel size `k`, stride `s`; see https://github.com/huawei-noah/ghostnet."""
                super().__init__()
                c_ = c2 // 2
                self.conv = nn.Sequential(
                    GhostConv(c1, c_, 1, 1),  # pw
                    DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(),  # dw
                    GhostConv(c_, c2, 1, 1, act=False),
                )  # pw-linear
                self.shortcut = (
                    nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
                )
        
            def forward(self, x):
                """Processes input through conv and shortcut layers, returning their summed output."""
                return self.conv(x) + self.shortcut(x)
        
        
        class Contract(nn.Module):
            # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
            def __init__(self, gain=2):
                """Initializes a layer to contract spatial dimensions (width-height) into channels, e.g., input shape
                (1,64,80,80) to (1,256,40,40).
                """
                super().__init__()
                self.gain = gain
        
            def forward(self, x):
                """Processes input tensor to expand channel dimensions by contracting spatial dimensions, yielding output shape
                `(b, c*s*s, h//s, w//s)`.
                """
                b, c, h, w = x.size()  # assert (h / s == 0) and (W / s == 0), 'Indivisible gain'
                s = self.gain
                x = x.view(b, c, h // s, s, w // s, s)  # x(1,64,40,2,40,2)
                x = x.permute(0, 3, 5, 1, 2, 4).contiguous()  # x(1,2,2,64,40,40)
                return x.view(b, c * s * s, h // s, w // s)  # x(1,256,40,40)
        
        
        class Expand(nn.Module):
            # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)
            def __init__(self, gain=2):
                """
                Initializes the Expand module to increase spatial dimensions by redistributing channels, with an optional gain
                factor.
        
                Example: x(1,64,80,80) to x(1,16,160,160).
                """
                super().__init__()
                self.gain = gain
        
            def forward(self, x):
                """Processes input tensor x to expand spatial dimensions by redistributing channels, requiring C / gain^2 ==
                0.
                """
                b, c, h, w = x.size()  # assert C / s ** 2 == 0, 'Indivisible gain'
                s = self.gain
                x = x.view(b, s, s, c // s**2, h, w)  # x(1,2,2,16,80,80)
                x = x.permute(0, 3, 4, 1, 5, 2).contiguous()  # x(1,16,80,2,80,2)
                return x.view(b, c // s**2, h * s, w * s)  # x(1,16,160,160)
        
        
        class Concat(nn.Module):
            # Concatenate a list of tensors along dimension
            def __init__(self, dimension=1):
                """Initializes a Concat module to concatenate tensors along a specified dimension."""
                super().__init__()
                self.d = dimension
        
            def forward(self, x):
                """Concatenates a list of tensors along a specified dimension; `x` is a list of tensors, `dimension` is an
                int.
                """
                return torch.cat(x, self.d)
        
        
        class DetectMultiBackend(nn.Module):
            # YOLOv5 MultiBackend class for python inference on various backends
            def __init__(self, weights="yolov5s.pt", device=torch.device("cpu"), dnn=False, data=None, fp16=False, fuse=True):
                """Initializes DetectMultiBackend with support for various inference backends, including PyTorch and ONNX."""
                #   PyTorch:              weights = *.pt
                #   TorchScript:                    *.torchscript
                #   ONNX Runtime:                   *.onnx
                #   ONNX OpenCV DNN:                *.onnx --dnn
                #   OpenVINO:                       *_openvino_model
                #   CoreML:                         *.mlmodel
                #   TensorRT:                       *.engine
                #   TensorFlow SavedModel:          *_saved_model
                #   TensorFlow GraphDef:            *.pb
                #   TensorFlow Lite:                *.tflite
                #   TensorFlow Edge TPU:            *_edgetpu.tflite
                #   PaddlePaddle:                   *_paddle_model
                from models.experimental import attempt_download, attempt_load  # scoped to avoid circular import
        
                super().__init__()
                w = str(weights[0] if isinstance(weights, list) else weights)
                pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, triton = self._model_type(w)
                fp16 &= pt or jit or onnx or engine or triton  # FP16
                nhwc = coreml or saved_model or pb or tflite or edgetpu  # BHWC formats (vs torch BCWH)
                stride = 32  # default stride
                cuda = torch.cuda.is_available() and device.type != "cpu"  # use CUDA
                if not (pt or triton):
                    w = attempt_download(w)  # download if not local
        
                if pt:  # PyTorch
                    model = attempt_load(weights if isinstance(weights, list) else w, device=device, inplace=True, fuse=fuse)
                    stride = max(int(model.stride.max()), 32)  # model stride
                    names = model.module.names if hasattr(model, "module") else model.names  # get class names
                    model.half() if fp16 else model.float()
                    self.model = model  # explicitly assign for to(), cpu(), cuda(), half()
                elif jit:  # TorchScript
                    LOGGER.info(f"Loading {w} for TorchScript inference...")
                    extra_files = {"config.txt": ""}  # model metadata
                    model = torch.jit.load(w, _extra_files=extra_files, map_location=device)
                    model.half() if fp16 else model.float()
                    if extra_files["config.txt"]:  # load metadata dict
                        d = json.loads(
                            extra_files["config.txt"],
                            object_hook=lambda d: {int(k) if k.isdigit() else k: v for k, v in d.items()},
                        )
                        stride, names = int(d["stride"]), d["names"]
                elif dnn:  # ONNX OpenCV DNN
                    LOGGER.info(f"Loading {w} for ONNX OpenCV DNN inference...")
                    check_requirements("opencv-python>=4.5.4")
                    net = cv2.dnn.readNetFromONNX(w)
                elif onnx:  # ONNX Runtime
                    LOGGER.info(f"Loading {w} for ONNX Runtime inference...")
                    check_requirements(("onnx", "onnxruntime-gpu" if cuda else "onnxruntime"))
                    import onnxruntime
        
                    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if cuda else ["CPUExecutionProvider"]
                    session = onnxruntime.InferenceSession(w, providers=providers)
                    output_names = [x.name for x in session.get_outputs()]
                    meta = session.get_modelmeta().custom_metadata_map  # metadata
                    if "stride" in meta:
                        stride, names = int(meta["stride"]), eval(meta["names"])
                elif xml:  # OpenVINO
                    LOGGER.info(f"Loading {w} for OpenVINO inference...")
                    check_requirements("openvino>=2023.0")  # requires openvino-dev: https://pypi.org/project/openvino-dev/
                    from openvino.runtime import Core, Layout, get_batch
        
                    core = Core()
                    if not Path(w).is_file():  # if not *.xml
                        w = next(Path(w).glob("*.xml"))  # get *.xml file from *_openvino_model dir
                    ov_model = core.read_model(model=w, weights=Path(w).with_suffix(".bin"))
                    if ov_model.get_parameters()[0].get_layout().empty:
                        ov_model.get_parameters()[0].set_layout(Layout("NCHW"))
                    batch_dim = get_batch(ov_model)
                    if batch_dim.is_static:
                        batch_size = batch_dim.get_length()
                    ov_compiled_model = core.compile_model(ov_model, device_name="AUTO")  # AUTO selects best available device
                    stride, names = self._load_metadata(Path(w).with_suffix(".yaml"))  # load metadata
                elif engine:  # TensorRT
                    LOGGER.info(f"Loading {w} for TensorRT inference...")
                    import tensorrt as trt  # https://developer.nvidia.com/nvidia-tensorrt-download
        
                    check_version(trt.__version__, "7.0.0", hard=True)  # require tensorrt>=7.0.0
                    if device.type == "cpu":
                        device = torch.device("cuda:0")
                    Binding = namedtuple("Binding", ("name", "dtype", "shape", "data", "ptr"))
                    logger = trt.Logger(trt.Logger.INFO)
                    with open(w, "rb") as f, trt.Runtime(logger) as runtime:
                        model = runtime.deserialize_cuda_engine(f.read())
                    context = model.create_execution_context()
                    bindings = OrderedDict()
                    output_names = []
                    fp16 = False  # default updated below
                    dynamic = False
                    is_trt10 = not hasattr(model, "num_bindings")
                    num = range(model.num_io_tensors) if is_trt10 else range(model.num_bindings)
                    for i in num:
                        if is_trt10:
                            name = model.get_tensor_name(i)
                            dtype = trt.nptype(model.get_tensor_dtype(name))
                            is_input = model.get_tensor_mode(name) == trt.TensorIOMode.INPUT
                            if is_input:
                                if -1 in tuple(model.get_tensor_shape(name)):  # dynamic
                                    dynamic = True
                                    context.set_input_shape(name, tuple(model.get_profile_shape(name, 0)[2]))
                                if dtype == np.float16:
                                    fp16 = True
                            else:  # output
                                output_names.append(name)
                            shape = tuple(context.get_tensor_shape(name))
                        else:
                            name = model.get_binding_name(i)
                            dtype = trt.nptype(model.get_binding_dtype(i))
                            if model.binding_is_input(i):
                                if -1 in tuple(model.get_binding_shape(i)):  # dynamic
                                    dynamic = True
                                    context.set_binding_shape(i, tuple(model.get_profile_shape(0, i)[2]))
                                if dtype == np.float16:
                                    fp16 = True
                            else:  # output
                                output_names.append(name)
                            shape = tuple(context.get_binding_shape(i))
                        im = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device)
                        bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr()))
                    binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
                    batch_size = bindings["images"].shape[0]  # if dynamic, this is instead max batch size
                elif coreml:  # CoreML
                    LOGGER.info(f"Loading {w} for CoreML inference...")
                    import coremltools as ct
        
                    model = ct.models.MLModel(w)
                elif saved_model:  # TF SavedModel
                    LOGGER.info(f"Loading {w} for TensorFlow SavedModel inference...")
                    import tensorflow as tf
        
                    keras = False  # assume TF1 saved_model
                    model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w)
                elif pb:  # GraphDef https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
                    LOGGER.info(f"Loading {w} for TensorFlow GraphDef inference...")
                    import tensorflow as tf
        
                    def wrap_frozen_graph(gd, inputs, outputs):
                        """Wraps a TensorFlow GraphDef for inference, returning a pruned function."""
                        x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), [])  # wrapped
                        ge = x.graph.as_graph_element
                        return x.prune(tf.nest.map_structure(ge, inputs), tf.nest.map_structure(ge, outputs))
        
                    def gd_outputs(gd):
                        """Generates a sorted list of graph outputs excluding NoOp nodes and inputs, formatted as '<name>:0'."""
                        name_list, input_list = [], []
                        for node in gd.node:  # tensorflow.core.framework.node_def_pb2.NodeDef
                            name_list.append(node.name)
                            input_list.extend(node.input)
                        return sorted(f"{x}:0" for x in list(set(name_list) - set(input_list)) if not x.startswith("NoOp"))
        
                    gd = tf.Graph().as_graph_def()  # TF GraphDef
                    with open(w, "rb") as f:
                        gd.ParseFromString(f.read())
                    frozen_func = wrap_frozen_graph(gd, inputs="x:0", outputs=gd_outputs(gd))
                elif tflite or edgetpu:  # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
                    try:  # https://coral.ai/docs/edgetpu/tflite-python/#update-existing-tf-lite-code-for-the-edge-tpu
                        from tflite_runtime.interpreter import Interpreter, load_delegate
                    except ImportError:
                        import tensorflow as tf
        
                        Interpreter, load_delegate = (
                            tf.lite.Interpreter,
                            tf.lite.experimental.load_delegate,
                        )
                    if edgetpu:  # TF Edge TPU https://coral.ai/software/#edgetpu-runtime
                        LOGGER.info(f"Loading {w} for TensorFlow Lite Edge TPU inference...")
                        delegate = {"Linux": "libedgetpu.so.1", "Darwin": "libedgetpu.1.dylib", "Windows": "edgetpu.dll"}[
                            platform.system()
                        ]
                        interpreter = Interpreter(model_path=w, experimental_delegates=[load_delegate(delegate)])
                    else:  # TFLite
                        LOGGER.info(f"Loading {w} for TensorFlow Lite inference...")
                        interpreter = Interpreter(model_path=w)  # load TFLite model
                    interpreter.allocate_tensors()  # allocate
                    input_details = interpreter.get_input_details()  # inputs
                    output_details = interpreter.get_output_details()  # outputs
                    # load metadata
                    with contextlib.suppress(zipfile.BadZipFile):
                        with zipfile.ZipFile(w, "r") as model:
                            meta_file = model.namelist()[0]
                            meta = ast.literal_eval(model.read(meta_file).decode("utf-8"))
                            stride, names = int(meta["stride"]), meta["names"]
                elif tfjs:  # TF.js
                    raise NotImplementedError("ERROR: YOLOv5 TF.js inference is not supported")
                elif paddle:  # PaddlePaddle
                    LOGGER.info(f"Loading {w} for PaddlePaddle inference...")
                    check_requirements("paddlepaddle-gpu" if cuda else "paddlepaddle")
                    import paddle.inference as pdi
        
                    if not Path(w).is_file():  # if not *.pdmodel
                        w = next(Path(w).rglob("*.pdmodel"))  # get *.pdmodel file from *_paddle_model dir
                    weights = Path(w).with_suffix(".pdiparams")
                    config = pdi.Config(str(w), str(weights))
                    if cuda:
                        config.enable_use_gpu(memory_pool_init_size_mb=2048, device_id=0)
                    predictor = pdi.create_predictor(config)
                    input_handle = predictor.get_input_handle(predictor.get_input_names()[0])
                    output_names = predictor.get_output_names()
                elif triton:  # NVIDIA Triton Inference Server
                    LOGGER.info(f"Using {w} as Triton Inference Server...")
                    check_requirements("tritonclient[all]")
                    from utils.triton import TritonRemoteModel
        
                    model = TritonRemoteModel(url=w)
                    nhwc = model.runtime.startswith("tensorflow")
                else:
                    raise NotImplementedError(f"ERROR: {w} is not a supported format")
        
                # class names
                if "names" not in locals():
                    names = yaml_load(data)["names"] if data else {i: f"class{i}" for i in range(999)}
                if names[0] == "n01440764" and len(names) == 1000:  # ImageNet
                    names = yaml_load(ROOT / "data/ImageNet.yaml")["names"]  # human-readable names
        
                self.__dict__.update(locals())  # assign all variables to self
        
            def forward(self, im, augment=False, visualize=False):
                """Performs YOLOv5 inference on input images with options for augmentation and visualization."""
                b, ch, h, w = im.shape  # batch, channel, height, width
                if self.fp16 and im.dtype != torch.float16:
                    im = im.half()  # to FP16
                if self.nhwc:
                    im = im.permute(0, 2, 3, 1)  # torch BCHW to numpy BHWC shape(1,320,192,3)
        
                if self.pt:  # PyTorch
                    y = self.model(im, augment=augment, visualize=visualize) if augment or visualize else self.model(im)
                elif self.jit:  # TorchScript
                    y = self.model(im)
                elif self.dnn:  # ONNX OpenCV DNN
                    im = im.cpu().numpy()  # torch to numpy
                    self.net.setInput(im)
                    y = self.net.forward()
                elif self.onnx:  # ONNX Runtime
                    im = im.cpu().numpy()  # torch to numpy
                    y = self.session.run(self.output_names, {self.session.get_inputs()[0].name: im})
                elif self.xml:  # OpenVINO
                    im = im.cpu().numpy()  # FP32
                    y = list(self.ov_compiled_model(im).values())
                elif self.engine:  # TensorRT
                    if self.dynamic and im.shape != self.bindings["images"].shape:
                        i = self.model.get_binding_index("images")
                        self.context.set_binding_shape(i, im.shape)  # reshape if dynamic
                        self.bindings["images"] = self.bindings["images"]._replace(shape=im.shape)
                        for name in self.output_names:
                            i = self.model.get_binding_index(name)
                            self.bindings[name].data.resize_(tuple(self.context.get_binding_shape(i)))
                    s = self.bindings["images"].shape
                    assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"
                    self.binding_addrs["images"] = int(im.data_ptr())
                    self.context.execute_v2(list(self.binding_addrs.values()))
                    y = [self.bindings[x].data for x in sorted(self.output_names)]
                elif self.coreml:  # CoreML
                    im = im.cpu().numpy()
                    im = Image.fromarray((im[0] * 255).astype("uint8"))
                    # im = im.resize((192, 320), Image.BILINEAR)
                    y = self.model.predict({"image": im})  # coordinates are xywh normalized
                    if "confidence" in y:
                        box = xywh2xyxy(y["coordinates"] * [[w, h, w, h]])  # xyxy pixels
                        conf, cls = y["confidence"].max(1), y["confidence"].argmax(1).astype(np.float)
                        y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
                    else:
                        y = list(reversed(y.values()))  # reversed for segmentation models (pred, proto)
                elif self.paddle:  # PaddlePaddle
                    im = im.cpu().numpy().astype(np.float32)
                    self.input_handle.copy_from_cpu(im)
                    self.predictor.run()
                    y = [self.predictor.get_output_handle(x).copy_to_cpu() for x in self.output_names]
                elif self.triton:  # NVIDIA Triton Inference Server
                    y = self.model(im)
                else:  # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
                    im = im.cpu().numpy()
                    if self.saved_model:  # SavedModel
                        y = self.model(im, training=False) if self.keras else self.model(im)
                    elif self.pb:  # GraphDef
                        y = self.frozen_func(x=self.tf.constant(im))
                    else:  # Lite or Edge TPU
                        input = self.input_details[0]
                        int8 = input["dtype"] == np.uint8  # is TFLite quantized uint8 model
                        if int8:
                            scale, zero_point = input["quantization"]
                            im = (im / scale + zero_point).astype(np.uint8)  # de-scale
                        self.interpreter.set_tensor(input["index"], im)
                        self.interpreter.invoke()
                        y = []
                        for output in self.output_details:
                            x = self.interpreter.get_tensor(output["index"])
                            if int8:
                                scale, zero_point = output["quantization"]
                                x = (x.astype(np.float32) - zero_point) * scale  # re-scale
                            y.append(x)
                    y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y]
                    y[0][..., :4] *= [w, h, w, h]  # xywh normalized to pixels
        
                if isinstance(y, (list, tuple)):
                    return self.from_numpy(y[0]) if len(y) == 1 else [self.from_numpy(x) for x in y]
                else:
                    return self.from_numpy(y)
        
            def from_numpy(self, x):
                """Converts a NumPy array to a torch tensor, maintaining device compatibility."""
                return torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x
        
            def warmup(self, imgsz=(1, 3, 640, 640)):
                """Performs a single inference warmup to initialize model weights, accepting an `imgsz` tuple for image size."""
                warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton
                if any(warmup_types) and (self.device.type != "cpu" or self.triton):
                    im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device)  # input
                    for _ in range(2 if self.jit else 1):  #
                        self.forward(im)  # warmup
        
            @staticmethod
            def _model_type(p="path/to/model.pt"):
                """
                Determines model type from file path or URL, supporting various export formats.
        
                Example: path='path/to/model.onnx' -> type=onnx
                """
                # types = [pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle]
                from export import export_formats
                from utils.downloads import is_url
        
                sf = list(export_formats().Suffix)  # export suffixes
                if not is_url(p, check=False):
                    check_suffix(p, sf)  # checks
                url = urlparse(p)  # if url may be Triton inference server
                types = [s in Path(p).name for s in sf]
                types[8] &= not types[9]  # tflite &= not edgetpu
                triton = not any(types) and all([any(s in url.scheme for s in ["http", "grpc"]), url.netloc])
                return types + [triton]
        
            @staticmethod
            def _load_metadata(f=Path("path/to/meta.yaml")):
                """Loads metadata from a YAML file, returning strides and names if the file exists, otherwise `None`."""
                if f.exists():
                    d = yaml_load(f)
                    return d["stride"], d["names"]  # assign stride, names
                return None, None
        
        
        class AutoShape(nn.Module):
            # YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
            conf = 0.25  # NMS confidence threshold
            iou = 0.45  # NMS IoU threshold
            agnostic = False  # NMS class-agnostic
            multi_label = False  # NMS multiple labels per box
            classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
            max_det = 1000  # maximum number of detections per image
            amp = False  # Automatic Mixed Precision (AMP) inference
        
            def __init__(self, model, verbose=True):
                """Initializes YOLOv5 model for inference, setting up attributes and preparing model for evaluation."""
                super().__init__()
                if verbose:
                    LOGGER.info("Adding AutoShape... ")
                copy_attr(self, model, include=("yaml", "nc", "hyp", "names", "stride", "abc"), exclude=())  # copy attributes
                self.dmb = isinstance(model, DetectMultiBackend)  # DetectMultiBackend() instance
                self.pt = not self.dmb or model.pt  # PyTorch model
                self.model = model.eval()
                if self.pt:
                    m = self.model.model.model[-1] if self.dmb else self.model.model[-1]  # Detect()
                    m.inplace = False  # Detect.inplace=False for safe multithread inference
                    m.export = True  # do not output loss values
        
            def _apply(self, fn):
                """
                Applies to(), cpu(), cuda(), half() etc.
        
                to model tensors excluding parameters or registered buffers.
                """
                self = super()._apply(fn)
                if self.pt:
                    m = self.model.model.model[-1] if self.dmb else self.model.model[-1]  # Detect()
                    m.stride = fn(m.stride)
                    m.grid = list(map(fn, m.grid))
                    if isinstance(m.anchor_grid, list):
                        m.anchor_grid = list(map(fn, m.anchor_grid))
                return self
        
            @smart_inference_mode()
            def forward(self, ims, size=640, augment=False, profile=False):
                """
                Performs inference on inputs with optional augment & profiling.
        
                Supports various formats including file, URI, OpenCV, PIL, numpy, torch.
                """
                # For size(height=640, width=1280), RGB images example inputs are:
                #   file:        ims = 'data/images/zidane.jpg'  # str or PosixPath
                #   URI:             = 'https://ultralytics.com/images/zidane.jpg'
                #   OpenCV:          = cv2.imread('image.jpg')[:,:,::-1]  # HWC BGR to RGB x(640,1280,3)
                #   PIL:             = Image.open('image.jpg') or ImageGrab.grab()  # HWC x(640,1280,3)
                #   numpy:           = np.zeros((640,1280,3))  # HWC
                #   torch:           = torch.zeros(16,3,320,640)  # BCHW (scaled to size=640, 0-1 values)
                #   multiple:        = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...]  # list of images
        
                dt = (Profile(), Profile(), Profile())
                with dt[0]:
                    if isinstance(size, int):  # expand
                        size = (size, size)
                    p = next(self.model.parameters()) if self.pt else torch.empty(1, device=self.model.device)  # param
                    autocast = self.amp and (p.device.type != "cpu")  # Automatic Mixed Precision (AMP) inference
                    if isinstance(ims, torch.Tensor):  # torch
                        with amp.autocast(autocast):
                            return self.model(ims.to(p.device).type_as(p), augment=augment)  # inference
        
                    # Pre-process
                    n, ims = (len(ims), list(ims)) if isinstance(ims, (list, tuple)) else (1, [ims])  # number, list of images
                    shape0, shape1, files = [], [], []  # image and inference shapes, filenames
                    for i, im in enumerate(ims):
                        f = f"image{i}"  # filename
                        if isinstance(im, (str, Path)):  # filename or uri
                            im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith("http") else im), im
                            im = np.asarray(exif_transpose(im))
                        elif isinstance(im, Image.Image):  # PIL Image
                            im, f = np.asarray(exif_transpose(im)), getattr(im, "filename", f) or f
                        files.append(Path(f).with_suffix(".jpg").name)
                        if im.shape[0] < 5:  # image in CHW
                            im = im.transpose((1, 2, 0))  # reverse dataloader .transpose(2, 0, 1)
                        im = im[..., :3] if im.ndim == 3 else cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)  # enforce 3ch input
                        s = im.shape[:2]  # HWC
                        shape0.append(s)  # image shape
                        g = max(size) / max(s)  # gain
                        shape1.append([int(y * g) for y in s])
                        ims[i] = im if im.data.contiguous else np.ascontiguousarray(im)  # update
                    shape1 = [make_divisible(x, self.stride) for x in np.array(shape1).max(0)]  # inf shape
                    x = [letterbox(im, shape1, auto=False)[0] for im in ims]  # pad
                    x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2)))  # stack and BHWC to BCHW
                    x = torch.from_numpy(x).to(p.device).type_as(p) / 255  # uint8 to fp16/32
        
                with amp.autocast(autocast):
                    # Inference
                    with dt[1]:
                        y = self.model(x, augment=augment)  # forward
        
                    # Post-process
                    with dt[2]:
                        y = non_max_suppression(
                            y if self.dmb else y[0],
                            self.conf,
                            self.iou,
                            self.classes,
                            self.agnostic,
                            self.multi_label,
                            max_det=self.max_det,
                        )  # NMS
                        for i in range(n):
                            scale_boxes(shape1, y[i][:, :4], shape0[i])
        
                    return Detections(ims, y, files, dt, self.names, x.shape)
        
        
        class Detections:
            # YOLOv5 detections class for inference results
            def __init__(self, ims, pred, files, times=(0, 0, 0), names=None, shape=None):
                """Initializes the YOLOv5 Detections class with image info, predictions, filenames, timing and normalization."""
                super().__init__()
                d = pred[0].device  # device
                gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1, 1], device=d) for im in ims]  # normalizations
                self.ims = ims  # list of images as numpy arrays
                self.pred = pred  # list of tensors pred[0] = (xyxy, conf, cls)
                self.names = names  # class names
                self.files = files  # image filenames
                self.times = times  # profiling times
                self.xyxy = pred  # xyxy pixels
                self.xywh = [xyxy2xywh(x) for x in pred]  # xywh pixels
                self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)]  # xyxy normalized
                self.xywhn = [x / g for x, g in zip(self.xywh, gn)]  # xywh normalized
                self.n = len(self.pred)  # number of images (batch size)
                self.t = tuple(x.t / self.n * 1e3 for x in times)  # timestamps (ms)
                self.s = tuple(shape)  # inference BCHW shape
        
            def _run(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path("")):
                """Executes model predictions, displaying and/or saving outputs with optional crops and labels."""
                s, crops = "", []
                for i, (im, pred) in enumerate(zip(self.ims, self.pred)):
                    s += f"\nimage {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} "  # string
                    if pred.shape[0]:
                        for c in pred[:, -1].unique():
                            n = (pred[:, -1] == c).sum()  # detections per class
                            s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "  # add to string
                        s = s.rstrip(", ")
                        if show or save or render or crop:
                            annotator = Annotator(im, example=str(self.names))
                            for *box, conf, cls in reversed(pred):  # xyxy, confidence, class
                                label = f"{self.names[int(cls)]} {conf:.2f}"
                                if crop:
                                    file = save_dir / "crops" / self.names[int(cls)] / self.files[i] if save else None
                                    crops.append(
                                        {
                                            "box": box,
                                            "conf": conf,
                                            "cls": cls,
                                            "label": label,
                                            "im": save_one_box(box, im, file=file, save=save),
                                        }
                                    )
                                else:  # all others
                                    annotator.box_label(box, label if labels else "", color=colors(cls))
                            im = annotator.im
                    else:
                        s += "(no detections)"
        
                    im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im  # from np
                    if show:
                        if is_jupyter():
                            from IPython.display import display
        
                            display(im)
                        else:
                            im.show(self.files[i])
                    if save:
                        f = self.files[i]
                        im.save(save_dir / f)  # save
                        if i == self.n - 1:
                            LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
                    if render:
                        self.ims[i] = np.asarray(im)
                if pprint:
                    s = s.lstrip("\n")
                    return f"{s}\nSpeed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {self.s}" % self.t
                if crop:
                    if save:
                        LOGGER.info(f"Saved results to {save_dir}\n")
                    return crops
        
            @TryExcept("Showing images is not supported in this environment")
            def show(self, labels=True):
                """
                Displays detection results with optional labels.
        
                Usage: show(labels=True)
                """
                self._run(show=True, labels=labels)  # show results
        
            def save(self, labels=True, save_dir="runs/detect/exp", exist_ok=False):
                """
                Saves detection results with optional labels to a specified directory.
        
                Usage: save(labels=True, save_dir='runs/detect/exp', exist_ok=False)
                """
                save_dir = increment_path(save_dir, exist_ok, mkdir=True)  # increment save_dir
                self._run(save=True, labels=labels, save_dir=save_dir)  # save results
        
            def crop(self, save=True, save_dir="runs/detect/exp", exist_ok=False):
                """
                Crops detection results, optionally saves them to a directory.
        
                Args: save (bool), save_dir (str), exist_ok (bool).
                """
                save_dir = increment_path(save_dir, exist_ok, mkdir=True) if save else None
                return self._run(crop=True, save=save, save_dir=save_dir)  # crop results
        
            def render(self, labels=True):
                """Renders detection results with optional labels on images; args: labels (bool) indicating label inclusion."""
                self._run(render=True, labels=labels)  # render results
                return self.ims
        
            def pandas(self):
                """
                Returns detections as pandas DataFrames for various box formats (xyxy, xyxyn, xywh, xywhn).
        
                Example: print(results.pandas().xyxy[0]).
                """
                new = copy(self)  # return copy
                ca = "xmin", "ymin", "xmax", "ymax", "confidence", "class", "name"  # xyxy columns
                cb = "xcenter", "ycenter", "width", "height", "confidence", "class", "name"  # xywh columns
                for k, c in zip(["xyxy", "xyxyn", "xywh", "xywhn"], [ca, ca, cb, cb]):
                    a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)]  # update
                    setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])
                return new
        
            def tolist(self):
                """
                Converts a Detections object into a list of individual detection results for iteration.
        
                Example: for result in results.tolist():
                """
                r = range(self.n)  # iterable
                return [
                    Detections(
                        [self.ims[i]],
                        [self.pred[i]],
                        [self.files[i]],
                        self.times,
                        self.names,
                        self.s,
                    )
                    for i in r
                ]
        
            def print(self):
                """Logs the string representation of the current object's state via the LOGGER."""
                LOGGER.info(self.__str__())
        
            def __len__(self):
                """Returns the number of results stored, overrides the default len(results)."""
                return self.n
        
            def __str__(self):
                """Returns a string representation of the model's results, suitable for printing, overrides default
                print(results).
                """
                return self._run(pprint=True)  # print results
        
            def __repr__(self):
                """Returns a string representation of the YOLOv5 object, including its class and formatted results."""
                return f"YOLOv5 {self.__class__} instance\n" + self.__str__()
        
        
        class Proto(nn.Module):
            # YOLOv5 mask Proto module for segmentation models
            def __init__(self, c1, c_=256, c2=32):
                """Initializes YOLOv5 Proto module for segmentation with input, proto, and mask channels configuration."""
                super().__init__()
                self.cv1 = Conv(c1, c_, k=3)
                self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
                self.cv2 = Conv(c_, c_, k=3)
                self.cv3 = Conv(c_, c2)
        
            def forward(self, x):
                """Performs a forward pass using convolutional layers and upsampling on input tensor `x`."""
                return self.cv3(self.cv2(self.upsample(self.cv1(x))))
        
        
        class Classify(nn.Module):
            # YOLOv5 classification head, i.e. x(b,c1,20,20) to x(b,c2)
            def __init__(
                self, c1, c2, k=1, s=1, p=None, g=1, dropout_p=0.0
            ):  # ch_in, ch_out, kernel, stride, padding, groups, dropout probability
                super().__init__()
                c_ = 1280  # efficientnet_b0 size
                self.conv = Conv(c1, c_, k, s, autopad(k, p), g)
                self.pool = nn.AdaptiveAvgPool2d(1)  # to x(b,c_,1,1)
                self.drop = nn.Dropout(p=dropout_p, inplace=True)
                self.linear = nn.Linear(c_, c2)  # to x(b,c2)
        
            def forward(self, x):
                """Processes input through conv, pool, drop, and linear layers; supports list concatenation input."""
                if isinstance(x, list):
                    x = torch.cat(x, 1)
                return self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))
        
      • 📄 experimental.py 5.1KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Experimental modules."""
        
        import math
        
        import numpy as np
        import torch
        import torch.nn as nn
        
        from utils.downloads import attempt_download
        
        
        class Sum(nn.Module):
            """Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070."""
        
            def __init__(self, n, weight=False):
                """Initializes a module to sum outputs of layers with number of inputs `n` and optional weighting, supporting 2+
                inputs.
                """
                super().__init__()
                self.weight = weight  # apply weights boolean
                self.iter = range(n - 1)  # iter object
                if weight:
                    self.w = nn.Parameter(-torch.arange(1.0, n) / 2, requires_grad=True)  # layer weights
        
            def forward(self, x):
                """Processes input through a customizable weighted sum of `n` inputs, optionally applying learned weights."""
                y = x[0]  # no weight
                if self.weight:
                    w = torch.sigmoid(self.w) * 2
                    for i in self.iter:
                        y = y + x[i + 1] * w[i]
                else:
                    for i in self.iter:
                        y = y + x[i + 1]
                return y
        
        
        class MixConv2d(nn.Module):
            """Mixed Depth-wise Conv https://arxiv.org/abs/1907.09595."""
        
            def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
                """Initializes MixConv2d with mixed depth-wise convolutional layers, taking input and output channels (c1, c2),
                kernel sizes (k), stride (s), and channel distribution strategy (equal_ch).
                """
                super().__init__()
                n = len(k)  # number of convolutions
                if equal_ch:  # equal c_ per group
                    i = torch.linspace(0, n - 1e-6, c2).floor()  # c2 indices
                    c_ = [(i == g).sum() for g in range(n)]  # intermediate channels
                else:  # equal weight.numel() per group
                    b = [c2] + [0] * n
                    a = np.eye(n + 1, n, k=-1)
                    a -= np.roll(a, 1, axis=1)
                    a *= np.array(k) ** 2
                    a[0] = 1
                    c_ = np.linalg.lstsq(a, b, rcond=None)[0].round()  # solve for equal weight indices, ax = b
        
                self.m = nn.ModuleList(
                    [nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)]
                )
                self.bn = nn.BatchNorm2d(c2)
                self.act = nn.SiLU()
        
            def forward(self, x):
                """Performs forward pass by applying SiLU activation on batch-normalized concatenated convolutional layer
                outputs.
                """
                return self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))
        
        
        class Ensemble(nn.ModuleList):
            """Ensemble of models."""
        
            def __init__(self):
                """Initializes an ensemble of models to be used for aggregated predictions."""
                super().__init__()
        
            def forward(self, x, augment=False, profile=False, visualize=False):
                """Performs forward pass aggregating outputs from an ensemble of models.."""
                y = [module(x, augment, profile, visualize)[0] for module in self]
                # y = torch.stack(y).max(0)[0]  # max ensemble
                # y = torch.stack(y).mean(0)  # mean ensemble
                y = torch.cat(y, 1)  # nms ensemble
                return y, None  # inference, train output
        
        
        def attempt_load(weights, device=None, inplace=True, fuse=True):
            """
            Loads and fuses an ensemble or single YOLOv5 model from weights, handling device placement and model adjustments.
        
            Example inputs: weights=[a,b,c] or a single model weights=[a] or weights=a.
            """
            from models.yolo import Detect, Model
        
            model = Ensemble()
            for w in weights if isinstance(weights, list) else [weights]:
                ckpt = torch.load(attempt_download(w), map_location="cpu",weights_only=False)  # load
                ckpt = (ckpt.get("ema") or ckpt["model"]).to(device).float()  # FP32 model
        
                # Model compatibility updates
                if not hasattr(ckpt, "stride"):
                    ckpt.stride = torch.tensor([32.0])
                if hasattr(ckpt, "names") and isinstance(ckpt.names, (list, tuple)):
                    ckpt.names = dict(enumerate(ckpt.names))  # convert to dict
        
                model.append(ckpt.fuse().eval() if fuse and hasattr(ckpt, "fuse") else ckpt.eval())  # model in eval mode
        
            # Module updates
            for m in model.modules():
                t = type(m)
                if t in (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Model):
                    m.inplace = inplace
                    if t is Detect and not isinstance(m.anchor_grid, list):
                        delattr(m, "anchor_grid")
                        setattr(m, "anchor_grid", [torch.zeros(1)] * m.nl)
                elif t is nn.Upsample and not hasattr(m, "recompute_scale_factor"):
                    m.recompute_scale_factor = None  # torch 1.11.0 compatibility
        
            # Return model
            if len(model) == 1:
                return model[-1]
        
            # Return detection ensemble
            print(f"Ensemble created with {weights}\n")
            for k in "names", "nc", "yaml":
                setattr(model, k, getattr(model[0], k))
            model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride  # max stride
            assert all(model[0].nc == m.nc for m in model), f"Models have different class counts: {[m.nc for m in model]}"
            return model
        
      • 📄 tf.py 31.2KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """
        TensorFlow, Keras and TFLite versions of YOLOv5
        Authored by https://github.com/zldrobit in PR https://github.com/ultralytics/yolov5/pull/1127
        
        Usage:
            $ python models/tf.py --weights yolov5s.pt
        
        Export:
            $ python export.py --weights yolov5s.pt --include saved_model pb tflite tfjs
        """
        
        import argparse
        import sys
        from copy import deepcopy
        from pathlib import Path
        
        FILE = Path(__file__).resolve()
        ROOT = FILE.parents[1]  # YOLOv5 root directory
        if str(ROOT) not in sys.path:
            sys.path.append(str(ROOT))  # add ROOT to PATH
        # ROOT = ROOT.relative_to(Path.cwd())  # relative
        
        import numpy as np
        import tensorflow as tf
        import torch
        import torch.nn as nn
        from tensorflow import keras
        
        from models.common import (
            C3,
            SPP,
            SPPF,
            Bottleneck,
            BottleneckCSP,
            C3x,
            Concat,
            Conv,
            CrossConv,
            DWConv,
            DWConvTranspose2d,
            Focus,
            autopad,
        )
        from models.experimental import MixConv2d, attempt_load
        from models.yolo import Detect, Segment
        from utils.activations import SiLU
        from utils.general import LOGGER, make_divisible, print_args
        
        
        class TFBN(keras.layers.Layer):
            # TensorFlow BatchNormalization wrapper
            def __init__(self, w=None):
                """Initializes a TensorFlow BatchNormalization layer with optional pretrained weights."""
                super().__init__()
                self.bn = keras.layers.BatchNormalization(
                    beta_initializer=keras.initializers.Constant(w.bias.numpy()),
                    gamma_initializer=keras.initializers.Constant(w.weight.numpy()),
                    moving_mean_initializer=keras.initializers.Constant(w.running_mean.numpy()),
                    moving_variance_initializer=keras.initializers.Constant(w.running_var.numpy()),
                    epsilon=w.eps,
                )
        
            def call(self, inputs):
                """Applies batch normalization to the inputs."""
                return self.bn(inputs)
        
        
        class TFPad(keras.layers.Layer):
            # Pad inputs in spatial dimensions 1 and 2
            def __init__(self, pad):
                """
                Initializes a padding layer for spatial dimensions 1 and 2 with specified padding, supporting both int and tuple
                inputs.
        
                Inputs are
                """
                super().__init__()
                if isinstance(pad, int):
                    self.pad = tf.constant([[0, 0], [pad, pad], [pad, pad], [0, 0]])
                else:  # tuple/list
                    self.pad = tf.constant([[0, 0], [pad[0], pad[0]], [pad[1], pad[1]], [0, 0]])
        
            def call(self, inputs):
                """Pads input tensor with zeros using specified padding, suitable for int and tuple pad dimensions."""
                return tf.pad(inputs, self.pad, mode="constant", constant_values=0)
        
        
        class TFConv(keras.layers.Layer):
            # Standard convolution
            def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None):
                """
                Initializes a standard convolution layer with optional batch normalization and activation; supports only
                group=1.
        
                Inputs are ch_in, ch_out, weights, kernel, stride, padding, groups.
                """
                super().__init__()
                assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument"
                # TensorFlow convolution padding is inconsistent with PyTorch (e.g. k=3 s=2 'SAME' padding)
                # see https://stackoverflow.com/questions/52975843/comparing-conv2d-with-padding-between-tensorflow-and-pytorch
                conv = keras.layers.Conv2D(
                    filters=c2,
                    kernel_size=k,
                    strides=s,
                    padding="SAME" if s == 1 else "VALID",
                    use_bias=not hasattr(w, "bn"),
                    kernel_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy()),
                    bias_initializer="zeros" if hasattr(w, "bn") else keras.initializers.Constant(w.conv.bias.numpy()),
                )
                self.conv = conv if s == 1 else keras.Sequential([TFPad(autopad(k, p)), conv])
                self.bn = TFBN(w.bn) if hasattr(w, "bn") else tf.identity
                self.act = activations(w.act) if act else tf.identity
        
            def call(self, inputs):
                """Applies convolution, batch normalization, and activation function to input tensors."""
                return self.act(self.bn(self.conv(inputs)))
        
        
        class TFDWConv(keras.layers.Layer):
            # Depthwise convolution
            def __init__(self, c1, c2, k=1, s=1, p=None, act=True, w=None):
                """
                Initializes a depthwise convolution layer with optional batch normalization and activation for TensorFlow
                models.
        
                Input are ch_in, ch_out, weights, kernel, stride, padding, groups.
                """
                super().__init__()
                assert c2 % c1 == 0, f"TFDWConv() output={c2} must be a multiple of input={c1} channels"
                conv = keras.layers.DepthwiseConv2D(
                    kernel_size=k,
                    depth_multiplier=c2 // c1,
                    strides=s,
                    padding="SAME" if s == 1 else "VALID",
                    use_bias=not hasattr(w, "bn"),
                    depthwise_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy()),
                    bias_initializer="zeros" if hasattr(w, "bn") else keras.initializers.Constant(w.conv.bias.numpy()),
                )
                self.conv = conv if s == 1 else keras.Sequential([TFPad(autopad(k, p)), conv])
                self.bn = TFBN(w.bn) if hasattr(w, "bn") else tf.identity
                self.act = activations(w.act) if act else tf.identity
        
            def call(self, inputs):
                """Applies convolution, batch normalization, and activation function to input tensors."""
                return self.act(self.bn(self.conv(inputs)))
        
        
        class TFDWConvTranspose2d(keras.layers.Layer):
            # Depthwise ConvTranspose2d
            def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0, w=None):
                """
                Initializes depthwise ConvTranspose2D layer with specific channel, kernel, stride, and padding settings.
        
                Inputs are ch_in, ch_out, weights, kernel, stride, padding, groups.
                """
                super().__init__()
                assert c1 == c2, f"TFDWConv() output={c2} must be equal to input={c1} channels"
                assert k == 4 and p1 == 1, "TFDWConv() only valid for k=4 and p1=1"
                weight, bias = w.weight.permute(2, 3, 1, 0).numpy(), w.bias.numpy()
                self.c1 = c1
                self.conv = [
                    keras.layers.Conv2DTranspose(
                        filters=1,
                        kernel_size=k,
                        strides=s,
                        padding="VALID",
                        output_padding=p2,
                        use_bias=True,
                        kernel_initializer=keras.initializers.Constant(weight[..., i : i + 1]),
                        bias_initializer=keras.initializers.Constant(bias[i]),
                    )
                    for i in range(c1)
                ]
        
            def call(self, inputs):
                """Processes input through parallel convolutions and concatenates results, trimming border pixels."""
                return tf.concat([m(x) for m, x in zip(self.conv, tf.split(inputs, self.c1, 3))], 3)[:, 1:-1, 1:-1]
        
        
        class TFFocus(keras.layers.Layer):
            # Focus wh information into c-space
            def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None):
                """
                Initializes TFFocus layer to focus width and height information into channel space with custom convolution
                parameters.
        
                Inputs are ch_in, ch_out, kernel, stride, padding, groups.
                """
                super().__init__()
                self.conv = TFConv(c1 * 4, c2, k, s, p, g, act, w.conv)
        
            def call(self, inputs):
                """
                Performs pixel shuffling and convolution on input tensor, downsampling by 2 and expanding channels by 4.
        
                Example x(b,w,h,c) -> y(b,w/2,h/2,4c).
                """
                inputs = [inputs[:, ::2, ::2, :], inputs[:, 1::2, ::2, :], inputs[:, ::2, 1::2, :], inputs[:, 1::2, 1::2, :]]
                return self.conv(tf.concat(inputs, 3))
        
        
        class TFBottleneck(keras.layers.Layer):
            # Standard bottleneck
            def __init__(self, c1, c2, shortcut=True, g=1, e=0.5, w=None):
                """
                Initializes a standard bottleneck layer for TensorFlow models, expanding and contracting channels with optional
                shortcut.
        
                Arguments are ch_in, ch_out, shortcut, groups, expansion.
                """
                super().__init__()
                c_ = int(c2 * e)  # hidden channels
                self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
                self.cv2 = TFConv(c_, c2, 3, 1, g=g, w=w.cv2)
                self.add = shortcut and c1 == c2
        
            def call(self, inputs):
                """Performs forward pass; if shortcut is True & input/output channels match, adds input to the convolution
                result.
                """
                return inputs + self.cv2(self.cv1(inputs)) if self.add else self.cv2(self.cv1(inputs))
        
        
        class TFCrossConv(keras.layers.Layer):
            # Cross Convolution
            def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False, w=None):
                """Initializes cross convolution layer with optional expansion, grouping, and shortcut addition capabilities."""
                super().__init__()
                c_ = int(c2 * e)  # hidden channels
                self.cv1 = TFConv(c1, c_, (1, k), (1, s), w=w.cv1)
                self.cv2 = TFConv(c_, c2, (k, 1), (s, 1), g=g, w=w.cv2)
                self.add = shortcut and c1 == c2
        
            def call(self, inputs):
                """Passes input through two convolutions optionally adding the input if channel dimensions match."""
                return inputs + self.cv2(self.cv1(inputs)) if self.add else self.cv2(self.cv1(inputs))
        
        
        class TFConv2d(keras.layers.Layer):
            # Substitution for PyTorch nn.Conv2D
            def __init__(self, c1, c2, k, s=1, g=1, bias=True, w=None):
                """Initializes a TensorFlow 2D convolution layer, mimicking PyTorch's nn.Conv2D functionality for given filter
                sizes and stride.
                """
                super().__init__()
                assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument"
                self.conv = keras.layers.Conv2D(
                    filters=c2,
                    kernel_size=k,
                    strides=s,
                    padding="VALID",
                    use_bias=bias,
                    kernel_initializer=keras.initializers.Constant(w.weight.permute(2, 3, 1, 0).numpy()),
                    bias_initializer=keras.initializers.Constant(w.bias.numpy()) if bias else None,
                )
        
            def call(self, inputs):
                """Applies a convolution operation to the inputs and returns the result."""
                return self.conv(inputs)
        
        
        class TFBottleneckCSP(keras.layers.Layer):
            # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
            def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None):
                """
                Initializes CSP bottleneck layer with specified channel sizes, count, shortcut option, groups, and expansion
                ratio.
        
                Inputs are ch_in, ch_out, number, shortcut, groups, expansion.
                """
                super().__init__()
                c_ = int(c2 * e)  # hidden channels
                self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
                self.cv2 = TFConv2d(c1, c_, 1, 1, bias=False, w=w.cv2)
                self.cv3 = TFConv2d(c_, c_, 1, 1, bias=False, w=w.cv3)
                self.cv4 = TFConv(2 * c_, c2, 1, 1, w=w.cv4)
                self.bn = TFBN(w.bn)
                self.act = lambda x: keras.activations.swish(x)
                self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)])
        
            def call(self, inputs):
                """Processes input through the model layers, concatenates, normalizes, activates, and reduces the output
                dimensions.
                """
                y1 = self.cv3(self.m(self.cv1(inputs)))
                y2 = self.cv2(inputs)
                return self.cv4(self.act(self.bn(tf.concat((y1, y2), axis=3))))
        
        
        class TFC3(keras.layers.Layer):
            # CSP Bottleneck with 3 convolutions
            def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None):
                """
                Initializes CSP Bottleneck with 3 convolutions, supporting optional shortcuts and group convolutions.
        
                Inputs are ch_in, ch_out, number, shortcut, groups, expansion.
                """
                super().__init__()
                c_ = int(c2 * e)  # hidden channels
                self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
                self.cv2 = TFConv(c1, c_, 1, 1, w=w.cv2)
                self.cv3 = TFConv(2 * c_, c2, 1, 1, w=w.cv3)
                self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)])
        
            def call(self, inputs):
                """
                Processes input through a sequence of transformations for object detection (YOLOv5).
        
                See https://github.com/ultralytics/yolov5.
                """
                return self.cv3(tf.concat((self.m(self.cv1(inputs)), self.cv2(inputs)), axis=3))
        
        
        class TFC3x(keras.layers.Layer):
            # 3 module with cross-convolutions
            def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None):
                """
                Initializes layer with cross-convolutions for enhanced feature extraction in object detection models.
        
                Inputs are ch_in, ch_out, number, shortcut, groups, expansion.
                """
                super().__init__()
                c_ = int(c2 * e)  # hidden channels
                self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
                self.cv2 = TFConv(c1, c_, 1, 1, w=w.cv2)
                self.cv3 = TFConv(2 * c_, c2, 1, 1, w=w.cv3)
                self.m = keras.Sequential(
                    [TFCrossConv(c_, c_, k=3, s=1, g=g, e=1.0, shortcut=shortcut, w=w.m[j]) for j in range(n)]
                )
        
            def call(self, inputs):
                """Processes input through cascaded convolutions and merges features, returning the final tensor output."""
                return self.cv3(tf.concat((self.m(self.cv1(inputs)), self.cv2(inputs)), axis=3))
        
        
        class TFSPP(keras.layers.Layer):
            # Spatial pyramid pooling layer used in YOLOv3-SPP
            def __init__(self, c1, c2, k=(5, 9, 13), w=None):
                """Initializes a YOLOv3-SPP layer with specific input/output channels and kernel sizes for pooling."""
                super().__init__()
                c_ = c1 // 2  # hidden channels
                self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
                self.cv2 = TFConv(c_ * (len(k) + 1), c2, 1, 1, w=w.cv2)
                self.m = [keras.layers.MaxPool2D(pool_size=x, strides=1, padding="SAME") for x in k]
        
            def call(self, inputs):
                """Processes input through two TFConv layers and concatenates with max-pooled outputs at intermediate stage."""
                x = self.cv1(inputs)
                return self.cv2(tf.concat([x] + [m(x) for m in self.m], 3))
        
        
        class TFSPPF(keras.layers.Layer):
            # Spatial pyramid pooling-Fast layer
            def __init__(self, c1, c2, k=5, w=None):
                """Initializes a fast spatial pyramid pooling layer with customizable in/out channels, kernel size, and
                weights.
                """
                super().__init__()
                c_ = c1 // 2  # hidden channels
                self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
                self.cv2 = TFConv(c_ * 4, c2, 1, 1, w=w.cv2)
                self.m = keras.layers.MaxPool2D(pool_size=k, strides=1, padding="SAME")
        
            def call(self, inputs):
                """Executes the model's forward pass, concatenating input features with three max-pooled versions before final
                convolution.
                """
                x = self.cv1(inputs)
                y1 = self.m(x)
                y2 = self.m(y1)
                return self.cv2(tf.concat([x, y1, y2, self.m(y2)], 3))
        
        
        class TFDetect(keras.layers.Layer):
            # TF YOLOv5 Detect layer
            def __init__(self, nc=80, anchors=(), ch=(), imgsz=(640, 640), w=None):
                """Initializes YOLOv5 detection layer for TensorFlow with configurable classes, anchors, channels, and image
                size.
                """
                super().__init__()
                self.stride = tf.convert_to_tensor(w.stride.numpy(), dtype=tf.float32)
                self.nc = nc  # number of classes
                self.no = nc + 5  # number of outputs per anchor
                self.nl = len(anchors)  # number of detection layers
                self.na = len(anchors[0]) // 2  # number of anchors
                self.grid = [tf.zeros(1)] * self.nl  # init grid
                self.anchors = tf.convert_to_tensor(w.anchors.numpy(), dtype=tf.float32)
                self.anchor_grid = tf.reshape(self.anchors * tf.reshape(self.stride, [self.nl, 1, 1]), [self.nl, 1, -1, 1, 2])
                self.m = [TFConv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in enumerate(ch)]
                self.training = False  # set to False after building model
                self.imgsz = imgsz
                for i in range(self.nl):
                    ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1] // self.stride[i]
                    self.grid[i] = self._make_grid(nx, ny)
        
            def call(self, inputs):
                """Performs forward pass through the model layers to predict object bounding boxes and classifications."""
                z = []  # inference output
                x = []
                for i in range(self.nl):
                    x.append(self.m[i](inputs[i]))
                    # x(bs,20,20,255) to x(bs,3,20,20,85)
                    ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1] // self.stride[i]
                    x[i] = tf.reshape(x[i], [-1, ny * nx, self.na, self.no])
        
                    if not self.training:  # inference
                        y = x[i]
                        grid = tf.transpose(self.grid[i], [0, 2, 1, 3]) - 0.5
                        anchor_grid = tf.transpose(self.anchor_grid[i], [0, 2, 1, 3]) * 4
                        xy = (tf.sigmoid(y[..., 0:2]) * 2 + grid) * self.stride[i]  # xy
                        wh = tf.sigmoid(y[..., 2:4]) ** 2 * anchor_grid
                        # Normalize xywh to 0-1 to reduce calibration error
                        xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
                        wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
                        y = tf.concat([xy, wh, tf.sigmoid(y[..., 4 : 5 + self.nc]), y[..., 5 + self.nc :]], -1)
                        z.append(tf.reshape(y, [-1, self.na * ny * nx, self.no]))
        
                return tf.transpose(x, [0, 2, 1, 3]) if self.training else (tf.concat(z, 1),)
        
            @staticmethod
            def _make_grid(nx=20, ny=20):
                """Generates a 2D grid of coordinates in (x, y) format with shape [1, 1, ny*nx, 2]."""
                # return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
                xv, yv = tf.meshgrid(tf.range(nx), tf.range(ny))
                return tf.cast(tf.reshape(tf.stack([xv, yv], 2), [1, 1, ny * nx, 2]), dtype=tf.float32)
        
        
        class TFSegment(TFDetect):
            # YOLOv5 Segment head for segmentation models
            def __init__(self, nc=80, anchors=(), nm=32, npr=256, ch=(), imgsz=(640, 640), w=None):
                """Initializes YOLOv5 Segment head with specified channel depths, anchors, and input size for segmentation
                models.
                """
                super().__init__(nc, anchors, ch, imgsz, w)
                self.nm = nm  # number of masks
                self.npr = npr  # number of protos
                self.no = 5 + nc + self.nm  # number of outputs per anchor
                self.m = [TFConv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in enumerate(ch)]  # output conv
                self.proto = TFProto(ch[0], self.npr, self.nm, w=w.proto)  # protos
                self.detect = TFDetect.call
        
            def call(self, x):
                """Applies detection and proto layers on input, returning detections and optionally protos if training."""
                p = self.proto(x[0])
                # p = TFUpsample(None, scale_factor=4, mode='nearest')(self.proto(x[0]))  # (optional) full-size protos
                p = tf.transpose(p, [0, 3, 1, 2])  # from shape(1,160,160,32) to shape(1,32,160,160)
                x = self.detect(self, x)
                return (x, p) if self.training else (x[0], p)
        
        
        class TFProto(keras.layers.Layer):
            def __init__(self, c1, c_=256, c2=32, w=None):
                """Initializes TFProto layer with convolutional and upsampling layers for feature extraction and
                transformation.
                """
                super().__init__()
                self.cv1 = TFConv(c1, c_, k=3, w=w.cv1)
                self.upsample = TFUpsample(None, scale_factor=2, mode="nearest")
                self.cv2 = TFConv(c_, c_, k=3, w=w.cv2)
                self.cv3 = TFConv(c_, c2, w=w.cv3)
        
            def call(self, inputs):
                """Performs forward pass through the model, applying convolutions and upscaling on input tensor."""
                return self.cv3(self.cv2(self.upsample(self.cv1(inputs))))
        
        
        class TFUpsample(keras.layers.Layer):
            # TF version of torch.nn.Upsample()
            def __init__(self, size, scale_factor, mode, w=None):
                """
                Initializes a TensorFlow upsampling layer with specified size, scale_factor, and mode, ensuring scale_factor is
                even.
        
                Warning: all arguments needed including 'w'
                """
                super().__init__()
                assert scale_factor % 2 == 0, "scale_factor must be multiple of 2"
                self.upsample = lambda x: tf.image.resize(x, (x.shape[1] * scale_factor, x.shape[2] * scale_factor), mode)
                # self.upsample = keras.layers.UpSampling2D(size=scale_factor, interpolation=mode)
                # with default arguments: align_corners=False, half_pixel_centers=False
                # self.upsample = lambda x: tf.raw_ops.ResizeNearestNeighbor(images=x,
                #                                                            size=(x.shape[1] * 2, x.shape[2] * 2))
        
            def call(self, inputs):
                """Applies upsample operation to inputs using nearest neighbor interpolation."""
                return self.upsample(inputs)
        
        
        class TFConcat(keras.layers.Layer):
            # TF version of torch.concat()
            def __init__(self, dimension=1, w=None):
                """Initializes a TensorFlow layer for NCHW to NHWC concatenation, requiring dimension=1."""
                super().__init__()
                assert dimension == 1, "convert only NCHW to NHWC concat"
                self.d = 3
        
            def call(self, inputs):
                """Concatenates a list of tensors along the last dimension, used for NCHW to NHWC conversion."""
                return tf.concat(inputs, self.d)
        
        
        def parse_model(d, ch, model, imgsz):
            """Parses a model definition dict `d` to create YOLOv5 model layers, including dynamic channel adjustments."""
            LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10}  {'module':<40}{'arguments':<30}")
            anchors, nc, gd, gw, ch_mul = (
                d["anchors"],
                d["nc"],
                d["depth_multiple"],
                d["width_multiple"],
                d.get("channel_multiple"),
            )
            na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # number of anchors
            no = na * (nc + 5)  # number of outputs = anchors * (classes + 5)
            if not ch_mul:
                ch_mul = 8
        
            layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
            for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):  # from, number, module, args
                m_str = m
                m = eval(m) if isinstance(m, str) else m  # eval strings
                for j, a in enumerate(args):
                    try:
                        args[j] = eval(a) if isinstance(a, str) else a  # eval strings
                    except NameError:
                        pass
        
                n = max(round(n * gd), 1) if n > 1 else n  # depth gain
                if m in [
                    nn.Conv2d,
                    Conv,
                    DWConv,
                    DWConvTranspose2d,
                    Bottleneck,
                    SPP,
                    SPPF,
                    MixConv2d,
                    Focus,
                    CrossConv,
                    BottleneckCSP,
                    C3,
                    C3x,
                ]:
                    c1, c2 = ch[f], args[0]
                    c2 = make_divisible(c2 * gw, ch_mul) if c2 != no else c2
        
                    args = [c1, c2, *args[1:]]
                    if m in [BottleneckCSP, C3, C3x]:
                        args.insert(2, n)
                        n = 1
                elif m is nn.BatchNorm2d:
                    args = [ch[f]]
                elif m is Concat:
                    c2 = sum(ch[-1 if x == -1 else x + 1] for x in f)
                elif m in [Detect, Segment]:
                    args.append([ch[x + 1] for x in f])
                    if isinstance(args[1], int):  # number of anchors
                        args[1] = [list(range(args[1] * 2))] * len(f)
                    if m is Segment:
                        args[3] = make_divisible(args[3] * gw, ch_mul)
                    args.append(imgsz)
                else:
                    c2 = ch[f]
        
                tf_m = eval("TF" + m_str.replace("nn.", ""))
                m_ = (
                    keras.Sequential([tf_m(*args, w=model.model[i][j]) for j in range(n)])
                    if n > 1
                    else tf_m(*args, w=model.model[i])
                )  # module
        
                torch_m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
                t = str(m)[8:-2].replace("__main__.", "")  # module type
                np = sum(x.numel() for x in torch_m_.parameters())  # number params
                m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
                LOGGER.info(f"{i:>3}{str(f):>18}{str(n):>3}{np:>10}  {t:<40}{str(args):<30}")  # print
                save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
                layers.append(m_)
                ch.append(c2)
            return keras.Sequential(layers), sorted(save)
        
        
        class TFModel:
            # TF YOLOv5 model
            def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, model=None, imgsz=(640, 640)):
                """Initializes TF YOLOv5 model with specified configuration, channels, classes, model instance, and input
                size.
                """
                super().__init__()
                if isinstance(cfg, dict):
                    self.yaml = cfg  # model dict
                else:  # is *.yaml
                    import yaml  # for torch hub
        
                    self.yaml_file = Path(cfg).name
                    with open(cfg) as f:
                        self.yaml = yaml.load(f, Loader=yaml.FullLoader)  # model dict
        
                # Define model
                if nc and nc != self.yaml["nc"]:
                    LOGGER.info(f"Overriding {cfg} nc={self.yaml['nc']} with nc={nc}")
                    self.yaml["nc"] = nc  # override yaml value
                self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch], model=model, imgsz=imgsz)
        
            def predict(
                self,
                inputs,
                tf_nms=False,
                agnostic_nms=False,
                topk_per_class=100,
                topk_all=100,
                iou_thres=0.45,
                conf_thres=0.25,
            ):
                y = []  # outputs
                x = inputs
                for m in self.model.layers:
                    if m.f != -1:  # if not from previous layer
                        x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers
        
                    x = m(x)  # run
                    y.append(x if m.i in self.savelist else None)  # save output
        
                # Add TensorFlow NMS
                if tf_nms:
                    boxes = self._xywh2xyxy(x[0][..., :4])
                    probs = x[0][:, :, 4:5]
                    classes = x[0][:, :, 5:]
                    scores = probs * classes
                    if agnostic_nms:
                        nms = AgnosticNMS()((boxes, classes, scores), topk_all, iou_thres, conf_thres)
                    else:
                        boxes = tf.expand_dims(boxes, 2)
                        nms = tf.image.combined_non_max_suppression(
                            boxes, scores, topk_per_class, topk_all, iou_thres, conf_thres, clip_boxes=False
                        )
                    return (nms,)
                return x  # output [1,6300,85] = [xywh, conf, class0, class1, ...]
                # x = x[0]  # [x(1,6300,85), ...] to x(6300,85)
                # xywh = x[..., :4]  # x(6300,4) boxes
                # conf = x[..., 4:5]  # x(6300,1) confidences
                # cls = tf.reshape(tf.cast(tf.argmax(x[..., 5:], axis=1), tf.float32), (-1, 1))  # x(6300,1)  classes
                # return tf.concat([conf, cls, xywh], 1)
        
            @staticmethod
            def _xywh2xyxy(xywh):
                """Converts bounding box format from [x, y, w, h] to [x1, y1, x2, y2], where xy1=top-left and xy2=bottom-
                right.
                """
                x, y, w, h = tf.split(xywh, num_or_size_splits=4, axis=-1)
                return tf.concat([x - w / 2, y - h / 2, x + w / 2, y + h / 2], axis=-1)
        
        
        class AgnosticNMS(keras.layers.Layer):
            # TF Agnostic NMS
            def call(self, input, topk_all, iou_thres, conf_thres):
                """Performs agnostic NMS on input tensors using given thresholds and top-K selection."""
                return tf.map_fn(
                    lambda x: self._nms(x, topk_all, iou_thres, conf_thres),
                    input,
                    fn_output_signature=(tf.float32, tf.float32, tf.float32, tf.int32),
                    name="agnostic_nms",
                )
        
            @staticmethod
            def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25):
                """Performs agnostic non-maximum suppression (NMS) on detected objects, filtering based on IoU and confidence
                thresholds.
                """
                boxes, classes, scores = x
                class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32)
                scores_inp = tf.reduce_max(scores, -1)
                selected_inds = tf.image.non_max_suppression(
                    boxes, scores_inp, max_output_size=topk_all, iou_threshold=iou_thres, score_threshold=conf_thres
                )
                selected_boxes = tf.gather(boxes, selected_inds)
                padded_boxes = tf.pad(
                    selected_boxes,
                    paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]],
                    mode="CONSTANT",
                    constant_values=0.0,
                )
                selected_scores = tf.gather(scores_inp, selected_inds)
                padded_scores = tf.pad(
                    selected_scores,
                    paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
                    mode="CONSTANT",
                    constant_values=-1.0,
                )
                selected_classes = tf.gather(class_inds, selected_inds)
                padded_classes = tf.pad(
                    selected_classes,
                    paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
                    mode="CONSTANT",
                    constant_values=-1.0,
                )
                valid_detections = tf.shape(selected_inds)[0]
                return padded_boxes, padded_scores, padded_classes, valid_detections
        
        
        def activations(act=nn.SiLU):
            """Converts PyTorch activations to TensorFlow equivalents, supporting LeakyReLU, Hardswish, and SiLU/Swish."""
            if isinstance(act, nn.LeakyReLU):
                return lambda x: keras.activations.relu(x, alpha=0.1)
            elif isinstance(act, nn.Hardswish):
                return lambda x: x * tf.nn.relu6(x + 3) * 0.166666667
            elif isinstance(act, (nn.SiLU, SiLU)):
                return lambda x: keras.activations.swish(x)
            else:
                raise Exception(f"no matching TensorFlow activation found for PyTorch activation {act}")
        
        
        def representative_dataset_gen(dataset, ncalib=100):
            """Generates a representative dataset for calibration by yielding transformed numpy arrays from the input
            dataset.
            """
            for n, (path, img, im0s, vid_cap, string) in enumerate(dataset):
                im = np.transpose(img, [1, 2, 0])
                im = np.expand_dims(im, axis=0).astype(np.float32)
                im /= 255
                yield [im]
                if n >= ncalib:
                    break
        
        
        def run(
            weights=ROOT / "yolov5s.pt",  # weights path
            imgsz=(640, 640),  # inference size h,w
            batch_size=1,  # batch size
            dynamic=False,  # dynamic batch size
        ):
            # PyTorch model
            im = torch.zeros((batch_size, 3, *imgsz))  # BCHW image
            model = attempt_load(weights, device=torch.device("cpu"), inplace=True, fuse=False)
            _ = model(im)  # inference
            model.info()
        
            # TensorFlow model
            im = tf.zeros((batch_size, *imgsz, 3))  # BHWC image
            tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
            _ = tf_model.predict(im)  # inference
        
            # Keras model
            im = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size)
            keras_model = keras.Model(inputs=im, outputs=tf_model.predict(im))
            keras_model.summary()
        
            LOGGER.info("PyTorch, TensorFlow and Keras models successfully verified.\nUse export.py for TF model export.")
        
        
        def parse_opt():
            """Parses and returns command-line options for model inference, including weights path, image size, batch size, and
            dynamic batching.
            """
            parser = argparse.ArgumentParser()
            parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="weights path")
            parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640], help="inference size h,w")
            parser.add_argument("--batch-size", type=int, default=1, help="batch size")
            parser.add_argument("--dynamic", action="store_true", help="dynamic batch size")
            opt = parser.parse_args()
            opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
            print_args(vars(opt))
            return opt
        
        
        def main(opt):
            """Executes the YOLOv5 model run function with parsed command line options."""
            run(**vars(opt))
        
        
        if __name__ == "__main__":
            opt = parse_opt()
            main(opt)
        
      • 📄 yolo.py 20.1KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """
        YOLO-specific modules.
        
        Usage:
            $ python models/yolo.py --cfg yolov5s.yaml
        """
        
        import argparse
        import contextlib
        import math
        import os
        import platform
        import sys
        from copy import deepcopy
        from pathlib import Path
        
        import torch
        import torch.nn as nn
        
        FILE = Path(__file__).resolve()
        ROOT = FILE.parents[1]  # YOLOv5 root directory
        if str(ROOT) not in sys.path:
            sys.path.append(str(ROOT))  # add ROOT to PATH
        if platform.system() != "Windows":
            ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
        
        from models.common import (
            C3,
            C3SPP,
            C3TR,
            SPP,
            SPPF,
            Bottleneck,
            BottleneckCSP,
            C3Ghost,
            C3x,
            Classify,
            Concat,
            Contract,
            Conv,
            CrossConv,
            DetectMultiBackend,
            DWConv,
            DWConvTranspose2d,
            Expand,
            Focus,
            GhostBottleneck,
            GhostConv,
            Proto,
        )
        from models.experimental import MixConv2d
        from utils.autoanchor import check_anchor_order
        from utils.general import LOGGER, check_version, check_yaml, colorstr, make_divisible, print_args
        from utils.plots import feature_visualization
        from utils.torch_utils import (
            fuse_conv_and_bn,
            initialize_weights,
            model_info,
            profile,
            scale_img,
            select_device,
            time_sync,
        )
        
        try:
            import thop  # for FLOPs computation
        except ImportError:
            thop = None
        
        
        class Detect(nn.Module):
            # YOLOv5 Detect head for detection models
            stride = None  # strides computed during build
            dynamic = False  # force grid reconstruction
            export = False  # export mode
        
            def __init__(self, nc=80, anchors=(), ch=(), inplace=True):
                """Initializes YOLOv5 detection layer with specified classes, anchors, channels, and inplace operations."""
                super().__init__()
                self.nc = nc  # number of classes
                self.no = nc + 5  # number of outputs per anchor
                self.nl = len(anchors)  # number of detection layers
                self.na = len(anchors[0]) // 2  # number of anchors
                self.grid = [torch.empty(0) for _ in range(self.nl)]  # init grid
                self.anchor_grid = [torch.empty(0) for _ in range(self.nl)]  # init anchor grid
                self.register_buffer("anchors", torch.tensor(anchors).float().view(self.nl, -1, 2))  # shape(nl,na,2)
                self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # output conv
                self.inplace = inplace  # use inplace ops (e.g. slice assignment)
        
            def forward(self, x):
                """Processes input through YOLOv5 layers, altering shape for detection: `x(bs, 3, ny, nx, 85)`."""
                z = []  # inference output
                for i in range(self.nl):
                    x[i] = self.m[i](x[i])  # conv
                    bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
                    x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
        
                    if not self.training:  # inference
                        if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
                            self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
        
                        if isinstance(self, Segment):  # (boxes + masks)
                            xy, wh, conf, mask = x[i].split((2, 2, self.nc + 1, self.no - self.nc - 5), 4)
                            xy = (xy.sigmoid() * 2 + self.grid[i]) * self.stride[i]  # xy
                            wh = (wh.sigmoid() * 2) ** 2 * self.anchor_grid[i]  # wh
                            y = torch.cat((xy, wh, conf.sigmoid(), mask), 4)
                        else:  # Detect (boxes only)
                            xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4)
                            xy = (xy * 2 + self.grid[i]) * self.stride[i]  # xy
                            wh = (wh * 2) ** 2 * self.anchor_grid[i]  # wh
                            y = torch.cat((xy, wh, conf), 4)
                        z.append(y.view(bs, self.na * nx * ny, self.no))
        
                return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
        
            def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, "1.10.0")):
                """Generates a mesh grid for anchor boxes with optional compatibility for torch versions < 1.10."""
                d = self.anchors[i].device
                t = self.anchors[i].dtype
                shape = 1, self.na, ny, nx, 2  # grid shape
                y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
                yv, xv = torch.meshgrid(y, x, indexing="ij") if torch_1_10 else torch.meshgrid(y, x)  # torch>=0.7 compatibility
                grid = torch.stack((xv, yv), 2).expand(shape) - 0.5  # add grid offset, i.e. y = 2.0 * x - 0.5
                anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
                return grid, anchor_grid
        
        
        class Segment(Detect):
            # YOLOv5 Segment head for segmentation models
            def __init__(self, nc=80, anchors=(), nm=32, npr=256, ch=(), inplace=True):
                """Initializes YOLOv5 Segment head with options for mask count, protos, and channel adjustments."""
                super().__init__(nc, anchors, ch, inplace)
                self.nm = nm  # number of masks
                self.npr = npr  # number of protos
                self.no = 5 + nc + self.nm  # number of outputs per anchor
                self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # output conv
                self.proto = Proto(ch[0], self.npr, self.nm)  # protos
                self.detect = Detect.forward
        
            def forward(self, x):
                """Processes input through the network, returning detections and prototypes; adjusts output based on
                training/export mode.
                """
                p = self.proto(x[0])
                x = self.detect(self, x)
                return (x, p) if self.training else (x[0], p) if self.export else (x[0], p, x[1])
        
        
        class BaseModel(nn.Module):
            """YOLOv5 base model."""
        
            def forward(self, x, profile=False, visualize=False):
                """Executes a single-scale inference or training pass on the YOLOv5 base model, with options for profiling and
                visualization.
                """
                return self._forward_once(x, profile, visualize)  # single-scale inference, train
        
            def _forward_once(self, x, profile=False, visualize=False):
                """Performs a forward pass on the YOLOv5 model, enabling profiling and feature visualization options."""
                y, dt = [], []  # outputs
                for m in self.model:
                    if m.f != -1:  # if not from previous layer
                        x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers
                    if profile:
                        self._profile_one_layer(m, x, dt)
                    x = m(x)  # run
                    y.append(x if m.i in self.save else None)  # save output
                    if visualize:
                        feature_visualization(x, m.type, m.i, save_dir=visualize)
                return x
        
            def _profile_one_layer(self, m, x, dt):
                """Profiles a single layer's performance by computing GFLOPs, execution time, and parameters."""
                c = m == self.model[-1]  # is final layer, copy input as inplace fix
                o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1e9 * 2 if thop else 0  # FLOPs
                t = time_sync()
                for _ in range(10):
                    m(x.copy() if c else x)
                dt.append((time_sync() - t) * 100)
                if m == self.model[0]:
                    LOGGER.info(f"{'time (ms)':>10s} {'GFLOPs':>10s} {'params':>10s}  module")
                LOGGER.info(f"{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f}  {m.type}")
                if c:
                    LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s}  Total")
        
            def fuse(self):
                """Fuses Conv2d() and BatchNorm2d() layers in the model to improve inference speed."""
                LOGGER.info("Fusing layers... ")
                for m in self.model.modules():
                    if isinstance(m, (Conv, DWConv)) and hasattr(m, "bn"):
                        m.conv = fuse_conv_and_bn(m.conv, m.bn)  # update conv
                        delattr(m, "bn")  # remove batchnorm
                        m.forward = m.forward_fuse  # update forward
                self.info()
                return self
        
            def info(self, verbose=False, img_size=640):
                """Prints model information given verbosity and image size, e.g., `info(verbose=True, img_size=640)`."""
                model_info(self, verbose, img_size)
        
            def _apply(self, fn):
                """Applies transformations like to(), cpu(), cuda(), half() to model tensors excluding parameters or registered
                buffers.
                """
                self = super()._apply(fn)
                m = self.model[-1]  # Detect()
                if isinstance(m, (Detect, Segment)):
                    m.stride = fn(m.stride)
                    m.grid = list(map(fn, m.grid))
                    if isinstance(m.anchor_grid, list):
                        m.anchor_grid = list(map(fn, m.anchor_grid))
                return self
        
        
        class DetectionModel(BaseModel):
            # YOLOv5 detection model
            def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, anchors=None):
                """Initializes YOLOv5 model with configuration file, input channels, number of classes, and custom anchors."""
                super().__init__()
                if isinstance(cfg, dict):
                    self.yaml = cfg  # model dict
                else:  # is *.yaml
                    import yaml  # for torch hub
        
                    self.yaml_file = Path(cfg).name
                    with open(cfg, encoding="ascii", errors="ignore") as f:
                        self.yaml = yaml.safe_load(f)  # model dict
        
                # Define model
                ch = self.yaml["ch"] = self.yaml.get("ch", ch)  # input channels
                if nc and nc != self.yaml["nc"]:
                    LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
                    self.yaml["nc"] = nc  # override yaml value
                if anchors:
                    LOGGER.info(f"Overriding model.yaml anchors with anchors={anchors}")
                    self.yaml["anchors"] = round(anchors)  # override yaml value
                self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])  # model, savelist
                self.names = [str(i) for i in range(self.yaml["nc"])]  # default names
                self.inplace = self.yaml.get("inplace", True)
        
                # Build strides, anchors
                m = self.model[-1]  # Detect()
                if isinstance(m, (Detect, Segment)):
                    s = 256  # 2x min stride
                    m.inplace = self.inplace
                    forward = lambda x: self.forward(x)[0] if isinstance(m, Segment) else self.forward(x)
                    m.stride = torch.tensor([s / x.shape[-2] for x in forward(torch.zeros(1, ch, s, s))])  # forward
                    check_anchor_order(m)
                    m.anchors /= m.stride.view(-1, 1, 1)
                    self.stride = m.stride
                    self._initialize_biases()  # only run once
        
                # Init weights, biases
                initialize_weights(self)
                self.info()
                LOGGER.info("")
        
            def forward(self, x, augment=False, profile=False, visualize=False):
                """Performs single-scale or augmented inference and may include profiling or visualization."""
                if augment:
                    return self._forward_augment(x)  # augmented inference, None
                return self._forward_once(x, profile, visualize)  # single-scale inference, train
        
            def _forward_augment(self, x):
                """Performs augmented inference across different scales and flips, returning combined detections."""
                img_size = x.shape[-2:]  # height, width
                s = [1, 0.83, 0.67]  # scales
                f = [None, 3, None]  # flips (2-ud, 3-lr)
                y = []  # outputs
                for si, fi in zip(s, f):
                    xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
                    yi = self._forward_once(xi)[0]  # forward
                    # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1])  # save
                    yi = self._descale_pred(yi, fi, si, img_size)
                    y.append(yi)
                y = self._clip_augmented(y)  # clip augmented tails
                return torch.cat(y, 1), None  # augmented inference, train
        
            def _descale_pred(self, p, flips, scale, img_size):
                """De-scales predictions from augmented inference, adjusting for flips and image size."""
                if self.inplace:
                    p[..., :4] /= scale  # de-scale
                    if flips == 2:
                        p[..., 1] = img_size[0] - p[..., 1]  # de-flip ud
                    elif flips == 3:
                        p[..., 0] = img_size[1] - p[..., 0]  # de-flip lr
                else:
                    x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale  # de-scale
                    if flips == 2:
                        y = img_size[0] - y  # de-flip ud
                    elif flips == 3:
                        x = img_size[1] - x  # de-flip lr
                    p = torch.cat((x, y, wh, p[..., 4:]), -1)
                return p
        
            def _clip_augmented(self, y):
                """Clips augmented inference tails for YOLOv5 models, affecting first and last tensors based on grid points and
                layer counts.
                """
                nl = self.model[-1].nl  # number of detection layers (P3-P5)
                g = sum(4**x for x in range(nl))  # grid points
                e = 1  # exclude layer count
                i = (y[0].shape[1] // g) * sum(4**x for x in range(e))  # indices
                y[0] = y[0][:, :-i]  # large
                i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e))  # indices
                y[-1] = y[-1][:, i:]  # small
                return y
        
            def _initialize_biases(self, cf=None):
                """
                Initializes biases for YOLOv5's Detect() module, optionally using class frequencies (cf).
        
                For details see https://arxiv.org/abs/1708.02002 section 3.3.
                """
                # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
                m = self.model[-1]  # Detect() module
                for mi, s in zip(m.m, m.stride):  # from
                    b = mi.bias.view(m.na, -1)  # conv.bias(255) to (3,85)
                    b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
                    b.data[:, 5 : 5 + m.nc] += (
                        math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
                    )  # cls
                    mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
        
        
        Model = DetectionModel  # retain YOLOv5 'Model' class for backwards compatibility
        
        
        class SegmentationModel(DetectionModel):
            # YOLOv5 segmentation model
            def __init__(self, cfg="yolov5s-seg.yaml", ch=3, nc=None, anchors=None):
                """Initializes a YOLOv5 segmentation model with configurable params: cfg (str) for configuration, ch (int) for channels, nc (int) for num classes, anchors (list)."""
                super().__init__(cfg, ch, nc, anchors)
        
        
        class ClassificationModel(BaseModel):
            # YOLOv5 classification model
            def __init__(self, cfg=None, model=None, nc=1000, cutoff=10):
                """Initializes YOLOv5 model with config file `cfg`, input channels `ch`, number of classes `nc`, and `cuttoff`
                index.
                """
                super().__init__()
                self._from_detection_model(model, nc, cutoff) if model is not None else self._from_yaml(cfg)
        
            def _from_detection_model(self, model, nc=1000, cutoff=10):
                """Creates a classification model from a YOLOv5 detection model, slicing at `cutoff` and adding a classification
                layer.
                """
                if isinstance(model, DetectMultiBackend):
                    model = model.model  # unwrap DetectMultiBackend
                model.model = model.model[:cutoff]  # backbone
                m = model.model[-1]  # last layer
                ch = m.conv.in_channels if hasattr(m, "conv") else m.cv1.conv.in_channels  # ch into module
                c = Classify(ch, nc)  # Classify()
                c.i, c.f, c.type = m.i, m.f, "models.common.Classify"  # index, from, type
                model.model[-1] = c  # replace
                self.model = model.model
                self.stride = model.stride
                self.save = []
                self.nc = nc
        
            def _from_yaml(self, cfg):
                """Creates a YOLOv5 classification model from a specified *.yaml configuration file."""
                self.model = None
        
        
        def parse_model(d, ch):
            """Parses a YOLOv5 model from a dict `d`, configuring layers based on input channels `ch` and model architecture."""
            LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10}  {'module':<40}{'arguments':<30}")
            anchors, nc, gd, gw, act, ch_mul = (
                d["anchors"],
                d["nc"],
                d["depth_multiple"],
                d["width_multiple"],
                d.get("activation"),
                d.get("channel_multiple"),
            )
            if act:
                Conv.default_act = eval(act)  # redefine default activation, i.e. Conv.default_act = nn.SiLU()
                LOGGER.info(f"{colorstr('activation:')} {act}")  # print
            if not ch_mul:
                ch_mul = 8
            na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # number of anchors
            no = na * (nc + 5)  # number of outputs = anchors * (classes + 5)
        
            layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
            for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):  # from, number, module, args
                m = eval(m) if isinstance(m, str) else m  # eval strings
                for j, a in enumerate(args):
                    with contextlib.suppress(NameError):
                        args[j] = eval(a) if isinstance(a, str) else a  # eval strings
        
                n = n_ = max(round(n * gd), 1) if n > 1 else n  # depth gain
                if m in {
                    Conv,
                    GhostConv,
                    Bottleneck,
                    GhostBottleneck,
                    SPP,
                    SPPF,
                    DWConv,
                    MixConv2d,
                    Focus,
                    CrossConv,
                    BottleneckCSP,
                    C3,
                    C3TR,
                    C3SPP,
                    C3Ghost,
                    nn.ConvTranspose2d,
                    DWConvTranspose2d,
                    C3x,
                }:
                    c1, c2 = ch[f], args[0]
                    if c2 != no:  # if not output
                        c2 = make_divisible(c2 * gw, ch_mul)
        
                    args = [c1, c2, *args[1:]]
                    if m in {BottleneckCSP, C3, C3TR, C3Ghost, C3x}:
                        args.insert(2, n)  # number of repeats
                        n = 1
                elif m is nn.BatchNorm2d:
                    args = [ch[f]]
                elif m is Concat:
                    c2 = sum(ch[x] for x in f)
                # TODO: channel, gw, gd
                elif m in {Detect, Segment}:
                    args.append([ch[x] for x in f])
                    if isinstance(args[1], int):  # number of anchors
                        args[1] = [list(range(args[1] * 2))] * len(f)
                    if m is Segment:
                        args[3] = make_divisible(args[3] * gw, ch_mul)
                elif m is Contract:
                    c2 = ch[f] * args[0] ** 2
                elif m is Expand:
                    c2 = ch[f] // args[0] ** 2
                else:
                    c2 = ch[f]
        
                m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
                t = str(m)[8:-2].replace("__main__.", "")  # module type
                np = sum(x.numel() for x in m_.parameters())  # number params
                m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
                LOGGER.info(f"{i:>3}{str(f):>18}{n_:>3}{np:10.0f}  {t:<40}{str(args):<30}")  # print
                save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
                layers.append(m_)
                if i == 0:
                    ch = []
                ch.append(c2)
            return nn.Sequential(*layers), sorted(save)
        
        
        if __name__ == "__main__":
            parser = argparse.ArgumentParser()
            parser.add_argument("--cfg", type=str, default="yolov5s.yaml", help="model.yaml")
            parser.add_argument("--batch-size", type=int, default=1, help="total batch size for all GPUs")
            parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
            parser.add_argument("--profile", action="store_true", help="profile model speed")
            parser.add_argument("--line-profile", action="store_true", help="profile model speed layer by layer")
            parser.add_argument("--test", action="store_true", help="test all yolo*.yaml")
            opt = parser.parse_args()
            opt.cfg = check_yaml(opt.cfg)  # check YAML
            print_args(vars(opt))
            device = select_device(opt.device)
        
            # Create model
            im = torch.rand(opt.batch_size, 3, 640, 640).to(device)
            model = Model(opt.cfg).to(device)
        
            # Options
            if opt.line_profile:  # profile layer by layer
                model(im, profile=True)
        
            elif opt.profile:  # profile forward-backward
                results = profile(input=im, ops=[model], n=3)
        
            elif opt.test:  # test all models
                for cfg in Path(ROOT / "models").rglob("yolo*.yaml"):
                    try:
                        _ = Model(cfg)
                    except Exception as e:
                        print(f"Error in {cfg}: {e}")
        
            else:  # report fused model summary
                model.fuse()
        
      • 📄 yolov5l.yaml 1.4KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        
        # Parameters
        nc: 80 # number of classes
        depth_multiple: 1.0 # model depth multiple
        width_multiple: 1.0 # layer channel multiple
        anchors:
          - [10, 13, 16, 30, 33, 23] # P3/8
          - [30, 61, 62, 45, 59, 119] # P4/16
          - [116, 90, 156, 198, 373, 326] # P5/32
        
        # YOLOv5 v6.0 backbone
        backbone:
          # [from, number, module, args]
          [
            [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
            [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
            [-1, 3, C3, [128]],
            [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
            [-1, 6, C3, [256]],
            [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
            [-1, 9, C3, [512]],
            [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
            [-1, 3, C3, [1024]],
            [-1, 1, SPPF, [1024, 5]], # 9
          ]
        
        # YOLOv5 v6.0 head
        head: [
            [-1, 1, Conv, [512, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 6], 1, Concat, [1]], # cat backbone P4
            [-1, 3, C3, [512, False]], # 13
        
            [-1, 1, Conv, [256, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 4], 1, Concat, [1]], # cat backbone P3
            [-1, 3, C3, [256, False]], # 17 (P3/8-small)
        
            [-1, 1, Conv, [256, 3, 2]],
            [[-1, 14], 1, Concat, [1]], # cat head P4
            [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
        
            [-1, 1, Conv, [512, 3, 2]],
            [[-1, 10], 1, Concat, [1]], # cat head P5
            [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
        
            [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
          ]
        
      • 📄 yolov5m.yaml 1.4KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        
        # Parameters
        nc: 80 # number of classes
        depth_multiple: 0.67 # model depth multiple
        width_multiple: 0.75 # layer channel multiple
        anchors:
          - [10, 13, 16, 30, 33, 23] # P3/8
          - [30, 61, 62, 45, 59, 119] # P4/16
          - [116, 90, 156, 198, 373, 326] # P5/32
        
        # YOLOv5 v6.0 backbone
        backbone:
          # [from, number, module, args]
          [
            [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
            [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
            [-1, 3, C3, [128]],
            [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
            [-1, 6, C3, [256]],
            [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
            [-1, 9, C3, [512]],
            [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
            [-1, 3, C3, [1024]],
            [-1, 1, SPPF, [1024, 5]], # 9
          ]
        
        # YOLOv5 v6.0 head
        head: [
            [-1, 1, Conv, [512, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 6], 1, Concat, [1]], # cat backbone P4
            [-1, 3, C3, [512, False]], # 13
        
            [-1, 1, Conv, [256, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 4], 1, Concat, [1]], # cat backbone P3
            [-1, 3, C3, [256, False]], # 17 (P3/8-small)
        
            [-1, 1, Conv, [256, 3, 2]],
            [[-1, 14], 1, Concat, [1]], # cat head P4
            [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
        
            [-1, 1, Conv, [512, 3, 2]],
            [[-1, 10], 1, Concat, [1]], # cat head P5
            [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
        
            [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
          ]
        
      • 📄 yolov5n.yaml 1.4KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        
        # Parameters
        nc: 80 # number of classes
        depth_multiple: 0.33 # model depth multiple
        width_multiple: 0.25 # layer channel multiple
        anchors:
          - [10, 13, 16, 30, 33, 23] # P3/8
          - [30, 61, 62, 45, 59, 119] # P4/16
          - [116, 90, 156, 198, 373, 326] # P5/32
        
        # YOLOv5 v6.0 backbone
        backbone:
          # [from, number, module, args]
          [
            [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
            [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
            [-1, 3, C3, [128]],
            [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
            [-1, 6, C3, [256]],
            [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
            [-1, 9, C3, [512]],
            [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
            [-1, 3, C3, [1024]],
            [-1, 1, SPPF, [1024, 5]], # 9
          ]
        
        # YOLOv5 v6.0 head
        head: [
            [-1, 1, Conv, [512, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 6], 1, Concat, [1]], # cat backbone P4
            [-1, 3, C3, [512, False]], # 13
        
            [-1, 1, Conv, [256, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 4], 1, Concat, [1]], # cat backbone P3
            [-1, 3, C3, [256, False]], # 17 (P3/8-small)
        
            [-1, 1, Conv, [256, 3, 2]],
            [[-1, 14], 1, Concat, [1]], # cat head P4
            [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
        
            [-1, 1, Conv, [512, 3, 2]],
            [[-1, 10], 1, Concat, [1]], # cat head P5
            [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
        
            [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
          ]
        
      • 📄 yolov5s.yaml 1.4KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        
        # Parameters
        nc: 80 # number of classes
        depth_multiple: 0.33 # model depth multiple
        width_multiple: 0.50 # layer channel multiple
        anchors:
          - [10, 13, 16, 30, 33, 23] # P3/8
          - [30, 61, 62, 45, 59, 119] # P4/16
          - [116, 90, 156, 198, 373, 326] # P5/32
        
        # YOLOv5 v6.0 backbone
        backbone:
          # [from, number, module, args]
          [
            [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
            [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
            [-1, 3, C3, [128]],  # 2
            [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
            [-1, 6, C3, [256]], # 4
            [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
            [-1, 9, C3, [512]], # 6
            [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
            [-1, 3, C3, [1024]],
            [-1, 1, SPPF, [1024, 5]], # 9
          ]
        
        # YOLOv5 v6.0 head
        head: [
            [-1, 1, Conv, [512, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 6], 1, Concat, [1]], # cat backbone P4
            [-1, 3, C3, [512, False]], # 13
        
            [-1, 1, Conv, [256, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 4], 1, Concat, [1]], # cat backbone P3
            [-1, 3, C3, [256, False]], # 17 (P3/8-small)
        
            [-1, 1, Conv, [256, 3, 2]],
            [[-1, 14], 1, Concat, [1]], # cat head P4
            [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
        
            [-1, 1, Conv, [512, 3, 2]],
            [[-1, 10], 1, Concat, [1]], # cat head P5
            [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
        
            [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
          ]
        
      • 📄 yolov5x.yaml 1.4KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        
        # Parameters
        nc: 80 # number of classes
        depth_multiple: 1.33 # model depth multiple
        width_multiple: 1.25 # layer channel multiple
        anchors:
          - [10, 13, 16, 30, 33, 23] # P3/8
          - [30, 61, 62, 45, 59, 119] # P4/16
          - [116, 90, 156, 198, 373, 326] # P5/32
        
        # YOLOv5 v6.0 backbone
        backbone:
          # [from, number, module, args]
          [
            [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
            [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
            [-1, 3, C3, [128]],
            [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
            [-1, 6, C3, [256]],
            [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
            [-1, 9, C3, [512]],
            [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
            [-1, 3, C3, [1024]],
            [-1, 1, SPPF, [1024, 5]], # 9
          ]
        
        # YOLOv5 v6.0 head
        head: [
            [-1, 1, Conv, [512, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 6], 1, Concat, [1]], # cat backbone P4
            [-1, 3, C3, [512, False]], # 13
        
            [-1, 1, Conv, [256, 1, 1]],
            [-1, 1, nn.Upsample, [None, 2, "nearest"]],
            [[-1, 4], 1, Concat, [1]], # cat backbone P3
            [-1, 3, C3, [256, False]], # 17 (P3/8-small)
        
            [-1, 1, Conv, [256, 3, 2]],
            [[-1, 14], 1, Concat, [1]], # cat head P4
            [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
        
            [-1, 1, Conv, [512, 3, 2]],
            [[-1, 10], 1, Concat, [1]], # cat head P5
            [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
        
            [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
          ]
        
    • 📁 segment/
      • 📄 predict.py 15.8KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """
        Run YOLOv5 segmentation inference on images, videos, directories, streams, etc.
        
        Usage - sources:
            $ python segment/predict.py --weights yolov5s-seg.pt --source 0                               # webcam
                                                                          img.jpg                         # image
                                                                          vid.mp4                         # video
                                                                          screen                          # screenshot
                                                                          path/                           # directory
                                                                          list.txt                        # list of images
                                                                          list.streams                    # list of streams
                                                                          'path/*.jpg'                    # glob
                                                                          'https://youtu.be/LNwODJXcvt4'  # YouTube
                                                                          'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
        
        Usage - formats:
            $ python segment/predict.py --weights yolov5s-seg.pt                 # PyTorch
                                                  yolov5s-seg.torchscript        # TorchScript
                                                  yolov5s-seg.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                                  yolov5s-seg_openvino_model     # OpenVINO
                                                  yolov5s-seg.engine             # TensorRT
                                                  yolov5s-seg.mlmodel            # CoreML (macOS-only)
                                                  yolov5s-seg_saved_model        # TensorFlow SavedModel
                                                  yolov5s-seg.pb                 # TensorFlow GraphDef
                                                  yolov5s-seg.tflite             # TensorFlow Lite
                                                  yolov5s-seg_edgetpu.tflite     # TensorFlow Edge TPU
                                                  yolov5s-seg_paddle_model       # PaddlePaddle
        """
        
        import argparse
        import os
        import platform
        import sys
        from pathlib import Path
        
        import torch
        
        FILE = Path(__file__).resolve()
        ROOT = FILE.parents[1]  # YOLOv5 root directory
        if str(ROOT) not in sys.path:
            sys.path.append(str(ROOT))  # add ROOT to PATH
        ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
        
        from ultralytics.utils.plotting import Annotator, colors, save_one_box
        
        from models.common import DetectMultiBackend
        from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
        from utils.general import (
            LOGGER,
            Profile,
            check_file,
            check_img_size,
            check_imshow,
            check_requirements,
            colorstr,
            cv2,
            increment_path,
            non_max_suppression,
            print_args,
            scale_boxes,
            scale_segments,
            strip_optimizer,
        )
        from utils.segment.general import masks2segments, process_mask, process_mask_native
        from utils.torch_utils import select_device, smart_inference_mode
        
        
        @smart_inference_mode()
        def run(
            weights=ROOT / "yolov5s-seg.pt",  # model.pt path(s)
            source=ROOT / "data/images",  # file/dir/URL/glob/screen/0(webcam)
            data=ROOT / "data/coco128.yaml",  # dataset.yaml path
            imgsz=(640, 640),  # inference size (height, width)
            conf_thres=0.25,  # confidence threshold
            iou_thres=0.45,  # NMS IOU threshold
            max_det=1000,  # maximum detections per image
            device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
            view_img=False,  # show results
            save_txt=False,  # save results to *.txt
            save_conf=False,  # save confidences in --save-txt labels
            save_crop=False,  # save cropped prediction boxes
            nosave=False,  # do not save images/videos
            classes=None,  # filter by class: --class 0, or --class 0 2 3
            agnostic_nms=False,  # class-agnostic NMS
            augment=False,  # augmented inference
            visualize=False,  # visualize features
            update=False,  # update all models
            project=ROOT / "runs/predict-seg",  # save results to project/name
            name="exp",  # save results to project/name
            exist_ok=False,  # existing project/name ok, do not increment
            line_thickness=3,  # bounding box thickness (pixels)
            hide_labels=False,  # hide labels
            hide_conf=False,  # hide confidences
            half=False,  # use FP16 half-precision inference
            dnn=False,  # use OpenCV DNN for ONNX inference
            vid_stride=1,  # video frame-rate stride
            retina_masks=False,
        ):
            source = str(source)
            save_img = not nosave and not source.endswith(".txt")  # save inference images
            is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
            is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
            webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
            screenshot = source.lower().startswith("screen")
            if is_url and is_file:
                source = check_file(source)  # download
        
            # Directories
            save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
            (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
        
            # Load model
            device = select_device(device)
            model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
            stride, names, pt = model.stride, model.names, model.pt
            imgsz = check_img_size(imgsz, s=stride)  # check image size
        
            # Dataloader
            bs = 1  # batch_size
            if webcam:
                view_img = check_imshow(warn=True)
                dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
                bs = len(dataset)
            elif screenshot:
                dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
            else:
                dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
            vid_path, vid_writer = [None] * bs, [None] * bs
        
            # Run inference
            model.warmup(imgsz=(1 if pt else bs, 3, *imgsz))  # warmup
            seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))
            for path, im, im0s, vid_cap, s in dataset:
                with dt[0]:
                    im = torch.from_numpy(im).to(model.device)
                    im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
                    im /= 255  # 0 - 255 to 0.0 - 1.0
                    if len(im.shape) == 3:
                        im = im[None]  # expand for batch dim
        
                # Inference
                with dt[1]:
                    visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
                    pred, proto = model(im, augment=augment, visualize=visualize)[:2]
        
                # NMS
                with dt[2]:
                    pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det, nm=32)
        
                # Second-stage classifier (optional)
                # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
        
                # Process predictions
                for i, det in enumerate(pred):  # per image
                    seen += 1
                    if webcam:  # batch_size >= 1
                        p, im0, frame = path[i], im0s[i].copy(), dataset.count
                        s += f"{i}: "
                    else:
                        p, im0, frame = path, im0s.copy(), getattr(dataset, "frame", 0)
        
                    p = Path(p)  # to Path
                    save_path = str(save_dir / p.name)  # im.jpg
                    txt_path = str(save_dir / "labels" / p.stem) + ("" if dataset.mode == "image" else f"_{frame}")  # im.txt
                    s += "%gx%g " % im.shape[2:]  # print string
                    imc = im0.copy() if save_crop else im0  # for save_crop
                    annotator = Annotator(im0, line_width=line_thickness, example=str(names))
                    if len(det):
                        if retina_masks:
                            # scale bbox first the crop masks
                            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()  # rescale boxes to im0 size
                            masks = process_mask_native(proto[i], det[:, 6:], det[:, :4], im0.shape[:2])  # HWC
                        else:
                            masks = process_mask(proto[i], det[:, 6:], det[:, :4], im.shape[2:], upsample=True)  # HWC
                            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()  # rescale boxes to im0 size
        
                        # Segments
                        if save_txt:
                            segments = [
                                scale_segments(im0.shape if retina_masks else im.shape[2:], x, im0.shape, normalize=True)
                                for x in reversed(masks2segments(masks))
                            ]
        
                        # Print results
                        for c in det[:, 5].unique():
                            n = (det[:, 5] == c).sum()  # detections per class
                            s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
        
                        # Mask plotting
                        annotator.masks(
                            masks,
                            colors=[colors(x, True) for x in det[:, 5]],
                            im_gpu=torch.as_tensor(im0, dtype=torch.float16).to(device).permute(2, 0, 1).flip(0).contiguous()
                            / 255
                            if retina_masks
                            else im[i],
                        )
        
                        # Write results
                        for j, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])):
                            if save_txt:  # Write to file
                                seg = segments[j].reshape(-1)  # (n,2) to (n*2)
                                line = (cls, *seg, conf) if save_conf else (cls, *seg)  # label format
                                with open(f"{txt_path}.txt", "a") as f:
                                    f.write(("%g " * len(line)).rstrip() % line + "\n")
        
                            if save_img or save_crop or view_img:  # Add bbox to image
                                c = int(cls)  # integer class
                                label = None if hide_labels else (names[c] if hide_conf else f"{names[c]} {conf:.2f}")
                                annotator.box_label(xyxy, label, color=colors(c, True))
                                # annotator.draw.polygon(segments[j], outline=colors(c, True), width=3)
                            if save_crop:
                                save_one_box(xyxy, imc, file=save_dir / "crops" / names[c] / f"{p.stem}.jpg", BGR=True)
        
                    # Stream results
                    im0 = annotator.result()
                    if view_img:
                        if platform.system() == "Linux" and p not in windows:
                            windows.append(p)
                            cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
                            cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
                        cv2.imshow(str(p), im0)
                        if cv2.waitKey(1) == ord("q"):  # 1 millisecond
                            exit()
        
                    # Save results (image with detections)
                    if save_img:
                        if dataset.mode == "image":
                            cv2.imwrite(save_path, im0)
                        else:  # 'video' or 'stream'
                            if vid_path[i] != save_path:  # new video
                                vid_path[i] = save_path
                                if isinstance(vid_writer[i], cv2.VideoWriter):
                                    vid_writer[i].release()  # release previous video writer
                                if vid_cap:  # video
                                    fps = vid_cap.get(cv2.CAP_PROP_FPS)
                                    w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                                    h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                                else:  # stream
                                    fps, w, h = 30, im0.shape[1], im0.shape[0]
                                save_path = str(Path(save_path).with_suffix(".mp4"))  # force *.mp4 suffix on results videos
                                vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
                            vid_writer[i].write(im0)
        
                # Print time (inference-only)
                LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
        
            # Print results
            t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
            LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}" % t)
            if save_txt or save_img:
                s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
                LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
            if update:
                strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
        
        
        def parse_opt():
            """Parses command-line options for YOLOv5 inference including model paths, data sources, inference settings, and
            output preferences.
            """
            parser = argparse.ArgumentParser()
            parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s-seg.pt", help="model path(s)")
            parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
            parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="(optional) dataset.yaml path")
            parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640], help="inference size h,w")
            parser.add_argument("--conf-thres", type=float, default=0.25, help="confidence threshold")
            parser.add_argument("--iou-thres", type=float, default=0.45, help="NMS IoU threshold")
            parser.add_argument("--max-det", type=int, default=1000, help="maximum detections per image")
            parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
            parser.add_argument("--view-img", action="store_true", help="show results")
            parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
            parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels")
            parser.add_argument("--save-crop", action="store_true", help="save cropped prediction boxes")
            parser.add_argument("--nosave", action="store_true", help="do not save images/videos")
            parser.add_argument("--classes", nargs="+", type=int, help="filter by class: --classes 0, or --classes 0 2 3")
            parser.add_argument("--agnostic-nms", action="store_true", help="class-agnostic NMS")
            parser.add_argument("--augment", action="store_true", help="augmented inference")
            parser.add_argument("--visualize", action="store_true", help="visualize features")
            parser.add_argument("--update", action="store_true", help="update all models")
            parser.add_argument("--project", default=ROOT / "runs/predict-seg", help="save results to project/name")
            parser.add_argument("--name", default="exp", help="save results to project/name")
            parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
            parser.add_argument("--line-thickness", default=3, type=int, help="bounding box thickness (pixels)")
            parser.add_argument("--hide-labels", default=False, action="store_true", help="hide labels")
            parser.add_argument("--hide-conf", default=False, action="store_true", help="hide confidences")
            parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
            parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
            parser.add_argument("--vid-stride", type=int, default=1, help="video frame-rate stride")
            parser.add_argument("--retina-masks", action="store_true", help="whether to plot masks in native resolution")
            opt = parser.parse_args()
            opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
            print_args(vars(opt))
            return opt
        
        
        def main(opt):
            """Executes YOLOv5 model inference with given options, checking for requirements before launching."""
            check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
            run(**vars(opt))
        
        
        if __name__ == "__main__":
            opt = parse_opt()
            main(opt)
        
      • 📄 train.py 34.2KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """
        Train a YOLOv5 segment model on a segment dataset Models and datasets download automatically from the latest YOLOv5
        release.
        
        Usage - Single-GPU training:
            $ python segment/train.py --data coco128-seg.yaml --weights yolov5s-seg.pt --img 640  # from pretrained (recommended)
            $ python segment/train.py --data coco128-seg.yaml --weights '' --cfg yolov5s-seg.yaml --img 640  # from scratch
        
        Usage - Multi-GPU DDP training:
            $ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 segment/train.py --data coco128-seg.yaml --weights yolov5s-seg.pt --img 640 --device 0,1,2,3
        
        Models:     https://github.com/ultralytics/yolov5/tree/master/models
        Datasets:   https://github.com/ultralytics/yolov5/tree/master/data
        Tutorial:   https://docs.ultralytics.com/yolov5/tutorials/train_custom_data
        """
        
        import argparse
        import math
        import os
        import random
        import subprocess
        import sys
        import time
        from copy import deepcopy
        from datetime import datetime
        from pathlib import Path
        
        import numpy as np
        import torch
        import torch.distributed as dist
        import torch.nn as nn
        import yaml
        from torch.optim import lr_scheduler
        from tqdm import tqdm
        
        FILE = Path(__file__).resolve()
        ROOT = FILE.parents[1]  # YOLOv5 root directory
        if str(ROOT) not in sys.path:
            sys.path.append(str(ROOT))  # add ROOT to PATH
        ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
        
        import segment.val as validate  # for end-of-epoch mAP
        from models.experimental import attempt_load
        from models.yolo import SegmentationModel
        from utils.autoanchor import check_anchors
        from utils.autobatch import check_train_batch_size
        from utils.callbacks import Callbacks
        from utils.downloads import attempt_download, is_url
        from utils.general import (
            LOGGER,
            TQDM_BAR_FORMAT,
            check_amp,
            check_dataset,
            check_file,
            check_git_info,
            check_git_status,
            check_img_size,
            check_requirements,
            check_suffix,
            check_yaml,
            colorstr,
            get_latest_run,
            increment_path,
            init_seeds,
            intersect_dicts,
            labels_to_class_weights,
            labels_to_image_weights,
            one_cycle,
            print_args,
            print_mutation,
            strip_optimizer,
            yaml_save,
        )
        from utils.loggers import GenericLogger
        from utils.plots import plot_evolve, plot_labels
        from utils.segment.dataloaders import create_dataloader
        from utils.segment.loss import ComputeLoss
        from utils.segment.metrics import KEYS, fitness
        from utils.segment.plots import plot_images_and_masks, plot_results_with_masks
        from utils.torch_utils import (
            EarlyStopping,
            ModelEMA,
            de_parallel,
            select_device,
            smart_DDP,
            smart_optimizer,
            smart_resume,
            torch_distributed_zero_first,
        )
        
        LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))  # https://pytorch.org/docs/stable/elastic/run.html
        RANK = int(os.getenv("RANK", -1))
        WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
        GIT_INFO = check_git_info()
        
        
        def train(hyp, opt, device, callbacks):
            """
            Trains the YOLOv5 model on a dataset, managing hyperparameters, model optimization, logging, and validation.
        
            `hyp` is path/to/hyp.yaml or hyp dictionary.
            """
            (
                save_dir,
                epochs,
                batch_size,
                weights,
                single_cls,
                evolve,
                data,
                cfg,
                resume,
                noval,
                nosave,
                workers,
                freeze,
                mask_ratio,
            ) = (
                Path(opt.save_dir),
                opt.epochs,
                opt.batch_size,
                opt.weights,
                opt.single_cls,
                opt.evolve,
                opt.data,
                opt.cfg,
                opt.resume,
                opt.noval,
                opt.nosave,
                opt.workers,
                opt.freeze,
                opt.mask_ratio,
            )
            # callbacks.run('on_pretrain_routine_start')
        
            # Directories
            w = save_dir / "weights"  # weights dir
            (w.parent if evolve else w).mkdir(parents=True, exist_ok=True)  # make dir
            last, best = w / "last.pt", w / "best.pt"
        
            # Hyperparameters
            if isinstance(hyp, str):
                with open(hyp, errors="ignore") as f:
                    hyp = yaml.safe_load(f)  # load hyps dict
            LOGGER.info(colorstr("hyperparameters: ") + ", ".join(f"{k}={v}" for k, v in hyp.items()))
            opt.hyp = hyp.copy()  # for saving hyps to checkpoints
        
            # Save run settings
            if not evolve:
                yaml_save(save_dir / "hyp.yaml", hyp)
                yaml_save(save_dir / "opt.yaml", vars(opt))
        
            # Loggers
            data_dict = None
            if RANK in {-1, 0}:
                logger = GenericLogger(opt=opt, console_logger=LOGGER)
        
            # Config
            plots = not evolve and not opt.noplots  # create plots
            overlap = not opt.no_overlap
            cuda = device.type != "cpu"
            init_seeds(opt.seed + 1 + RANK, deterministic=True)
            with torch_distributed_zero_first(LOCAL_RANK):
                data_dict = data_dict or check_dataset(data)  # check if None
            train_path, val_path = data_dict["train"], data_dict["val"]
            nc = 1 if single_cls else int(data_dict["nc"])  # number of classes
            names = {0: "item"} if single_cls and len(data_dict["names"]) != 1 else data_dict["names"]  # class names
            is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt")  # COCO dataset
        
            # Model
            check_suffix(weights, ".pt")  # check weights
            pretrained = weights.endswith(".pt")
            if pretrained:
                with torch_distributed_zero_first(LOCAL_RANK):
                    weights = attempt_download(weights)  # download if not found locally
                ckpt = torch.load(weights, map_location="cpu")  # load checkpoint to CPU to avoid CUDA memory leak
                model = SegmentationModel(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)
                exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else []  # exclude keys
                csd = ckpt["model"].float().state_dict()  # checkpoint state_dict as FP32
                csd = intersect_dicts(csd, model.state_dict(), exclude=exclude)  # intersect
                model.load_state_dict(csd, strict=False)  # load
                LOGGER.info(f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}")  # report
            else:
                model = SegmentationModel(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)  # create
            amp = check_amp(model)  # check AMP
        
            # Freeze
            freeze = [f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0]))]  # layers to freeze
            for k, v in model.named_parameters():
                v.requires_grad = True  # train all layers
                # v.register_hook(lambda x: torch.nan_to_num(x))  # NaN to 0 (commented for erratic training results)
                if any(x in k for x in freeze):
                    LOGGER.info(f"freezing {k}")
                    v.requires_grad = False
        
            # Image size
            gs = max(int(model.stride.max()), 32)  # grid size (max stride)
            imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2)  # verify imgsz is gs-multiple
        
            # Batch size
            if RANK == -1 and batch_size == -1:  # single-GPU only, estimate best batch size
                batch_size = check_train_batch_size(model, imgsz, amp)
                logger.update_params({"batch_size": batch_size})
                # loggers.on_params_update({"batch_size": batch_size})
        
            # Optimizer
            nbs = 64  # nominal batch size
            accumulate = max(round(nbs / batch_size), 1)  # accumulate loss before optimizing
            hyp["weight_decay"] *= batch_size * accumulate / nbs  # scale weight_decay
            optimizer = smart_optimizer(model, opt.optimizer, hyp["lr0"], hyp["momentum"], hyp["weight_decay"])
        
            # Scheduler
            if opt.cos_lr:
                lf = one_cycle(1, hyp["lrf"], epochs)  # cosine 1->hyp['lrf']
            else:
                lf = lambda x: (1 - x / epochs) * (1.0 - hyp["lrf"]) + hyp["lrf"]  # linear
            scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)  # plot_lr_scheduler(optimizer, scheduler, epochs)
        
            # EMA
            ema = ModelEMA(model) if RANK in {-1, 0} else None
        
            # Resume
            best_fitness, start_epoch = 0.0, 0
            if pretrained:
                if resume:
                    best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
                del ckpt, csd
        
            # DP mode
            if cuda and RANK == -1 and torch.cuda.device_count() > 1:
                LOGGER.warning(
                    "WARNING ⚠️ DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n"
                    "See Multi-GPU Tutorial at https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training to get started."
                )
                model = torch.nn.DataParallel(model)
        
            # SyncBatchNorm
            if opt.sync_bn and cuda and RANK != -1:
                model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
                LOGGER.info("Using SyncBatchNorm()")
        
            # Trainloader
            train_loader, dataset = create_dataloader(
                train_path,
                imgsz,
                batch_size // WORLD_SIZE,
                gs,
                single_cls,
                hyp=hyp,
                augment=True,
                cache=None if opt.cache == "val" else opt.cache,
                rect=opt.rect,
                rank=LOCAL_RANK,
                workers=workers,
                image_weights=opt.image_weights,
                quad=opt.quad,
                prefix=colorstr("train: "),
                shuffle=True,
                mask_downsample_ratio=mask_ratio,
                overlap_mask=overlap,
            )
            labels = np.concatenate(dataset.labels, 0)
            mlc = int(labels[:, 0].max())  # max label class
            assert mlc < nc, f"Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}"
        
            # Process 0
            if RANK in {-1, 0}:
                val_loader = create_dataloader(
                    val_path,
                    imgsz,
                    batch_size // WORLD_SIZE * 2,
                    gs,
                    single_cls,
                    hyp=hyp,
                    cache=None if noval else opt.cache,
                    rect=True,
                    rank=-1,
                    workers=workers * 2,
                    pad=0.5,
                    mask_downsample_ratio=mask_ratio,
                    overlap_mask=overlap,
                    prefix=colorstr("val: "),
                )[0]
        
                if not resume:
                    if not opt.noautoanchor:
                        check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz)  # run AutoAnchor
                    model.half().float()  # pre-reduce anchor precision
        
                    if plots:
                        plot_labels(labels, names, save_dir)
                # callbacks.run('on_pretrain_routine_end', labels, names)
        
            # DDP mode
            if cuda and RANK != -1:
                model = smart_DDP(model)
        
            # Model attributes
            nl = de_parallel(model).model[-1].nl  # number of detection layers (to scale hyps)
            hyp["box"] *= 3 / nl  # scale to layers
            hyp["cls"] *= nc / 80 * 3 / nl  # scale to classes and layers
            hyp["obj"] *= (imgsz / 640) ** 2 * 3 / nl  # scale to image size and layers
            hyp["label_smoothing"] = opt.label_smoothing
            model.nc = nc  # attach number of classes to model
            model.hyp = hyp  # attach hyperparameters to model
            model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc  # attach class weights
            model.names = names
        
            # Start training
            t0 = time.time()
            nb = len(train_loader)  # number of batches
            nw = max(round(hyp["warmup_epochs"] * nb), 100)  # number of warmup iterations, max(3 epochs, 100 iterations)
            # nw = min(nw, (epochs - start_epoch) / 2 * nb)  # limit warmup to < 1/2 of training
            last_opt_step = -1
            maps = np.zeros(nc)  # mAP per class
            results = (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)  # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
            scheduler.last_epoch = start_epoch - 1  # do not move
            scaler = torch.cuda.amp.GradScaler(enabled=amp)
            stopper, stop = EarlyStopping(patience=opt.patience), False
            compute_loss = ComputeLoss(model, overlap=overlap)  # init loss class
            # callbacks.run('on_train_start')
            LOGGER.info(
                f'Image sizes {imgsz} train, {imgsz} val\n'
                f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'
                f"Logging results to {colorstr('bold', save_dir)}\n"
                f'Starting training for {epochs} epochs...'
            )
            for epoch in range(start_epoch, epochs):  # epoch ------------------------------------------------------------------
                # callbacks.run('on_train_epoch_start')
                model.train()
        
                # Update image weights (optional, single-GPU only)
                if opt.image_weights:
                    cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc  # class weights
                    iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)  # image weights
                    dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)  # rand weighted idx
        
                # Update mosaic border (optional)
                # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
                # dataset.mosaic_border = [b - imgsz, -b]  # height, width borders
        
                mloss = torch.zeros(4, device=device)  # mean losses
                if RANK != -1:
                    train_loader.sampler.set_epoch(epoch)
                pbar = enumerate(train_loader)
                LOGGER.info(
                    ("\n" + "%11s" * 8)
                    % ("Epoch", "GPU_mem", "box_loss", "seg_loss", "obj_loss", "cls_loss", "Instances", "Size")
                )
                if RANK in {-1, 0}:
                    pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT)  # progress bar
                optimizer.zero_grad()
                for i, (imgs, targets, paths, _, masks) in pbar:  # batch ------------------------------------------------------
                    # callbacks.run('on_train_batch_start')
                    ni = i + nb * epoch  # number integrated batches (since train start)
                    imgs = imgs.to(device, non_blocking=True).float() / 255  # uint8 to float32, 0-255 to 0.0-1.0
        
                    # Warmup
                    if ni <= nw:
                        xi = [0, nw]  # x interp
                        # compute_loss.gr = np.interp(ni, xi, [0.0, 1.0])  # iou loss ratio (obj_loss = 1.0 or iou)
                        accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
                        for j, x in enumerate(optimizer.param_groups):
                            # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
                            x["lr"] = np.interp(ni, xi, [hyp["warmup_bias_lr"] if j == 0 else 0.0, x["initial_lr"] * lf(epoch)])
                            if "momentum" in x:
                                x["momentum"] = np.interp(ni, xi, [hyp["warmup_momentum"], hyp["momentum"]])
        
                    # Multi-scale
                    if opt.multi_scale:
                        sz = random.randrange(int(imgsz * 0.5), int(imgsz * 1.5) + gs) // gs * gs  # size
                        sf = sz / max(imgs.shape[2:])  # scale factor
                        if sf != 1:
                            ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to gs-multiple)
                            imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
        
                    # Forward
                    with torch.cuda.amp.autocast(amp):
                        pred = model(imgs)  # forward
                        loss, loss_items = compute_loss(pred, targets.to(device), masks=masks.to(device).float())
                        if RANK != -1:
                            loss *= WORLD_SIZE  # gradient averaged between devices in DDP mode
                        if opt.quad:
                            loss *= 4.0
        
                    # Backward
                    scaler.scale(loss).backward()
        
                    # Optimize - https://pytorch.org/docs/master/notes/amp_examples.html
                    if ni - last_opt_step >= accumulate:
                        scaler.unscale_(optimizer)  # unscale gradients
                        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)  # clip gradients
                        scaler.step(optimizer)  # optimizer.step
                        scaler.update()
                        optimizer.zero_grad()
                        if ema:
                            ema.update(model)
                        last_opt_step = ni
        
                    # Log
                    if RANK in {-1, 0}:
                        mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses
                        mem = f"{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G"  # (GB)
                        pbar.set_description(
                            ("%11s" * 2 + "%11.4g" * 6)
                            % (f"{epoch}/{epochs - 1}", mem, *mloss, targets.shape[0], imgs.shape[-1])
                        )
                        # callbacks.run('on_train_batch_end', model, ni, imgs, targets, paths)
                        # if callbacks.stop_training:
                        #    return
        
                        # Mosaic plots
                        if plots:
                            if ni < 3:
                                plot_images_and_masks(imgs, targets, masks, paths, save_dir / f"train_batch{ni}.jpg")
                            if ni == 10:
                                files = sorted(save_dir.glob("train*.jpg"))
                                logger.log_images(files, "Mosaics", epoch)
                    # end batch ------------------------------------------------------------------------------------------------
        
                # Scheduler
                lr = [x["lr"] for x in optimizer.param_groups]  # for loggers
                scheduler.step()
        
                if RANK in {-1, 0}:
                    # mAP
                    # callbacks.run('on_train_epoch_end', epoch=epoch)
                    ema.update_attr(model, include=["yaml", "nc", "hyp", "names", "stride", "class_weights"])
                    final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
                    if not noval or final_epoch:  # Calculate mAP
                        results, maps, _ = validate.run(
                            data_dict,
                            batch_size=batch_size // WORLD_SIZE * 2,
                            imgsz=imgsz,
                            half=amp,
                            model=ema.ema,
                            single_cls=single_cls,
                            dataloader=val_loader,
                            save_dir=save_dir,
                            plots=False,
                            callbacks=callbacks,
                            compute_loss=compute_loss,
                            mask_downsample_ratio=mask_ratio,
                            overlap=overlap,
                        )
        
                    # Update best mAP
                    fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
                    stop = stopper(epoch=epoch, fitness=fi)  # early stop check
                    if fi > best_fitness:
                        best_fitness = fi
                    log_vals = list(mloss) + list(results) + lr
                    # callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
                    # Log val metrics and media
                    metrics_dict = dict(zip(KEYS, log_vals))
                    logger.log_metrics(metrics_dict, epoch)
        
                    # Save model
                    if (not nosave) or (final_epoch and not evolve):  # if save
                        ckpt = {
                            "epoch": epoch,
                            "best_fitness": best_fitness,
                            "model": deepcopy(de_parallel(model)).half(),
                            "ema": deepcopy(ema.ema).half(),
                            "updates": ema.updates,
                            "optimizer": optimizer.state_dict(),
                            "opt": vars(opt),
                            "git": GIT_INFO,  # {remote, branch, commit} if a git repo
                            "date": datetime.now().isoformat(),
                        }
        
                        # Save last, best and delete
                        torch.save(ckpt, last)
                        if best_fitness == fi:
                            torch.save(ckpt, best)
                        if opt.save_period > 0 and epoch % opt.save_period == 0:
                            torch.save(ckpt, w / f"epoch{epoch}.pt")
                            logger.log_model(w / f"epoch{epoch}.pt")
                        del ckpt
                        # callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
        
                # EarlyStopping
                if RANK != -1:  # if DDP training
                    broadcast_list = [stop if RANK == 0 else None]
                    dist.broadcast_object_list(broadcast_list, 0)  # broadcast 'stop' to all ranks
                    if RANK != 0:
                        stop = broadcast_list[0]
                if stop:
                    break  # must break all DDP ranks
        
                # end epoch ----------------------------------------------------------------------------------------------------
            # end training -----------------------------------------------------------------------------------------------------
            if RANK in {-1, 0}:
                LOGGER.info(f"\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.")
                for f in last, best:
                    if f.exists():
                        strip_optimizer(f)  # strip optimizers
                        if f is best:
                            LOGGER.info(f"\nValidating {f}...")
                            results, _, _ = validate.run(
                                data_dict,
                                batch_size=batch_size // WORLD_SIZE * 2,
                                imgsz=imgsz,
                                model=attempt_load(f, device).half(),
                                iou_thres=0.65 if is_coco else 0.60,  # best pycocotools at iou 0.65
                                single_cls=single_cls,
                                dataloader=val_loader,
                                save_dir=save_dir,
                                save_json=is_coco,
                                verbose=True,
                                plots=plots,
                                callbacks=callbacks,
                                compute_loss=compute_loss,
                                mask_downsample_ratio=mask_ratio,
                                overlap=overlap,
                            )  # val best model with plots
                            if is_coco:
                                # callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
                                metrics_dict = dict(zip(KEYS, list(mloss) + list(results) + lr))
                                logger.log_metrics(metrics_dict, epoch)
        
                # callbacks.run('on_train_end', last, best, epoch, results)
                # on train end callback using genericLogger
                logger.log_metrics(dict(zip(KEYS[4:16], results)), epochs)
                if not opt.evolve:
                    logger.log_model(best, epoch)
                if plots:
                    plot_results_with_masks(file=save_dir / "results.csv")  # save results.png
                    files = ["results.png", "confusion_matrix.png", *(f"{x}_curve.png" for x in ("F1", "PR", "P", "R"))]
                    files = [(save_dir / f) for f in files if (save_dir / f).exists()]  # filter
                    LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
                    logger.log_images(files, "Results", epoch + 1)
                    logger.log_images(sorted(save_dir.glob("val*.jpg")), "Validation", epoch + 1)
            torch.cuda.empty_cache()
            return results
        
        
        def parse_opt(known=False):
            """
            Parses command line arguments for training configurations, returning parsed arguments.
        
            Supports both known and unknown args.
            """
            parser = argparse.ArgumentParser()
            parser.add_argument("--weights", type=str, default=ROOT / "yolov5s-seg.pt", help="initial weights path")
            parser.add_argument("--cfg", type=str, default="", help="model.yaml path")
            parser.add_argument("--data", type=str, default=ROOT / "data/coco128-seg.yaml", help="dataset.yaml path")
            parser.add_argument("--hyp", type=str, default=ROOT / "data/hyps/hyp.scratch-low.yaml", help="hyperparameters path")
            parser.add_argument("--epochs", type=int, default=100, help="total training epochs")
            parser.add_argument("--batch-size", type=int, default=16, help="total batch size for all GPUs, -1 for autobatch")
            parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="train, val image size (pixels)")
            parser.add_argument("--rect", action="store_true", help="rectangular training")
            parser.add_argument("--resume", nargs="?", const=True, default=False, help="resume most recent training")
            parser.add_argument("--nosave", action="store_true", help="only save final checkpoint")
            parser.add_argument("--noval", action="store_true", help="only validate final epoch")
            parser.add_argument("--noautoanchor", action="store_true", help="disable AutoAnchor")
            parser.add_argument("--noplots", action="store_true", help="save no plot files")
            parser.add_argument("--evolve", type=int, nargs="?", const=300, help="evolve hyperparameters for x generations")
            parser.add_argument("--bucket", type=str, default="", help="gsutil bucket")
            parser.add_argument("--cache", type=str, nargs="?", const="ram", help="image --cache ram/disk")
            parser.add_argument("--image-weights", action="store_true", help="use weighted image selection for training")
            parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
            parser.add_argument("--multi-scale", action="store_true", help="vary img-size +/- 50%%")
            parser.add_argument("--single-cls", action="store_true", help="train multi-class data as single-class")
            parser.add_argument("--optimizer", type=str, choices=["SGD", "Adam", "AdamW"], default="SGD", help="optimizer")
            parser.add_argument("--sync-bn", action="store_true", help="use SyncBatchNorm, only available in DDP mode")
            parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
            parser.add_argument("--project", default=ROOT / "runs/train-seg", help="save to project/name")
            parser.add_argument("--name", default="exp", help="save to project/name")
            parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
            parser.add_argument("--quad", action="store_true", help="quad dataloader")
            parser.add_argument("--cos-lr", action="store_true", help="cosine LR scheduler")
            parser.add_argument("--label-smoothing", type=float, default=0.0, help="Label smoothing epsilon")
            parser.add_argument("--patience", type=int, default=100, help="EarlyStopping patience (epochs without improvement)")
            parser.add_argument("--freeze", nargs="+", type=int, default=[0], help="Freeze layers: backbone=10, first3=0 1 2")
            parser.add_argument("--save-period", type=int, default=-1, help="Save checkpoint every x epochs (disabled if < 1)")
            parser.add_argument("--seed", type=int, default=0, help="Global training seed")
            parser.add_argument("--local_rank", type=int, default=-1, help="Automatic DDP Multi-GPU argument, do not modify")
        
            # Instance Segmentation Args
            parser.add_argument("--mask-ratio", type=int, default=4, help="Downsample the truth masks to saving memory")
            parser.add_argument("--no-overlap", action="store_true", help="Overlap masks train faster at slightly less mAP")
        
            return parser.parse_known_args()[0] if known else parser.parse_args()
        
        
        def main(opt, callbacks=Callbacks()):
            """Initializes training or evolution of YOLOv5 models based on provided configuration and options."""
            if RANK in {-1, 0}:
                print_args(vars(opt))
                check_git_status()
                check_requirements(ROOT / "requirements.txt")
        
            # Resume
            if opt.resume and not opt.evolve:  # resume from specified or most recent last.pt
                last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
                opt_yaml = last.parent.parent / "opt.yaml"  # train options yaml
                opt_data = opt.data  # original dataset
                if opt_yaml.is_file():
                    with open(opt_yaml, errors="ignore") as f:
                        d = yaml.safe_load(f)
                else:
                    d = torch.load(last, map_location="cpu")["opt"]
                opt = argparse.Namespace(**d)  # replace
                opt.cfg, opt.weights, opt.resume = "", str(last), True  # reinstate
                if is_url(opt_data):
                    opt.data = check_file(opt_data)  # avoid HUB resume auth timeout
            else:
                opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = (
                    check_file(opt.data),
                    check_yaml(opt.cfg),
                    check_yaml(opt.hyp),
                    str(opt.weights),
                    str(opt.project),
                )  # checks
                assert len(opt.cfg) or len(opt.weights), "either --cfg or --weights must be specified"
                if opt.evolve:
                    if opt.project == str(ROOT / "runs/train-seg"):  # if default project name, rename to runs/evolve-seg
                        opt.project = str(ROOT / "runs/evolve-seg")
                    opt.exist_ok, opt.resume = opt.resume, False  # pass resume to exist_ok and disable resume
                if opt.name == "cfg":
                    opt.name = Path(opt.cfg).stem  # use model.yaml as name
                opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
        
            # DDP mode
            device = select_device(opt.device, batch_size=opt.batch_size)
            if LOCAL_RANK != -1:
                msg = "is not compatible with YOLOv5 Multi-GPU DDP training"
                assert not opt.image_weights, f"--image-weights {msg}"
                assert not opt.evolve, f"--evolve {msg}"
                assert opt.batch_size != -1, f"AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size"
                assert opt.batch_size % WORLD_SIZE == 0, f"--batch-size {opt.batch_size} must be multiple of WORLD_SIZE"
                assert torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command"
                torch.cuda.set_device(LOCAL_RANK)
                device = torch.device("cuda", LOCAL_RANK)
                dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
        
            # Train
            if not opt.evolve:
                train(opt.hyp, opt, device, callbacks)
        
            # Evolve hyperparameters (optional)
            else:
                # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
                meta = {
                    "lr0": (1, 1e-5, 1e-1),  # initial learning rate (SGD=1E-2, Adam=1E-3)
                    "lrf": (1, 0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)
                    "momentum": (0.3, 0.6, 0.98),  # SGD momentum/Adam beta1
                    "weight_decay": (1, 0.0, 0.001),  # optimizer weight decay
                    "warmup_epochs": (1, 0.0, 5.0),  # warmup epochs (fractions ok)
                    "warmup_momentum": (1, 0.0, 0.95),  # warmup initial momentum
                    "warmup_bias_lr": (1, 0.0, 0.2),  # warmup initial bias lr
                    "box": (1, 0.02, 0.2),  # box loss gain
                    "cls": (1, 0.2, 4.0),  # cls loss gain
                    "cls_pw": (1, 0.5, 2.0),  # cls BCELoss positive_weight
                    "obj": (1, 0.2, 4.0),  # obj loss gain (scale with pixels)
                    "obj_pw": (1, 0.5, 2.0),  # obj BCELoss positive_weight
                    "iou_t": (0, 0.1, 0.7),  # IoU training threshold
                    "anchor_t": (1, 2.0, 8.0),  # anchor-multiple threshold
                    "anchors": (2, 2.0, 10.0),  # anchors per output grid (0 to ignore)
                    "fl_gamma": (0, 0.0, 2.0),  # focal loss gamma (efficientDet default gamma=1.5)
                    "hsv_h": (1, 0.0, 0.1),  # image HSV-Hue augmentation (fraction)
                    "hsv_s": (1, 0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
                    "hsv_v": (1, 0.0, 0.9),  # image HSV-Value augmentation (fraction)
                    "degrees": (1, 0.0, 45.0),  # image rotation (+/- deg)
                    "translate": (1, 0.0, 0.9),  # image translation (+/- fraction)
                    "scale": (1, 0.0, 0.9),  # image scale (+/- gain)
                    "shear": (1, 0.0, 10.0),  # image shear (+/- deg)
                    "perspective": (0, 0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001
                    "flipud": (1, 0.0, 1.0),  # image flip up-down (probability)
                    "fliplr": (0, 0.0, 1.0),  # image flip left-right (probability)
                    "mosaic": (1, 0.0, 1.0),  # image mixup (probability)
                    "mixup": (1, 0.0, 1.0),  # image mixup (probability)
                    "copy_paste": (1, 0.0, 1.0),
                }  # segment copy-paste (probability)
        
                with open(opt.hyp, errors="ignore") as f:
                    hyp = yaml.safe_load(f)  # load hyps dict
                    if "anchors" not in hyp:  # anchors commented in hyp.yaml
                        hyp["anchors"] = 3
                if opt.noautoanchor:
                    del hyp["anchors"], meta["anchors"]
                opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir)  # only val/save final epoch
                # ei = [isinstance(x, (int, float)) for x in hyp.values()]  # evolvable indices
                evolve_yaml, evolve_csv = save_dir / "hyp_evolve.yaml", save_dir / "evolve.csv"
                if opt.bucket:
                    # download evolve.csv if exists
                    subprocess.run(
                        [
                            "gsutil",
                            "cp",
                            f"gs://{opt.bucket}/evolve.csv",
                            str(evolve_csv),
                        ]
                    )
        
                for _ in range(opt.evolve):  # generations to evolve
                    if evolve_csv.exists():  # if evolve.csv exists: select best hyps and mutate
                        # Select parent(s)
                        parent = "single"  # parent selection method: 'single' or 'weighted'
                        x = np.loadtxt(evolve_csv, ndmin=2, delimiter=",", skiprows=1)
                        n = min(5, len(x))  # number of previous results to consider
                        x = x[np.argsort(-fitness(x))][:n]  # top n mutations
                        w = fitness(x) - fitness(x).min() + 1e-6  # weights (sum > 0)
                        if parent == "single" or len(x) == 1:
                            # x = x[random.randint(0, n - 1)]  # random selection
                            x = x[random.choices(range(n), weights=w)[0]]  # weighted selection
                        elif parent == "weighted":
                            x = (x * w.reshape(n, 1)).sum(0) / w.sum()  # weighted combination
        
                        # Mutate
                        mp, s = 0.8, 0.2  # mutation probability, sigma
                        npr = np.random
                        npr.seed(int(time.time()))
                        g = np.array([meta[k][0] for k in hyp.keys()])  # gains 0-1
                        ng = len(meta)
                        v = np.ones(ng)
                        while all(v == 1):  # mutate until a change occurs (prevent duplicates)
                            v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)
                        for i, k in enumerate(hyp.keys()):  # plt.hist(v.ravel(), 300)
                            hyp[k] = float(x[i + 12] * v[i])  # mutate
        
                    # Constrain to limits
                    for k, v in meta.items():
                        hyp[k] = max(hyp[k], v[1])  # lower limit
                        hyp[k] = min(hyp[k], v[2])  # upper limit
                        hyp[k] = round(hyp[k], 5)  # significant digits
        
                    # Train mutation
                    results = train(hyp.copy(), opt, device, callbacks)
                    callbacks = Callbacks()
                    # Write mutation results
                    print_mutation(KEYS[4:16], results, hyp.copy(), save_dir, opt.bucket)
        
                # Plot results
                plot_evolve(evolve_csv)
                LOGGER.info(
                    f'Hyperparameter evolution finished {opt.evolve} generations\n'
                    f"Results saved to {colorstr('bold', save_dir)}\n"
                    f'Usage example: $ python train.py --hyp {evolve_yaml}'
                )
        
        
        def run(**kwargs):
            """
            Executes YOLOv5 training with given parameters, altering options programmatically; returns updated options.
        
            Example: mport train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
            """
            opt = parse_opt(True)
            for k, v in kwargs.items():
                setattr(opt, k, v)
            main(opt)
            return opt
        
        
        if __name__ == "__main__":
            opt = parse_opt()
            main(opt)
        
      • 📄 tutorial.ipynb 42.4KB

        Run on Gradient Open In Colab Open In Kaggle
        This YOLOv5 🚀 notebook by Ultralytics presents simple train, validate and predict examples to help start your AI adventure.
        See GitHub for community support or contact us for professional support.

        Setup

        Clone GitHub repository, install dependencies and check PyTorch and GPU.

        !git clone https://github.com/ultralytics/yolov5  # clone
        %cd yolov5
        %pip install -qr requirements.txt comet_ml  # install
        
        import torch
        import utils
        display = utils.notebook_init()  # checks
        
        YOLOv5 🚀 v7.0-2-gc9d47ae Python-3.7.15 torch-1.12.1+cu113 CUDA:0 (Tesla T4, 15110MiB)
        Setup complete ✅ (2 CPUs, 12.7 GB RAM, 22.6/78.2 GB disk)
        

        1. Predict

        segment/predict.py runs YOLOv5 instance segmentation inference on a variety of sources, downloading models automatically from the latest YOLOv5 release, and saving results to runs/predict. Example inference sources are:

        python segment/predict.py --source 0  # webcam
                                     img.jpg  # image 
                                     vid.mp4  # video
                                     screen  # screenshot
                                     path/  # directory
                                     'path/*.jpg'  # glob
                                     'https://youtu.be/LNwODJXcvt4'  # YouTube
                                     'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
        
        !python segment/predict.py --weights yolov5s-seg.pt --img 640 --conf 0.25 --source data/images
        #display.Image(filename='runs/predict-seg/exp/zidane.jpg', width=600)
        
        segment/predict: weights=['yolov5s-seg.pt'], source=data/images, data=data/coco128.yaml, imgsz=[640, 640], conf_thres=0.25, iou_thres=0.45, max_det=1000, device=, view_img=False, save_txt=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=runs/predict-seg, name=exp, exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, half=False, dnn=False, vid_stride=1, retina_masks=False
        YOLOv5 🚀 v7.0-2-gc9d47ae Python-3.7.15 torch-1.12.1+cu113 CUDA:0 (Tesla T4, 15110MiB)
        
        Downloading https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s-seg.pt to yolov5s-seg.pt...
        100% 14.9M/14.9M [00:01<00:00, 12.0MB/s]
        
        Fusing layers... 
        YOLOv5s-seg summary: 224 layers, 7611485 parameters, 0 gradients, 26.4 GFLOPs
        image 1/2 /content/yolov5/data/images/bus.jpg: 640x480 4 persons, 1 bus, 18.2ms
        image 2/2 /content/yolov5/data/images/zidane.jpg: 384x640 2 persons, 1 tie, 13.4ms
        Speed: 0.5ms pre-process, 15.8ms inference, 18.5ms NMS per image at shape (1, 3, 640, 640)
        Results saved to runs/predict-seg/exp
        

                

        2. Validate

        Validate a model's accuracy on the COCO dataset's val or test splits. Models are downloaded automatically from the latest YOLOv5 release. To show results by class use the --verbose flag.

        # Download COCO val
        !bash data/scripts/get_coco.sh --val --segments  # download (780M - 5000 images)
        
        Downloading https://github.com/ultralytics/yolov5/releases/download/v1.0/coco2017labels-segments.zip  ...
        Downloading http://images.cocodataset.org/zips/val2017.zip ...
        ######################################################################## 100.0%
        ######################################################################## 100.0%
        
        # Validate YOLOv5s-seg on COCO val
        !python segment/val.py --weights yolov5s-seg.pt --data coco.yaml --img 640 --half
        
        segment/val: data=/content/yolov5/data/coco.yaml, weights=['yolov5s-seg.pt'], batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.6, max_det=300, task=val, device=, workers=8, single_cls=False, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=runs/val-seg, name=exp, exist_ok=False, half=True, dnn=False
        YOLOv5 🚀 v7.0-2-gc9d47ae Python-3.7.15 torch-1.12.1+cu113 CUDA:0 (Tesla T4, 15110MiB)
        
        Fusing layers... 
        YOLOv5s-seg summary: 224 layers, 7611485 parameters, 0 gradients, 26.4 GFLOPs
        val: Scanning /content/datasets/coco/val2017... 4952 images, 48 backgrounds, 0 corrupt: 100% 5000/5000 [00:03<00:00, 1361.31it/s]
        val: New cache created: /content/datasets/coco/val2017.cache
                         Class     Images  Instances      Box(P          R      mAP50  mAP50-95)     Mask(P          R      mAP50  mAP50-95): 100% 157/157 [01:54<00:00,  1.37it/s]
                           all       5000      36335      0.673      0.517      0.566      0.373      0.672       0.49      0.532      0.319
        Speed: 0.6ms pre-process, 4.4ms inference, 2.9ms NMS per image at shape (32, 3, 640, 640)
        Results saved to runs/val-seg/exp
        

        3. Train

        Close the active learning loop by sampling images from your inference conditions with the roboflow pip package


        Train a YOLOv5s-seg model on the COCO128 dataset with --data coco128-seg.yaml, starting from pretrained --weights yolov5s-seg.pt, or from randomly initialized --weights '' --cfg yolov5s-seg.yaml.

        A Mosaic Dataloader is used for training which combines 4 images into 1 mosaic.

        Train on Custom Data with Roboflow 🌟 NEW

        Roboflow enables you to easily organize, label, and prepare a high quality dataset with your own custom data. Roboflow also makes it easy to establish an active learning pipeline, collaborate with your team on dataset improvement, and integrate directly into your model building workflow with the roboflow pip package.

        Label images lightning fast (including with model-assisted labeling)

        #@title Select YOLOv5 🚀 logger {run: 'auto'}
        logger = 'Comet' #@param ['Comet', 'ClearML', 'TensorBoard']
        
        if logger == 'Comet':
          %pip install -q comet_ml
          import comet_ml; comet_ml.init()
        elif logger == 'ClearML':
          %pip install -q clearml
          import clearml; clearml.browser_login()
        elif logger == 'TensorBoard':
          %load_ext tensorboard
          %tensorboard --logdir runs/train
        
        # Train YOLOv5s on COCO128 for 3 epochs
        !python segment/train.py --img 640 --batch 16 --epochs 3 --data coco128-seg.yaml --weights yolov5s-seg.pt --cache
        
        segment/train: weights=yolov5s-seg.pt, cfg=, data=coco128-seg.yaml, hyp=data/hyps/hyp.scratch-low.yaml, epochs=3, batch_size=16, imgsz=640, rect=False, resume=False, nosave=False, noval=False, noautoanchor=False, noplots=False, evolve=None, bucket=, cache=ram, image_weights=False, device=, multi_scale=False, single_cls=False, optimizer=SGD, sync_bn=False, workers=8, project=runs/train-seg, name=exp, exist_ok=False, quad=False, cos_lr=False, label_smoothing=0.0, patience=100, freeze=[0], save_period=-1, seed=0, local_rank=-1, mask_ratio=4, no_overlap=False
        github: up to date with https://github.com/ultralytics/yolov5 ✅
        YOLOv5 🚀 v7.0-2-gc9d47ae Python-3.7.15 torch-1.12.1+cu113 CUDA:0 (Tesla T4, 15110MiB)
        
        hyperparameters: lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=0.05, cls=0.5, cls_pw=1.0, obj=1.0, obj_pw=1.0, iou_t=0.2, anchor_t=4.0, fl_gamma=0.0, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, mosaic=1.0, mixup=0.0, copy_paste=0.0
        TensorBoard: Start with 'tensorboard --logdir runs/train-seg', view at http://localhost:6006/
        
        Dataset not found ⚠️, missing paths ['/content/datasets/coco128-seg/images/train2017']
        Downloading https://ultralytics.com/assets/coco128-seg.zip to coco128-seg.zip...
        100% 6.79M/6.79M [00:01<00:00, 6.73MB/s]
        Dataset download success ✅ (1.9s), saved to /content/datasets
        
                         from  n    params  module                                  arguments                     
          0                -1  1      3520  models.common.Conv                      [3, 32, 6, 2, 2]              
          1                -1  1     18560  models.common.Conv                      [32, 64, 3, 2]                
          2                -1  1     18816  models.common.C3                        [64, 64, 1]                   
          3                -1  1     73984  models.common.Conv                      [64, 128, 3, 2]               
          4                -1  2    115712  models.common.C3                        [128, 128, 2]                 
          5                -1  1    295424  models.common.Conv                      [128, 256, 3, 2]              
          6                -1  3    625152  models.common.C3                        [256, 256, 3]                 
          7                -1  1   1180672  models.common.Conv                      [256, 512, 3, 2]              
          8                -1  1   1182720  models.common.C3                        [512, 512, 1]                 
          9                -1  1    656896  models.common.SPPF                      [512, 512, 5]                 
         10                -1  1    131584  models.common.Conv                      [512, 256, 1, 1]              
         11                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
         12           [-1, 6]  1         0  models.common.Concat                    [1]                           
         13                -1  1    361984  models.common.C3                        [512, 256, 1, False]          
         14                -1  1     33024  models.common.Conv                      [256, 128, 1, 1]              
         15                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
         16           [-1, 4]  1         0  models.common.Concat                    [1]                           
         17                -1  1     90880  models.common.C3                        [256, 128, 1, False]          
         18                -1  1    147712  models.common.Conv                      [128, 128, 3, 2]              
         19          [-1, 14]  1         0  models.common.Concat                    [1]                           
         20                -1  1    296448  models.common.C3                        [256, 256, 1, False]          
         21                -1  1    590336  models.common.Conv                      [256, 256, 3, 2]              
         22          [-1, 10]  1         0  models.common.Concat                    [1]                           
         23                -1  1   1182720  models.common.C3                        [512, 512, 1, False]          
         24      [17, 20, 23]  1    615133  models.yolo.Segment                     [80, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], 32, 128, [128, 256, 512]]
        Model summary: 225 layers, 7621277 parameters, 7621277 gradients, 26.6 GFLOPs
        
        Transferred 367/367 items from yolov5s-seg.pt
        AMP: checks passed ✅
        optimizer: SGD(lr=0.01) with parameter groups 60 weight(decay=0.0), 63 weight(decay=0.0005), 63 bias
        albumentations: Blur(p=0.01, blur_limit=(3, 7)), MedianBlur(p=0.01, blur_limit=(3, 7)), ToGray(p=0.01), CLAHE(p=0.01, clip_limit=(1, 4.0), tile_grid_size=(8, 8))
        train: Scanning /content/datasets/coco128-seg/labels/train2017... 126 images, 2 backgrounds, 0 corrupt: 100% 128/128 [00:00<00:00, 1389.59it/s]
        train: New cache created: /content/datasets/coco128-seg/labels/train2017.cache
        train: Caching images (0.1GB ram): 100% 128/128 [00:00<00:00, 238.86it/s]
        val: Scanning /content/datasets/coco128-seg/labels/train2017.cache... 126 images, 2 backgrounds, 0 corrupt: 100% 128/128 [00:00<?, ?it/s]
        val: Caching images (0.1GB ram): 100% 128/128 [00:01<00:00, 98.90it/s]
        
        AutoAnchor: 4.27 anchors/target, 0.994 Best Possible Recall (BPR). Current anchors are a good fit to dataset ✅
        Plotting labels to runs/train-seg/exp/labels.jpg... 
        Image sizes 640 train, 640 val
        Using 2 dataloader workers
        Logging results to runs/train-seg/exp
        Starting training for 3 epochs...
        
              Epoch    GPU_mem   box_loss   seg_loss   obj_loss   cls_loss  Instances       Size
                0/2      4.92G     0.0417    0.04646    0.06066    0.02126        192        640: 100% 8/8 [00:08<00:00,  1.10s/it]
                         Class     Images  Instances      Box(P          R      mAP50  mAP50-95)     Mask(P          R      mAP50  mAP50-95): 100% 4/4 [00:02<00:00,  1.81it/s]
                           all        128        929      0.737      0.649      0.715      0.492      0.719      0.617      0.658      0.408
        
              Epoch    GPU_mem   box_loss   seg_loss   obj_loss   cls_loss  Instances       Size
                1/2      6.29G    0.04157    0.04503    0.05772    0.01777        208        640: 100% 8/8 [00:09<00:00,  1.21s/it]
                         Class     Images  Instances      Box(P          R      mAP50  mAP50-95)     Mask(P          R      mAP50  mAP50-95): 100% 4/4 [00:02<00:00,  1.87it/s]
                           all        128        929      0.756      0.674      0.738      0.506      0.725       0.64       0.68      0.422
        
              Epoch    GPU_mem   box_loss   seg_loss   obj_loss   cls_loss  Instances       Size
                2/2      6.29G     0.0425    0.04793    0.06784    0.01863        161        640: 100% 8/8 [00:03<00:00,  2.02it/s]
                         Class     Images  Instances      Box(P          R      mAP50  mAP50-95)     Mask(P          R      mAP50  mAP50-95): 100% 4/4 [00:02<00:00,  1.88it/s]
                           all        128        929      0.736      0.694      0.747      0.522      0.769      0.622      0.683      0.427
        
        3 epochs completed in 0.009 hours.
        Optimizer stripped from runs/train-seg/exp/weights/last.pt, 15.6MB
        Optimizer stripped from runs/train-seg/exp/weights/best.pt, 15.6MB
        
        Validating runs/train-seg/exp/weights/best.pt...
        Fusing layers... 
        Model summary: 165 layers, 7611485 parameters, 0 gradients, 26.4 GFLOPs
                         Class     Images  Instances      Box(P          R      mAP50  mAP50-95)     Mask(P          R      mAP50  mAP50-95): 100% 4/4 [00:06<00:00,  1.59s/it]
                           all        128        929      0.738      0.694      0.746      0.522      0.759      0.625      0.682      0.426
                        person        128        254      0.845      0.756      0.836       0.55      0.861      0.669      0.759      0.407
                       bicycle        128          6      0.475      0.333      0.549      0.341      0.711      0.333      0.526      0.322
                           car        128         46      0.612      0.565      0.539      0.257      0.555      0.435      0.477      0.171
                    motorcycle        128          5       0.73        0.8      0.752      0.571      0.747        0.8      0.752       0.42
                      airplane        128          6          1      0.943      0.995      0.732       0.92      0.833      0.839      0.555
                           bus        128          7      0.677      0.714      0.722      0.653      0.711      0.714      0.722      0.593
                         train        128          3          1      0.951      0.995      0.551          1      0.884      0.995      0.781
                         truck        128         12      0.555      0.417      0.457      0.285      0.624      0.417      0.397      0.277
                          boat        128          6      0.624        0.5      0.584      0.186          1      0.326      0.412      0.133
                 traffic light        128         14      0.513      0.302      0.411      0.247      0.435      0.214      0.376      0.251
                     stop sign        128          2      0.824          1      0.995      0.796      0.906          1      0.995      0.747
                         bench        128          9       0.75      0.667      0.763      0.367      0.724      0.585      0.698      0.209
                          bird        128         16      0.961          1      0.995      0.686      0.918      0.938       0.91      0.525
                           cat        128          4      0.771      0.857      0.945      0.752       0.76        0.8      0.945      0.728
                           dog        128          9      0.987      0.778      0.963      0.681          1      0.705       0.89      0.574
                         horse        128          2      0.703          1      0.995      0.697      0.759          1      0.995      0.249
                      elephant        128         17      0.916      0.882       0.93      0.691      0.811      0.765      0.829      0.537
                          bear        128          1      0.664          1      0.995      0.995      0.701          1      0.995      0.895
                         zebra        128          4      0.864          1      0.995      0.921      0.879          1      0.995      0.804
                       giraffe        128          9      0.883      0.889       0.94      0.683      0.845      0.778       0.78      0.463
                      backpack        128          6          1       0.59      0.701      0.372          1      0.474       0.52      0.252
                      umbrella        128         18      0.654      0.839      0.887       0.52      0.517      0.556      0.427      0.229
                       handbag        128         19       0.54      0.211      0.408      0.221      0.796      0.206      0.396      0.196
                           tie        128          7      0.864      0.857      0.857      0.577      0.925      0.857      0.857      0.534
                      suitcase        128          4      0.716          1      0.945      0.647      0.767          1      0.945      0.634
                       frisbee        128          5      0.708        0.8      0.761      0.643      0.737        0.8      0.761      0.501
                          skis        128          1      0.691          1      0.995      0.796      0.761          1      0.995      0.199
                     snowboard        128          7      0.918      0.857      0.904      0.604       0.32      0.286      0.235      0.137
                   sports ball        128          6      0.902      0.667      0.701      0.466      0.727        0.5      0.497      0.471
                          kite        128         10      0.586        0.4      0.511      0.231      0.663      0.394      0.417      0.139
                  baseball bat        128          4      0.359        0.5      0.401      0.169      0.631        0.5      0.526      0.133
                baseball glove        128          7          1      0.519       0.58      0.327      0.687      0.286      0.455      0.328
                    skateboard        128          5      0.729        0.8      0.862      0.631      0.599        0.6      0.604      0.379
                 tennis racket        128          7       0.57      0.714      0.645      0.448      0.608      0.714      0.645      0.412
                        bottle        128         18      0.469      0.393      0.537      0.357      0.661      0.389      0.543      0.349
                    wine glass        128         16      0.677      0.938      0.866      0.441       0.53      0.625       0.67      0.334
                           cup        128         36      0.777      0.722      0.812      0.466      0.725      0.583      0.762      0.467
                          fork        128          6      0.948      0.333      0.425       0.27      0.527      0.167       0.18      0.102
                         knife        128         16      0.757      0.587      0.669      0.458       0.79        0.5      0.552       0.34
                         spoon        128         22       0.74      0.364      0.559      0.269      0.925      0.364      0.513      0.213
                          bowl        128         28      0.766      0.714      0.725      0.559      0.803      0.584      0.665      0.353
                        banana        128          1      0.408          1      0.995      0.398      0.539          1      0.995      0.497
                      sandwich        128          2          1          0      0.695      0.536          1          0      0.498      0.448
                        orange        128          4      0.467          1      0.995      0.693      0.518          1      0.995      0.663
                      broccoli        128         11      0.462      0.455      0.383      0.259      0.548      0.455      0.384      0.256
                        carrot        128         24      0.631      0.875       0.77      0.533      0.757      0.909      0.853      0.499
                       hot dog        128          2      0.555          1      0.995      0.995      0.578          1      0.995      0.796
                         pizza        128          5       0.89        0.8      0.962      0.796          1      0.778      0.962      0.766
                         donut        128         14      0.695          1      0.893      0.772      0.704          1      0.893      0.696
                          cake        128          4      0.826          1      0.995       0.92      0.862          1      0.995      0.846
                         chair        128         35       0.53      0.571      0.613      0.336       0.67        0.6      0.538      0.271
                         couch        128          6      0.972      0.667      0.833      0.627          1       0.62      0.696      0.394
                  potted plant        128         14        0.7      0.857      0.883      0.552      0.836      0.857      0.883      0.473
                           bed        128          3      0.979      0.667       0.83      0.366          1          0       0.83      0.373
                  dining table        128         13      0.775      0.308      0.505      0.364      0.644      0.231       0.25     0.0804
                        toilet        128          2      0.836          1      0.995      0.846      0.887          1      0.995      0.797
                            tv        128          2        0.6          1      0.995      0.846      0.655          1      0.995      0.896
                        laptop        128          3      0.822      0.333      0.445      0.307          1          0      0.392       0.12
                         mouse        128          2          1          0          0          0          1          0          0          0
                        remote        128          8      0.745        0.5       0.62      0.459      0.821        0.5      0.624      0.449
                    cell phone        128          8      0.686      0.375      0.502      0.272      0.488       0.25       0.28      0.132
                     microwave        128          3      0.831          1      0.995      0.722      0.867          1      0.995      0.592
                          oven        128          5      0.439        0.4      0.435      0.294      0.823        0.6      0.645      0.418
                          sink        128          6      0.677        0.5      0.565      0.448      0.722        0.5       0.46      0.362
                  refrigerator        128          5      0.533        0.8      0.783      0.524      0.558        0.8      0.783      0.527
                          book        128         29      0.732      0.379      0.423      0.196       0.69      0.207       0.38      0.131
                         clock        128          9      0.889      0.778      0.917      0.677      0.908      0.778      0.875      0.604
                          vase        128          2      0.375          1      0.995      0.995      0.455          1      0.995      0.796
                      scissors        128          1          1          0     0.0166    0.00166          1          0          0          0
                    teddy bear        128         21      0.813      0.829      0.841      0.457      0.826      0.678      0.786      0.422
                    toothbrush        128          5      0.806          1      0.995      0.733      0.991          1      0.995      0.628
        Results saved to runs/train-seg/exp
        

        4. Visualize

        Comet Logging and Visualization 🌟 NEW

        Comet is now fully integrated with YOLOv5. Track and visualize model metrics in real time, save your hyperparameters, datasets, and model checkpoints, and visualize your model predictions with Comet Custom Panels! Comet makes sure you never lose track of your work and makes it easy to share results and collaborate across teams of all sizes!

        Getting started is easy:

        pip install comet_ml  # 1. install
        export COMET_API_KEY=<Your API Key>  # 2. paste API key
        python train.py --img 640 --epochs 3 --data coco128.yaml --weights yolov5s.pt  # 3. train
        

        To learn more about all of the supported Comet features for this integration, check out the Comet Tutorial. If you'd like to learn more about Comet, head over to our documentation. Get started by trying out the Comet Colab Notebook:
        Open In Colab


        Comet Dashboard

        ClearML Logging and Automation 🌟 NEW

        ClearML is completely integrated into YOLOv5 to track your experimentation, manage dataset versions and even remotely execute training runs. To enable ClearML (check cells above):

        You'll get all the great expected features from an experiment manager: live updates, model upload, experiment comparison etc. but ClearML also tracks uncommitted changes and installed packages for example. Thanks to that ClearML Tasks (which is what we call experiments) are also reproducible on different machines! With only 1 extra line, we can schedule a YOLOv5 training task on a queue to be executed by any number of ClearML Agents (workers).

        You can use ClearML Data to version your dataset and then pass it to YOLOv5 simply using its unique ID. This will help you keep track of your data without adding extra hassle. Explore the ClearML Tutorial for details!


        ClearML Experiment Management UI

        Local Logging

        Training results are automatically logged with Tensorboard and CSV loggers to runs/train, with a new experiment directory created for each new training as runs/train/exp2, runs/train/exp3, etc.

        This directory contains train and val statistics, mosaics, labels, predictions and augmentated mosaics, as well as metrics and charts including precision-recall (PR) curves and confusion matrices.

        Local logging results

        Environments

        YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

        Status

        YOLOv5 CI

        If this badge is green, all YOLOv5 GitHub Actions Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training (train.py), testing (val.py), inference (detect.py) and export (export.py) on macOS, Windows, and Ubuntu every 24 hours and on every commit.

        Appendix

        Additional content below.

        # YOLOv5 PyTorch HUB Inference (DetectionModels only)
        import torch
        
        model = torch.hub.load('ultralytics/yolov5', 'yolov5s-seg', force_reload=True, trust_repo=True)  # or yolov5n - yolov5x6 or custom
        im = 'https://ultralytics.com/images/zidane.jpg'  # file, Path, PIL.Image, OpenCV, nparray, list
        results = model(im)  # inference
        results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
        
      • 📄 val.py 23.6KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """
        Validate a trained YOLOv5 segment model on a segment dataset.
        
        Usage:
            $ bash data/scripts/get_coco.sh --val --segments  # download COCO-segments val split (1G, 5000 images)
            $ python segment/val.py --weights yolov5s-seg.pt --data coco.yaml --img 640  # validate COCO-segments
        
        Usage - formats:
            $ python segment/val.py --weights yolov5s-seg.pt                 # PyTorch
                                              yolov5s-seg.torchscript        # TorchScript
                                              yolov5s-seg.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                              yolov5s-seg_openvino_label     # OpenVINO
                                              yolov5s-seg.engine             # TensorRT
                                              yolov5s-seg.mlmodel            # CoreML (macOS-only)
                                              yolov5s-seg_saved_model        # TensorFlow SavedModel
                                              yolov5s-seg.pb                 # TensorFlow GraphDef
                                              yolov5s-seg.tflite             # TensorFlow Lite
                                              yolov5s-seg_edgetpu.tflite     # TensorFlow Edge TPU
                                              yolov5s-seg_paddle_model       # PaddlePaddle
        """
        
        import argparse
        import json
        import os
        import subprocess
        import sys
        from multiprocessing.pool import ThreadPool
        from pathlib import Path
        
        import numpy as np
        import torch
        from tqdm import tqdm
        
        FILE = Path(__file__).resolve()
        ROOT = FILE.parents[1]  # YOLOv5 root directory
        if str(ROOT) not in sys.path:
            sys.path.append(str(ROOT))  # add ROOT to PATH
        ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
        
        import torch.nn.functional as F
        
        from models.common import DetectMultiBackend
        from models.yolo import SegmentationModel
        from utils.callbacks import Callbacks
        from utils.general import (
            LOGGER,
            NUM_THREADS,
            TQDM_BAR_FORMAT,
            Profile,
            check_dataset,
            check_img_size,
            check_requirements,
            check_yaml,
            coco80_to_coco91_class,
            colorstr,
            increment_path,
            non_max_suppression,
            print_args,
            scale_boxes,
            xywh2xyxy,
            xyxy2xywh,
        )
        from utils.metrics import ConfusionMatrix, box_iou
        from utils.plots import output_to_target, plot_val_study
        from utils.segment.dataloaders import create_dataloader
        from utils.segment.general import mask_iou, process_mask, process_mask_native, scale_image
        from utils.segment.metrics import Metrics, ap_per_class_box_and_mask
        from utils.segment.plots import plot_images_and_masks
        from utils.torch_utils import de_parallel, select_device, smart_inference_mode
        
        
        def save_one_txt(predn, save_conf, shape, file):
            """Saves detection results in txt format; includes class, xywh (normalized), optionally confidence if `save_conf` is
            True.
            """
            gn = torch.tensor(shape)[[1, 0, 1, 0]]  # normalization gain whwh
            for *xyxy, conf, cls in predn.tolist():
                xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
                with open(file, "a") as f:
                    f.write(("%g " * len(line)).rstrip() % line + "\n")
        
        
        def save_one_json(predn, jdict, path, class_map, pred_masks):
            """
            Saves a JSON file with detection results including bounding boxes, category IDs, scores, and segmentation masks.
        
            Example JSON result: {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}.
            """
            from pycocotools.mask import encode
        
            def single_encode(x):
                rle = encode(np.asarray(x[:, :, None], order="F", dtype="uint8"))[0]
                rle["counts"] = rle["counts"].decode("utf-8")
                return rle
        
            image_id = int(path.stem) if path.stem.isnumeric() else path.stem
            box = xyxy2xywh(predn[:, :4])  # xywh
            box[:, :2] -= box[:, 2:] / 2  # xy center to top-left corner
            pred_masks = np.transpose(pred_masks, (2, 0, 1))
            with ThreadPool(NUM_THREADS) as pool:
                rles = pool.map(single_encode, pred_masks)
            for i, (p, b) in enumerate(zip(predn.tolist(), box.tolist())):
                jdict.append(
                    {
                        "image_id": image_id,
                        "category_id": class_map[int(p[5])],
                        "bbox": [round(x, 3) for x in b],
                        "score": round(p[4], 5),
                        "segmentation": rles[i],
                    }
                )
        
        
        def process_batch(detections, labels, iouv, pred_masks=None, gt_masks=None, overlap=False, masks=False):
            """
            Return correct prediction matrix
            Arguments:
                detections (array[N, 6]), x1, y1, x2, y2, conf, class
                labels (array[M, 5]), class, x1, y1, x2, y2
            Returns:
                correct (array[N, 10]), for 10 IoU levels
            """
            if masks:
                if overlap:
                    nl = len(labels)
                    index = torch.arange(nl, device=gt_masks.device).view(nl, 1, 1) + 1
                    gt_masks = gt_masks.repeat(nl, 1, 1)  # shape(1,640,640) -> (n,640,640)
                    gt_masks = torch.where(gt_masks == index, 1.0, 0.0)
                if gt_masks.shape[1:] != pred_masks.shape[1:]:
                    gt_masks = F.interpolate(gt_masks[None], pred_masks.shape[1:], mode="bilinear", align_corners=False)[0]
                    gt_masks = gt_masks.gt_(0.5)
                iou = mask_iou(gt_masks.view(gt_masks.shape[0], -1), pred_masks.view(pred_masks.shape[0], -1))
            else:  # boxes
                iou = box_iou(labels[:, 1:], detections[:, :4])
        
            correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
            correct_class = labels[:, 0:1] == detections[:, 5]
            for i in range(len(iouv)):
                x = torch.where((iou >= iouv[i]) & correct_class)  # IoU > threshold and classes match
                if x[0].shape[0]:
                    matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()  # [label, detect, iou]
                    if x[0].shape[0] > 1:
                        matches = matches[matches[:, 2].argsort()[::-1]]
                        matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
                        # matches = matches[matches[:, 2].argsort()[::-1]]
                        matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
                    correct[matches[:, 1].astype(int), i] = True
            return torch.tensor(correct, dtype=torch.bool, device=iouv.device)
        
        
        @smart_inference_mode()
        def run(
            data,
            weights=None,  # model.pt path(s)
            batch_size=32,  # batch size
            imgsz=640,  # inference size (pixels)
            conf_thres=0.001,  # confidence threshold
            iou_thres=0.6,  # NMS IoU threshold
            max_det=300,  # maximum detections per image
            task="val",  # train, val, test, speed or study
            device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
            workers=8,  # max dataloader workers (per RANK in DDP mode)
            single_cls=False,  # treat as single-class dataset
            augment=False,  # augmented inference
            verbose=False,  # verbose output
            save_txt=False,  # save results to *.txt
            save_hybrid=False,  # save label+prediction hybrid results to *.txt
            save_conf=False,  # save confidences in --save-txt labels
            save_json=False,  # save a COCO-JSON results file
            project=ROOT / "runs/val-seg",  # save to project/name
            name="exp",  # save to project/name
            exist_ok=False,  # existing project/name ok, do not increment
            half=True,  # use FP16 half-precision inference
            dnn=False,  # use OpenCV DNN for ONNX inference
            model=None,
            dataloader=None,
            save_dir=Path(""),
            plots=True,
            overlap=False,
            mask_downsample_ratio=1,
            compute_loss=None,
            callbacks=Callbacks(),
        ):
            if save_json:
                check_requirements("pycocotools>=2.0.6")
                process = process_mask_native  # more accurate
            else:
                process = process_mask  # faster
        
            # Initialize/load model and set device
            training = model is not None
            if training:  # called by train.py
                device, pt, jit, engine = next(model.parameters()).device, True, False, False  # get model device, PyTorch model
                half &= device.type != "cpu"  # half precision only supported on CUDA
                model.half() if half else model.float()
                nm = de_parallel(model).model[-1].nm  # number of masks
            else:  # called directly
                device = select_device(device, batch_size=batch_size)
        
                # Directories
                save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
                (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
        
                # Load model
                model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
                stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
                imgsz = check_img_size(imgsz, s=stride)  # check image size
                half = model.fp16  # FP16 supported on limited backends with CUDA
                nm = de_parallel(model).model.model[-1].nm if isinstance(model, SegmentationModel) else 32  # number of masks
                if engine:
                    batch_size = model.batch_size
                else:
                    device = model.device
                    if not (pt or jit):
                        batch_size = 1  # export.py models default to batch-size 1
                        LOGGER.info(f"Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models")
        
                # Data
                data = check_dataset(data)  # check
        
            # Configure
            model.eval()
            cuda = device.type != "cpu"
            is_coco = isinstance(data.get("val"), str) and data["val"].endswith(f"coco{os.sep}val2017.txt")  # COCO dataset
            nc = 1 if single_cls else int(data["nc"])  # number of classes
            iouv = torch.linspace(0.5, 0.95, 10, device=device)  # iou vector for mAP@0.5:0.95
            niou = iouv.numel()
        
            # Dataloader
            if not training:
                if pt and not single_cls:  # check --weights are trained on --data
                    ncm = model.model.nc
                    assert ncm == nc, (
                        f"{weights} ({ncm} classes) trained on different --data than what you passed ({nc} "
                        f"classes). Pass correct combination of --weights and --data that are trained together."
                    )
                model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz))  # warmup
                pad, rect = (0.0, False) if task == "speed" else (0.5, pt)  # square inference for benchmarks
                task = task if task in ("train", "val", "test") else "val"  # path to train/val/test images
                dataloader = create_dataloader(
                    data[task],
                    imgsz,
                    batch_size,
                    stride,
                    single_cls,
                    pad=pad,
                    rect=rect,
                    workers=workers,
                    prefix=colorstr(f"{task}: "),
                    overlap_mask=overlap,
                    mask_downsample_ratio=mask_downsample_ratio,
                )[0]
        
            seen = 0
            confusion_matrix = ConfusionMatrix(nc=nc)
            names = model.names if hasattr(model, "names") else model.module.names  # get class names
            if isinstance(names, (list, tuple)):  # old format
                names = dict(enumerate(names))
            class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
            s = ("%22s" + "%11s" * 10) % (
                "Class",
                "Images",
                "Instances",
                "Box(P",
                "R",
                "mAP50",
                "mAP50-95)",
                "Mask(P",
                "R",
                "mAP50",
                "mAP50-95)",
            )
            dt = Profile(device=device), Profile(device=device), Profile(device=device)
            metrics = Metrics()
            loss = torch.zeros(4, device=device)
            jdict, stats = [], []
            # callbacks.run('on_val_start')
            pbar = tqdm(dataloader, desc=s, bar_format=TQDM_BAR_FORMAT)  # progress bar
            for batch_i, (im, targets, paths, shapes, masks) in enumerate(pbar):
                # callbacks.run('on_val_batch_start')
                with dt[0]:
                    if cuda:
                        im = im.to(device, non_blocking=True)
                        targets = targets.to(device)
                        masks = masks.to(device)
                    masks = masks.float()
                    im = im.half() if half else im.float()  # uint8 to fp16/32
                    im /= 255  # 0 - 255 to 0.0 - 1.0
                    nb, _, height, width = im.shape  # batch size, channels, height, width
        
                # Inference
                with dt[1]:
                    preds, protos, train_out = model(im) if compute_loss else (*model(im, augment=augment)[:2], None)
        
                # Loss
                if compute_loss:
                    loss += compute_loss((train_out, protos), targets, masks)[1]  # box, obj, cls
        
                # NMS
                targets[:, 2:] *= torch.tensor((width, height, width, height), device=device)  # to pixels
                lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else []  # for autolabelling
                with dt[2]:
                    preds = non_max_suppression(
                        preds, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls, max_det=max_det, nm=nm
                    )
        
                # Metrics
                plot_masks = []  # masks for plotting
                for si, (pred, proto) in enumerate(zip(preds, protos)):
                    labels = targets[targets[:, 0] == si, 1:]
                    nl, npr = labels.shape[0], pred.shape[0]  # number of labels, predictions
                    path, shape = Path(paths[si]), shapes[si][0]
                    correct_masks = torch.zeros(npr, niou, dtype=torch.bool, device=device)  # init
                    correct_bboxes = torch.zeros(npr, niou, dtype=torch.bool, device=device)  # init
                    seen += 1
        
                    if npr == 0:
                        if nl:
                            stats.append((correct_masks, correct_bboxes, *torch.zeros((2, 0), device=device), labels[:, 0]))
                            if plots:
                                confusion_matrix.process_batch(detections=None, labels=labels[:, 0])
                        continue
        
                    # Masks
                    midx = [si] if overlap else targets[:, 0] == si
                    gt_masks = masks[midx]
                    pred_masks = process(proto, pred[:, 6:], pred[:, :4], shape=im[si].shape[1:])
        
                    # Predictions
                    if single_cls:
                        pred[:, 5] = 0
                    predn = pred.clone()
                    scale_boxes(im[si].shape[1:], predn[:, :4], shape, shapes[si][1])  # native-space pred
        
                    # Evaluate
                    if nl:
                        tbox = xywh2xyxy(labels[:, 1:5])  # target boxes
                        scale_boxes(im[si].shape[1:], tbox, shape, shapes[si][1])  # native-space labels
                        labelsn = torch.cat((labels[:, 0:1], tbox), 1)  # native-space labels
                        correct_bboxes = process_batch(predn, labelsn, iouv)
                        correct_masks = process_batch(predn, labelsn, iouv, pred_masks, gt_masks, overlap=overlap, masks=True)
                        if plots:
                            confusion_matrix.process_batch(predn, labelsn)
                    stats.append((correct_masks, correct_bboxes, pred[:, 4], pred[:, 5], labels[:, 0]))  # (conf, pcls, tcls)
        
                    pred_masks = torch.as_tensor(pred_masks, dtype=torch.uint8)
                    if plots and batch_i < 3:
                        plot_masks.append(pred_masks[:15])  # filter top 15 to plot
        
                    # Save/log
                    if save_txt:
                        save_one_txt(predn, save_conf, shape, file=save_dir / "labels" / f"{path.stem}.txt")
                    if save_json:
                        pred_masks = scale_image(
                            im[si].shape[1:], pred_masks.permute(1, 2, 0).contiguous().cpu().numpy(), shape, shapes[si][1]
                        )
                        save_one_json(predn, jdict, path, class_map, pred_masks)  # append to COCO-JSON dictionary
                    # callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
        
                # Plot images
                if plots and batch_i < 3:
                    if len(plot_masks):
                        plot_masks = torch.cat(plot_masks, dim=0)
                    plot_images_and_masks(im, targets, masks, paths, save_dir / f"val_batch{batch_i}_labels.jpg", names)
                    plot_images_and_masks(
                        im,
                        output_to_target(preds, max_det=15),
                        plot_masks,
                        paths,
                        save_dir / f"val_batch{batch_i}_pred.jpg",
                        names,
                    )  # pred
        
                # callbacks.run('on_val_batch_end')
        
            # Compute metrics
            stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)]  # to numpy
            if len(stats) and stats[0].any():
                results = ap_per_class_box_and_mask(*stats, plot=plots, save_dir=save_dir, names=names)
                metrics.update(results)
            nt = np.bincount(stats[4].astype(int), minlength=nc)  # number of targets per class
        
            # Print results
            pf = "%22s" + "%11i" * 2 + "%11.3g" * 8  # print format
            LOGGER.info(pf % ("all", seen, nt.sum(), *metrics.mean_results()))
            if nt.sum() == 0:
                LOGGER.warning(f"WARNING ⚠️ no labels found in {task} set, can not compute metrics without labels")
        
            # Print results per class
            if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
                for i, c in enumerate(metrics.ap_class_index):
                    LOGGER.info(pf % (names[c], seen, nt[c], *metrics.class_result(i)))
        
            # Print speeds
            t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
            if not training:
                shape = (batch_size, 3, imgsz, imgsz)
                LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}" % t)
        
            # Plots
            if plots:
                confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
            # callbacks.run('on_val_end')
        
            mp_bbox, mr_bbox, map50_bbox, map_bbox, mp_mask, mr_mask, map50_mask, map_mask = metrics.mean_results()
        
            # Save JSON
            if save_json and len(jdict):
                w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else ""  # weights
                anno_json = str(Path("../datasets/coco/annotations/instances_val2017.json"))  # annotations
                pred_json = str(save_dir / f"{w}_predictions.json")  # predictions
                LOGGER.info(f"\nEvaluating pycocotools mAP... saving {pred_json}...")
                with open(pred_json, "w") as f:
                    json.dump(jdict, f)
        
                try:  # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
                    from pycocotools.coco import COCO
                    from pycocotools.cocoeval import COCOeval
        
                    anno = COCO(anno_json)  # init annotations api
                    pred = anno.loadRes(pred_json)  # init predictions api
                    results = []
                    for eval in COCOeval(anno, pred, "bbox"), COCOeval(anno, pred, "segm"):
                        if is_coco:
                            eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.im_files]  # img ID to evaluate
                        eval.evaluate()
                        eval.accumulate()
                        eval.summarize()
                        results.extend(eval.stats[:2])  # update results (mAP@0.5:0.95, mAP@0.5)
                    map_bbox, map50_bbox, map_mask, map50_mask = results
                except Exception as e:
                    LOGGER.info(f"pycocotools unable to run: {e}")
        
            # Return results
            model.float()  # for training
            if not training:
                s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
                LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
            final_metric = mp_bbox, mr_bbox, map50_bbox, map_bbox, mp_mask, mr_mask, map50_mask, map_mask
            return (*final_metric, *(loss.cpu() / len(dataloader)).tolist()), metrics.get_maps(nc), t
        
        
        def parse_opt():
            """Parses command line arguments for configuring YOLOv5 options like dataset path, weights, batch size, and
            inference settings.
            """
            parser = argparse.ArgumentParser()
            parser.add_argument("--data", type=str, default=ROOT / "data/coco128-seg.yaml", help="dataset.yaml path")
            parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s-seg.pt", help="model path(s)")
            parser.add_argument("--batch-size", type=int, default=32, help="batch size")
            parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="inference size (pixels)")
            parser.add_argument("--conf-thres", type=float, default=0.001, help="confidence threshold")
            parser.add_argument("--iou-thres", type=float, default=0.6, help="NMS IoU threshold")
            parser.add_argument("--max-det", type=int, default=300, help="maximum detections per image")
            parser.add_argument("--task", default="val", help="train, val, test, speed or study")
            parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
            parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
            parser.add_argument("--single-cls", action="store_true", help="treat as single-class dataset")
            parser.add_argument("--augment", action="store_true", help="augmented inference")
            parser.add_argument("--verbose", action="store_true", help="report mAP by class")
            parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
            parser.add_argument("--save-hybrid", action="store_true", help="save label+prediction hybrid results to *.txt")
            parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels")
            parser.add_argument("--save-json", action="store_true", help="save a COCO-JSON results file")
            parser.add_argument("--project", default=ROOT / "runs/val-seg", help="save results to project/name")
            parser.add_argument("--name", default="exp", help="save to project/name")
            parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
            parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
            parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
            opt = parser.parse_args()
            opt.data = check_yaml(opt.data)  # check YAML
            # opt.save_json |= opt.data.endswith('coco.yaml')
            opt.save_txt |= opt.save_hybrid
            print_args(vars(opt))
            return opt
        
        
        def main(opt):
            """Executes YOLOv5 tasks including training, validation, testing, speed, and study with configurable options."""
            check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
        
            if opt.task in ("train", "val", "test"):  # run normally
                if opt.conf_thres > 0.001:  # https://github.com/ultralytics/yolov5/issues/1466
                    LOGGER.warning(f"WARNING ⚠️ confidence threshold {opt.conf_thres} > 0.001 produces invalid results")
                if opt.save_hybrid:
                    LOGGER.warning("WARNING ⚠️ --save-hybrid returns high mAP from hybrid labels, not from predictions alone")
                run(**vars(opt))
        
            else:
                weights = opt.weights if isinstance(opt.weights, list) else [opt.weights]
                opt.half = torch.cuda.is_available() and opt.device != "cpu"  # FP16 for fastest results
                if opt.task == "speed":  # speed benchmarks
                    # python val.py --task speed --data coco.yaml --batch 1 --weights yolov5n.pt yolov5s.pt...
                    opt.conf_thres, opt.iou_thres, opt.save_json = 0.25, 0.45, False
                    for opt.weights in weights:
                        run(**vars(opt), plots=False)
        
                elif opt.task == "study":  # speed vs mAP benchmarks
                    # python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n.pt yolov5s.pt...
                    for opt.weights in weights:
                        f = f"study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt"  # filename to save to
                        x, y = list(range(256, 1536 + 128, 128)), []  # x axis (image sizes), y axis
                        for opt.imgsz in x:  # img-size
                            LOGGER.info(f"\nRunning {f} --imgsz {opt.imgsz}...")
                            r, _, t = run(**vars(opt), plots=False)
                            y.append(r + t)  # results and times
                        np.savetxt(f, y, fmt="%10.4g")  # save
                    subprocess.run(["zip", "-r", "study.zip", "study_*.txt"])
                    plot_val_study(x=x)  # plot
                else:
                    raise NotImplementedError(f'--task {opt.task} not in ("train", "val", "test", "speed", "study")')
        
        
        if __name__ == "__main__":
            opt = parse_opt()
            main(opt)
        
    • 📁 utils/
      • 📁 aws/
        • 📄 __init__.py 0B
          
          
        • 📄 mime.sh 780B
          # AWS EC2 instance startup 'MIME' script https://aws.amazon.com/premiumsupport/knowledge-center/execute-user-data-ec2/
          # This script will run on every instance restart, not only on first start
          # --- DO NOT COPY ABOVE COMMENTS WHEN PASTING INTO USERDATA ---
          
          Content-Type: multipart/mixed; boundary="//"
          MIME-Version: 1.0
          
          --//
          Content-Type: text/cloud-config; charset="us-ascii"
          MIME-Version: 1.0
          Content-Transfer-Encoding: 7bit
          Content-Disposition: attachment; filename="cloud-config.txt"
          
          #cloud-config
          cloud_final_modules:
          - [scripts-user, always]
          
          --//
          Content-Type: text/x-shellscript; charset="us-ascii"
          MIME-Version: 1.0
          Content-Transfer-Encoding: 7bit
          Content-Disposition: attachment; filename="userdata.txt"
          
          #!/bin/bash
          # --- paste contents of userdata.sh here ---
          --//
          
        • 📄 resume.py 1.2KB
          # Resume all interrupted trainings in yolov5/ dir including DDP trainings
          # Usage: $ python utils/aws/resume.py
          
          import os
          import sys
          from pathlib import Path
          
          import torch
          import yaml
          
          FILE = Path(__file__).resolve()
          ROOT = FILE.parents[2]  # YOLOv5 root directory
          if str(ROOT) not in sys.path:
              sys.path.append(str(ROOT))  # add ROOT to PATH
          
          port = 0  # --master_port
          path = Path("").resolve()
          for last in path.rglob("*/**/last.pt"):
              ckpt = torch.load(last)
              if ckpt["optimizer"] is None:
                  continue
          
              # Load opt.yaml
              with open(last.parent.parent / "opt.yaml", errors="ignore") as f:
                  opt = yaml.safe_load(f)
          
              # Get device count
              d = opt["device"].split(",")  # devices
              nd = len(d)  # number of devices
              ddp = nd > 1 or (nd == 0 and torch.cuda.device_count() > 1)  # distributed data parallel
          
              if ddp:  # multi-GPU
                  port += 1
                  cmd = f"python -m torch.distributed.run --nproc_per_node {nd} --master_port {port} train.py --resume {last}"
              else:  # single-GPU
                  cmd = f"python train.py --resume {last}"
          
              cmd += " > /dev/null 2>&1 &"  # redirect output to dev/null and run in daemon thread
              print(cmd)
              os.system(cmd)
          
        • 📄 userdata.sh 1.2KB
          #!/bin/bash
          # AWS EC2 instance startup script https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/user-data.html
          # This script will run only once on first instance start (for a re-start script see mime.sh)
          # /home/ubuntu (ubuntu) or /home/ec2-user (amazon-linux) is working dir
          # Use >300 GB SSD
          
          cd home/ubuntu
          if [ ! -d yolov5 ]; then
            echo "Running first-time script." # install dependencies, download COCO, pull Docker
            git clone https://github.com/ultralytics/yolov5 -b master && sudo chmod -R 777 yolov5
            cd yolov5
            bash data/scripts/get_coco.sh && echo "COCO done." &
            sudo docker pull ultralytics/yolov5:latest && echo "Docker done." &
            python -m pip install --upgrade pip && pip install -r requirements.txt && python detect.py && echo "Requirements done." &
            wait && echo "All tasks done." # finish background tasks
          else
            echo "Running re-start script." # resume interrupted runs
            i=0
            list=$(sudo docker ps -qa) # container list i.e. $'one\ntwo\nthree\nfour'
            while IFS= read -r id; do
              ((i++))
              echo "restarting container $i: $id"
              sudo docker start $id
              # sudo docker exec -it $id python train.py --resume # single-GPU
              sudo docker exec -d $id python utils/aws/resume.py # multi-scenario
            done <<<"$list"
          fi
          
      • 📁 docker/
        • 📄 Dockerfile 2.5KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          # Builds ultralytics/yolov5:latest image on DockerHub https://hub.docker.com/r/ultralytics/yolov5
          # Image is CUDA-optimized for YOLOv5 single/multi-GPU training and inference
          
          # Start FROM PyTorch image https://hub.docker.com/r/pytorch/pytorch
          FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime
          
          # Downloads to user config dir
          ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Arial.Unicode.ttf /root/.config/Ultralytics/
          
          # Install linux packages
          ENV DEBIAN_FRONTEND noninteractive
          RUN apt update
          RUN TZ=Etc/UTC apt install -y tzdata
          RUN apt install --no-install-recommends -y gcc git zip curl htop libgl1 libglib2.0-0 libpython3-dev gnupg
          # RUN alias python=python3
          
          # Security updates
          # https://security.snyk.io/vuln/SNYK-UBUNTU1804-OPENSSL-3314796
          RUN apt upgrade --no-install-recommends -y openssl
          
          # Create working directory
          RUN rm -rf /usr/src/app && mkdir -p /usr/src/app
          WORKDIR /usr/src/app
          
          # Copy contents
          COPY . /usr/src/app
          
          # Install pip packages
          COPY requirements.txt .
          RUN python3 -m pip install --upgrade pip wheel
          RUN pip install --no-cache -r requirements.txt albumentations comet gsutil notebook \
              coremltools onnx onnx-simplifier onnxruntime 'openvino-dev>=2023.0'
              # tensorflow tensorflowjs \
          
          # Set environment variables
          ENV OMP_NUM_THREADS=1
          
          # Cleanup
          ENV DEBIAN_FRONTEND teletype
          
          
          # Usage Examples -------------------------------------------------------------------------------------------------------
          
          # Build and Push
          # t=ultralytics/yolov5:latest && sudo docker build -f utils/docker/Dockerfile -t $t . && sudo docker push $t
          
          # Pull and Run
          # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all $t
          
          # Pull and Run with local directory access
          # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all -v "$(pwd)"/datasets:/usr/src/datasets $t
          
          # Kill all
          # sudo docker kill $(sudo docker ps -q)
          
          # Kill all image-based
          # sudo docker kill $(sudo docker ps -qa --filter ancestor=ultralytics/yolov5:latest)
          
          # DockerHub tag update
          # t=ultralytics/yolov5:latest tnew=ultralytics/yolov5:v6.2 && sudo docker pull $t && sudo docker tag $t $tnew && sudo docker push $tnew
          
          # Clean up
          # sudo docker system prune -a --volumes
          
          # Update Ubuntu drivers
          # https://www.maketecheasier.com/install-nvidia-drivers-ubuntu/
          
          # DDP test
          # python -m torch.distributed.run --nproc_per_node 2 --master_port 1 train.py --epochs 3
          
          # GCP VM from Image
          # docker.io/ultralytics/yolov5:latest
          
        • ⚙️ Dockerfile-arm64 1.5KB
        • ⚙️ Dockerfile-cpu 1.8KB
      • 📁 flask_rest_api/
        • 📄 README.md 1.7KB
          # Flask REST API
          
          [REST](https://en.wikipedia.org/wiki/Representational_state_transfer) [API](https://en.wikipedia.org/wiki/API)s are commonly used to expose Machine Learning (ML) models to other services. This folder contains an example REST API created using Flask to expose the YOLOv5s model from [PyTorch Hub](https://pytorch.org/hub/ultralytics_yolov5/).
          
          ## Requirements
          
          [Flask](https://palletsprojects.com/p/flask/) is required. Install with:
          
          ```shell
          $ pip install Flask
          ```
          
          ## Run
          
          After Flask installation run:
          
          ```shell
          $ python3 restapi.py --port 5000
          ```
          
          Then use [curl](https://curl.se/) to perform a request:
          
          ```shell
          $ curl -X POST -F image=@zidane.jpg 'http://localhost:5000/v1/object-detection/yolov5s'
          ```
          
          The model inference results are returned as a JSON response:
          
          ```json
          [
            {
              "class": 0,
              "confidence": 0.8900438547,
              "height": 0.9318675399,
              "name": "person",
              "width": 0.3264600933,
              "xcenter": 0.7438579798,
              "ycenter": 0.5207948685
            },
            {
              "class": 0,
              "confidence": 0.8440024257,
              "height": 0.7155083418,
              "name": "person",
              "width": 0.6546785235,
              "xcenter": 0.427829951,
              "ycenter": 0.6334488392
            },
            {
              "class": 27,
              "confidence": 0.3771208823,
              "height": 0.3902671337,
              "name": "tie",
              "width": 0.0696444362,
              "xcenter": 0.3675483763,
              "ycenter": 0.7991207838
            },
            {
              "class": 27,
              "confidence": 0.3527112305,
              "height": 0.1540903747,
              "name": "tie",
              "width": 0.0336618312,
              "xcenter": 0.7814827561,
              "ycenter": 0.5065554976
            }
          ]
          ```
          
          An example python script to perform inference using [requests](https://docs.python-requests.org/en/master/) is given in `example_request.py`
          
        • 📄 example_request.py 368B
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          """Perform test request."""
          
          import pprint
          
          import requests
          
          DETECTION_URL = "http://localhost:5000/v1/object-detection/yolov5s"
          IMAGE = "zidane.jpg"
          
          # Read image
          with open(IMAGE, "rb") as f:
              image_data = f.read()
          
          response = requests.post(DETECTION_URL, files={"image": image_data}).json()
          
          pprint.pprint(response)
          
        • 📄 restapi.py 1.5KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          """Run a Flask REST API exposing one or more YOLOv5s models."""
          
          import argparse
          import io
          
          import torch
          from flask import Flask, request
          from PIL import Image
          
          app = Flask(__name__)
          models = {}
          
          DETECTION_URL = "/v1/object-detection/<model>"
          
          
          @app.route(DETECTION_URL, methods=["POST"])
          def predict(model):
              """Predict and return object detections in JSON format given an image and model name via a Flask REST API POST
              request.
              """
              if request.method != "POST":
                  return
          
              if request.files.get("image"):
                  # Method 1
                  # with request.files["image"] as f:
                  #     im = Image.open(io.BytesIO(f.read()))
          
                  # Method 2
                  im_file = request.files["image"]
                  im_bytes = im_file.read()
                  im = Image.open(io.BytesIO(im_bytes))
          
                  if model in models:
                      results = models[model](im, size=640)  # reduce size=320 for faster inference
                      return results.pandas().xyxy[0].to_json(orient="records")
          
          
          if __name__ == "__main__":
              parser = argparse.ArgumentParser(description="Flask API exposing YOLOv5 model")
              parser.add_argument("--port", default=5000, type=int, help="port number")
              parser.add_argument("--model", nargs="+", default=["yolov5s"], help="model(s) to run, i.e. --model yolov5n yolov5s")
              opt = parser.parse_args()
          
              for m in opt.model:
                  models[m] = torch.hub.load("ultralytics/yolov5", m, force_reload=True, skip_validation=True)
          
              app.run(host="0.0.0.0", port=opt.port)  # debug=True causes Restarting with stat
          
      • 📁 google_app_engine/
        • 📄 Dockerfile 821B
          FROM gcr.io/google-appengine/python
          
          # Create a virtualenv for dependencies. This isolates these packages from
          # system-level packages.
          # Use -p python3 or -p python3.7 to select python version. Default is version 2.
          RUN virtualenv /env -p python3
          
          # Setting these environment variables are the same as running
          # source /env/bin/activate.
          ENV VIRTUAL_ENV /env
          ENV PATH /env/bin:$PATH
          
          RUN apt-get update && apt-get install -y python-opencv
          
          # Copy the application's requirements.txt and run pip to install all
          # dependencies into the virtualenv.
          ADD requirements.txt /app/requirements.txt
          RUN pip install -r /app/requirements.txt
          
          # Add the application source code.
          ADD . /app
          
          # Run a WSGI server to serve the application. gunicorn must be declared as
          # a dependency in requirements.txt.
          CMD gunicorn -b :$PORT main:app
          
        • 📄 additional_requirements.txt 186B
          # add these requirements in your app on top of the existing ones
          pip==23.3
          Flask==2.3.2
          gunicorn==22.0.0
          werkzeug>=3.0.1 # not directly required, pinned by Snyk to avoid a vulnerability
          
        • 📄 app.yaml 174B
          runtime: custom
          env: flex
          
          service: yolov5app
          
          liveness_check:
            initial_delay_sec: 600
          
          manual_scaling:
            instances: 1
          resources:
            cpu: 1
            memory_gb: 4
            disk_size_gb: 20
          
      • 📁 loggers/
        • 📄 __init__.py 19.6KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          """Logging utils."""
          
          import json
          import os
          import warnings
          from pathlib import Path
          
          import pkg_resources as pkg
          import torch
          
          from utils.general import LOGGER, colorstr, cv2
          from utils.loggers.clearml.clearml_utils import ClearmlLogger
          from utils.loggers.wandb.wandb_utils import WandbLogger
          from utils.plots import plot_images, plot_labels, plot_results
          from utils.torch_utils import de_parallel
          
          LOGGERS = ("csv", "tb", "wandb", "clearml", "comet")  # *.csv, TensorBoard, Weights & Biases, ClearML
          RANK = int(os.getenv("RANK", -1))
          
          try:
              from torch.utils.tensorboard import SummaryWriter
          except ImportError:
              SummaryWriter = lambda *args: None  # None = SummaryWriter(str)
          
          try:
              import wandb
          
              assert hasattr(wandb, "__version__")  # verify package import not local dir
              if pkg.parse_version(wandb.__version__) >= pkg.parse_version("0.12.2") and RANK in {0, -1}:
                  try:
                      wandb_login_success = wandb.login(timeout=30)
                  except wandb.errors.UsageError:  # known non-TTY terminal issue
                      wandb_login_success = False
                  if not wandb_login_success:
                      wandb = None
          except (ImportError, AssertionError):
              wandb = None
          
          try:
              import clearml
          
              assert hasattr(clearml, "__version__")  # verify package import not local dir
          except (ImportError, AssertionError):
              clearml = None
          
          try:
              if RANK in {0, -1}:
                  import comet_ml
          
                  assert hasattr(comet_ml, "__version__")  # verify package import not local dir
                  from utils.loggers.comet import CometLogger
          
              else:
                  comet_ml = None
          except (ImportError, AssertionError):
              comet_ml = None
          
          
          def _json_default(value):
              """
              Format `value` for JSON serialization (e.g. unwrap tensors).
          
              Fall back to strings.
              """
              if isinstance(value, torch.Tensor):
                  try:
                      value = value.item()
                  except ValueError:  # "only one element tensors can be converted to Python scalars"
                      pass
              return value if isinstance(value, float) else str(value)
          
          
          class Loggers:
              # YOLOv5 Loggers class
              def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, include=LOGGERS):
                  """Initializes loggers for YOLOv5 training and validation metrics, paths, and options."""
                  self.save_dir = save_dir
                  self.weights = weights
                  self.opt = opt
                  self.hyp = hyp
                  self.plots = not opt.noplots  # plot results
                  self.logger = logger  # for printing results to console
                  self.include = include
                  self.keys = [
                      "train/box_loss",
                      "train/obj_loss",
                      "train/cls_loss",  # train loss
                      "metrics/precision",
                      "metrics/recall",
                      "metrics/mAP_0.5",
                      "metrics/mAP_0.5:0.95",  # metrics
                      "val/box_loss",
                      "val/obj_loss",
                      "val/cls_loss",  # val loss
                      "x/lr0",
                      "x/lr1",
                      "x/lr2",
                  ]  # params
                  self.best_keys = ["best/epoch", "best/precision", "best/recall", "best/mAP_0.5", "best/mAP_0.5:0.95"]
                  for k in LOGGERS:
                      setattr(self, k, None)  # init empty logger dictionary
                  self.csv = True  # always log to csv
                  self.ndjson_console = "ndjson_console" in self.include  # log ndjson to console
                  self.ndjson_file = "ndjson_file" in self.include  # log ndjson to file
          
                  # Messages
                  if not comet_ml:
                      prefix = colorstr("Comet: ")
                      s = f"{prefix}run 'pip install comet_ml' to automatically track and visualize YOLOv5 🚀 runs in Comet"
                      self.logger.info(s)
                  # TensorBoard
                  s = self.save_dir
                  if "tb" in self.include and not self.opt.evolve:
                      prefix = colorstr("TensorBoard: ")
                      self.logger.info(f"{prefix}Start with 'tensorboard --logdir {s.parent}', view at http://localhost:6006/")
                      self.tb = SummaryWriter(str(s))
          
                  # W&B
                  if wandb and "wandb" in self.include:
                      self.opt.hyp = self.hyp  # add hyperparameters
                      self.wandb = WandbLogger(self.opt)
                  else:
                      self.wandb = None
          
                  # ClearML
                  if clearml and "clearml" in self.include:
                      try:
                          self.clearml = ClearmlLogger(self.opt, self.hyp)
                      except Exception:
                          self.clearml = None
                          prefix = colorstr("ClearML: ")
                          LOGGER.warning(
                              f"{prefix}WARNING ⚠️ ClearML is installed but not configured, skipping ClearML logging."
                              f" See https://docs.ultralytics.com/yolov5/tutorials/clearml_logging_integration#readme"
                          )
          
                  else:
                      self.clearml = None
          
                  # Comet
                  if comet_ml and "comet" in self.include:
                      if isinstance(self.opt.resume, str) and self.opt.resume.startswith("comet://"):
                          run_id = self.opt.resume.split("/")[-1]
                          self.comet_logger = CometLogger(self.opt, self.hyp, run_id=run_id)
          
                      else:
                          self.comet_logger = CometLogger(self.opt, self.hyp)
          
                  else:
                      self.comet_logger = None
          
              @property
              def remote_dataset(self):
                  """Fetches dataset dictionary from remote logging services like ClearML, Weights & Biases, or Comet ML."""
                  data_dict = None
                  if self.clearml:
                      data_dict = self.clearml.data_dict
                  if self.wandb:
                      data_dict = self.wandb.data_dict
                  if self.comet_logger:
                      data_dict = self.comet_logger.data_dict
          
                  return data_dict
          
              def on_train_start(self):
                  """Initializes the training process for Comet ML logger if it's configured."""
                  if self.comet_logger:
                      self.comet_logger.on_train_start()
          
              def on_pretrain_routine_start(self):
                  """Invokes pre-training routine start hook for Comet ML logger if available."""
                  if self.comet_logger:
                      self.comet_logger.on_pretrain_routine_start()
          
              def on_pretrain_routine_end(self, labels, names):
                  """Callback that runs at the end of pre-training routine, logging label plots if enabled."""
                  if self.plots:
                      plot_labels(labels, names, self.save_dir)
                      paths = self.save_dir.glob("*labels*.jpg")  # training labels
                      if self.wandb:
                          self.wandb.log({"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]})
                      if self.comet_logger:
                          self.comet_logger.on_pretrain_routine_end(paths)
                      if self.clearml:
                          for path in paths:
                              self.clearml.log_plot(title=path.stem, plot_path=path)
          
              def on_train_batch_end(self, model, ni, imgs, targets, paths, vals):
                  """Logs training batch end events, plots images, and updates external loggers with batch-end data."""
                  log_dict = dict(zip(self.keys[:3], vals))
                  # Callback runs on train batch end
                  # ni: number integrated batches (since train start)
                  if self.plots:
                      if ni < 3:
                          f = self.save_dir / f"train_batch{ni}.jpg"  # filename
                          plot_images(imgs, targets, paths, f)
                          if ni == 0 and self.tb and not self.opt.sync_bn:
                              log_tensorboard_graph(self.tb, model, imgsz=(self.opt.imgsz, self.opt.imgsz))
                      if ni == 10 and (self.wandb or self.clearml):
                          files = sorted(self.save_dir.glob("train*.jpg"))
                          if self.wandb:
                              self.wandb.log({"Mosaics": [wandb.Image(str(f), caption=f.name) for f in files if f.exists()]})
                          if self.clearml:
                              self.clearml.log_debug_samples(files, title="Mosaics")
          
                  if self.comet_logger:
                      self.comet_logger.on_train_batch_end(log_dict, step=ni)
          
              def on_train_epoch_end(self, epoch):
                  """Callback that updates the current epoch in Weights & Biases at the end of a training epoch."""
                  if self.wandb:
                      self.wandb.current_epoch = epoch + 1
          
                  if self.comet_logger:
                      self.comet_logger.on_train_epoch_end(epoch)
          
              def on_val_start(self):
                  """Callback that signals the start of a validation phase to the Comet logger."""
                  if self.comet_logger:
                      self.comet_logger.on_val_start()
          
              def on_val_image_end(self, pred, predn, path, names, im):
                  """Callback that logs a validation image and its predictions to WandB or ClearML."""
                  if self.wandb:
                      self.wandb.val_one_image(pred, predn, path, names, im)
                  if self.clearml:
                      self.clearml.log_image_with_boxes(path, pred, names, im)
          
              def on_val_batch_end(self, batch_i, im, targets, paths, shapes, out):
                  """Logs validation batch results to Comet ML during training at the end of each validation batch."""
                  if self.comet_logger:
                      self.comet_logger.on_val_batch_end(batch_i, im, targets, paths, shapes, out)
          
              def on_val_end(self, nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix):
                  """Logs validation results to WandB or ClearML at the end of the validation process."""
                  if self.wandb or self.clearml:
                      files = sorted(self.save_dir.glob("val*.jpg"))
                  if self.wandb:
                      self.wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in files]})
                  if self.clearml:
                      self.clearml.log_debug_samples(files, title="Validation")
          
                  if self.comet_logger:
                      self.comet_logger.on_val_end(nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix)
          
              def on_fit_epoch_end(self, vals, epoch, best_fitness, fi):
                  """Callback that logs metrics and saves them to CSV or NDJSON at the end of each fit (train+val) epoch."""
                  x = dict(zip(self.keys, vals))
                  if self.csv:
                      file = self.save_dir / "results.csv"
                      n = len(x) + 1  # number of cols
                      s = "" if file.exists() else (("%20s," * n % tuple(["epoch"] + self.keys)).rstrip(",") + "\n")  # add header
                      with open(file, "a") as f:
                          f.write(s + ("%20.5g," * n % tuple([epoch] + vals)).rstrip(",") + "\n")
                  if self.ndjson_console or self.ndjson_file:
                      json_data = json.dumps(dict(epoch=epoch, **x), default=_json_default)
                  if self.ndjson_console:
                      print(json_data)
                  if self.ndjson_file:
                      file = self.save_dir / "results.ndjson"
                      with open(file, "a") as f:
                          print(json_data, file=f)
          
                  if self.tb:
                      for k, v in x.items():
                          self.tb.add_scalar(k, v, epoch)
                  elif self.clearml:  # log to ClearML if TensorBoard not used
                      self.clearml.log_scalars(x, epoch)
          
                  if self.wandb:
                      if best_fitness == fi:
                          best_results = [epoch] + vals[3:7]
                          for i, name in enumerate(self.best_keys):
                              self.wandb.wandb_run.summary[name] = best_results[i]  # log best results in the summary
                      self.wandb.log(x)
                      self.wandb.end_epoch()
          
                  if self.clearml:
                      self.clearml.current_epoch_logged_images = set()  # reset epoch image limit
                      self.clearml.current_epoch += 1
          
                  if self.comet_logger:
                      self.comet_logger.on_fit_epoch_end(x, epoch=epoch)
          
              def on_model_save(self, last, epoch, final_epoch, best_fitness, fi):
                  """Callback that handles model saving events, logging to Weights & Biases or ClearML if enabled."""
                  if (epoch + 1) % self.opt.save_period == 0 and not final_epoch and self.opt.save_period != -1:
                      if self.wandb:
                          self.wandb.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)
                      if self.clearml:
                          self.clearml.task.update_output_model(
                              model_path=str(last), model_name="Latest Model", auto_delete_file=False
                          )
          
                  if self.comet_logger:
                      self.comet_logger.on_model_save(last, epoch, final_epoch, best_fitness, fi)
          
              def on_train_end(self, last, best, epoch, results):
                  """Callback that runs at the end of training to save plots and log results."""
                  if self.plots:
                      plot_results(file=self.save_dir / "results.csv")  # save results.png
                  files = ["results.png", "confusion_matrix.png", *(f"{x}_curve.png" for x in ("F1", "PR", "P", "R"))]
                  files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()]  # filter
                  self.logger.info(f"Results saved to {colorstr('bold', self.save_dir)}")
          
                  if self.tb and not self.clearml:  # These images are already captured by ClearML by now, we don't want doubles
                      for f in files:
                          self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats="HWC")
          
                  if self.wandb:
                      self.wandb.log(dict(zip(self.keys[3:10], results)))
                      self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]})
                      # Calling wandb.log. TODO: Refactor this into WandbLogger.log_model
                      if not self.opt.evolve:
                          wandb.log_artifact(
                              str(best if best.exists() else last),
                              type="model",
                              name=f"run_{self.wandb.wandb_run.id}_model",
                              aliases=["latest", "best", "stripped"],
                          )
                      self.wandb.finish_run()
          
                  if self.clearml and not self.opt.evolve:
                      self.clearml.log_summary(dict(zip(self.keys[3:10], results)))
                      [self.clearml.log_plot(title=f.stem, plot_path=f) for f in files]
                      self.clearml.log_model(
                          str(best if best.exists() else last), "Best Model" if best.exists() else "Last Model", epoch
                      )
          
                  if self.comet_logger:
                      final_results = dict(zip(self.keys[3:10], results))
                      self.comet_logger.on_train_end(files, self.save_dir, last, best, epoch, final_results)
          
              def on_params_update(self, params: dict):
                  """Updates experiment hyperparameters or configurations in WandB, Comet, or ClearML."""
                  if self.wandb:
                      self.wandb.wandb_run.config.update(params, allow_val_change=True)
                  if self.comet_logger:
                      self.comet_logger.on_params_update(params)
                  if self.clearml:
                      self.clearml.task.connect(params)
          
          
          class GenericLogger:
              """
              YOLOv5 General purpose logger for non-task specific logging
              Usage: from utils.loggers import GenericLogger; logger = GenericLogger(...)
              Arguments
                  opt:             Run arguments
                  console_logger:  Console logger
                  include:         loggers to include
              """
          
              def __init__(self, opt, console_logger, include=("tb", "wandb", "clearml")):
                  """Initializes a generic logger with optional TensorBoard, W&B, and ClearML support."""
                  self.save_dir = Path(opt.save_dir)
                  self.include = include
                  self.console_logger = console_logger
                  self.csv = self.save_dir / "results.csv"  # CSV logger
                  if "tb" in self.include:
                      prefix = colorstr("TensorBoard: ")
                      self.console_logger.info(
                          f"{prefix}Start with 'tensorboard --logdir {self.save_dir.parent}', view at http://localhost:6006/"
                      )
                      self.tb = SummaryWriter(str(self.save_dir))
          
                  if wandb and "wandb" in self.include:
                      self.wandb = wandb.init(
                          project=web_project_name(str(opt.project)), name=None if opt.name == "exp" else opt.name, config=opt
                      )
                  else:
                      self.wandb = None
          
                  if clearml and "clearml" in self.include:
                      try:
                          # Hyp is not available in classification mode
                          hyp = {} if "hyp" not in opt else opt.hyp
                          self.clearml = ClearmlLogger(opt, hyp)
                      except Exception:
                          self.clearml = None
                          prefix = colorstr("ClearML: ")
                          LOGGER.warning(
                              f"{prefix}WARNING ⚠️ ClearML is installed but not configured, skipping ClearML logging."
                              f" See https://docs.ultralytics.com/yolov5/tutorials/clearml_logging_integration"
                          )
                  else:
                      self.clearml = None
          
              def log_metrics(self, metrics, epoch):
                  """Logs metrics to CSV, TensorBoard, W&B, and ClearML; `metrics` is a dict, `epoch` is an int."""
                  if self.csv:
                      keys, vals = list(metrics.keys()), list(metrics.values())
                      n = len(metrics) + 1  # number of cols
                      s = "" if self.csv.exists() else (("%23s," * n % tuple(["epoch"] + keys)).rstrip(",") + "\n")  # header
                      with open(self.csv, "a") as f:
                          f.write(s + ("%23.5g," * n % tuple([epoch] + vals)).rstrip(",") + "\n")
          
                  if self.tb:
                      for k, v in metrics.items():
                          self.tb.add_scalar(k, v, epoch)
          
                  if self.wandb:
                      self.wandb.log(metrics, step=epoch)
          
                  if self.clearml:
                      self.clearml.log_scalars(metrics, epoch)
          
              def log_images(self, files, name="Images", epoch=0):
                  """Logs images to all loggers with optional naming and epoch specification."""
                  files = [Path(f) for f in (files if isinstance(files, (tuple, list)) else [files])]  # to Path
                  files = [f for f in files if f.exists()]  # filter by exists
          
                  if self.tb:
                      for f in files:
                          self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats="HWC")
          
                  if self.wandb:
                      self.wandb.log({name: [wandb.Image(str(f), caption=f.name) for f in files]}, step=epoch)
          
                  if self.clearml:
                      if name == "Results":
                          [self.clearml.log_plot(f.stem, f) for f in files]
                      else:
                          self.clearml.log_debug_samples(files, title=name)
          
              def log_graph(self, model, imgsz=(640, 640)):
                  """Logs model graph to all configured loggers with specified input image size."""
                  if self.tb:
                      log_tensorboard_graph(self.tb, model, imgsz)
          
              def log_model(self, model_path, epoch=0, metadata=None):
                  """Logs the model to all configured loggers with optional epoch and metadata."""
                  if metadata is None:
                      metadata = {}
                  # Log model to all loggers
                  if self.wandb:
                      art = wandb.Artifact(name=f"run_{wandb.run.id}_model", type="model", metadata=metadata)
                      art.add_file(str(model_path))
                      wandb.log_artifact(art)
                  if self.clearml:
                      self.clearml.log_model(model_path=model_path, model_name=model_path.stem)
          
              def update_params(self, params):
                  """Updates logged parameters in WandB and/or ClearML if enabled."""
                  if self.wandb:
                      wandb.run.config.update(params, allow_val_change=True)
                  if self.clearml:
                      self.clearml.task.connect(params)
          
          
          def log_tensorboard_graph(tb, model, imgsz=(640, 640)):
              """Logs the model graph to TensorBoard with specified image size and model."""
              try:
                  p = next(model.parameters())  # for device, type
                  imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else imgsz  # expand
                  im = torch.zeros((1, 3, *imgsz)).to(p.device).type_as(p)  # input image (WARNING: must be zeros, not empty)
                  with warnings.catch_warnings():
                      warnings.simplefilter("ignore")  # suppress jit trace warning
                      tb.add_graph(torch.jit.trace(de_parallel(model), im, strict=False), [])
              except Exception as e:
                  LOGGER.warning(f"WARNING ⚠️ TensorBoard graph visualization failure {e}")
          
          
          def web_project_name(project):
              """Converts a local project name to a standardized web project name with optional suffixes."""
              if not project.startswith("runs/train"):
                  return project
              suffix = "-Classify" if project.endswith("-cls") else "-Segment" if project.endswith("-seg") else ""
              return f"YOLOv5{suffix}"
          
      • 📁 segment/
        • 📄 __init__.py 0B
          
          
        • 📄 augmentations.py 3.6KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          """Image augmentation functions."""
          
          import math
          import random
          
          import cv2
          import numpy as np
          
          from ..augmentations import box_candidates
          from ..general import resample_segments, segment2box
          
          
          def mixup(im, labels, segments, im2, labels2, segments2):
              """
              Applies MixUp augmentation blending two images, labels, and segments with a random ratio.
          
              See https://arxiv.org/pdf/1710.09412.pdf
              """
              r = np.random.beta(32.0, 32.0)  # mixup ratio, alpha=beta=32.0
              im = (im * r + im2 * (1 - r)).astype(np.uint8)
              labels = np.concatenate((labels, labels2), 0)
              segments = np.concatenate((segments, segments2), 0)
              return im, labels, segments
          
          
          def random_perspective(
              im, targets=(), segments=(), degrees=10, translate=0.1, scale=0.1, shear=10, perspective=0.0, border=(0, 0)
          ):
              # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))
              # targets = [cls, xyxy]
          
              height = im.shape[0] + border[0] * 2  # shape(h,w,c)
              width = im.shape[1] + border[1] * 2
          
              # Center
              C = np.eye(3)
              C[0, 2] = -im.shape[1] / 2  # x translation (pixels)
              C[1, 2] = -im.shape[0] / 2  # y translation (pixels)
          
              # Perspective
              P = np.eye(3)
              P[2, 0] = random.uniform(-perspective, perspective)  # x perspective (about y)
              P[2, 1] = random.uniform(-perspective, perspective)  # y perspective (about x)
          
              # Rotation and Scale
              R = np.eye(3)
              a = random.uniform(-degrees, degrees)
              # a += random.choice([-180, -90, 0, 90])  # add 90deg rotations to small rotations
              s = random.uniform(1 - scale, 1 + scale)
              # s = 2 ** random.uniform(-scale, scale)
              R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)
          
              # Shear
              S = np.eye(3)
              S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # x shear (deg)
              S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # y shear (deg)
          
              # Translation
              T = np.eye(3)
              T[0, 2] = random.uniform(0.5 - translate, 0.5 + translate) * width  # x translation (pixels)
              T[1, 2] = random.uniform(0.5 - translate, 0.5 + translate) * height  # y translation (pixels)
          
              # Combined rotation matrix
              M = T @ S @ R @ P @ C  # order of operations (right to left) is IMPORTANT
              if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any():  # image changed
                  if perspective:
                      im = cv2.warpPerspective(im, M, dsize=(width, height), borderValue=(114, 114, 114))
                  else:  # affine
                      im = cv2.warpAffine(im, M[:2], dsize=(width, height), borderValue=(114, 114, 114))
          
              # Visualize
              # import matplotlib.pyplot as plt
              # ax = plt.subplots(1, 2, figsize=(12, 6))[1].ravel()
              # ax[0].imshow(im[:, :, ::-1])  # base
              # ax[1].imshow(im2[:, :, ::-1])  # warped
          
              # Transform label coordinates
              n = len(targets)
              new_segments = []
              if n:
                  new = np.zeros((n, 4))
                  segments = resample_segments(segments)  # upsample
                  for i, segment in enumerate(segments):
                      xy = np.ones((len(segment), 3))
                      xy[:, :2] = segment
                      xy = xy @ M.T  # transform
                      xy = xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]  # perspective rescale or affine
          
                      # clip
                      new[i] = segment2box(xy, width, height)
                      new_segments.append(xy)
          
                  # filter candidates
                  i = box_candidates(box1=targets[:, 1:5].T * s, box2=new.T, area_thr=0.01)
                  targets = targets[i]
                  targets[:, 1:5] = new[i]
                  new_segments = np.array(new_segments)[i]
          
              return im, targets, new_segments
          
        • 📄 dataloaders.py 13.0KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          """Dataloaders."""
          
          import os
          import random
          
          import cv2
          import numpy as np
          import torch
          from torch.utils.data import DataLoader, distributed
          
          from ..augmentations import augment_hsv, copy_paste, letterbox
          from ..dataloaders import InfiniteDataLoader, LoadImagesAndLabels, SmartDistributedSampler, seed_worker
          from ..general import LOGGER, xyn2xy, xywhn2xyxy, xyxy2xywhn
          from ..torch_utils import torch_distributed_zero_first
          from .augmentations import mixup, random_perspective
          
          RANK = int(os.getenv("RANK", -1))
          
          
          def create_dataloader(
              path,
              imgsz,
              batch_size,
              stride,
              single_cls=False,
              hyp=None,
              augment=False,
              cache=False,
              pad=0.0,
              rect=False,
              rank=-1,
              workers=8,
              image_weights=False,
              quad=False,
              prefix="",
              shuffle=False,
              mask_downsample_ratio=1,
              overlap_mask=False,
              seed=0,
          ):
              if rect and shuffle:
                  LOGGER.warning("WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False")
                  shuffle = False
              with torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDP
                  dataset = LoadImagesAndLabelsAndMasks(
                      path,
                      imgsz,
                      batch_size,
                      augment=augment,  # augmentation
                      hyp=hyp,  # hyperparameters
                      rect=rect,  # rectangular batches
                      cache_images=cache,
                      single_cls=single_cls,
                      stride=int(stride),
                      pad=pad,
                      image_weights=image_weights,
                      prefix=prefix,
                      downsample_ratio=mask_downsample_ratio,
                      overlap=overlap_mask,
                      rank=rank,
                  )
          
              batch_size = min(batch_size, len(dataset))
              nd = torch.cuda.device_count()  # number of CUDA devices
              nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers])  # number of workers
              sampler = None if rank == -1 else SmartDistributedSampler(dataset, shuffle=shuffle)
              loader = DataLoader if image_weights else InfiniteDataLoader  # only DataLoader allows for attribute updates
              generator = torch.Generator()
              generator.manual_seed(6148914691236517205 + seed + RANK)
              return loader(
                  dataset,
                  batch_size=batch_size,
                  shuffle=shuffle and sampler is None,
                  num_workers=nw,
                  sampler=sampler,
                  pin_memory=True,
                  collate_fn=LoadImagesAndLabelsAndMasks.collate_fn4 if quad else LoadImagesAndLabelsAndMasks.collate_fn,
                  worker_init_fn=seed_worker,
                  generator=generator,
              ), dataset
          
          
          class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
              def __init__(
                  self,
                  path,
                  img_size=640,
                  batch_size=16,
                  augment=False,
                  hyp=None,
                  rect=False,
                  image_weights=False,
                  cache_images=False,
                  single_cls=False,
                  stride=32,
                  pad=0,
                  min_items=0,
                  prefix="",
                  downsample_ratio=1,
                  overlap=False,
                  rank=-1,
                  seed=0,
              ):
                  super().__init__(
                      path,
                      img_size,
                      batch_size,
                      augment,
                      hyp,
                      rect,
                      image_weights,
                      cache_images,
                      single_cls,
                      stride,
                      pad,
                      min_items,
                      prefix,
                      rank,
                      seed,
                  )
                  self.downsample_ratio = downsample_ratio
                  self.overlap = overlap
          
              def __getitem__(self, index):
                  """Returns a transformed item from the dataset at the specified index, handling indexing and image weighting."""
                  index = self.indices[index]  # linear, shuffled, or image_weights
          
                  hyp = self.hyp
                  mosaic = self.mosaic and random.random() < hyp["mosaic"]
                  masks = []
                  if mosaic:
                      # Load mosaic
                      img, labels, segments = self.load_mosaic(index)
                      shapes = None
          
                      # MixUp augmentation
                      if random.random() < hyp["mixup"]:
                          img, labels, segments = mixup(img, labels, segments, *self.load_mosaic(random.randint(0, self.n - 1)))
          
                  else:
                      # Load image
                      img, (h0, w0), (h, w) = self.load_image(index)
          
                      # Letterbox
                      shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size  # final letterboxed shape
                      img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
                      shapes = (h0, w0), ((h / h0, w / w0), pad)  # for COCO mAP rescaling
          
                      labels = self.labels[index].copy()
                      # [array, array, ....], array.shape=(num_points, 2), xyxyxyxy
                      segments = self.segments[index].copy()
                      if len(segments):
                          for i_s in range(len(segments)):
                              segments[i_s] = xyn2xy(
                                  segments[i_s],
                                  ratio[0] * w,
                                  ratio[1] * h,
                                  padw=pad[0],
                                  padh=pad[1],
                              )
                      if labels.size:  # normalized xywh to pixel xyxy format
                          labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
          
                      if self.augment:
                          img, labels, segments = random_perspective(
                              img,
                              labels,
                              segments=segments,
                              degrees=hyp["degrees"],
                              translate=hyp["translate"],
                              scale=hyp["scale"],
                              shear=hyp["shear"],
                              perspective=hyp["perspective"],
                          )
          
                  nl = len(labels)  # number of labels
                  if nl:
                      labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1e-3)
                      if self.overlap:
                          masks, sorted_idx = polygons2masks_overlap(
                              img.shape[:2], segments, downsample_ratio=self.downsample_ratio
                          )
                          masks = masks[None]  # (640, 640) -> (1, 640, 640)
                          labels = labels[sorted_idx]
                      else:
                          masks = polygons2masks(img.shape[:2], segments, color=1, downsample_ratio=self.downsample_ratio)
          
                  masks = (
                      torch.from_numpy(masks)
                      if len(masks)
                      else torch.zeros(
                          1 if self.overlap else nl, img.shape[0] // self.downsample_ratio, img.shape[1] // self.downsample_ratio
                      )
                  )
                  # TODO: albumentations support
                  if self.augment:
                      # Albumentations
                      # there are some augmentation that won't change boxes and masks,
                      # so just be it for now.
                      img, labels = self.albumentations(img, labels)
                      nl = len(labels)  # update after albumentations
          
                      # HSV color-space
                      augment_hsv(img, hgain=hyp["hsv_h"], sgain=hyp["hsv_s"], vgain=hyp["hsv_v"])
          
                      # Flip up-down
                      if random.random() < hyp["flipud"]:
                          img = np.flipud(img)
                          if nl:
                              labels[:, 2] = 1 - labels[:, 2]
                              masks = torch.flip(masks, dims=[1])
          
                      # Flip left-right
                      if random.random() < hyp["fliplr"]:
                          img = np.fliplr(img)
                          if nl:
                              labels[:, 1] = 1 - labels[:, 1]
                              masks = torch.flip(masks, dims=[2])
          
                      # Cutouts  # labels = cutout(img, labels, p=0.5)
          
                  labels_out = torch.zeros((nl, 6))
                  if nl:
                      labels_out[:, 1:] = torch.from_numpy(labels)
          
                  # Convert
                  img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
                  img = np.ascontiguousarray(img)
          
                  return (torch.from_numpy(img), labels_out, self.im_files[index], shapes, masks)
          
              def load_mosaic(self, index):
                  """Loads 1 image + 3 random images into a 4-image YOLOv5 mosaic, adjusting labels and segments accordingly."""
                  labels4, segments4 = [], []
                  s = self.img_size
                  yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border)  # mosaic center x, y
          
                  # 3 additional image indices
                  indices = [index] + random.choices(self.indices, k=3)  # 3 additional image indices
                  for i, index in enumerate(indices):
                      # Load image
                      img, _, (h, w) = self.load_image(index)
          
                      # place img in img4
                      if i == 0:  # top left
                          img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
                          x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)
                          x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)
                      elif i == 1:  # top right
                          x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
                          x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
                      elif i == 2:  # bottom left
                          x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
                          x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
                      elif i == 3:  # bottom right
                          x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
                          x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)
          
                      img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]
                      padw = x1a - x1b
                      padh = y1a - y1b
          
                      labels, segments = self.labels[index].copy(), self.segments[index].copy()
          
                      if labels.size:
                          labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh)  # normalized xywh to pixel xyxy format
                          segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
                      labels4.append(labels)
                      segments4.extend(segments)
          
                  # Concat/clip labels
                  labels4 = np.concatenate(labels4, 0)
                  for x in (labels4[:, 1:], *segments4):
                      np.clip(x, 0, 2 * s, out=x)  # clip when using random_perspective()
                  # img4, labels4 = replicate(img4, labels4)  # replicate
          
                  # Augment
                  img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp["copy_paste"])
                  img4, labels4, segments4 = random_perspective(
                      img4,
                      labels4,
                      segments4,
                      degrees=self.hyp["degrees"],
                      translate=self.hyp["translate"],
                      scale=self.hyp["scale"],
                      shear=self.hyp["shear"],
                      perspective=self.hyp["perspective"],
                      border=self.mosaic_border,
                  )  # border to remove
                  return img4, labels4, segments4
          
              @staticmethod
              def collate_fn(batch):
                  """Custom collation function for DataLoader, batches images, labels, paths, shapes, and segmentation masks."""
                  img, label, path, shapes, masks = zip(*batch)  # transposed
                  batched_masks = torch.cat(masks, 0)
                  for i, l in enumerate(label):
                      l[:, 0] = i  # add target image index for build_targets()
                  return torch.stack(img, 0), torch.cat(label, 0), path, shapes, batched_masks
          
          
          def polygon2mask(img_size, polygons, color=1, downsample_ratio=1):
              """
              Args:
                  img_size (tuple): The image size.
                  polygons (np.ndarray): [N, M], N is the number of polygons,
                      M is the number of points(Be divided by 2).
              """
              mask = np.zeros(img_size, dtype=np.uint8)
              polygons = np.asarray(polygons)
              polygons = polygons.astype(np.int32)
              shape = polygons.shape
              polygons = polygons.reshape(shape[0], -1, 2)
              cv2.fillPoly(mask, polygons, color=color)
              nh, nw = (img_size[0] // downsample_ratio, img_size[1] // downsample_ratio)
              # NOTE: fillPoly firstly then resize is trying the keep the same way
              # of loss calculation when mask-ratio=1.
              mask = cv2.resize(mask, (nw, nh))
              return mask
          
          
          def polygons2masks(img_size, polygons, color, downsample_ratio=1):
              """
              Args:
                  img_size (tuple): The image size.
                  polygons (list[np.ndarray]): each polygon is [N, M],
                      N is the number of polygons,
                      M is the number of points(Be divided by 2).
              """
              masks = []
              for si in range(len(polygons)):
                  mask = polygon2mask(img_size, [polygons[si].reshape(-1)], color, downsample_ratio)
                  masks.append(mask)
              return np.array(masks)
          
          
          def polygons2masks_overlap(img_size, segments, downsample_ratio=1):
              """Return a (640, 640) overlap mask."""
              masks = np.zeros(
                  (img_size[0] // downsample_ratio, img_size[1] // downsample_ratio),
                  dtype=np.int32 if len(segments) > 255 else np.uint8,
              )
              areas = []
              ms = []
              for si in range(len(segments)):
                  mask = polygon2mask(
                      img_size,
                      [segments[si].reshape(-1)],
                      downsample_ratio=downsample_ratio,
                      color=1,
                  )
                  ms.append(mask)
                  areas.append(mask.sum())
              areas = np.asarray(areas)
              index = np.argsort(-areas)
              ms = np.array(ms)[index]
              for i in range(len(segments)):
                  mask = ms[i] * (i + 1)
                  masks = masks + mask
                  masks = np.clip(masks, a_min=0, a_max=i + 1)
              return masks, index
          
        • 📄 general.py 5.8KB
          import cv2
          import numpy as np
          import torch
          import torch.nn.functional as F
          
          
          def crop_mask(masks, boxes):
              """
              "Crop" predicted masks by zeroing out everything not in the predicted bbox. Vectorized by Chong (thanks Chong).
          
              Args:
                  - masks should be a size [n, h, w] tensor of masks
                  - boxes should be a size [n, 4] tensor of bbox coords in relative point form
              """
          
              n, h, w = masks.shape
              x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1)  # x1 shape(1,1,n)
              r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :]  # rows shape(1,w,1)
              c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None]  # cols shape(h,1,1)
          
              return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
          
          
          def process_mask_upsample(protos, masks_in, bboxes, shape):
              """
              Crop after upsample.
              protos: [mask_dim, mask_h, mask_w]
              masks_in: [n, mask_dim], n is number of masks after nms
              bboxes: [n, 4], n is number of masks after nms
              shape: input_image_size, (h, w)
          
              return: h, w, n
              """
          
              c, mh, mw = protos.shape  # CHW
              masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
              masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
              masks = crop_mask(masks, bboxes)  # CHW
              return masks.gt_(0.5)
          
          
          def process_mask(protos, masks_in, bboxes, shape, upsample=False):
              """
              Crop before upsample.
              proto_out: [mask_dim, mask_h, mask_w]
              out_masks: [n, mask_dim], n is number of masks after nms
              bboxes: [n, 4], n is number of masks after nms
              shape:input_image_size, (h, w)
          
              return: h, w, n
              """
          
              c, mh, mw = protos.shape  # CHW
              ih, iw = shape
              masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)  # CHW
          
              downsampled_bboxes = bboxes.clone()
              downsampled_bboxes[:, 0] *= mw / iw
              downsampled_bboxes[:, 2] *= mw / iw
              downsampled_bboxes[:, 3] *= mh / ih
              downsampled_bboxes[:, 1] *= mh / ih
          
              masks = crop_mask(masks, downsampled_bboxes)  # CHW
              if upsample:
                  masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
              return masks.gt_(0.5)
          
          
          def process_mask_native(protos, masks_in, bboxes, shape):
              """
              Crop after upsample.
              protos: [mask_dim, mask_h, mask_w]
              masks_in: [n, mask_dim], n is number of masks after nms
              bboxes: [n, 4], n is number of masks after nms
              shape: input_image_size, (h, w)
          
              return: h, w, n
              """
              c, mh, mw = protos.shape  # CHW
              masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
              gain = min(mh / shape[0], mw / shape[1])  # gain  = old / new
              pad = (mw - shape[1] * gain) / 2, (mh - shape[0] * gain) / 2  # wh padding
              top, left = int(pad[1]), int(pad[0])  # y, x
              bottom, right = int(mh - pad[1]), int(mw - pad[0])
              masks = masks[:, top:bottom, left:right]
          
              masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
              masks = crop_mask(masks, bboxes)  # CHW
              return masks.gt_(0.5)
          
          
          def scale_image(im1_shape, masks, im0_shape, ratio_pad=None):
              """
              img1_shape: model input shape, [h, w]
              img0_shape: origin pic shape, [h, w, 3]
              masks: [h, w, num]
              """
              # Rescale coordinates (xyxy) from im1_shape to im0_shape
              if ratio_pad is None:  # calculate from im0_shape
                  gain = min(im1_shape[0] / im0_shape[0], im1_shape[1] / im0_shape[1])  # gain  = old / new
                  pad = (im1_shape[1] - im0_shape[1] * gain) / 2, (im1_shape[0] - im0_shape[0] * gain) / 2  # wh padding
              else:
                  pad = ratio_pad[1]
              top, left = int(pad[1]), int(pad[0])  # y, x
              bottom, right = int(im1_shape[0] - pad[1]), int(im1_shape[1] - pad[0])
          
              if len(masks.shape) < 2:
                  raise ValueError(f'"len of masks shape" should be 2 or 3, but got {len(masks.shape)}')
              masks = masks[top:bottom, left:right]
              # masks = masks.permute(2, 0, 1).contiguous()
              # masks = F.interpolate(masks[None], im0_shape[:2], mode='bilinear', align_corners=False)[0]
              # masks = masks.permute(1, 2, 0).contiguous()
              masks = cv2.resize(masks, (im0_shape[1], im0_shape[0]))
          
              if len(masks.shape) == 2:
                  masks = masks[:, :, None]
              return masks
          
          
          def mask_iou(mask1, mask2, eps=1e-7):
              """
              mask1: [N, n] m1 means number of predicted objects
              mask2: [M, n] m2 means number of gt objects
              Note: n means image_w x image_h
          
              return: masks iou, [N, M]
              """
              intersection = torch.matmul(mask1, mask2.t()).clamp(0)
              union = (mask1.sum(1)[:, None] + mask2.sum(1)[None]) - intersection  # (area1 + area2) - intersection
              return intersection / (union + eps)
          
          
          def masks_iou(mask1, mask2, eps=1e-7):
              """
              mask1: [N, n] m1 means number of predicted objects
              mask2: [N, n] m2 means number of gt objects
              Note: n means image_w x image_h
          
              return: masks iou, (N, )
              """
              intersection = (mask1 * mask2).sum(1).clamp(0)  # (N, )
              union = (mask1.sum(1) + mask2.sum(1))[None] - intersection  # (area1 + area2) - intersection
              return intersection / (union + eps)
          
          
          def masks2segments(masks, strategy="largest"):
              """Converts binary (n,160,160) masks to polygon segments with options for concatenation or selecting the largest
              segment.
              """
              segments = []
              for x in masks.int().cpu().numpy().astype("uint8"):
                  c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
                  if c:
                      if strategy == "concat":  # concatenate all segments
                          c = np.concatenate([x.reshape(-1, 2) for x in c])
                      elif strategy == "largest":  # select largest segment
                          c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)
                  else:
                      c = np.zeros((0, 2))  # no segments found
                  segments.append(c.astype("float32"))
              return segments
          
        • 📄 loss.py 8.9KB
          import torch
          import torch.nn as nn
          import torch.nn.functional as F
          
          from ..general import xywh2xyxy
          from ..loss import FocalLoss, smooth_BCE
          from ..metrics import bbox_iou
          from ..torch_utils import de_parallel
          from .general import crop_mask
          
          
          class ComputeLoss:
              # Compute losses
              def __init__(self, model, autobalance=False, overlap=False):
                  """Initializes the compute loss function for YOLOv5 models with options for autobalancing and overlap
                  handling.
                  """
                  self.sort_obj_iou = False
                  self.overlap = overlap
                  device = next(model.parameters()).device  # get model device
                  h = model.hyp  # hyperparameters
          
                  # Define criteria
                  BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["cls_pw"]], device=device))
                  BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["obj_pw"]], device=device))
          
                  # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
                  self.cp, self.cn = smooth_BCE(eps=h.get("label_smoothing", 0.0))  # positive, negative BCE targets
          
                  # Focal loss
                  g = h["fl_gamma"]  # focal loss gamma
                  if g > 0:
                      BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
          
                  m = de_parallel(model).model[-1]  # Detect() module
                  self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl, [4.0, 1.0, 0.25, 0.06, 0.02])  # P3-P7
                  self.ssi = list(m.stride).index(16) if autobalance else 0  # stride 16 index
                  self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance
                  self.na = m.na  # number of anchors
                  self.nc = m.nc  # number of classes
                  self.nl = m.nl  # number of layers
                  self.nm = m.nm  # number of masks
                  self.anchors = m.anchors
                  self.device = device
          
              def __call__(self, preds, targets, masks):  # predictions, targets, model
                  """Evaluates YOLOv5 model's loss for given predictions, targets, and masks; returns total loss components."""
                  p, proto = preds
                  bs, nm, mask_h, mask_w = proto.shape  # batch size, number of masks, mask height, mask width
                  lcls = torch.zeros(1, device=self.device)
                  lbox = torch.zeros(1, device=self.device)
                  lobj = torch.zeros(1, device=self.device)
                  lseg = torch.zeros(1, device=self.device)
                  tcls, tbox, indices, anchors, tidxs, xywhn = self.build_targets(p, targets)  # targets
          
                  # Losses
                  for i, pi in enumerate(p):  # layer index, layer predictions
                      b, a, gj, gi = indices[i]  # image, anchor, gridy, gridx
                      tobj = torch.zeros(pi.shape[:4], dtype=pi.dtype, device=self.device)  # target obj
          
                      n = b.shape[0]  # number of targets
                      if n:
                          pxy, pwh, _, pcls, pmask = pi[b, a, gj, gi].split((2, 2, 1, self.nc, nm), 1)  # subset of predictions
          
                          # Box regression
                          pxy = pxy.sigmoid() * 2 - 0.5
                          pwh = (pwh.sigmoid() * 2) ** 2 * anchors[i]
                          pbox = torch.cat((pxy, pwh), 1)  # predicted box
                          iou = bbox_iou(pbox, tbox[i], CIoU=True).squeeze()  # iou(prediction, target)
                          lbox += (1.0 - iou).mean()  # iou loss
          
                          # Objectness
                          iou = iou.detach().clamp(0).type(tobj.dtype)
                          if self.sort_obj_iou:
                              j = iou.argsort()
                              b, a, gj, gi, iou = b[j], a[j], gj[j], gi[j], iou[j]
                          if self.gr < 1:
                              iou = (1.0 - self.gr) + self.gr * iou
                          tobj[b, a, gj, gi] = iou  # iou ratio
          
                          # Classification
                          if self.nc > 1:  # cls loss (only if multiple classes)
                              t = torch.full_like(pcls, self.cn, device=self.device)  # targets
                              t[range(n), tcls[i]] = self.cp
                              lcls += self.BCEcls(pcls, t)  # BCE
          
                          # Mask regression
                          if tuple(masks.shape[-2:]) != (mask_h, mask_w):  # downsample
                              masks = F.interpolate(masks[None], (mask_h, mask_w), mode="nearest")[0]
                          marea = xywhn[i][:, 2:].prod(1)  # mask width, height normalized
                          mxyxy = xywh2xyxy(xywhn[i] * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=self.device))
                          for bi in b.unique():
                              j = b == bi  # matching index
                              if self.overlap:
                                  mask_gti = torch.where(masks[bi][None] == tidxs[i][j].view(-1, 1, 1), 1.0, 0.0)
                              else:
                                  mask_gti = masks[tidxs[i]][j]
                              lseg += self.single_mask_loss(mask_gti, pmask[j], proto[bi], mxyxy[j], marea[j])
          
                      obji = self.BCEobj(pi[..., 4], tobj)
                      lobj += obji * self.balance[i]  # obj loss
                      if self.autobalance:
                          self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()
          
                  if self.autobalance:
                      self.balance = [x / self.balance[self.ssi] for x in self.balance]
                  lbox *= self.hyp["box"]
                  lobj *= self.hyp["obj"]
                  lcls *= self.hyp["cls"]
                  lseg *= self.hyp["box"] / bs
          
                  loss = lbox + lobj + lcls + lseg
                  return loss * bs, torch.cat((lbox, lseg, lobj, lcls)).detach()
          
              def single_mask_loss(self, gt_mask, pred, proto, xyxy, area):
                  """Calculates and normalizes single mask loss for YOLOv5 between predicted and ground truth masks."""
                  pred_mask = (pred @ proto.view(self.nm, -1)).view(-1, *proto.shape[1:])  # (n,32) @ (32,80,80) -> (n,80,80)
                  loss = F.binary_cross_entropy_with_logits(pred_mask, gt_mask, reduction="none")
                  return (crop_mask(loss, xyxy).mean(dim=(1, 2)) / area).mean()
          
              def build_targets(self, p, targets):
                  """Prepares YOLOv5 targets for loss computation; inputs targets (image, class, x, y, w, h), output target
                  classes/boxes.
                  """
                  na, nt = self.na, targets.shape[0]  # number of anchors, targets
                  tcls, tbox, indices, anch, tidxs, xywhn = [], [], [], [], [], []
                  gain = torch.ones(8, device=self.device)  # normalized to gridspace gain
                  ai = torch.arange(na, device=self.device).float().view(na, 1).repeat(1, nt)  # same as .repeat_interleave(nt)
                  if self.overlap:
                      batch = p[0].shape[0]
                      ti = []
                      for i in range(batch):
                          num = (targets[:, 0] == i).sum()  # find number of targets of each image
                          ti.append(torch.arange(num, device=self.device).float().view(1, num).repeat(na, 1) + 1)  # (na, num)
                      ti = torch.cat(ti, 1)  # (na, nt)
                  else:
                      ti = torch.arange(nt, device=self.device).float().view(1, nt).repeat(na, 1)
                  targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None], ti[..., None]), 2)  # append anchor indices
          
                  g = 0.5  # bias
                  off = (
                      torch.tensor(
                          [
                              [0, 0],
                              [1, 0],
                              [0, 1],
                              [-1, 0],
                              [0, -1],  # j,k,l,m
                              # [1, 1], [1, -1], [-1, 1], [-1, -1],  # jk,jm,lk,lm
                          ],
                          device=self.device,
                      ).float()
                      * g
                  )  # offsets
          
                  for i in range(self.nl):
                      anchors, shape = self.anchors[i], p[i].shape
                      gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]]  # xyxy gain
          
                      # Match targets to anchors
                      t = targets * gain  # shape(3,n,7)
                      if nt:
                          # Matches
                          r = t[..., 4:6] / anchors[:, None]  # wh ratio
                          j = torch.max(r, 1 / r).max(2)[0] < self.hyp["anchor_t"]  # compare
                          # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']  # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
                          t = t[j]  # filter
          
                          # Offsets
                          gxy = t[:, 2:4]  # grid xy
                          gxi = gain[[2, 3]] - gxy  # inverse
                          j, k = ((gxy % 1 < g) & (gxy > 1)).T
                          l, m = ((gxi % 1 < g) & (gxi > 1)).T
                          j = torch.stack((torch.ones_like(j), j, k, l, m))
                          t = t.repeat((5, 1, 1))[j]
                          offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
                      else:
                          t = targets[0]
                          offsets = 0
          
                      # Define
                      bc, gxy, gwh, at = t.chunk(4, 1)  # (image, class), grid xy, grid wh, anchors
                      (a, tidx), (b, c) = at.long().T, bc.long().T  # anchors, image, class
                      gij = (gxy - offsets).long()
                      gi, gj = gij.T  # grid indices
          
                      # Append
                      indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1)))  # image, anchor, grid
                      tbox.append(torch.cat((gxy - gij, gwh), 1))  # box
                      anch.append(anchors[a])  # anchors
                      tcls.append(c)  # class
                      tidxs.append(tidx)
                      xywhn.append(torch.cat((gxy, gwh), 1) / gain[2:6])  # xywh normalized
          
                  return tcls, tbox, indices, anch, tidxs, xywhn
          
        • 📄 metrics.py 5.5KB
          # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
          """Model validation metrics."""
          
          import numpy as np
          
          from ..metrics import ap_per_class
          
          
          def fitness(x):
              """Evaluates model fitness by a weighted sum of 8 metrics, `x`: [N,8] array, weights: [0.1, 0.9] for mAP and F1."""
              w = [0.0, 0.0, 0.1, 0.9, 0.0, 0.0, 0.1, 0.9]
              return (x[:, :8] * w).sum(1)
          
          
          def ap_per_class_box_and_mask(
              tp_m,
              tp_b,
              conf,
              pred_cls,
              target_cls,
              plot=False,
              save_dir=".",
              names=(),
          ):
              """
              Args:
                  tp_b: tp of boxes.
                  tp_m: tp of masks.
                  other arguments see `func: ap_per_class`.
              """
              results_boxes = ap_per_class(
                  tp_b, conf, pred_cls, target_cls, plot=plot, save_dir=save_dir, names=names, prefix="Box"
              )[2:]
              results_masks = ap_per_class(
                  tp_m, conf, pred_cls, target_cls, plot=plot, save_dir=save_dir, names=names, prefix="Mask"
              )[2:]
          
              return {
                  "boxes": {
                      "p": results_boxes[0],
                      "r": results_boxes[1],
                      "ap": results_boxes[3],
                      "f1": results_boxes[2],
                      "ap_class": results_boxes[4],
                  },
                  "masks": {
                      "p": results_masks[0],
                      "r": results_masks[1],
                      "ap": results_masks[3],
                      "f1": results_masks[2],
                      "ap_class": results_masks[4],
                  },
              }
          
          
          class Metric:
              def __init__(self) -> None:
                  self.p = []  # (nc, )
                  self.r = []  # (nc, )
                  self.f1 = []  # (nc, )
                  self.all_ap = []  # (nc, 10)
                  self.ap_class_index = []  # (nc, )
          
              @property
              def ap50(self):
                  """
                  AP@0.5 of all classes.
          
                  Return:
                      (nc, ) or [].
                  """
                  return self.all_ap[:, 0] if len(self.all_ap) else []
          
              @property
              def ap(self):
                  """AP@0.5:0.95
                  Return:
                      (nc, ) or [].
                  """
                  return self.all_ap.mean(1) if len(self.all_ap) else []
          
              @property
              def mp(self):
                  """
                  Mean precision of all classes.
          
                  Return:
                      float.
                  """
                  return self.p.mean() if len(self.p) else 0.0
          
              @property
              def mr(self):
                  """
                  Mean recall of all classes.
          
                  Return:
                      float.
                  """
                  return self.r.mean() if len(self.r) else 0.0
          
              @property
              def map50(self):
                  """
                  Mean AP@0.5 of all classes.
          
                  Return:
                      float.
                  """
                  return self.all_ap[:, 0].mean() if len(self.all_ap) else 0.0
          
              @property
              def map(self):
                  """
                  Mean AP@0.5:0.95 of all classes.
          
                  Return:
                      float.
                  """
                  return self.all_ap.mean() if len(self.all_ap) else 0.0
          
              def mean_results(self):
                  """Mean of results, return mp, mr, map50, map."""
                  return (self.mp, self.mr, self.map50, self.map)
          
              def class_result(self, i):
                  """Class-aware result, return p[i], r[i], ap50[i], ap[i]"""
                  return (self.p[i], self.r[i], self.ap50[i], self.ap[i])
          
              def get_maps(self, nc):
                  """Calculates and returns mean Average Precision (mAP) for each class given number of classes `nc`."""
                  maps = np.zeros(nc) + self.map
                  for i, c in enumerate(self.ap_class_index):
                      maps[c] = self.ap[i]
                  return maps
          
              def update(self, results):
                  """
                  Args:
                      results: tuple(p, r, ap, f1, ap_class)
                  """
                  p, r, all_ap, f1, ap_class_index = results
                  self.p = p
                  self.r = r
                  self.all_ap = all_ap
                  self.f1 = f1
                  self.ap_class_index = ap_class_index
          
          
          class Metrics:
              """Metric for boxes and masks."""
          
              def __init__(self) -> None:
                  self.metric_box = Metric()
                  self.metric_mask = Metric()
          
              def update(self, results):
                  """
                  Args:
                      results: Dict{'boxes': Dict{}, 'masks': Dict{}}
                  """
                  self.metric_box.update(list(results["boxes"].values()))
                  self.metric_mask.update(list(results["masks"].values()))
          
              def mean_results(self):
                  """Computes and returns the mean results for both box and mask metrics by summing their individual means."""
                  return self.metric_box.mean_results() + self.metric_mask.mean_results()
          
              def class_result(self, i):
                  """Returns the sum of box and mask metric results for a specified class index `i`."""
                  return self.metric_box.class_result(i) + self.metric_mask.class_result(i)
          
              def get_maps(self, nc):
                  """Calculates and returns the sum of mean average precisions (mAPs) for both box and mask metrics for `nc`
                  classes.
                  """
                  return self.metric_box.get_maps(nc) + self.metric_mask.get_maps(nc)
          
              @property
              def ap_class_index(self):
                  """Returns the class index for average precision, shared by both box and mask metrics."""
                  return self.metric_box.ap_class_index
          
          
          KEYS = [
              "train/box_loss",
              "train/seg_loss",  # train loss
              "train/obj_loss",
              "train/cls_loss",
              "metrics/precision(B)",
              "metrics/recall(B)",
              "metrics/mAP_0.5(B)",
              "metrics/mAP_0.5:0.95(B)",  # metrics
              "metrics/precision(M)",
              "metrics/recall(M)",
              "metrics/mAP_0.5(M)",
              "metrics/mAP_0.5:0.95(M)",  # metrics
              "val/box_loss",
              "val/seg_loss",  # val loss
              "val/obj_loss",
              "val/cls_loss",
              "x/lr0",
              "x/lr1",
              "x/lr2",
          ]
          
          BEST_KEYS = [
              "best/epoch",
              "best/precision(B)",
              "best/recall(B)",
              "best/mAP_0.5(B)",
              "best/mAP_0.5:0.95(B)",
              "best/precision(M)",
              "best/recall(M)",
              "best/mAP_0.5(M)",
              "best/mAP_0.5:0.95(M)",
          ]
          
        • 📄 plots.py 6.5KB
          import contextlib
          import math
          from pathlib import Path
          
          import cv2
          import matplotlib.pyplot as plt
          import numpy as np
          import pandas as pd
          import torch
          
          from .. import threaded
          from ..general import xywh2xyxy
          from ..plots import Annotator, colors
          
          
          @threaded
          def plot_images_and_masks(images, targets, masks, paths=None, fname="images.jpg", names=None):
              """Plots a grid of images, their labels, and masks with optional resizing and annotations, saving to fname."""
              if isinstance(images, torch.Tensor):
                  images = images.cpu().float().numpy()
              if isinstance(targets, torch.Tensor):
                  targets = targets.cpu().numpy()
              if isinstance(masks, torch.Tensor):
                  masks = masks.cpu().numpy().astype(int)
          
              max_size = 1920  # max image size
              max_subplots = 16  # max image subplots, i.e. 4x4
              bs, _, h, w = images.shape  # batch size, _, height, width
              bs = min(bs, max_subplots)  # limit plot images
              ns = np.ceil(bs**0.5)  # number of subplots (square)
              if np.max(images[0]) <= 1:
                  images *= 255  # de-normalise (optional)
          
              # Build Image
              mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8)  # init
              for i, im in enumerate(images):
                  if i == max_subplots:  # if last batch has fewer images than we expect
                      break
                  x, y = int(w * (i // ns)), int(h * (i % ns))  # block origin
                  im = im.transpose(1, 2, 0)
                  mosaic[y : y + h, x : x + w, :] = im
          
              # Resize (optional)
              scale = max_size / ns / max(h, w)
              if scale < 1:
                  h = math.ceil(scale * h)
                  w = math.ceil(scale * w)
                  mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h)))
          
              # Annotate
              fs = int((h + w) * ns * 0.01)  # font size
              annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names)
              for i in range(i + 1):
                  x, y = int(w * (i // ns)), int(h * (i % ns))  # block origin
                  annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2)  # borders
                  if paths:
                      annotator.text([x + 5, y + 5], text=Path(paths[i]).name[:40], txt_color=(220, 220, 220))  # filenames
                  if len(targets) > 0:
                      idx = targets[:, 0] == i
                      ti = targets[idx]  # image targets
          
                      boxes = xywh2xyxy(ti[:, 2:6]).T
                      classes = ti[:, 1].astype("int")
                      labels = ti.shape[1] == 6  # labels if no conf column
                      conf = None if labels else ti[:, 6]  # check for confidence presence (label vs pred)
          
                      if boxes.shape[1]:
                          if boxes.max() <= 1.01:  # if normalized with tolerance 0.01
                              boxes[[0, 2]] *= w  # scale to pixels
                              boxes[[1, 3]] *= h
                          elif scale < 1:  # absolute coords need scale if image scales
                              boxes *= scale
                      boxes[[0, 2]] += x
                      boxes[[1, 3]] += y
                      for j, box in enumerate(boxes.T.tolist()):
                          cls = classes[j]
                          color = colors(cls)
                          cls = names[cls] if names else cls
                          if labels or conf[j] > 0.25:  # 0.25 conf thresh
                              label = f"{cls}" if labels else f"{cls} {conf[j]:.1f}"
                              annotator.box_label(box, label, color=color)
          
                      # Plot masks
                      if len(masks):
                          if masks.max() > 1.0:  # mean that masks are overlap
                              image_masks = masks[[i]]  # (1, 640, 640)
                              nl = len(ti)
                              index = np.arange(nl).reshape(nl, 1, 1) + 1
                              image_masks = np.repeat(image_masks, nl, axis=0)
                              image_masks = np.where(image_masks == index, 1.0, 0.0)
                          else:
                              image_masks = masks[idx]
          
                          im = np.asarray(annotator.im).copy()
                          for j, box in enumerate(boxes.T.tolist()):
                              if labels or conf[j] > 0.25:  # 0.25 conf thresh
                                  color = colors(classes[j])
                                  mh, mw = image_masks[j].shape
                                  if mh != h or mw != w:
                                      mask = image_masks[j].astype(np.uint8)
                                      mask = cv2.resize(mask, (w, h))
                                      mask = mask.astype(bool)
                                  else:
                                      mask = image_masks[j].astype(bool)
                                  with contextlib.suppress(Exception):
                                      im[y : y + h, x : x + w, :][mask] = (
                                          im[y : y + h, x : x + w, :][mask] * 0.4 + np.array(color) * 0.6
                                      )
                          annotator.fromarray(im)
              annotator.im.save(fname)  # save
          
          
          def plot_results_with_masks(file="path/to/results.csv", dir="", best=True):
              """
              Plots training results from CSV files, plotting best or last result highlights based on `best` parameter.
          
              Example: from utils.plots import *; plot_results('path/to/results.csv')
              """
              save_dir = Path(file).parent if file else Path(dir)
              fig, ax = plt.subplots(2, 8, figsize=(18, 6), tight_layout=True)
              ax = ax.ravel()
              files = list(save_dir.glob("results*.csv"))
              assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot."
              for f in files:
                  try:
                      data = pd.read_csv(f)
                      index = np.argmax(
                          0.9 * data.values[:, 8] + 0.1 * data.values[:, 7] + 0.9 * data.values[:, 12] + 0.1 * data.values[:, 11]
                      )
                      s = [x.strip() for x in data.columns]
                      x = data.values[:, 0]
                      for i, j in enumerate([1, 2, 3, 4, 5, 6, 9, 10, 13, 14, 15, 16, 7, 8, 11, 12]):
                          y = data.values[:, j]
                          # y[y == 0] = np.nan  # don't show zero values
                          ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=2)
                          if best:
                              # best
                              ax[i].scatter(index, y[index], color="r", label=f"best:{index}", marker="*", linewidth=3)
                              ax[i].set_title(s[j] + f"\n{round(y[index], 5)}")
                          else:
                              # last
                              ax[i].scatter(x[-1], y[-1], color="r", label="last", marker="*", linewidth=3)
                              ax[i].set_title(s[j] + f"\n{round(y[-1], 5)}")
                          # if j in [8, 9, 10]:  # share train and val loss y axes
                          #     ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
                  except Exception as e:
                      print(f"Warning: Plotting error for {f}: {e}")
              ax[1].legend()
              fig.savefig(save_dir / "results.png", dpi=200)
              plt.close()
          
      • 📄 __init__.py 3.1KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """utils/initialization."""
        
        import contextlib
        import platform
        import threading
        
        
        def emojis(str=""):
            """Returns an emoji-safe version of a string, stripped of emojis on Windows platforms."""
            return str.encode().decode("ascii", "ignore") if platform.system() == "Windows" else str
        
        
        class TryExcept(contextlib.ContextDecorator):
            # YOLOv5 TryExcept class. Usage: @TryExcept() decorator or 'with TryExcept():' context manager
            def __init__(self, msg=""):
                """Initializes TryExcept with an optional message, used as a decorator or context manager for error handling."""
                self.msg = msg
        
            def __enter__(self):
                """Enter the runtime context related to this object for error handling with an optional message."""
                pass
        
            def __exit__(self, exc_type, value, traceback):
                """Context manager exit method that prints an error message with emojis if an exception occurred, always returns
                True.
                """
                if value:
                    print(emojis(f"{self.msg}{': ' if self.msg else ''}{value}"))
                return True
        
        
        def threaded(func):
            """Decorator @threaded to run a function in a separate thread, returning the thread instance."""
        
            def wrapper(*args, **kwargs):
                thread = threading.Thread(target=func, args=args, kwargs=kwargs, daemon=True)
                thread.start()
                return thread
        
            return wrapper
        
        
        def join_threads(verbose=False):
            """
            Joins all daemon threads, optionally printing their names if verbose is True.
        
            Example: atexit.register(lambda: join_threads())
            """
            main_thread = threading.current_thread()
            for t in threading.enumerate():
                if t is not main_thread:
                    if verbose:
                        print(f"Joining thread {t.name}")
                    t.join()
        
        
        def notebook_init(verbose=True):
            """Initializes notebook environment by checking requirements, cleaning up, and displaying system info."""
            print("Checking setup...")
        
            import os
            import shutil
        
            from ultralytics.utils.checks import check_requirements
        
            from utils.general import check_font, is_colab
            from utils.torch_utils import select_device  # imports
        
            check_font()
        
            import psutil
        
            if check_requirements("wandb", install=False):
                os.system("pip uninstall -y wandb")  # eliminate unexpected account creation prompt with infinite hang
            if is_colab():
                shutil.rmtree("/content/sample_data", ignore_errors=True)  # remove colab /sample_data directory
        
            # System info
            display = None
            if verbose:
                gb = 1 << 30  # bytes to GiB (1024 ** 3)
                ram = psutil.virtual_memory().total
                total, used, free = shutil.disk_usage("/")
                with contextlib.suppress(Exception):  # clear display if ipython is installed
                    from IPython import display
        
                    display.clear_output()
                s = f"({os.cpu_count()} CPUs, {ram / gb:.1f} GB RAM, {(total - free) / gb:.1f}/{total / gb:.1f} GB disk)"
            else:
                s = ""
        
            select_device(newline=False)
            print(emojis(f"Setup complete ✅ {s}"))
            return display
        
      • 📄 activations.py 4.5KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Activation functions."""
        
        import torch
        import torch.nn as nn
        import torch.nn.functional as F
        
        
        class SiLU(nn.Module):
            @staticmethod
            def forward(x):
                """
                Applies the Sigmoid-weighted Linear Unit (SiLU) activation function.
        
                https://arxiv.org/pdf/1606.08415.pdf.
                """
                return x * torch.sigmoid(x)
        
        
        class Hardswish(nn.Module):
            @staticmethod
            def forward(x):
                """
                Applies the Hardswish activation function, compatible with TorchScript, CoreML, and ONNX.
        
                Equivalent to x * F.hardsigmoid(x)
                """
                return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0  # for TorchScript, CoreML and ONNX
        
        
        class Mish(nn.Module):
            """Mish activation https://github.com/digantamisra98/Mish."""
        
            @staticmethod
            def forward(x):
                """Applies the Mish activation function, a smooth alternative to ReLU."""
                return x * F.softplus(x).tanh()
        
        
        class MemoryEfficientMish(nn.Module):
            class F(torch.autograd.Function):
                @staticmethod
                def forward(ctx, x):
                    """Applies the Mish activation function, a smooth ReLU alternative, to the input tensor `x`."""
                    ctx.save_for_backward(x)
                    return x.mul(torch.tanh(F.softplus(x)))  # x * tanh(ln(1 + exp(x)))
        
                @staticmethod
                def backward(ctx, grad_output):
                    """Computes the gradient of the Mish activation function with respect to input `x`."""
                    x = ctx.saved_tensors[0]
                    sx = torch.sigmoid(x)
                    fx = F.softplus(x).tanh()
                    return grad_output * (fx + x * sx * (1 - fx * fx))
        
            def forward(self, x):
                """Applies the Mish activation function to the input tensor `x`."""
                return self.F.apply(x)
        
        
        class FReLU(nn.Module):
            """FReLU activation https://arxiv.org/abs/2007.11824."""
        
            def __init__(self, c1, k=3):  # ch_in, kernel
                """Initializes FReLU activation with channel `c1` and kernel size `k`."""
                super().__init__()
                self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1, bias=False)
                self.bn = nn.BatchNorm2d(c1)
        
            def forward(self, x):
                """
                Applies FReLU activation with max operation between input and BN-convolved input.
        
                https://arxiv.org/abs/2007.11824
                """
                return torch.max(x, self.bn(self.conv(x)))
        
        
        class AconC(nn.Module):
            """
            ACON activation (activate or not) function.
        
            AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter
            See "Activate or Not: Learning Customized Activation" https://arxiv.org/pdf/2009.04759.pdf.
            """
        
            def __init__(self, c1):
                """Initializes AconC with learnable parameters p1, p2, and beta for channel-wise activation control."""
                super().__init__()
                self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1))
                self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1))
                self.beta = nn.Parameter(torch.ones(1, c1, 1, 1))
        
            def forward(self, x):
                """Applies AconC activation function with learnable parameters for channel-wise control on input tensor x."""
                dpx = (self.p1 - self.p2) * x
                return dpx * torch.sigmoid(self.beta * dpx) + self.p2 * x
        
        
        class MetaAconC(nn.Module):
            """
            ACON activation (activate or not) function.
        
            AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter
            See "Activate or Not: Learning Customized Activation" https://arxiv.org/pdf/2009.04759.pdf.
            """
        
            def __init__(self, c1, k=1, s=1, r=16):
                """Initializes MetaAconC with params: channel_in (c1), kernel size (k=1), stride (s=1), reduction (r=16)."""
                super().__init__()
                c2 = max(r, c1 // r)
                self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1))
                self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1))
                self.fc1 = nn.Conv2d(c1, c2, k, s, bias=True)
                self.fc2 = nn.Conv2d(c2, c1, k, s, bias=True)
                # self.bn1 = nn.BatchNorm2d(c2)
                # self.bn2 = nn.BatchNorm2d(c1)
        
            def forward(self, x):
                """Applies a forward pass transforming input `x` using learnable parameters and sigmoid activation."""
                y = x.mean(dim=2, keepdims=True).mean(dim=3, keepdims=True)
                # batch-size 1 bug/instabilities https://github.com/ultralytics/yolov5/issues/2891
                # beta = torch.sigmoid(self.bn2(self.fc2(self.bn1(self.fc1(y)))))  # bug/unstable
                beta = torch.sigmoid(self.fc2(self.fc1(y)))  # bug patch BN layers removed
                dpx = (self.p1 - self.p2) * x
                return dpx * torch.sigmoid(beta * dpx) + self.p2 * x
        
      • 📄 augmentations.py 18.2KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Image augmentation functions."""
        
        import math
        import random
        
        import cv2
        import numpy as np
        import torch
        import torchvision.transforms as T
        import torchvision.transforms.functional as TF
        
        from utils.general import LOGGER, check_version, colorstr, resample_segments, segment2box, xywhn2xyxy
        from utils.metrics import bbox_ioa
        
        IMAGENET_MEAN = 0.485, 0.456, 0.406  # RGB mean
        IMAGENET_STD = 0.229, 0.224, 0.225  # RGB standard deviation
        
        
        class Albumentations:
            # YOLOv5 Albumentations class (optional, only used if package is installed)
            def __init__(self, size=640):
                """Initializes Albumentations class for optional data augmentation in YOLOv5 with specified input size."""
                self.transform = None
                prefix = colorstr("albumentations: ")
                try:
                    import albumentations as A
        
                    check_version(A.__version__, "1.0.3", hard=True)  # version requirement
        
                    T = [
                        A.RandomResizedCrop(height=size, width=size, scale=(0.8, 1.0), ratio=(0.9, 1.11), p=0.0),
                        A.Blur(p=0.01),
                        A.MedianBlur(p=0.01),
                        A.ToGray(p=0.01),
                        A.CLAHE(p=0.01),
                        A.RandomBrightnessContrast(p=0.0),
                        A.RandomGamma(p=0.0),
                        A.ImageCompression(quality_lower=75, p=0.0),
                    ]  # transforms
                    self.transform = A.Compose(T, bbox_params=A.BboxParams(format="yolo", label_fields=["class_labels"]))
        
                    LOGGER.info(prefix + ", ".join(f"{x}".replace("always_apply=False, ", "") for x in T if x.p))
                except ImportError:  # package not installed, skip
                    pass
                except Exception as e:
                    LOGGER.info(f"{prefix}{e}")
        
            def __call__(self, im, labels, p=1.0):
                """Applies transformations to an image and labels with probability `p`, returning updated image and labels."""
                if self.transform and random.random() < p:
                    new = self.transform(image=im, bboxes=labels[:, 1:], class_labels=labels[:, 0])  # transformed
                    im, labels = new["image"], np.array([[c, *b] for c, b in zip(new["class_labels"], new["bboxes"])])
                return im, labels
        
        
        def normalize(x, mean=IMAGENET_MEAN, std=IMAGENET_STD, inplace=False):
            """
            Applies ImageNet normalization to RGB images in BCHW format, modifying them in-place if specified.
        
            Example: y = (x - mean) / std
            """
            return TF.normalize(x, mean, std, inplace=inplace)
        
        
        def denormalize(x, mean=IMAGENET_MEAN, std=IMAGENET_STD):
            """Reverses ImageNet normalization for BCHW format RGB images by applying `x = x * std + mean`."""
            for i in range(3):
                x[:, i] = x[:, i] * std[i] + mean[i]
            return x
        
        
        def augment_hsv(im, hgain=0.5, sgain=0.5, vgain=0.5):
            """Applies HSV color-space augmentation to an image with random gains for hue, saturation, and value."""
            if hgain or sgain or vgain:
                r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains
                hue, sat, val = cv2.split(cv2.cvtColor(im, cv2.COLOR_BGR2HSV))
                dtype = im.dtype  # uint8
        
                x = np.arange(0, 256, dtype=r.dtype)
                lut_hue = ((x * r[0]) % 180).astype(dtype)
                lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
                lut_val = np.clip(x * r[2], 0, 255).astype(dtype)
        
                im_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
                cv2.cvtColor(im_hsv, cv2.COLOR_HSV2BGR, dst=im)  # no return needed
        
        
        def hist_equalize(im, clahe=True, bgr=False):
            """Equalizes image histogram, with optional CLAHE, for BGR or RGB image with shape (n,m,3) and range 0-255."""
            yuv = cv2.cvtColor(im, cv2.COLOR_BGR2YUV if bgr else cv2.COLOR_RGB2YUV)
            if clahe:
                c = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
                yuv[:, :, 0] = c.apply(yuv[:, :, 0])
            else:
                yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0])  # equalize Y channel histogram
            return cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR if bgr else cv2.COLOR_YUV2RGB)  # convert YUV image to RGB
        
        
        def replicate(im, labels):
            """
            Replicates half of the smallest object labels in an image for data augmentation.
        
            Returns augmented image and labels.
            """
            h, w = im.shape[:2]
            boxes = labels[:, 1:].astype(int)
            x1, y1, x2, y2 = boxes.T
            s = ((x2 - x1) + (y2 - y1)) / 2  # side length (pixels)
            for i in s.argsort()[: round(s.size * 0.5)]:  # smallest indices
                x1b, y1b, x2b, y2b = boxes[i]
                bh, bw = y2b - y1b, x2b - x1b
                yc, xc = int(random.uniform(0, h - bh)), int(random.uniform(0, w - bw))  # offset x, y
                x1a, y1a, x2a, y2a = [xc, yc, xc + bw, yc + bh]
                im[y1a:y2a, x1a:x2a] = im[y1b:y2b, x1b:x2b]  # im4[ymin:ymax, xmin:xmax]
                labels = np.append(labels, [[labels[i, 0], x1a, y1a, x2a, y2a]], axis=0)
        
            return im, labels
        
        
        def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
            """Resizes and pads image to new_shape with stride-multiple constraints, returns resized image, ratio, padding."""
            shape = im.shape[:2]  # current shape [height, width]
            if isinstance(new_shape, int):
                new_shape = (new_shape, new_shape)
        
            # Scale ratio (new / old)
            r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
            if not scaleup:  # only scale down, do not scale up (for better val mAP)
                r = min(r, 1.0)
        
            # Compute padding
            ratio = r, r  # width, height ratios
            new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
            dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
            if auto:  # minimum rectangle
                dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
            elif scaleFill:  # stretch
                dw, dh = 0.0, 0.0
                new_unpad = (new_shape[1], new_shape[0])
                ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios
        
            dw /= 2  # divide padding into 2 sides
            dh /= 2
        
            if shape[::-1] != new_unpad:  # resize
                im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
            top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
            left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
            im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
            return im, ratio, (dw, dh)
        
        
        def random_perspective(
            im, targets=(), segments=(), degrees=10, translate=0.1, scale=0.1, shear=10, perspective=0.0, border=(0, 0)
        ):
            # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10))
            # targets = [cls, xyxy]
        
            height = im.shape[0] + border[0] * 2  # shape(h,w,c)
            width = im.shape[1] + border[1] * 2
        
            # Center
            C = np.eye(3)
            C[0, 2] = -im.shape[1] / 2  # x translation (pixels)
            C[1, 2] = -im.shape[0] / 2  # y translation (pixels)
        
            # Perspective
            P = np.eye(3)
            P[2, 0] = random.uniform(-perspective, perspective)  # x perspective (about y)
            P[2, 1] = random.uniform(-perspective, perspective)  # y perspective (about x)
        
            # Rotation and Scale
            R = np.eye(3)
            a = random.uniform(-degrees, degrees)
            # a += random.choice([-180, -90, 0, 90])  # add 90deg rotations to small rotations
            s = random.uniform(1 - scale, 1 + scale)
            # s = 2 ** random.uniform(-scale, scale)
            R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)
        
            # Shear
            S = np.eye(3)
            S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # x shear (deg)
            S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # y shear (deg)
        
            # Translation
            T = np.eye(3)
            T[0, 2] = random.uniform(0.5 - translate, 0.5 + translate) * width  # x translation (pixels)
            T[1, 2] = random.uniform(0.5 - translate, 0.5 + translate) * height  # y translation (pixels)
        
            # Combined rotation matrix
            M = T @ S @ R @ P @ C  # order of operations (right to left) is IMPORTANT
            if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any():  # image changed
                if perspective:
                    im = cv2.warpPerspective(im, M, dsize=(width, height), borderValue=(114, 114, 114))
                else:  # affine
                    im = cv2.warpAffine(im, M[:2], dsize=(width, height), borderValue=(114, 114, 114))
        
            # Visualize
            # import matplotlib.pyplot as plt
            # ax = plt.subplots(1, 2, figsize=(12, 6))[1].ravel()
            # ax[0].imshow(im[:, :, ::-1])  # base
            # ax[1].imshow(im2[:, :, ::-1])  # warped
        
            # Transform label coordinates
            n = len(targets)
            if n:
                use_segments = any(x.any() for x in segments) and len(segments) == n
                new = np.zeros((n, 4))
                if use_segments:  # warp segments
                    segments = resample_segments(segments)  # upsample
                    for i, segment in enumerate(segments):
                        xy = np.ones((len(segment), 3))
                        xy[:, :2] = segment
                        xy = xy @ M.T  # transform
                        xy = xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]  # perspective rescale or affine
        
                        # clip
                        new[i] = segment2box(xy, width, height)
        
                else:  # warp boxes
                    xy = np.ones((n * 4, 3))
                    xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2)  # x1y1, x2y2, x1y2, x2y1
                    xy = xy @ M.T  # transform
                    xy = (xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]).reshape(n, 8)  # perspective rescale or affine
        
                    # create new boxes
                    x = xy[:, [0, 2, 4, 6]]
                    y = xy[:, [1, 3, 5, 7]]
                    new = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
        
                    # clip
                    new[:, [0, 2]] = new[:, [0, 2]].clip(0, width)
                    new[:, [1, 3]] = new[:, [1, 3]].clip(0, height)
        
                # filter candidates
                i = box_candidates(box1=targets[:, 1:5].T * s, box2=new.T, area_thr=0.01 if use_segments else 0.10)
                targets = targets[i]
                targets[:, 1:5] = new[i]
        
            return im, targets
        
        
        def copy_paste(im, labels, segments, p=0.5):
            """
            Applies Copy-Paste augmentation by flipping and merging segments and labels on an image.
        
            Details at https://arxiv.org/abs/2012.07177.
            """
            n = len(segments)
            if p and n:
                h, w, c = im.shape  # height, width, channels
                im_new = np.zeros(im.shape, np.uint8)
                for j in random.sample(range(n), k=round(p * n)):
                    l, s = labels[j], segments[j]
                    box = w - l[3], l[2], w - l[1], l[4]
                    ioa = bbox_ioa(box, labels[:, 1:5])  # intersection over area
                    if (ioa < 0.30).all():  # allow 30% obscuration of existing labels
                        labels = np.concatenate((labels, [[l[0], *box]]), 0)
                        segments.append(np.concatenate((w - s[:, 0:1], s[:, 1:2]), 1))
                        cv2.drawContours(im_new, [segments[j].astype(np.int32)], -1, (1, 1, 1), cv2.FILLED)
        
                result = cv2.flip(im, 1)  # augment segments (flip left-right)
                i = cv2.flip(im_new, 1).astype(bool)
                im[i] = result[i]  # cv2.imwrite('debug.jpg', im)  # debug
        
            return im, labels, segments
        
        
        def cutout(im, labels, p=0.5):
            """
            Applies cutout augmentation to an image with optional label adjustment, using random masks of varying sizes.
        
            Details at https://arxiv.org/abs/1708.04552.
            """
            if random.random() < p:
                h, w = im.shape[:2]
                scales = [0.5] * 1 + [0.25] * 2 + [0.125] * 4 + [0.0625] * 8 + [0.03125] * 16  # image size fraction
                for s in scales:
                    mask_h = random.randint(1, int(h * s))  # create random masks
                    mask_w = random.randint(1, int(w * s))
        
                    # box
                    xmin = max(0, random.randint(0, w) - mask_w // 2)
                    ymin = max(0, random.randint(0, h) - mask_h // 2)
                    xmax = min(w, xmin + mask_w)
                    ymax = min(h, ymin + mask_h)
        
                    # apply random color mask
                    im[ymin:ymax, xmin:xmax] = [random.randint(64, 191) for _ in range(3)]
        
                    # return unobscured labels
                    if len(labels) and s > 0.03:
                        box = np.array([xmin, ymin, xmax, ymax], dtype=np.float32)
                        ioa = bbox_ioa(box, xywhn2xyxy(labels[:, 1:5], w, h))  # intersection over area
                        labels = labels[ioa < 0.60]  # remove >60% obscured labels
        
            return labels
        
        
        def mixup(im, labels, im2, labels2):
            """
            Applies MixUp augmentation by blending images and labels.
        
            See https://arxiv.org/pdf/1710.09412.pdf for details.
            """
            r = np.random.beta(32.0, 32.0)  # mixup ratio, alpha=beta=32.0
            im = (im * r + im2 * (1 - r)).astype(np.uint8)
            labels = np.concatenate((labels, labels2), 0)
            return im, labels
        
        
        def box_candidates(box1, box2, wh_thr=2, ar_thr=100, area_thr=0.1, eps=1e-16):
            """
            Filters bounding box candidates by minimum width-height threshold `wh_thr` (pixels), aspect ratio threshold
            `ar_thr`, and area ratio threshold `area_thr`.
        
            box1(4,n) is before augmentation, box2(4,n) is after augmentation.
            """
            w1, h1 = box1[2] - box1[0], box1[3] - box1[1]
            w2, h2 = box2[2] - box2[0], box2[3] - box2[1]
            ar = np.maximum(w2 / (h2 + eps), h2 / (w2 + eps))  # aspect ratio
            return (w2 > wh_thr) & (h2 > wh_thr) & (w2 * h2 / (w1 * h1 + eps) > area_thr) & (ar < ar_thr)  # candidates
        
        
        def classify_albumentations(
            augment=True,
            size=224,
            scale=(0.08, 1.0),
            ratio=(0.75, 1.0 / 0.75),  # 0.75, 1.33
            hflip=0.5,
            vflip=0.0,
            jitter=0.4,
            mean=IMAGENET_MEAN,
            std=IMAGENET_STD,
            auto_aug=False,
        ):
            # YOLOv5 classification Albumentations (optional, only used if package is installed)
            prefix = colorstr("albumentations: ")
            try:
                import albumentations as A
                from albumentations.pytorch import ToTensorV2
        
                check_version(A.__version__, "1.0.3", hard=True)  # version requirement
                if augment:  # Resize and crop
                    T = [A.RandomResizedCrop(height=size, width=size, scale=scale, ratio=ratio)]
                    if auto_aug:
                        # TODO: implement AugMix, AutoAug & RandAug in albumentation
                        LOGGER.info(f"{prefix}auto augmentations are currently not supported")
                    else:
                        if hflip > 0:
                            T += [A.HorizontalFlip(p=hflip)]
                        if vflip > 0:
                            T += [A.VerticalFlip(p=vflip)]
                        if jitter > 0:
                            color_jitter = (float(jitter),) * 3  # repeat value for brightness, contrast, satuaration, 0 hue
                            T += [A.ColorJitter(*color_jitter, 0)]
                else:  # Use fixed crop for eval set (reproducibility)
                    T = [A.SmallestMaxSize(max_size=size), A.CenterCrop(height=size, width=size)]
                T += [A.Normalize(mean=mean, std=std), ToTensorV2()]  # Normalize and convert to Tensor
                LOGGER.info(prefix + ", ".join(f"{x}".replace("always_apply=False, ", "") for x in T if x.p))
                return A.Compose(T)
        
            except ImportError:  # package not installed, skip
                LOGGER.warning(f"{prefix}⚠️ not found, install with `pip install albumentations` (recommended)")
            except Exception as e:
                LOGGER.info(f"{prefix}{e}")
        
        
        def classify_transforms(size=224):
            """Applies a series of transformations including center crop, ToTensor, and normalization for classification."""
            assert isinstance(size, int), f"ERROR: classify_transforms size {size} must be integer, not (list, tuple)"
            # T.Compose([T.ToTensor(), T.Resize(size), T.CenterCrop(size), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
            return T.Compose([CenterCrop(size), ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
        
        
        class LetterBox:
            # YOLOv5 LetterBox class for image preprocessing, i.e. T.Compose([LetterBox(size), ToTensor()])
            def __init__(self, size=(640, 640), auto=False, stride=32):
                """Initializes a LetterBox object for YOLOv5 image preprocessing with optional auto sizing and stride
                adjustment.
                """
                super().__init__()
                self.h, self.w = (size, size) if isinstance(size, int) else size
                self.auto = auto  # pass max size integer, automatically solve for short side using stride
                self.stride = stride  # used with auto
        
            def __call__(self, im):
                """
                Resizes and pads input image `im` (HWC format) to specified dimensions, maintaining aspect ratio.
        
                im = np.array HWC
                """
                imh, imw = im.shape[:2]
                r = min(self.h / imh, self.w / imw)  # ratio of new/old
                h, w = round(imh * r), round(imw * r)  # resized image
                hs, ws = (math.ceil(x / self.stride) * self.stride for x in (h, w)) if self.auto else self.h, self.w
                top, left = round((hs - h) / 2 - 0.1), round((ws - w) / 2 - 0.1)
                im_out = np.full((self.h, self.w, 3), 114, dtype=im.dtype)
                im_out[top : top + h, left : left + w] = cv2.resize(im, (w, h), interpolation=cv2.INTER_LINEAR)
                return im_out
        
        
        class CenterCrop:
            # YOLOv5 CenterCrop class for image preprocessing, i.e. T.Compose([CenterCrop(size), ToTensor()])
            def __init__(self, size=640):
                """Initializes CenterCrop for image preprocessing, accepting single int or tuple for size, defaults to 640."""
                super().__init__()
                self.h, self.w = (size, size) if isinstance(size, int) else size
        
            def __call__(self, im):
                """
                Applies center crop to the input image and resizes it to a specified size, maintaining aspect ratio.
        
                im = np.array HWC
                """
                imh, imw = im.shape[:2]
                m = min(imh, imw)  # min dimension
                top, left = (imh - m) // 2, (imw - m) // 2
                return cv2.resize(im[top : top + m, left : left + m], (self.w, self.h), interpolation=cv2.INTER_LINEAR)
        
        
        class ToTensor:
            # YOLOv5 ToTensor class for image preprocessing, i.e. T.Compose([LetterBox(size), ToTensor()])
            def __init__(self, half=False):
                """Initializes ToTensor for YOLOv5 image preprocessing, with optional half precision (half=True for FP16)."""
                super().__init__()
                self.half = half
        
            def __call__(self, im):
                """
                Converts BGR np.array image from HWC to RGB CHW format, and normalizes to [0, 1], with support for FP16 if
                `half=True`.
        
                im = np.array HWC in BGR order
                """
                im = np.ascontiguousarray(im.transpose((2, 0, 1))[::-1])  # HWC to CHW -> BGR to RGB -> contiguous
                im = torch.from_numpy(im)  # to torch
                im = im.half() if self.half else im.float()  # uint8 to fp16/32
                im /= 255.0  # 0-255 to 0.0-1.0
                return im
        
      • 📄 autoanchor.py 7.3KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """AutoAnchor utils."""
        
        import random
        
        import numpy as np
        import torch
        import yaml
        from tqdm import tqdm
        
        from utils import TryExcept
        from utils.general import LOGGER, TQDM_BAR_FORMAT, colorstr
        
        PREFIX = colorstr("AutoAnchor: ")
        
        
        def check_anchor_order(m):
            """Checks and corrects anchor order against stride in YOLOv5 Detect() module if necessary."""
            a = m.anchors.prod(-1).mean(-1).view(-1)  # mean anchor area per output layer
            da = a[-1] - a[0]  # delta a
            ds = m.stride[-1] - m.stride[0]  # delta s
            if da and (da.sign() != ds.sign()):  # same order
                LOGGER.info(f"{PREFIX}Reversing anchor order")
                m.anchors[:] = m.anchors.flip(0)
        
        
        @TryExcept(f"{PREFIX}ERROR")
        def check_anchors(dataset, model, thr=4.0, imgsz=640):
            """Evaluates anchor fit to dataset and adjusts if necessary, supporting customizable threshold and image size."""
            m = model.module.model[-1] if hasattr(model, "module") else model.model[-1]  # Detect()
            shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True)
            scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1))  # augment scale
            wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float()  # wh
        
            def metric(k):  # compute metric
                r = wh[:, None] / k[None]
                x = torch.min(r, 1 / r).min(2)[0]  # ratio metric
                best = x.max(1)[0]  # best_x
                aat = (x > 1 / thr).float().sum(1).mean()  # anchors above threshold
                bpr = (best > 1 / thr).float().mean()  # best possible recall
                return bpr, aat
        
            stride = m.stride.to(m.anchors.device).view(-1, 1, 1)  # model strides
            anchors = m.anchors.clone() * stride  # current anchors
            bpr, aat = metric(anchors.cpu().view(-1, 2))
            s = f"\n{PREFIX}{aat:.2f} anchors/target, {bpr:.3f} Best Possible Recall (BPR). "
            if bpr > 0.98:  # threshold to recompute
                LOGGER.info(f"{s}Current anchors are a good fit to dataset ✅")
            else:
                LOGGER.info(f"{s}Anchors are a poor fit to dataset ⚠️, attempting to improve...")
                na = m.anchors.numel() // 2  # number of anchors
                anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
                new_bpr = metric(anchors)[0]
                if new_bpr > bpr:  # replace anchors
                    anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors)
                    m.anchors[:] = anchors.clone().view_as(m.anchors)
                    check_anchor_order(m)  # must be in pixel-space (not grid-space)
                    m.anchors /= stride
                    s = f"{PREFIX}Done ✅ (optional: update model *.yaml to use these anchors in the future)"
                else:
                    s = f"{PREFIX}Done ⚠️ (original anchors better than new anchors, proceeding with original anchors)"
                LOGGER.info(s)
        
        
        def kmean_anchors(dataset="./data/coco128.yaml", n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
            """
            Creates kmeans-evolved anchors from training dataset.
        
            Arguments:
                dataset: path to data.yaml, or a loaded dataset
                n: number of anchors
                img_size: image size used for training
                thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0
                gen: generations to evolve anchors using genetic algorithm
                verbose: print all results
        
            Return:
                k: kmeans evolved anchors
        
            Usage:
                from utils.autoanchor import *; _ = kmean_anchors()
            """
            from scipy.cluster.vq import kmeans
        
            npr = np.random
            thr = 1 / thr
        
            def metric(k, wh):  # compute metrics
                r = wh[:, None] / k[None]
                x = torch.min(r, 1 / r).min(2)[0]  # ratio metric
                # x = wh_iou(wh, torch.tensor(k))  # iou metric
                return x, x.max(1)[0]  # x, best_x
        
            def anchor_fitness(k):  # mutation fitness
                _, best = metric(torch.tensor(k, dtype=torch.float32), wh)
                return (best * (best > thr).float()).mean()  # fitness
        
            def print_results(k, verbose=True):
                k = k[np.argsort(k.prod(1))]  # sort small to large
                x, best = metric(k, wh0)
                bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n  # best possible recall, anch > thr
                s = (
                    f"{PREFIX}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr\n"
                    f"{PREFIX}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, "
                    f"past_thr={x[x > thr].mean():.3f}-mean: "
                )
                for x in k:
                    s += "%i,%i, " % (round(x[0]), round(x[1]))
                if verbose:
                    LOGGER.info(s[:-2])
                return k
        
            if isinstance(dataset, str):  # *.yaml file
                with open(dataset, errors="ignore") as f:
                    data_dict = yaml.safe_load(f)  # model dict
                from utils.dataloaders import LoadImagesAndLabels
        
                dataset = LoadImagesAndLabels(data_dict["train"], augment=True, rect=True)
        
            # Get label wh
            shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True)
            wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)])  # wh
        
            # Filter
            i = (wh0 < 3.0).any(1).sum()
            if i:
                LOGGER.info(f"{PREFIX}WARNING ⚠️ Extremely small objects found: {i} of {len(wh0)} labels are <3 pixels in size")
            wh = wh0[(wh0 >= 2.0).any(1)].astype(np.float32)  # filter > 2 pixels
            # wh = wh * (npr.rand(wh.shape[0], 1) * 0.9 + 0.1)  # multiply by random scale 0-1
        
            # Kmeans init
            try:
                LOGGER.info(f"{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...")
                assert n <= len(wh)  # apply overdetermined constraint
                s = wh.std(0)  # sigmas for whitening
                k = kmeans(wh / s, n, iter=30)[0] * s  # points
                assert n == len(k)  # kmeans may return fewer points than requested if wh is insufficient or too similar
            except Exception:
                LOGGER.warning(f"{PREFIX}WARNING ⚠️ switching strategies from kmeans to random init")
                k = np.sort(npr.rand(n * 2)).reshape(n, 2) * img_size  # random init
            wh, wh0 = (torch.tensor(x, dtype=torch.float32) for x in (wh, wh0))
            k = print_results(k, verbose=False)
        
            # Plot
            # k, d = [None] * 20, [None] * 20
            # for i in tqdm(range(1, 21)):
            #     k[i-1], d[i-1] = kmeans(wh / s, i)  # points, mean distance
            # fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True)
            # ax = ax.ravel()
            # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.')
            # fig, ax = plt.subplots(1, 2, figsize=(14, 7))  # plot wh
            # ax[0].hist(wh[wh[:, 0]<100, 0],400)
            # ax[1].hist(wh[wh[:, 1]<100, 1],400)
            # fig.savefig('wh.png', dpi=200)
        
            # Evolve
            f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1  # fitness, generations, mutation prob, sigma
            pbar = tqdm(range(gen), bar_format=TQDM_BAR_FORMAT)  # progress bar
            for _ in pbar:
                v = np.ones(sh)
                while (v == 1).all():  # mutate until a change occurs (prevent duplicates)
                    v = ((npr.random(sh) < mp) * random.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0)
                kg = (k.copy() * v).clip(min=2.0)
                fg = anchor_fitness(kg)
                if fg > f:
                    f, k = fg, kg.copy()
                    pbar.desc = f"{PREFIX}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}"
                    if verbose:
                        print_results(k, verbose)
        
            return print_results(k).astype(np.float32)
        
      • 📄 autobatch.py 3.0KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Auto-batch utils."""
        
        from copy import deepcopy
        
        import numpy as np
        import torch
        
        from utils.general import LOGGER, colorstr
        from utils.torch_utils import profile
        
        
        def check_train_batch_size(model, imgsz=640, amp=True):
            """Checks and computes optimal training batch size for YOLOv5 model, given image size and AMP setting."""
            with torch.cuda.amp.autocast(amp):
                return autobatch(deepcopy(model).train(), imgsz)  # compute optimal batch size
        
        
        def autobatch(model, imgsz=640, fraction=0.8, batch_size=16):
            """Estimates optimal YOLOv5 batch size using `fraction` of CUDA memory."""
            # Usage:
            #     import torch
            #     from utils.autobatch import autobatch
            #     model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False)
            #     print(autobatch(model))
        
            # Check device
            prefix = colorstr("AutoBatch: ")
            LOGGER.info(f"{prefix}Computing optimal batch size for --imgsz {imgsz}")
            device = next(model.parameters()).device  # get model device
            if device.type == "cpu":
                LOGGER.info(f"{prefix}CUDA not detected, using default CPU batch-size {batch_size}")
                return batch_size
            if torch.backends.cudnn.benchmark:
                LOGGER.info(f"{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}")
                return batch_size
        
            # Inspect CUDA memory
            gb = 1 << 30  # bytes to GiB (1024 ** 3)
            d = str(device).upper()  # 'CUDA:0'
            properties = torch.cuda.get_device_properties(device)  # device properties
            t = properties.total_memory / gb  # GiB total
            r = torch.cuda.memory_reserved(device) / gb  # GiB reserved
            a = torch.cuda.memory_allocated(device) / gb  # GiB allocated
            f = t - (r + a)  # GiB free
            LOGGER.info(f"{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free")
        
            # Profile batch sizes
            batch_sizes = [1, 2, 4, 8, 16]
            try:
                img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes]
                results = profile(img, model, n=3, device=device)
            except Exception as e:
                LOGGER.warning(f"{prefix}{e}")
        
            # Fit a solution
            y = [x[2] for x in results if x]  # memory [2]
            p = np.polyfit(batch_sizes[: len(y)], y, deg=1)  # first degree polynomial fit
            b = int((f * fraction - p[1]) / p[0])  # y intercept (optimal batch size)
            if None in results:  # some sizes failed
                i = results.index(None)  # first fail index
                if b >= batch_sizes[i]:  # y intercept above failure point
                    b = batch_sizes[max(i - 1, 0)]  # select prior safe point
            if b < 1 or b > 1024:  # b outside of safe range
                b = batch_size
                LOGGER.warning(f"{prefix}WARNING ⚠️ CUDA anomaly detected, recommend restart environment and retry command.")
        
            fraction = (np.polyval(p, b) + r + a) / t  # actual fraction predicted
            LOGGER.info(f"{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅")
            return b
        
      • 📄 callbacks.py 2.7KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Callback utils."""
        
        import threading
        
        
        class Callbacks:
            """Handles all registered callbacks for YOLOv5 Hooks."""
        
            def __init__(self):
                """Initializes a Callbacks object to manage registered YOLOv5 training event hooks."""
                self._callbacks = {
                    "on_pretrain_routine_start": [],
                    "on_pretrain_routine_end": [],
                    "on_train_start": [],
                    "on_train_epoch_start": [],
                    "on_train_batch_start": [],
                    "optimizer_step": [],
                    "on_before_zero_grad": [],
                    "on_train_batch_end": [],
                    "on_train_epoch_end": [],
                    "on_val_start": [],
                    "on_val_batch_start": [],
                    "on_val_image_end": [],
                    "on_val_batch_end": [],
                    "on_val_end": [],
                    "on_fit_epoch_end": [],  # fit = train + val
                    "on_model_save": [],
                    "on_train_end": [],
                    "on_params_update": [],
                    "teardown": [],
                }
                self.stop_training = False  # set True to interrupt training
        
            def register_action(self, hook, name="", callback=None):
                """
                Register a new action to a callback hook.
        
                Args:
                    hook: The callback hook name to register the action to
                    name: The name of the action for later reference
                    callback: The callback to fire
                """
                assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
                assert callable(callback), f"callback '{callback}' is not callable"
                self._callbacks[hook].append({"name": name, "callback": callback})
        
            def get_registered_actions(self, hook=None):
                """
                Returns all the registered actions by callback hook.
        
                Args:
                    hook: The name of the hook to check, defaults to all
                """
                return self._callbacks[hook] if hook else self._callbacks
        
            def run(self, hook, *args, thread=False, **kwargs):
                """
                Loop through the registered actions and fire all callbacks on main thread.
        
                Args:
                    hook: The name of the hook to check, defaults to all
                    args: Arguments to receive from YOLOv5
                    thread: (boolean) Run callbacks in daemon thread
                    kwargs: Keyword Arguments to receive from YOLOv5
                """
        
                assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
                for logger in self._callbacks[hook]:
                    if thread:
                        threading.Thread(target=logger["callback"], args=args, kwargs=kwargs, daemon=True).start()
                    else:
                        logger["callback"](*args, **kwargs)
        
      • 📄 dataloaders.py 58.6KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Dataloaders and dataset utils."""
        
        import contextlib
        import glob
        import hashlib
        import json
        import math
        import os
        import random
        import shutil
        import time
        from itertools import repeat
        from multiprocessing.pool import Pool, ThreadPool
        from pathlib import Path
        from threading import Thread
        from urllib.parse import urlparse
        
        import numpy as np
        import psutil
        import torch
        import torch.nn.functional as F
        import torchvision
        import yaml
        from PIL import ExifTags, Image, ImageOps
        from torch.utils.data import DataLoader, Dataset, dataloader, distributed
        from tqdm import tqdm
        
        from utils.augmentations import (
            Albumentations,
            augment_hsv,
            classify_albumentations,
            classify_transforms,
            copy_paste,
            letterbox,
            mixup,
            random_perspective,
        )
        from utils.general import (
            DATASETS_DIR,
            LOGGER,
            NUM_THREADS,
            TQDM_BAR_FORMAT,
            check_dataset,
            check_requirements,
            check_yaml,
            clean_str,
            cv2,
            is_colab,
            is_kaggle,
            segments2boxes,
            unzip_file,
            xyn2xy,
            xywh2xyxy,
            xywhn2xyxy,
            xyxy2xywhn,
        )
        from utils.torch_utils import torch_distributed_zero_first
        
        # Parameters
        HELP_URL = "See https://docs.ultralytics.com/yolov5/tutorials/train_custom_data"
        IMG_FORMATS = "bmp", "dng", "jpeg", "jpg", "mpo", "png", "tif", "tiff", "webp", "pfm"  # include image suffixes
        VID_FORMATS = "asf", "avi", "gif", "m4v", "mkv", "mov", "mp4", "mpeg", "mpg", "ts", "wmv"  # include video suffixes
        LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))  # https://pytorch.org/docs/stable/elastic/run.html
        RANK = int(os.getenv("RANK", -1))
        WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
        PIN_MEMORY = str(os.getenv("PIN_MEMORY", True)).lower() == "true"  # global pin_memory for dataloaders
        
        # Get orientation exif tag
        for orientation in ExifTags.TAGS.keys():
            if ExifTags.TAGS[orientation] == "Orientation":
                break
        
        
        def get_hash(paths):
            """Generates a single SHA256 hash for a list of file or directory paths by combining their sizes and paths."""
            size = sum(os.path.getsize(p) for p in paths if os.path.exists(p))  # sizes
            h = hashlib.sha256(str(size).encode())  # hash sizes
            h.update("".join(paths).encode())  # hash paths
            return h.hexdigest()  # return hash
        
        
        def exif_size(img):
            """Returns corrected PIL image size (width, height) considering EXIF orientation."""
            s = img.size  # (width, height)
            with contextlib.suppress(Exception):
                rotation = dict(img._getexif().items())[orientation]
                if rotation in [6, 8]:  # rotation 270 or 90
                    s = (s[1], s[0])
            return s
        
        
        def exif_transpose(image):
            """
            Transpose a PIL image accordingly if it has an EXIF Orientation tag.
            Inplace version of https://github.com/python-pillow/Pillow/blob/master/src/PIL/ImageOps.py exif_transpose()
        
            :param image: The image to transpose.
            :return: An image.
            """
            exif = image.getexif()
            orientation = exif.get(0x0112, 1)  # default 1
            if orientation > 1:
                method = {
                    2: Image.FLIP_LEFT_RIGHT,
                    3: Image.ROTATE_180,
                    4: Image.FLIP_TOP_BOTTOM,
                    5: Image.TRANSPOSE,
                    6: Image.ROTATE_270,
                    7: Image.TRANSVERSE,
                    8: Image.ROTATE_90,
                }.get(orientation)
                if method is not None:
                    image = image.transpose(method)
                    del exif[0x0112]
                    image.info["exif"] = exif.tobytes()
            return image
        
        
        def seed_worker(worker_id):
            """
            Sets the seed for a dataloader worker to ensure reproducibility, based on PyTorch's randomness notes.
        
            See https://pytorch.org/docs/stable/notes/randomness.html#dataloader.
            """
            worker_seed = torch.initial_seed() % 2**32
            np.random.seed(worker_seed)
            random.seed(worker_seed)
        
        
        # Inherit from DistributedSampler and override iterator
        # https://github.com/pytorch/pytorch/blob/master/torch/utils/data/distributed.py
        class SmartDistributedSampler(distributed.DistributedSampler):
            def __iter__(self):
                """Yields indices for distributed data sampling, shuffled deterministically based on epoch and seed."""
                g = torch.Generator()
                g.manual_seed(self.seed + self.epoch)
        
                # determine the the eventual size (n) of self.indices (DDP indices)
                n = int((len(self.dataset) - self.rank - 1) / self.num_replicas) + 1  # num_replicas == WORLD_SIZE
                idx = torch.randperm(n, generator=g)
                if not self.shuffle:
                    idx = idx.sort()[0]
        
                idx = idx.tolist()
                if self.drop_last:
                    idx = idx[: self.num_samples]
                else:
                    padding_size = self.num_samples - len(idx)
                    if padding_size <= len(idx):
                        idx += idx[:padding_size]
                    else:
                        idx += (idx * math.ceil(padding_size / len(idx)))[:padding_size]
        
                return iter(idx)
        
        
        def create_dataloader(
            path,
            imgsz,
            batch_size,
            stride,
            single_cls=False,
            hyp=None,
            augment=False,
            cache=False,
            pad=0.0,
            rect=False,
            rank=-1,
            workers=8,
            image_weights=False,
            quad=False,
            prefix="",
            shuffle=False,
            seed=0,
        ):
            if rect and shuffle:
                LOGGER.warning("WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False")
                shuffle = False
            with torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDP
                dataset = LoadImagesAndLabels(
                    path,
                    imgsz,
                    batch_size,
                    augment=augment,  # augmentation
                    hyp=hyp,  # hyperparameters
                    rect=rect,  # rectangular batches
                    cache_images=cache,
                    single_cls=single_cls,
                    stride=int(stride),
                    pad=pad,
                    image_weights=image_weights,
                    prefix=prefix,
                    rank=rank,
                )
        
            batch_size = min(batch_size, len(dataset))
            nd = torch.cuda.device_count()  # number of CUDA devices
            nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers])  # number of workers
            sampler = None if rank == -1 else SmartDistributedSampler(dataset, shuffle=shuffle)
            loader = DataLoader if image_weights else InfiniteDataLoader  # only DataLoader allows for attribute updates
            generator = torch.Generator()
            generator.manual_seed(6148914691236517205 + seed + RANK)
            return loader(
                dataset,
                batch_size=batch_size,
                shuffle=shuffle and sampler is None,
                num_workers=nw,
                sampler=sampler,
                pin_memory=PIN_MEMORY,
                collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn,
                worker_init_fn=seed_worker,
                generator=generator,
            ), dataset
        
        
        class InfiniteDataLoader(dataloader.DataLoader):
            """
            Dataloader that reuses workers.
        
            Uses same syntax as vanilla DataLoader
            """
        
            def __init__(self, *args, **kwargs):
                """Initializes an InfiniteDataLoader that reuses workers with standard DataLoader syntax, augmenting with a
                repeating sampler.
                """
                super().__init__(*args, **kwargs)
                object.__setattr__(self, "batch_sampler", _RepeatSampler(self.batch_sampler))
                self.iterator = super().__iter__()
        
            def __len__(self):
                """Returns the length of the batch sampler's sampler in the InfiniteDataLoader."""
                return len(self.batch_sampler.sampler)
        
            def __iter__(self):
                """Yields batches of data indefinitely in a loop by resetting the sampler when exhausted."""
                for _ in range(len(self)):
                    yield next(self.iterator)
        
        
        class _RepeatSampler:
            """
            Sampler that repeats forever.
        
            Args:
                sampler (Sampler)
            """
        
            def __init__(self, sampler):
                """Initializes a perpetual sampler wrapping a provided `Sampler` instance for endless data iteration."""
                self.sampler = sampler
        
            def __iter__(self):
                """Returns an infinite iterator over the dataset by repeatedly yielding from the given sampler."""
                while True:
                    yield from iter(self.sampler)
        
        
        class LoadScreenshots:
            # YOLOv5 screenshot dataloader, i.e. `python detect.py --source "screen 0 100 100 512 256"`
            def __init__(self, source, img_size=640, stride=32, auto=True, transforms=None):
                """
                Initializes a screenshot dataloader for YOLOv5 with specified source region, image size, stride, auto, and
                transforms.
        
                Source = [screen_number left top width height] (pixels)
                """
                check_requirements("mss")
                import mss
        
                source, *params = source.split()
                self.screen, left, top, width, height = 0, None, None, None, None  # default to full screen 0
                if len(params) == 1:
                    self.screen = int(params[0])
                elif len(params) == 4:
                    left, top, width, height = (int(x) for x in params)
                elif len(params) == 5:
                    self.screen, left, top, width, height = (int(x) for x in params)
                self.img_size = img_size
                self.stride = stride
                self.transforms = transforms
                self.auto = auto
                self.mode = "stream"
                self.frame = 0
                self.sct = mss.mss()
        
                # Parse monitor shape
                monitor = self.sct.monitors[self.screen]
                self.top = monitor["top"] if top is None else (monitor["top"] + top)
                self.left = monitor["left"] if left is None else (monitor["left"] + left)
                self.width = width or monitor["width"]
                self.height = height or monitor["height"]
                self.monitor = {"left": self.left, "top": self.top, "width": self.width, "height": self.height}
        
            def __iter__(self):
                """Iterates over itself, enabling use in loops and iterable contexts."""
                return self
        
            def __next__(self):
                """Captures and returns the next screen frame as a BGR numpy array, cropping to only the first three channels
                from BGRA.
                """
                im0 = np.array(self.sct.grab(self.monitor))[:, :, :3]  # [:, :, :3] BGRA to BGR
                s = f"screen {self.screen} (LTWH): {self.left},{self.top},{self.width},{self.height}: "
        
                if self.transforms:
                    im = self.transforms(im0)  # transforms
                else:
                    im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0]  # padded resize
                    im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
                    im = np.ascontiguousarray(im)  # contiguous
                self.frame += 1
                return str(self.screen), im, im0, None, s  # screen, img, original img, im0s, s
        
        
        class LoadImages:
            """YOLOv5 image/video dataloader, i.e. `python detect.py --source image.jpg/vid.mp4`"""
        
            def __init__(self, path, img_size=640, stride=32, auto=True, transforms=None, vid_stride=1):
                """Initializes YOLOv5 loader for images/videos, supporting glob patterns, directories, and lists of paths."""
                if isinstance(path, str) and Path(path).suffix == ".txt":  # *.txt file with img/vid/dir on each line
                    path = Path(path).read_text().rsplit()
                files = []
                for p in sorted(path) if isinstance(path, (list, tuple)) else [path]:
                    p = str(Path(p).resolve())
                    if "*" in p:
                        files.extend(sorted(glob.glob(p, recursive=True)))  # glob
                    elif os.path.isdir(p):
                        files.extend(sorted(glob.glob(os.path.join(p, "*.*"))))  # dir
                    elif os.path.isfile(p):
                        files.append(p)  # files
                    else:
                        raise FileNotFoundError(f"{p} does not exist")
        
                images = [x for x in files if x.split(".")[-1].lower() in IMG_FORMATS]
                videos = [x for x in files if x.split(".")[-1].lower() in VID_FORMATS]
                ni, nv = len(images), len(videos)
        
                self.img_size = img_size
                self.stride = stride
                self.files = images + videos
                self.nf = ni + nv  # number of files
                self.video_flag = [False] * ni + [True] * nv
                self.mode = "image"
                self.auto = auto
                self.transforms = transforms  # optional
                self.vid_stride = vid_stride  # video frame-rate stride
                if any(videos):
                    self._new_video(videos[0])  # new video
                else:
                    self.cap = None
                assert self.nf > 0, (
                    f"No images or videos found in {p}. "
                    f"Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}"
                )
        
            def __iter__(self):
                """Initializes iterator by resetting count and returns the iterator object itself."""
                self.count = 0
                return self
        
            def __next__(self):
                """Advances to the next file in the dataset, raising StopIteration if at the end."""
                if self.count == self.nf:
                    raise StopIteration
                path = self.files[self.count]
        
                if self.video_flag[self.count]:
                    # Read video
                    self.mode = "video"
                    for _ in range(self.vid_stride):
                        self.cap.grab()
                    ret_val, im0 = self.cap.retrieve()
                    while not ret_val:
                        self.count += 1
                        self.cap.release()
                        if self.count == self.nf:  # last video
                            raise StopIteration
                        path = self.files[self.count]
                        self._new_video(path)
                        ret_val, im0 = self.cap.read()
        
                    self.frame += 1
                    # im0 = self._cv2_rotate(im0)  # for use if cv2 autorotation is False
                    s = f"video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: "
        
                else:
                    # Read image
                    self.count += 1
                    im0 = cv2.imread(path)  # BGR
                    assert im0 is not None, f"Image Not Found {path}"
                    s = f"image {self.count}/{self.nf} {path}: "
        
                if self.transforms:
                    im = self.transforms(im0)  # transforms
                else:
                    im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0]  # padded resize
                    im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
                    im = np.ascontiguousarray(im)  # contiguous
        
                return path, im, im0, self.cap, s
        
            def _new_video(self, path):
                """Initializes a new video capture object with path, frame count adjusted by stride, and orientation
                metadata.
                """
                self.frame = 0
                self.cap = cv2.VideoCapture(path)
                self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT) / self.vid_stride)
                self.orientation = int(self.cap.get(cv2.CAP_PROP_ORIENTATION_META))  # rotation degrees
                # self.cap.set(cv2.CAP_PROP_ORIENTATION_AUTO, 0)  # disable https://github.com/ultralytics/yolov5/issues/8493
        
            def _cv2_rotate(self, im):
                """Rotates a cv2 image based on its orientation; supports 0, 90, and 180 degrees rotations."""
                if self.orientation == 0:
                    return cv2.rotate(im, cv2.ROTATE_90_CLOCKWISE)
                elif self.orientation == 180:
                    return cv2.rotate(im, cv2.ROTATE_90_COUNTERCLOCKWISE)
                elif self.orientation == 90:
                    return cv2.rotate(im, cv2.ROTATE_180)
                return im
        
            def __len__(self):
                """Returns the number of files in the dataset."""
                return self.nf  # number of files
        
        
        class LoadStreams:
            # YOLOv5 streamloader, i.e. `python detect.py --source 'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP streams`
            def __init__(self, sources="file.streams", img_size=640, stride=32, auto=True, transforms=None, vid_stride=1):
                """Initializes a stream loader for processing video streams with YOLOv5, supporting various sources including
                YouTube.
                """
                torch.backends.cudnn.benchmark = True  # faster for fixed-size inference
                self.mode = "stream"
                self.img_size = img_size
                self.stride = stride
                self.vid_stride = vid_stride  # video frame-rate stride
                sources = Path(sources).read_text().rsplit() if os.path.isfile(sources) else [sources]
                n = len(sources)
                self.sources = [clean_str(x) for x in sources]  # clean source names for later
                self.imgs, self.fps, self.frames, self.threads = [None] * n, [0] * n, [0] * n, [None] * n
                for i, s in enumerate(sources):  # index, source
                    # Start thread to read frames from video stream
                    st = f"{i + 1}/{n}: {s}... "
                    if urlparse(s).hostname in ("www.youtube.com", "youtube.com", "youtu.be"):  # if source is YouTube video
                        # YouTube format i.e. 'https://www.youtube.com/watch?v=Zgi9g1ksQHc' or 'https://youtu.be/LNwODJXcvt4'
                        check_requirements(("pafy", "youtube_dl==2020.12.2"))
                        import pafy
        
                        s = pafy.new(s).getbest(preftype="mp4").url  # YouTube URL
                    s = eval(s) if s.isnumeric() else s  # i.e. s = '0' local webcam
                    if s == 0:
                        assert not is_colab(), "--source 0 webcam unsupported on Colab. Rerun command in a local environment."
                        assert not is_kaggle(), "--source 0 webcam unsupported on Kaggle. Rerun command in a local environment."
                    cap = cv2.VideoCapture(s)
                    assert cap.isOpened(), f"{st}Failed to open {s}"
                    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                    fps = cap.get(cv2.CAP_PROP_FPS)  # warning: may return 0 or nan
                    self.frames[i] = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float("inf")  # infinite stream fallback
                    self.fps[i] = max((fps if math.isfinite(fps) else 0) % 100, 0) or 30  # 30 FPS fallback
        
                    _, self.imgs[i] = cap.read()  # guarantee first frame
                    self.threads[i] = Thread(target=self.update, args=([i, cap, s]), daemon=True)
                    LOGGER.info(f"{st} Success ({self.frames[i]} frames {w}x{h} at {self.fps[i]:.2f} FPS)")
                    self.threads[i].start()
                LOGGER.info("")  # newline
        
                # check for common shapes
                s = np.stack([letterbox(x, img_size, stride=stride, auto=auto)[0].shape for x in self.imgs])
                self.rect = np.unique(s, axis=0).shape[0] == 1  # rect inference if all shapes equal
                self.auto = auto and self.rect
                self.transforms = transforms  # optional
                if not self.rect:
                    LOGGER.warning("WARNING ⚠️ Stream shapes differ. For optimal performance supply similarly-shaped streams.")
        
            def update(self, i, cap, stream):
                """Reads frames from stream `i`, updating imgs array; handles stream reopening on signal loss."""
                n, f = 0, self.frames[i]  # frame number, frame array
                while cap.isOpened() and n < f:
                    n += 1
                    cap.grab()  # .read() = .grab() followed by .retrieve()
                    if n % self.vid_stride == 0:
                        success, im = cap.retrieve()
                        if success:
                            self.imgs[i] = im
                        else:
                            LOGGER.warning("WARNING ⚠️ Video stream unresponsive, please check your IP camera connection.")
                            self.imgs[i] = np.zeros_like(self.imgs[i])
                            cap.open(stream)  # re-open stream if signal was lost
                    time.sleep(0.0)  # wait time
        
            def __iter__(self):
                """Resets and returns the iterator for iterating over video frames or images in a dataset."""
                self.count = -1
                return self
        
            def __next__(self):
                """Iterates over video frames or images, halting on thread stop or 'q' key press, raising `StopIteration` when
                done.
                """
                self.count += 1
                if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord("q"):  # q to quit
                    cv2.destroyAllWindows()
                    raise StopIteration
        
                im0 = self.imgs.copy()
                if self.transforms:
                    im = np.stack([self.transforms(x) for x in im0])  # transforms
                else:
                    im = np.stack([letterbox(x, self.img_size, stride=self.stride, auto=self.auto)[0] for x in im0])  # resize
                    im = im[..., ::-1].transpose((0, 3, 1, 2))  # BGR to RGB, BHWC to BCHW
                    im = np.ascontiguousarray(im)  # contiguous
        
                return self.sources, im, im0, None, ""
        
            def __len__(self):
                """Returns the number of sources in the dataset, supporting up to 32 streams at 30 FPS over 30 years."""
                return len(self.sources)  # 1E12 frames = 32 streams at 30 FPS for 30 years
        
        
        def img2label_paths(img_paths):
            """Generates label file paths from corresponding image file paths by replacing `/images/` with `/labels/` and
            extension with `.txt`.
            """
            sa, sb = f"{os.sep}images{os.sep}", f"{os.sep}labels{os.sep}"  # /images/, /labels/ substrings
            return [sb.join(x.rsplit(sa, 1)).rsplit(".", 1)[0] + ".txt" for x in img_paths]
        
        
        class LoadImagesAndLabels(Dataset):
            # YOLOv5 train_loader/val_loader, loads images and labels for training and validation
            cache_version = 0.6  # dataset labels *.cache version
            rand_interp_methods = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4]
        
            def __init__(
                self,
                path,
                img_size=640,
                batch_size=16,
                augment=False,
                hyp=None,
                rect=False,
                image_weights=False,
                cache_images=False,
                single_cls=False,
                stride=32,
                pad=0.0,
                min_items=0,
                prefix="",
                rank=-1,
                seed=0,
            ):
                self.img_size = img_size
                self.augment = augment
                self.hyp = hyp
                self.image_weights = image_weights
                self.rect = False if image_weights else rect
                self.mosaic = self.augment and not self.rect  # load 4 images at a time into a mosaic (only during training)
                self.mosaic_border = [-img_size // 2, -img_size // 2]
                self.stride = stride
                self.path = path
                self.albumentations = Albumentations(size=img_size) if augment else None
        
                try:
                    f = []  # image files
                    for p in path if isinstance(path, list) else [path]:
                        p = Path(p)  # os-agnostic
                        if p.is_dir():  # dir
                            f += glob.glob(str(p / "**" / "*.*"), recursive=True)
                            # f = list(p.rglob('*.*'))  # pathlib
                        elif p.is_file():  # file
                            with open(p) as t:
                                t = t.read().strip().splitlines()
                                parent = str(p.parent) + os.sep
                                f += [x.replace("./", parent, 1) if x.startswith("./") else x for x in t]  # to global path
                                # f += [p.parent / x.lstrip(os.sep) for x in t]  # to global path (pathlib)
                        else:
                            raise FileNotFoundError(f"{prefix}{p} does not exist")
                    self.im_files = sorted(x.replace("/", os.sep) for x in f if x.split(".")[-1].lower() in IMG_FORMATS)
                    # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS])  # pathlib
                    assert self.im_files, f"{prefix}No images found"
                except Exception as e:
                    raise Exception(f"{prefix}Error loading data from {path}: {e}\n{HELP_URL}") from e
        
                # Check cache
                self.label_files = img2label_paths(self.im_files)  # labels
                cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix(".cache")
                try:
                    cache, exists = np.load(cache_path, allow_pickle=True).item(), True  # load dict
                    assert cache["version"] == self.cache_version  # matches current version
                    assert cache["hash"] == get_hash(self.label_files + self.im_files)  # identical hash
                except Exception:
                    cache, exists = self.cache_labels(cache_path, prefix), False  # run cache ops
        
                # Display cache
                nf, nm, ne, nc, n = cache.pop("results")  # found, missing, empty, corrupt, total
                if exists and LOCAL_RANK in {-1, 0}:
                    d = f"Scanning {cache_path}... {nf} images, {nm + ne} backgrounds, {nc} corrupt"
                    tqdm(None, desc=prefix + d, total=n, initial=n, bar_format=TQDM_BAR_FORMAT)  # display cache results
                    if cache["msgs"]:
                        LOGGER.info("\n".join(cache["msgs"]))  # display warnings
                assert nf > 0 or not augment, f"{prefix}No labels found in {cache_path}, can not start training. {HELP_URL}"
        
                # Read cache
                [cache.pop(k) for k in ("hash", "version", "msgs")]  # remove items
                labels, shapes, self.segments = zip(*cache.values())
                nl = len(np.concatenate(labels, 0))  # number of labels
                assert nl > 0 or not augment, f"{prefix}All labels empty in {cache_path}, can not start training. {HELP_URL}"
                self.labels = list(labels)
                self.shapes = np.array(shapes)
                self.im_files = list(cache.keys())  # update
                self.label_files = img2label_paths(cache.keys())  # update
        
                # Filter images
                if min_items:
                    include = np.array([len(x) >= min_items for x in self.labels]).nonzero()[0].astype(int)
                    LOGGER.info(f"{prefix}{n - len(include)}/{n} images filtered from dataset")
                    self.im_files = [self.im_files[i] for i in include]
                    self.label_files = [self.label_files[i] for i in include]
                    self.labels = [self.labels[i] for i in include]
                    self.segments = [self.segments[i] for i in include]
                    self.shapes = self.shapes[include]  # wh
        
                # Create indices
                n = len(self.shapes)  # number of images
                bi = np.floor(np.arange(n) / batch_size).astype(int)  # batch index
                nb = bi[-1] + 1  # number of batches
                self.batch = bi  # batch index of image
                self.n = n
                self.indices = np.arange(n)
                if rank > -1:  # DDP indices (see: SmartDistributedSampler)
                    # force each rank (i.e. GPU process) to sample the same subset of data on every epoch
                    self.indices = self.indices[np.random.RandomState(seed=seed).permutation(n) % WORLD_SIZE == RANK]
        
                # Update labels
                include_class = []  # filter labels to include only these classes (optional)
                self.segments = list(self.segments)
                include_class_array = np.array(include_class).reshape(1, -1)
                for i, (label, segment) in enumerate(zip(self.labels, self.segments)):
                    if include_class:
                        j = (label[:, 0:1] == include_class_array).any(1)
                        self.labels[i] = label[j]
                        if segment:
                            self.segments[i] = [segment[idx] for idx, elem in enumerate(j) if elem]
                    if single_cls:  # single-class training, merge all classes into 0
                        self.labels[i][:, 0] = 0
        
                # Rectangular Training
                if self.rect:
                    # Sort by aspect ratio
                    s = self.shapes  # wh
                    ar = s[:, 1] / s[:, 0]  # aspect ratio
                    irect = ar.argsort()
                    self.im_files = [self.im_files[i] for i in irect]
                    self.label_files = [self.label_files[i] for i in irect]
                    self.labels = [self.labels[i] for i in irect]
                    self.segments = [self.segments[i] for i in irect]
                    self.shapes = s[irect]  # wh
                    ar = ar[irect]
        
                    # Set training image shapes
                    shapes = [[1, 1]] * nb
                    for i in range(nb):
                        ari = ar[bi == i]
                        mini, maxi = ari.min(), ari.max()
                        if maxi < 1:
                            shapes[i] = [maxi, 1]
                        elif mini > 1:
                            shapes[i] = [1, 1 / mini]
        
                    self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(int) * stride
        
                # Cache images into RAM/disk for faster training
                if cache_images == "ram" and not self.check_cache_ram(prefix=prefix):
                    cache_images = False
                self.ims = [None] * n
                self.npy_files = [Path(f).with_suffix(".npy") for f in self.im_files]
                if cache_images:
                    b, gb = 0, 1 << 30  # bytes of cached images, bytes per gigabytes
                    self.im_hw0, self.im_hw = [None] * n, [None] * n
                    fcn = self.cache_images_to_disk if cache_images == "disk" else self.load_image
                    results = ThreadPool(NUM_THREADS).imap(lambda i: (i, fcn(i)), self.indices)
                    pbar = tqdm(results, total=len(self.indices), bar_format=TQDM_BAR_FORMAT, disable=LOCAL_RANK > 0)
                    for i, x in pbar:
                        if cache_images == "disk":
                            b += self.npy_files[i].stat().st_size
                        else:  # 'ram'
                            self.ims[i], self.im_hw0[i], self.im_hw[i] = x  # im, hw_orig, hw_resized = load_image(self, i)
                            b += self.ims[i].nbytes * WORLD_SIZE
                        pbar.desc = f"{prefix}Caching images ({b / gb:.1f}GB {cache_images})"
                    pbar.close()
        
            def check_cache_ram(self, safety_margin=0.1, prefix=""):
                """Checks if available RAM is sufficient for caching images, adjusting for a safety margin."""
                b, gb = 0, 1 << 30  # bytes of cached images, bytes per gigabytes
                n = min(self.n, 30)  # extrapolate from 30 random images
                for _ in range(n):
                    im = cv2.imread(random.choice(self.im_files))  # sample image
                    ratio = self.img_size / max(im.shape[0], im.shape[1])  # max(h, w)  # ratio
                    b += im.nbytes * ratio**2
                mem_required = b * self.n / n  # GB required to cache dataset into RAM
                mem = psutil.virtual_memory()
                cache = mem_required * (1 + safety_margin) < mem.available  # to cache or not to cache, that is the question
                if not cache:
                    LOGGER.info(
                        f'{prefix}{mem_required / gb:.1f}GB RAM required, '
                        f'{mem.available / gb:.1f}/{mem.total / gb:.1f}GB available, '
                        f"{'caching images ✅' if cache else 'not caching images ⚠️'}"
                    )
                return cache
        
            def cache_labels(self, path=Path("./labels.cache"), prefix=""):
                """Caches dataset labels, verifies images, reads shapes, and tracks dataset integrity."""
                x = {}  # dict
                nm, nf, ne, nc, msgs = 0, 0, 0, 0, []  # number missing, found, empty, corrupt, messages
                desc = f"{prefix}Scanning {path.parent / path.stem}..."
                with Pool(NUM_THREADS) as pool:
                    pbar = tqdm(
                        pool.imap(verify_image_label, zip(self.im_files, self.label_files, repeat(prefix))),
                        desc=desc,
                        total=len(self.im_files),
                        bar_format=TQDM_BAR_FORMAT,
                    )
                    for im_file, lb, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar:
                        nm += nm_f
                        nf += nf_f
                        ne += ne_f
                        nc += nc_f
                        if im_file:
                            x[im_file] = [lb, shape, segments]
                        if msg:
                            msgs.append(msg)
                        pbar.desc = f"{desc} {nf} images, {nm + ne} backgrounds, {nc} corrupt"
        
                pbar.close()
                if msgs:
                    LOGGER.info("\n".join(msgs))
                if nf == 0:
                    LOGGER.warning(f"{prefix}WARNING ⚠️ No labels found in {path}. {HELP_URL}")
                x["hash"] = get_hash(self.label_files + self.im_files)
                x["results"] = nf, nm, ne, nc, len(self.im_files)
                x["msgs"] = msgs  # warnings
                x["version"] = self.cache_version  # cache version
                try:
                    np.save(path, x)  # save cache for next time
                    path.with_suffix(".cache.npy").rename(path)  # remove .npy suffix
                    LOGGER.info(f"{prefix}New cache created: {path}")
                except Exception as e:
                    LOGGER.warning(f"{prefix}WARNING ⚠️ Cache directory {path.parent} is not writeable: {e}")  # not writeable
                return x
        
            def __len__(self):
                """Returns the number of images in the dataset."""
                return len(self.im_files)
        
            # def __iter__(self):
            #     self.count = -1
            #     print('ran dataset iter')
            #     #self.shuffled_vector = np.random.permutation(self.nF) if self.augment else np.arange(self.nF)
            #     return self
        
            def __getitem__(self, index):
                """Fetches the dataset item at the given index, considering linear, shuffled, or weighted sampling."""
                index = self.indices[index]  # linear, shuffled, or image_weights
        
                hyp = self.hyp
                mosaic = self.mosaic and random.random() < hyp["mosaic"]
                if mosaic:
                    # Load mosaic
                    img, labels = self.load_mosaic(index)
                    shapes = None
        
                    # MixUp augmentation
                    if random.random() < hyp["mixup"]:
                        img, labels = mixup(img, labels, *self.load_mosaic(random.choice(self.indices)))
        
                else:
                    # Load image
                    img, (h0, w0), (h, w) = self.load_image(index)
        
                    # Letterbox
                    shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size  # final letterboxed shape
                    img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
                    shapes = (h0, w0), ((h / h0, w / w0), pad)  # for COCO mAP rescaling
        
                    labels = self.labels[index].copy()
                    if labels.size:  # normalized xywh to pixel xyxy format
                        labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
        
                    if self.augment:
                        img, labels = random_perspective(
                            img,
                            labels,
                            degrees=hyp["degrees"],
                            translate=hyp["translate"],
                            scale=hyp["scale"],
                            shear=hyp["shear"],
                            perspective=hyp["perspective"],
                        )
        
                nl = len(labels)  # number of labels
                if nl:
                    labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1e-3)
        
                if self.augment:
                    # Albumentations
                    img, labels = self.albumentations(img, labels)
                    nl = len(labels)  # update after albumentations
        
                    # HSV color-space
                    augment_hsv(img, hgain=hyp["hsv_h"], sgain=hyp["hsv_s"], vgain=hyp["hsv_v"])
        
                    # Flip up-down
                    if random.random() < hyp["flipud"]:
                        img = np.flipud(img)
                        if nl:
                            labels[:, 2] = 1 - labels[:, 2]
        
                    # Flip left-right
                    if random.random() < hyp["fliplr"]:
                        img = np.fliplr(img)
                        if nl:
                            labels[:, 1] = 1 - labels[:, 1]
        
                    # Cutouts
                    # labels = cutout(img, labels, p=0.5)
                    # nl = len(labels)  # update after cutout
        
                labels_out = torch.zeros((nl, 6))
                if nl:
                    labels_out[:, 1:] = torch.from_numpy(labels)
        
                # Convert
                img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
                img = np.ascontiguousarray(img)
        
                return torch.from_numpy(img), labels_out, self.im_files[index], shapes
        
            def load_image(self, i):
                """
                Loads an image by index, returning the image, its original dimensions, and resized dimensions.
        
                Returns (im, original hw, resized hw)
                """
                im, f, fn = (
                    self.ims[i],
                    self.im_files[i],
                    self.npy_files[i],
                )
                if im is None:  # not cached in RAM
                    if fn.exists():  # load npy
                        im = np.load(fn)
                    else:  # read image
                        im = cv2.imread(f)  # BGR
                        assert im is not None, f"Image Not Found {f}"
                    h0, w0 = im.shape[:2]  # orig hw
                    r = self.img_size / max(h0, w0)  # ratio
                    if r != 1:  # if sizes are not equal
                        interp = cv2.INTER_LINEAR if (self.augment or r > 1) else cv2.INTER_AREA
                        im = cv2.resize(im, (math.ceil(w0 * r), math.ceil(h0 * r)), interpolation=interp)
                    return im, (h0, w0), im.shape[:2]  # im, hw_original, hw_resized
                return self.ims[i], self.im_hw0[i], self.im_hw[i]  # im, hw_original, hw_resized
        
            def cache_images_to_disk(self, i):
                """Saves an image to disk as an *.npy file for quicker loading, identified by index `i`."""
                f = self.npy_files[i]
                if not f.exists():
                    np.save(f.as_posix(), cv2.imread(self.im_files[i]))
        
            def load_mosaic(self, index):
                """Loads a 4-image mosaic for YOLOv5, combining 1 selected and 3 random images, with labels and segments."""
                labels4, segments4 = [], []
                s = self.img_size
                yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border)  # mosaic center x, y
                indices = [index] + random.choices(self.indices, k=3)  # 3 additional image indices
                random.shuffle(indices)
                for i, index in enumerate(indices):
                    # Load image
                    img, _, (h, w) = self.load_image(index)
        
                    # place img in img4
                    if i == 0:  # top left
                        img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
                        x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)
                        x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)
                    elif i == 1:  # top right
                        x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
                        x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
                    elif i == 2:  # bottom left
                        x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
                        x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
                    elif i == 3:  # bottom right
                        x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
                        x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)
        
                    img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]
                    padw = x1a - x1b
                    padh = y1a - y1b
        
                    # Labels
                    labels, segments = self.labels[index].copy(), self.segments[index].copy()
                    if labels.size:
                        labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh)  # normalized xywh to pixel xyxy format
                        segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
                    labels4.append(labels)
                    segments4.extend(segments)
        
                # Concat/clip labels
                labels4 = np.concatenate(labels4, 0)
                for x in (labels4[:, 1:], *segments4):
                    np.clip(x, 0, 2 * s, out=x)  # clip when using random_perspective()
                # img4, labels4 = replicate(img4, labels4)  # replicate
        
                # Augment
                img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp["copy_paste"])
                img4, labels4 = random_perspective(
                    img4,
                    labels4,
                    segments4,
                    degrees=self.hyp["degrees"],
                    translate=self.hyp["translate"],
                    scale=self.hyp["scale"],
                    shear=self.hyp["shear"],
                    perspective=self.hyp["perspective"],
                    border=self.mosaic_border,
                )  # border to remove
        
                return img4, labels4
        
            def load_mosaic9(self, index):
                """Loads 1 image + 8 random images into a 9-image mosaic for augmented YOLOv5 training, returning labels and
                segments.
                """
                labels9, segments9 = [], []
                s = self.img_size
                indices = [index] + random.choices(self.indices, k=8)  # 8 additional image indices
                random.shuffle(indices)
                hp, wp = -1, -1  # height, width previous
                for i, index in enumerate(indices):
                    # Load image
                    img, _, (h, w) = self.load_image(index)
        
                    # place img in img9
                    if i == 0:  # center
                        img9 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
                        h0, w0 = h, w
                        c = s, s, s + w, s + h  # xmin, ymin, xmax, ymax (base) coordinates
                    elif i == 1:  # top
                        c = s, s - h, s + w, s
                    elif i == 2:  # top right
                        c = s + wp, s - h, s + wp + w, s
                    elif i == 3:  # right
                        c = s + w0, s, s + w0 + w, s + h
                    elif i == 4:  # bottom right
                        c = s + w0, s + hp, s + w0 + w, s + hp + h
                    elif i == 5:  # bottom
                        c = s + w0 - w, s + h0, s + w0, s + h0 + h
                    elif i == 6:  # bottom left
                        c = s + w0 - wp - w, s + h0, s + w0 - wp, s + h0 + h
                    elif i == 7:  # left
                        c = s - w, s + h0 - h, s, s + h0
                    elif i == 8:  # top left
                        c = s - w, s + h0 - hp - h, s, s + h0 - hp
        
                    padx, pady = c[:2]
                    x1, y1, x2, y2 = (max(x, 0) for x in c)  # allocate coords
        
                    # Labels
                    labels, segments = self.labels[index].copy(), self.segments[index].copy()
                    if labels.size:
                        labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padx, pady)  # normalized xywh to pixel xyxy format
                        segments = [xyn2xy(x, w, h, padx, pady) for x in segments]
                    labels9.append(labels)
                    segments9.extend(segments)
        
                    # Image
                    img9[y1:y2, x1:x2] = img[y1 - pady :, x1 - padx :]  # img9[ymin:ymax, xmin:xmax]
                    hp, wp = h, w  # height, width previous
        
                # Offset
                yc, xc = (int(random.uniform(0, s)) for _ in self.mosaic_border)  # mosaic center x, y
                img9 = img9[yc : yc + 2 * s, xc : xc + 2 * s]
        
                # Concat/clip labels
                labels9 = np.concatenate(labels9, 0)
                labels9[:, [1, 3]] -= xc
                labels9[:, [2, 4]] -= yc
                c = np.array([xc, yc])  # centers
                segments9 = [x - c for x in segments9]
        
                for x in (labels9[:, 1:], *segments9):
                    np.clip(x, 0, 2 * s, out=x)  # clip when using random_perspective()
                # img9, labels9 = replicate(img9, labels9)  # replicate
        
                # Augment
                img9, labels9, segments9 = copy_paste(img9, labels9, segments9, p=self.hyp["copy_paste"])
                img9, labels9 = random_perspective(
                    img9,
                    labels9,
                    segments9,
                    degrees=self.hyp["degrees"],
                    translate=self.hyp["translate"],
                    scale=self.hyp["scale"],
                    shear=self.hyp["shear"],
                    perspective=self.hyp["perspective"],
                    border=self.mosaic_border,
                )  # border to remove
        
                return img9, labels9
        
            @staticmethod
            def collate_fn(batch):
                """Batches images, labels, paths, and shapes, assigning unique indices to targets in merged label tensor."""
                im, label, path, shapes = zip(*batch)  # transposed
                for i, lb in enumerate(label):
                    lb[:, 0] = i  # add target image index for build_targets()
                return torch.stack(im, 0), torch.cat(label, 0), path, shapes
        
            @staticmethod
            def collate_fn4(batch):
                """Bundles a batch's data by quartering the number of shapes and paths, preparing it for model input."""
                im, label, path, shapes = zip(*batch)  # transposed
                n = len(shapes) // 4
                im4, label4, path4, shapes4 = [], [], path[:n], shapes[:n]
        
                ho = torch.tensor([[0.0, 0, 0, 1, 0, 0]])
                wo = torch.tensor([[0.0, 0, 1, 0, 0, 0]])
                s = torch.tensor([[1, 1, 0.5, 0.5, 0.5, 0.5]])  # scale
                for i in range(n):  # zidane torch.zeros(16,3,720,1280)  # BCHW
                    i *= 4
                    if random.random() < 0.5:
                        im1 = F.interpolate(im[i].unsqueeze(0).float(), scale_factor=2.0, mode="bilinear", align_corners=False)[
                            0
                        ].type(im[i].type())
                        lb = label[i]
                    else:
                        im1 = torch.cat((torch.cat((im[i], im[i + 1]), 1), torch.cat((im[i + 2], im[i + 3]), 1)), 2)
                        lb = torch.cat((label[i], label[i + 1] + ho, label[i + 2] + wo, label[i + 3] + ho + wo), 0) * s
                    im4.append(im1)
                    label4.append(lb)
        
                for i, lb in enumerate(label4):
                    lb[:, 0] = i  # add target image index for build_targets()
        
                return torch.stack(im4, 0), torch.cat(label4, 0), path4, shapes4
        
        
        # Ancillary functions --------------------------------------------------------------------------------------------------
        def flatten_recursive(path=DATASETS_DIR / "coco128"):
            """Flattens a directory by copying all files from subdirectories to a new top-level directory, preserving
            filenames.
            """
            new_path = Path(f"{str(path)}_flat")
            if os.path.exists(new_path):
                shutil.rmtree(new_path)  # delete output folder
            os.makedirs(new_path)  # make new output folder
            for file in tqdm(glob.glob(f"{str(Path(path))}/**/*.*", recursive=True)):
                shutil.copyfile(file, new_path / Path(file).name)
        
        
        def extract_boxes(path=DATASETS_DIR / "coco128"):
            """
            Converts a detection dataset to a classification dataset, creating a directory for each class and extracting
            bounding boxes.
        
            Example: from utils.dataloaders import *; extract_boxes()
            """
            path = Path(path)  # images dir
            shutil.rmtree(path / "classification") if (path / "classification").is_dir() else None  # remove existing
            files = list(path.rglob("*.*"))
            n = len(files)  # number of files
            for im_file in tqdm(files, total=n):
                if im_file.suffix[1:] in IMG_FORMATS:
                    # image
                    im = cv2.imread(str(im_file))[..., ::-1]  # BGR to RGB
                    h, w = im.shape[:2]
        
                    # labels
                    lb_file = Path(img2label_paths([str(im_file)])[0])
                    if Path(lb_file).exists():
                        with open(lb_file) as f:
                            lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32)  # labels
        
                        for j, x in enumerate(lb):
                            c = int(x[0])  # class
                            f = (path / "classification") / f"{c}" / f"{path.stem}_{im_file.stem}_{j}.jpg"  # new filename
                            if not f.parent.is_dir():
                                f.parent.mkdir(parents=True)
        
                            b = x[1:] * [w, h, w, h]  # box
                            # b[2:] = b[2:].max()  # rectangle to square
                            b[2:] = b[2:] * 1.2 + 3  # pad
                            b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(int)
        
                            b[[0, 2]] = np.clip(b[[0, 2]], 0, w)  # clip boxes outside of image
                            b[[1, 3]] = np.clip(b[[1, 3]], 0, h)
                            assert cv2.imwrite(str(f), im[b[1] : b[3], b[0] : b[2]]), f"box failure in {f}"
        
        
        def autosplit(path=DATASETS_DIR / "coco128/images", weights=(0.9, 0.1, 0.0), annotated_only=False):
            """Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files
            Usage: from utils.dataloaders import *; autosplit()
            Arguments
                path:            Path to images directory
                weights:         Train, val, test weights (list, tuple)
                annotated_only:  Only use images with an annotated txt file
            """
            path = Path(path)  # images dir
            files = sorted(x for x in path.rglob("*.*") if x.suffix[1:].lower() in IMG_FORMATS)  # image files only
            n = len(files)  # number of files
            random.seed(0)  # for reproducibility
            indices = random.choices([0, 1, 2], weights=weights, k=n)  # assign each image to a split
        
            txt = ["autosplit_train.txt", "autosplit_val.txt", "autosplit_test.txt"]  # 3 txt files
            for x in txt:
                if (path.parent / x).exists():
                    (path.parent / x).unlink()  # remove existing
        
            print(f"Autosplitting images from {path}" + ", using *.txt labeled images only" * annotated_only)
            for i, img in tqdm(zip(indices, files), total=n):
                if not annotated_only or Path(img2label_paths([str(img)])[0]).exists():  # check label
                    with open(path.parent / txt[i], "a") as f:
                        f.write(f"./{img.relative_to(path.parent).as_posix()}" + "\n")  # add image to txt file
        
        
        def verify_image_label(args):
            """Verifies a single image-label pair, ensuring image format, size, and legal label values."""
            im_file, lb_file, prefix = args
            nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, "", []  # number (missing, found, empty, corrupt), message, segments
            try:
                # verify images
                im = Image.open(im_file)
                im.verify()  # PIL verify
                shape = exif_size(im)  # image size
                assert (shape[0] > 9) & (shape[1] > 9), f"image size {shape} <10 pixels"
                assert im.format.lower() in IMG_FORMATS, f"invalid image format {im.format}"
                if im.format.lower() in ("jpg", "jpeg"):
                    with open(im_file, "rb") as f:
                        f.seek(-2, 2)
                        if f.read() != b"\xff\xd9":  # corrupt JPEG
                            ImageOps.exif_transpose(Image.open(im_file)).save(im_file, "JPEG", subsampling=0, quality=100)
                            msg = f"{prefix}WARNING ⚠️ {im_file}: corrupt JPEG restored and saved"
        
                # verify labels
                if os.path.isfile(lb_file):
                    nf = 1  # label found
                    with open(lb_file) as f:
                        lb = [x.split() for x in f.read().strip().splitlines() if len(x)]
                        if any(len(x) > 6 for x in lb):  # is segment
                            classes = np.array([x[0] for x in lb], dtype=np.float32)
                            segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in lb]  # (cls, xy1...)
                            lb = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1)  # (cls, xywh)
                        lb = np.array(lb, dtype=np.float32)
                    nl = len(lb)
                    if nl:
                        assert lb.shape[1] == 5, f"labels require 5 columns, {lb.shape[1]} columns detected"
                        assert (lb >= 0).all(), f"negative label values {lb[lb < 0]}"
                        assert (lb[:, 1:] <= 1).all(), f"non-normalized or out of bounds coordinates {lb[:, 1:][lb[:, 1:] > 1]}"
                        _, i = np.unique(lb, axis=0, return_index=True)
                        if len(i) < nl:  # duplicate row check
                            lb = lb[i]  # remove duplicates
                            if segments:
                                segments = [segments[x] for x in i]
                            msg = f"{prefix}WARNING ⚠️ {im_file}: {nl - len(i)} duplicate labels removed"
                    else:
                        ne = 1  # label empty
                        lb = np.zeros((0, 5), dtype=np.float32)
                else:
                    nm = 1  # label missing
                    lb = np.zeros((0, 5), dtype=np.float32)
                return im_file, lb, shape, segments, nm, nf, ne, nc, msg
            except Exception as e:
                nc = 1
                msg = f"{prefix}WARNING ⚠️ {im_file}: ignoring corrupt image/label: {e}"
                return [None, None, None, None, nm, nf, ne, nc, msg]
        
        
        class HUBDatasetStats:
            """
            Class for generating HUB dataset JSON and `-hub` dataset directory.
        
            Arguments
                path:           Path to data.yaml or data.zip (with data.yaml inside data.zip)
                autodownload:   Attempt to download dataset if not found locally
        
            Usage
                from utils.dataloaders import HUBDatasetStats
                stats = HUBDatasetStats('coco128.yaml', autodownload=True)  # usage 1
                stats = HUBDatasetStats('path/to/coco128.zip')  # usage 2
                stats.get_json(save=False)
                stats.process_images()
            """
        
            def __init__(self, path="coco128.yaml", autodownload=False):
                """Initializes HUBDatasetStats with optional auto-download for datasets, given a path to dataset YAML or ZIP
                file.
                """
                zipped, data_dir, yaml_path = self._unzip(Path(path))
                try:
                    with open(check_yaml(yaml_path), errors="ignore") as f:
                        data = yaml.safe_load(f)  # data dict
                        if zipped:
                            data["path"] = data_dir
                except Exception as e:
                    raise Exception("error/HUB/dataset_stats/yaml_load") from e
        
                check_dataset(data, autodownload)  # download dataset if missing
                self.hub_dir = Path(data["path"] + "-hub")
                self.im_dir = self.hub_dir / "images"
                self.im_dir.mkdir(parents=True, exist_ok=True)  # makes /images
                self.stats = {"nc": data["nc"], "names": list(data["names"].values())}  # statistics dictionary
                self.data = data
        
            @staticmethod
            def _find_yaml(dir):
                """Finds and returns the path to a single '.yaml' file in the specified directory, preferring files that match
                the directory name.
                """
                files = list(dir.glob("*.yaml")) or list(dir.rglob("*.yaml"))  # try root level first and then recursive
                assert files, f"No *.yaml file found in {dir}"
                if len(files) > 1:
                    files = [f for f in files if f.stem == dir.stem]  # prefer *.yaml files that match dir name
                    assert files, f"Multiple *.yaml files found in {dir}, only 1 *.yaml file allowed"
                assert len(files) == 1, f"Multiple *.yaml files found: {files}, only 1 *.yaml file allowed in {dir}"
                return files[0]
        
            def _unzip(self, path):
                """Unzips a .zip file at 'path', returning success status, unzipped directory, and path to YAML file within."""
                if not str(path).endswith(".zip"):  # path is data.yaml
                    return False, None, path
                assert Path(path).is_file(), f"Error unzipping {path}, file not found"
                unzip_file(path, path=path.parent)
                dir = path.with_suffix("")  # dataset directory == zip name
                assert dir.is_dir(), f"Error unzipping {path}, {dir} not found. path/to/abc.zip MUST unzip to path/to/abc/"
                return True, str(dir), self._find_yaml(dir)  # zipped, data_dir, yaml_path
        
            def _hub_ops(self, f, max_dim=1920):
                """Resizes and saves an image at reduced quality for web/app viewing, supporting both PIL and OpenCV."""
                f_new = self.im_dir / Path(f).name  # dataset-hub image filename
                try:  # use PIL
                    im = Image.open(f)
                    r = max_dim / max(im.height, im.width)  # ratio
                    if r < 1.0:  # image too large
                        im = im.resize((int(im.width * r), int(im.height * r)))
                    im.save(f_new, "JPEG", quality=50, optimize=True)  # save
                except Exception as e:  # use OpenCV
                    LOGGER.info(f"WARNING ⚠️ HUB ops PIL failure {f}: {e}")
                    im = cv2.imread(f)
                    im_height, im_width = im.shape[:2]
                    r = max_dim / max(im_height, im_width)  # ratio
                    if r < 1.0:  # image too large
                        im = cv2.resize(im, (int(im_width * r), int(im_height * r)), interpolation=cv2.INTER_AREA)
                    cv2.imwrite(str(f_new), im)
        
            def get_json(self, save=False, verbose=False):
                """Generates dataset JSON for Ultralytics HUB, optionally saves or prints it; save=bool, verbose=bool."""
        
                def _round(labels):
                    """Rounds class labels to integers and coordinates to 4 decimal places for improved label accuracy."""
                    return [[int(c), *(round(x, 4) for x in points)] for c, *points in labels]
        
                for split in "train", "val", "test":
                    if self.data.get(split) is None:
                        self.stats[split] = None  # i.e. no test set
                        continue
                    dataset = LoadImagesAndLabels(self.data[split])  # load dataset
                    x = np.array(
                        [
                            np.bincount(label[:, 0].astype(int), minlength=self.data["nc"])
                            for label in tqdm(dataset.labels, total=dataset.n, desc="Statistics")
                        ]
                    )  # shape(128x80)
                    self.stats[split] = {
                        "instance_stats": {"total": int(x.sum()), "per_class": x.sum(0).tolist()},
                        "image_stats": {
                            "total": dataset.n,
                            "unlabelled": int(np.all(x == 0, 1).sum()),
                            "per_class": (x > 0).sum(0).tolist(),
                        },
                        "labels": [{str(Path(k).name): _round(v.tolist())} for k, v in zip(dataset.im_files, dataset.labels)],
                    }
        
                # Save, print and return
                if save:
                    stats_path = self.hub_dir / "stats.json"
                    print(f"Saving {stats_path.resolve()}...")
                    with open(stats_path, "w") as f:
                        json.dump(self.stats, f)  # save stats.json
                if verbose:
                    print(json.dumps(self.stats, indent=2, sort_keys=False))
                return self.stats
        
            def process_images(self):
                """Compresses images for Ultralytics HUB across 'train', 'val', 'test' splits and saves to specified
                directory.
                """
                for split in "train", "val", "test":
                    if self.data.get(split) is None:
                        continue
                    dataset = LoadImagesAndLabels(self.data[split])  # load dataset
                    desc = f"{split} images"
                    for _ in tqdm(ThreadPool(NUM_THREADS).imap(self._hub_ops, dataset.im_files), total=dataset.n, desc=desc):
                        pass
                print(f"Done. All images saved to {self.im_dir}")
                return self.im_dir
        
        
        # Classification dataloaders -------------------------------------------------------------------------------------------
        class ClassificationDataset(torchvision.datasets.ImageFolder):
            """
            YOLOv5 Classification Dataset.
        
            Arguments
                root:  Dataset path
                transform:  torchvision transforms, used by default
                album_transform: Albumentations transforms, used if installed
            """
        
            def __init__(self, root, augment, imgsz, cache=False):
                """Initializes YOLOv5 Classification Dataset with optional caching, augmentations, and transforms for image
                classification.
                """
                super().__init__(root=root)
                self.torch_transforms = classify_transforms(imgsz)
                self.album_transforms = classify_albumentations(augment, imgsz) if augment else None
                self.cache_ram = cache is True or cache == "ram"
                self.cache_disk = cache == "disk"
                self.samples = [list(x) + [Path(x[0]).with_suffix(".npy"), None] for x in self.samples]  # file, index, npy, im
        
            def __getitem__(self, i):
                """Fetches and transforms an image sample by index, supporting RAM/disk caching and Augmentations."""
                f, j, fn, im = self.samples[i]  # filename, index, filename.with_suffix('.npy'), image
                if self.cache_ram and im is None:
                    im = self.samples[i][3] = cv2.imread(f)
                elif self.cache_disk:
                    if not fn.exists():  # load npy
                        np.save(fn.as_posix(), cv2.imread(f))
                    im = np.load(fn)
                else:  # read image
                    im = cv2.imread(f)  # BGR
                if self.album_transforms:
                    sample = self.album_transforms(image=cv2.cvtColor(im, cv2.COLOR_BGR2RGB))["image"]
                else:
                    sample = self.torch_transforms(im)
                return sample, j
        
        
        def create_classification_dataloader(
            path, imgsz=224, batch_size=16, augment=True, cache=False, rank=-1, workers=8, shuffle=True
        ):
            # Returns Dataloader object to be used with YOLOv5 Classifier
            with torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDP
                dataset = ClassificationDataset(root=path, imgsz=imgsz, augment=augment, cache=cache)
            batch_size = min(batch_size, len(dataset))
            nd = torch.cuda.device_count()
            nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers])
            sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)
            generator = torch.Generator()
            generator.manual_seed(6148914691236517205 + RANK)
            return InfiniteDataLoader(
                dataset,
                batch_size=batch_size,
                shuffle=shuffle and sampler is None,
                num_workers=nw,
                sampler=sampler,
                pin_memory=PIN_MEMORY,
                worker_init_fn=seed_worker,
                generator=generator,
            )  # or DataLoader(persistent_workers=True)
        
      • 📄 downloads.py 5.1KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Download utils."""
        
        import logging
        import subprocess
        import urllib
        from pathlib import Path
        
        import requests
        import torch
        
        
        def is_url(url, check=True):
            """Determines if a string is a URL and optionally checks its existence online, returning a boolean."""
            try:
                url = str(url)
                result = urllib.parse.urlparse(url)
                assert all([result.scheme, result.netloc])  # check if is url
                return (urllib.request.urlopen(url).getcode() == 200) if check else True  # check if exists online
            except (AssertionError, urllib.request.HTTPError):
                return False
        
        
        def gsutil_getsize(url=""):
            """
            Returns the size in bytes of a file at a Google Cloud Storage URL using `gsutil du`.
        
            Returns 0 if the command fails or output is empty.
            """
            output = subprocess.check_output(["gsutil", "du", url], shell=True, encoding="utf-8")
            return int(output.split()[0]) if output else 0
        
        
        def url_getsize(url="https://ultralytics.com/images/bus.jpg"):
            """Returns the size in bytes of a downloadable file at a given URL; defaults to -1 if not found."""
            response = requests.head(url, allow_redirects=True)
            return int(response.headers.get("content-length", -1))
        
        
        def curl_download(url, filename, *, silent: bool = False) -> bool:
            """Download a file from a url to a filename using curl."""
            silent_option = "sS" if silent else ""  # silent
            proc = subprocess.run(
                [
                    "curl",
                    "-#",
                    f"-{silent_option}L",
                    url,
                    "--output",
                    filename,
                    "--retry",
                    "9",
                    "-C",
                    "-",
                ]
            )
            return proc.returncode == 0
        
        
        def safe_download(file, url, url2=None, min_bytes=1e0, error_msg=""):
            """
            Downloads a file from a URL (or alternate URL) to a specified path if file is above a minimum size.
        
            Removes incomplete downloads.
            """
            from utils.general import LOGGER
        
            file = Path(file)
            assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}"
            try:  # url1
                LOGGER.info(f"Downloading {url} to {file}...")
                torch.hub.download_url_to_file(url, str(file), progress=LOGGER.level <= logging.INFO)
                assert file.exists() and file.stat().st_size > min_bytes, assert_msg  # check
            except Exception as e:  # url2
                if file.exists():
                    file.unlink()  # remove partial downloads
                LOGGER.info(f"ERROR: {e}\nRe-attempting {url2 or url} to {file}...")
                # curl download, retry and resume on fail
                curl_download(url2 or url, file)
            finally:
                if not file.exists() or file.stat().st_size < min_bytes:  # check
                    if file.exists():
                        file.unlink()  # remove partial downloads
                    LOGGER.info(f"ERROR: {assert_msg}\n{error_msg}")
                LOGGER.info("")
        
        
        def attempt_download(file, repo="ultralytics/yolov5", release="v7.0"):
            """Downloads a file from GitHub release assets or via direct URL if not found locally, supporting backup
            versions.
            """
            from utils.general import LOGGER
        
            def github_assets(repository, version="latest"):
                # Return GitHub repo tag (i.e. 'v7.0') and assets (i.e. ['yolov5s.pt', 'yolov5m.pt', ...])
                if version != "latest":
                    version = f"tags/{version}"  # i.e. tags/v7.0
                response = requests.get(f"https://api.github.com/repos/{repository}/releases/{version}").json()  # github api
                return response["tag_name"], [x["name"] for x in response["assets"]]  # tag, assets
        
            file = Path(str(file).strip().replace("'", ""))
            if not file.exists():
                # URL specified
                name = Path(urllib.parse.unquote(str(file))).name  # decode '%2F' to '/' etc.
                if str(file).startswith(("http:/", "https:/")):  # download
                    url = str(file).replace(":/", "://")  # Pathlib turns :// -> :/
                    file = name.split("?")[0]  # parse authentication https://url.com/file.txt?auth...
                    if Path(file).is_file():
                        LOGGER.info(f"Found {url} locally at {file}")  # file already exists
                    else:
                        safe_download(file=file, url=url, min_bytes=1e5)
                    return file
        
                # GitHub assets
                assets = [f"yolov5{size}{suffix}.pt" for size in "nsmlx" for suffix in ("", "6", "-cls", "-seg")]  # default
                try:
                    tag, assets = github_assets(repo, release)
                except Exception:
                    try:
                        tag, assets = github_assets(repo)  # latest release
                    except Exception:
                        try:
                            tag = subprocess.check_output("git tag", shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
                        except Exception:
                            tag = release
        
                if name in assets:
                    file.parent.mkdir(parents=True, exist_ok=True)  # make parent dir (if required)
                    safe_download(
                        file,
                        url=f"https://github.com/{repo}/releases/download/{tag}/{name}",
                        min_bytes=1e5,
                        error_msg=f"{file} missing, try downloading from https://github.com/{repo}/releases/{tag}",
                    )
        
            return str(file)
        
      • 📄 general.py 49.9KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """General utils."""
        
        import contextlib
        import glob
        import inspect
        import logging
        import logging.config
        import math
        import os
        import platform
        import random
        import re
        import signal
        import subprocess
        import sys
        import time
        import urllib
        from copy import deepcopy
        from datetime import datetime
        from itertools import repeat
        from multiprocessing.pool import ThreadPool
        from pathlib import Path
        from subprocess import check_output
        from tarfile import is_tarfile
        from typing import Optional
        from zipfile import ZipFile, is_zipfile
        
        import cv2
        import numpy as np
        import pandas as pd
        import pkg_resources as pkg
        import torch
        import torchvision
        import yaml
        
        # Import 'ultralytics' package or install if missing
        try:
            import ultralytics
        
            assert hasattr(ultralytics, "__version__")  # verify package is not directory
        except (ImportError, AssertionError):
            os.system("pip install -U ultralytics")
            import ultralytics
        
        from ultralytics.utils.checks import check_requirements
        
        from utils import TryExcept, emojis
        from utils.downloads import curl_download, gsutil_getsize
        from utils.metrics import box_iou, fitness
        
        FILE = Path(__file__).resolve()
        ROOT = FILE.parents[1]  # YOLOv5 root directory
        RANK = int(os.getenv("RANK", -1))
        
        # Settings
        NUM_THREADS = min(8, max(1, os.cpu_count() - 1))  # number of YOLOv5 multiprocessing threads
        DATASETS_DIR = Path(os.getenv("YOLOv5_DATASETS_DIR", ROOT.parent / "datasets"))  # global datasets directory
        AUTOINSTALL = str(os.getenv("YOLOv5_AUTOINSTALL", True)).lower() == "true"  # global auto-install mode
        VERBOSE = str(os.getenv("YOLOv5_VERBOSE", True)).lower() == "true"  # global verbose mode
        TQDM_BAR_FORMAT = "{l_bar}{bar:10}{r_bar}"  # tqdm bar format
        FONT = "Arial.ttf"  # https://ultralytics.com/assets/Arial.ttf
        
        torch.set_printoptions(linewidth=320, precision=5, profile="long")
        np.set_printoptions(linewidth=320, formatter={"float_kind": "{:11.5g}".format})  # format short g, %precision=5
        pd.options.display.max_columns = 10
        cv2.setNumThreads(0)  # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
        os.environ["NUMEXPR_MAX_THREADS"] = str(NUM_THREADS)  # NumExpr max threads
        os.environ["OMP_NUM_THREADS"] = "1" if platform.system() == "darwin" else str(NUM_THREADS)  # OpenMP (PyTorch and SciPy)
        os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # suppress verbose TF compiler warnings in Colab
        os.environ["TORCH_CPP_LOG_LEVEL"] = "ERROR"  # suppress "NNPACK.cpp could not initialize NNPACK" warnings
        os.environ["KINETO_LOG_LEVEL"] = "5"  # suppress verbose PyTorch profiler output when computing FLOPs
        
        
        def is_ascii(s=""):
            """Checks if input string `s` contains only ASCII characters; returns `True` if so, otherwise `False`."""
            s = str(s)  # convert list, tuple, None, etc. to str
            return len(s.encode().decode("ascii", "ignore")) == len(s)
        
        
        def is_chinese(s="人工智能"):
            """Determines if a string `s` contains any Chinese characters; returns `True` if so, otherwise `False`."""
            return bool(re.search("[\u4e00-\u9fff]", str(s)))
        
        
        def is_colab():
            """Checks if the current environment is a Google Colab instance; returns `True` for Colab, otherwise `False`."""
            return "google.colab" in sys.modules
        
        
        def is_jupyter():
            """
            Check if the current script is running inside a Jupyter Notebook. Verified on Colab, Jupyterlab, Kaggle, Paperspace.
        
            Returns:
                bool: True if running inside a Jupyter Notebook, False otherwise.
            """
            with contextlib.suppress(Exception):
                from IPython import get_ipython
        
                return get_ipython() is not None
            return False
        
        
        def is_kaggle():
            """Checks if the current environment is a Kaggle Notebook by validating environment variables."""
            return os.environ.get("PWD") == "/kaggle/working" and os.environ.get("KAGGLE_URL_BASE") == "https://www.kaggle.com"
        
        
        def is_docker() -> bool:
            """Check if the process runs inside a docker container."""
            # if Path("/.dockerenv").exists():
            if Path("/workspace").exists():
                return True
            try:  # check if docker is in control groups
                with open("/proc/self/cgroup") as file:
                    return any("docker" in line for line in file)
            except OSError:
                return False
        
        
        def is_writeable(dir, test=False):
            """Checks if a directory is writable, optionally testing by creating a temporary file if `test=True`."""
            if not test:
                return os.access(dir, os.W_OK)  # possible issues on Windows
            file = Path(dir) / "tmp.txt"
            try:
                with open(file, "w"):  # open file with write permissions
                    pass
                file.unlink()  # remove file
                return True
            except OSError:
                return False
        
        
        LOGGING_NAME = "yolov5"
        
        
        def set_logging(name=LOGGING_NAME, verbose=True):
            """Configures logging with specified verbosity; `name` sets the logger's name, `verbose` controls logging level."""
            rank = int(os.getenv("RANK", -1))  # rank in world for Multi-GPU trainings
            level = logging.INFO if verbose and rank in {-1, 0} else logging.ERROR
            logging.config.dictConfig(
                {
                    "version": 1,
                    "disable_existing_loggers": False,
                    "formatters": {name: {"format": "%(message)s"}},
                    "handlers": {
                        name: {
                            "class": "logging.StreamHandler",
                            "formatter": name,
                            "level": level,
                        }
                    },
                    "loggers": {
                        name: {
                            "level": level,
                            "handlers": [name],
                            "propagate": False,
                        }
                    },
                }
            )
        
        
        set_logging(LOGGING_NAME)  # run before defining LOGGER
        LOGGER = logging.getLogger(LOGGING_NAME)  # define globally (used in train.py, val.py, detect.py, etc.)
        if platform.system() == "Windows":
            for fn in LOGGER.info, LOGGER.warning:
                setattr(LOGGER, fn.__name__, lambda x: fn(emojis(x)))  # emoji safe logging
        
        
        def user_config_dir(dir="Ultralytics", env_var="YOLOV5_CONFIG_DIR"):
            """Returns user configuration directory path, preferring environment variable `YOLOV5_CONFIG_DIR` if set, else OS-
            specific.
            """
            env = os.getenv(env_var)
            if env:
                path = Path(env)  # use environment variable
            else:
                cfg = {"Windows": "AppData/Roaming", "Linux": ".config", "Darwin": "Library/Application Support"}  # 3 OS dirs
                path = Path.home() / cfg.get(platform.system(), "")  # OS-specific config dir
                path = (path if is_writeable(path) else Path("/tmp")) / dir  # GCP and AWS lambda fix, only /tmp is writeable
            path.mkdir(exist_ok=True)  # make if required
            return path
        
        
        CONFIG_DIR = user_config_dir()  # Ultralytics settings dir
        
        
        class Profile(contextlib.ContextDecorator):
            # YOLOv5 Profile class. Usage: @Profile() decorator or 'with Profile():' context manager
            def __init__(self, t=0.0, device: torch.device = None):
                """Initializes a profiling context for YOLOv5 with optional timing threshold and device specification."""
                self.t = t
                self.device = device
                self.cuda = bool(device and str(device).startswith("cuda"))
        
            def __enter__(self):
                """Initializes timing at the start of a profiling context block for performance measurement."""
                self.start = self.time()
                return self
        
            def __exit__(self, type, value, traceback):
                """Concludes timing, updating duration for profiling upon exiting a context block."""
                self.dt = self.time() - self.start  # delta-time
                self.t += self.dt  # accumulate dt
        
            def time(self):
                """Measures and returns the current time, synchronizing CUDA operations if `cuda` is True."""
                if self.cuda:
                    torch.cuda.synchronize(self.device)
                return time.time()
        
        
        class Timeout(contextlib.ContextDecorator):
            # YOLOv5 Timeout class. Usage: @Timeout(seconds) decorator or 'with Timeout(seconds):' context manager
            def __init__(self, seconds, *, timeout_msg="", suppress_timeout_errors=True):
                """Initializes a timeout context/decorator with defined seconds, optional message, and error suppression."""
                self.seconds = int(seconds)
                self.timeout_message = timeout_msg
                self.suppress = bool(suppress_timeout_errors)
        
            def _timeout_handler(self, signum, frame):
                """Raises a TimeoutError with a custom message when a timeout event occurs."""
                raise TimeoutError(self.timeout_message)
        
            def __enter__(self):
                """Initializes timeout mechanism on non-Windows platforms, starting a countdown to raise TimeoutError."""
                if platform.system() != "Windows":  # not supported on Windows
                    signal.signal(signal.SIGALRM, self._timeout_handler)  # Set handler for SIGALRM
                    signal.alarm(self.seconds)  # start countdown for SIGALRM to be raised
        
            def __exit__(self, exc_type, exc_val, exc_tb):
                """Disables active alarm on non-Windows systems and optionally suppresses TimeoutError if set."""
                if platform.system() != "Windows":
                    signal.alarm(0)  # Cancel SIGALRM if it's scheduled
                    if self.suppress and exc_type is TimeoutError:  # Suppress TimeoutError
                        return True
        
        
        class WorkingDirectory(contextlib.ContextDecorator):
            # Usage: @WorkingDirectory(dir) decorator or 'with WorkingDirectory(dir):' context manager
            def __init__(self, new_dir):
                """Initializes a context manager/decorator to temporarily change the working directory."""
                self.dir = new_dir  # new dir
                self.cwd = Path.cwd().resolve()  # current dir
        
            def __enter__(self):
                """Temporarily changes the working directory within a 'with' statement context."""
                os.chdir(self.dir)
        
            def __exit__(self, exc_type, exc_val, exc_tb):
                """Restores the original working directory upon exiting a 'with' statement context."""
                os.chdir(self.cwd)
        
        
        def methods(instance):
            """Returns list of method names for a class/instance excluding dunder methods."""
            return [f for f in dir(instance) if callable(getattr(instance, f)) and not f.startswith("__")]
        
        
        def print_args(args: Optional[dict] = None, show_file=True, show_func=False):
            """Logs the arguments of the calling function, with options to include the filename and function name."""
            x = inspect.currentframe().f_back  # previous frame
            file, _, func, _, _ = inspect.getframeinfo(x)
            if args is None:  # get args automatically
                args, _, _, frm = inspect.getargvalues(x)
                args = {k: v for k, v in frm.items() if k in args}
            try:
                file = Path(file).resolve().relative_to(ROOT).with_suffix("")
            except ValueError:
                file = Path(file).stem
            s = (f"{file}: " if show_file else "") + (f"{func}: " if show_func else "")
            LOGGER.info(colorstr(s) + ", ".join(f"{k}={v}" for k, v in args.items()))
        
        
        def init_seeds(seed=0, deterministic=False):
            """
            Initializes RNG seeds and sets deterministic options if specified.
        
            See https://pytorch.org/docs/stable/notes/randomness.html
            """
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)  # for Multi-GPU, exception safe
            # torch.backends.cudnn.benchmark = True  # AutoBatch problem https://github.com/ultralytics/yolov5/issues/9287
            if deterministic and check_version(torch.__version__, "1.12.0"):  # https://github.com/ultralytics/yolov5/pull/8213
                torch.use_deterministic_algorithms(True)
                torch.backends.cudnn.deterministic = True
                os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
                os.environ["PYTHONHASHSEED"] = str(seed)
        
        
        def intersect_dicts(da, db, exclude=()):
            """Returns intersection of `da` and `db` dicts with matching keys and shapes, excluding `exclude` keys; uses `da`
            values.
            """
            return {k: v for k, v in da.items() if k in db and all(x not in k for x in exclude) and v.shape == db[k].shape}
        
        
        def get_default_args(func):
            """Returns a dict of `func` default arguments by inspecting its signature."""
            signature = inspect.signature(func)
            return {k: v.default for k, v in signature.parameters.items() if v.default is not inspect.Parameter.empty}
        
        
        def get_latest_run(search_dir="."):
            """Returns the path to the most recent 'last.pt' file in /runs to resume from, searches in `search_dir`."""
            last_list = glob.glob(f"{search_dir}/**/last*.pt", recursive=True)
            return max(last_list, key=os.path.getctime) if last_list else ""
        
        
        def file_age(path=__file__):
            """Calculates and returns the age of a file in days based on its last modification time."""
            dt = datetime.now() - datetime.fromtimestamp(Path(path).stat().st_mtime)  # delta
            return dt.days  # + dt.seconds / 86400  # fractional days
        
        
        def file_date(path=__file__):
            """Returns a human-readable file modification date in 'YYYY-M-D' format, given a file path."""
            t = datetime.fromtimestamp(Path(path).stat().st_mtime)
            return f"{t.year}-{t.month}-{t.day}"
        
        
        def file_size(path):
            """Returns file or directory size in megabytes (MB) for a given path, where directories are recursively summed."""
            mb = 1 << 20  # bytes to MiB (1024 ** 2)
            path = Path(path)
            if path.is_file():
                return path.stat().st_size / mb
            elif path.is_dir():
                return sum(f.stat().st_size for f in path.glob("**/*") if f.is_file()) / mb
            else:
                return 0.0
        
        
        def check_online():
            """Checks internet connectivity by attempting to create a connection to "1.1.1.1" on port 443, retries once if the
            first attempt fails.
            """
            import socket
        
            def run_once():
                # Check once
                try:
                    socket.create_connection(("1.1.1.1", 443), 5)  # check host accessibility
                    return True
                except OSError:
                    return False
        
            return run_once() or run_once()  # check twice to increase robustness to intermittent connectivity issues
        
        
        def git_describe(path=ROOT):
            """
            Returns a human-readable git description of the repository at `path`, or an empty string on failure.
        
            Example output is 'fv5.0-5-g3e25f1e'. See https://git-scm.com/docs/git-describe.
            """
            try:
                assert (Path(path) / ".git").is_dir()
                return check_output(f"git -C {path} describe --tags --long --always", shell=True).decode()[:-1]
            except Exception:
                return ""
        
        
        @TryExcept()
        @WorkingDirectory(ROOT)
        def check_git_status(repo="ultralytics/yolov5", branch="master"):
            """Checks if YOLOv5 code is up-to-date with the repository, advising 'git pull' if behind; errors return informative
            messages.
            """
            url = f"https://github.com/{repo}"
            msg = f", for updates see {url}"
            s = colorstr("github: ")  # string
            assert Path(".git").exists(), s + "skipping check (not a git repository)" + msg
            assert check_online(), s + "skipping check (offline)" + msg
        
            splits = re.split(pattern=r"\s", string=check_output("git remote -v", shell=True).decode())
            matches = [repo in s for s in splits]
            if any(matches):
                remote = splits[matches.index(True) - 1]
            else:
                remote = "ultralytics"
                check_output(f"git remote add {remote} {url}", shell=True)
            check_output(f"git fetch {remote}", shell=True, timeout=5)  # git fetch
            local_branch = check_output("git rev-parse --abbrev-ref HEAD", shell=True).decode().strip()  # checked out
            n = int(check_output(f"git rev-list {local_branch}..{remote}/{branch} --count", shell=True))  # commits behind
            if n > 0:
                pull = "git pull" if remote == "origin" else f"git pull {remote} {branch}"
                s += f"⚠️ YOLOv5 is out of date by {n} commit{'s' * (n > 1)}. Use '{pull}' or 'git clone {url}' to update."
            else:
                s += f"up to date with {url} ✅"
            LOGGER.info(s)
        
        
        @WorkingDirectory(ROOT)
        def check_git_info(path="."):
            """Checks YOLOv5 git info, returning a dict with remote URL, branch name, and commit hash."""
            check_requirements("gitpython")
            import git
        
            try:
                repo = git.Repo(path)
                remote = repo.remotes.origin.url.replace(".git", "")  # i.e. 'https://github.com/ultralytics/yolov5'
                commit = repo.head.commit.hexsha  # i.e. '3134699c73af83aac2a481435550b968d5792c0d'
                try:
                    branch = repo.active_branch.name  # i.e. 'main'
                except TypeError:  # not on any branch
                    branch = None  # i.e. 'detached HEAD' state
                return {"remote": remote, "branch": branch, "commit": commit}
            except git.exc.InvalidGitRepositoryError:  # path is not a git dir
                return {"remote": None, "branch": None, "commit": None}
        
        
        def check_python(minimum="3.8.0"):
            """Checks if current Python version meets the minimum required version, exits if not."""
            check_version(platform.python_version(), minimum, name="Python ", hard=True)
        
        
        def check_version(current="0.0.0", minimum="0.0.0", name="version ", pinned=False, hard=False, verbose=False):
            """Checks if the current version meets the minimum required version, exits or warns based on parameters."""
            current, minimum = (pkg.parse_version(x) for x in (current, minimum))
            result = (current == minimum) if pinned else (current >= minimum)  # bool
            s = f"WARNING ⚠️ {name}{minimum} is required by YOLOv5, but {name}{current} is currently installed"  # string
            if hard:
                assert result, emojis(s)  # assert min requirements met
            if verbose and not result:
                LOGGER.warning(s)
            return result
        
        
        def check_img_size(imgsz, s=32, floor=0):
            """Adjusts image size to be divisible by stride `s`, supports int or list/tuple input, returns adjusted size."""
            if isinstance(imgsz, int):  # integer i.e. img_size=640
                new_size = max(make_divisible(imgsz, int(s)), floor)
            else:  # list i.e. img_size=[640, 480]
                imgsz = list(imgsz)  # convert to list if tuple
                new_size = [max(make_divisible(x, int(s)), floor) for x in imgsz]
            if new_size != imgsz:
                LOGGER.warning(f"WARNING ⚠️ --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}")
            return new_size
        
        
        def check_imshow(warn=False):
            """Checks environment support for image display; warns on failure if `warn=True`."""
            try:
                assert not is_jupyter()
                assert not is_docker()
                cv2.imshow("test", np.zeros((1, 1, 3)))
                cv2.waitKey(1)
                cv2.destroyAllWindows()
                cv2.waitKey(1)
                return True
            except Exception as e:
                if warn:
                    LOGGER.warning(f"WARNING ⚠️ Environment does not support cv2.imshow() or PIL Image.show()\n{e}")
                return False
        
        
        def check_suffix(file="yolov5s.pt", suffix=(".pt",), msg=""):
            """Validates if a file or files have an acceptable suffix, raising an error if not."""
            if file and suffix:
                if isinstance(suffix, str):
                    suffix = [suffix]
                for f in file if isinstance(file, (list, tuple)) else [file]:
                    s = Path(f).suffix.lower()  # file suffix
                    if len(s):
                        assert s in suffix, f"{msg}{f} acceptable suffix is {suffix}"
        
        
        def check_yaml(file, suffix=(".yaml", ".yml")):
            """Searches/downloads a YAML file, verifies its suffix (.yaml or .yml), and returns the file path."""
            return check_file(file, suffix)
        
        
        def check_file(file, suffix=""):
            """Searches/downloads a file, checks its suffix (if provided), and returns the file path."""
            check_suffix(file, suffix)  # optional
            file = str(file)  # convert to str()
            if os.path.isfile(file) or not file:  # exists
                return file
            elif file.startswith(("http:/", "https:/")):  # download
                url = file  # warning: Pathlib turns :// -> :/
                file = Path(urllib.parse.unquote(file).split("?")[0]).name  # '%2F' to '/', split https://url.com/file.txt?auth
                if os.path.isfile(file):
                    LOGGER.info(f"Found {url} locally at {file}")  # file already exists
                else:
                    LOGGER.info(f"Downloading {url} to {file}...")
                    torch.hub.download_url_to_file(url, file)
                    assert Path(file).exists() and Path(file).stat().st_size > 0, f"File download failed: {url}"  # check
                return file
            elif file.startswith("clearml://"):  # ClearML Dataset ID
                assert (
                    "clearml" in sys.modules
                ), "ClearML is not installed, so cannot use ClearML dataset. Try running 'pip install clearml'."
                return file
            else:  # search
                files = []
                for d in "data", "models", "utils":  # search directories
                    files.extend(glob.glob(str(ROOT / d / "**" / file), recursive=True))  # find file
                assert len(files), f"File not found: {file}"  # assert file was found
                assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}"  # assert unique
                return files[0]  # return file
        
        
        def check_font(font=FONT, progress=False):
            """Ensures specified font exists or downloads it from Ultralytics assets, optionally displaying progress."""
            font = Path(font)
            file = CONFIG_DIR / font.name
            if not font.exists() and not file.exists():
                url = f"https://ultralytics.com/assets/{font.name}"
                LOGGER.info(f"Downloading {url} to {file}...")
                torch.hub.download_url_to_file(url, str(file), progress=progress)
        
        
        def check_dataset(data, autodownload=True):
            """Validates and/or auto-downloads a dataset, returning its configuration as a dictionary."""
        
            # Download (optional)
            extract_dir = ""
            if isinstance(data, (str, Path)) and (is_zipfile(data) or is_tarfile(data)):
                download(data, dir=f"{DATASETS_DIR}/{Path(data).stem}", unzip=True, delete=False, curl=False, threads=1)
                data = next((DATASETS_DIR / Path(data).stem).rglob("*.yaml"))
                extract_dir, autodownload = data.parent, False
        
            # Read yaml (optional)
            if isinstance(data, (str, Path)):
                data = yaml_load(data)  # dictionary
        
            # Checks
            for k in "train", "val", "names":
                assert k in data, emojis(f"data.yaml '{k}:' field missing ❌")
            if isinstance(data["names"], (list, tuple)):  # old array format
                data["names"] = dict(enumerate(data["names"]))  # convert to dict
            assert all(isinstance(k, int) for k in data["names"].keys()), "data.yaml names keys must be integers, i.e. 2: car"
            data["nc"] = len(data["names"])
        
            # Resolve paths
            path = Path(extract_dir or data.get("path") or "")  # optional 'path' default to '.'
            if not path.is_absolute():
                path = (ROOT / path).resolve()
                data["path"] = path  # download scripts
            for k in "train", "val", "test":
                if data.get(k):  # prepend path
                    if isinstance(data[k], str):
                        x = (path / data[k]).resolve()
                        if not x.exists() and data[k].startswith("../"):
                            x = (path / data[k][3:]).resolve()
                        data[k] = str(x)
                    else:
                        data[k] = [str((path / x).resolve()) for x in data[k]]
        
            # Parse yaml
            train, val, test, s = (data.get(x) for x in ("train", "val", "test", "download"))
            if val:
                val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])]  # val path
                if not all(x.exists() for x in val):
                    LOGGER.info("\nDataset not found ⚠️, missing paths %s" % [str(x) for x in val if not x.exists()])
                    if not s or not autodownload:
                        raise Exception("Dataset not found ❌")
                    t = time.time()
                    if s.startswith("http") and s.endswith(".zip"):  # URL
                        f = Path(s).name  # filename
                        LOGGER.info(f"Downloading {s} to {f}...")
                        torch.hub.download_url_to_file(s, f)
                        Path(DATASETS_DIR).mkdir(parents=True, exist_ok=True)  # create root
                        unzip_file(f, path=DATASETS_DIR)  # unzip
                        Path(f).unlink()  # remove zip
                        r = None  # success
                    elif s.startswith("bash "):  # bash script
                        LOGGER.info(f"Running {s} ...")
                        r = subprocess.run(s, shell=True)
                    else:  # python script
                        r = exec(s, {"yaml": data})  # return None
                    dt = f"({round(time.time() - t, 1)}s)"
                    s = f"success ✅ {dt}, saved to {colorstr('bold', DATASETS_DIR)}" if r in (0, None) else f"failure {dt} ❌"
                    LOGGER.info(f"Dataset download {s}")
            check_font("Arial.ttf" if is_ascii(data["names"]) else "Arial.Unicode.ttf", progress=True)  # download fonts
            return data  # dictionary
        
        
        def check_amp(model):
            """Checks PyTorch AMP functionality for a model, returns True if AMP operates correctly, otherwise False."""
            from models.common import AutoShape, DetectMultiBackend
        
            def amp_allclose(model, im):
                # All close FP32 vs AMP results
                m = AutoShape(model, verbose=False)  # model
                a = m(im).xywhn[0]  # FP32 inference
                m.amp = True
                b = m(im).xywhn[0]  # AMP inference
                return a.shape == b.shape and torch.allclose(a, b, atol=0.1)  # close to 10% absolute tolerance
        
            prefix = colorstr("AMP: ")
            device = next(model.parameters()).device  # get model device
            if device.type in ("cpu", "mps"):
                return False  # AMP only used on CUDA devices
            f = ROOT / "data" / "images" / "bus.jpg"  # image to check
            im = f if f.exists() else "https://ultralytics.com/images/bus.jpg" if check_online() else np.ones((640, 640, 3))
            try:
                assert amp_allclose(deepcopy(model), im) or amp_allclose(DetectMultiBackend("yolov5n.pt", device), im)
                LOGGER.info(f"{prefix}checks passed ✅")
                return True
            except Exception:
                help_url = "https://github.com/ultralytics/yolov5/issues/7908"
                LOGGER.warning(f"{prefix}checks failed ❌, disabling Automatic Mixed Precision. See {help_url}")
                return False
        
        
        def yaml_load(file="data.yaml"):
            """Safely loads and returns the contents of a YAML file specified by `file` argument."""
            with open(file, errors="ignore") as f:
                return yaml.safe_load(f)
        
        
        def yaml_save(file="data.yaml", data={}):
            """Safely saves `data` to a YAML file specified by `file`, converting `Path` objects to strings; `data` is a
            dictionary.
            """
            with open(file, "w") as f:
                yaml.safe_dump({k: str(v) if isinstance(v, Path) else v for k, v in data.items()}, f, sort_keys=False)
        
        
        def unzip_file(file, path=None, exclude=(".DS_Store", "__MACOSX")):
            """Unzips `file` to `path` (default: file's parent), excluding filenames containing any in `exclude` (`.DS_Store`,
            `__MACOSX`).
            """
            if path is None:
                path = Path(file).parent  # default path
            with ZipFile(file) as zipObj:
                for f in zipObj.namelist():  # list all archived filenames in the zip
                    if all(x not in f for x in exclude):
                        zipObj.extract(f, path=path)
        
        
        def url2file(url):
            """
            Converts a URL string to a valid filename by stripping protocol, domain, and any query parameters.
        
            Example https://url.com/file.txt?auth -> file.txt
            """
            url = str(Path(url)).replace(":/", "://")  # Pathlib turns :// -> :/
            return Path(urllib.parse.unquote(url)).name.split("?")[0]  # '%2F' to '/', split https://url.com/file.txt?auth
        
        
        def download(url, dir=".", unzip=True, delete=True, curl=False, threads=1, retry=3):
            """Downloads and optionally unzips files concurrently, supporting retries and curl fallback."""
        
            def download_one(url, dir):
                # Download 1 file
                success = True
                if os.path.isfile(url):
                    f = Path(url)  # filename
                else:  # does not exist
                    f = dir / Path(url).name
                    LOGGER.info(f"Downloading {url} to {f}...")
                    for i in range(retry + 1):
                        if curl:
                            success = curl_download(url, f, silent=(threads > 1))
                        else:
                            torch.hub.download_url_to_file(url, f, progress=threads == 1)  # torch download
                            success = f.is_file()
                        if success:
                            break
                        elif i < retry:
                            LOGGER.warning(f"⚠️ Download failure, retrying {i + 1}/{retry} {url}...")
                        else:
                            LOGGER.warning(f"❌ Failed to download {url}...")
        
                if unzip and success and (f.suffix == ".gz" or is_zipfile(f) or is_tarfile(f)):
                    LOGGER.info(f"Unzipping {f}...")
                    if is_zipfile(f):
                        unzip_file(f, dir)  # unzip
                    elif is_tarfile(f):
                        subprocess.run(["tar", "xf", f, "--directory", f.parent], check=True)  # unzip
                    elif f.suffix == ".gz":
                        subprocess.run(["tar", "xfz", f, "--directory", f.parent], check=True)  # unzip
                    if delete:
                        f.unlink()  # remove zip
        
            dir = Path(dir)
            dir.mkdir(parents=True, exist_ok=True)  # make directory
            if threads > 1:
                pool = ThreadPool(threads)
                pool.imap(lambda x: download_one(*x), zip(url, repeat(dir)))  # multithreaded
                pool.close()
                pool.join()
            else:
                for u in [url] if isinstance(url, (str, Path)) else url:
                    download_one(u, dir)
        
        
        def make_divisible(x, divisor):
            """Adjusts `x` to be divisible by `divisor`, returning the nearest greater or equal value."""
            if isinstance(divisor, torch.Tensor):
                divisor = int(divisor.max())  # to int
            return math.ceil(x / divisor) * divisor
        
        
        def clean_str(s):
            """Cleans a string by replacing special characters with underscore, e.g., `clean_str('#example!')` returns
            '_example_'.
            """
            return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s)
        
        
        def one_cycle(y1=0.0, y2=1.0, steps=100):
            """
            Generates a lambda for a sinusoidal ramp from y1 to y2 over 'steps'.
        
            See https://arxiv.org/pdf/1812.01187.pdf for details.
            """
            return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1
        
        
        def colorstr(*input):
            """
            Colors a string using ANSI escape codes, e.g., colorstr('blue', 'hello world').
        
            See https://en.wikipedia.org/wiki/ANSI_escape_code.
            """
            *args, string = input if len(input) > 1 else ("blue", "bold", input[0])  # color arguments, string
            colors = {
                "black": "\033[30m",  # basic colors
                "red": "\033[31m",
                "green": "\033[32m",
                "yellow": "\033[33m",
                "blue": "\033[34m",
                "magenta": "\033[35m",
                "cyan": "\033[36m",
                "white": "\033[37m",
                "bright_black": "\033[90m",  # bright colors
                "bright_red": "\033[91m",
                "bright_green": "\033[92m",
                "bright_yellow": "\033[93m",
                "bright_blue": "\033[94m",
                "bright_magenta": "\033[95m",
                "bright_cyan": "\033[96m",
                "bright_white": "\033[97m",
                "end": "\033[0m",  # misc
                "bold": "\033[1m",
                "underline": "\033[4m",
            }
            return "".join(colors[x] for x in args) + f"{string}" + colors["end"]
        
        
        def labels_to_class_weights(labels, nc=80):
            """Calculates class weights from labels to handle class imbalance in training; input shape: (n, 5)."""
            if labels[0] is None:  # no labels loaded
                return torch.Tensor()
        
            labels = np.concatenate(labels, 0)  # labels.shape = (866643, 5) for COCO
            classes = labels[:, 0].astype(int)  # labels = [class xywh]
            weights = np.bincount(classes, minlength=nc)  # occurrences per class
        
            # Prepend gridpoint count (for uCE training)
            # gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum()  # gridpoints per image
            # weights = np.hstack([gpi * len(labels)  - weights.sum() * 9, weights * 9]) ** 0.5  # prepend gridpoints to start
        
            weights[weights == 0] = 1  # replace empty bins with 1
            weights = 1 / weights  # number of targets per class
            weights /= weights.sum()  # normalize
            return torch.from_numpy(weights).float()
        
        
        def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):
            """Calculates image weights from labels using class weights for weighted sampling."""
            # Usage: index = random.choices(range(n), weights=image_weights, k=1)  # weighted image sample
            class_counts = np.array([np.bincount(x[:, 0].astype(int), minlength=nc) for x in labels])
            return (class_weights.reshape(1, nc) * class_counts).sum(1)
        
        
        def coco80_to_coco91_class():
            """
            Converts COCO 80-class index to COCO 91-class index used in the paper.
        
            Reference: https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
            """
            # a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n')
            # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
            # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)]  # darknet to coco
            # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)]  # coco to darknet
            return [
                1,
                2,
                3,
                4,
                5,
                6,
                7,
                8,
                9,
                10,
                11,
                13,
                14,
                15,
                16,
                17,
                18,
                19,
                20,
                21,
                22,
                23,
                24,
                25,
                27,
                28,
                31,
                32,
                33,
                34,
                35,
                36,
                37,
                38,
                39,
                40,
                41,
                42,
                43,
                44,
                46,
                47,
                48,
                49,
                50,
                51,
                52,
                53,
                54,
                55,
                56,
                57,
                58,
                59,
                60,
                61,
                62,
                63,
                64,
                65,
                67,
                70,
                72,
                73,
                74,
                75,
                76,
                77,
                78,
                79,
                80,
                81,
                82,
                84,
                85,
                86,
                87,
                88,
                89,
                90,
            ]
        
        
        def xyxy2xywh(x):
            """Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right."""
            y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
            y[..., 0] = (x[..., 0] + x[..., 2]) / 2  # x center
            y[..., 1] = (x[..., 1] + x[..., 3]) / 2  # y center
            y[..., 2] = x[..., 2] - x[..., 0]  # width
            y[..., 3] = x[..., 3] - x[..., 1]  # height
            return y
        
        
        def xywh2xyxy(x):
            """Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right."""
            y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
            y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
            y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
            y[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right x
            y[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right y
            return y
        
        
        def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
            """Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right."""
            y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
            y[..., 0] = w * (x[..., 0] - x[..., 2] / 2) + padw  # top left x
            y[..., 1] = h * (x[..., 1] - x[..., 3] / 2) + padh  # top left y
            y[..., 2] = w * (x[..., 0] + x[..., 2] / 2) + padw  # bottom right x
            y[..., 3] = h * (x[..., 1] + x[..., 3] / 2) + padh  # bottom right y
            return y
        
        
        def xyxy2xywhn(x, w=640, h=640, clip=False, eps=0.0):
            """Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] normalized where xy1=top-left, xy2=bottom-right."""
            if clip:
                clip_boxes(x, (h - eps, w - eps))  # warning: inplace clip
            y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
            y[..., 0] = ((x[..., 0] + x[..., 2]) / 2) / w  # x center
            y[..., 1] = ((x[..., 1] + x[..., 3]) / 2) / h  # y center
            y[..., 2] = (x[..., 2] - x[..., 0]) / w  # width
            y[..., 3] = (x[..., 3] - x[..., 1]) / h  # height
            return y
        
        
        def xyn2xy(x, w=640, h=640, padw=0, padh=0):
            """Convert normalized segments into pixel segments, shape (n,2)."""
            y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
            y[..., 0] = w * x[..., 0] + padw  # top left x
            y[..., 1] = h * x[..., 1] + padh  # top left y
            return y
        
        
        def segment2box(segment, width=640, height=640):
            """Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy)."""
            x, y = segment.T  # segment xy
            inside = (x >= 0) & (y >= 0) & (x <= width) & (y <= height)
            (
                x,
                y,
            ) = x[inside], y[inside]
            return np.array([x.min(), y.min(), x.max(), y.max()]) if any(x) else np.zeros((1, 4))  # xyxy
        
        
        def segments2boxes(segments):
            """Convert segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh)."""
            boxes = []
            for s in segments:
                x, y = s.T  # segment xy
                boxes.append([x.min(), y.min(), x.max(), y.max()])  # cls, xyxy
            return xyxy2xywh(np.array(boxes))  # cls, xywh
        
        
        def resample_segments(segments, n=1000):
            """Resamples an (n,2) segment to a fixed number of points for consistent representation."""
            for i, s in enumerate(segments):
                s = np.concatenate((s, s[0:1, :]), axis=0)
                x = np.linspace(0, len(s) - 1, n)
                xp = np.arange(len(s))
                segments[i] = np.concatenate([np.interp(x, xp, s[:, i]) for i in range(2)]).reshape(2, -1).T  # segment xy
            return segments
        
        
        def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None):
            """Rescales (xyxy) bounding boxes from img1_shape to img0_shape, optionally using provided `ratio_pad`."""
            if ratio_pad is None:  # calculate from img0_shape
                gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
                pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
            else:
                gain = ratio_pad[0][0]
                pad = ratio_pad[1]
        
            boxes[..., [0, 2]] -= pad[0]  # x padding
            boxes[..., [1, 3]] -= pad[1]  # y padding
            boxes[..., :4] /= gain
            clip_boxes(boxes, img0_shape)
            return boxes
        
        
        def scale_segments(img1_shape, segments, img0_shape, ratio_pad=None, normalize=False):
            """Rescales segment coordinates from img1_shape to img0_shape, optionally normalizing them with custom padding."""
            if ratio_pad is None:  # calculate from img0_shape
                gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
                pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
            else:
                gain = ratio_pad[0][0]
                pad = ratio_pad[1]
        
            segments[:, 0] -= pad[0]  # x padding
            segments[:, 1] -= pad[1]  # y padding
            segments /= gain
            clip_segments(segments, img0_shape)
            if normalize:
                segments[:, 0] /= img0_shape[1]  # width
                segments[:, 1] /= img0_shape[0]  # height
            return segments
        
        
        def clip_boxes(boxes, shape):
            """Clips bounding box coordinates (xyxy) to fit within the specified image shape (height, width)."""
            if isinstance(boxes, torch.Tensor):  # faster individually
                boxes[..., 0].clamp_(0, shape[1])  # x1
                boxes[..., 1].clamp_(0, shape[0])  # y1
                boxes[..., 2].clamp_(0, shape[1])  # x2
                boxes[..., 3].clamp_(0, shape[0])  # y2
            else:  # np.array (faster grouped)
                boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, shape[1])  # x1, x2
                boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0])  # y1, y2
        
        
        def clip_segments(segments, shape):
            """Clips segment coordinates (xy1, xy2, ...) to an image's boundaries given its shape (height, width)."""
            if isinstance(segments, torch.Tensor):  # faster individually
                segments[:, 0].clamp_(0, shape[1])  # x
                segments[:, 1].clamp_(0, shape[0])  # y
            else:  # np.array (faster grouped)
                segments[:, 0] = segments[:, 0].clip(0, shape[1])  # x
                segments[:, 1] = segments[:, 1].clip(0, shape[0])  # y
        
        
        def non_max_suppression(
            prediction,
            conf_thres=0.25,
            iou_thres=0.45,
            classes=None,
            agnostic=False,
            multi_label=False,
            labels=(),
            max_det=300,
            nm=0,  # number of masks
        ):
            """
            Non-Maximum Suppression (NMS) on inference results to reject overlapping detections.
        
            Returns:
                 list of detections, on (n,6) tensor per image [xyxy, conf, cls]
            """
        
            # Checks
            assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0"
            assert 0 <= iou_thres <= 1, f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0"
            if isinstance(prediction, (list, tuple)):  # YOLOv5 model in validation model, output = (inference_out, loss_out)
                prediction = prediction[0]  # select only inference output
        
            device = prediction.device
            mps = "mps" in device.type  # Apple MPS
            if mps:  # MPS not fully supported yet, convert tensors to CPU before NMS
                prediction = prediction.cpu()
            bs = prediction.shape[0]  # batch size
            nc = prediction.shape[2] - nm - 5  # number of classes
            xc = prediction[..., 4] > conf_thres  # candidates
        
            # Settings
            # min_wh = 2  # (pixels) minimum box width and height
            max_wh = 7680  # (pixels) maximum box width and height
            max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
            time_limit = 0.5 + 0.05 * bs  # seconds to quit after
            redundant = True  # require redundant detections
            multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
            merge = False  # use merge-NMS
        
            t = time.time()
            mi = 5 + nc  # mask start index
            output = [torch.zeros((0, 6 + nm), device=prediction.device)] * bs
            for xi, x in enumerate(prediction):  # image index, image inference
                # Apply constraints
                # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
                x = x[xc[xi]]  # confidence
        
                # Cat apriori labels if autolabelling
                if labels and len(labels[xi]):
                    lb = labels[xi]
                    v = torch.zeros((len(lb), nc + nm + 5), device=x.device)
                    v[:, :4] = lb[:, 1:5]  # box
                    v[:, 4] = 1.0  # conf
                    v[range(len(lb)), lb[:, 0].long() + 5] = 1.0  # cls
                    x = torch.cat((x, v), 0)
        
                # If none remain process next image
                if not x.shape[0]:
                    continue
        
                # Compute conf
                x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
        
                # Box/Mask
                box = xywh2xyxy(x[:, :4])  # center_x, center_y, width, height) to (x1, y1, x2, y2)
                mask = x[:, mi:]  # zero columns if no masks
        
                # Detections matrix nx6 (xyxy, conf, cls)
                if multi_label:
                    i, j = (x[:, 5:mi] > conf_thres).nonzero(as_tuple=False).T
                    x = torch.cat((box[i], x[i, 5 + j, None], j[:, None].float(), mask[i]), 1)
                else:  # best class only
                    conf, j = x[:, 5:mi].max(1, keepdim=True)
                    x = torch.cat((box, conf, j.float(), mask), 1)[conf.view(-1) > conf_thres]
        
                # Filter by class
                if classes is not None:
                    x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
        
                # Apply finite constraint
                # if not torch.isfinite(x).all():
                #     x = x[torch.isfinite(x).all(1)]
        
                # Check shape
                n = x.shape[0]  # number of boxes
                if not n:  # no boxes
                    continue
                x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence and remove excess boxes
        
                # Batched NMS
                c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
                boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
                i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
                i = i[:max_det]  # limit detections
                if merge and (1 < n < 3e3):  # Merge NMS (boxes merged using weighted mean)
                    # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
                    iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
                    weights = iou * scores[None]  # box weights
                    x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
                    if redundant:
                        i = i[iou.sum(1) > 1]  # require redundancy
        
                output[xi] = x[i]
                if mps:
                    output[xi] = output[xi].to(device)
                if (time.time() - t) > time_limit:
                    LOGGER.warning(f"WARNING ⚠️ NMS time limit {time_limit:.3f}s exceeded")
                    break  # time limit exceeded
        
            return output
        
        
        def strip_optimizer(f="best.pt", s=""):
            """
            Strips optimizer and optionally saves checkpoint to finalize training; arguments are file path 'f' and save path
            's'.
        
            Example: from utils.general import *; strip_optimizer()
            """
            x = torch.load(f, map_location=torch.device("cpu"),weights_only=False)
            if x.get("ema"):
                x["model"] = x["ema"]  # replace model with ema
            for k in "optimizer", "best_fitness", "ema", "updates":  # keys
                x[k] = None
            x["epoch"] = -1
            x["model"].half()  # to FP16
            for p in x["model"].parameters():
                p.requires_grad = False
            torch.save(x, s or f)
            mb = os.path.getsize(s or f) / 1e6  # filesize
            LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB")
        
        
        def print_mutation(keys, results, hyp, save_dir, bucket, prefix=colorstr("evolve: ")):
            """Logs evolution results and saves to CSV and YAML in `save_dir`, optionally syncs with `bucket`."""
            evolve_csv = save_dir / "evolve.csv"
            evolve_yaml = save_dir / "hyp_evolve.yaml"
            keys = tuple(keys) + tuple(hyp.keys())  # [results + hyps]
            keys = tuple(x.strip() for x in keys)
            vals = results + tuple(hyp.values())
            n = len(keys)
        
            # Download (optional)
            if bucket:
                url = f"gs://{bucket}/evolve.csv"
                if gsutil_getsize(url) > (evolve_csv.stat().st_size if evolve_csv.exists() else 0):
                    subprocess.run(["gsutil", "cp", f"{url}", f"{save_dir}"])  # download evolve.csv if larger than local
        
            # Log to evolve.csv
            s = "" if evolve_csv.exists() else (("%20s," * n % keys).rstrip(",") + "\n")  # add header
            with open(evolve_csv, "a") as f:
                f.write(s + ("%20.5g," * n % vals).rstrip(",") + "\n")
        
            # Save yaml
            with open(evolve_yaml, "w") as f:
                data = pd.read_csv(evolve_csv, skipinitialspace=True)
                data = data.rename(columns=lambda x: x.strip())  # strip keys
                i = np.argmax(fitness(data.values[:, :4]))  #
                generations = len(data)
                f.write(
                    "# YOLOv5 Hyperparameter Evolution Results\n"
                    + f"# Best generation: {i}\n"
                    + f"# Last generation: {generations - 1}\n"
                    + "# "
                    + ", ".join(f"{x.strip():>20s}" for x in keys[:7])
                    + "\n"
                    + "# "
                    + ", ".join(f"{x:>20.5g}" for x in data.values[i, :7])
                    + "\n\n"
                )
                yaml.safe_dump(data.loc[i][7:].to_dict(), f, sort_keys=False)
        
            # Print to screen
            LOGGER.info(
                prefix
                + f"{generations} generations finished, current result:\n"
                + prefix
                + ", ".join(f"{x.strip():>20s}" for x in keys)
                + "\n"
                + prefix
                + ", ".join(f"{x:20.5g}" for x in vals)
                + "\n\n"
            )
        
            if bucket:
                subprocess.run(["gsutil", "cp", f"{evolve_csv}", f"{evolve_yaml}", f"gs://{bucket}"])  # upload
        
        
        def apply_classifier(x, model, img, im0):
            """Applies second-stage classifier to YOLO outputs, filtering detections by class match."""
            # Example model = torchvision.models.__dict__['efficientnet_b0'](pretrained=True).to(device).eval()
            im0 = [im0] if isinstance(im0, np.ndarray) else im0
            for i, d in enumerate(x):  # per image
                if d is not None and len(d):
                    d = d.clone()
        
                    # Reshape and pad cutouts
                    b = xyxy2xywh(d[:, :4])  # boxes
                    b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1)  # rectangle to square
                    b[:, 2:] = b[:, 2:] * 1.3 + 30  # pad
                    d[:, :4] = xywh2xyxy(b).long()
        
                    # Rescale boxes from img_size to im0 size
                    scale_boxes(img.shape[2:], d[:, :4], im0[i].shape)
        
                    # Classes
                    pred_cls1 = d[:, 5].long()
                    ims = []
                    for a in d:
                        cutout = im0[i][int(a[1]) : int(a[3]), int(a[0]) : int(a[2])]
                        im = cv2.resize(cutout, (224, 224))  # BGR
        
                        im = im[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
                        im = np.ascontiguousarray(im, dtype=np.float32)  # uint8 to float32
                        im /= 255  # 0 - 255 to 0.0 - 1.0
                        ims.append(im)
        
                    pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1)  # classifier prediction
                    x[i] = x[i][pred_cls1 == pred_cls2]  # retain matching class detections
        
            return x
        
        
        def increment_path(path, exist_ok=False, sep="", mkdir=False):
            """
            Generates an incremented file or directory path if it exists, with optional mkdir; args: path, exist_ok=False,
            sep="", mkdir=False.
        
            Example: runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc
            """
            path = Path(path)  # os-agnostic
            if path.exists() and not exist_ok:
                path, suffix = (path.with_suffix(""), path.suffix) if path.is_file() else (path, "")
        
                # Method 1
                for n in range(2, 9999):
                    p = f"{path}{sep}{n}{suffix}"  # increment path
                    if not os.path.exists(p):  #
                        break
                path = Path(p)
        
                # Method 2 (deprecated)
                # dirs = glob.glob(f"{path}{sep}*")  # similar paths
                # matches = [re.search(rf"{path.stem}{sep}(\d+)", d) for d in dirs]
                # i = [int(m.groups()[0]) for m in matches if m]  # indices
                # n = max(i) + 1 if i else 2  # increment number
                # path = Path(f"{path}{sep}{n}{suffix}")  # increment path
        
            if mkdir:
                path.mkdir(parents=True, exist_ok=True)  # make directory
        
            return path
        
        
        # OpenCV Multilanguage-friendly functions ------------------------------------------------------------------------------------
        imshow_ = cv2.imshow  # copy to avoid recursion errors
        
        
        def imread(filename, flags=cv2.IMREAD_COLOR):
            """Reads an image from a file and returns it as a numpy array, using OpenCV's imdecode to support multilanguage
            paths.
            """
            return cv2.imdecode(np.fromfile(filename, np.uint8), flags)
        
        
        def imwrite(filename, img):
            """Writes an image to a file, returns True on success and False on failure, supports multilanguage paths."""
            try:
                cv2.imencode(Path(filename).suffix, img)[1].tofile(filename)
                return True
            except Exception:
                return False
        
        
        def imshow(path, im):
            """Displays an image using Unicode path, requires encoded path and image matrix as input."""
            imshow_(path.encode("unicode_escape").decode(), im)
        
        
        if Path(inspect.stack()[0].filename).parent.parent.as_posix() in inspect.stack()[-1].filename:
            cv2.imread, cv2.imwrite, cv2.imshow = imread, imwrite, imshow  # redefine
        
        # Variables ------------------------------------------------------------------------------------------------------------
        
      • 📄 loss.py 10.9KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Loss functions."""
        
        import torch
        import torch.nn as nn
        
        from utils.metrics import bbox_iou
        from utils.torch_utils import de_parallel
        
        
        def smooth_BCE(eps=0.1):
            """Returns label smoothing BCE targets for reducing overfitting; pos: `1.0 - 0.5*eps`, neg: `0.5*eps`. For details see https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441"""
            return 1.0 - 0.5 * eps, 0.5 * eps
        
        
        class BCEBlurWithLogitsLoss(nn.Module):
            # BCEwithLogitLoss() with reduced missing label effects.
            def __init__(self, alpha=0.05):
                """Initializes a modified BCEWithLogitsLoss with reduced missing label effects, taking optional alpha smoothing
                parameter.
                """
                super().__init__()
                self.loss_fcn = nn.BCEWithLogitsLoss(reduction="none")  # must be nn.BCEWithLogitsLoss()
                self.alpha = alpha
        
            def forward(self, pred, true):
                """Computes modified BCE loss for YOLOv5 with reduced missing label effects, taking pred and true tensors,
                returns mean loss.
                """
                loss = self.loss_fcn(pred, true)
                pred = torch.sigmoid(pred)  # prob from logits
                dx = pred - true  # reduce only missing label effects
                # dx = (pred - true).abs()  # reduce missing label and false label effects
                alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4))
                loss *= alpha_factor
                return loss.mean()
        
        
        class FocalLoss(nn.Module):
            # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
            def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
                """Initializes FocalLoss with specified loss function, gamma, and alpha values; modifies loss reduction to
                'none'.
                """
                super().__init__()
                self.loss_fcn = loss_fcn  # must be nn.BCEWithLogitsLoss()
                self.gamma = gamma
                self.alpha = alpha
                self.reduction = loss_fcn.reduction
                self.loss_fcn.reduction = "none"  # required to apply FL to each element
        
            def forward(self, pred, true):
                """Calculates the focal loss between predicted and true labels using a modified BCEWithLogitsLoss."""
                loss = self.loss_fcn(pred, true)
                # p_t = torch.exp(-loss)
                # loss *= self.alpha * (1.000001 - p_t) ** self.gamma  # non-zero power for gradient stability
        
                # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
                pred_prob = torch.sigmoid(pred)  # prob from logits
                p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
                alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
                modulating_factor = (1.0 - p_t) ** self.gamma
                loss *= alpha_factor * modulating_factor
        
                if self.reduction == "mean":
                    return loss.mean()
                elif self.reduction == "sum":
                    return loss.sum()
                else:  # 'none'
                    return loss
        
        
        class QFocalLoss(nn.Module):
            # Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
            def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
                """Initializes Quality Focal Loss with given loss function, gamma, alpha; modifies reduction to 'none'."""
                super().__init__()
                self.loss_fcn = loss_fcn  # must be nn.BCEWithLogitsLoss()
                self.gamma = gamma
                self.alpha = alpha
                self.reduction = loss_fcn.reduction
                self.loss_fcn.reduction = "none"  # required to apply FL to each element
        
            def forward(self, pred, true):
                """Computes the focal loss between `pred` and `true` using BCEWithLogitsLoss, adjusting for imbalance with
                `gamma` and `alpha`.
                """
                loss = self.loss_fcn(pred, true)
        
                pred_prob = torch.sigmoid(pred)  # prob from logits
                alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
                modulating_factor = torch.abs(true - pred_prob) ** self.gamma
                loss *= alpha_factor * modulating_factor
        
                if self.reduction == "mean":
                    return loss.mean()
                elif self.reduction == "sum":
                    return loss.sum()
                else:  # 'none'
                    return loss
        
        
        class ComputeLoss:
            sort_obj_iou = False
        
            # Compute losses
            def __init__(self, model, autobalance=False):
                """Initializes ComputeLoss with model and autobalance option, autobalances losses if True."""
                device = next(model.parameters()).device  # get model device
                h = model.hyp  # hyperparameters
        
                # Define criteria
                BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["cls_pw"]], device=device))
                BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["obj_pw"]], device=device))
        
                # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
                self.cp, self.cn = smooth_BCE(eps=h.get("label_smoothing", 0.0))  # positive, negative BCE targets
        
                # Focal loss
                g = h["fl_gamma"]  # focal loss gamma
                if g > 0:
                    BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
        
                m = de_parallel(model).model[-1]  # Detect() module
                self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl, [4.0, 1.0, 0.25, 0.06, 0.02])  # P3-P7
                self.ssi = list(m.stride).index(16) if autobalance else 0  # stride 16 index
                self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance
                self.na = m.na  # number of anchors
                self.nc = m.nc  # number of classes
                self.nl = m.nl  # number of layers
                self.anchors = m.anchors
                self.device = device
        
            def __call__(self, p, targets):  # predictions, targets
                """Performs forward pass, calculating class, box, and object loss for given predictions and targets."""
                lcls = torch.zeros(1, device=self.device)  # class loss
                lbox = torch.zeros(1, device=self.device)  # box loss
                lobj = torch.zeros(1, device=self.device)  # object loss
                tcls, tbox, indices, anchors = self.build_targets(p, targets)  # targets
        
                # Losses
                for i, pi in enumerate(p):  # layer index, layer predictions
                    b, a, gj, gi = indices[i]  # image, anchor, gridy, gridx
                    tobj = torch.zeros(pi.shape[:4], dtype=pi.dtype, device=self.device)  # target obj
        
                    n = b.shape[0]  # number of targets
                    if n:
                        # pxy, pwh, _, pcls = pi[b, a, gj, gi].tensor_split((2, 4, 5), dim=1)  # faster, requires torch 1.8.0
                        pxy, pwh, _, pcls = pi[b, a, gj, gi].split((2, 2, 1, self.nc), 1)  # target-subset of predictions
        
                        # Regression
                        pxy = pxy.sigmoid() * 2 - 0.5
                        pwh = (pwh.sigmoid() * 2) ** 2 * anchors[i]
                        pbox = torch.cat((pxy, pwh), 1)  # predicted box
                        iou = bbox_iou(pbox, tbox[i], CIoU=True).squeeze()  # iou(prediction, target)
                        lbox += (1.0 - iou).mean()  # iou loss
        
                        # Objectness
                        iou = iou.detach().clamp(0).type(tobj.dtype)
                        if self.sort_obj_iou:
                            j = iou.argsort()
                            b, a, gj, gi, iou = b[j], a[j], gj[j], gi[j], iou[j]
                        if self.gr < 1:
                            iou = (1.0 - self.gr) + self.gr * iou
                        tobj[b, a, gj, gi] = iou  # iou ratio
        
                        # Classification
                        if self.nc > 1:  # cls loss (only if multiple classes)
                            t = torch.full_like(pcls, self.cn, device=self.device)  # targets
                            t[range(n), tcls[i]] = self.cp
                            lcls += self.BCEcls(pcls, t)  # BCE
        
                        # Append targets to text file
                        # with open('targets.txt', 'a') as file:
                        #     [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
        
                    obji = self.BCEobj(pi[..., 4], tobj)
                    lobj += obji * self.balance[i]  # obj loss
                    if self.autobalance:
                        self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()
        
                if self.autobalance:
                    self.balance = [x / self.balance[self.ssi] for x in self.balance]
                lbox *= self.hyp["box"]
                lobj *= self.hyp["obj"]
                lcls *= self.hyp["cls"]
                bs = tobj.shape[0]  # batch size
        
                return (lbox + lobj + lcls) * bs, torch.cat((lbox, lobj, lcls)).detach()
        
            def build_targets(self, p, targets):
                """Prepares model targets from input targets (image,class,x,y,w,h) for loss computation, returning class, box,
                indices, and anchors.
                """
                na, nt = self.na, targets.shape[0]  # number of anchors, targets
                tcls, tbox, indices, anch = [], [], [], []
                gain = torch.ones(7, device=self.device)  # normalized to gridspace gain
                ai = torch.arange(na, device=self.device).float().view(na, 1).repeat(1, nt)  # same as .repeat_interleave(nt)
                targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None]), 2)  # append anchor indices
        
                g = 0.5  # bias
                off = (
                    torch.tensor(
                        [
                            [0, 0],
                            [1, 0],
                            [0, 1],
                            [-1, 0],
                            [0, -1],  # j,k,l,m
                            # [1, 1], [1, -1], [-1, 1], [-1, -1],  # jk,jm,lk,lm
                        ],
                        device=self.device,
                    ).float()
                    * g
                )  # offsets
        
                for i in range(self.nl):
                    anchors, shape = self.anchors[i], p[i].shape
                    gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]]  # xyxy gain
        
                    # Match targets to anchors
                    t = targets * gain  # shape(3,n,7)
                    if nt:
                        # Matches
                        r = t[..., 4:6] / anchors[:, None]  # wh ratio
                        j = torch.max(r, 1 / r).max(2)[0] < self.hyp["anchor_t"]  # compare
                        # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']  # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
                        t = t[j]  # filter
        
                        # Offsets
                        gxy = t[:, 2:4]  # grid xy
                        gxi = gain[[2, 3]] - gxy  # inverse
                        j, k = ((gxy % 1 < g) & (gxy > 1)).T
                        l, m = ((gxi % 1 < g) & (gxi > 1)).T
                        j = torch.stack((torch.ones_like(j), j, k, l, m))
                        t = t.repeat((5, 1, 1))[j]
                        offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
                    else:
                        t = targets[0]
                        offsets = 0
        
                    # Define
                    bc, gxy, gwh, a = t.chunk(4, 1)  # (image, class), grid xy, grid wh, anchors
                    a, (b, c) = a.long().view(-1), bc.long().T  # anchors, image, class
                    gij = (gxy - offsets).long()
                    gi, gj = gij.T  # grid indices
        
                    # Append
                    indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1)))  # image, anchor, grid
                    tbox.append(torch.cat((gxy - gij, gwh), 1))  # box
                    anch.append(anchors[a])  # anchors
                    tcls.append(c)  # class
        
                return tcls, tbox, indices, anch
        
      • 📄 metrics.py 15.1KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Model validation metrics."""
        
        import math
        import warnings
        from pathlib import Path
        
        import matplotlib.pyplot as plt
        import numpy as np
        import torch
        
        from utils import TryExcept, threaded
        
        
        def fitness(x):
            """Calculates fitness of a model using weighted sum of metrics P, R, mAP@0.5, mAP@0.5:0.95."""
            w = [0.0, 0.0, 0.1, 0.9]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
            return (x[:, :4] * w).sum(1)
        
        
        def smooth(y, f=0.05):
            """Applies box filter smoothing to array `y` with fraction `f`, yielding a smoothed array."""
            nf = round(len(y) * f * 2) // 2 + 1  # number of filter elements (must be odd)
            p = np.ones(nf // 2)  # ones padding
            yp = np.concatenate((p * y[0], y, p * y[-1]), 0)  # y padded
            return np.convolve(yp, np.ones(nf) / nf, mode="valid")  # y-smoothed
        
        
        def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir=".", names=(), eps=1e-16, prefix=""):
            """
            Compute the average precision, given the recall and precision curves.
        
            Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
            # Arguments
                tp:  True positives (nparray, nx1 or nx10).
                conf:  Objectness value from 0-1 (nparray).
                pred_cls:  Predicted object classes (nparray).
                target_cls:  True object classes (nparray).
                plot:  Plot precision-recall curve at mAP@0.5
                save_dir:  Plot save directory
            # Returns
                The average precision as computed in py-faster-rcnn.
            """
        
            # Sort by objectness
            i = np.argsort(-conf)
            tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]
        
            # Find unique classes
            unique_classes, nt = np.unique(target_cls, return_counts=True)
            nc = unique_classes.shape[0]  # number of classes, number of detections
        
            # Create Precision-Recall curve and compute AP for each class
            px, py = np.linspace(0, 1, 1000), []  # for plotting
            ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000))
            for ci, c in enumerate(unique_classes):
                i = pred_cls == c
                n_l = nt[ci]  # number of labels
                n_p = i.sum()  # number of predictions
                if n_p == 0 or n_l == 0:
                    continue
        
                # Accumulate FPs and TPs
                fpc = (1 - tp[i]).cumsum(0)
                tpc = tp[i].cumsum(0)
        
                # Recall
                recall = tpc / (n_l + eps)  # recall curve
                r[ci] = np.interp(-px, -conf[i], recall[:, 0], left=0)  # negative x, xp because xp decreases
        
                # Precision
                precision = tpc / (tpc + fpc)  # precision curve
                p[ci] = np.interp(-px, -conf[i], precision[:, 0], left=1)  # p at pr_score
        
                # AP from recall-precision curve
                for j in range(tp.shape[1]):
                    ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j])
                    if plot and j == 0:
                        py.append(np.interp(px, mrec, mpre))  # precision at mAP@0.5
        
            # Compute F1 (harmonic mean of precision and recall)
            f1 = 2 * p * r / (p + r + eps)
            names = [v for k, v in names.items() if k in unique_classes]  # list: only classes that have data
            names = dict(enumerate(names))  # to dict
            if plot:
                plot_pr_curve(px, py, ap, Path(save_dir) / f"{prefix}PR_curve.png", names)
                plot_mc_curve(px, f1, Path(save_dir) / f"{prefix}F1_curve.png", names, ylabel="F1")
                plot_mc_curve(px, p, Path(save_dir) / f"{prefix}P_curve.png", names, ylabel="Precision")
                plot_mc_curve(px, r, Path(save_dir) / f"{prefix}R_curve.png", names, ylabel="Recall")
        
            i = smooth(f1.mean(0), 0.1).argmax()  # max F1 index
            p, r, f1 = p[:, i], r[:, i], f1[:, i]
            tp = (r * nt).round()  # true positives
            fp = (tp / (p + eps) - tp).round()  # false positives
            return tp, fp, p, r, f1, ap, unique_classes.astype(int)
        
        
        def compute_ap(recall, precision):
            """Compute the average precision, given the recall and precision curves
            # Arguments
                recall:    The recall curve (list)
                precision: The precision curve (list)
            # Returns
                Average precision, precision curve, recall curve
            """
        
            # Append sentinel values to beginning and end
            mrec = np.concatenate(([0.0], recall, [1.0]))
            mpre = np.concatenate(([1.0], precision, [0.0]))
        
            # Compute the precision envelope
            mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
        
            # Integrate area under curve
            method = "interp"  # methods: 'continuous', 'interp'
            if method == "interp":
                x = np.linspace(0, 1, 101)  # 101-point interp (COCO)
                ap = np.trapz(np.interp(x, mrec, mpre), x)  # integrate
            else:  # 'continuous'
                i = np.where(mrec[1:] != mrec[:-1])[0]  # points where x axis (recall) changes
                ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])  # area under curve
        
            return ap, mpre, mrec
        
        
        class ConfusionMatrix:
            # Updated version of https://github.com/kaanakan/object_detection_confusion_matrix
            def __init__(self, nc, conf=0.25, iou_thres=0.45):
                """Initializes ConfusionMatrix with given number of classes, confidence, and IoU threshold."""
                self.matrix = np.zeros((nc + 1, nc + 1))
                self.nc = nc  # number of classes
                self.conf = conf
                self.iou_thres = iou_thres
        
            def process_batch(self, detections, labels):
                """
                Return intersection-over-union (Jaccard index) of boxes.
        
                Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
                Arguments:
                    detections (Array[N, 6]), x1, y1, x2, y2, conf, class
                    labels (Array[M, 5]), class, x1, y1, x2, y2
                Returns:
                    None, updates confusion matrix accordingly
                """
                if detections is None:
                    gt_classes = labels.int()
                    for gc in gt_classes:
                        self.matrix[self.nc, gc] += 1  # background FN
                    return
        
                detections = detections[detections[:, 4] > self.conf]
                gt_classes = labels[:, 0].int()
                detection_classes = detections[:, 5].int()
                iou = box_iou(labels[:, 1:], detections[:, :4])
        
                x = torch.where(iou > self.iou_thres)
                if x[0].shape[0]:
                    matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()
                    if x[0].shape[0] > 1:
                        matches = matches[matches[:, 2].argsort()[::-1]]
                        matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
                        matches = matches[matches[:, 2].argsort()[::-1]]
                        matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
                else:
                    matches = np.zeros((0, 3))
        
                n = matches.shape[0] > 0
                m0, m1, _ = matches.transpose().astype(int)
                for i, gc in enumerate(gt_classes):
                    j = m0 == i
                    if n and sum(j) == 1:
                        self.matrix[detection_classes[m1[j]], gc] += 1  # correct
                    else:
                        self.matrix[self.nc, gc] += 1  # true background
        
                if n:
                    for i, dc in enumerate(detection_classes):
                        if not any(m1 == i):
                            self.matrix[dc, self.nc] += 1  # predicted background
        
            def tp_fp(self):
                """Calculates true positives (tp) and false positives (fp) excluding the background class from the confusion
                matrix.
                """
                tp = self.matrix.diagonal()  # true positives
                fp = self.matrix.sum(1) - tp  # false positives
                # fn = self.matrix.sum(0) - tp  # false negatives (missed detections)
                return tp[:-1], fp[:-1]  # remove background class
        
            @TryExcept("WARNING ⚠️ ConfusionMatrix plot failure")
            def plot(self, normalize=True, save_dir="", names=()):
                """Plots confusion matrix using seaborn, optional normalization; can save plot to specified directory."""
                import seaborn as sn
        
                array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1e-9) if normalize else 1)  # normalize columns
                array[array < 0.005] = np.nan  # don't annotate (would appear as 0.00)
        
                fig, ax = plt.subplots(1, 1, figsize=(12, 9), tight_layout=True)
                nc, nn = self.nc, len(names)  # number of classes, names
                sn.set(font_scale=1.0 if nc < 50 else 0.8)  # for label size
                labels = (0 < nn < 99) and (nn == nc)  # apply names to ticklabels
                ticklabels = (names + ["background"]) if labels else "auto"
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")  # suppress empty matrix RuntimeWarning: All-NaN slice encountered
                    sn.heatmap(
                        array,
                        ax=ax,
                        annot=nc < 30,
                        annot_kws={"size": 8},
                        cmap="Blues",
                        fmt=".2f",
                        square=True,
                        vmin=0.0,
                        xticklabels=ticklabels,
                        yticklabels=ticklabels,
                    ).set_facecolor((1, 1, 1))
                ax.set_xlabel("True")
                ax.set_ylabel("Predicted")
                ax.set_title("Confusion Matrix")
                fig.savefig(Path(save_dir) / "confusion_matrix.png", dpi=250)
                plt.close(fig)
        
            def print(self):
                """Prints the confusion matrix row-wise, with each class and its predictions separated by spaces."""
                for i in range(self.nc + 1):
                    print(" ".join(map(str, self.matrix[i])))
        
        
        def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
            """
            Calculates IoU, GIoU, DIoU, or CIoU between two boxes, supporting xywh/xyxy formats.
        
            Input shapes are box1(1,4) to box2(n,4).
            """
        
            # Get the coordinates of bounding boxes
            if xywh:  # transform from xywh to xyxy
                (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1)
                w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
                b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
                b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
            else:  # x1, y1, x2, y2 = box1
                b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1)
                b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1)
                w1, h1 = b1_x2 - b1_x1, (b1_y2 - b1_y1).clamp(eps)
                w2, h2 = b2_x2 - b2_x1, (b2_y2 - b2_y1).clamp(eps)
        
            # Intersection area
            inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp(0) * (
                b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)
            ).clamp(0)
        
            # Union Area
            union = w1 * h1 + w2 * h2 - inter + eps
        
            # IoU
            iou = inter / union
            if CIoU or DIoU or GIoU:
                cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1)  # convex (smallest enclosing box) width
                ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1)  # convex height
                if CIoU or DIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
                    c2 = cw**2 + ch**2 + eps  # convex diagonal squared
                    rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center dist ** 2
                    if CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
                        v = (4 / math.pi**2) * (torch.atan(w2 / h2) - torch.atan(w1 / h1)).pow(2)
                        with torch.no_grad():
                            alpha = v / (v - iou + (1 + eps))
                        return iou - (rho2 / c2 + v * alpha)  # CIoU
                    return iou - rho2 / c2  # DIoU
                c_area = cw * ch + eps  # convex area
                return iou - (c_area - union) / c_area  # GIoU https://arxiv.org/pdf/1902.09630.pdf
            return iou  # IoU
        
        
        def box_iou(box1, box2, eps=1e-7):
            # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
            """
            Return intersection-over-union (Jaccard index) of boxes.
        
            Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
            Arguments:
                box1 (Tensor[N, 4])
                box2 (Tensor[M, 4])
            Returns:
                iou (Tensor[N, M]): the NxM matrix containing the pairwise
                    IoU values for every element in boxes1 and boxes2
            """
        
            # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
            (a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2)
            inter = (torch.min(a2, b2) - torch.max(a1, b1)).clamp(0).prod(2)
        
            # IoU = inter / (area1 + area2 - inter)
            return inter / ((a2 - a1).prod(2) + (b2 - b1).prod(2) - inter + eps)
        
        
        def bbox_ioa(box1, box2, eps=1e-7):
            """
            Returns the intersection over box2 area given box1, box2.
        
            Boxes are x1y1x2y2
            box1:       np.array of shape(4)
            box2:       np.array of shape(nx4)
            returns:    np.array of shape(n)
            """
        
            # Get the coordinates of bounding boxes
            b1_x1, b1_y1, b1_x2, b1_y2 = box1
            b2_x1, b2_y1, b2_x2, b2_y2 = box2.T
        
            # Intersection area
            inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * (
                np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)
            ).clip(0)
        
            # box2 area
            box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + eps
        
            # Intersection over box2 area
            return inter_area / box2_area
        
        
        def wh_iou(wh1, wh2, eps=1e-7):
            """Calculates the Intersection over Union (IoU) for two sets of widths and heights; `wh1` and `wh2` should be nx2
            and mx2 tensors.
            """
            wh1 = wh1[:, None]  # [N,1,2]
            wh2 = wh2[None]  # [1,M,2]
            inter = torch.min(wh1, wh2).prod(2)  # [N,M]
            return inter / (wh1.prod(2) + wh2.prod(2) - inter + eps)  # iou = inter / (area1 + area2 - inter)
        
        
        # Plots ----------------------------------------------------------------------------------------------------------------
        
        
        @threaded
        def plot_pr_curve(px, py, ap, save_dir=Path("pr_curve.png"), names=()):
            """Plots precision-recall curve, optionally per class, saving to `save_dir`; `px`, `py` are lists, `ap` is Nx2
            array, `names` optional.
            """
            fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
            py = np.stack(py, axis=1)
        
            if 0 < len(names) < 21:  # display per-class legend if < 21 classes
                for i, y in enumerate(py.T):
                    ax.plot(px, y, linewidth=1, label=f"{names[i]} {ap[i, 0]:.3f}")  # plot(recall, precision)
            else:
                ax.plot(px, py, linewidth=1, color="grey")  # plot(recall, precision)
        
            ax.plot(px, py.mean(1), linewidth=3, color="blue", label="all classes %.3f mAP@0.5" % ap[:, 0].mean())
            ax.set_xlabel("Recall")
            ax.set_ylabel("Precision")
            ax.set_xlim(0, 1)
            ax.set_ylim(0, 1)
            ax.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
            ax.set_title("Precision-Recall Curve")
            fig.savefig(save_dir, dpi=250)
            plt.close(fig)
        
        
        @threaded
        def plot_mc_curve(px, py, save_dir=Path("mc_curve.png"), names=(), xlabel="Confidence", ylabel="Metric"):
            """Plots a metric-confidence curve for model predictions, supporting per-class visualization and smoothing."""
            fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
        
            if 0 < len(names) < 21:  # display per-class legend if < 21 classes
                for i, y in enumerate(py):
                    ax.plot(px, y, linewidth=1, label=f"{names[i]}")  # plot(confidence, metric)
            else:
                ax.plot(px, py.T, linewidth=1, color="grey")  # plot(confidence, metric)
        
            y = smooth(py.mean(0), 0.05)
            ax.plot(px, y, linewidth=3, color="blue", label=f"all classes {y.max():.2f} at {px[y.argmax()]:.3f}")
            ax.set_xlabel(xlabel)
            ax.set_ylabel(ylabel)
            ax.set_xlim(0, 1)
            ax.set_ylim(0, 1)
            ax.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
            ax.set_title(f"{ylabel}-Confidence Curve")
            fig.savefig(save_dir, dpi=250)
            plt.close(fig)
        
      • 📄 plots.py 20.0KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Plotting utils."""
        
        import contextlib
        import math
        import os
        from copy import copy
        from pathlib import Path
        
        import cv2
        import matplotlib
        import matplotlib.pyplot as plt
        import numpy as np
        import pandas as pd
        import seaborn as sn
        import torch
        from PIL import Image, ImageDraw
        from scipy.ndimage.filters import gaussian_filter1d
        from ultralytics.utils.plotting import Annotator
        
        from utils import TryExcept, threaded
        from utils.general import LOGGER, clip_boxes, increment_path, xywh2xyxy, xyxy2xywh
        from utils.metrics import fitness
        
        # Settings
        RANK = int(os.getenv("RANK", -1))
        matplotlib.rc("font", **{"size": 11})
        matplotlib.use("Agg")  # for writing to files only
        
        
        class Colors:
            # Ultralytics color palette https://ultralytics.com/
            def __init__(self):
                """
                Initializes the Colors class with a palette derived from Ultralytics color scheme, converting hex codes to RGB.
        
                Colors derived from `hex = matplotlib.colors.TABLEAU_COLORS.values()`.
                """
                hexs = (
                    "FF3838",
                    "FF9D97",
                    "FF701F",
                    "FFB21D",
                    "CFD231",
                    "48F90A",
                    "92CC17",
                    "3DDB86",
                    "1A9334",
                    "00D4BB",
                    "2C99A8",
                    "00C2FF",
                    "344593",
                    "6473FF",
                    "0018EC",
                    "8438FF",
                    "520085",
                    "CB38FF",
                    "FF95C8",
                    "FF37C7",
                )
                self.palette = [self.hex2rgb(f"#{c}") for c in hexs]
                self.n = len(self.palette)
        
            def __call__(self, i, bgr=False):
                """Returns color from palette by index `i`, in BGR format if `bgr=True`, else RGB; `i` is an integer index."""
                c = self.palette[int(i) % self.n]
                return (c[2], c[1], c[0]) if bgr else c
        
            @staticmethod
            def hex2rgb(h):
                """Converts hexadecimal color `h` to an RGB tuple (PIL-compatible) with order (R, G, B)."""
                return tuple(int(h[1 + i : 1 + i + 2], 16) for i in (0, 2, 4))
        
        
        colors = Colors()  # create instance for 'from utils.plots import colors'
        
        
        def feature_visualization(x, module_type, stage, n=32, save_dir=Path("runs/detect/exp")):
            """
            x:              Features to be visualized
            module_type:    Module type
            stage:          Module stage within model
            n:              Maximum number of feature maps to plot
            save_dir:       Directory to save results
            """
            if ("Detect" not in module_type) and (
                "Segment" not in module_type
            ):  # 'Detect' for Object Detect task,'Segment' for Segment task
                batch, channels, height, width = x.shape  # batch, channels, height, width
                if height > 1 and width > 1:
                    f = save_dir / f"stage{stage}_{module_type.split('.')[-1]}_features.png"  # filename
        
                    blocks = torch.chunk(x[0].cpu(), channels, dim=0)  # select batch index 0, block by channels
                    n = min(n, channels)  # number of plots
                    fig, ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True)  # 8 rows x n/8 cols
                    ax = ax.ravel()
                    plt.subplots_adjust(wspace=0.05, hspace=0.05)
                    for i in range(n):
                        ax[i].imshow(blocks[i].squeeze())  # cmap='gray'
                        ax[i].axis("off")
        
                    LOGGER.info(f"Saving {f}... ({n}/{channels})")
                    plt.savefig(f, dpi=300, bbox_inches="tight")
                    plt.close()
                    np.save(str(f.with_suffix(".npy")), x[0].cpu().numpy())  # npy save
        
        
        def hist2d(x, y, n=100):
            """
            Generates a logarithmic 2D histogram, useful for visualizing label or evolution distributions.
        
            Used in used in labels.png and evolve.png.
            """
            xedges, yedges = np.linspace(x.min(), x.max(), n), np.linspace(y.min(), y.max(), n)
            hist, xedges, yedges = np.histogram2d(x, y, (xedges, yedges))
            xidx = np.clip(np.digitize(x, xedges) - 1, 0, hist.shape[0] - 1)
            yidx = np.clip(np.digitize(y, yedges) - 1, 0, hist.shape[1] - 1)
            return np.log(hist[xidx, yidx])
        
        
        def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
            """Applies a low-pass Butterworth filter to `data` with specified `cutoff`, `fs`, and `order`."""
            from scipy.signal import butter, filtfilt
        
            # https://stackoverflow.com/questions/28536191/how-to-filter-smooth-with-scipy-numpy
            def butter_lowpass(cutoff, fs, order):
                nyq = 0.5 * fs
                normal_cutoff = cutoff / nyq
                return butter(order, normal_cutoff, btype="low", analog=False)
        
            b, a = butter_lowpass(cutoff, fs, order=order)
            return filtfilt(b, a, data)  # forward-backward filter
        
        
        def output_to_target(output, max_det=300):
            """Converts YOLOv5 model output to [batch_id, class_id, x, y, w, h, conf] format for plotting, limiting detections
            to `max_det`.
            """
            targets = []
            for i, o in enumerate(output):
                box, conf, cls = o[:max_det, :6].cpu().split((4, 1, 1), 1)
                j = torch.full((conf.shape[0], 1), i)
                targets.append(torch.cat((j, cls, xyxy2xywh(box), conf), 1))
            return torch.cat(targets, 0).numpy()
        
        
        @threaded
        def plot_images(images, targets, paths=None, fname="images.jpg", names=None):
            """Plots an image grid with labels from YOLOv5 predictions or targets, saving to `fname`."""
            if isinstance(images, torch.Tensor):
                images = images.cpu().float().numpy()
            if isinstance(targets, torch.Tensor):
                targets = targets.cpu().numpy()
        
            max_size = 1920  # max image size
            max_subplots = 16  # max image subplots, i.e. 4x4
            bs, _, h, w = images.shape  # batch size, _, height, width
            bs = min(bs, max_subplots)  # limit plot images
            ns = np.ceil(bs**0.5)  # number of subplots (square)
            if np.max(images[0]) <= 1:
                images *= 255  # de-normalise (optional)
        
            # Build Image
            mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8)  # init
            for i, im in enumerate(images):
                if i == max_subplots:  # if last batch has fewer images than we expect
                    break
                x, y = int(w * (i // ns)), int(h * (i % ns))  # block origin
                im = im.transpose(1, 2, 0)
                mosaic[y : y + h, x : x + w, :] = im
        
            # Resize (optional)
            scale = max_size / ns / max(h, w)
            if scale < 1:
                h = math.ceil(scale * h)
                w = math.ceil(scale * w)
                mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h)))
        
            # Annotate
            fs = int((h + w) * ns * 0.01)  # font size
            annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names)
            for i in range(i + 1):
                x, y = int(w * (i // ns)), int(h * (i % ns))  # block origin
                annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2)  # borders
                if paths:
                    annotator.text([x + 5, y + 5], text=Path(paths[i]).name[:40], txt_color=(220, 220, 220))  # filenames
                if len(targets) > 0:
                    ti = targets[targets[:, 0] == i]  # image targets
                    boxes = xywh2xyxy(ti[:, 2:6]).T
                    classes = ti[:, 1].astype("int")
                    labels = ti.shape[1] == 6  # labels if no conf column
                    conf = None if labels else ti[:, 6]  # check for confidence presence (label vs pred)
        
                    if boxes.shape[1]:
                        if boxes.max() <= 1.01:  # if normalized with tolerance 0.01
                            boxes[[0, 2]] *= w  # scale to pixels
                            boxes[[1, 3]] *= h
                        elif scale < 1:  # absolute coords need scale if image scales
                            boxes *= scale
                    boxes[[0, 2]] += x
                    boxes[[1, 3]] += y
                    for j, box in enumerate(boxes.T.tolist()):
                        cls = classes[j]
                        color = colors(cls)
                        cls = names[cls] if names else cls
                        if labels or conf[j] > 0.25:  # 0.25 conf thresh
                            label = f"{cls}" if labels else f"{cls} {conf[j]:.1f}"
                            annotator.box_label(box, label, color=color)
            annotator.im.save(fname)  # save
        
        
        def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=""):
            """Plots learning rate schedule for given optimizer and scheduler, saving plot to `save_dir`."""
            optimizer, scheduler = copy(optimizer), copy(scheduler)  # do not modify originals
            y = []
            for _ in range(epochs):
                scheduler.step()
                y.append(optimizer.param_groups[0]["lr"])
            plt.plot(y, ".-", label="LR")
            plt.xlabel("epoch")
            plt.ylabel("LR")
            plt.grid()
            plt.xlim(0, epochs)
            plt.ylim(0)
            plt.savefig(Path(save_dir) / "LR.png", dpi=200)
            plt.close()
        
        
        def plot_val_txt():
            """
            Plots 2D and 1D histograms of bounding box centers from 'val.txt' using matplotlib, saving as 'hist2d.png' and
            'hist1d.png'.
        
            Example: from utils.plots import *; plot_val()
            """
            x = np.loadtxt("val.txt", dtype=np.float32)
            box = xyxy2xywh(x[:, :4])
            cx, cy = box[:, 0], box[:, 1]
        
            fig, ax = plt.subplots(1, 1, figsize=(6, 6), tight_layout=True)
            ax.hist2d(cx, cy, bins=600, cmax=10, cmin=0)
            ax.set_aspect("equal")
            plt.savefig("hist2d.png", dpi=300)
        
            fig, ax = plt.subplots(1, 2, figsize=(12, 6), tight_layout=True)
            ax[0].hist(cx, bins=600)
            ax[1].hist(cy, bins=600)
            plt.savefig("hist1d.png", dpi=200)
        
        
        def plot_targets_txt():
            """
            Plots histograms of object detection targets from 'targets.txt', saving the figure as 'targets.jpg'.
        
            Example: from utils.plots import *; plot_targets_txt()
            """
            x = np.loadtxt("targets.txt", dtype=np.float32).T
            s = ["x targets", "y targets", "width targets", "height targets"]
            fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)
            ax = ax.ravel()
            for i in range(4):
                ax[i].hist(x[i], bins=100, label=f"{x[i].mean():.3g} +/- {x[i].std():.3g}")
                ax[i].legend()
                ax[i].set_title(s[i])
            plt.savefig("targets.jpg", dpi=200)
        
        
        def plot_val_study(file="", dir="", x=None):
            """
            Plots validation study results from 'study*.txt' files in a directory or a specific file, comparing model
            performance and speed.
        
            Example: from utils.plots import *; plot_val_study()
            """
            save_dir = Path(file).parent if file else Path(dir)
            plot2 = False  # plot additional results
            if plot2:
                ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)[1].ravel()
        
            fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
            # for f in [save_dir / f'study_coco_{x}.txt' for x in ['yolov5n6', 'yolov5s6', 'yolov5m6', 'yolov5l6', 'yolov5x6']]:
            for f in sorted(save_dir.glob("study*.txt")):
                y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
                x = np.arange(y.shape[1]) if x is None else np.array(x)
                if plot2:
                    s = ["P", "R", "mAP@.5", "mAP@.5:.95", "t_preprocess (ms/img)", "t_inference (ms/img)", "t_NMS (ms/img)"]
                    for i in range(7):
                        ax[i].plot(x, y[i], ".-", linewidth=2, markersize=8)
                        ax[i].set_title(s[i])
        
                j = y[3].argmax() + 1
                ax2.plot(
                    y[5, 1:j],
                    y[3, 1:j] * 1e2,
                    ".-",
                    linewidth=2,
                    markersize=8,
                    label=f.stem.replace("study_coco_", "").replace("yolo", "YOLO"),
                )
        
            ax2.plot(
                1e3 / np.array([209, 140, 97, 58, 35, 18]),
                [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
                "k.-",
                linewidth=2,
                markersize=8,
                alpha=0.25,
                label="EfficientDet",
            )
        
            ax2.grid(alpha=0.2)
            ax2.set_yticks(np.arange(20, 60, 5))
            ax2.set_xlim(0, 57)
            ax2.set_ylim(25, 55)
            ax2.set_xlabel("GPU Speed (ms/img)")
            ax2.set_ylabel("COCO AP val")
            ax2.legend(loc="lower right")
            f = save_dir / "study.png"
            print(f"Saving {f}...")
            plt.savefig(f, dpi=300)
        
        
        @TryExcept()  # known issue https://github.com/ultralytics/yolov5/issues/5395
        def plot_labels(labels, names=(), save_dir=Path("")):
            """Plots dataset labels, saving correlogram and label images, handles classes, and visualizes bounding boxes."""
            LOGGER.info(f"Plotting labels to {save_dir / 'labels.jpg'}... ")
            c, b = labels[:, 0], labels[:, 1:].transpose()  # classes, boxes
            nc = int(c.max() + 1)  # number of classes
            x = pd.DataFrame(b.transpose(), columns=["x", "y", "width", "height"])
        
            # seaborn correlogram
            sn.pairplot(x, corner=True, diag_kind="auto", kind="hist", diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9))
            plt.savefig(save_dir / "labels_correlogram.jpg", dpi=200)
            plt.close()
        
            # matplotlib labels
            matplotlib.use("svg")  # faster
            ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel()
            y = ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8)
            with contextlib.suppress(Exception):  # color histogram bars by class
                [y[2].patches[i].set_color([x / 255 for x in colors(i)]) for i in range(nc)]  # known issue #3195
            ax[0].set_ylabel("instances")
            if 0 < len(names) < 30:
                ax[0].set_xticks(range(len(names)))
                ax[0].set_xticklabels(list(names.values()), rotation=90, fontsize=10)
            else:
                ax[0].set_xlabel("classes")
            sn.histplot(x, x="x", y="y", ax=ax[2], bins=50, pmax=0.9)
            sn.histplot(x, x="width", y="height", ax=ax[3], bins=50, pmax=0.9)
        
            # rectangles
            labels[:, 1:3] = 0.5  # center
            labels[:, 1:] = xywh2xyxy(labels[:, 1:]) * 2000
            img = Image.fromarray(np.ones((2000, 2000, 3), dtype=np.uint8) * 255)
            for cls, *box in labels[:1000]:
                ImageDraw.Draw(img).rectangle(box, width=1, outline=colors(cls))  # plot
            ax[1].imshow(img)
            ax[1].axis("off")
        
            for a in [0, 1, 2, 3]:
                for s in ["top", "right", "left", "bottom"]:
                    ax[a].spines[s].set_visible(False)
        
            plt.savefig(save_dir / "labels.jpg", dpi=200)
            matplotlib.use("Agg")
            plt.close()
        
        
        def imshow_cls(im, labels=None, pred=None, names=None, nmax=25, verbose=False, f=Path("images.jpg")):
            """Displays a grid of images with optional labels and predictions, saving to a file."""
            from utils.augmentations import denormalize
        
            names = names or [f"class{i}" for i in range(1000)]
            blocks = torch.chunk(
                denormalize(im.clone()).cpu().float(), len(im), dim=0
            )  # select batch index 0, block by channels
            n = min(len(blocks), nmax)  # number of plots
            m = min(8, round(n**0.5))  # 8 x 8 default
            fig, ax = plt.subplots(math.ceil(n / m), m)  # 8 rows x n/8 cols
            ax = ax.ravel() if m > 1 else [ax]
            # plt.subplots_adjust(wspace=0.05, hspace=0.05)
            for i in range(n):
                ax[i].imshow(blocks[i].squeeze().permute((1, 2, 0)).numpy().clip(0.0, 1.0))
                ax[i].axis("off")
                if labels is not None:
                    s = names[labels[i]] + (f"—{names[pred[i]]}" if pred is not None else "")
                    ax[i].set_title(s, fontsize=8, verticalalignment="top")
            plt.savefig(f, dpi=300, bbox_inches="tight")
            plt.close()
            if verbose:
                LOGGER.info(f"Saving {f}")
                if labels is not None:
                    LOGGER.info("True:     " + " ".join(f"{names[i]:3s}" for i in labels[:nmax]))
                if pred is not None:
                    LOGGER.info("Predicted:" + " ".join(f"{names[i]:3s}" for i in pred[:nmax]))
            return f
        
        
        def plot_evolve(evolve_csv="path/to/evolve.csv"):
            """
            Plots hyperparameter evolution results from a given CSV, saving the plot and displaying best results.
        
            Example: from utils.plots import *; plot_evolve()
            """
            evolve_csv = Path(evolve_csv)
            data = pd.read_csv(evolve_csv)
            keys = [x.strip() for x in data.columns]
            x = data.values
            f = fitness(x)
            j = np.argmax(f)  # max fitness index
            plt.figure(figsize=(10, 12), tight_layout=True)
            matplotlib.rc("font", **{"size": 8})
            print(f"Best results from row {j} of {evolve_csv}:")
            for i, k in enumerate(keys[7:]):
                v = x[:, 7 + i]
                mu = v[j]  # best single result
                plt.subplot(6, 5, i + 1)
                plt.scatter(v, f, c=hist2d(v, f, 20), cmap="viridis", alpha=0.8, edgecolors="none")
                plt.plot(mu, f.max(), "k+", markersize=15)
                plt.title(f"{k} = {mu:.3g}", fontdict={"size": 9})  # limit to 40 characters
                if i % 5 != 0:
                    plt.yticks([])
                print(f"{k:>15}: {mu:.3g}")
            f = evolve_csv.with_suffix(".png")  # filename
            plt.savefig(f, dpi=200)
            plt.close()
            print(f"Saved {f}")
        
        
        def plot_results(file="path/to/results.csv", dir=""):
            """
            Plots training results from a 'results.csv' file; accepts file path and directory as arguments.
        
            Example: from utils.plots import *; plot_results('path/to/results.csv')
            """
            save_dir = Path(file).parent if file else Path(dir)
            fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True)
            ax = ax.ravel()
            files = list(save_dir.glob("results*.csv"))
            assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot."
            for f in files:
                try:
                    data = pd.read_csv(f)
                    s = [x.strip() for x in data.columns]
                    x = data.values[:, 0]
                    for i, j in enumerate([1, 2, 3, 4, 5, 8, 9, 10, 6, 7]):
                        y = data.values[:, j].astype("float")
                        # y[y == 0] = np.nan  # don't show zero values
                        ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=8)  # actual results
                        ax[i].plot(x, gaussian_filter1d(y, sigma=3), ":", label="smooth", linewidth=2)  # smoothing line
                        ax[i].set_title(s[j], fontsize=12)
                        # if j in [8, 9, 10]:  # share train and val loss y axes
                        #     ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
                except Exception as e:
                    LOGGER.info(f"Warning: Plotting error for {f}: {e}")
            ax[1].legend()
            fig.savefig(save_dir / "results.png", dpi=200)
            plt.close()
        
        
        def profile_idetection(start=0, stop=0, labels=(), save_dir=""):
            """
            Plots per-image iDetection logs, comparing metrics like storage and performance over time.
        
            Example: from utils.plots import *; profile_idetection()
            """
            ax = plt.subplots(2, 4, figsize=(12, 6), tight_layout=True)[1].ravel()
            s = ["Images", "Free Storage (GB)", "RAM Usage (GB)", "Battery", "dt_raw (ms)", "dt_smooth (ms)", "real-world FPS"]
            files = list(Path(save_dir).glob("frames*.txt"))
            for fi, f in enumerate(files):
                try:
                    results = np.loadtxt(f, ndmin=2).T[:, 90:-30]  # clip first and last rows
                    n = results.shape[1]  # number of rows
                    x = np.arange(start, min(stop, n) if stop else n)
                    results = results[:, x]
                    t = results[0] - results[0].min()  # set t0=0s
                    results[0] = x
                    for i, a in enumerate(ax):
                        if i < len(results):
                            label = labels[fi] if len(labels) else f.stem.replace("frames_", "")
                            a.plot(t, results[i], marker=".", label=label, linewidth=1, markersize=5)
                            a.set_title(s[i])
                            a.set_xlabel("time (s)")
                            # if fi == len(files) - 1:
                            #     a.set_ylim(bottom=0)
                            for side in ["top", "right"]:
                                a.spines[side].set_visible(False)
                        else:
                            a.remove()
                except Exception as e:
                    print(f"Warning: Plotting error for {f}; {e}")
            ax[1].legend()
            plt.savefig(Path(save_dir) / "idetection_profile.png", dpi=200)
        
        
        def save_one_box(xyxy, im, file=Path("im.jpg"), gain=1.02, pad=10, square=False, BGR=False, save=True):
            """Crops and saves an image from bounding box `xyxy`, applied with `gain` and `pad`, optionally squares and adjusts
            for BGR.
            """
            xyxy = torch.tensor(xyxy).view(-1, 4)
            b = xyxy2xywh(xyxy)  # boxes
            if square:
                b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1)  # attempt rectangle to square
            b[:, 2:] = b[:, 2:] * gain + pad  # box wh * gain + pad
            xyxy = xywh2xyxy(b).long()
            clip_boxes(xyxy, im.shape)
            crop = im[int(xyxy[0, 1]) : int(xyxy[0, 3]), int(xyxy[0, 0]) : int(xyxy[0, 2]), :: (1 if BGR else -1)]
            if save:
                file.parent.mkdir(parents=True, exist_ok=True)  # make directory
                f = str(increment_path(file).with_suffix(".jpg"))
                # cv2.imwrite(f, crop)  # save BGR, https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue
                Image.fromarray(crop[..., ::-1]).save(f, quality=95, subsampling=0)  # save RGB
            return crop
        
      • 📄 torch_utils.py 20.9KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """PyTorch utils."""
        
        import math
        import os
        import platform
        import subprocess
        import time
        import warnings
        from contextlib import contextmanager
        from copy import deepcopy
        from pathlib import Path
        
        import torch
        import torch.distributed as dist
        import torch.nn as nn
        import torch.nn.functional as F
        from torch.nn.parallel import DistributedDataParallel as DDP
        
        from utils.general import LOGGER, check_version, colorstr, file_date, git_describe
        
        LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))  # https://pytorch.org/docs/stable/elastic/run.html
        RANK = int(os.getenv("RANK", -1))
        WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
        
        try:
            import thop  # for FLOPs computation
        except ImportError:
            thop = None
        
        # Suppress PyTorch warnings
        warnings.filterwarnings("ignore", message="User provided device_type of 'cuda', but CUDA is not available. Disabling")
        warnings.filterwarnings("ignore", category=UserWarning)
        
        
        def smart_inference_mode(torch_1_9=check_version(torch.__version__, "1.9.0")):
            """Applies torch.inference_mode() if torch>=1.9.0, else torch.no_grad() as a decorator for functions."""
        
            def decorate(fn):
                return (torch.inference_mode if torch_1_9 else torch.no_grad)()(fn)
        
            return decorate
        
        
        def smartCrossEntropyLoss(label_smoothing=0.0):
            """Returns a CrossEntropyLoss with optional label smoothing for torch>=1.10.0; warns if smoothing on lower
            versions.
            """
            if check_version(torch.__version__, "1.10.0"):
                return nn.CrossEntropyLoss(label_smoothing=label_smoothing)
            if label_smoothing > 0:
                LOGGER.warning(f"WARNING ⚠️ label smoothing {label_smoothing} requires torch>=1.10.0")
            return nn.CrossEntropyLoss()
        
        
        def smart_DDP(model):
            """Initializes DistributedDataParallel (DDP) for model training, respecting torch version constraints."""
            assert not check_version(torch.__version__, "1.12.0", pinned=True), (
                "torch==1.12.0 torchvision==0.13.0 DDP training is not supported due to a known issue. "
                "Please upgrade or downgrade torch to use DDP. See https://github.com/ultralytics/yolov5/issues/8395"
            )
            if check_version(torch.__version__, "1.11.0"):
                return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK, static_graph=True)
            else:
                return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
        
        
        def reshape_classifier_output(model, n=1000):
            """Reshapes last layer of model to match class count 'n', supporting Classify, Linear, Sequential types."""
            from models.common import Classify
        
            name, m = list((model.model if hasattr(model, "model") else model).named_children())[-1]  # last module
            if isinstance(m, Classify):  # YOLOv5 Classify() head
                if m.linear.out_features != n:
                    m.linear = nn.Linear(m.linear.in_features, n)
            elif isinstance(m, nn.Linear):  # ResNet, EfficientNet
                if m.out_features != n:
                    setattr(model, name, nn.Linear(m.in_features, n))
            elif isinstance(m, nn.Sequential):
                types = [type(x) for x in m]
                if nn.Linear in types:
                    i = types.index(nn.Linear)  # nn.Linear index
                    if m[i].out_features != n:
                        m[i] = nn.Linear(m[i].in_features, n)
                elif nn.Conv2d in types:
                    i = types.index(nn.Conv2d)  # nn.Conv2d index
                    if m[i].out_channels != n:
                        m[i] = nn.Conv2d(m[i].in_channels, n, m[i].kernel_size, m[i].stride, bias=m[i].bias is not None)
        
        
        @contextmanager
        def torch_distributed_zero_first(local_rank: int):
            """Context manager ensuring ordered operations in distributed training by making all processes wait for the leading
            process.
            """
            if local_rank not in [-1, 0]:
                dist.barrier(device_ids=[local_rank])
            yield
            if local_rank == 0:
                dist.barrier(device_ids=[0])
        
        
        def device_count():
            """Returns the number of available CUDA devices; works on Linux and Windows by invoking `nvidia-smi`."""
            assert platform.system() in ("Linux", "Windows"), "device_count() only supported on Linux or Windows"
            try:
                cmd = "nvidia-smi -L | wc -l" if platform.system() == "Linux" else 'nvidia-smi -L | find /c /v ""'  # Windows
                return int(subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1])
            except Exception:
                return 0
        
        
        def select_device(device="", batch_size=0, newline=True):
            """Selects computing device (CPU, CUDA GPU, MPS) for YOLOv5 model deployment, logging device info."""
            s = f"YOLOv5 🚀 {git_describe() or file_date()} Python-{platform.python_version()} torch-{torch.__version__} "
            device = str(device).strip().lower().replace("cuda:", "").replace("none", "")  # to string, 'cuda:0' to '0'
            cpu = device == "cpu"
            mps = device == "mps"  # Apple Metal Performance Shaders (MPS)
            if cpu or mps:
                os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # force torch.cuda.is_available() = False
            elif device:  # non-cpu device requested
                os.environ["CUDA_VISIBLE_DEVICES"] = device  # set environment variable - must be before assert is_available()
                assert torch.cuda.is_available() and torch.cuda.device_count() >= len(
                    device.replace(",", "")
                ), f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"
        
            if not cpu and not mps and torch.cuda.is_available():  # prefer GPU if available
                devices = device.split(",") if device else "0"  # range(torch.cuda.device_count())  # i.e. 0,1,6,7
                n = len(devices)  # device count
                if n > 1 and batch_size > 0:  # check batch_size is divisible by device_count
                    assert batch_size % n == 0, f"batch-size {batch_size} not multiple of GPU count {n}"
                space = " " * (len(s) + 1)
                for i, d in enumerate(devices):
                    p = torch.cuda.get_device_properties(i)
                    s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n"  # bytes to MB
                arg = "cuda:0"
            elif mps and getattr(torch, "has_mps", False) and torch.backends.mps.is_available():  # prefer MPS if available
                s += "MPS\n"
                arg = "mps"
            else:  # revert to CPU
                s += "CPU\n"
                arg = "cpu"
        
            if not newline:
                s = s.rstrip()
            LOGGER.info(s)
            return torch.device(arg)
        
        
        def time_sync():
            """Synchronizes PyTorch for accurate timing, leveraging CUDA if available, and returns the current time."""
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            return time.time()
        
        
        def profile(input, ops, n=10, device=None):
            """YOLOv5 speed/memory/FLOPs profiler
            Usage:
                input = torch.randn(16, 3, 640, 640)
                m1 = lambda x: x * torch.sigmoid(x)
                m2 = nn.SiLU()
                profile(input, [m1, m2], n=100)  # profile over 100 iterations
            """
            results = []
            if not isinstance(device, torch.device):
                device = select_device(device)
            print(
                f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"
                f"{'input':>24s}{'output':>24s}"
            )
        
            for x in input if isinstance(input, list) else [input]:
                x = x.to(device)
                x.requires_grad = True
                for m in ops if isinstance(ops, list) else [ops]:
                    m = m.to(device) if hasattr(m, "to") else m  # device
                    m = m.half() if hasattr(m, "half") and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m
                    tf, tb, t = 0, 0, [0, 0, 0]  # dt forward, backward
                    try:
                        flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1e9 * 2  # GFLOPs
                    except Exception:
                        flops = 0
        
                    try:
                        for _ in range(n):
                            t[0] = time_sync()
                            y = m(x)
                            t[1] = time_sync()
                            try:
                                _ = (sum(yi.sum() for yi in y) if isinstance(y, list) else y).sum().backward()
                                t[2] = time_sync()
                            except Exception:  # no backward method
                                # print(e)  # for debug
                                t[2] = float("nan")
                            tf += (t[1] - t[0]) * 1000 / n  # ms per op forward
                            tb += (t[2] - t[1]) * 1000 / n  # ms per op backward
                        mem = torch.cuda.memory_reserved() / 1e9 if torch.cuda.is_available() else 0  # (GB)
                        s_in, s_out = (tuple(x.shape) if isinstance(x, torch.Tensor) else "list" for x in (x, y))  # shapes
                        p = sum(x.numel() for x in m.parameters()) if isinstance(m, nn.Module) else 0  # parameters
                        print(f"{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}")
                        results.append([p, flops, mem, tf, tb, s_in, s_out])
                    except Exception as e:
                        print(e)
                        results.append(None)
                    torch.cuda.empty_cache()
            return results
        
        
        def is_parallel(model):
            """Checks if the model is using Data Parallelism (DP) or Distributed Data Parallelism (DDP)."""
            return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
        
        
        def de_parallel(model):
            """Returns a single-GPU model by removing Data Parallelism (DP) or Distributed Data Parallelism (DDP) if applied."""
            return model.module if is_parallel(model) else model
        
        
        def initialize_weights(model):
            """Initializes weights of Conv2d, BatchNorm2d, and activations (Hardswish, LeakyReLU, ReLU, ReLU6, SiLU) in the
            model.
            """
            for m in model.modules():
                t = type(m)
                if t is nn.Conv2d:
                    pass  # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                elif t is nn.BatchNorm2d:
                    m.eps = 1e-3
                    m.momentum = 0.03
                elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
                    m.inplace = True
        
        
        def find_modules(model, mclass=nn.Conv2d):
            """Finds and returns list of layer indices in `model.module_list` matching the specified `mclass`."""
            return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]
        
        
        def sparsity(model):
            """Calculates and returns the global sparsity of a model as the ratio of zero-valued parameters to total
            parameters.
            """
            a, b = 0, 0
            for p in model.parameters():
                a += p.numel()
                b += (p == 0).sum()
            return b / a
        
        
        def prune(model, amount=0.3):
            """Prunes Conv2d layers in a model to a specified sparsity using L1 unstructured pruning."""
            import torch.nn.utils.prune as prune
        
            for name, m in model.named_modules():
                if isinstance(m, nn.Conv2d):
                    prune.l1_unstructured(m, name="weight", amount=amount)  # prune
                    prune.remove(m, "weight")  # make permanent
            LOGGER.info(f"Model pruned to {sparsity(model):.3g} global sparsity")
        
        
        def fuse_conv_and_bn(conv, bn):
            """
            Fuses Conv2d and BatchNorm2d layers into a single Conv2d layer.
        
            See https://tehnokv.com/posts/fusing-batchnorm-and-conv/.
            """
            fusedconv = (
                nn.Conv2d(
                    conv.in_channels,
                    conv.out_channels,
                    kernel_size=conv.kernel_size,
                    stride=conv.stride,
                    padding=conv.padding,
                    dilation=conv.dilation,
                    groups=conv.groups,
                    bias=True,
                )
                .requires_grad_(False)
                .to(conv.weight.device)
            )
        
            # Prepare filters
            w_conv = conv.weight.clone().view(conv.out_channels, -1)
            w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
            fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))
        
            # Prepare spatial bias
            b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
            b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
            fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
        
            return fusedconv
        
        
        def model_info(model, verbose=False, imgsz=640):
            """
            Prints model summary including layers, parameters, gradients, and FLOPs; imgsz may be int or list.
        
            Example: img_size=640 or img_size=[640, 320]
            """
            n_p = sum(x.numel() for x in model.parameters())  # number parameters
            n_g = sum(x.numel() for x in model.parameters() if x.requires_grad)  # number gradients
            if verbose:
                print(f"{'layer':>5} {'name':>40} {'gradient':>9} {'parameters':>12} {'shape':>20} {'mu':>10} {'sigma':>10}")
                for i, (name, p) in enumerate(model.named_parameters()):
                    name = name.replace("module_list.", "")
                    print(
                        "%5g %40s %9s %12g %20s %10.3g %10.3g"
                        % (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())
                    )
        
            try:  # FLOPs
                p = next(model.parameters())
                stride = max(int(model.stride.max()), 32) if hasattr(model, "stride") else 32  # max stride
                im = torch.empty((1, p.shape[1], stride, stride), device=p.device)  # input image in BCHW format
                flops = thop.profile(deepcopy(model), inputs=(im,), verbose=False)[0] / 1e9 * 2  # stride GFLOPs
                imgsz = imgsz if isinstance(imgsz, list) else [imgsz, imgsz]  # expand if int/float
                fs = f", {flops * imgsz[0] / stride * imgsz[1] / stride:.1f} GFLOPs"  # 640x640 GFLOPs
            except Exception:
                fs = ""
        
            name = Path(model.yaml_file).stem.replace("yolov5", "YOLOv5") if hasattr(model, "yaml_file") else "Model"
            LOGGER.info(f"{name} summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
        
        
        def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)
            """Scales an image tensor `img` of shape (bs,3,y,x) by `ratio`, optionally maintaining the original shape, padded to
            multiples of `gs`.
            """
            if ratio == 1.0:
                return img
            h, w = img.shape[2:]
            s = (int(h * ratio), int(w * ratio))  # new size
            img = F.interpolate(img, size=s, mode="bilinear", align_corners=False)  # resize
            if not same_shape:  # pad/crop img
                h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))
            return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean
        
        
        def copy_attr(a, b, include=(), exclude=()):
            """Copies attributes from object b to a, optionally filtering with include and exclude lists."""
            for k, v in b.__dict__.items():
                if (len(include) and k not in include) or k.startswith("_") or k in exclude:
                    continue
                else:
                    setattr(a, k, v)
        
        
        def smart_optimizer(model, name="Adam", lr=0.001, momentum=0.9, decay=1e-5):
            """
            Initializes YOLOv5 smart optimizer with 3 parameter groups for different decay configurations.
        
            Groups are 0) weights with decay, 1) weights no decay, 2) biases no decay.
            """
            g = [], [], []  # optimizer parameter groups
            bn = tuple(v for k, v in nn.__dict__.items() if "Norm" in k)  # normalization layers, i.e. BatchNorm2d()
            for v in model.modules():
                for p_name, p in v.named_parameters(recurse=0):
                    if p_name == "bias":  # bias (no decay)
                        g[2].append(p)
                    elif p_name == "weight" and isinstance(v, bn):  # weight (no decay)
                        g[1].append(p)
                    else:
                        g[0].append(p)  # weight (with decay)
        
            if name == "Adam":
                optimizer = torch.optim.Adam(g[2], lr=lr, betas=(momentum, 0.999))  # adjust beta1 to momentum
            elif name == "AdamW":
                optimizer = torch.optim.AdamW(g[2], lr=lr, betas=(momentum, 0.999), weight_decay=0.0)
            elif name == "RMSProp":
                optimizer = torch.optim.RMSprop(g[2], lr=lr, momentum=momentum)
            elif name == "SGD":
                optimizer = torch.optim.SGD(g[2], lr=lr, momentum=momentum, nesterov=True)
            else:
                raise NotImplementedError(f"Optimizer {name} not implemented.")
        
            optimizer.add_param_group({"params": g[0], "weight_decay": decay})  # add g0 with weight_decay
            optimizer.add_param_group({"params": g[1], "weight_decay": 0.0})  # add g1 (BatchNorm2d weights)
            LOGGER.info(
                f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}) with parameter groups "
                f'{len(g[1])} weight(decay=0.0), {len(g[0])} weight(decay={decay}), {len(g[2])} bias'
            )
            return optimizer
        
        
        def smart_hub_load(repo="ultralytics/yolov5", model="yolov5s", **kwargs):
            """YOLOv5 torch.hub.load() wrapper with smart error handling, adjusting torch arguments for compatibility."""
            if check_version(torch.__version__, "1.9.1"):
                kwargs["skip_validation"] = True  # validation causes GitHub API rate limit errors
            if check_version(torch.__version__, "1.12.0"):
                kwargs["trust_repo"] = True  # argument required starting in torch 0.12
            try:
                return torch.hub.load(repo, model, **kwargs)
            except Exception:
                return torch.hub.load(repo, model, force_reload=True, **kwargs)
        
        
        def smart_resume(ckpt, optimizer, ema=None, weights="yolov5s.pt", epochs=300, resume=True):
            """Resumes training from a checkpoint, updating optimizer, ema, and epochs, with optional resume verification."""
            best_fitness = 0.0
            start_epoch = ckpt["epoch"] + 1
            if ckpt["optimizer"] is not None:
                optimizer.load_state_dict(ckpt["optimizer"])  # optimizer
                best_fitness = ckpt["best_fitness"]
            if ema and ckpt.get("ema"):
                ema.ema.load_state_dict(ckpt["ema"].float().state_dict())  # EMA
                ema.updates = ckpt["updates"]
            if resume:
                assert start_epoch > 0, (
                    f"{weights} training to {epochs} epochs is finished, nothing to resume.\n"
                    f"Start a new training without --resume, i.e. 'python train.py --weights {weights}'"
                )
                LOGGER.info(f"Resuming training from {weights} from epoch {start_epoch} to {epochs} total epochs")
            if epochs < start_epoch:
                LOGGER.info(f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs.")
                epochs += ckpt["epoch"]  # finetune additional epochs
            return best_fitness, start_epoch, epochs
        
        
        class EarlyStopping:
            # YOLOv5 simple early stopper
            def __init__(self, patience=30):
                """Initializes simple early stopping mechanism for YOLOv5, with adjustable patience for non-improving epochs."""
                self.best_fitness = 0.0  # i.e. mAP
                self.best_epoch = 0
                self.patience = patience or float("inf")  # epochs to wait after fitness stops improving to stop
                self.possible_stop = False  # possible stop may occur next epoch
        
            def __call__(self, epoch, fitness):
                """Evaluates if training should stop based on fitness improvement and patience, returning a boolean."""
                if fitness >= self.best_fitness:  # >= 0 to allow for early zero-fitness stage of training
                    self.best_epoch = epoch
                    self.best_fitness = fitness
                delta = epoch - self.best_epoch  # epochs without improvement
                self.possible_stop = delta >= (self.patience - 1)  # possible stop may occur next epoch
                stop = delta >= self.patience  # stop training if patience exceeded
                if stop:
                    LOGGER.info(
                        f"Stopping training early as no improvement observed in last {self.patience} epochs. "
                        f"Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n"
                        f"To update EarlyStopping(patience={self.patience}) pass a new patience value, "
                        f"i.e. `python train.py --patience 300` or use `--patience 0` to disable EarlyStopping."
                    )
                return stop
        
        
        class ModelEMA:
            """Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-models
            Keeps a moving average of everything in the model state_dict (parameters and buffers)
            For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
            """
        
            def __init__(self, model, decay=0.9999, tau=2000, updates=0):
                """Initializes EMA with model parameters, decay rate, tau for decay adjustment, and update count; sets model to
                evaluation mode.
                """
                self.ema = deepcopy(de_parallel(model)).eval()  # FP32 EMA
                self.updates = updates  # number of EMA updates
                self.decay = lambda x: decay * (1 - math.exp(-x / tau))  # decay exponential ramp (to help early epochs)
                for p in self.ema.parameters():
                    p.requires_grad_(False)
        
            def update(self, model):
                """Updates the Exponential Moving Average (EMA) parameters based on the current model's parameters."""
                self.updates += 1
                d = self.decay(self.updates)
        
                msd = de_parallel(model).state_dict()  # model state_dict
                for k, v in self.ema.state_dict().items():
                    if v.dtype.is_floating_point:  # true for FP16 and FP32
                        v *= d
                        v += (1 - d) * msd[k].detach()
                # assert v.dtype == msd[k].dtype == torch.float32, f'{k}: EMA {v.dtype} and model {msd[k].dtype} must be FP32'
        
            def update_attr(self, model, include=(), exclude=("process_group", "reducer")):
                """Updates EMA attributes by copying specified attributes from model to EMA, excluding certain attributes by
                default.
                """
                copy_attr(self.ema, model, include, exclude)
        
      • 📄 triton.py 3.7KB
        # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
        """Utils to interact with the Triton Inference Server."""
        
        import typing
        from urllib.parse import urlparse
        
        import torch
        
        
        class TritonRemoteModel:
            """
            A wrapper over a model served by the Triton Inference Server.
        
            It can be configured to communicate over GRPC or HTTP. It accepts Torch Tensors as input and returns them as
            outputs.
            """
        
            def __init__(self, url: str):
                """
                Keyword arguments:
                url: Fully qualified address of the Triton server - for e.g. grpc://localhost:8000
                """
        
                parsed_url = urlparse(url)
                if parsed_url.scheme == "grpc":
                    from tritonclient.grpc import InferenceServerClient, InferInput
        
                    self.client = InferenceServerClient(parsed_url.netloc)  # Triton GRPC client
                    model_repository = self.client.get_model_repository_index()
                    self.model_name = model_repository.models[0].name
                    self.metadata = self.client.get_model_metadata(self.model_name, as_json=True)
        
                    def create_input_placeholders() -> typing.List[InferInput]:
                        return [
                            InferInput(i["name"], [int(s) for s in i["shape"]], i["datatype"]) for i in self.metadata["inputs"]
                        ]
        
                else:
                    from tritonclient.http import InferenceServerClient, InferInput
        
                    self.client = InferenceServerClient(parsed_url.netloc)  # Triton HTTP client
                    model_repository = self.client.get_model_repository_index()
                    self.model_name = model_repository[0]["name"]
                    self.metadata = self.client.get_model_metadata(self.model_name)
        
                    def create_input_placeholders() -> typing.List[InferInput]:
                        return [
                            InferInput(i["name"], [int(s) for s in i["shape"]], i["datatype"]) for i in self.metadata["inputs"]
                        ]
        
                self._create_input_placeholders_fn = create_input_placeholders
        
            @property
            def runtime(self):
                """Returns the model runtime."""
                return self.metadata.get("backend", self.metadata.get("platform"))
        
            def __call__(self, *args, **kwargs) -> typing.Union[torch.Tensor, typing.Tuple[torch.Tensor, ...]]:
                """
                Invokes the model.
        
                Parameters can be provided via args or kwargs. args, if provided, are assumed to match the order of inputs of
                the model. kwargs are matched with the model input names.
                """
                inputs = self._create_inputs(*args, **kwargs)
                response = self.client.infer(model_name=self.model_name, inputs=inputs)
                result = []
                for output in self.metadata["outputs"]:
                    tensor = torch.as_tensor(response.as_numpy(output["name"]))
                    result.append(tensor)
                return result[0] if len(result) == 1 else result
        
            def _create_inputs(self, *args, **kwargs):
                """Creates input tensors from args or kwargs, not both; raises error if none or both are provided."""
                args_len, kwargs_len = len(args), len(kwargs)
                if not args_len and not kwargs_len:
                    raise RuntimeError("No inputs provided.")
                if args_len and kwargs_len:
                    raise RuntimeError("Cannot specify args and kwargs at the same time")
        
                placeholders = self._create_input_placeholders_fn()
                if args_len:
                    if args_len != len(placeholders):
                        raise RuntimeError(f"Expected {len(placeholders)} inputs, got {args_len}.")
                    for input, value in zip(placeholders, args):
                        input.set_data_from_numpy(value.cpu().numpy())
                else:
                    for input in placeholders:
                        value = kwargs[input.name]
                        input.set_data_from_numpy(value.cpu().numpy())
                return placeholders
        
    • 📄 .dockerignore 3.6KB
      # Repo-specific DockerIgnore -------------------------------------------------------------------------------------------
      .git
      .cache
      .idea
      runs
      output
      coco
      storage.googleapis.com
      
      data/samples/*
      **/results*.csv
      *.jpg
      
      # Neural Network weights -----------------------------------------------------------------------------------------------
      **/*.pt
      **/*.pth
      **/*.onnx
      **/*.engine
      **/*.mlmodel
      **/*.torchscript
      **/*.torchscript.pt
      **/*.tflite
      **/*.h5
      **/*.pb
      *_saved_model/
      *_web_model/
      *_openvino_model/
      
      # Below Copied From .gitignore -----------------------------------------------------------------------------------------
      # Below Copied From .gitignore -----------------------------------------------------------------------------------------
      
      
      # GitHub Python GitIgnore ----------------------------------------------------------------------------------------------
      # Byte-compiled / optimized / DLL files
      __pycache__/
      *.py[cod]
      *$py.class
      
      # C extensions
      *.so
      
      # Distribution / packaging
      .Python
      env/
      build/
      develop-eggs/
      dist/
      downloads/
      eggs/
      .eggs/
      lib/
      lib64/
      parts/
      sdist/
      var/
      wheels/
      *.egg-info/
      wandb/
      .installed.cfg
      *.egg
      
      # PyInstaller
      #  Usually these files are written by a python script from a template
      #  before PyInstaller builds the exe, so as to inject date/other infos into it.
      *.manifest
      *.spec
      
      # Installer logs
      pip-log.txt
      pip-delete-this-directory.txt
      
      # Unit test / coverage reports
      htmlcov/
      .tox/
      .coverage
      .coverage.*
      .cache
      nosetests.xml
      coverage.xml
      *.cover
      .hypothesis/
      
      # Translations
      *.mo
      *.pot
      
      # Django stuff:
      *.log
      local_settings.py
      
      # Flask stuff:
      instance/
      .webassets-cache
      
      # Scrapy stuff:
      .scrapy
      
      # Sphinx documentation
      docs/_build/
      
      # PyBuilder
      target/
      
      # Jupyter Notebook
      .ipynb_checkpoints
      
      # pyenv
      .python-version
      
      # celery beat schedule file
      celerybeat-schedule
      
      # SageMath parsed files
      *.sage.py
      
      # dotenv
      .env
      
      # virtualenv
      .venv*
      venv*/
      ENV*/
      
      # Spyder project settings
      .spyderproject
      .spyproject
      
      # Rope project settings
      .ropeproject
      
      # mkdocs documentation
      /site
      
      # mypy
      .mypy_cache/
      
      
      # https://github.com/github/gitignore/blob/master/Global/macOS.gitignore -----------------------------------------------
      
      # General
      .DS_Store
      .AppleDouble
      .LSOverride
      
      # Icon must end with two \r
      Icon
      Icon?
      
      # Thumbnails
      ._*
      
      # Files that might appear in the root of a volume
      .DocumentRevisions-V100
      .fseventsd
      .Spotlight-V100
      .TemporaryItems
      .Trashes
      .VolumeIcon.icns
      .com.apple.timemachine.donotpresent
      
      # Directories potentially created on remote AFP share
      .AppleDB
      .AppleDesktop
      Network Trash Folder
      Temporary Items
      .apdisk
      
      
      # https://github.com/github/gitignore/blob/master/Global/JetBrains.gitignore
      # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm
      # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
      
      # User-specific stuff:
      .idea/*
      .idea/**/workspace.xml
      .idea/**/tasks.xml
      .idea/dictionaries
      .html  # Bokeh Plots
      .pg  # TensorFlow Frozen Graphs
      .avi # videos
      
      # Sensitive or high-churn files:
      .idea/**/dataSources/
      .idea/**/dataSources.ids
      .idea/**/dataSources.local.xml
      .idea/**/sqlDataSources.xml
      .idea/**/dynamic.xml
      .idea/**/uiDesigner.xml
      
      # Gradle:
      .idea/**/gradle.xml
      .idea/**/libraries
      
      # CMake
      cmake-build-debug/
      cmake-build-release/
      
      # Mongo Explorer plugin:
      .idea/**/mongoSettings.xml
      
      ## File-based project format:
      *.iws
      
      ## Plugin-specific files:
      
      # IntelliJ
      out/
      
      # mpeltonen/sbt-idea plugin
      .idea_modules/
      
      # JIRA plugin
      atlassian-ide-plugin.xml
      
      # Cursive Clojure plugin
      .idea/replstate.xml
      
      # Crashlytics plugin (for Android Studio and IntelliJ)
      com_crashlytics_export_strings.xml
      crashlytics.properties
      crashlytics-build.properties
      fabric.properties
      
    • ⚙️ .gitattributes 75B
    • 📄 .gitignore 3.9KB
      # Repo-specific GitIgnore ----------------------------------------------------------------------------------------------
      *.jpg
      *.jpeg
      *.png
      *.bmp
      *.tif
      *.tiff
      *.heic
      *.JPG
      *.JPEG
      *.PNG
      *.BMP
      *.TIF
      *.TIFF
      *.HEIC
      *.mp4
      *.mov
      *.MOV
      *.avi
      *.data
      *.json
      *.cfg
      !setup.cfg
      !cfg/yolov3*.cfg
      
      storage.googleapis.com
      runs/*
      data/*
      data/images/*
      !data/*.yaml
      !data/hyps
      !data/scripts
      !data/images
      !data/images/zidane.jpg
      !data/images/bus.jpg
      !data/*.sh
      
      results*.csv
      
      # Datasets -------------------------------------------------------------------------------------------------------------
      coco/
      coco128/
      VOC/
      
      # MATLAB GitIgnore -----------------------------------------------------------------------------------------------------
      *.m~
      *.mat
      !targets*.mat
      
      # Neural Network weights -----------------------------------------------------------------------------------------------
      *.weights
      *.pt
      *.pb
      *.onnx
      *.engine
      *.mlmodel
      *.torchscript
      *.tflite
      *.h5
      *_saved_model/
      *_web_model/
      *_openvino_model/
      *_paddle_model/
      darknet53.conv.74
      yolov3-tiny.conv.15
      
      # GitHub Python GitIgnore ----------------------------------------------------------------------------------------------
      # Byte-compiled / optimized / DLL files
      __pycache__/
      *.py[cod]
      *$py.class
      
      # C extensions
      *.so
      
      # Distribution / packaging
      .Python
      env/
      build/
      develop-eggs/
      dist/
      downloads/
      eggs/
      .eggs/
      lib/
      lib64/
      parts/
      sdist/
      var/
      wheels/
      *.egg-info/
      /wandb/
      .installed.cfg
      *.egg
      
      
      # PyInstaller
      #  Usually these files are written by a python script from a template
      #  before PyInstaller builds the exe, so as to inject date/other infos into it.
      *.manifest
      *.spec
      
      # Installer logs
      pip-log.txt
      pip-delete-this-directory.txt
      
      # Unit test / coverage reports
      htmlcov/
      .tox/
      .coverage
      .coverage.*
      .cache
      nosetests.xml
      coverage.xml
      *.cover
      .hypothesis/
      
      # Translations
      *.mo
      *.pot
      
      # Django stuff:
      *.log
      local_settings.py
      
      # Flask stuff:
      instance/
      .webassets-cache
      
      # Scrapy stuff:
      .scrapy
      
      # Sphinx documentation
      docs/_build/
      
      # PyBuilder
      target/
      
      # Jupyter Notebook
      .ipynb_checkpoints
      
      # pyenv
      .python-version
      
      # celery beat schedule file
      celerybeat-schedule
      
      # SageMath parsed files
      *.sage.py
      
      # dotenv
      .env
      
      # virtualenv
      .venv*
      venv*/
      ENV*/
      
      # Spyder project settings
      .spyderproject
      .spyproject
      
      # Rope project settings
      .ropeproject
      
      # mkdocs documentation
      /site
      
      # mypy
      .mypy_cache/
      
      
      # https://github.com/github/gitignore/blob/master/Global/macOS.gitignore -----------------------------------------------
      
      # General
      .DS_Store
      .AppleDouble
      .LSOverride
      
      # Icon must end with two \r
      Icon
      Icon?
      
      # Thumbnails
      ._*
      
      # Files that might appear in the root of a volume
      .DocumentRevisions-V100
      .fseventsd
      .Spotlight-V100
      .TemporaryItems
      .Trashes
      .VolumeIcon.icns
      .com.apple.timemachine.donotpresent
      
      # Directories potentially created on remote AFP share
      .AppleDB
      .AppleDesktop
      Network Trash Folder
      Temporary Items
      .apdisk
      
      
      # https://github.com/github/gitignore/blob/master/Global/JetBrains.gitignore
      # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm
      # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
      
      # User-specific stuff:
      .idea/*
      .idea/**/workspace.xml
      .idea/**/tasks.xml
      .idea/dictionaries
      .html  # Bokeh Plots
      .pg  # TensorFlow Frozen Graphs
      .avi # videos
      
      # Sensitive or high-churn files:
      .idea/**/dataSources/
      .idea/**/dataSources.ids
      .idea/**/dataSources.local.xml
      .idea/**/sqlDataSources.xml
      .idea/**/dynamic.xml
      .idea/**/uiDesigner.xml
      
      # Gradle:
      .idea/**/gradle.xml
      .idea/**/libraries
      
      # CMake
      cmake-build-debug/
      cmake-build-release/
      
      # Mongo Explorer plugin:
      .idea/**/mongoSettings.xml
      
      ## File-based project format:
      *.iws
      
      ## Plugin-specific files:
      
      # IntelliJ
      out/
      
      # mpeltonen/sbt-idea plugin
      .idea_modules/
      
      # JIRA plugin
      atlassian-ide-plugin.xml
      
      # Cursive Clojure plugin
      .idea/replstate.xml
      
      # Crashlytics plugin (for Android Studio and IntelliJ)
      com_crashlytics_export_strings.xml
      crashlytics.properties
      crashlytics-build.properties
      fabric.properties
      
    • ⚙️ CITATION.cff 393B
    • 📄 CONTRIBUTING.md 4.9KB
      ## Contributing to YOLOv5 🚀
      
      We love your input! We want to make contributing to YOLOv5 as easy and transparent as possible, whether it's:
      
      - Reporting a bug
      - Discussing the current state of the code
      - Submitting a fix
      - Proposing a new feature
      - Becoming a maintainer
      
      YOLOv5 works so well due to our combined community effort, and for every small improvement you contribute you will be helping push the frontiers of what's possible in AI 😃!
      
      ## Submitting a Pull Request (PR) 🛠️
      
      Submitting a PR is easy! This example shows how to submit a PR for updating `requirements.txt` in 4 steps:
      
      ### 1. Select File to Update
      
      Select `requirements.txt` to update by clicking on it in GitHub.
      
      <p align="center"><img width="800" alt="PR_step1" src="https://user-images.githubusercontent.com/26833433/122260847-08be2600-ced4-11eb-828b-8287ace4136c.png"></p>
      
      ### 2. Click 'Edit this file'
      
      The button is in the top-right corner.
      
      <p align="center"><img width="800" alt="PR_step2" src="https://user-images.githubusercontent.com/26833433/122260844-06f46280-ced4-11eb-9eec-b8a24be519ca.png"></p>
      
      ### 3. Make Changes
      
      Change the `matplotlib` version from `3.2.2` to `3.3`.
      
      <p align="center"><img width="800" alt="PR_step3" src="https://user-images.githubusercontent.com/26833433/122260853-0a87e980-ced4-11eb-9fd2-3650fb6e0842.png"></p>
      
      ### 4. Preview Changes and Submit PR
      
      Click on the **Preview changes** tab to verify your updates. At the bottom of the screen select 'Create a **new branch** for this commit', assign your branch a descriptive name such as `fix/matplotlib_version` and click the green **Propose changes** button. All done, your PR is now submitted to YOLOv5 for review and approval 😃!
      
      <p align="center"><img width="800" alt="PR_step4" src="https://user-images.githubusercontent.com/26833433/122260856-0b208000-ced4-11eb-8e8e-77b6151cbcc3.png"></p>
      
      ### PR recommendations
      
      To allow your work to be integrated as seamlessly as possible, we advise you to:
      
      - ✅ Verify your PR is **up-to-date** with `ultralytics/yolov5` `master` branch. If your PR is behind you can update your code by clicking the 'Update branch' button or by running `git pull` and `git merge master` locally.
      
      <p align="center"><img width="751" alt="Screenshot 2022-08-29 at 22 47 15" src="https://user-images.githubusercontent.com/26833433/187295893-50ed9f44-b2c9-4138-a614-de69bd1753d7.png"></p>
      
      - ✅ Verify all YOLOv5 Continuous Integration (CI) **checks are passing**.
      
      <p align="center"><img width="751" alt="Screenshot 2022-08-29 at 22 47 03" src="https://user-images.githubusercontent.com/26833433/187296922-545c5498-f64a-4d8c-8300-5fa764360da6.png"></p>
      
      - ✅ Reduce changes to the absolute **minimum** required for your bug fix or feature addition. _"It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is."_ — Bruce Lee
      
      ## Submitting a Bug Report 🐛
      
      If you spot a problem with YOLOv5 please submit a Bug Report!
      
      For us to start investigating a possible problem we need to be able to reproduce it ourselves first. We've created a few short guidelines below to help users provide what we need to get started.
      
      When asking a question, people will be better able to provide help if you provide **code** that they can easily understand and use to **reproduce** the problem. This is referred to by community members as creating a [minimum reproducible example](https://docs.ultralytics.com/help/minimum_reproducible_example/). Your code that reproduces the problem should be:
      
      - ✅ **Minimal** – Use as little code as possible that still produces the same problem
      - ✅ **Complete** – Provide **all** parts someone else needs to reproduce your problem in the question itself
      - ✅ **Reproducible** – Test the code you're about to provide to make sure it reproduces the problem
      
      In addition to the above requirements, for [Ultralytics](https://ultralytics.com/) to provide assistance your code should be:
      
      - ✅ **Current** – Verify that your code is up-to-date with the current GitHub [master](https://github.com/ultralytics/yolov5/tree/master), and if necessary `git pull` or `git clone` a new copy to ensure your problem has not already been resolved by previous commits.
      - ✅ **Unmodified** – Your problem must be reproducible without any modifications to the codebase in this repository. [Ultralytics](https://ultralytics.com/) does not provide support for custom code ⚠️.
      
      If you believe your problem meets all of the above criteria, please close this issue and raise a new one using the 🐛 **Bug Report** [template](https://github.com/ultralytics/yolov5/issues/new/choose) and provide a [minimum reproducible example](https://docs.ultralytics.com/help/minimum_reproducible_example/) to help us better understand and diagnose your problem.
      
      ## License
      
      By contributing, you agree that your contributions will be licensed under the [AGPL-3.0 license](https://choosealicense.com/licenses/agpl-3.0/)
      
    • 📄 LICENSE 33.7KB
                          GNU AFFERO GENERAL PUBLIC LICENSE
                             Version 3, 19 November 2007
      
       Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
       Everyone is permitted to copy and distribute verbatim copies
       of this license document, but changing it is not allowed.
      
                                  Preamble
      
        The GNU Affero General Public License is a free, copyleft license for
      software and other kinds of works, specifically designed to ensure
      cooperation with the community in the case of network server software.
      
        The licenses for most software and other practical works are designed
      to take away your freedom to share and change the works.  By contrast,
      our General Public Licenses are intended to guarantee your freedom to
      share and change all versions of a program--to make sure it remains free
      software for all its users.
      
        When we speak of free software, we are referring to freedom, not
      price.  Our General Public Licenses are designed to make sure that you
      have the freedom to distribute copies of free software (and charge for
      them if you wish), that you receive source code or can get it if you
      want it, that you can change the software or use pieces of it in new
      free programs, and that you know you can do these things.
      
        Developers that use our General Public Licenses protect your rights
      with two steps: (1) assert copyright on the software, and (2) offer
      you this License which gives you legal permission to copy, distribute
      and/or modify the software.
      
        A secondary benefit of defending all users' freedom is that
      improvements made in alternate versions of the program, if they
      receive widespread use, become available for other developers to
      incorporate.  Many developers of free software are heartened and
      encouraged by the resulting cooperation.  However, in the case of
      software used on network servers, this result may fail to come about.
      The GNU General Public License permits making a modified version and
      letting the public access it on a server without ever releasing its
      source code to the public.
      
        The GNU Affero General Public License is designed specifically to
      ensure that, in such cases, the modified source code becomes available
      to the community.  It requires the operator of a network server to
      provide the source code of the modified version running there to the
      users of that server.  Therefore, public use of a modified version, on
      a publicly accessible server, gives the public access to the source
      code of the modified version.
      
        An older license, called the Affero General Public License and
      published by Affero, was designed to accomplish similar goals.  This is
      a different license, not a version of the Affero GPL, but Affero has
      released a new version of the Affero GPL which permits relicensing under
      this license.
      
        The precise terms and conditions for copying, distribution and
      modification follow.
      
                             TERMS AND CONDITIONS
      
        0. Definitions.
      
        "This License" refers to version 3 of the GNU Affero General Public License.
      
        "Copyright" also means copyright-like laws that apply to other kinds of
      works, such as semiconductor masks.
      
        "The Program" refers to any copyrightable work licensed under this
      License.  Each licensee is addressed as "you".  "Licensees" and
      "recipients" may be individuals or organizations.
      
        To "modify" a work means to copy from or adapt all or part of the work
      in a fashion requiring copyright permission, other than the making of an
      exact copy.  The resulting work is called a "modified version" of the
      earlier work or a work "based on" the earlier work.
      
        A "covered work" means either the unmodified Program or a work based
      on the Program.
      
        To "propagate" a work means to do anything with it that, without
      permission, would make you directly or secondarily liable for
      infringement under applicable copyright law, except executing it on a
      computer or modifying a private copy.  Propagation includes copying,
      distribution (with or without modification), making available to the
      public, and in some countries other activities as well.
      
        To "convey" a work means any kind of propagation that enables other
      parties to make or receive copies.  Mere interaction with a user through
      a computer network, with no transfer of a copy, is not conveying.
      
        An interactive user interface displays "Appropriate Legal Notices"
      to the extent that it includes a convenient and prominently visible
      feature that (1) displays an appropriate copyright notice, and (2)
      tells the user that there is no warranty for the work (except to the
      extent that warranties are provided), that licensees may convey the
      work under this License, and how to view a copy of this License.  If
      the interface presents a list of user commands or options, such as a
      menu, a prominent item in the list meets this criterion.
      
        1. Source Code.
      
        The "source code" for a work means the preferred form of the work
      for making modifications to it.  "Object code" means any non-source
      form of a work.
      
        A "Standard Interface" means an interface that either is an official
      standard defined by a recognized standards body, or, in the case of
      interfaces specified for a particular programming language, one that
      is widely used among developers working in that language.
      
        The "System Libraries" of an executable work include anything, other
      than the work as a whole, that (a) is included in the normal form of
      packaging a Major Component, but which is not part of that Major
      Component, and (b) serves only to enable use of the work with that
      Major Component, or to implement a Standard Interface for which an
      implementation is available to the public in source code form.  A
      "Major Component", in this context, means a major essential component
      (kernel, window system, and so on) of the specific operating system
      (if any) on which the executable work runs, or a compiler used to
      produce the work, or an object code interpreter used to run it.
      
        The "Corresponding Source" for a work in object code form means all
      the source code needed to generate, install, and (for an executable
      work) run the object code and to modify the work, including scripts to
      control those activities.  However, it does not include the work's
      System Libraries, or general-purpose tools or generally available free
      programs which are used unmodified in performing those activities but
      which are not part of the work.  For example, Corresponding Source
      includes interface definition files associated with source files for
      the work, and the source code for shared libraries and dynamically
      linked subprograms that the work is specifically designed to require,
      such as by intimate data communication or control flow between those
      subprograms and other parts of the work.
      
        The Corresponding Source need not include anything that users
      can regenerate automatically from other parts of the Corresponding
      Source.
      
        The Corresponding Source for a work in source code form is that
      same work.
      
        2. Basic Permissions.
      
        All rights granted under this License are granted for the term of
      copyright on the Program, and are irrevocable provided the stated
      conditions are met.  This License explicitly affirms your unlimited
      permission to run the unmodified Program.  The output from running a
      covered work is covered by this License only if the output, given its
      content, constitutes a covered work.  This License acknowledges your
      rights of fair use or other equivalent, as provided by copyright law.
      
        You may make, run and propagate covered works that you do not
      convey, without conditions so long as your license otherwise remains
      in force.  You may convey covered works to others for the sole purpose
      of having them make modifications exclusively for you, or provide you
      with facilities for running those works, provided that you comply with
      the terms of this License in conveying all material for which you do
      not control copyright.  Those thus making or running the covered works
      for you must do so exclusively on your behalf, under your direction
      and control, on terms that prohibit them from making any copies of
      your copyrighted material outside their relationship with you.
      
        Conveying under any other circumstances is permitted solely under
      the conditions stated below.  Sublicensing is not allowed; section 10
      makes it unnecessary.
      
        3. Protecting Users' Legal Rights From Anti-Circumvention Law.
      
        No covered work shall be deemed part of an effective technological
      measure under any applicable law fulfilling obligations under article
      11 of the WIPO copyright treaty adopted on 20 December 1996, or
      similar laws prohibiting or restricting circumvention of such
      measures.
      
        When you convey a covered work, you waive any legal power to forbid
      circumvention of technological measures to the extent such circumvention
      is effected by exercising rights under this License with respect to
      the covered work, and you disclaim any intention to limit operation or
      modification of the work as a means of enforcing, against the work's
      users, your or third parties' legal rights to forbid circumvention of
      technological measures.
      
        4. Conveying Verbatim Copies.
      
        You may convey verbatim copies of the Program's source code as you
      receive it, in any medium, provided that you conspicuously and
      appropriately publish on each copy an appropriate copyright notice;
      keep intact all notices stating that this License and any
      non-permissive terms added in accord with section 7 apply to the code;
      keep intact all notices of the absence of any warranty; and give all
      recipients a copy of this License along with the Program.
      
        You may charge any price or no price for each copy that you convey,
      and you may offer support or warranty protection for a fee.
      
        5. Conveying Modified Source Versions.
      
        You may convey a work based on the Program, or the modifications to
      produce it from the Program, in the form of source code under the
      terms of section 4, provided that you also meet all of these conditions:
      
          a) The work must carry prominent notices stating that you modified
          it, and giving a relevant date.
      
          b) The work must carry prominent notices stating that it is
          released under this License and any conditions added under section
          7.  This requirement modifies the requirement in section 4 to
          "keep intact all notices".
      
          c) You must license the entire work, as a whole, under this
          License to anyone who comes into possession of a copy.  This
          License will therefore apply, along with any applicable section 7
          additional terms, to the whole of the work, and all its parts,
          regardless of how they are packaged.  This License gives no
          permission to license the work in any other way, but it does not
          invalidate such permission if you have separately received it.
      
          d) If the work has interactive user interfaces, each must display
          Appropriate Legal Notices; however, if the Program has interactive
          interfaces that do not display Appropriate Legal Notices, your
          work need not make them do so.
      
        A compilation of a covered work with other separate and independent
      works, which are not by their nature extensions of the covered work,
      and which are not combined with it such as to form a larger program,
      in or on a volume of a storage or distribution medium, is called an
      "aggregate" if the compilation and its resulting copyright are not
      used to limit the access or legal rights of the compilation's users
      beyond what the individual works permit.  Inclusion of a covered work
      in an aggregate does not cause this License to apply to the other
      parts of the aggregate.
      
        6. Conveying Non-Source Forms.
      
        You may convey a covered work in object code form under the terms
      of sections 4 and 5, provided that you also convey the
      machine-readable Corresponding Source under the terms of this License,
      in one of these ways:
      
          a) Convey the object code in, or embodied in, a physical product
          (including a physical distribution medium), accompanied by the
          Corresponding Source fixed on a durable physical medium
          customarily used for software interchange.
      
          b) Convey the object code in, or embodied in, a physical product
          (including a physical distribution medium), accompanied by a
          written offer, valid for at least three years and valid for as
          long as you offer spare parts or customer support for that product
          model, to give anyone who possesses the object code either (1) a
          copy of the Corresponding Source for all the software in the
          product that is covered by this License, on a durable physical
          medium customarily used for software interchange, for a price no
          more than your reasonable cost of physically performing this
          conveying of source, or (2) access to copy the
          Corresponding Source from a network server at no charge.
      
          c) Convey individual copies of the object code with a copy of the
          written offer to provide the Corresponding Source.  This
          alternative is allowed only occasionally and noncommercially, and
          only if you received the object code with such an offer, in accord
          with subsection 6b.
      
          d) Convey the object code by offering access from a designated
          place (gratis or for a charge), and offer equivalent access to the
          Corresponding Source in the same way through the same place at no
          further charge.  You need not require recipients to copy the
          Corresponding Source along with the object code.  If the place to
          copy the object code is a network server, the Corresponding Source
          may be on a different server (operated by you or a third party)
          that supports equivalent copying facilities, provided you maintain
          clear directions next to the object code saying where to find the
          Corresponding Source.  Regardless of what server hosts the
          Corresponding Source, you remain obligated to ensure that it is
          available for as long as needed to satisfy these requirements.
      
          e) Convey the object code using peer-to-peer transmission, provided
          you inform other peers where the object code and Corresponding
          Source of the work are being offered to the general public at no
          charge under subsection 6d.
      
        A separable portion of the object code, whose source code is excluded
      from the Corresponding Source as a System Library, need not be
      included in conveying the object code work.
      
        A "User Product" is either (1) a "consumer product", which means any
      tangible personal property which is normally used for personal, family,
      or household purposes, or (2) anything designed or sold for incorporation
      into a dwelling.  In determining whether a product is a consumer product,
      doubtful cases shall be resolved in favor of coverage.  For a particular
      product received by a particular user, "normally used" refers to a
      typical or common use of that class of product, regardless of the status
      of the particular user or of the way in which the particular user
      actually uses, or expects or is expected to use, the product.  A product
      is a consumer product regardless of whether the product has substantial
      commercial, industrial or non-consumer uses, unless such uses represent
      the only significant mode of use of the product.
      
        "Installation Information" for a User Product means any methods,
      procedures, authorization keys, or other information required to install
      and execute modified versions of a covered work in that User Product from
      a modified version of its Corresponding Source.  The information must
      suffice to ensure that the continued functioning of the modified object
      code is in no case prevented or interfered with solely because
      modification has been made.
      
        If you convey an object code work under this section in, or with, or
      specifically for use in, a User Product, and the conveying occurs as
      part of a transaction in which the right of possession and use of the
      User Product is transferred to the recipient in perpetuity or for a
      fixed term (regardless of how the transaction is characterized), the
      Corresponding Source conveyed under this section must be accompanied
      by the Installation Information.  But this requirement does not apply
      if neither you nor any third party retains the ability to install
      modified object code on the User Product (for example, the work has
      been installed in ROM).
      
        The requirement to provide Installation Information does not include a
      requirement to continue to provide support service, warranty, or updates
      for a work that has been modified or installed by the recipient, or for
      the User Product in which it has been modified or installed.  Access to a
      network may be denied when the modification itself materially and
      adversely affects the operation of the network or violates the rules and
      protocols for communication across the network.
      
        Corresponding Source conveyed, and Installation Information provided,
      in accord with this section must be in a format that is publicly
      documented (and with an implementation available to the public in
      source code form), and must require no special password or key for
      unpacking, reading or copying.
      
        7. Additional Terms.
      
        "Additional permissions" are terms that supplement the terms of this
      License by making exceptions from one or more of its conditions.
      Additional permissions that are applicable to the entire Program shall
      be treated as though they were included in this License, to the extent
      that they are valid under applicable law.  If additional permissions
      apply only to part of the Program, that part may be used separately
      under those permissions, but the entire Program remains governed by
      this License without regard to the additional permissions.
      
        When you convey a copy of a covered work, you may at your option
      remove any additional permissions from that copy, or from any part of
      it.  (Additional permissions may be written to require their own
      removal in certain cases when you modify the work.)  You may place
      additional permissions on material, added by you to a covered work,
      for which you have or can give appropriate copyright permission.
      
        Notwithstanding any other provision of this License, for material you
      add to a covered work, you may (if authorized by the copyright holders of
      that material) supplement the terms of this License with terms:
      
          a) Disclaiming warranty or limiting liability differently from the
          terms of sections 15 and 16 of this License; or
      
          b) Requiring preservation of specified reasonable legal notices or
          author attributions in that material or in the Appropriate Legal
          Notices displayed by works containing it; or
      
          c) Prohibiting misrepresentation of the origin of that material, or
          requiring that modified versions of such material be marked in
          reasonable ways as different from the original version; or
      
          d) Limiting the use for publicity purposes of names of licensors or
          authors of the material; or
      
          e) Declining to grant rights under trademark law for use of some
          trade names, trademarks, or service marks; or
      
          f) Requiring indemnification of licensors and authors of that
          material by anyone who conveys the material (or modified versions of
          it) with contractual assumptions of liability to the recipient, for
          any liability that these contractual assumptions directly impose on
          those licensors and authors.
      
        All other non-permissive additional terms are considered "further
      restrictions" within the meaning of section 10.  If the Program as you
      received it, or any part of it, contains a notice stating that it is
      governed by this License along with a term that is a further
      restriction, you may remove that term.  If a license document contains
      a further restriction but permits relicensing or conveying under this
      License, you may add to a covered work material governed by the terms
      of that license document, provided that the further restriction does
      not survive such relicensing or conveying.
      
        If you add terms to a covered work in accord with this section, you
      must place, in the relevant source files, a statement of the
      additional terms that apply to those files, or a notice indicating
      where to find the applicable terms.
      
        Additional terms, permissive or non-permissive, may be stated in the
      form of a separately written license, or stated as exceptions;
      the above requirements apply either way.
      
        8. Termination.
      
        You may not propagate or modify a covered work except as expressly
      provided under this License.  Any attempt otherwise to propagate or
      modify it is void, and will automatically terminate your rights under
      this License (including any patent licenses granted under the third
      paragraph of section 11).
      
        However, if you cease all violation of this License, then your
      license from a particular copyright holder is reinstated (a)
      provisionally, unless and until the copyright holder explicitly and
      finally terminates your license, and (b) permanently, if the copyright
      holder fails to notify you of the violation by some reasonable means
      prior to 60 days after the cessation.
      
        Moreover, your license from a particular copyright holder is
      reinstated permanently if the copyright holder notifies you of the
      violation by some reasonable means, this is the first time you have
      received notice of violation of this License (for any work) from that
      copyright holder, and you cure the violation prior to 30 days after
      your receipt of the notice.
      
        Termination of your rights under this section does not terminate the
      licenses of parties who have received copies or rights from you under
      this License.  If your rights have been terminated and not permanently
      reinstated, you do not qualify to receive new licenses for the same
      material under section 10.
      
        9. Acceptance Not Required for Having Copies.
      
        You are not required to accept this License in order to receive or
      run a copy of the Program.  Ancillary propagation of a covered work
      occurring solely as a consequence of using peer-to-peer transmission
      to receive a copy likewise does not require acceptance.  However,
      nothing other than this License grants you permission to propagate or
      modify any covered work.  These actions infringe copyright if you do
      not accept this License.  Therefore, by modifying or propagating a
      covered work, you indicate your acceptance of this License to do so.
      
        10. Automatic Licensing of Downstream Recipients.
      
        Each time you convey a covered work, the recipient automatically
      receives a license from the original licensors, to run, modify and
      propagate that work, subject to this License.  You are not responsible
      for enforcing compliance by third parties with this License.
      
        An "entity transaction" is a transaction transferring control of an
      organization, or substantially all assets of one, or subdividing an
      organization, or merging organizations.  If propagation of a covered
      work results from an entity transaction, each party to that
      transaction who receives a copy of the work also receives whatever
      licenses to the work the party's predecessor in interest had or could
      give under the previous paragraph, plus a right to possession of the
      Corresponding Source of the work from the predecessor in interest, if
      the predecessor has it or can get it with reasonable efforts.
      
        You may not impose any further restrictions on the exercise of the
      rights granted or affirmed under this License.  For example, you may
      not impose a license fee, royalty, or other charge for exercise of
      rights granted under this License, and you may not initiate litigation
      (including a cross-claim or counterclaim in a lawsuit) alleging that
      any patent claim is infringed by making, using, selling, offering for
      sale, or importing the Program or any portion of it.
      
        11. Patents.
      
        A "contributor" is a copyright holder who authorizes use under this
      License of the Program or a work on which the Program is based.  The
      work thus licensed is called the contributor's "contributor version".
      
        A contributor's "essential patent claims" are all patent claims
      owned or controlled by the contributor, whether already acquired or
      hereafter acquired, that would be infringed by some manner, permitted
      by this License, of making, using, or selling its contributor version,
      but do not include claims that would be infringed only as a
      consequence of further modification of the contributor version.  For
      purposes of this definition, "control" includes the right to grant
      patent sublicenses in a manner consistent with the requirements of
      this License.
      
        Each contributor grants you a non-exclusive, worldwide, royalty-free
      patent license under the contributor's essential patent claims, to
      make, use, sell, offer for sale, import and otherwise run, modify and
      propagate the contents of its contributor version.
      
        In the following three paragraphs, a "patent license" is any express
      agreement or commitment, however denominated, not to enforce a patent
      (such as an express permission to practice a patent or covenant not to
      sue for patent infringement).  To "grant" such a patent license to a
      party means to make such an agreement or commitment not to enforce a
      patent against the party.
      
        If you convey a covered work, knowingly relying on a patent license,
      and the Corresponding Source of the work is not available for anyone
      to copy, free of charge and under the terms of this License, through a
      publicly available network server or other readily accessible means,
      then you must either (1) cause the Corresponding Source to be so
      available, or (2) arrange to deprive yourself of the benefit of the
      patent license for this particular work, or (3) arrange, in a manner
      consistent with the requirements of this License, to extend the patent
      license to downstream recipients.  "Knowingly relying" means you have
      actual knowledge that, but for the patent license, your conveying the
      covered work in a country, or your recipient's use of the covered work
      in a country, would infringe one or more identifiable patents in that
      country that you have reason to believe are valid.
      
        If, pursuant to or in connection with a single transaction or
      arrangement, you convey, or propagate by procuring conveyance of, a
      covered work, and grant a patent license to some of the parties
      receiving the covered work authorizing them to use, propagate, modify
      or convey a specific copy of the covered work, then the patent license
      you grant is automatically extended to all recipients of the covered
      work and works based on it.
      
        A patent license is "discriminatory" if it does not include within
      the scope of its coverage, prohibits the exercise of, or is
      conditioned on the non-exercise of one or more of the rights that are
      specifically granted under this License.  You may not convey a covered
      work if you are a party to an arrangement with a third party that is
      in the business of distributing software, under which you make payment
      to the third party based on the extent of your activity of conveying
      the work, and under which the third party grants, to any of the
      parties who would receive the covered work from you, a discriminatory
      patent license (a) in connection with copies of the covered work
      conveyed by you (or copies made from those copies), or (b) primarily
      for and in connection with specific products or compilations that
      contain the covered work, unless you entered into that arrangement,
      or that patent license was granted, prior to 28 March 2007.
      
        Nothing in this License shall be construed as excluding or limiting
      any implied license or other defenses to infringement that may
      otherwise be available to you under applicable patent law.
      
        12. No Surrender of Others' Freedom.
      
        If conditions are imposed on you (whether by court order, agreement or
      otherwise) that contradict the conditions of this License, they do not
      excuse you from the conditions of this License.  If you cannot convey a
      covered work so as to satisfy simultaneously your obligations under this
      License and any other pertinent obligations, then as a consequence you may
      not convey it at all.  For example, if you agree to terms that obligate you
      to collect a royalty for further conveying from those to whom you convey
      the Program, the only way you could satisfy both those terms and this
      License would be to refrain entirely from conveying the Program.
      
        13. Remote Network Interaction; Use with the GNU General Public License.
      
        Notwithstanding any other provision of this License, if you modify the
      Program, your modified version must prominently offer all users
      interacting with it remotely through a computer network (if your version
      supports such interaction) an opportunity to receive the Corresponding
      Source of your version by providing access to the Corresponding Source
      from a network server at no charge, through some standard or customary
      means of facilitating copying of software.  This Corresponding Source
      shall include the Corresponding Source for any work covered by version 3
      of the GNU General Public License that is incorporated pursuant to the
      following paragraph.
      
        Notwithstanding any other provision of this License, you have
      permission to link or combine any covered work with a work licensed
      under version 3 of the GNU General Public License into a single
      combined work, and to convey the resulting work.  The terms of this
      License will continue to apply to the part which is the covered work,
      but the work with which it is combined will remain governed by version
      3 of the GNU General Public License.
      
        14. Revised Versions of this License.
      
        The Free Software Foundation may publish revised and/or new versions of
      the GNU Affero General Public License from time to time.  Such new versions
      will be similar in spirit to the present version, but may differ in detail to
      address new problems or concerns.
      
        Each version is given a distinguishing version number.  If the
      Program specifies that a certain numbered version of the GNU Affero General
      Public License "or any later version" applies to it, you have the
      option of following the terms and conditions either of that numbered
      version or of any later version published by the Free Software
      Foundation.  If the Program does not specify a version number of the
      GNU Affero General Public License, you may choose any version ever published
      by the Free Software Foundation.
      
        If the Program specifies that a proxy can decide which future
      versions of the GNU Affero General Public License can be used, that proxy's
      public statement of acceptance of a version permanently authorizes you
      to choose that version for the Program.
      
        Later license versions may give you additional or different
      permissions.  However, no additional obligations are imposed on any
      author or copyright holder as a result of your choosing to follow a
      later version.
      
        15. Disclaimer of Warranty.
      
        THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
      APPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
      HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
      OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
      THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
      PURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
      IS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
      ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
      
        16. Limitation of Liability.
      
        IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
      WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
      THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
      GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
      USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
      DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
      PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
      EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
      SUCH DAMAGES.
      
        17. Interpretation of Sections 15 and 16.
      
        If the disclaimer of warranty and limitation of liability provided
      above cannot be given local legal effect according to their terms,
      reviewing courts shall apply local law that most closely approximates
      an absolute waiver of all civil liability in connection with the
      Program, unless a warranty or assumption of liability accompanies a
      copy of the Program in return for a fee.
      
                           END OF TERMS AND CONDITIONS
      
                  How to Apply These Terms to Your New Programs
      
        If you develop a new program, and you want it to be of the greatest
      possible use to the public, the best way to achieve this is to make it
      free software which everyone can redistribute and change under these terms.
      
        To do so, attach the following notices to the program.  It is safest
      to attach them to the start of each source file to most effectively
      state the exclusion of warranty; and each file should have at least
      the "copyright" line and a pointer to where the full notice is found.
      
          <one line to give the program's name and a brief idea of what it does.>
          Copyright (C) <year>  <name of author>
      
          This program is free software: you can redistribute it and/or modify
          it under the terms of the GNU Affero General Public License as published by
          the Free Software Foundation, either version 3 of the License, or
          (at your option) any later version.
      
          This program is distributed in the hope that it will be useful,
          but WITHOUT ANY WARRANTY; without even the implied warranty of
          MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
          GNU Affero General Public License for more details.
      
          You should have received a copy of the GNU Affero General Public License
          along with this program.  If not, see <https://www.gnu.org/licenses/>.
      
      Also add information on how to contact you by electronic and paper mail.
      
        If your software can interact with users remotely through a computer
      network, you should also make sure that it provides a way for users to
      get its source.  For example, if your program is a web application, its
      interface could display a "Source" link that leads users to an archive
      of the code.  There are many ways you could offer source, and different
      solutions will be better for different programs; see section 13 for the
      specific requirements.
      
        You should also get your employer (if you work as a programmer) or school,
      if any, to sign a "copyright disclaimer" for the program, if necessary.
      For more information on this, and how to apply and follow the GNU AGPL, see
      <https://www.gnu.org/licenses/>.
      
    • 📄 README.md 41.4KB
      <div align="center">
        <p>
          <a href="http://www.ultralytics.com/blog/ultralytics-yolov8-turns-one-a-year-of-breakthroughs-and-innovations" target="_blank">
            <img width="100%" src="https://raw.githubusercontent.com/ultralytics/assets/main/yolov8/banner-yolov8.png"></a>
          <!--
          <a align="center" href="https://ultralytics.com/yolov5" target="_blank">
            <img width="100%" src="https://raw.githubusercontent.com/ultralytics/assets/main/yolov5/v70/splash.png"></a>
          -->
        </p>
      
      [中文](https://docs.ultralytics.com/zh/) | [한국어](https://docs.ultralytics.com/ko/) | [日本語](https://docs.ultralytics.com/ja/) | [Русский](https://docs.ultralytics.com/ru/) | [Deutsch](https://docs.ultralytics.com/de/) | [Français](https://docs.ultralytics.com/fr/) | [Español](https://docs.ultralytics.com/es/) | [Português](https://docs.ultralytics.com/pt/) | [Türkçe](https://docs.ultralytics.com/tr/) | [Tiếng Việt](https://docs.ultralytics.com/vi/) | [हिन्दी](https://docs.ultralytics.com/hi/) | [العربية](https://docs.ultralytics.com/ar/)
      
      <div>
          <a href="https://github.com/ultralytics/yolov5/actions/workflows/ci-testing.yml"><img src="https://github.com/ultralytics/yolov5/actions/workflows/ci-testing.yml/badge.svg" alt="YOLOv5 CI"></a>
          <a href="https://zenodo.org/badge/latestdoi/264818686"><img src="https://zenodo.org/badge/264818686.svg" alt="YOLOv5 Citation"></a>
          <a href="https://hub.docker.com/r/ultralytics/yolov5"><img src="https://img.shields.io/docker/pulls/ultralytics/yolov5?logo=docker" alt="Docker Pulls"></a>
          <a href="https://ultralytics.com/discord"><img alt="Discord" src="https://img.shields.io/discord/1089800235347353640?logo=discord&logoColor=white&label=Discord&color=blue"></a>
          <br>
          <a href="https://bit.ly/yolov5-paperspace-notebook"><img src="https://assets.paperspace.io/img/gradient-badge.svg" alt="Run on Gradient"></a>
          <a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
          <a href="https://www.kaggle.com/ultralytics/yolov5"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
        </div>
        <br>
      
      YOLOv5 🚀 is the world's most loved vision AI, representing <a href="https://ultralytics.com">Ultralytics</a> open-source research into future vision AI methods, incorporating lessons learned and best practices evolved over thousands of hours of research and development.
      
      We hope that the resources here will help you get the most out of YOLOv5. Please browse the YOLOv5 <a href="https://docs.ultralytics.com/yolov5">Docs</a> for details, raise an issue on <a href="https://github.com/ultralytics/yolov5/issues/new/choose">GitHub</a> for support, and join our <a href="https://ultralytics.com/discord">Discord</a> community for questions and discussions!
      
      To request an Enterprise License please complete the form at [Ultralytics Licensing](https://ultralytics.com/license).
      
      <div align="center">
        <a href="https://github.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-github.png" width="2%" alt="Ultralytics GitHub"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://www.linkedin.com/company/ultralytics/"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-linkedin.png" width="2%" alt="Ultralytics LinkedIn"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://twitter.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-twitter.png" width="2%" alt="Ultralytics Twitter"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://youtube.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-youtube.png" width="2%" alt="Ultralytics YouTube"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://www.tiktok.com/@ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-tiktok.png" width="2%" alt="Ultralytics TikTok"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://www.instagram.com/ultralytics/"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-instagram.png" width="2%" alt="Ultralytics Instagram"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://ultralytics.com/discord"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-discord.png" width="2%" alt="Ultralytics Discord"></a>
      </div>
      
      </div>
      <br>
      
      ## <div align="center">YOLOv8 🚀 NEW</div>
      
      We are thrilled to announce the launch of Ultralytics YOLOv8 🚀, our NEW cutting-edge, state-of-the-art (SOTA) model released at **[https://github.com/ultralytics/ultralytics](https://github.com/ultralytics/ultralytics)**. YOLOv8 is designed to be fast, accurate, and easy to use, making it an excellent choice for a wide range of object detection, image segmentation and image classification tasks.
      
      See the [YOLOv8 Docs](https://docs.ultralytics.com) for details and get started with:
      
      [![PyPI version](https://badge.fury.io/py/ultralytics.svg)](https://badge.fury.io/py/ultralytics) [![Downloads](https://static.pepy.tech/badge/ultralytics)](https://pepy.tech/project/ultralytics)
      
      ```bash
      pip install ultralytics
      ```
      
      <div align="center">
        <a href="https://ultralytics.com/yolov8" target="_blank">
        <img width="100%" src="https://raw.githubusercontent.com/ultralytics/assets/main/yolov8/yolo-comparison-plots.png"></a>
      </div>
      
      ## <div align="center">Documentation</div>
      
      See the [YOLOv5 Docs](https://docs.ultralytics.com/yolov5) for full documentation on training, testing and deployment. See below for quickstart examples.
      
      <details open>
      <summary>Install</summary>
      
      Clone repo and install [requirements.txt](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) in a [**Python>=3.8.0**](https://www.python.org/) environment, including [**PyTorch>=1.8**](https://pytorch.org/get-started/locally/).
      
      ```bash
      git clone https://github.com/ultralytics/yolov5  # clone
      cd yolov5
      pip install -r requirements.txt  # install
      ```
      
      </details>
      
      <details>
      <summary>Inference</summary>
      
      YOLOv5 [PyTorch Hub](https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading) inference. [Models](https://github.com/ultralytics/yolov5/tree/master/models) download automatically from the latest YOLOv5 [release](https://github.com/ultralytics/yolov5/releases).
      
      ```python
      import torch
      
      # Model
      model = torch.hub.load("ultralytics/yolov5", "yolov5s")  # or yolov5n - yolov5x6, custom
      
      # Images
      img = "https://ultralytics.com/images/zidane.jpg"  # or file, Path, PIL, OpenCV, numpy, list
      
      # Inference
      results = model(img)
      
      # Results
      results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
      ```
      
      </details>
      
      <details>
      <summary>Inference with detect.py</summary>
      
      `detect.py` runs inference on a variety of sources, downloading [models](https://github.com/ultralytics/yolov5/tree/master/models) automatically from the latest YOLOv5 [release](https://github.com/ultralytics/yolov5/releases) and saving results to `runs/detect`.
      
      ```bash
      python detect.py --weights yolov5s.pt --source 0                               # webcam
                                                     img.jpg                         # image
                                                     vid.mp4                         # video
                                                     screen                          # screenshot
                                                     path/                           # directory
                                                     list.txt                        # list of images
                                                     list.streams                    # list of streams
                                                     'path/*.jpg'                    # glob
                                                     'https://youtu.be/LNwODJXcvt4'  # YouTube
                                                     'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
      ```
      
      </details>
      
      <details>
      <summary>Training</summary>
      
      The commands below reproduce YOLOv5 [COCO](https://github.com/ultralytics/yolov5/blob/master/data/scripts/get_coco.sh) results. [Models](https://github.com/ultralytics/yolov5/tree/master/models) and [datasets](https://github.com/ultralytics/yolov5/tree/master/data) download automatically from the latest YOLOv5 [release](https://github.com/ultralytics/yolov5/releases). Training times for YOLOv5n/s/m/l/x are 1/2/4/6/8 days on a V100 GPU ([Multi-GPU](https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training) times faster). Use the largest `--batch-size` possible, or pass `--batch-size -1` for YOLOv5 [AutoBatch](https://github.com/ultralytics/yolov5/pull/5092). Batch sizes shown for V100-16GB.
      
      ```bash
      python train.py --data coco.yaml --epochs 300 --weights '' --cfg yolov5n.yaml  --batch-size 128
                                                                       yolov5s                    64
                                                                       yolov5m                    40
                                                                       yolov5l                    24
                                                                       yolov5x                    16
      ```
      
      <img width="800" src="https://user-images.githubusercontent.com/26833433/90222759-949d8800-ddc1-11ea-9fa1-1c97eed2b963.png">
      
      </details>
      
      <details open>
      <summary>Tutorials</summary>
      
      - [Train Custom Data](https://docs.ultralytics.com/yolov5/tutorials/train_custom_data) 🚀 RECOMMENDED
      - [Tips for Best Training Results](https://docs.ultralytics.com/yolov5/tutorials/tips_for_best_training_results) ☘️
      - [Multi-GPU Training](https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training)
      - [PyTorch Hub](https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading) 🌟 NEW
      - [TFLite, ONNX, CoreML, TensorRT Export](https://docs.ultralytics.com/yolov5/tutorials/model_export) 🚀
      - [NVIDIA Jetson platform Deployment](https://docs.ultralytics.com/yolov5/tutorials/running_on_jetson_nano) 🌟 NEW
      - [Test-Time Augmentation (TTA)](https://docs.ultralytics.com/yolov5/tutorials/test_time_augmentation)
      - [Model Ensembling](https://docs.ultralytics.com/yolov5/tutorials/model_ensembling)
      - [Model Pruning/Sparsity](https://docs.ultralytics.com/yolov5/tutorials/model_pruning_and_sparsity)
      - [Hyperparameter Evolution](https://docs.ultralytics.com/yolov5/tutorials/hyperparameter_evolution)
      - [Transfer Learning with Frozen Layers](https://docs.ultralytics.com/yolov5/tutorials/transfer_learning_with_frozen_layers)
      - [Architecture Summary](https://docs.ultralytics.com/yolov5/tutorials/architecture_description) 🌟 NEW
      - [Roboflow for Datasets, Labeling, and Active Learning](https://docs.ultralytics.com/yolov5/tutorials/roboflow_datasets_integration)
      - [ClearML Logging](https://docs.ultralytics.com/yolov5/tutorials/clearml_logging_integration) 🌟 NEW
      - [YOLOv5 with Neural Magic's Deepsparse](https://docs.ultralytics.com/yolov5/tutorials/neural_magic_pruning_quantization) 🌟 NEW
      - [Comet Logging](https://docs.ultralytics.com/yolov5/tutorials/comet_logging_integration) 🌟 NEW
      
      </details>
      
      ## <div align="center">Integrations</div>
      
      <br>
      <a align="center" href="https://bit.ly/ultralytics_hub" target="_blank">
      <img width="100%" src="https://github.com/ultralytics/assets/raw/main/im/integrations-loop.png"></a>
      <br>
      <br>
      
      <div align="center">
        <a href="https://roboflow.com/?ref=ultralytics">
          <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-roboflow.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
        <a href="https://cutt.ly/yolov5-readme-clearml">
          <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-clearml.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
        <a href="https://bit.ly/yolov5-readme-comet2">
          <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-comet.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
        <a href="https://bit.ly/yolov5-neuralmagic">
          <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-neuralmagic.png" width="10%" /></a>
      </div>
      
      |                                                           Roboflow                                                           |                                                            ClearML ⭐ NEW                                                            |                                                                        Comet ⭐ NEW                                                                         |                                           Neural Magic ⭐ NEW                                           |
      | :--------------------------------------------------------------------------------------------------------------------------: | :---------------------------------------------------------------------------------------------------------------------------------: | :--------------------------------------------------------------------------------------------------------------------------------------------------------: | :----------------------------------------------------------------------------------------------------: |
      | Label and export your custom datasets directly to YOLOv5 for training with [Roboflow](https://roboflow.com/?ref=ultralytics) | Automatically track, visualize and even remotely train YOLOv5 using [ClearML](https://cutt.ly/yolov5-readme-clearml) (open-source!) | Free forever, [Comet](https://bit.ly/yolov5-readme-comet2) lets you save YOLOv5 models, resume training, and interactively visualise and debug predictions | Run YOLOv5 inference up to 6x faster with [Neural Magic DeepSparse](https://bit.ly/yolov5-neuralmagic) |
      
      ## <div align="center">Ultralytics HUB</div>
      
      Experience seamless AI with [Ultralytics HUB](https://bit.ly/ultralytics_hub) ⭐, the all-in-one solution for data visualization, YOLOv5 and YOLOv8 🚀 model training and deployment, without any coding. Transform images into actionable insights and bring your AI visions to life with ease using our cutting-edge platform and user-friendly [Ultralytics App](https://ultralytics.com/app_install). Start your journey for **Free** now!
      
      <a align="center" href="https://bit.ly/ultralytics_hub" target="_blank">
      <img width="100%" src="https://github.com/ultralytics/assets/raw/main/im/ultralytics-hub.png"></a>
      
      ## <div align="center">Why YOLOv5</div>
      
      YOLOv5 has been designed to be super easy to get started and simple to learn. We prioritize real-world results.
      
      <p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/155040763-93c22a27-347c-4e3c-847a-8094621d3f4e.png"></p>
      <details>
        <summary>YOLOv5-P5 640 Figure</summary>
      
      <p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/155040757-ce0934a3-06a6-43dc-a979-2edbbd69ea0e.png"></p>
      </details>
      <details>
        <summary>Figure Notes</summary>
      
      - **COCO AP val** denotes mAP@0.5:0.95 metric measured on the 5000-image [COCO val2017](http://cocodataset.org) dataset over various inference sizes from 256 to 1536.
      - **GPU Speed** measures average inference time per image on [COCO val2017](http://cocodataset.org) dataset using a [AWS p3.2xlarge](https://aws.amazon.com/ec2/instance-types/p3/) V100 instance at batch-size 32.
      - **EfficientDet** data from [google/automl](https://github.com/google/automl) at batch size 8.
      - **Reproduce** by `python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n6.pt yolov5s6.pt yolov5m6.pt yolov5l6.pt yolov5x6.pt`
      
      </details>
      
      ### Pretrained Checkpoints
      
      | Model                                                                                           | size<br><sup>(pixels) | mAP<sup>val<br>50-95 | mAP<sup>val<br>50 | Speed<br><sup>CPU b1<br>(ms) | Speed<br><sup>V100 b1<br>(ms) | Speed<br><sup>V100 b32<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>@640 (B) |
      | ----------------------------------------------------------------------------------------------- | --------------------- | -------------------- | ----------------- | ---------------------------- | ----------------------------- | ------------------------------ | ------------------ | ---------------------- |
      | [YOLOv5n](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n.pt)              | 640                   | 28.0                 | 45.7              | **45**                       | **6.3**                       | **0.6**                        | **1.9**            | **4.5**                |
      | [YOLOv5s](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt)              | 640                   | 37.4                 | 56.8              | 98                           | 6.4                           | 0.9                            | 7.2                | 16.5                   |
      | [YOLOv5m](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m.pt)              | 640                   | 45.4                 | 64.1              | 224                          | 8.2                           | 1.7                            | 21.2               | 49.0                   |
      | [YOLOv5l](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5l.pt)              | 640                   | 49.0                 | 67.3              | 430                          | 10.1                          | 2.7                            | 46.5               | 109.1                  |
      | [YOLOv5x](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5x.pt)              | 640                   | 50.7                 | 68.9              | 766                          | 12.1                          | 4.8                            | 86.7               | 205.7                  |
      |                                                                                                 |                       |                      |                   |                              |                               |                                |                    |                        |
      | [YOLOv5n6](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n6.pt)            | 1280                  | 36.0                 | 54.4              | 153                          | 8.1                           | 2.1                            | 3.2                | 4.6                    |
      | [YOLOv5s6](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s6.pt)            | 1280                  | 44.8                 | 63.7              | 385                          | 8.2                           | 3.6                            | 12.6               | 16.8                   |
      | [YOLOv5m6](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m6.pt)            | 1280                  | 51.3                 | 69.3              | 887                          | 11.1                          | 6.8                            | 35.7               | 50.0                   |
      | [YOLOv5l6](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5l6.pt)            | 1280                  | 53.7                 | 71.3              | 1784                         | 15.8                          | 10.5                           | 76.8               | 111.4                  |
      | [YOLOv5x6](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5x6.pt)<br>+ [TTA] | 1280<br>1536          | 55.0<br>**55.8**     | 72.7<br>**72.7**  | 3136<br>-                    | 26.2<br>-                     | 19.4<br>-                      | 140.7<br>-         | 209.8<br>-             |
      
      <details>
        <summary>Table Notes</summary>
      
      - All checkpoints are trained to 300 epochs with default settings. Nano and Small models use [hyp.scratch-low.yaml](https://github.com/ultralytics/yolov5/blob/master/data/hyps/hyp.scratch-low.yaml) hyps, all others use [hyp.scratch-high.yaml](https://github.com/ultralytics/yolov5/blob/master/data/hyps/hyp.scratch-high.yaml).
      - **mAP<sup>val</sup>** values are for single-model single-scale on [COCO val2017](http://cocodataset.org) dataset.<br>Reproduce by `python val.py --data coco.yaml --img 640 --conf 0.001 --iou 0.65`
      - **Speed** averaged over COCO val images using a [AWS p3.2xlarge](https://aws.amazon.com/ec2/instance-types/p3/) instance. NMS times (~1 ms/img) not included.<br>Reproduce by `python val.py --data coco.yaml --img 640 --task speed --batch 1`
      - **TTA** [Test Time Augmentation](https://docs.ultralytics.com/yolov5/tutorials/test_time_augmentation) includes reflection and scale augmentations.<br>Reproduce by `python val.py --data coco.yaml --img 1536 --iou 0.7 --augment`
      
      </details>
      
      ## <div align="center">Segmentation</div>
      
      Our new YOLOv5 [release v7.0](https://github.com/ultralytics/yolov5/releases/v7.0) instance segmentation models are the fastest and most accurate in the world, beating all current [SOTA benchmarks](https://paperswithcode.com/sota/real-time-instance-segmentation-on-mscoco). We've made them super simple to train, validate and deploy. See full details in our [Release Notes](https://github.com/ultralytics/yolov5/releases/v7.0) and visit our [YOLOv5 Segmentation Colab Notebook](https://github.com/ultralytics/yolov5/blob/master/segment/tutorial.ipynb) for quickstart tutorials.
      
      <details>
        <summary>Segmentation Checkpoints</summary>
      
      <div align="center">
      <a align="center" href="https://ultralytics.com/yolov5" target="_blank">
      <img width="800" src="https://user-images.githubusercontent.com/61612323/204180385-84f3aca9-a5e9-43d8-a617-dda7ca12e54a.png"></a>
      </div>
      
      We trained YOLOv5 segmentations models on COCO for 300 epochs at image size 640 using A100 GPUs. We exported all models to ONNX FP32 for CPU speed tests and to TensorRT FP16 for GPU speed tests. We ran all speed tests on Google [Colab Pro](https://colab.research.google.com/signup) notebooks for easy reproducibility.
      
      | Model                                                                                      | size<br><sup>(pixels) | mAP<sup>box<br>50-95 | mAP<sup>mask<br>50-95 | Train time<br><sup>300 epochs<br>A100 (hours) | Speed<br><sup>ONNX CPU<br>(ms) | Speed<br><sup>TRT A100<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>@640 (B) |
      | ------------------------------------------------------------------------------------------ | --------------------- | -------------------- | --------------------- | --------------------------------------------- | ------------------------------ | ------------------------------ | ------------------ | ---------------------- |
      | [YOLOv5n-seg](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n-seg.pt) | 640                   | 27.6                 | 23.4                  | 80:17                                         | **62.7**                       | **1.2**                        | **2.0**            | **7.1**                |
      | [YOLOv5s-seg](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s-seg.pt) | 640                   | 37.6                 | 31.7                  | 88:16                                         | 173.3                          | 1.4                            | 7.6                | 26.4                   |
      | [YOLOv5m-seg](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m-seg.pt) | 640                   | 45.0                 | 37.1                  | 108:36                                        | 427.0                          | 2.2                            | 22.0               | 70.8                   |
      | [YOLOv5l-seg](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5l-seg.pt) | 640                   | 49.0                 | 39.9                  | 66:43 (2x)                                    | 857.4                          | 2.9                            | 47.9               | 147.7                  |
      | [YOLOv5x-seg](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5x-seg.pt) | 640                   | **50.7**             | **41.4**              | 62:56 (3x)                                    | 1579.2                         | 4.5                            | 88.8               | 265.7                  |
      
      - All checkpoints are trained to 300 epochs with SGD optimizer with `lr0=0.01` and `weight_decay=5e-5` at image size 640 and all default settings.<br>Runs logged to https://wandb.ai/glenn-jocher/YOLOv5_v70_official
      - **Accuracy** values are for single-model single-scale on COCO dataset.<br>Reproduce by `python segment/val.py --data coco.yaml --weights yolov5s-seg.pt`
      - **Speed** averaged over 100 inference images using a [Colab Pro](https://colab.research.google.com/signup) A100 High-RAM instance. Values indicate inference speed only (NMS adds about 1ms per image). <br>Reproduce by `python segment/val.py --data coco.yaml --weights yolov5s-seg.pt --batch 1`
      - **Export** to ONNX at FP32 and TensorRT at FP16 done with `export.py`. <br>Reproduce by `python export.py --weights yolov5s-seg.pt --include engine --device 0 --half`
      
      </details>
      
      <details>
        <summary>Segmentation Usage Examples &nbsp;<a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/segment/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a></summary>
      
      ### Train
      
      YOLOv5 segmentation training supports auto-download COCO128-seg segmentation dataset with `--data coco128-seg.yaml` argument and manual download of COCO-segments dataset with `bash data/scripts/get_coco.sh --train --val --segments` and then `python train.py --data coco.yaml`.
      
      ```bash
      # Single-GPU
      python segment/train.py --data coco128-seg.yaml --weights yolov5s-seg.pt --img 640
      
      # Multi-GPU DDP
      python -m torch.distributed.run --nproc_per_node 4 --master_port 1 segment/train.py --data coco128-seg.yaml --weights yolov5s-seg.pt --img 640 --device 0,1,2,3
      ```
      
      ### Val
      
      Validate YOLOv5s-seg mask mAP on COCO dataset:
      
      ```bash
      bash data/scripts/get_coco.sh --val --segments  # download COCO val segments split (780MB, 5000 images)
      python segment/val.py --weights yolov5s-seg.pt --data coco.yaml --img 640  # validate
      ```
      
      ### Predict
      
      Use pretrained YOLOv5m-seg.pt to predict bus.jpg:
      
      ```bash
      python segment/predict.py --weights yolov5m-seg.pt --source data/images/bus.jpg
      ```
      
      ```python
      model = torch.hub.load(
          "ultralytics/yolov5", "custom", "yolov5m-seg.pt"
      )  # load from PyTorch Hub (WARNING: inference not yet supported)
      ```
      
      | ![zidane](https://user-images.githubusercontent.com/26833433/203113421-decef4c4-183d-4a0a-a6c2-6435b33bc5d3.jpg) | ![bus](https://user-images.githubusercontent.com/26833433/203113416-11fe0025-69f7-4874-a0a6-65d0bfe2999a.jpg) |
      | ---------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------- |
      
      ### Export
      
      Export YOLOv5s-seg model to ONNX and TensorRT:
      
      ```bash
      python export.py --weights yolov5s-seg.pt --include onnx engine --img 640 --device 0
      ```
      
      </details>
      
      ## <div align="center">Classification</div>
      
      YOLOv5 [release v6.2](https://github.com/ultralytics/yolov5/releases) brings support for classification model training, validation and deployment! See full details in our [Release Notes](https://github.com/ultralytics/yolov5/releases/v6.2) and visit our [YOLOv5 Classification Colab Notebook](https://github.com/ultralytics/yolov5/blob/master/classify/tutorial.ipynb) for quickstart tutorials.
      
      <details>
        <summary>Classification Checkpoints</summary>
      
      <br>
      
      We trained YOLOv5-cls classification models on ImageNet for 90 epochs using a 4xA100 instance, and we trained ResNet and EfficientNet models alongside with the same default training settings to compare. We exported all models to ONNX FP32 for CPU speed tests and to TensorRT FP16 for GPU speed tests. We ran all speed tests on Google [Colab Pro](https://colab.research.google.com/signup) for easy reproducibility.
      
      | Model                                                                                              | size<br><sup>(pixels) | acc<br><sup>top1 | acc<br><sup>top5 | Training<br><sup>90 epochs<br>4xA100 (hours) | Speed<br><sup>ONNX CPU<br>(ms) | Speed<br><sup>TensorRT V100<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>@224 (B) |
      | -------------------------------------------------------------------------------------------------- | --------------------- | ---------------- | ---------------- | -------------------------------------------- | ------------------------------ | ----------------------------------- | ------------------ | ---------------------- |
      | [YOLOv5n-cls](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n-cls.pt)         | 224                   | 64.6             | 85.4             | 7:59                                         | **3.3**                        | **0.5**                             | **2.5**            | **0.5**                |
      | [YOLOv5s-cls](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s-cls.pt)         | 224                   | 71.5             | 90.2             | 8:09                                         | 6.6                            | 0.6                                 | 5.4                | 1.4                    |
      | [YOLOv5m-cls](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m-cls.pt)         | 224                   | 75.9             | 92.9             | 10:06                                        | 15.5                           | 0.9                                 | 12.9               | 3.9                    |
      | [YOLOv5l-cls](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5l-cls.pt)         | 224                   | 78.0             | 94.0             | 11:56                                        | 26.9                           | 1.4                                 | 26.5               | 8.5                    |
      | [YOLOv5x-cls](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5x-cls.pt)         | 224                   | **79.0**         | **94.4**         | 15:04                                        | 54.3                           | 1.8                                 | 48.1               | 15.9                   |
      |                                                                                                    |                       |                  |                  |                                              |                                |                                     |                    |                        |
      | [ResNet18](https://github.com/ultralytics/yolov5/releases/download/v7.0/resnet18.pt)               | 224                   | 70.3             | 89.5             | **6:47**                                     | 11.2                           | 0.5                                 | 11.7               | 3.7                    |
      | [ResNet34](https://github.com/ultralytics/yolov5/releases/download/v7.0/resnet34.pt)               | 224                   | 73.9             | 91.8             | 8:33                                         | 20.6                           | 0.9                                 | 21.8               | 7.4                    |
      | [ResNet50](https://github.com/ultralytics/yolov5/releases/download/v7.0/resnet50.pt)               | 224                   | 76.8             | 93.4             | 11:10                                        | 23.4                           | 1.0                                 | 25.6               | 8.5                    |
      | [ResNet101](https://github.com/ultralytics/yolov5/releases/download/v7.0/resnet101.pt)             | 224                   | 78.5             | 94.3             | 17:10                                        | 42.1                           | 1.9                                 | 44.5               | 15.9                   |
      |                                                                                                    |                       |                  |                  |                                              |                                |                                     |                    |                        |
      | [EfficientNet_b0](https://github.com/ultralytics/yolov5/releases/download/v7.0/efficientnet_b0.pt) | 224                   | 75.1             | 92.4             | 13:03                                        | 12.5                           | 1.3                                 | 5.3                | 1.0                    |
      | [EfficientNet_b1](https://github.com/ultralytics/yolov5/releases/download/v7.0/efficientnet_b1.pt) | 224                   | 76.4             | 93.2             | 17:04                                        | 14.9                           | 1.6                                 | 7.8                | 1.5                    |
      | [EfficientNet_b2](https://github.com/ultralytics/yolov5/releases/download/v7.0/efficientnet_b2.pt) | 224                   | 76.6             | 93.4             | 17:10                                        | 15.9                           | 1.6                                 | 9.1                | 1.7                    |
      | [EfficientNet_b3](https://github.com/ultralytics/yolov5/releases/download/v7.0/efficientnet_b3.pt) | 224                   | 77.7             | 94.0             | 19:19                                        | 18.9                           | 1.9                                 | 12.2               | 2.4                    |
      
      <details>
        <summary>Table Notes (click to expand)</summary>
      
      - All checkpoints are trained to 90 epochs with SGD optimizer with `lr0=0.001` and `weight_decay=5e-5` at image size 224 and all default settings.<br>Runs logged to https://wandb.ai/glenn-jocher/YOLOv5-Classifier-v6-2
      - **Accuracy** values are for single-model single-scale on [ImageNet-1k](https://www.image-net.org/index.php) dataset.<br>Reproduce by `python classify/val.py --data ../datasets/imagenet --img 224`
      - **Speed** averaged over 100 inference images using a Google [Colab Pro](https://colab.research.google.com/signup) V100 High-RAM instance.<br>Reproduce by `python classify/val.py --data ../datasets/imagenet --img 224 --batch 1`
      - **Export** to ONNX at FP32 and TensorRT at FP16 done with `export.py`. <br>Reproduce by `python export.py --weights yolov5s-cls.pt --include engine onnx --imgsz 224`
      
      </details>
      </details>
      
      <details>
        <summary>Classification Usage Examples &nbsp;<a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/classify/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a></summary>
      
      ### Train
      
      YOLOv5 classification training supports auto-download of MNIST, Fashion-MNIST, CIFAR10, CIFAR100, Imagenette, Imagewoof, and ImageNet datasets with the `--data` argument. To start training on MNIST for example use `--data mnist`.
      
      ```bash
      # Single-GPU
      python classify/train.py --model yolov5s-cls.pt --data cifar100 --epochs 5 --img 224 --batch 128
      
      # Multi-GPU DDP
      python -m torch.distributed.run --nproc_per_node 4 --master_port 1 classify/train.py --model yolov5s-cls.pt --data imagenet --epochs 5 --img 224 --device 0,1,2,3
      ```
      
      ### Val
      
      Validate YOLOv5m-cls accuracy on ImageNet-1k dataset:
      
      ```bash
      bash data/scripts/get_imagenet.sh --val  # download ImageNet val split (6.3G, 50000 images)
      python classify/val.py --weights yolov5m-cls.pt --data ../datasets/imagenet --img 224  # validate
      ```
      
      ### Predict
      
      Use pretrained YOLOv5s-cls.pt to predict bus.jpg:
      
      ```bash
      python classify/predict.py --weights yolov5s-cls.pt --source data/images/bus.jpg
      ```
      
      ```python
      model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s-cls.pt")  # load from PyTorch Hub
      ```
      
      ### Export
      
      Export a group of trained YOLOv5s-cls, ResNet and EfficientNet models to ONNX and TensorRT:
      
      ```bash
      python export.py --weights yolov5s-cls.pt resnet50.pt efficientnet_b0.pt --include onnx engine --img 224
      ```
      
      </details>
      
      ## <div align="center">Environments</div>
      
      Get started in seconds with our verified environments. Click each icon below for details.
      
      <div align="center">
        <a href="https://bit.ly/yolov5-paperspace-notebook">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-gradient.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="5%" alt="" />
        <a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-colab-small.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="5%" alt="" />
        <a href="https://www.kaggle.com/ultralytics/yolov5">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-kaggle-small.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="5%" alt="" />
        <a href="https://hub.docker.com/r/ultralytics/yolov5">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-docker-small.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="5%" alt="" />
        <a href="https://docs.ultralytics.com/yolov5/environments/aws_quickstart_tutorial/">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-aws-small.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="5%" alt="" />
        <a href="https://docs.ultralytics.com/yolov5/environments/google_cloud_quickstart_tutorial/">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-gcp-small.png" width="10%" /></a>
      </div>
      
      ## <div align="center">Contribute</div>
      
      We love your input! We want to make contributing to YOLOv5 as easy and transparent as possible. Please see our [Contributing Guide](https://docs.ultralytics.com/help/contributing/) to get started, and fill out the [YOLOv5 Survey](https://ultralytics.com/survey?utm_source=github&utm_medium=social&utm_campaign=Survey) to send us feedback on your experiences. Thank you to all our contributors!
      
      <!-- SVG image from https://opencollective.com/ultralytics/contributors.svg?width=990 -->
      
      <a href="https://github.com/ultralytics/yolov5/graphs/contributors">
      <img src="https://github.com/ultralytics/assets/raw/main/im/image-contributors.png" /></a>
      
      ## <div align="center">License</div>
      
      Ultralytics offers two licensing options to accommodate diverse use cases:
      
      - **AGPL-3.0 License**: This [OSI-approved](https://opensource.org/licenses/) open-source license is ideal for students and enthusiasts, promoting open collaboration and knowledge sharing. See the [LICENSE](https://github.com/ultralytics/yolov5/blob/master/LICENSE) file for more details.
      - **Enterprise License**: Designed for commercial use, this license permits seamless integration of Ultralytics software and AI models into commercial goods and services, bypassing the open-source requirements of AGPL-3.0. If your scenario involves embedding our solutions into a commercial offering, reach out through [Ultralytics Licensing](https://ultralytics.com/license).
      
      ## <div align="center">Contact</div>
      
      For YOLOv5 bug reports and feature requests please visit [GitHub Issues](https://github.com/ultralytics/yolov5/issues), and join our [Discord](https://ultralytics.com/discord) community for questions and discussions!
      
      <br>
      <div align="center">
        <a href="https://github.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-github.png" width="3%" alt="Ultralytics GitHub"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://www.linkedin.com/company/ultralytics/"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-linkedin.png" width="3%" alt="Ultralytics LinkedIn"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://twitter.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-twitter.png" width="3%" alt="Ultralytics Twitter"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://youtube.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-youtube.png" width="3%" alt="Ultralytics YouTube"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://www.tiktok.com/@ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-tiktok.png" width="3%" alt="Ultralytics TikTok"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://www.instagram.com/ultralytics/"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-instagram.png" width="3%" alt="Ultralytics Instagram"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://ultralytics.com/discord"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-discord.png" width="3%" alt="Ultralytics Discord"></a>
      </div>
      
      [tta]: https://docs.ultralytics.com/yolov5/tutorials/test_time_augmentation
      
    • 📄 README.zh-CN.md 41.3KB
      <div align="center">
        <p>
          <a href="http://www.ultralytics.com/blog/ultralytics-yolov8-turns-one-a-year-of-breakthroughs-and-innovations" target="_blank">
            <img width="100%" src="https://raw.githubusercontent.com/ultralytics/assets/main/yolov8/banner-yolov8.png"></a>
          <!--
          <a align="center" href="https://ultralytics.com/yolov5" target="_blank">
            <img width="100%" src="https://raw.githubusercontent.com/ultralytics/assets/main/yolov5/v70/splash.png"></a>
          -->
        </p>
      
      [中文](https://docs.ultralytics.com/zh/) | [한국어](https://docs.ultralytics.com/ko/) | [日本語](https://docs.ultralytics.com/ja/) | [Русский](https://docs.ultralytics.com/ru/) | [Deutsch](https://docs.ultralytics.com/de/) | [Français](https://docs.ultralytics.com/fr/) | [Español](https://docs.ultralytics.com/es/) | [Português](https://docs.ultralytics.com/pt/) | [Türkçe](https://docs.ultralytics.com/tr/) | [Tiếng Việt](https://docs.ultralytics.com/vi/) | [हिन्दी](https://docs.ultralytics.com/hi/) | [العربية](https://docs.ultralytics.com/ar/)
      
      <div>
          <a href="https://github.com/ultralytics/yolov5/actions/workflows/ci-testing.yml"><img src="https://github.com/ultralytics/yolov5/actions/workflows/ci-testing.yml/badge.svg" alt="YOLOv5 CI"></a>
          <a href="https://zenodo.org/badge/latestdoi/264818686"><img src="https://zenodo.org/badge/264818686.svg" alt="YOLOv5 Citation"></a>
          <a href="https://hub.docker.com/r/ultralytics/yolov5"><img src="https://img.shields.io/docker/pulls/ultralytics/yolov5?logo=docker" alt="Docker Pulls"></a>
          <br>
          <a href="https://bit.ly/yolov5-paperspace-notebook"><img src="https://assets.paperspace.io/img/gradient-badge.svg" alt="Run on Gradient"></a>
          <a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
          <a href="https://www.kaggle.com/ultralytics/yolov5"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
        </div>
        <br>
      
      YOLOv5 🚀 是世界上最受欢迎的视觉 AI,代表<a href="https://ultralytics.com"> Ultralytics </a>对未来视觉 AI 方法的开源研究,结合在数千小时的研究和开发中积累的经验教训和最佳实践。
      
      我们希望这里的资源能帮助您充分利用 YOLOv5。请浏览 YOLOv5 <a href="https://docs.ultralytics.com/yolov5/">文档</a> 了解详细信息,在 <a href="https://github.com/ultralytics/yolov5/issues/new/choose">GitHub</a> 上提交问题以获得支持,并加入我们的 <a href="https://ultralytics.com/discord">Discord</a> 社区进行问题和讨论!
      
      如需申请企业许可,请在 [Ultralytics Licensing](https://ultralytics.com/license) 处填写表格
      
      <div align="center">
        <a href="https://github.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-github.png" width="2%" alt="Ultralytics GitHub"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://www.linkedin.com/company/ultralytics/"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-linkedin.png" width="2%" alt="Ultralytics LinkedIn"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://twitter.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-twitter.png" width="2%" alt="Ultralytics Twitter"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://youtube.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-youtube.png" width="2%" alt="Ultralytics YouTube"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://www.tiktok.com/@ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-tiktok.png" width="2%" alt="Ultralytics TikTok"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://www.instagram.com/ultralytics/"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-instagram.png" width="2%" alt="Ultralytics Instagram"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="2%">
        <a href="https://ultralytics.com/discord"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-discord.png" width="2%" alt="Ultralytics Discord"></a>
      </div>
      </div>
      
      ## <div align="center">YOLOv8 🚀 新品</div>
      
      我们很高兴宣布 Ultralytics YOLOv8 🚀 的发布,这是我们新推出的领先水平、最先进的(SOTA)模型,发布于 **[https://github.com/ultralytics/ultralytics](https://github.com/ultralytics/ultralytics)**。 YOLOv8 旨在快速、准确且易于使用,使其成为广泛的物体检测、图像分割和图像分类任务的极佳选择。
      
      请查看 [YOLOv8 文档](https://docs.ultralytics.com)了解详细信息,并开始使用:
      
      [![PyPI 版本](https://badge.fury.io/py/ultralytics.svg)](https://badge.fury.io/py/ultralytics) [![下载量](https://static.pepy.tech/badge/ultralytics)](https://pepy.tech/project/ultralytics)
      
      ```commandline
      pip install ultralytics
      ```
      
      <div align="center">
        <a href="https://ultralytics.com/yolov8" target="_blank">
        <img width="100%" src="https://raw.githubusercontent.com/ultralytics/assets/main/yolov8/yolo-comparison-plots.png"></a>
      </div>
      
      ## <div align="center">文档</div>
      
      有关训练、测试和部署的完整文档见[YOLOv5 文档](https://docs.ultralytics.com/yolov5/)。请参阅下面的快速入门示例。
      
      <details open>
      <summary>安装</summary>
      
      克隆 repo,并要求在 [**Python>=3.8.0**](https://www.python.org/) 环境中安装 [requirements.txt](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) ,且要求 [**PyTorch>=1.8**](https://pytorch.org/get-started/locally/) 。
      
      ```bash
      git clone https://github.com/ultralytics/yolov5  # clone
      cd yolov5
      pip install -r requirements.txt  # install
      ```
      
      </details>
      
      <details>
      <summary>推理</summary>
      
      使用 YOLOv5 [PyTorch Hub](https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading) 推理。最新 [模型](https://github.com/ultralytics/yolov5/tree/master/models) 将自动的从 YOLOv5 [release](https://github.com/ultralytics/yolov5/releases) 中下载。
      
      ```python
      import torch
      
      # Model
      model = torch.hub.load("ultralytics/yolov5", "yolov5s")  # or yolov5n - yolov5x6, custom
      
      # Images
      img = "https://ultralytics.com/images/zidane.jpg"  # or file, Path, PIL, OpenCV, numpy, list
      
      # Inference
      results = model(img)
      
      # Results
      results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
      ```
      
      </details>
      
      <details>
      <summary>使用 detect.py 推理</summary>
      
      `detect.py` 在各种来源上运行推理, [模型](https://github.com/ultralytics/yolov5/tree/master/models) 自动从 最新的YOLOv5 [release](https://github.com/ultralytics/yolov5/releases) 中下载,并将结果保存到 `runs/detect` 。
      
      ```bash
      python detect.py --weights yolov5s.pt --source 0                               # webcam
                                                     img.jpg                         # image
                                                     vid.mp4                         # video
                                                     screen                          # screenshot
                                                     path/                           # directory
                                                     list.txt                        # list of images
                                                     list.streams                    # list of streams
                                                     'path/*.jpg'                    # glob
                                                     'https://youtu.be/LNwODJXcvt4'  # YouTube
                                                     'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
      ```
      
      </details>
      
      <details>
      <summary>训练</summary>
      
      下面的命令重现 YOLOv5 在 [COCO](https://github.com/ultralytics/yolov5/blob/master/data/scripts/get_coco.sh) 数据集上的结果。 最新的 [模型](https://github.com/ultralytics/yolov5/tree/master/models) 和 [数据集](https://github.com/ultralytics/yolov5/tree/master/data)
      将自动的从 YOLOv5 [release](https://github.com/ultralytics/yolov5/releases) 中下载。 YOLOv5n/s/m/l/x 在 V100 GPU 的训练时间为 1/2/4/6/8 天( [多GPU](https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training) 训练速度更快)。 尽可能使用更大的 `--batch-size` ,或通过 `--batch-size -1` 实现 YOLOv5 [自动批处理](https://github.com/ultralytics/yolov5/pull/5092) 。下方显示的 batchsize 适用于 V100-16GB。
      
      ```bash
      python train.py --data coco.yaml --epochs 300 --weights '' --cfg yolov5n.yaml  --batch-size 128
                                                                       yolov5s                    64
                                                                       yolov5m                    40
                                                                       yolov5l                    24
                                                                       yolov5x                    16
      ```
      
      <img width="800" src="https://user-images.githubusercontent.com/26833433/90222759-949d8800-ddc1-11ea-9fa1-1c97eed2b963.png">
      
      </details>
      
      <details open>
      <summary>教程</summary>
      
      - [训练自定义数据](https://docs.ultralytics.com/yolov5/tutorials/train_custom_data) 🚀 推荐
      - [获得最佳训练结果的技巧](https://docs.ultralytics.com/yolov5/tutorials/tips_for_best_training_results) ☘️
      - [多GPU训练](https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training)
      - [PyTorch Hub](https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading) 🌟 新
      - [TFLite,ONNX,CoreML,TensorRT导出](https://docs.ultralytics.com/yolov5/tutorials/model_export) 🚀
      - [NVIDIA Jetson平台部署](https://docs.ultralytics.com/yolov5/tutorials/running_on_jetson_nano) 🌟 新
      - [测试时增强 (TTA)](https://docs.ultralytics.com/yolov5/tutorials/test_time_augmentation)
      - [模型集成](https://docs.ultralytics.com/yolov5/tutorials/model_ensembling)
      - [模型剪枝/稀疏](https://docs.ultralytics.com/yolov5/tutorials/model_pruning_and_sparsity)
      - [超参数进化](https://docs.ultralytics.com/yolov5/tutorials/hyperparameter_evolution)
      - [冻结层的迁移学习](https://docs.ultralytics.com/yolov5/tutorials/transfer_learning_with_frozen_layers)
      - [架构概述](https://docs.ultralytics.com/yolov5/tutorials/architecture_description) 🌟 新
      - [Roboflow用于数据集、标注和主动学习](https://docs.ultralytics.com/yolov5/tutorials/roboflow_datasets_integration)
      - [ClearML日志记录](https://docs.ultralytics.com/yolov5/tutorials/clearml_logging_integration) 🌟 新
      - [使用Neural Magic的Deepsparse的YOLOv5](https://docs.ultralytics.com/yolov5/tutorials/neural_magic_pruning_quantization) 🌟 新
      - [Comet日志记录](https://docs.ultralytics.com/yolov5/tutorials/comet_logging_integration) 🌟 新
      
      </details>
      
      ## <div align="center">模块集成</div>
      
      <br>
      <a align="center" href="https://bit.ly/ultralytics_hub" target="_blank">
      <img width="100%" src="https://github.com/ultralytics/assets/raw/main/im/integrations-loop.png"></a>
      <br>
      <br>
      
      <div align="center">
        <a href="https://roboflow.com/?ref=ultralytics">
          <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-roboflow.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
        <a href="https://cutt.ly/yolov5-readme-clearml">
          <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-clearml.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
        <a href="https://bit.ly/yolov5-readme-comet2">
          <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-comet.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
        <a href="https://bit.ly/yolov5-neuralmagic">
          <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-neuralmagic.png" width="10%" /></a>
      </div>
      
      |                                                  Roboflow                                                  |                                         ClearML ⭐ 新                                          |                                                   Comet ⭐ 新                                                    |                                           Neural Magic ⭐ 新                                           |
      | :--------------------------------------------------------------------------------------------------------: | :--------------------------------------------------------------------------------------------: | :--------------------------------------------------------------------------------------------------------------: | :----------------------------------------------------------------------------------------------------: |
      | 将您的自定义数据集进行标注并直接导出到 YOLOv5 以进行训练 [Roboflow](https://roboflow.com/?ref=ultralytics) | 自动跟踪、可视化甚至远程训练 YOLOv5 [ClearML](https://cutt.ly/yolov5-readme-clearml)(开源!) | 永远免费,[Comet](https://bit.ly/yolov5-readme-comet2)可让您保存 YOLOv5 模型、恢复训练以及交互式可视化和调试预测 | 使用 [Neural Magic DeepSparse](https://bit.ly/yolov5-neuralmagic),运行 YOLOv5 推理的速度最高可提高6倍 |
      
      ## <div align="center">Ultralytics HUB</div>
      
      [Ultralytics HUB](https://bit.ly/ultralytics_hub) 是我们的⭐**新的**用于可视化数据集、训练 YOLOv5 🚀 模型并以无缝体验部署到现实世界的无代码解决方案。现在开始 **免费** 使用他!
      
      <a align="center" href="https://bit.ly/ultralytics_hub" target="_blank">
      <img width="100%" src="https://github.com/ultralytics/assets/raw/main/im/ultralytics-hub.png"></a>
      
      ## <div align="center">为什么选择 YOLOv5</div>
      
      YOLOv5 超级容易上手,简单易学。我们优先考虑现实世界的结果。
      
      <p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/155040763-93c22a27-347c-4e3c-847a-8094621d3f4e.png"></p>
      <details>
        <summary>YOLOv5-P5 640 图</summary>
      
      <p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/155040757-ce0934a3-06a6-43dc-a979-2edbbd69ea0e.png"></p>
      </details>
      <details>
        <summary>图表笔记</summary>
      
      - **COCO AP val** 表示 mAP@0.5:0.95 指标,在 [COCO val2017](http://cocodataset.org) 数据集的 5000 张图像上测得, 图像包含 256 到 1536 各种推理大小。
      - **显卡推理速度** 为在 [COCO val2017](http://cocodataset.org) 数据集上的平均推理时间,使用 [AWS p3.2xlarge](https://aws.amazon.com/ec2/instance-types/p3/) V100实例,batchsize 为 32 。
      - **EfficientDet** 数据来自 [google/automl](https://github.com/google/automl) , batchsize 为32。
      - **复现命令** 为 `python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n6.pt yolov5s6.pt yolov5m6.pt yolov5l6.pt yolov5x6.pt`
      
      </details>
      
      ### 预训练模型
      
      | 模型                                                                                           | 尺寸<br><sup>(像素) | mAP<sup>val<br>50-95 | mAP<sup>val<br>50 | 推理速度<br><sup>CPU b1<br>(ms) | 推理速度<br><sup>V100 b1<br>(ms) | 速度<br><sup>V100 b32<br>(ms) | 参数量<br><sup>(M) | FLOPs<br><sup>@640 (B) |
      | ---------------------------------------------------------------------------------------------- | --------------------- | -------------------- | ----------------- | --------------------------------- | ---------------------------------- | ------------------------------- | ------------------ | ---------------------- |
      | [YOLOv5n](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n.pt)             | 640                   | 28.0                 | 45.7              | **45**                            | **6.3**                            | **0.6**                         | **1.9**            | **4.5**                |
      | [YOLOv5s](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt)             | 640                   | 37.4                 | 56.8              | 98                                | 6.4                                | 0.9                             | 7.2                | 16.5                   |
      | [YOLOv5m](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m.pt)             | 640                   | 45.4                 | 64.1              | 224                               | 8.2                                | 1.7                             | 21.2               | 49.0                   |
      | [YOLOv5l](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5l.pt)             | 640                   | 49.0                 | 67.3              | 430                               | 10.1                               | 2.7                             | 46.5               | 109.1                  |
      | [YOLOv5x](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5x.pt)             | 640                   | 50.7                 | 68.9              | 766                               | 12.1                               | 4.8                             | 86.7               | 205.7                  |
      |                                                                                                |                       |                      |                   |                                   |                                    |                                 |                    |                        |
      | [YOLOv5n6](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n6.pt)           | 1280                  | 36.0                 | 54.4              | 153                               | 8.1                                | 2.1                             | 3.2                | 4.6                    |
      | [YOLOv5s6](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s6.pt)           | 1280                  | 44.8                 | 63.7              | 385                               | 8.2                                | 3.6                             | 12.6               | 16.8                   |
      | [YOLOv5m6](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m6.pt)           | 1280                  | 51.3                 | 69.3              | 887                               | 11.1                               | 6.8                             | 35.7               | 50.0                   |
      | [YOLOv5l6](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5l6.pt)           | 1280                  | 53.7                 | 71.3              | 1784                              | 15.8                               | 10.5                            | 76.8               | 111.4                  |
      | [YOLOv5x6](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5x6.pt)<br>+[TTA] | 1280<br>1536          | 55.0<br>**55.8**     | 72.7<br>**72.7**  | 3136<br>-                         | 26.2<br>-                          | 19.4<br>-                       | 140.7<br>-         | 209.8<br>-             |
      
      <details>
        <summary>笔记</summary>
      
      - 所有模型都使用默认配置,训练 300 epochs。n和s模型使用 [hyp.scratch-low.yaml](https://github.com/ultralytics/yolov5/blob/master/data/hyps/hyp.scratch-low.yaml) ,其他模型都使用 [hyp.scratch-high.yaml](https://github.com/ultralytics/yolov5/blob/master/data/hyps/hyp.scratch-high.yaml) 。
      - \*\*mAP<sup>val</sup>\*\*在单模型单尺度上计算,数据集使用 [COCO val2017](http://cocodataset.org) 。<br>复现命令 `python val.py --data coco.yaml --img 640 --conf 0.001 --iou 0.65`
      - **推理速度**在 COCO val 图像总体时间上进行平均得到,测试环境使用[AWS p3.2xlarge](https://aws.amazon.com/ec2/instance-types/p3/)实例。 NMS 时间 (大约 1 ms/img) 不包括在内。<br>复现命令 `python val.py --data coco.yaml --img 640 --task speed --batch 1`
      - **TTA** [测试时数据增强](https://docs.ultralytics.com/yolov5/tutorials/test_time_augmentation) 包括反射和尺度变换。<br>复现命令 `python val.py --data coco.yaml --img 1536 --iou 0.7 --augment`
      
      </details>
      
      ## <div align="center">实例分割模型 ⭐ 新</div>
      
      我们新的 YOLOv5 [release v7.0](https://github.com/ultralytics/yolov5/releases/v7.0) 实例分割模型是世界上最快和最准确的模型,击败所有当前 [SOTA 基准](https://paperswithcode.com/sota/real-time-instance-segmentation-on-mscoco)。我们使它非常易于训练、验证和部署。更多细节请查看 [发行说明](https://github.com/ultralytics/yolov5/releases/v7.0) 或访问我们的 [YOLOv5 分割 Colab 笔记本](https://github.com/ultralytics/yolov5/blob/master/segment/tutorial.ipynb) 以快速入门。
      
      <details>
        <summary>实例分割模型列表</summary>
      
      <br>
      
      <div align="center">
      <a align="center" href="https://ultralytics.com/yolov5" target="_blank">
      <img width="800" src="https://user-images.githubusercontent.com/61612323/204180385-84f3aca9-a5e9-43d8-a617-dda7ca12e54a.png"></a>
      </div>
      
      我们使用 A100 GPU 在 COCO 上以 640 图像大小训练了 300 epochs 得到 YOLOv5 分割模型。我们将所有模型导出到 ONNX FP32 以进行 CPU 速度测试,并导出到 TensorRT FP16 以进行 GPU 速度测试。为了便于再现,我们在 Google [Colab Pro](https://colab.research.google.com/signup) 上进行了所有速度测试。
      
      | 模型                                                                                       | 尺寸<br><sup>(像素) | mAP<sup>box<br>50-95 | mAP<sup>mask<br>50-95 | 训练时长<br><sup>300 epochs<br>A100 GPU(小时) | 推理速度<br><sup>ONNX CPU<br>(ms) | 推理速度<br><sup>TRT A100<br>(ms) | 参数量<br><sup>(M) | FLOPs<br><sup>@640 (B) |
      | ------------------------------------------------------------------------------------------ | --------------------- | -------------------- | --------------------- | ----------------------------------------------- | ----------------------------------- | ----------------------------------- | ------------------ | ---------------------- |
      | [YOLOv5n-seg](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n-seg.pt) | 640                   | 27.6                 | 23.4                  | 80:17                                           | **62.7**                            | **1.2**                             | **2.0**            | **7.1**                |
      | [YOLOv5s-seg](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s-seg.pt) | 640                   | 37.6                 | 31.7                  | 88:16                                           | 173.3                               | 1.4                                 | 7.6                | 26.4                   |
      | [YOLOv5m-seg](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m-seg.pt) | 640                   | 45.0                 | 37.1                  | 108:36                                          | 427.0                               | 2.2                                 | 22.0               | 70.8                   |
      | [YOLOv5l-seg](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5l-seg.pt) | 640                   | 49.0                 | 39.9                  | 66:43 (2x)                                      | 857.4                               | 2.9                                 | 47.9               | 147.7                  |
      | [YOLOv5x-seg](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5x-seg.pt) | 640                   | **50.7**             | **41.4**              | 62:56 (3x)                                      | 1579.2                              | 4.5                                 | 88.8               | 265.7                  |
      
      - 所有模型使用 SGD 优化器训练, 都使用 `lr0=0.01` 和 `weight_decay=5e-5` 参数, 图像大小为 640 。<br>训练 log 可以查看 https://wandb.ai/glenn-jocher/YOLOv5_v70_official
      - **准确性**结果都在 COCO 数据集上,使用单模型单尺度测试得到。<br>复现命令 `python segment/val.py --data coco.yaml --weights yolov5s-seg.pt`
      - **推理速度**是使用 100 张图像推理时间进行平均得到,测试环境使用 [Colab Pro](https://colab.research.google.com/signup) 上 A100 高 RAM 实例。结果仅表示推理速度(NMS 每张图像增加约 1 毫秒)。<br>复现命令 `python segment/val.py --data coco.yaml --weights yolov5s-seg.pt --batch 1`
      - **模型转换**到 FP32 的 ONNX 和 FP16 的 TensorRT 脚本为 `export.py`.<br>运行命令 `python export.py --weights yolov5s-seg.pt --include engine --device 0 --half`
      
      </details>
      
      <details>
        <summary>分割模型使用示例 &nbsp;<a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/segment/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a></summary>
      
      ### 训练
      
      YOLOv5分割训练支持自动下载 COCO128-seg 分割数据集,用户仅需在启动指令中包含 `--data coco128-seg.yaml` 参数。 若要手动下载,使用命令 `bash data/scripts/get_coco.sh --train --val --segments`, 在下载完毕后,使用命令 `python train.py --data coco.yaml` 开启训练。
      
      ```bash
      # 单 GPU
      python segment/train.py --data coco128-seg.yaml --weights yolov5s-seg.pt --img 640
      
      # 多 GPU, DDP 模式
      python -m torch.distributed.run --nproc_per_node 4 --master_port 1 segment/train.py --data coco128-seg.yaml --weights yolov5s-seg.pt --img 640 --device 0,1,2,3
      ```
      
      ### 验证
      
      在 COCO 数据集上验证 YOLOv5s-seg mask mAP:
      
      ```bash
      bash data/scripts/get_coco.sh --val --segments  # 下载 COCO val segments 数据集 (780MB, 5000 images)
      python segment/val.py --weights yolov5s-seg.pt --data coco.yaml --img 640  # 验证
      ```
      
      ### 预测
      
      使用预训练的 YOLOv5m-seg.pt 来预测 bus.jpg:
      
      ```bash
      python segment/predict.py --weights yolov5m-seg.pt --source data/images/bus.jpg
      ```
      
      ```python
      model = torch.hub.load(
          "ultralytics/yolov5", "custom", "yolov5m-seg.pt"
      )  # 从load from PyTorch Hub 加载模型 (WARNING: 推理暂未支持)
      ```
      
      | ![zidane](https://user-images.githubusercontent.com/26833433/203113421-decef4c4-183d-4a0a-a6c2-6435b33bc5d3.jpg) | ![bus](https://user-images.githubusercontent.com/26833433/203113416-11fe0025-69f7-4874-a0a6-65d0bfe2999a.jpg) |
      | ---------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------- |
      
      ### 模型导出
      
      将 YOLOv5s-seg 模型导出到 ONNX 和 TensorRT:
      
      ```bash
      python export.py --weights yolov5s-seg.pt --include onnx engine --img 640 --device 0
      ```
      
      </details>
      
      ## <div align="center">分类网络 ⭐ 新</div>
      
      YOLOv5 [release v6.2](https://github.com/ultralytics/yolov5/releases) 带来对分类模型训练、验证和部署的支持!详情请查看 [发行说明](https://github.com/ultralytics/yolov5/releases/v6.2) 或访问我们的 [YOLOv5 分类 Colab 笔记本](https://github.com/ultralytics/yolov5/blob/master/classify/tutorial.ipynb) 以快速入门。
      
      <details>
        <summary>分类网络模型</summary>
      
      <br>
      
      我们使用 4xA100 实例在 ImageNet 上训练了 90 个 epochs 得到 YOLOv5-cls 分类模型,我们训练了 ResNet 和 EfficientNet 模型以及相同的默认训练设置以进行比较。我们将所有模型导出到 ONNX FP32 以进行 CPU 速度测试,并导出到 TensorRT FP16 以进行 GPU 速度测试。为了便于重现,我们在 Google 上进行了所有速度测试 [Colab Pro](https://colab.research.google.com/signup) 。
      
      | 模型                                                                                               | 尺寸<br><sup>(像素) | acc<br><sup>top1 | acc<br><sup>top5 | 训练时长<br><sup>90 epochs<br>4xA100(小时) | 推理速度<br><sup>ONNX CPU<br>(ms) | 推理速度<br><sup>TensorRT V100<br>(ms) | 参数<br><sup>(M) | FLOPs<br><sup>@640 (B) |
      | -------------------------------------------------------------------------------------------------- | --------------------- | ---------------- | ---------------- | -------------------------------------------- | ----------------------------------- | ---------------------------------------- | ---------------- | ---------------------- |
      | [YOLOv5n-cls](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n-cls.pt)         | 224                   | 64.6             | 85.4             | 7:59                                         | **3.3**                             | **0.5**                                  | **2.5**          | **0.5**                |
      | [YOLOv5s-cls](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s-cls.pt)         | 224                   | 71.5             | 90.2             | 8:09                                         | 6.6                                 | 0.6                                      | 5.4              | 1.4                    |
      | [YOLOv5m-cls](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m-cls.pt)         | 224                   | 75.9             | 92.9             | 10:06                                        | 15.5                                | 0.9                                      | 12.9             | 3.9                    |
      | [YOLOv5l-cls](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5l-cls.pt)         | 224                   | 78.0             | 94.0             | 11:56                                        | 26.9                                | 1.4                                      | 26.5             | 8.5                    |
      | [YOLOv5x-cls](https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5x-cls.pt)         | 224                   | **79.0**         | **94.4**         | 15:04                                        | 54.3                                | 1.8                                      | 48.1             | 15.9                   |
      |                                                                                                    |                       |                  |                  |                                              |                                     |                                          |                  |                        |
      | [ResNet18](https://github.com/ultralytics/yolov5/releases/download/v7.0/resnet18.pt)               | 224                   | 70.3             | 89.5             | **6:47**                                     | 11.2                                | 0.5                                      | 11.7             | 3.7                    |
      | [Resnetzch](https://github.com/ultralytics/yolov5/releases/download/v7.0/resnet34.pt)              | 224                   | 73.9             | 91.8             | 8:33                                         | 20.6                                | 0.9                                      | 21.8             | 7.4                    |
      | [ResNet50](https://github.com/ultralytics/yolov5/releases/download/v7.0/resnet50.pt)               | 224                   | 76.8             | 93.4             | 11:10                                        | 23.4                                | 1.0                                      | 25.6             | 8.5                    |
      | [ResNet101](https://github.com/ultralytics/yolov5/releases/download/v7.0/resnet101.pt)             | 224                   | 78.5             | 94.3             | 17:10                                        | 42.1                                | 1.9                                      | 44.5             | 15.9                   |
      |                                                                                                    |                       |                  |                  |                                              |                                     |                                          |                  |                        |
      | [EfficientNet_b0](https://github.com/ultralytics/yolov5/releases/download/v7.0/efficientnet_b0.pt) | 224                   | 75.1             | 92.4             | 13:03                                        | 12.5                                | 1.3                                      | 5.3              | 1.0                    |
      | [EfficientNet_b1](https://github.com/ultralytics/yolov5/releases/download/v7.0/efficientnet_b1.pt) | 224                   | 76.4             | 93.2             | 17:04                                        | 14.9                                | 1.6                                      | 7.8              | 1.5                    |
      | [EfficientNet_b2](https://github.com/ultralytics/yolov5/releases/download/v7.0/efficientnet_b2.pt) | 224                   | 76.6             | 93.4             | 17:10                                        | 15.9                                | 1.6                                      | 9.1              | 1.7                    |
      | [EfficientNet_b3](https://github.com/ultralytics/yolov5/releases/download/v7.0/efficientnet_b3.pt) | 224                   | 77.7             | 94.0             | 19:19                                        | 18.9                                | 1.9                                      | 12.2             | 2.4                    |
      
      <details>
        <summary>Table Notes (点击以展开)</summary>
      
      - 所有模型都使用 SGD 优化器训练 90 个 epochs,都使用 `lr0=0.001` 和 `weight_decay=5e-5` 参数, 图像大小为 224 ,且都使用默认设置。<br>训练 log 可以查看 https://wandb.ai/glenn-jocher/YOLOv5-Classifier-v6-2
      - **准确性**都在单模型单尺度上计算,数据集使用 [ImageNet-1k](https://www.image-net.org/index.php) 。<br>复现命令 `python classify/val.py --data ../datasets/imagenet --img 224`
      - **推理速度**是使用 100 个推理图像进行平均得到,测试环境使用谷歌 [Colab Pro](https://colab.research.google.com/signup) V100 高 RAM 实例。<br>复现命令 `python classify/val.py --data ../datasets/imagenet --img 224 --batch 1`
      - **模型导出**到 FP32 的 ONNX 和 FP16 的 TensorRT 使用 `export.py` 。<br>复现命令 `python export.py --weights yolov5s-cls.pt --include engine onnx --imgsz 224`
        </details>
        </details>
      
      <details>
        <summary>分类训练示例 &nbsp;<a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/classify/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a></summary>
      
      ### 训练
      
      YOLOv5 分类训练支持自动下载 MNIST、Fashion-MNIST、CIFAR10、CIFAR100、Imagenette、Imagewoof 和 ImageNet 数据集,命令中使用 `--data` 即可。 MNIST 示例 `--data mnist` 。
      
      ```bash
      # 单 GPU
      python classify/train.py --model yolov5s-cls.pt --data cifar100 --epochs 5 --img 224 --batch 128
      
      # 多 GPU, DDP 模式
      python -m torch.distributed.run --nproc_per_node 4 --master_port 1 classify/train.py --model yolov5s-cls.pt --data imagenet --epochs 5 --img 224 --device 0,1,2,3
      ```
      
      ### 验证
      
      在 ImageNet-1k 数据集上验证 YOLOv5m-cls 的准确性:
      
      ```bash
      bash data/scripts/get_imagenet.sh --val  # download ImageNet val split (6.3G, 50000 images)
      python classify/val.py --weights yolov5m-cls.pt --data ../datasets/imagenet --img 224  # validate
      ```
      
      ### 预测
      
      使用预训练的 YOLOv5s-cls.pt 来预测 bus.jpg:
      
      ```bash
      python classify/predict.py --weights yolov5s-cls.pt --source data/images/bus.jpg
      ```
      
      ```python
      model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s-cls.pt")  # load from PyTorch Hub
      ```
      
      ### 模型导出
      
      将一组经过训练的 YOLOv5s-cls、ResNet 和 EfficientNet 模型导出到 ONNX 和 TensorRT:
      
      ```bash
      python export.py --weights yolov5s-cls.pt resnet50.pt efficientnet_b0.pt --include onnx engine --img 224
      ```
      
      </details>
      
      ## <div align="center">环境</div>
      
      使用下面我们经过验证的环境,在几秒钟内开始使用 YOLOv5 。单击下面的图标了解详细信息。
      
      <div align="center">
        <a href="https://bit.ly/yolov5-paperspace-notebook">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-gradient.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="5%" alt="" />
        <a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-colab-small.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="5%" alt="" />
        <a href="https://www.kaggle.com/ultralytics/yolov5">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-kaggle-small.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="5%" alt="" />
        <a href="https://hub.docker.com/r/ultralytics/yolov5">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-docker-small.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="5%" alt="" />
        <a href="https://docs.ultralytics.com/yolov5/environments/aws_quickstart_tutorial/">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-aws-small.png" width="10%" /></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="5%" alt="" />
        <a href="https://docs.ultralytics.com/yolov5/environments/google_cloud_quickstart_tutorial/">
          <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-gcp-small.png" width="10%" /></a>
      </div>
      
      ## <div align="center">贡献</div>
      
      我们喜欢您的意见或建议!我们希望尽可能简单和透明地为 YOLOv5 做出贡献。请看我们的 [投稿指南](https://docs.ultralytics.com/help/contributing/),并填写 [YOLOv5调查](https://ultralytics.com/survey?utm_source=github&utm_medium=social&utm_campaign=Survey) 向我们发送您的体验反馈。感谢我们所有的贡献者!
      
      <!-- SVG image from https://opencollective.com/ultralytics/contributors.svg?width=990 -->
      
      <a href="https://github.com/ultralytics/yolov5/graphs/contributors">
      <img src="https://github.com/ultralytics/assets/raw/main/im/image-contributors.png" /></a>
      
      ## <div align="center">许可证</div>
      
      Ultralytics 提供两种许可证选项以适应各种使用场景:
      
      - **AGPL-3.0 许可证**:这个[OSI 批准](https://opensource.org/licenses/)的开源许可证非常适合学生和爱好者,可以推动开放的协作和知识分享。请查看[LICENSE](https://github.com/ultralytics/yolov5/blob/master/LICENSE) 文件以了解更多细节。
      - **企业许可证**:专为商业用途设计,该许可证允许将 Ultralytics 的软件和 AI 模型无缝集成到商业产品和服务中,从而绕过 AGPL-3.0 的开源要求。如果您的场景涉及将我们的解决方案嵌入到商业产品中,请通过 [Ultralytics Licensing](https://ultralytics.com/license)与我们联系。
      
      ## <div align="center">联系方式</div>
      
      对于 Ultralytics 的错误报告和功能请求,请访问 [GitHub Issues](https://github.com/ultralytics/yolov5/issues),并加入我们的 [Discord](https://ultralytics.com/discord) 社区进行问题和讨论!
      
      <br>
      <div align="center">
        <a href="https://github.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-github.png" width="3%" alt="Ultralytics GitHub"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://www.linkedin.com/company/ultralytics/"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-linkedin.png" width="3%" alt="Ultralytics LinkedIn"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://twitter.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-twitter.png" width="3%" alt="Ultralytics Twitter"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://youtube.com/ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-youtube.png" width="3%" alt="Ultralytics YouTube"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://www.tiktok.com/@ultralytics"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-tiktok.png" width="3%" alt="Ultralytics TikTok"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://www.instagram.com/ultralytics/"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-instagram.png" width="3%" alt="Ultralytics Instagram"></a>
        <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="3%">
        <a href="https://ultralytics.com/discord"><img src="https://github.com/ultralytics/assets/raw/main/social/logo-social-discord.png" width="3%" alt="Ultralytics Discord"></a>
      </div>
      
      [tta]: https://docs.ultralytics.com/yolov5/tutorials/test_time_augmentation
      
    • 📄 YOLOv5_demo.py 1.8KB
      # 1.目标检测Demo
      # # 导入预测工具
      # import detect
      #
      # # 运行预测程序:指定要预测的图片,预训练好的模型及是否显示图片
      # detect.run(source="data/images/bus.jpg", weights='yolov5s.pt', view_img=True)
      # detect.run(source=0, weights='yolov5s.pt', device='cpu')
      
      # ------------------------------------------------------------------------------------------------------------------ #
      
      # # 2.网络结构
      # # 2.1 输入端
      # from utils.augmentations import letterbox
      # import matplotlib.pyplot as plt
      # import cv2
      
      # cat = plt.imread("data/images/bus.jpg")
      # print(cat.shape)
      # img = letterbox(cat, new_shape=(640, 640), auto=False, scaleFill=False, scaleup=True, stride=32)
      # print(img[0].shape)
      # plt.imshow(img[0])
      # plt.show()
      
      # ------------------------------------------------------------------------------------------------------------------ #
      
      # 3.模型训练
      ## 导入模型训练
      import train
      
      if __name__ == '__main__':
          # 模型训练:设置相应的参数:模型配置信息
          train.run(cfg='models/yolov5s.yaml',  # 模型结构的配置文件
                    data='data/coco128.yaml',  # 数据的配置文件
                    imgsz=640,  # 图像大小
                    batch_size=2,  # 批次大小
                    weights='yolov5s.pt',  # 预训练模型
                    epochs=1,  # 训练轮次
                    worker=1,  # 数据加载的线程数,根据硬件资源进行设置
                    lr=0.01
                    )
      
      # ------------------------------------------------------------------------------------------------------------------ #
      
      ## 4.模型预测
      ## 导入预测工具
      # import detect
      
      # # 运行预测程序:指定要预测的图片,预训练好的模型及是否显示图片
      # detect.run(source="data/images/bus.jpg", weights='runs/train/exp/weights/best.pt', view_img=True)
      
    • 📄 benchmarks.py 7.9KB
      # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
      """
      Run YOLOv5 benchmarks on all supported export formats.
      
      Format                      | `export.py --include`         | Model
      ---                         | ---                           | ---
      PyTorch                     | -                             | yolov5s.pt
      TorchScript                 | `torchscript`                 | yolov5s.torchscript
      ONNX                        | `onnx`                        | yolov5s.onnx
      OpenVINO                    | `openvino`                    | yolov5s_openvino_model/
      TensorRT                    | `engine`                      | yolov5s.engine
      CoreML                      | `coreml`                      | yolov5s.mlmodel
      TensorFlow SavedModel       | `saved_model`                 | yolov5s_saved_model/
      TensorFlow GraphDef         | `pb`                          | yolov5s.pb
      TensorFlow Lite             | `tflite`                      | yolov5s.tflite
      TensorFlow Edge TPU         | `edgetpu`                     | yolov5s_edgetpu.tflite
      TensorFlow.js               | `tfjs`                        | yolov5s_web_model/
      
      Requirements:
          $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu  # CPU
          $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow  # GPU
          $ pip install -U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com  # TensorRT
      
      Usage:
          $ python benchmarks.py --weights yolov5s.pt --img 640
      """
      
      import argparse
      import platform
      import sys
      import time
      from pathlib import Path
      
      import pandas as pd
      
      FILE = Path(__file__).resolve()
      ROOT = FILE.parents[0]  # YOLOv5 root directory
      if str(ROOT) not in sys.path:
          sys.path.append(str(ROOT))  # add ROOT to PATH
      # ROOT = ROOT.relative_to(Path.cwd())  # relative
      
      import export
      from models.experimental import attempt_load
      from models.yolo import SegmentationModel
      from segment.val import run as val_seg
      from utils import notebook_init
      from utils.general import LOGGER, check_yaml, file_size, print_args
      from utils.torch_utils import select_device
      from val import run as val_det
      
      
      def run(
          weights=ROOT / "yolov5s.pt",  # weights path
          imgsz=640,  # inference size (pixels)
          batch_size=1,  # batch size
          data=ROOT / "data/coco128.yaml",  # dataset.yaml path
          device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
          half=False,  # use FP16 half-precision inference
          test=False,  # test exports only
          pt_only=False,  # test PyTorch only
          hard_fail=False,  # throw error on benchmark failure
      ):
          y, t = [], time.time()
          device = select_device(device)
          model_type = type(attempt_load(weights, fuse=False))  # DetectionModel, SegmentationModel, etc.
          for i, (name, f, suffix, cpu, gpu) in export.export_formats().iterrows():  # index, (name, file, suffix, CPU, GPU)
              try:
                  assert i not in (9, 10), "inference not supported"  # Edge TPU and TF.js are unsupported
                  assert i != 5 or platform.system() == "Darwin", "inference only supported on macOS>=10.13"  # CoreML
                  if "cpu" in device.type:
                      assert cpu, "inference not supported on CPU"
                  if "cuda" in device.type:
                      assert gpu, "inference not supported on GPU"
      
                  # Export
                  if f == "-":
                      w = weights  # PyTorch format
                  else:
                      w = export.run(
                          weights=weights, imgsz=[imgsz], include=[f], batch_size=batch_size, device=device, half=half
                      )[-1]  # all others
                  assert suffix in str(w), "export failed"
      
                  # Validate
                  if model_type == SegmentationModel:
                      result = val_seg(data, w, batch_size, imgsz, plots=False, device=device, task="speed", half=half)
                      metric = result[0][7]  # (box(p, r, map50, map), mask(p, r, map50, map), *loss(box, obj, cls))
                  else:  # DetectionModel:
                      result = val_det(data, w, batch_size, imgsz, plots=False, device=device, task="speed", half=half)
                      metric = result[0][3]  # (p, r, map50, map, *loss(box, obj, cls))
                  speed = result[2][1]  # times (preprocess, inference, postprocess)
                  y.append([name, round(file_size(w), 1), round(metric, 4), round(speed, 2)])  # MB, mAP, t_inference
              except Exception as e:
                  if hard_fail:
                      assert type(e) is AssertionError, f"Benchmark --hard-fail for {name}: {e}"
                  LOGGER.warning(f"WARNING ⚠️ Benchmark failure for {name}: {e}")
                  y.append([name, None, None, None])  # mAP, t_inference
              if pt_only and i == 0:
                  break  # break after PyTorch
      
          # Print results
          LOGGER.info("\n")
          parse_opt()
          notebook_init()  # print system info
          c = ["Format", "Size (MB)", "mAP50-95", "Inference time (ms)"] if map else ["Format", "Export", "", ""]
          py = pd.DataFrame(y, columns=c)
          LOGGER.info(f"\nBenchmarks complete ({time.time() - t:.2f}s)")
          LOGGER.info(str(py if map else py.iloc[:, :2]))
          if hard_fail and isinstance(hard_fail, str):
              metrics = py["mAP50-95"].array  # values to compare to floor
              floor = eval(hard_fail)  # minimum metric floor to pass, i.e. = 0.29 mAP for YOLOv5n
              assert all(x > floor for x in metrics if pd.notna(x)), f"HARD FAIL: mAP50-95 < floor {floor}"
          return py
      
      
      def test(
          weights=ROOT / "yolov5s.pt",  # weights path
          imgsz=640,  # inference size (pixels)
          batch_size=1,  # batch size
          data=ROOT / "data/coco128.yaml",  # dataset.yaml path
          device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
          half=False,  # use FP16 half-precision inference
          test=False,  # test exports only
          pt_only=False,  # test PyTorch only
          hard_fail=False,  # throw error on benchmark failure
      ):
          y, t = [], time.time()
          device = select_device(device)
          for i, (name, f, suffix, gpu) in export.export_formats().iterrows():  # index, (name, file, suffix, gpu-capable)
              try:
                  w = (
                      weights
                      if f == "-"
                      else export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1]
                  )  # weights
                  assert suffix in str(w), "export failed"
                  y.append([name, True])
              except Exception:
                  y.append([name, False])  # mAP, t_inference
      
          # Print results
          LOGGER.info("\n")
          parse_opt()
          notebook_init()  # print system info
          py = pd.DataFrame(y, columns=["Format", "Export"])
          LOGGER.info(f"\nExports complete ({time.time() - t:.2f}s)")
          LOGGER.info(str(py))
          return py
      
      
      def parse_opt():
          """Parses command-line arguments for YOLOv5 model inference configuration."""
          parser = argparse.ArgumentParser()
          parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="weights path")
          parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="inference size (pixels)")
          parser.add_argument("--batch-size", type=int, default=1, help="batch size")
          parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
          parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
          parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
          parser.add_argument("--test", action="store_true", help="test exports only")
          parser.add_argument("--pt-only", action="store_true", help="test PyTorch only")
          parser.add_argument("--hard-fail", nargs="?", const=True, default=False, help="Exception on error or < min metric")
          opt = parser.parse_args()
          opt.data = check_yaml(opt.data)  # check YAML
          print_args(vars(opt))
          return opt
      
      
      def main(opt):
          """Executes a test run if `opt.test` is True, otherwise starts training or inference with provided options."""
          test(**vars(opt)) if opt.test else run(**vars(opt))
      
      
      if __name__ == "__main__":
          opt = parse_opt()
          main(opt)
      
    • 📄 detect.py 15.8KB
      # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
      """
      Run YOLOv5 detection inference on images, videos, directories, globs, YouTube, webcam, streams, etc.
      
      Usage - sources:
          $ python detect.py --weights yolov5s.pt --source 0                               # webcam
                                                           img.jpg                         # image
                                                           vid.mp4                         # video
                                                           screen                          # screenshot
                                                           path/                           # directory
                                                           list.txt                        # list of images
                                                           list.streams                    # list of streams
                                                           'path/*.jpg'                    # glob
                                                           'https://youtu.be/LNwODJXcvt4'  # YouTube
                                                           'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
      
      Usage - formats:
          $ python detect.py --weights yolov5s.pt                 # PyTorch
                                       yolov5s.torchscript        # TorchScript
                                       yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                       yolov5s_openvino_model     # OpenVINO
                                       yolov5s.engine             # TensorRT
                                       yolov5s.mlmodel            # CoreML (macOS-only)
                                       yolov5s_saved_model        # TensorFlow SavedModel
                                       yolov5s.pb                 # TensorFlow GraphDef
                                       yolov5s.tflite             # TensorFlow Lite
                                       yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
                                       yolov5s_paddle_model       # PaddlePaddle
      """
      
      import argparse
      import csv
      import os
      import platform
      import sys
      from pathlib import Path
      
      import torch
      
      FILE = Path(__file__).resolve()
      ROOT = FILE.parents[0]  # YOLOv5 root directory
      if str(ROOT) not in sys.path:
          sys.path.append(str(ROOT))  # add ROOT to PATH
      ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
      
      from ultralytics.utils.plotting import Annotator, colors, save_one_box
      
      from models.common import DetectMultiBackend
      from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
      from utils.general import (
          LOGGER,
          Profile,
          check_file,
          check_img_size,
          check_imshow,
          check_requirements,
          colorstr,
          cv2,
          increment_path,
          non_max_suppression,
          print_args,
          scale_boxes,
          strip_optimizer,
          xyxy2xywh,
      )
      from utils.torch_utils import select_device, smart_inference_mode
      
      
      @smart_inference_mode()
      def run(
          weights=ROOT / "yolov5s.pt",  # model path or triton URL
          source=ROOT / "data/images",  # file/dir/URL/glob/screen/0(webcam)
          data=ROOT / "data/coco128.yaml",  # dataset.yaml path
          imgsz=(640, 640),  # inference size (height, width)
          conf_thres=0.25,  # confidence threshold
          iou_thres=0.45,  # NMS IOU threshold
          max_det=1000,  # maximum detections per image
          device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
          view_img=False,  # show results
          save_txt=False,  # save results to *.txt
          save_csv=False,  # save results in CSV format
          save_conf=False,  # save confidences in --save-txt labels
          save_crop=False,  # save cropped prediction boxes
          nosave=False,  # do not save images/videos
          classes=None,  # filter by class: --class 0, or --class 0 2 3
          agnostic_nms=False,  # class-agnostic NMS
          augment=False,  # augmented inference
          visualize=False,  # visualize features
          update=False,  # update all models
          project=ROOT / "runs/detect",  # save results to project/name
          name="exp",  # save results to project/name
          exist_ok=False,  # existing project/name ok, do not increment
          line_thickness=3,  # bounding box thickness (pixels)
          hide_labels=False,  # hide labels
          hide_conf=False,  # hide confidences
          half=False,  # use FP16 half-precision inference
          dnn=False,  # use OpenCV DNN for ONNX inference
          vid_stride=1,  # video frame-rate stride
      ):
          source = str(source)
          save_img = not nosave and not source.endswith(".txt")  # save inference images
          is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
          is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
          webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
          screenshot = source.lower().startswith("screen")
          if is_url and is_file:
              source = check_file(source)  # download
      
          # Directories
          save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
          (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
      
          # Load model
          device = select_device(device)
          model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
          stride, names, pt = model.stride, model.names, model.pt
          imgsz = check_img_size(imgsz, s=stride)  # check image size
      
          # Dataloader
          bs = 1  # batch_size
          if webcam:
              view_img = check_imshow(warn=True)
              dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
              bs = len(dataset)
          elif screenshot:
              dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
          else:
              dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
          vid_path, vid_writer = [None] * bs, [None] * bs
      
          # Run inference
          model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
          seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))
          for path, im, im0s, vid_cap, s in dataset:
              with dt[0]:
                  im = torch.from_numpy(im).to(model.device)
                  im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
                  im /= 255  # 0 - 255 to 0.0 - 1.0
                  if len(im.shape) == 3:
                      im = im[None]  # expand for batch dim
                  if model.xml and im.shape[0] > 1:
                      ims = torch.chunk(im, im.shape[0], 0)
      
              # Inference
              with dt[1]:
                  visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
                  if model.xml and im.shape[0] > 1:
                      pred = None
                      for image in ims:
                          if pred is None:
                              pred = model(image, augment=augment, visualize=visualize).unsqueeze(0)
                          else:
                              pred = torch.cat((pred, model(image, augment=augment, visualize=visualize).unsqueeze(0)), dim=0)
                      pred = [pred, None]
                  else:
                      pred = model(im, augment=augment, visualize=visualize)
              # NMS
              with dt[2]:
                  pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
      
              # Second-stage classifier (optional)
              # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
      
              # Define the path for the CSV file
              csv_path = save_dir / "predictions.csv"
      
              # Create or append to the CSV file
              def write_to_csv(image_name, prediction, confidence):
                  """Writes prediction data for an image to a CSV file, appending if the file exists."""
                  data = {"Image Name": image_name, "Prediction": prediction, "Confidence": confidence}
                  with open(csv_path, mode="a", newline="") as f:
                      writer = csv.DictWriter(f, fieldnames=data.keys())
                      if not csv_path.is_file():
                          writer.writeheader()
                      writer.writerow(data)
      
              # Process predictions
              for i, det in enumerate(pred):  # per image
                  seen += 1
                  if webcam:  # batch_size >= 1
                      p, im0, frame = path[i], im0s[i].copy(), dataset.count
                      s += f"{i}: "
                  else:
                      p, im0, frame = path, im0s.copy(), getattr(dataset, "frame", 0)
      
                  p = Path(p)  # to Path
                  save_path = str(save_dir / p.name)  # im.jpg
                  txt_path = str(save_dir / "labels" / p.stem) + ("" if dataset.mode == "image" else f"_{frame}")  # im.txt
                  s += "%gx%g " % im.shape[2:]  # print string
                  gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
                  imc = im0.copy() if save_crop else im0  # for save_crop
                  annotator = Annotator(im0, line_width=line_thickness, example=str(names))
                  if len(det):
                      # Rescale boxes from img_size to im0 size
                      det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
      
                      # Print results
                      for c in det[:, 5].unique():
                          n = (det[:, 5] == c).sum()  # detections per class
                          s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
      
                      # Write results
                      for *xyxy, conf, cls in reversed(det):
                          c = int(cls)  # integer class
                          label = names[c] if hide_conf else f"{names[c]}"
                          confidence = float(conf)
                          confidence_str = f"{confidence:.2f}"
      
                          if save_csv:
                              write_to_csv(p.name, label, confidence_str)
      
                          if save_txt:  # Write to file
                              xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                              line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
                              with open(f"{txt_path}.txt", "a") as f:
                                  f.write(("%g " * len(line)).rstrip() % line + "\n")
      
                          if save_img or save_crop or view_img:  # Add bbox to image
                              c = int(cls)  # integer class
                              label = None if hide_labels else (names[c] if hide_conf else f"{names[c]} {conf:.2f}")
                              annotator.box_label(xyxy, label, color=colors(c, True))
                          if save_crop:
                              save_one_box(xyxy, imc, file=save_dir / "crops" / names[c] / f"{p.stem}.jpg", BGR=True)
      
                  # Stream results
                  im0 = annotator.result()
                  if view_img:
                      if platform.system() == "Linux" and p not in windows:
                          windows.append(p)
                          cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
                          cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
                      cv2.imshow(str(p), im0)
                      cv2.waitKey(1)  # 1 millisecond
      
                  # Save results (image with detections)
                  if save_img:
                      if dataset.mode == "image":
                          cv2.imwrite(save_path, im0)
                      else:  # 'video' or 'stream'
                          if vid_path[i] != save_path:  # new video
                              vid_path[i] = save_path
                              if isinstance(vid_writer[i], cv2.VideoWriter):
                                  vid_writer[i].release()  # release previous video writer
                              if vid_cap:  # video
                                  fps = vid_cap.get(cv2.CAP_PROP_FPS)
                                  w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                                  h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                              else:  # stream
                                  fps, w, h = 30, im0.shape[1], im0.shape[0]
                              save_path = str(Path(save_path).with_suffix(".mp4"))  # force *.mp4 suffix on results videos
                              vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
                          vid_writer[i].write(im0)
      
              # Print time (inference-only)
              LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
      
          # Print results
          t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
          LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}" % t)
          if save_txt or save_img:
              s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
              LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
          if update:
              strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
      
      
      def parse_opt():
          """Parses command-line arguments for YOLOv5 detection, setting inference options and model configurations."""
          parser = argparse.ArgumentParser()
          parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model path or triton URL")
          parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
          parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="(optional) dataset.yaml path")
          parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640], help="inference size h,w")
          parser.add_argument("--conf-thres", type=float, default=0.25, help="confidence threshold")
          parser.add_argument("--iou-thres", type=float, default=0.45, help="NMS IoU threshold")
          parser.add_argument("--max-det", type=int, default=1000, help="maximum detections per image")
          parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
          parser.add_argument("--view-img", action="store_true", help="show results")
          parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
          parser.add_argument("--save-csv", action="store_true", help="save results in CSV format")
          parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels")
          parser.add_argument("--save-crop", action="store_true", help="save cropped prediction boxes")
          parser.add_argument("--nosave", action="store_true", help="do not save images/videos")
          parser.add_argument("--classes", nargs="+", type=int, help="filter by class: --classes 0, or --classes 0 2 3")
          parser.add_argument("--agnostic-nms", action="store_true", help="class-agnostic NMS")
          parser.add_argument("--augment", action="store_true", help="augmented inference")
          parser.add_argument("--visualize", action="store_true", help="visualize features")
          parser.add_argument("--update", action="store_true", help="update all models")
          parser.add_argument("--project", default=ROOT / "runs/detect", help="save results to project/name")
          parser.add_argument("--name", default="exp", help="save results to project/name")
          parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
          parser.add_argument("--line-thickness", default=3, type=int, help="bounding box thickness (pixels)")
          parser.add_argument("--hide-labels", default=False, action="store_true", help="hide labels")
          parser.add_argument("--hide-conf", default=False, action="store_true", help="hide confidences")
          parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
          parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
          parser.add_argument("--vid-stride", type=int, default=1, help="video frame-rate stride")
          opt = parser.parse_args()
          opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
          print_args(vars(opt))
          return opt
      
      
      def main(opt):
          """Executes YOLOv5 model inference with given options, checking requirements before running the model."""
          check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
          run(**vars(opt))
      
      
      if __name__ == "__main__":
          opt = parse_opt()
          main(opt)
      
    • 📄 export.py 41.6KB
      # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
      """
      Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit
      
      Format                      | `export.py --include`         | Model
      ---                         | ---                           | ---
      PyTorch                     | -                             | yolov5s.pt
      TorchScript                 | `torchscript`                 | yolov5s.torchscript
      ONNX                        | `onnx`                        | yolov5s.onnx
      OpenVINO                    | `openvino`                    | yolov5s_openvino_model/
      TensorRT                    | `engine`                      | yolov5s.engine
      CoreML                      | `coreml`                      | yolov5s.mlmodel
      TensorFlow SavedModel       | `saved_model`                 | yolov5s_saved_model/
      TensorFlow GraphDef         | `pb`                          | yolov5s.pb
      TensorFlow Lite             | `tflite`                      | yolov5s.tflite
      TensorFlow Edge TPU         | `edgetpu`                     | yolov5s_edgetpu.tflite
      TensorFlow.js               | `tfjs`                        | yolov5s_web_model/
      PaddlePaddle                | `paddle`                      | yolov5s_paddle_model/
      
      Requirements:
          $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu  # CPU
          $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow  # GPU
      
      Usage:
          $ python export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...
      
      Inference:
          $ python detect.py --weights yolov5s.pt                 # PyTorch
                                       yolov5s.torchscript        # TorchScript
                                       yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                       yolov5s_openvino_model     # OpenVINO
                                       yolov5s.engine             # TensorRT
                                       yolov5s.mlmodel            # CoreML (macOS-only)
                                       yolov5s_saved_model        # TensorFlow SavedModel
                                       yolov5s.pb                 # TensorFlow GraphDef
                                       yolov5s.tflite             # TensorFlow Lite
                                       yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
                                       yolov5s_paddle_model       # PaddlePaddle
      
      TensorFlow.js:
          $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
          $ npm install
          $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
          $ npm start
      """
      
      import argparse
      import contextlib
      import json
      import os
      import platform
      import re
      import subprocess
      import sys
      import time
      import warnings
      from pathlib import Path
      
      import pandas as pd
      import torch
      from torch.utils.mobile_optimizer import optimize_for_mobile
      
      FILE = Path(__file__).resolve()
      ROOT = FILE.parents[0]  # YOLOv5 root directory
      if str(ROOT) not in sys.path:
          sys.path.append(str(ROOT))  # add ROOT to PATH
      if platform.system() != "Windows":
          ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
      
      from models.experimental import attempt_load
      from models.yolo import ClassificationModel, Detect, DetectionModel, SegmentationModel
      from utils.dataloaders import LoadImages
      from utils.general import (
          LOGGER,
          Profile,
          check_dataset,
          check_img_size,
          check_requirements,
          check_version,
          check_yaml,
          colorstr,
          file_size,
          get_default_args,
          print_args,
          url2file,
          yaml_save,
      )
      from utils.torch_utils import select_device, smart_inference_mode
      
      MACOS = platform.system() == "Darwin"  # macOS environment
      
      
      class iOSModel(torch.nn.Module):
          def __init__(self, model, im):
              """Initializes an iOS compatible model with normalization based on image dimensions."""
              super().__init__()
              b, c, h, w = im.shape  # batch, channel, height, width
              self.model = model
              self.nc = model.nc  # number of classes
              if w == h:
                  self.normalize = 1.0 / w
              else:
                  self.normalize = torch.tensor([1.0 / w, 1.0 / h, 1.0 / w, 1.0 / h])  # broadcast (slower, smaller)
                  # np = model(im)[0].shape[1]  # number of points
                  # self.normalize = torch.tensor([1. / w, 1. / h, 1. / w, 1. / h]).expand(np, 4)  # explicit (faster, larger)
      
          def forward(self, x):
              """Runs forward pass on the input tensor, returning class confidences and normalized coordinates."""
              xywh, conf, cls = self.model(x)[0].squeeze().split((4, 1, self.nc), 1)
              return cls * conf, xywh * self.normalize  # confidence (3780, 80), coordinates (3780, 4)
      
      
      def export_formats():
          """Returns a DataFrame of supported YOLOv5 model export formats and their properties."""
          x = [
              ["PyTorch", "-", ".pt", True, True],
              ["TorchScript", "torchscript", ".torchscript", True, True],
              ["ONNX", "onnx", ".onnx", True, True],
              ["OpenVINO", "openvino", "_openvino_model", True, False],
              ["TensorRT", "engine", ".engine", False, True],
              ["CoreML", "coreml", ".mlmodel", True, False],
              ["TensorFlow SavedModel", "saved_model", "_saved_model", True, True],
              ["TensorFlow GraphDef", "pb", ".pb", True, True],
              ["TensorFlow Lite", "tflite", ".tflite", True, False],
              ["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", False, False],
              ["TensorFlow.js", "tfjs", "_web_model", False, False],
              ["PaddlePaddle", "paddle", "_paddle_model", True, True],
          ]
          return pd.DataFrame(x, columns=["Format", "Argument", "Suffix", "CPU", "GPU"])
      
      
      def try_export(inner_func):
          """Decorator @try_export for YOLOv5 model export functions that logs success/failure, time taken, and file size."""
          inner_args = get_default_args(inner_func)
      
          def outer_func(*args, **kwargs):
              prefix = inner_args["prefix"]
              try:
                  with Profile() as dt:
                      f, model = inner_func(*args, **kwargs)
                  LOGGER.info(f"{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)")
                  return f, model
              except Exception as e:
                  LOGGER.info(f"{prefix} export failure ❌ {dt.t:.1f}s: {e}")
                  return None, None
      
          return outer_func
      
      
      @try_export
      def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")):
          """Exports YOLOv5 model to TorchScript format, optionally optimized for mobile, with image shape and stride
          metadata.
          """
          LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...")
          f = file.with_suffix(".torchscript")
      
          ts = torch.jit.trace(model, im, strict=False)
          d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
          extra_files = {"config.txt": json.dumps(d)}  # torch._C.ExtraFilesMap()
          if optimize:  # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
              optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
          else:
              ts.save(str(f), _extra_files=extra_files)
          return f, None
      
      
      @try_export
      def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX:")):
          """Exports a YOLOv5 model to ONNX format with dynamic axes and optional simplification."""
          check_requirements("onnx>=1.12.0")
          import onnx
      
          LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__}...")
          f = str(file.with_suffix(".onnx"))
      
          output_names = ["output0", "output1"] if isinstance(model, SegmentationModel) else ["output0"]
          if dynamic:
              dynamic = {"images": {0: "batch", 2: "height", 3: "width"}}  # shape(1,3,640,640)
              if isinstance(model, SegmentationModel):
                  dynamic["output0"] = {0: "batch", 1: "anchors"}  # shape(1,25200,85)
                  dynamic["output1"] = {0: "batch", 2: "mask_height", 3: "mask_width"}  # shape(1,32,160,160)
              elif isinstance(model, DetectionModel):
                  dynamic["output0"] = {0: "batch", 1: "anchors"}  # shape(1,25200,85)
      
          torch.onnx.export(
              model.cpu() if dynamic else model,  # --dynamic only compatible with cpu
              im.cpu() if dynamic else im,
              f,
              verbose=False,
              opset_version=opset,
              do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
              input_names=["images"],
              output_names=output_names,
              dynamic_axes=dynamic or None,
          )
      
          # Checks
          model_onnx = onnx.load(f)  # load onnx model
          onnx.checker.check_model(model_onnx)  # check onnx model
      
          # Metadata
          d = {"stride": int(max(model.stride)), "names": model.names}
          for k, v in d.items():
              meta = model_onnx.metadata_props.add()
              meta.key, meta.value = k, str(v)
          onnx.save(model_onnx, f)
      
          # Simplify
          if simplify:
              try:
                  cuda = torch.cuda.is_available()
                  check_requirements(("onnxruntime-gpu" if cuda else "onnxruntime", "onnx-simplifier>=0.4.1"))
                  import onnxsim
      
                  LOGGER.info(f"{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...")
                  model_onnx, check = onnxsim.simplify(model_onnx)
                  assert check, "assert check failed"
                  onnx.save(model_onnx, f)
              except Exception as e:
                  LOGGER.info(f"{prefix} simplifier failure: {e}")
          return f, model_onnx
      
      
      @try_export
      def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")):
          # YOLOv5 OpenVINO export
          check_requirements("openvino-dev>=2023.0")  # requires openvino-dev: https://pypi.org/project/openvino-dev/
          import openvino.runtime as ov  # noqa
          from openvino.tools import mo  # noqa
      
          LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...")
          f = str(file).replace(file.suffix, f"_{'int8_' if int8 else ''}openvino_model{os.sep}")
          f_onnx = file.with_suffix(".onnx")
          f_ov = str(Path(f) / file.with_suffix(".xml").name)
      
          ov_model = mo.convert_model(f_onnx, model_name=file.stem, framework="onnx", compress_to_fp16=half)  # export
      
          if int8:
              check_requirements("nncf>=2.5.0")  # requires at least version 2.5.0 to use the post-training quantization
              import nncf
              import numpy as np
      
              from utils.dataloaders import create_dataloader
      
              def gen_dataloader(yaml_path, task="train", imgsz=640, workers=4):
                  data_yaml = check_yaml(yaml_path)
                  data = check_dataset(data_yaml)
                  dataloader = create_dataloader(
                      data[task], imgsz=imgsz, batch_size=1, stride=32, pad=0.5, single_cls=False, rect=False, workers=workers
                  )[0]
                  return dataloader
      
              # noqa: F811
      
              def transform_fn(data_item):
                  """
                  Quantization transform function.
      
                  Extracts and preprocess input data from dataloader item for quantization.
                  Parameters:
                     data_item: Tuple with data item produced by DataLoader during iteration
                  Returns:
                      input_tensor: Input data for quantization
                  """
                  assert data_item[0].dtype == torch.uint8, "input image must be uint8 for the quantization preprocessing"
      
                  img = data_item[0].numpy().astype(np.float32)  # uint8 to fp16/32
                  img /= 255.0  # 0 - 255 to 0.0 - 1.0
                  return np.expand_dims(img, 0) if img.ndim == 3 else img
      
              ds = gen_dataloader(data)
              quantization_dataset = nncf.Dataset(ds, transform_fn)
              ov_model = nncf.quantize(ov_model, quantization_dataset, preset=nncf.QuantizationPreset.MIXED)
      
          ov.serialize(ov_model, f_ov)  # save
          yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata)  # add metadata.yaml
          return f, None
      
      
      @try_export
      def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
          """Exports a YOLOv5 model to PaddlePaddle format using X2Paddle, saving to `save_dir` and adding a metadata.yaml
          file.
          """
          check_requirements(("paddlepaddle", "x2paddle"))
          import x2paddle
          from x2paddle.convert import pytorch2paddle
      
          LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...")
          f = str(file).replace(".pt", f"_paddle_model{os.sep}")
      
          pytorch2paddle(module=model, save_dir=f, jit_type="trace", input_examples=[im])  # export
          yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata)  # add metadata.yaml
          return f, None
      
      
      @try_export
      def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")):
          """Exports YOLOv5 model to CoreML format with optional NMS, INT8, and FP16 support; requires coremltools."""
          check_requirements("coremltools")
          import coremltools as ct
      
          LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
          f = file.with_suffix(".mlmodel")
      
          if nms:
              model = iOSModel(model, im)
          ts = torch.jit.trace(model, im, strict=False)  # TorchScript model
          ct_model = ct.convert(ts, inputs=[ct.ImageType("image", shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
          bits, mode = (8, "kmeans_lut") if int8 else (16, "linear") if half else (32, None)
          if bits < 32:
              if MACOS:  # quantization only supported on macOS
                  with warnings.catch_warnings():
                      warnings.filterwarnings("ignore", category=DeprecationWarning)  # suppress numpy==1.20 float warning
                      ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
              else:
                  print(f"{prefix} quantization only supported on macOS, skipping...")
          ct_model.save(f)
          return f, ct_model
      
      
      @try_export
      def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr("TensorRT:")):
          """
          Exports a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0.
      
          https://developer.nvidia.com/tensorrt
          """
          assert im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. `python export.py --device 0`"
          try:
              import tensorrt as trt
          except Exception:
              if platform.system() == "Linux":
                  check_requirements("nvidia-tensorrt", cmds="-U --index-url https://pypi.ngc.nvidia.com")
              import tensorrt as trt
      
          if trt.__version__[0] == "7":  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
              grid = model.model[-1].anchor_grid
              model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
              export_onnx(model, im, file, 12, dynamic, simplify)  # opset 12
              model.model[-1].anchor_grid = grid
          else:  # TensorRT >= 8
              check_version(trt.__version__, "8.0.0", hard=True)  # require tensorrt>=8.0.0
              export_onnx(model, im, file, 12, dynamic, simplify)  # opset 12
          onnx = file.with_suffix(".onnx")
      
          LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...")
          is_trt10 = int(trt.__version__.split(".")[0]) >= 10  # is TensorRT >= 10
          assert onnx.exists(), f"failed to export ONNX file: {onnx}"
          f = file.with_suffix(".engine")  # TensorRT engine file
          logger = trt.Logger(trt.Logger.INFO)
          if verbose:
              logger.min_severity = trt.Logger.Severity.VERBOSE
      
          builder = trt.Builder(logger)
          config = builder.create_builder_config()
          if is_trt10:
              config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30)
          else:  # TensorRT versions 7, 8
              config.max_workspace_size = workspace * 1 << 30
          flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
          network = builder.create_network(flag)
          parser = trt.OnnxParser(network, logger)
          if not parser.parse_from_file(str(onnx)):
              raise RuntimeError(f"failed to load ONNX file: {onnx}")
      
          inputs = [network.get_input(i) for i in range(network.num_inputs)]
          outputs = [network.get_output(i) for i in range(network.num_outputs)]
          for inp in inputs:
              LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}')
          for out in outputs:
              LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}')
      
          if dynamic:
              if im.shape[0] <= 1:
                  LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument")
              profile = builder.create_optimization_profile()
              for inp in inputs:
                  profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)
              config.add_optimization_profile(profile)
      
          LOGGER.info(f"{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}")
          if builder.platform_has_fast_fp16 and half:
              config.set_flag(trt.BuilderFlag.FP16)
      
          build = builder.build_serialized_network if is_trt10 else builder.build_engine
          with build(network, config) as engine, open(f, "wb") as t:
              t.write(engine if is_trt10 else engine.serialize())
          return f, None
      
      
      @try_export
      def export_saved_model(
          model,
          im,
          file,
          dynamic,
          tf_nms=False,
          agnostic_nms=False,
          topk_per_class=100,
          topk_all=100,
          iou_thres=0.45,
          conf_thres=0.25,
          keras=False,
          prefix=colorstr("TensorFlow SavedModel:"),
      ):
          # YOLOv5 TensorFlow SavedModel export
          try:
              import tensorflow as tf
          except Exception:
              check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}<=2.15.1")
      
              import tensorflow as tf
          from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
      
          from models.tf import TFModel
      
          LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
          if tf.__version__ > "2.13.1":
              helper_url = "https://github.com/ultralytics/yolov5/issues/12489"
              LOGGER.info(
                  f"WARNING ⚠️ using Tensorflow {tf.__version__} > 2.13.1 might cause issue when exporting the model to tflite {helper_url}"
              )  # handling issue https://github.com/ultralytics/yolov5/issues/12489
          f = str(file).replace(".pt", "_saved_model")
          batch_size, ch, *imgsz = list(im.shape)  # BCHW
      
          tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
          im = tf.zeros((batch_size, *imgsz, ch))  # BHWC order for TensorFlow
          _ = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
          inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if dynamic else batch_size)
          outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
          keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)
          keras_model.trainable = False
          keras_model.summary()
          if keras:
              keras_model.save(f, save_format="tf")
          else:
              spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)
              m = tf.function(lambda x: keras_model(x))  # full model
              m = m.get_concrete_function(spec)
              frozen_func = convert_variables_to_constants_v2(m)
              tfm = tf.Module()
              tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x), [spec])
              tfm.__call__(im)
              tf.saved_model.save(
                  tfm,
                  f,
                  options=tf.saved_model.SaveOptions(experimental_custom_gradients=False)
                  if check_version(tf.__version__, "2.6")
                  else tf.saved_model.SaveOptions(),
              )
          return f, keras_model
      
      
      @try_export
      def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")):
          """Exports YOLOv5 model to TensorFlow GraphDef *.pb format; see https://github.com/leimao/Frozen_Graph_TensorFlow for details."""
          import tensorflow as tf
          from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
      
          LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
          f = file.with_suffix(".pb")
      
          m = tf.function(lambda x: keras_model(x))  # full model
          m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
          frozen_func = convert_variables_to_constants_v2(m)
          frozen_func.graph.as_graph_def()
          tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
          return f, None
      
      
      @try_export
      def export_tflite(
          keras_model, im, file, int8, per_tensor, data, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:")
      ):
          # YOLOv5 TensorFlow Lite export
          import tensorflow as tf
      
          LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
          batch_size, ch, *imgsz = list(im.shape)  # BCHW
          f = str(file).replace(".pt", "-fp16.tflite")
      
          converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
          converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
          converter.target_spec.supported_types = [tf.float16]
          converter.optimizations = [tf.lite.Optimize.DEFAULT]
          if int8:
              from models.tf import representative_dataset_gen
      
              dataset = LoadImages(check_dataset(check_yaml(data))["train"], img_size=imgsz, auto=False)
              converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)
              converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
              converter.target_spec.supported_types = []
              converter.inference_input_type = tf.uint8  # or tf.int8
              converter.inference_output_type = tf.uint8  # or tf.int8
              converter.experimental_new_quantizer = True
              if per_tensor:
                  converter._experimental_disable_per_channel = True
              f = str(file).replace(".pt", "-int8.tflite")
          if nms or agnostic_nms:
              converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS)
      
          tflite_model = converter.convert()
          open(f, "wb").write(tflite_model)
          return f, None
      
      
      @try_export
      def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
          """
          Exports a YOLOv5 model to Edge TPU compatible TFLite format; requires Linux and Edge TPU compiler.
      
          https://coral.ai/docs/edgetpu/models-intro/
          """
          cmd = "edgetpu_compiler --version"
          help_url = "https://coral.ai/docs/edgetpu/compiler/"
          assert platform.system() == "Linux", f"export only supported on Linux. See {help_url}"
          if subprocess.run(f"{cmd} > /dev/null 2>&1", shell=True).returncode != 0:
              LOGGER.info(f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}")
              sudo = subprocess.run("sudo --version >/dev/null", shell=True).returncode == 0  # sudo installed on system
              for c in (
                  "curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -",
                  'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',
                  "sudo apt-get update",
                  "sudo apt-get install edgetpu-compiler",
              ):
                  subprocess.run(c if sudo else c.replace("sudo ", ""), shell=True, check=True)
          ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]
      
          LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
          f = str(file).replace(".pt", "-int8_edgetpu.tflite")  # Edge TPU model
          f_tfl = str(file).replace(".pt", "-int8.tflite")  # TFLite model
      
          subprocess.run(
              [
                  "edgetpu_compiler",
                  "-s",
                  "-d",
                  "-k",
                  "10",
                  "--out_dir",
                  str(file.parent),
                  f_tfl,
              ],
              check=True,
          )
          return f, None
      
      
      @try_export
      def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
          """Exports a YOLOv5 model to TensorFlow.js format, optionally with uint8 quantization."""
          check_requirements("tensorflowjs")
          import tensorflowjs as tfjs
      
          LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
          f = str(file).replace(".pt", "_web_model")  # js dir
          f_pb = file.with_suffix(".pb")  # *.pb path
          f_json = f"{f}/model.json"  # *.json path
      
          args = [
              "tensorflowjs_converter",
              "--input_format=tf_frozen_model",
              "--quantize_uint8" if int8 else "",
              "--output_node_names=Identity,Identity_1,Identity_2,Identity_3",
              str(f_pb),
              f,
          ]
          subprocess.run([arg for arg in args if arg], check=True)
      
          json = Path(f_json).read_text()
          with open(f_json, "w") as j:  # sort JSON Identity_* in ascending order
              subst = re.sub(
                  r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
                  r'"Identity.?.?": {"name": "Identity.?.?"}, '
                  r'"Identity.?.?": {"name": "Identity.?.?"}, '
                  r'"Identity.?.?": {"name": "Identity.?.?"}}}',
                  r'{"outputs": {"Identity": {"name": "Identity"}, '
                  r'"Identity_1": {"name": "Identity_1"}, '
                  r'"Identity_2": {"name": "Identity_2"}, '
                  r'"Identity_3": {"name": "Identity_3"}}}',
                  json,
              )
              j.write(subst)
          return f, None
      
      
      def add_tflite_metadata(file, metadata, num_outputs):
          """
          Adds TFLite metadata to a model file, supporting multiple outputs, as specified by TensorFlow guidelines.
      
          https://www.tensorflow.org/lite/models/convert/metadata
          """
          with contextlib.suppress(ImportError):
              # check_requirements('tflite_support')
              from tflite_support import flatbuffers
              from tflite_support import metadata as _metadata
              from tflite_support import metadata_schema_py_generated as _metadata_fb
      
              tmp_file = Path("/tmp/meta.txt")
              with open(tmp_file, "w") as meta_f:
                  meta_f.write(str(metadata))
      
              model_meta = _metadata_fb.ModelMetadataT()
              label_file = _metadata_fb.AssociatedFileT()
              label_file.name = tmp_file.name
              model_meta.associatedFiles = [label_file]
      
              subgraph = _metadata_fb.SubGraphMetadataT()
              subgraph.inputTensorMetadata = [_metadata_fb.TensorMetadataT()]
              subgraph.outputTensorMetadata = [_metadata_fb.TensorMetadataT()] * num_outputs
              model_meta.subgraphMetadata = [subgraph]
      
              b = flatbuffers.Builder(0)
              b.Finish(model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
              metadata_buf = b.Output()
      
              populator = _metadata.MetadataPopulator.with_model_file(file)
              populator.load_metadata_buffer(metadata_buf)
              populator.load_associated_files([str(tmp_file)])
              populator.populate()
              tmp_file.unlink()
      
      
      def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")):
          """Converts a PyTorch YOLOv5 model to CoreML format with NMS, handling different input/output shapes and saving the
          model.
          """
          import coremltools as ct
          from PIL import Image
      
          print(f"{prefix} starting pipeline with coremltools {ct.__version__}...")
          batch_size, ch, h, w = list(im.shape)  # BCHW
          t = time.time()
      
          # YOLOv5 Output shapes
          spec = model.get_spec()
          out0, out1 = iter(spec.description.output)
          if platform.system() == "Darwin":
              img = Image.new("RGB", (w, h))  # img(192 width, 320 height)
              # img = torch.zeros((*opt.img_size, 3)).numpy()  # img size(320,192,3) iDetection
              out = model.predict({"image": img})
              out0_shape, out1_shape = out[out0.name].shape, out[out1.name].shape
          else:  # linux and windows can not run model.predict(), get sizes from pytorch output y
              s = tuple(y[0].shape)
              out0_shape, out1_shape = (s[1], s[2] - 5), (s[1], 4)  # (3780, 80), (3780, 4)
      
          # Checks
          nx, ny = spec.description.input[0].type.imageType.width, spec.description.input[0].type.imageType.height
          na, nc = out0_shape
          # na, nc = out0.type.multiArrayType.shape  # number anchors, classes
          assert len(names) == nc, f"{len(names)} names found for nc={nc}"  # check
      
          # Define output shapes (missing)
          out0.type.multiArrayType.shape[:] = out0_shape  # (3780, 80)
          out1.type.multiArrayType.shape[:] = out1_shape  # (3780, 4)
          # spec.neuralNetwork.preprocessing[0].featureName = '0'
      
          # Flexible input shapes
          # from coremltools.models.neural_network import flexible_shape_utils
          # s = [] # shapes
          # s.append(flexible_shape_utils.NeuralNetworkImageSize(320, 192))
          # s.append(flexible_shape_utils.NeuralNetworkImageSize(640, 384))  # (height, width)
          # flexible_shape_utils.add_enumerated_image_sizes(spec, feature_name='image', sizes=s)
          # r = flexible_shape_utils.NeuralNetworkImageSizeRange()  # shape ranges
          # r.add_height_range((192, 640))
          # r.add_width_range((192, 640))
          # flexible_shape_utils.update_image_size_range(spec, feature_name='image', size_range=r)
      
          # Print
          print(spec.description)
      
          # Model from spec
          model = ct.models.MLModel(spec)
      
          # 3. Create NMS protobuf
          nms_spec = ct.proto.Model_pb2.Model()
          nms_spec.specificationVersion = 5
          for i in range(2):
              decoder_output = model._spec.description.output[i].SerializeToString()
              nms_spec.description.input.add()
              nms_spec.description.input[i].ParseFromString(decoder_output)
              nms_spec.description.output.add()
              nms_spec.description.output[i].ParseFromString(decoder_output)
      
          nms_spec.description.output[0].name = "confidence"
          nms_spec.description.output[1].name = "coordinates"
      
          output_sizes = [nc, 4]
          for i in range(2):
              ma_type = nms_spec.description.output[i].type.multiArrayType
              ma_type.shapeRange.sizeRanges.add()
              ma_type.shapeRange.sizeRanges[0].lowerBound = 0
              ma_type.shapeRange.sizeRanges[0].upperBound = -1
              ma_type.shapeRange.sizeRanges.add()
              ma_type.shapeRange.sizeRanges[1].lowerBound = output_sizes[i]
              ma_type.shapeRange.sizeRanges[1].upperBound = output_sizes[i]
              del ma_type.shape[:]
      
          nms = nms_spec.nonMaximumSuppression
          nms.confidenceInputFeatureName = out0.name  # 1x507x80
          nms.coordinatesInputFeatureName = out1.name  # 1x507x4
          nms.confidenceOutputFeatureName = "confidence"
          nms.coordinatesOutputFeatureName = "coordinates"
          nms.iouThresholdInputFeatureName = "iouThreshold"
          nms.confidenceThresholdInputFeatureName = "confidenceThreshold"
          nms.iouThreshold = 0.45
          nms.confidenceThreshold = 0.25
          nms.pickTop.perClass = True
          nms.stringClassLabels.vector.extend(names.values())
          nms_model = ct.models.MLModel(nms_spec)
      
          # 4. Pipeline models together
          pipeline = ct.models.pipeline.Pipeline(
              input_features=[
                  ("image", ct.models.datatypes.Array(3, ny, nx)),
                  ("iouThreshold", ct.models.datatypes.Double()),
                  ("confidenceThreshold", ct.models.datatypes.Double()),
              ],
              output_features=["confidence", "coordinates"],
          )
          pipeline.add_model(model)
          pipeline.add_model(nms_model)
      
          # Correct datatypes
          pipeline.spec.description.input[0].ParseFromString(model._spec.description.input[0].SerializeToString())
          pipeline.spec.description.output[0].ParseFromString(nms_model._spec.description.output[0].SerializeToString())
          pipeline.spec.description.output[1].ParseFromString(nms_model._spec.description.output[1].SerializeToString())
      
          # Update metadata
          pipeline.spec.specificationVersion = 5
          pipeline.spec.description.metadata.versionString = "https://github.com/ultralytics/yolov5"
          pipeline.spec.description.metadata.shortDescription = "https://github.com/ultralytics/yolov5"
          pipeline.spec.description.metadata.author = "glenn.jocher@ultralytics.com"
          pipeline.spec.description.metadata.license = "https://github.com/ultralytics/yolov5/blob/master/LICENSE"
          pipeline.spec.description.metadata.userDefined.update(
              {
                  "classes": ",".join(names.values()),
                  "iou_threshold": str(nms.iouThreshold),
                  "confidence_threshold": str(nms.confidenceThreshold),
              }
          )
      
          # Save the model
          f = file.with_suffix(".mlmodel")  # filename
          model = ct.models.MLModel(pipeline.spec)
          model.input_description["image"] = "Input image"
          model.input_description["iouThreshold"] = f"(optional) IOU Threshold override (default: {nms.iouThreshold})"
          model.input_description["confidenceThreshold"] = (
              f"(optional) Confidence Threshold override (default: {nms.confidenceThreshold})"
          )
          model.output_description["confidence"] = 'Boxes × Class confidence (see user-defined metadata "classes")'
          model.output_description["coordinates"] = "Boxes × [x, y, width, height] (relative to image size)"
          model.save(f)  # pipelined
          print(f"{prefix} pipeline success ({time.time() - t:.2f}s), saved as {f} ({file_size(f):.1f} MB)")
      
      
      @smart_inference_mode()
      def run(
          data=ROOT / "data/coco128.yaml",  # 'dataset.yaml path'
          weights=ROOT / "yolov5s.pt",  # weights path
          imgsz=(640, 640),  # image (height, width)
          batch_size=1,  # batch size
          device="cpu",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
          include=("torchscript", "onnx"),  # include formats
          half=False,  # FP16 half-precision export
          inplace=False,  # set YOLOv5 Detect() inplace=True
          keras=False,  # use Keras
          optimize=False,  # TorchScript: optimize for mobile
          int8=False,  # CoreML/TF INT8 quantization
          per_tensor=False,  # TF per tensor quantization
          dynamic=False,  # ONNX/TF/TensorRT: dynamic axes
          simplify=False,  # ONNX: simplify model
          opset=12,  # ONNX: opset version
          verbose=False,  # TensorRT: verbose log
          workspace=4,  # TensorRT: workspace size (GB)
          nms=False,  # TF: add NMS to model
          agnostic_nms=False,  # TF: add agnostic NMS to model
          topk_per_class=100,  # TF.js NMS: topk per class to keep
          topk_all=100,  # TF.js NMS: topk for all classes to keep
          iou_thres=0.45,  # TF.js NMS: IoU threshold
          conf_thres=0.25,  # TF.js NMS: confidence threshold
      ):
          t = time.time()
          include = [x.lower() for x in include]  # to lowercase
          fmts = tuple(export_formats()["Argument"][1:])  # --include arguments
          flags = [x in include for x in fmts]
          assert sum(flags) == len(include), f"ERROR: Invalid --include {include}, valid --include arguments are {fmts}"
          jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags  # export booleans
          file = Path(url2file(weights) if str(weights).startswith(("http:/", "https:/")) else weights)  # PyTorch weights
      
          # Load PyTorch model
          device = select_device(device)
          if half:
              assert device.type != "cpu" or coreml, "--half only compatible with GPU export, i.e. use --device 0"
              assert not dynamic, "--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both"
          model = attempt_load(weights, device=device, inplace=True, fuse=True)  # load FP32 model
      
          # Checks
          imgsz *= 2 if len(imgsz) == 1 else 1  # expand
          if optimize:
              assert device.type == "cpu", "--optimize not compatible with cuda devices, i.e. use --device cpu"
      
          # Input
          gs = int(max(model.stride))  # grid size (max stride)
          imgsz = [check_img_size(x, gs) for x in imgsz]  # verify img_size are gs-multiples
          im = torch.zeros(batch_size, 3, *imgsz).to(device)  # image size(1,3,320,192) BCHW iDetection
      
          # Update model
          model.eval()
          for k, m in model.named_modules():
              if isinstance(m, Detect):
                  m.inplace = inplace
                  m.dynamic = dynamic
                  m.export = True
      
          for _ in range(2):
              y = model(im)  # dry runs
          if half and not coreml:
              im, model = im.half(), model.half()  # to FP16
          shape = tuple((y[0] if isinstance(y, tuple) else y).shape)  # model output shape
          metadata = {"stride": int(max(model.stride)), "names": model.names}  # model metadata
          LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")
      
          # Exports
          f = [""] * len(fmts)  # exported filenames
          warnings.filterwarnings(action="ignore", category=torch.jit.TracerWarning)  # suppress TracerWarning
          if jit:  # TorchScript
              f[0], _ = export_torchscript(model, im, file, optimize)
          if engine:  # TensorRT required before ONNX
              f[1], _ = export_engine(model, im, file, half, dynamic, simplify, workspace, verbose)
          if onnx or xml:  # OpenVINO requires ONNX
              f[2], _ = export_onnx(model, im, file, opset, dynamic, simplify)
          if xml:  # OpenVINO
              f[3], _ = export_openvino(file, metadata, half, int8, data)
          if coreml:  # CoreML
              f[4], ct_model = export_coreml(model, im, file, int8, half, nms)
              if nms:
                  pipeline_coreml(ct_model, im, file, model.names, y)
          if any((saved_model, pb, tflite, edgetpu, tfjs)):  # TensorFlow formats
              assert not tflite or not tfjs, "TFLite and TF.js models must be exported separately, please pass only one type."
              assert not isinstance(model, ClassificationModel), "ClassificationModel export to TF formats not yet supported."
              f[5], s_model = export_saved_model(
                  model.cpu(),
                  im,
                  file,
                  dynamic,
                  tf_nms=nms or agnostic_nms or tfjs,
                  agnostic_nms=agnostic_nms or tfjs,
                  topk_per_class=topk_per_class,
                  topk_all=topk_all,
                  iou_thres=iou_thres,
                  conf_thres=conf_thres,
                  keras=keras,
              )
              if pb or tfjs:  # pb prerequisite to tfjs
                  f[6], _ = export_pb(s_model, file)
              if tflite or edgetpu:
                  f[7], _ = export_tflite(
                      s_model, im, file, int8 or edgetpu, per_tensor, data=data, nms=nms, agnostic_nms=agnostic_nms
                  )
                  if edgetpu:
                      f[8], _ = export_edgetpu(file)
                  add_tflite_metadata(f[8] or f[7], metadata, num_outputs=len(s_model.outputs))
              if tfjs:
                  f[9], _ = export_tfjs(file, int8)
          if paddle:  # PaddlePaddle
              f[10], _ = export_paddle(model, im, file, metadata)
      
          # Finish
          f = [str(x) for x in f if x]  # filter out '' and None
          if any(f):
              cls, det, seg = (isinstance(model, x) for x in (ClassificationModel, DetectionModel, SegmentationModel))  # type
              det &= not seg  # segmentation models inherit from SegmentationModel(DetectionModel)
              dir = Path("segment" if seg else "classify" if cls else "")
              h = "--half" if half else ""  # --half FP16 inference arg
              s = (
                  "# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference"
                  if cls
                  else "# WARNING ⚠️ SegmentationModel not yet supported for PyTorch Hub AutoShape inference"
                  if seg
                  else ""
              )
              LOGGER.info(
                  f'\nExport complete ({time.time() - t:.1f}s)'
                  f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
                  f"\nDetect:          python {dir / ('detect.py' if det else 'predict.py')} --weights {f[-1]} {h}"
                  f"\nValidate:        python {dir / 'val.py'} --weights {f[-1]} {h}"
                  f"\nPyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}')  {s}"
                  f'\nVisualize:       https://netron.app'
              )
          return f  # return list of exported files/dirs
      
      
      def parse_opt(known=False):
          """Parses command-line arguments for YOLOv5 model export configurations, returning the parsed options."""
          parser = argparse.ArgumentParser()
          parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
          parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model.pt path(s)")
          parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640, 640], help="image (h, w)")
          parser.add_argument("--batch-size", type=int, default=1, help="batch size")
          parser.add_argument("--device", default="cpu", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
          parser.add_argument("--half", action="store_true", help="FP16 half-precision export")
          parser.add_argument("--inplace", action="store_true", help="set YOLOv5 Detect() inplace=True")
          parser.add_argument("--keras", action="store_true", help="TF: use Keras")
          parser.add_argument("--optimize", action="store_true", help="TorchScript: optimize for mobile")
          parser.add_argument("--int8", action="store_true", help="CoreML/TF/OpenVINO INT8 quantization")
          parser.add_argument("--per-tensor", action="store_true", help="TF per-tensor quantization")
          parser.add_argument("--dynamic", action="store_true", help="ONNX/TF/TensorRT: dynamic axes")
          parser.add_argument("--simplify", action="store_true", help="ONNX: simplify model")
          parser.add_argument("--opset", type=int, default=17, help="ONNX: opset version")
          parser.add_argument("--verbose", action="store_true", help="TensorRT: verbose log")
          parser.add_argument("--workspace", type=int, default=4, help="TensorRT: workspace size (GB)")
          parser.add_argument("--nms", action="store_true", help="TF: add NMS to model")
          parser.add_argument("--agnostic-nms", action="store_true", help="TF: add agnostic NMS to model")
          parser.add_argument("--topk-per-class", type=int, default=100, help="TF.js NMS: topk per class to keep")
          parser.add_argument("--topk-all", type=int, default=100, help="TF.js NMS: topk for all classes to keep")
          parser.add_argument("--iou-thres", type=float, default=0.45, help="TF.js NMS: IoU threshold")
          parser.add_argument("--conf-thres", type=float, default=0.25, help="TF.js NMS: confidence threshold")
          parser.add_argument(
              "--include",
              nargs="+",
              default=["torchscript"],
              help="torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle",
          )
          opt = parser.parse_known_args()[0] if known else parser.parse_args()
          print_args(vars(opt))
          return opt
      
      
      def main(opt):
          """Executes the YOLOv5 model inference or export with specified weights and options."""
          for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
              run(**vars(opt))
      
      
      if __name__ == "__main__":
          opt = parse_opt()
          main(opt)
      
    • 📄 hubconf.py 8.6KB
      # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
      """
      PyTorch Hub models https://pytorch.org/hub/ultralytics_yolov5
      
      Usage:
          import torch
          model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # official model
          model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s')  # from branch
          model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt')  # custom/local model
          model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local')  # local repo
      """
      
      import torch
      
      
      def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
          """
          Creates or loads a YOLOv5 model.
      
          Arguments:
              name (str): model name 'yolov5s' or path 'path/to/best.pt'
              pretrained (bool): load pretrained weights into the model
              channels (int): number of input channels
              classes (int): number of model classes
              autoshape (bool): apply YOLOv5 .autoshape() wrapper to model
              verbose (bool): print all information to screen
              device (str, torch.device, None): device to use for model parameters
      
          Returns:
              YOLOv5 model
          """
          from pathlib import Path
      
          from models.common import AutoShape, DetectMultiBackend
          from models.experimental import attempt_load
          from models.yolo import ClassificationModel, DetectionModel, SegmentationModel
          from utils.downloads import attempt_download
          from utils.general import LOGGER, ROOT, check_requirements, intersect_dicts, logging
          from utils.torch_utils import select_device
      
          if not verbose:
              LOGGER.setLevel(logging.WARNING)
          check_requirements(ROOT / "requirements.txt", exclude=("opencv-python", "tensorboard", "thop"))
          name = Path(name)
          path = name.with_suffix(".pt") if name.suffix == "" and not name.is_dir() else name  # checkpoint path
          try:
              device = select_device(device)
              if pretrained and channels == 3 and classes == 80:
                  try:
                      model = DetectMultiBackend(path, device=device, fuse=autoshape)  # detection model
                      if autoshape:
                          if model.pt and isinstance(model.model, ClassificationModel):
                              LOGGER.warning(
                                  "WARNING ⚠️ YOLOv5 ClassificationModel is not yet AutoShape compatible. "
                                  "You must pass torch tensors in BCHW to this model, i.e. shape(1,3,224,224)."
                              )
                          elif model.pt and isinstance(model.model, SegmentationModel):
                              LOGGER.warning(
                                  "WARNING ⚠️ YOLOv5 SegmentationModel is not yet AutoShape compatible. "
                                  "You will not be able to run inference with this model."
                              )
                          else:
                              model = AutoShape(model)  # for file/URI/PIL/cv2/np inputs and NMS
                  except Exception:
                      model = attempt_load(path, device=device, fuse=False)  # arbitrary model
              else:
                  cfg = list((Path(__file__).parent / "models").rglob(f"{path.stem}.yaml"))[0]  # model.yaml path
                  model = DetectionModel(cfg, channels, classes)  # create model
                  if pretrained:
                      ckpt = torch.load(attempt_download(path), map_location=device)  # load
                      csd = ckpt["model"].float().state_dict()  # checkpoint state_dict as FP32
                      csd = intersect_dicts(csd, model.state_dict(), exclude=["anchors"])  # intersect
                      model.load_state_dict(csd, strict=False)  # load
                      if len(ckpt["model"].names) == classes:
                          model.names = ckpt["model"].names  # set class names attribute
              if not verbose:
                  LOGGER.setLevel(logging.INFO)  # reset to default
              return model.to(device)
      
          except Exception as e:
              help_url = "https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading"
              s = f"{e}. Cache may be out of date, try `force_reload=True` or see {help_url} for help."
              raise Exception(s) from e
      
      
      def custom(path="path/to/model.pt", autoshape=True, _verbose=True, device=None):
          """Loads a custom or local YOLOv5 model from a given path with optional autoshaping and device specification."""
          return _create(path, autoshape=autoshape, verbose=_verbose, device=device)
      
      
      def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
          """Instantiates the YOLOv5-nano model with options for pretraining, input channels, class count, autoshaping,
          verbosity, and device.
          """
          return _create("yolov5n", pretrained, channels, classes, autoshape, _verbose, device)
      
      
      def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
          """Creates YOLOv5-small model with options for pretraining, input channels, class count, autoshaping, verbosity, and
          device.
          """
          return _create("yolov5s", pretrained, channels, classes, autoshape, _verbose, device)
      
      
      def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
          """Instantiates the YOLOv5-medium model with customizable pretraining, channel count, class count, autoshaping,
          verbosity, and device.
          """
          return _create("yolov5m", pretrained, channels, classes, autoshape, _verbose, device)
      
      
      def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
          """Creates YOLOv5-large model with options for pretraining, channels, classes, autoshaping, verbosity, and device
          selection.
          """
          return _create("yolov5l", pretrained, channels, classes, autoshape, _verbose, device)
      
      
      def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
          """Instantiates the YOLOv5-xlarge model with customizable pretraining, channel count, class count, autoshaping,
          verbosity, and device.
          """
          return _create("yolov5x", pretrained, channels, classes, autoshape, _verbose, device)
      
      
      def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
          """Creates YOLOv5-nano-P6 model with options for pretraining, channels, classes, autoshaping, verbosity, and
          device.
          """
          return _create("yolov5n6", pretrained, channels, classes, autoshape, _verbose, device)
      
      
      def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
          """Instantiate YOLOv5-small-P6 model with options for pretraining, input channels, number of classes, autoshaping,
          verbosity, and device selection.
          """
          return _create("yolov5s6", pretrained, channels, classes, autoshape, _verbose, device)
      
      
      def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
          """Creates YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity,
          and device.
          """
          return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device)
      
      
      def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
          """Instantiates the YOLOv5-large-P6 model with customizable pretraining, channel and class counts, autoshaping,
          verbosity, and device selection.
          """
          return _create("yolov5l6", pretrained, channels, classes, autoshape, _verbose, device)
      
      
      def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
          """Creates YOLOv5-xlarge-P6 model with options for pretraining, channels, classes, autoshaping, verbosity, and
          device.
          """
          return _create("yolov5x6", pretrained, channels, classes, autoshape, _verbose, device)
      
      
      if __name__ == "__main__":
          import argparse
          from pathlib import Path
      
          import numpy as np
          from PIL import Image
      
          from utils.general import cv2, print_args
      
          # Argparser
          parser = argparse.ArgumentParser()
          parser.add_argument("--model", type=str, default="yolov5s", help="model name")
          opt = parser.parse_args()
          print_args(vars(opt))
      
          # Model
          model = _create(name=opt.model, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True)
          # model = custom(path='path/to/model.pt')  # custom
      
          # Images
          imgs = [
              "data/images/zidane.jpg",  # filename
              Path("data/images/zidane.jpg"),  # Path
              "https://ultralytics.com/images/zidane.jpg",  # URI
              cv2.imread("data/images/bus.jpg")[:, :, ::-1],  # OpenCV
              Image.open("data/images/bus.jpg"),  # PIL
              np.zeros((320, 640, 3)),
          ]  # numpy
      
          # Inference
          results = model(imgs, size=320)  # batched inference
      
          # Results
          results.print()
          results.save()
      
    • 📄 main.sh 121B
      python train.py --data coco.yaml --weights 'yolov5s.pt' --cfg yolov5s.yaml --img 640 --batch 16 --epochs 50 --patience 5
      
    • 📄 pyproject.toml 5.3KB
      # Ultralyticsv5 YOLO 🚀, AGPL-3.0 license
      
      # Overview:
      # This pyproject.toml file manages the build, packaging, and distribution of the Ultralytics library.
      # It defines essential project metadata, dependencies, and settings used to develop and deploy the library.
      
      # Key Sections:
      # - [build-system]: Specifies the build requirements and backend (e.g., setuptools, wheel).
      # - [project]: Includes details like name, version, description, authors, dependencies and more.
      # - [project.optional-dependencies]: Provides additional, optional packages for extended features.
      # - [tool.*]: Configures settings for various tools (pytest, yapf, etc.) used in the project.
      
      # Installation:
      # The Ultralytics library can be installed using the command: 'pip install ultralytics'
      # For development purposes, you can install the package in editable mode with: 'pip install -e .'
      # This approach allows for real-time code modifications without the need for re-installation.
      
      # Documentation:
      # For comprehensive documentation and usage instructions, visit: https://docs.ultralytics.com
      
      [build-system]
      requires = ["setuptools>=43.0.0", "wheel"]
      build-backend = "setuptools.build_meta"
      
      # Project settings -----------------------------------------------------------------------------------------------------
      [project]
      version = "7.0.0"
      name = "YOLOv5"
      description = "Ultralytics YOLOv5 for SOTA object detection, instance segmentation and image classification."
      readme = "README.md"
      requires-python = ">=3.8"
      license = { "text" = "AGPL-3.0" }
      keywords = ["machine-learning", "deep-learning", "computer-vision", "ML", "DL", "AI", "YOLO", "YOLOv3", "YOLOv5", "YOLOv8", "HUB", "Ultralytics"]
      authors = [
          { name = "Glenn Jocher" },
          { name = "Ayush Chaurasia" },
          { name = "Jing Qiu" }
      ]
      maintainers = [
          { name = "Glenn Jocher" },
          { name = "Ayush Chaurasia" },
          { name = "Jing Qiu" }
      ]
      classifiers = [
          "Development Status :: 4 - Beta",
          "Intended Audience :: Developers",
          "Intended Audience :: Education",
          "Intended Audience :: Science/Research",
          "License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)",
          "Programming Language :: Python :: 3",
          "Programming Language :: Python :: 3.8",
          "Programming Language :: Python :: 3.9",
          "Programming Language :: Python :: 3.10",
          "Programming Language :: Python :: 3.11",
          "Topic :: Software Development",
          "Topic :: Scientific/Engineering",
          "Topic :: Scientific/Engineering :: Artificial Intelligence",
          "Topic :: Scientific/Engineering :: Image Recognition",
          "Operating System :: POSIX :: Linux",
          "Operating System :: MacOS",
          "Operating System :: Microsoft :: Windows",
      ]
      
      # Required dependencies ------------------------------------------------------------------------------------------------
      dependencies = [
          "matplotlib>=3.3.0",
          "numpy>=1.22.2",
          "opencv-python>=4.6.0",
          "pillow>=7.1.2",
          "pyyaml>=5.3.1",
          "requests>=2.23.0",
          "scipy>=1.4.1",
          "torch>=1.8.0",
          "torchvision>=0.9.0",
          "tqdm>=4.64.0", # progress bars
          "psutil", # system utilization
          "py-cpuinfo", # display CPU info
          "thop>=0.1.1", # FLOPs computation
          "pandas>=1.1.4",
          "seaborn>=0.11.0", # plotting
          "ultralytics>=8.1.47"
      ]
      
      # Optional dependencies ------------------------------------------------------------------------------------------------
      [project.optional-dependencies]
      dev = [
          "ipython",
          "check-manifest",
          "pre-commit",
          "pytest",
          "pytest-cov",
          "coverage[toml]",
          "mkdocs-material",
          "mkdocstrings[python]",
          "mkdocs-redirects", # for 301 redirects
          "mkdocs-ultralytics-plugin>=0.0.34", # for meta descriptions and images, dates and authors
      ]
      export = [
          "onnx>=1.12.0", # ONNX export
          "coremltools>=7.0; platform_system != 'Windows'", # CoreML only supported on macOS and Linux
          "openvino-dev>=2023.0", # OpenVINO export
          "tensorflow<=2.13.1", # TF bug https://github.com/ultralytics/ultralytics/issues/5161
          "tensorflowjs>=3.9.0", # TF.js export, automatically installs tensorflow
      ]
      # tensorflow>=2.4.1,<=2.13.1  # TF exports (-cpu, -aarch64, -macos)
      # tflite-support  # for TFLite model metadata
      # scikit-learn==0.19.2  # CoreML quantization
      # nvidia-pyindex  # TensorRT export
      # nvidia-tensorrt  # TensorRT export
      logging = [
          "comet", # https://docs.ultralytics.com/integrations/comet/
          "tensorboard>=2.13.0",
          "dvclive>=2.12.0",
      ]
      extra = [
          "ipython", # interactive notebook
          "albumentations>=1.0.3", # training augmentations
          "pycocotools>=2.0.6", # COCO mAP
      ]
      
      [project.urls]
      "Bug Reports" = "https://github.com/ultralytics/yolov5/issues"
      "Funding" = "https://ultralytics.com"
      "Source" = "https://github.com/ultralytics/yolov5/"
      
      # Tools settings -------------------------------------------------------------------------------------------------------
      [tool.pytest]
      norecursedirs = [".git", "dist", "build"]
      addopts = "--doctest-modules --durations=30 --color=yes"
      
      [tool.isort]
      line_length = 120
      multi_line_output = 0
      
      [tool.ruff]
      line-length = 120
      
      [tool.docformatter]
      wrap-summaries = 120
      wrap-descriptions = 120
      in-place = true
      pre-summary-newline = true
      close-quotes-on-newline = true
      
      [tool.codespell]
      ignore-words-list = "crate,nd,strack,dota,ane,segway,fo,gool,winn,commend"
      skip = '*.csv,*venv*,docs/??/,docs/mkdocs_??.yml'
      
    • 📄 requirements.txt 1.6KB
      # YOLOv5 requirements
      # Usage: pip install -r requirements.txt
      
      # Base ------------------------------------------------------------------------
      gitpython>=3.1.30
      matplotlib>=3.3
      numpy>=1.23.5
      opencv-python>=4.1.1
      pillow>=10.3.0
      psutil  # system resources
      PyYAML>=5.3.1
      requests>=2.23.0
      scipy>=1.4.1
      thop>=0.1.1  # FLOPs computation
      torch>=1.8.0  # see https://pytorch.org/get-started/locally (recommended)
      torchvision>=0.9.0
      tqdm>=4.64.0
      ultralytics>=8.0.232
      # protobuf<=3.20.1  # https://github.com/ultralytics/yolov5/issues/8012
      
      # Logging ---------------------------------------------------------------------
      # tensorboard>=2.4.1
      # clearml>=1.2.0
      # comet
      
      # Plotting --------------------------------------------------------------------
      pandas>=1.1.4
      seaborn>=0.11.0
      
      # Export ----------------------------------------------------------------------
      # coremltools>=6.0  # CoreML export
      # onnx>=1.10.0  # ONNX export
      # onnx-simplifier>=0.4.1  # ONNX simplifier
      # nvidia-pyindex  # TensorRT export
      # nvidia-tensorrt  # TensorRT export
      # scikit-learn<=1.1.2  # CoreML quantization
      # tensorflow>=2.4.0,<=2.13.1  # TF exports (-cpu, -aarch64, -macos)
      # tensorflowjs>=3.9.0  # TF.js export
      # openvino-dev>=2023.0  # OpenVINO export
      
      # Deploy ----------------------------------------------------------------------
      setuptools>=65.5.1 # Snyk vulnerability fix
      # tritonclient[all]~=2.24.0
      
      # Extras ----------------------------------------------------------------------
      # ipython  # interactive notebook
      # mss  # screenshots
      # albumentations>=1.0.3
      # pycocotools>=2.0.6  # COCO mAP
      wheel>=0.38.0 # not directly required, pinned by Snyk to avoid a vulnerability
      
    • 📄 train.py 38.9KB
      # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
      """
      Train a YOLOv5 model on a custom dataset. Models and datasets download automatically from the latest YOLOv5 release.
      
      Usage - Single-GPU training:
          $ python train.py --data coco128.yaml --weights yolov5s.pt --img 640  # from pretrained (recommended)
          $ python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640  # from scratch
      
      Usage - Multi-GPU DDP training:
          $ python -m torch.distributed.run --nproc_per_node 3 --master_port 1 train.py --data coco128.yaml --weights yolov5s.pt --img 640 --device 0,2,3
      
      Models:     https://github.com/ultralytics/yolov5/tree/master/models
      Datasets:   https://github.com/ultralytics/yolov5/tree/master/data
      Tutorial:   https://docs.ultralytics.com/yolov5/tutorials/train_custom_data
      """
      
      import argparse
      import math
      import os
      import random
      import subprocess
      import sys
      import time
      from copy import deepcopy
      from datetime import datetime, timedelta
      from pathlib import Path
      
      try:
          import comet_ml  # must be imported before torch (if installed)
      except ImportError:
          comet_ml = None
      
      import numpy as np
      import torch
      import torch.distributed as dist
      import torch.nn as nn
      import yaml
      from torch.optim import lr_scheduler
      from tqdm import tqdm
      
      FILE = Path(__file__).resolve()
      ROOT = FILE.parents[0]  # YOLOv5 root directory
      if str(ROOT) not in sys.path:
          sys.path.append(str(ROOT))  # add ROOT to PATH
      ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
      
      import val as validate  # for end-of-epoch mAP
      from models.experimental import attempt_load
      from models.yolo import Model
      from utils.autoanchor import check_anchors
      from utils.autobatch import check_train_batch_size
      from utils.callbacks import Callbacks
      from utils.dataloaders import create_dataloader
      from utils.downloads import attempt_download, is_url
      from utils.general import (
          LOGGER,
          TQDM_BAR_FORMAT,
          check_amp,
          check_dataset,
          check_file,
          check_git_info,
          check_git_status,
          check_img_size,
          check_requirements,
          check_suffix,
          check_yaml,
          colorstr,
          get_latest_run,
          increment_path,
          init_seeds,
          intersect_dicts,
          labels_to_class_weights,
          labels_to_image_weights,
          methods,
          one_cycle,
          print_args,
          print_mutation,
          strip_optimizer,
          yaml_save,
      )
      from utils.loggers import LOGGERS, Loggers
      from utils.loggers.comet.comet_utils import check_comet_resume
      from utils.loss import ComputeLoss
      from utils.metrics import fitness
      from utils.plots import plot_evolve
      from utils.torch_utils import (
          EarlyStopping,
          ModelEMA,
          de_parallel,
          select_device,
          smart_DDP,
          smart_optimizer,
          smart_resume,
          torch_distributed_zero_first,
      )
      
      LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))  # https://pytorch.org/docs/stable/elastic/run.html
      RANK = int(os.getenv("RANK", -1))
      WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
      GIT_INFO = check_git_info()
      
      
      def train(hyp, opt, device, callbacks):
          """
          Trains YOLOv5 model with given hyperparameters, options, and device, managing datasets, model architecture, loss
          computation, and optimizer steps.
      
          `hyp` argument is path/to/hyp.yaml or hyp dictionary.
          """
          save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = (
              Path(opt.save_dir),
              opt.epochs,
              opt.batch_size,
              opt.weights,
              opt.single_cls,
              opt.evolve,
              opt.data,
              opt.cfg,
              opt.resume,
              opt.noval,
              opt.nosave,
              opt.workers,
              opt.freeze,
          )
          callbacks.run("on_pretrain_routine_start")
      
          # Directories
          w = save_dir / "weights"  # weights dir
          (w.parent if evolve else w).mkdir(parents=True, exist_ok=True)  # make dir
          last, best = w / "last.pt", w / "best.pt"
      
          # Hyperparameters
          if isinstance(hyp, str):
              with open(hyp, errors="ignore") as f:
                  hyp = yaml.safe_load(f)  # load hyps dict
          LOGGER.info(colorstr("hyperparameters: ") + ", ".join(f"{k}={v}" for k, v in hyp.items()))
          opt.hyp = hyp.copy()  # for saving hyps to checkpoints
      
          # Save run settings
          if not evolve:
              yaml_save(save_dir / "hyp.yaml", hyp)
              yaml_save(save_dir / "opt.yaml", vars(opt))
      
          # Loggers
          data_dict = None
          if RANK in {-1, 0}:
              include_loggers = list(LOGGERS)
              if getattr(opt, "ndjson_console", False):
                  include_loggers.append("ndjson_console")
              if getattr(opt, "ndjson_file", False):
                  include_loggers.append("ndjson_file")
      
              loggers = Loggers(
                  save_dir=save_dir,
                  weights=weights,
                  opt=opt,
                  hyp=hyp,
                  logger=LOGGER,
                  include=tuple(include_loggers),
              )
      
              # Register actions
              for k in methods(loggers):
                  callbacks.register_action(k, callback=getattr(loggers, k))
      
              # Process custom dataset artifact link
              data_dict = loggers.remote_dataset
              if resume:  # If resuming runs from remote artifact
                  weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
      
          # Config
          plots = not evolve and not opt.noplots  # create plots
          cuda = device.type != "cpu"
          init_seeds(opt.seed + 1 + RANK, deterministic=True)
          with torch_distributed_zero_first(LOCAL_RANK):
              data_dict = data_dict or check_dataset(data)  # check if None
          train_path, val_path = data_dict["train"], data_dict["val"]
          nc = 1 if single_cls else int(data_dict["nc"])  # number of classes
          names = {0: "item"} if single_cls and len(data_dict["names"]) != 1 else data_dict["names"]  # class names
          is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt")  # COCO dataset
      
          # Model
          check_suffix(weights, ".pt")  # check weights
          pretrained = weights.endswith(".pt")
          if pretrained:
              with torch_distributed_zero_first(LOCAL_RANK):
                  weights = attempt_download(weights)  # download if not found locally
              ckpt = torch.load(weights, map_location="cpu",weights_only=False)  # load checkpoint to CPU to avoid CUDA memory leak
              model = Model(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)  # create
              exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else []  # exclude keys
              csd = ckpt["model"].float().state_dict()  # checkpoint state_dict as FP32
              csd = intersect_dicts(csd, model.state_dict(), exclude=exclude)  # intersect
              model.load_state_dict(csd, strict=False)  # load
              LOGGER.info(f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}")  # report
          else:
              model = Model(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)  # create
          amp = check_amp(model)  # check AMP
      
          # Freeze
          freeze = [f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0]))]  # layers to freeze
          for k, v in model.named_parameters():
              v.requires_grad = True  # train all layers
              # v.register_hook(lambda x: torch.nan_to_num(x))  # NaN to 0 (commented for erratic training results)
              if any(x in k for x in freeze):
                  LOGGER.info(f"freezing {k}")
                  v.requires_grad = False
      
          # Image size
          gs = max(int(model.stride.max()), 32)  # grid size (max stride)
          imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2)  # verify imgsz is gs-multiple
      
          # Batch size
          if RANK == -1 and batch_size == -1:  # single-GPU only, estimate best batch size
              batch_size = check_train_batch_size(model, imgsz, amp)
              loggers.on_params_update({"batch_size": batch_size})
      
          # Optimizer
          nbs = 64  # nominal batch size
          accumulate = max(round(nbs / batch_size), 1)  # accumulate loss before optimizing
          hyp["weight_decay"] *= batch_size * accumulate / nbs  # scale weight_decay
          optimizer = smart_optimizer(model, opt.optimizer, hyp["lr0"], hyp["momentum"], hyp["weight_decay"])
      
          # Scheduler
          if opt.cos_lr:
              lf = one_cycle(1, hyp["lrf"], epochs)  # cosine 1->hyp['lrf']
          else:
              lf = lambda x: (1 - x / epochs) * (1.0 - hyp["lrf"]) + hyp["lrf"]  # linear
          scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)  # plot_lr_scheduler(optimizer, scheduler, epochs)
      
          # EMA
          ema = ModelEMA(model) if RANK in {-1, 0} else None
      
          # Resume
          best_fitness, start_epoch = 0.0, 0
          if pretrained:
              if resume:
                  best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
              del ckpt, csd
      
          # DP mode
          if cuda and RANK == -1 and torch.cuda.device_count() > 1:
              LOGGER.warning(
                  "WARNING ⚠️ DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n"
                  "See Multi-GPU Tutorial at https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training to get started."
              )
              model = torch.nn.DataParallel(model)
      
          # SyncBatchNorm
          if opt.sync_bn and cuda and RANK != -1:
              model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
              LOGGER.info("Using SyncBatchNorm()")
      
          # Trainloader
          train_loader, dataset = create_dataloader(
              train_path,
              imgsz,
              batch_size // WORLD_SIZE,
              gs,
              single_cls,
              hyp=hyp,
              augment=True,
              cache=None if opt.cache == "val" else opt.cache,
              rect=opt.rect,
              rank=LOCAL_RANK,
              workers=workers,
              image_weights=opt.image_weights,
              quad=opt.quad,
              prefix=colorstr("train: "),
              shuffle=True,
              seed=opt.seed,
          )
          labels = np.concatenate(dataset.labels, 0)
          mlc = int(labels[:, 0].max())  # max label class
          assert mlc < nc, f"Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}"
      
          # Process 0
          if RANK in {-1, 0}:
              val_loader = create_dataloader(
                  val_path,
                  imgsz,
                  batch_size // WORLD_SIZE * 2,
                  gs,
                  single_cls,
                  hyp=hyp,
                  cache=None if noval else opt.cache,
                  rect=True,
                  rank=-1,
                  workers=workers * 2,
                  pad=0.5,
                  prefix=colorstr("val: "),
              )[0]
      
              if not resume:
                  if not opt.noautoanchor:
                      check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz)  # run AutoAnchor
                  model.half().float()  # pre-reduce anchor precision
      
              callbacks.run("on_pretrain_routine_end", labels, names)
      
          # DDP mode
          if cuda and RANK != -1:
              model = smart_DDP(model)
      
          # Model attributes
          nl = de_parallel(model).model[-1].nl  # number of detection layers (to scale hyps)
          hyp["box"] *= 3 / nl  # scale to layers
          hyp["cls"] *= nc / 80 * 3 / nl  # scale to classes and layers
          hyp["obj"] *= (imgsz / 640) ** 2 * 3 / nl  # scale to image size and layers
          hyp["label_smoothing"] = opt.label_smoothing
          model.nc = nc  # attach number of classes to model
          model.hyp = hyp  # attach hyperparameters to model
          model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc  # attach class weights
          model.names = names
      
          # Start training
          t0 = time.time()
          nb = len(train_loader)  # number of batches
          nw = max(round(hyp["warmup_epochs"] * nb), 100)  # number of warmup iterations, max(3 epochs, 100 iterations)
          # nw = min(nw, (epochs - start_epoch) / 2 * nb)  # limit warmup to < 1/2 of training
          last_opt_step = -1
          maps = np.zeros(nc)  # mAP per class
          results = (0, 0, 0, 0, 0, 0, 0)  # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
          scheduler.last_epoch = start_epoch - 1  # do not move
          scaler = torch.cuda.amp.GradScaler(enabled=amp)
          stopper, stop = EarlyStopping(patience=opt.patience), False
          compute_loss = ComputeLoss(model)  # init loss class
          callbacks.run("on_train_start")
          LOGGER.info(
              f'Image sizes {imgsz} train, {imgsz} val\n'
              f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'
              f"Logging results to {colorstr('bold', save_dir)}\n"
              f'Starting training for {epochs} epochs...'
          )
          for epoch in range(start_epoch, epochs):  # epoch ------------------------------------------------------------------
              callbacks.run("on_train_epoch_start")
              model.train()
      
              # Update image weights (optional, single-GPU only)
              if opt.image_weights:
                  cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc  # class weights
                  iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)  # image weights
                  dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)  # rand weighted idx
      
              # Update mosaic border (optional)
              # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
              # dataset.mosaic_border = [b - imgsz, -b]  # height, width borders
      
              mloss = torch.zeros(3, device=device)  # mean losses
              if RANK != -1:
                  train_loader.sampler.set_epoch(epoch)
              pbar = enumerate(train_loader)
              LOGGER.info(("\n" + "%11s" * 7) % ("Epoch", "GPU_mem", "box_loss", "obj_loss", "cls_loss", "Instances", "Size"))
              if RANK in {-1, 0}:
                  pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT)  # progress bar
              optimizer.zero_grad()
              for i, (imgs, targets, paths, _) in pbar:  # batch -------------------------------------------------------------
                  callbacks.run("on_train_batch_start")
                  ni = i + nb * epoch  # number integrated batches (since train start)
                  imgs = imgs.to(device, non_blocking=True).float() / 255  # uint8 to float32, 0-255 to 0.0-1.0
      
                  # Warmup
                  if ni <= nw:
                      xi = [0, nw]  # x interp
                      # compute_loss.gr = np.interp(ni, xi, [0.0, 1.0])  # iou loss ratio (obj_loss = 1.0 or iou)
                      accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
                      for j, x in enumerate(optimizer.param_groups):
                          # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
                          x["lr"] = np.interp(ni, xi, [hyp["warmup_bias_lr"] if j == 0 else 0.0, x["initial_lr"] * lf(epoch)])
                          if "momentum" in x:
                              x["momentum"] = np.interp(ni, xi, [hyp["warmup_momentum"], hyp["momentum"]])
      
                  # Multi-scale
                  if opt.multi_scale:
                      sz = random.randrange(int(imgsz * 0.5), int(imgsz * 1.5) + gs) // gs * gs  # size
                      sf = sz / max(imgs.shape[2:])  # scale factor
                      if sf != 1:
                          ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to gs-multiple)
                          imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
      
                  # Forward
                  with torch.cuda.amp.autocast(amp):
                      pred = model(imgs)  # forward
                      loss, loss_items = compute_loss(pred, targets.to(device))  # loss scaled by batch_size
                      if RANK != -1:
                          loss *= WORLD_SIZE  # gradient averaged between devices in DDP mode
                      if opt.quad:
                          loss *= 4.0
      
                  # Backward
                  scaler.scale(loss).backward()
      
                  # Optimize - https://pytorch.org/docs/master/notes/amp_examples.html
                  if ni - last_opt_step >= accumulate:
                      scaler.unscale_(optimizer)  # unscale gradients
                      torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)  # clip gradients
                      scaler.step(optimizer)  # optimizer.step
                      scaler.update()
                      optimizer.zero_grad()
                      if ema:
                          ema.update(model)
                      last_opt_step = ni
      
                  # Log
                  if RANK in {-1, 0}:
                      mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses
                      mem = f"{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G"  # (GB)
                      pbar.set_description(
                          ("%11s" * 2 + "%11.4g" * 5)
                          % (f"{epoch}/{epochs - 1}", mem, *mloss, targets.shape[0], imgs.shape[-1])
                      )
                      callbacks.run("on_train_batch_end", model, ni, imgs, targets, paths, list(mloss))
                      if callbacks.stop_training:
                          return
                  # end batch ------------------------------------------------------------------------------------------------
      
              # Scheduler
              lr = [x["lr"] for x in optimizer.param_groups]  # for loggers
              scheduler.step()
      
              if RANK in {-1, 0}:
                  # mAP
                  callbacks.run("on_train_epoch_end", epoch=epoch)
                  ema.update_attr(model, include=["yaml", "nc", "hyp", "names", "stride", "class_weights"])
                  final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
                  if not noval or final_epoch:  # Calculate mAP
                      results, maps, _ = validate.run(
                          data_dict,
                          batch_size=batch_size // WORLD_SIZE * 2,
                          imgsz=imgsz,
                          half=amp,
                          model=ema.ema,
                          single_cls=single_cls,
                          dataloader=val_loader,
                          save_dir=save_dir,
                          plots=False,
                          callbacks=callbacks,
                          compute_loss=compute_loss,
                      )
      
                  # Update best mAP
                  fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
                  stop = stopper(epoch=epoch, fitness=fi)  # early stop check
                  if fi > best_fitness:
                      best_fitness = fi
                  log_vals = list(mloss) + list(results) + lr
                  callbacks.run("on_fit_epoch_end", log_vals, epoch, best_fitness, fi)
      
                  # Save model
                  if (not nosave) or (final_epoch and not evolve):  # if save
                      ckpt = {
                          "epoch": epoch,
                          "best_fitness": best_fitness,
                          "model": deepcopy(de_parallel(model)).half(),
                          "ema": deepcopy(ema.ema).half(),
                          "updates": ema.updates,
                          "optimizer": optimizer.state_dict(),
                          "opt": vars(opt),
                          "git": GIT_INFO,  # {remote, branch, commit} if a git repo
                          "date": datetime.now().isoformat(),
                      }
      
                      # Save last, best and delete
                      torch.save(ckpt, last)
                      if best_fitness == fi:
                          torch.save(ckpt, best)
                      if opt.save_period > 0 and epoch % opt.save_period == 0:
                          torch.save(ckpt, w / f"epoch{epoch}.pt")
                      del ckpt
                      callbacks.run("on_model_save", last, epoch, final_epoch, best_fitness, fi)
      
              # EarlyStopping
              if RANK != -1:  # if DDP training
                  broadcast_list = [stop if RANK == 0 else None]
                  dist.broadcast_object_list(broadcast_list, 0)  # broadcast 'stop' to all ranks
                  if RANK != 0:
                      stop = broadcast_list[0]
              if stop:
                  break  # must break all DDP ranks
      
              # end epoch ----------------------------------------------------------------------------------------------------
          # end training -----------------------------------------------------------------------------------------------------
          if RANK in {-1, 0}:
              LOGGER.info(f"\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.")
              for f in last, best:
                  if f.exists():
                      strip_optimizer(f)  # strip optimizers
                      if f is best:
                          LOGGER.info(f"\nValidating {f}...")
                          results, _, _ = validate.run(
                              data_dict,
                              batch_size=batch_size // WORLD_SIZE * 2,
                              imgsz=imgsz,
                              model=attempt_load(f, device).half(),
                              iou_thres=0.65 if is_coco else 0.60,  # best pycocotools at iou 0.65
                              single_cls=single_cls,
                              dataloader=val_loader,
                              save_dir=save_dir,
                              save_json=is_coco,
                              verbose=True,
                              plots=plots,
                              callbacks=callbacks,
                              compute_loss=compute_loss,
                          )  # val best model with plots
                          if is_coco:
                              callbacks.run("on_fit_epoch_end", list(mloss) + list(results) + lr, epoch, best_fitness, fi)
      
              callbacks.run("on_train_end", last, best, epoch, results)
      
          torch.cuda.empty_cache()
          return results
      
      
      def parse_opt(known=False):
          """Parses command-line arguments for YOLOv5 training, validation, and testing."""
          parser = argparse.ArgumentParser()
          parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")
          parser.add_argument("--cfg", type=str, default="", help="model.yaml path")
          parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
          parser.add_argument("--hyp", type=str, default=ROOT / "data/hyps/hyp.scratch-low.yaml", help="hyperparameters path")
          parser.add_argument("--epochs", type=int, default=100, help="total training epochs")
          parser.add_argument("--batch-size", type=int, default=16, help="total batch size for all GPUs, -1 for autobatch")
          parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="train, val image size (pixels)")
          parser.add_argument("--rect", action="store_true", help="rectangular training")
          parser.add_argument("--resume", nargs="?", const=True, default=False, help="resume most recent training")
          parser.add_argument("--nosave", action="store_true", help="only save final checkpoint")
          parser.add_argument("--noval", action="store_true", help="only validate final epoch")
          parser.add_argument("--noautoanchor", action="store_true", help="disable AutoAnchor")
          parser.add_argument("--noplots", action="store_true", help="save no plot files")
          parser.add_argument("--evolve", type=int, nargs="?", const=300, help="evolve hyperparameters for x generations")
          parser.add_argument(
              "--evolve_population", type=str, default=ROOT / "data/hyps", help="location for loading population"
          )
          parser.add_argument("--resume_evolve", type=str, default=None, help="resume evolve from last generation")
          parser.add_argument("--bucket", type=str, default="", help="gsutil bucket")
          parser.add_argument("--cache", type=str, nargs="?", const="ram", help="image --cache ram/disk")
          parser.add_argument("--image-weights", action="store_true", help="use weighted image selection for training")
          parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
          parser.add_argument("--multi-scale", action="store_true", help="vary img-size +/- 50%%")
          parser.add_argument("--single-cls", action="store_true", help="train multi-class data as single-class")
          parser.add_argument("--optimizer", type=str, choices=["SGD", "Adam", "AdamW"], default="SGD", help="optimizer")
          parser.add_argument("--sync-bn", action="store_true", help="use SyncBatchNorm, only available in DDP mode")
          parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
          parser.add_argument("--project", default=ROOT / "runs/train", help="save to project/name")
          parser.add_argument("--name", default="exp", help="save to project/name")
          parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
          parser.add_argument("--quad", action="store_true", help="quad dataloader")
          parser.add_argument("--cos-lr", action="store_true", help="cosine LR scheduler")
          parser.add_argument("--label-smoothing", type=float, default=0.0, help="Label smoothing epsilon")
          parser.add_argument("--patience", type=int, default=100, help="EarlyStopping patience (epochs without improvement)")
          parser.add_argument("--freeze", nargs="+", type=int, default=[0], help="Freeze layers: backbone=10, first3=0 1 2")
          parser.add_argument("--save-period", type=int, default=-1, help="Save checkpoint every x epochs (disabled if < 1)")
          parser.add_argument("--seed", type=int, default=0, help="Global training seed")
          parser.add_argument("--local_rank", type=int, default=-1, help="Automatic DDP Multi-GPU argument, do not modify")
      
          # Logger arguments
          parser.add_argument("--entity", default=None, help="Entity")
          parser.add_argument("--upload_dataset", nargs="?", const=True, default=False, help='Upload data, "val" option')
          parser.add_argument("--bbox_interval", type=int, default=-1, help="Set bounding-box image logging interval")
          parser.add_argument("--artifact_alias", type=str, default="latest", help="Version of dataset artifact to use")
      
          # NDJSON logging
          parser.add_argument("--ndjson-console", action="store_true", help="Log ndjson to console")
          parser.add_argument("--ndjson-file", action="store_true", help="Log ndjson to file")
      
          return parser.parse_known_args()[0] if known else parser.parse_args()
      
      
      def main(opt, callbacks=Callbacks()):
          """Runs training or hyperparameter evolution with specified options and optional callbacks."""
          if RANK in {-1, 0}:
              print_args(vars(opt))
              check_git_status()
              check_requirements(ROOT / "requirements.txt")
      
          # Resume (from specified or most recent last.pt)
          if opt.resume and not check_comet_resume(opt) and not opt.evolve:
              last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
              opt_yaml = last.parent.parent / "opt.yaml"  # train options yaml
              opt_data = opt.data  # original dataset
              if opt_yaml.is_file():
                  with open(opt_yaml, errors="ignore") as f:
                      d = yaml.safe_load(f)
              else:
                  d = torch.load(last, map_location="cpu")["opt"]
              opt = argparse.Namespace(**d)  # replace
              opt.cfg, opt.weights, opt.resume = "", str(last), True  # reinstate
              if is_url(opt_data):
                  opt.data = check_file(opt_data)  # avoid HUB resume auth timeout
          else:
              opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = (
                  check_file(opt.data),
                  check_yaml(opt.cfg),
                  check_yaml(opt.hyp),
                  str(opt.weights),
                  str(opt.project),
              )  # checks
              assert len(opt.cfg) or len(opt.weights), "either --cfg or --weights must be specified"
              if opt.evolve:
                  if opt.project == str(ROOT / "runs/train"):  # if default project name, rename to runs/evolve
                      opt.project = str(ROOT / "runs/evolve")
                  opt.exist_ok, opt.resume = opt.resume, False  # pass resume to exist_ok and disable resume
              if opt.name == "cfg":
                  opt.name = Path(opt.cfg).stem  # use model.yaml as name
              opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
      
          # DDP mode
          device = select_device(opt.device, batch_size=opt.batch_size)
          if LOCAL_RANK != -1:
              msg = "is not compatible with YOLOv5 Multi-GPU DDP training"
              assert not opt.image_weights, f"--image-weights {msg}"
              assert not opt.evolve, f"--evolve {msg}"
              assert opt.batch_size != -1, f"AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size"
              assert opt.batch_size % WORLD_SIZE == 0, f"--batch-size {opt.batch_size} must be multiple of WORLD_SIZE"
              assert torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command"
              torch.cuda.set_device(LOCAL_RANK)
              device = torch.device("cuda", LOCAL_RANK)
              dist.init_process_group(
                  backend="nccl" if dist.is_nccl_available() else "gloo", timeout=timedelta(seconds=10800)
              )
      
          # Train
          if not opt.evolve:
              train(opt.hyp, opt, device, callbacks)
      
          # Evolve hyperparameters (optional)
          else:
              # Hyperparameter evolution metadata (including this hyperparameter True-False, lower_limit, upper_limit)
              meta = {
                  "lr0": (False, 1e-5, 1e-1),  # initial learning rate (SGD=1E-2, Adam=1E-3)
                  "lrf": (False, 0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)
                  "momentum": (False, 0.6, 0.98),  # SGD momentum/Adam beta1
                  "weight_decay": (False, 0.0, 0.001),  # optimizer weight decay
                  "warmup_epochs": (False, 0.0, 5.0),  # warmup epochs (fractions ok)
                  "warmup_momentum": (False, 0.0, 0.95),  # warmup initial momentum
                  "warmup_bias_lr": (False, 0.0, 0.2),  # warmup initial bias lr
                  "box": (False, 0.02, 0.2),  # box loss gain
                  "cls": (False, 0.2, 4.0),  # cls loss gain
                  "cls_pw": (False, 0.5, 2.0),  # cls BCELoss positive_weight
                  "obj": (False, 0.2, 4.0),  # obj loss gain (scale with pixels)
                  "obj_pw": (False, 0.5, 2.0),  # obj BCELoss positive_weight
                  "iou_t": (False, 0.1, 0.7),  # IoU training threshold
                  "anchor_t": (False, 2.0, 8.0),  # anchor-multiple threshold
                  "anchors": (False, 2.0, 10.0),  # anchors per output grid (0 to ignore)
                  "fl_gamma": (False, 0.0, 2.0),  # focal loss gamma (efficientDet default gamma=1.5)
                  "hsv_h": (True, 0.0, 0.1),  # image HSV-Hue augmentation (fraction)
                  "hsv_s": (True, 0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
                  "hsv_v": (True, 0.0, 0.9),  # image HSV-Value augmentation (fraction)
                  "degrees": (True, 0.0, 45.0),  # image rotation (+/- deg)
                  "translate": (True, 0.0, 0.9),  # image translation (+/- fraction)
                  "scale": (True, 0.0, 0.9),  # image scale (+/- gain)
                  "shear": (True, 0.0, 10.0),  # image shear (+/- deg)
                  "perspective": (True, 0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001
                  "flipud": (True, 0.0, 1.0),  # image flip up-down (probability)
                  "fliplr": (True, 0.0, 1.0),  # image flip left-right (probability)
                  "mosaic": (True, 0.0, 1.0),  # image mixup (probability)
                  "mixup": (True, 0.0, 1.0),  # image mixup (probability)
                  "copy_paste": (True, 0.0, 1.0),
              }  # segment copy-paste (probability)
      
              # GA configs
              pop_size = 50
              mutation_rate_min = 0.01
              mutation_rate_max = 0.5
              crossover_rate_min = 0.5
              crossover_rate_max = 1
              min_elite_size = 2
              max_elite_size = 5
              tournament_size_min = 2
              tournament_size_max = 10
      
              with open(opt.hyp, errors="ignore") as f:
                  hyp = yaml.safe_load(f)  # load hyps dict
                  if "anchors" not in hyp:  # anchors commented in hyp.yaml
                      hyp["anchors"] = 3
              if opt.noautoanchor:
                  del hyp["anchors"], meta["anchors"]
              opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir)  # only val/save final epoch
              # ei = [isinstance(x, (int, float)) for x in hyp.values()]  # evolvable indices
              evolve_yaml, evolve_csv = save_dir / "hyp_evolve.yaml", save_dir / "evolve.csv"
              if opt.bucket:
                  # download evolve.csv if exists
                  subprocess.run(
                      [
                          "gsutil",
                          "cp",
                          f"gs://{opt.bucket}/evolve.csv",
                          str(evolve_csv),
                      ]
                  )
      
              # Delete the items in meta dictionary whose first value is False
              del_ = [item for item, value_ in meta.items() if value_[0] is False]
              hyp_GA = hyp.copy()  # Make a copy of hyp dictionary
              for item in del_:
                  del meta[item]  # Remove the item from meta dictionary
                  del hyp_GA[item]  # Remove the item from hyp_GA dictionary
      
              # Set lower_limit and upper_limit arrays to hold the search space boundaries
              lower_limit = np.array([meta[k][1] for k in hyp_GA.keys()])
              upper_limit = np.array([meta[k][2] for k in hyp_GA.keys()])
      
              # Create gene_ranges list to hold the range of values for each gene in the population
              gene_ranges = [(lower_limit[i], upper_limit[i]) for i in range(len(upper_limit))]
      
              # Initialize the population with initial_values or random values
              initial_values = []
      
              # If resuming evolution from a previous checkpoint
              if opt.resume_evolve is not None:
                  assert os.path.isfile(ROOT / opt.resume_evolve), "evolve population path is wrong!"
                  with open(ROOT / opt.resume_evolve, errors="ignore") as f:
                      evolve_population = yaml.safe_load(f)
                      for value in evolve_population.values():
                          value = np.array([value[k] for k in hyp_GA.keys()])
                          initial_values.append(list(value))
      
              # If not resuming from a previous checkpoint, generate initial values from .yaml files in opt.evolve_population
              else:
                  yaml_files = [f for f in os.listdir(opt.evolve_population) if f.endswith(".yaml")]
                  for file_name in yaml_files:
                      with open(os.path.join(opt.evolve_population, file_name)) as yaml_file:
                          value = yaml.safe_load(yaml_file)
                          value = np.array([value[k] for k in hyp_GA.keys()])
                          initial_values.append(list(value))
      
              # Generate random values within the search space for the rest of the population
              if initial_values is None:
                  population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size)]
              elif pop_size > 1:
                  population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size - len(initial_values))]
                  for initial_value in initial_values:
                      population = [initial_value] + population
      
              # Run the genetic algorithm for a fixed number of generations
              list_keys = list(hyp_GA.keys())
              for generation in range(opt.evolve):
                  if generation >= 1:
                      save_dict = {}
                      for i in range(len(population)):
                          little_dict = {list_keys[j]: float(population[i][j]) for j in range(len(population[i]))}
                          save_dict[f"gen{str(generation)}number{str(i)}"] = little_dict
      
                      with open(save_dir / "evolve_population.yaml", "w") as outfile:
                          yaml.dump(save_dict, outfile, default_flow_style=False)
      
                  # Adaptive elite size
                  elite_size = min_elite_size + int((max_elite_size - min_elite_size) * (generation / opt.evolve))
                  # Evaluate the fitness of each individual in the population
                  fitness_scores = []
                  for individual in population:
                      for key, value in zip(hyp_GA.keys(), individual):
                          hyp_GA[key] = value
                      hyp.update(hyp_GA)
                      results = train(hyp.copy(), opt, device, callbacks)
                      callbacks = Callbacks()
                      # Write mutation results
                      keys = (
                          "metrics/precision",
                          "metrics/recall",
                          "metrics/mAP_0.5",
                          "metrics/mAP_0.5:0.95",
                          "val/box_loss",
                          "val/obj_loss",
                          "val/cls_loss",
                      )
                      print_mutation(keys, results, hyp.copy(), save_dir, opt.bucket)
                      fitness_scores.append(results[2])
      
                  # Select the fittest individuals for reproduction using adaptive tournament selection
                  selected_indices = []
                  for _ in range(pop_size - elite_size):
                      # Adaptive tournament size
                      tournament_size = max(
                          max(2, tournament_size_min),
                          int(min(tournament_size_max, pop_size) - (generation / (opt.evolve / 10))),
                      )
                      # Perform tournament selection to choose the best individual
                      tournament_indices = random.sample(range(pop_size), tournament_size)
                      tournament_fitness = [fitness_scores[j] for j in tournament_indices]
                      winner_index = tournament_indices[tournament_fitness.index(max(tournament_fitness))]
                      selected_indices.append(winner_index)
      
                  # Add the elite individuals to the selected indices
                  elite_indices = [i for i in range(pop_size) if fitness_scores[i] in sorted(fitness_scores)[-elite_size:]]
                  selected_indices.extend(elite_indices)
                  # Create the next generation through crossover and mutation
                  next_generation = []
                  for _ in range(pop_size):
                      parent1_index = selected_indices[random.randint(0, pop_size - 1)]
                      parent2_index = selected_indices[random.randint(0, pop_size - 1)]
                      # Adaptive crossover rate
                      crossover_rate = max(
                          crossover_rate_min, min(crossover_rate_max, crossover_rate_max - (generation / opt.evolve))
                      )
                      if random.uniform(0, 1) < crossover_rate:
                          crossover_point = random.randint(1, len(hyp_GA) - 1)
                          child = population[parent1_index][:crossover_point] + population[parent2_index][crossover_point:]
                      else:
                          child = population[parent1_index]
                      # Adaptive mutation rate
                      mutation_rate = max(
                          mutation_rate_min, min(mutation_rate_max, mutation_rate_max - (generation / opt.evolve))
                      )
                      for j in range(len(hyp_GA)):
                          if random.uniform(0, 1) < mutation_rate:
                              child[j] += random.uniform(-0.1, 0.1)
                              child[j] = min(max(child[j], gene_ranges[j][0]), gene_ranges[j][1])
                      next_generation.append(child)
                  # Replace the old population with the new generation
                  population = next_generation
              # Print the best solution found
              best_index = fitness_scores.index(max(fitness_scores))
              best_individual = population[best_index]
              print("Best solution found:", best_individual)
              # Plot results
              plot_evolve(evolve_csv)
              LOGGER.info(
                  f'Hyperparameter evolution finished {opt.evolve} generations\n'
                  f"Results saved to {colorstr('bold', save_dir)}\n"
                  f'Usage example: $ python train.py --hyp {evolve_yaml}'
              )
      
      
      def generate_individual(input_ranges, individual_length):
          """Generates a list of random values within specified input ranges for each gene in the individual."""
          individual = []
          for i in range(individual_length):
              lower_bound, upper_bound = input_ranges[i]
              individual.append(random.uniform(lower_bound, upper_bound))
          return individual
      
      
      def run(**kwargs):
          """
          Executes YOLOv5 training with given options, overriding with any kwargs provided.
      
          Example: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
          """
          opt = parse_opt(True)
          for k, v in kwargs.items():
              setattr(opt, k, v)
          main(opt)
          return opt
      
      
      if __name__ == "__main__":
          opt = parse_opt()
          main(opt)
      
    • 📄 tutorial.ipynb 40.4KB
      [中文](https://docs.ultralytics.com/zh/) | [한국어](https://docs.ultralytics.com/ko/) | [日本語](https://docs.ultralytics.com/ja/) | [Русский](https://docs.ultralytics.com/ru/) | [Deutsch](https://docs.ultralytics.com/de/) | [Français](https://docs.ultralytics.com/fr/) | [Español](https://docs.ultralytics.com/es/) | [Português](https://docs.ultralytics.com/pt/) | [हिन्दी](https://docs.ultralytics.com/hi/) | [العربية](https://docs.ultralytics.com/ar/) Run on Gradient Open In Colab Open In Kaggle This YOLOv5 🚀 notebook by Ultralytics presents simple train, validate and predict examples to help start your AI adventure.
      We hope that the resources in this notebook will help you get the most out of YOLOv5. Please browse the YOLOv5 Docs for details, raise an issue on GitHub for support, and join our Discord community for questions and discussions!

      Setup

      Clone GitHub repository, install dependencies and check PyTorch and GPU.

      !git clone https://github.com/ultralytics/yolov5  # clone
      %cd yolov5
      %pip install -qr requirements.txt comet_ml  # install
      
      import torch
      import utils
      display = utils.notebook_init()  # checks
      
      YOLOv5 🚀 v7.0-136-g71244ae Python-3.9.16 torch-2.0.0+cu118 CUDA:0 (Tesla T4, 15102MiB)
      Setup complete ✅ (2 CPUs, 12.7 GB RAM, 23.3/166.8 GB disk)
      

      1. Detect

      detect.py runs YOLOv5 inference on a variety of sources, downloading models automatically from the latest YOLOv5 release, and saving results to runs/detect. Example inference sources are:

      python detect.py --source 0  # webcam
                                img.jpg  # image
                                vid.mp4  # video
                                screen  # screenshot
                                path/  # directory
                               'path/*.jpg'  # glob
                               'https://youtu.be/LNwODJXcvt4'  # YouTube
                               'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
      
      !python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source data/images
      # display.Image(filename='runs/detect/exp/zidane.jpg', width=600)
      
      detect: weights=['yolov5s.pt'], source=data/images, data=data/coco128.yaml, imgsz=[640, 640], conf_thres=0.25, iou_thres=0.45, max_det=1000, device=, view_img=False, save_txt=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=runs/detect, name=exp, exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, half=False, dnn=False, vid_stride=1
      YOLOv5 🚀 v7.0-136-g71244ae Python-3.9.16 torch-2.0.0+cu118 CUDA:0 (Tesla T4, 15102MiB)
      
      Downloading https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt to yolov5s.pt...
      100% 14.1M/14.1M [00:00<00:00, 24.5MB/s]
      
      Fusing layers... 
      YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients
      image 1/2 /content/yolov5/data/images/bus.jpg: 640x480 4 persons, 1 bus, 41.5ms
      image 2/2 /content/yolov5/data/images/zidane.jpg: 384x640 2 persons, 2 ties, 60.0ms
      Speed: 0.5ms pre-process, 50.8ms inference, 37.7ms NMS per image at shape (1, 3, 640, 640)
      Results saved to runs/detect/exp
      

              

      2. Validate

      Validate a model's accuracy on the COCO dataset's val or test splits. Models are downloaded automatically from the latest YOLOv5 release. To show results by class use the --verbose flag.

      # Download COCO val
      torch.hub.download_url_to_file('https://ultralytics.com/assets/coco2017val.zip', 'tmp.zip')  # download (780M - 5000 images)
      !unzip -q tmp.zip -d ../datasets && rm tmp.zip  # unzip
      
      100%|██████████| 780M/780M [00:12<00:00, 66.6MB/s]
      
      # Validate YOLOv5s on COCO val
      !python val.py --weights yolov5s.pt --data coco.yaml --img 640 --half
      
      val: data=/content/yolov5/data/coco.yaml, weights=['yolov5s.pt'], batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.6, max_det=300, task=val, device=, workers=8, single_cls=False, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=True, project=runs/val, name=exp, exist_ok=False, half=True, dnn=False
      YOLOv5 🚀 v7.0-136-g71244ae Python-3.9.16 torch-2.0.0+cu118 CUDA:0 (Tesla T4, 15102MiB)
      
      Fusing layers... 
      YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients
      val: Scanning /content/datasets/coco/val2017... 4952 images, 48 backgrounds, 0 corrupt: 100% 5000/5000 [00:02<00:00, 2024.59it/s]
      val: New cache created: /content/datasets/coco/val2017.cache
                       Class     Images  Instances          P          R      mAP50   mAP50-95: 100% 157/157 [01:25<00:00,  1.84it/s]
                         all       5000      36335      0.671      0.519      0.566      0.371
      Speed: 0.1ms pre-process, 3.1ms inference, 2.3ms NMS per image at shape (32, 3, 640, 640)
      
      Evaluating pycocotools mAP... saving runs/val/exp/yolov5s_predictions.json...
      loading annotations into memory...
      Done (t=0.43s)
      creating index...
      index created!
      Loading and preparing results...
      DONE (t=5.32s)
      creating index...
      index created!
      Running per image evaluation...
      Evaluate annotation type *bbox*
      DONE (t=78.89s).
      Accumulating evaluation results...
      DONE (t=14.51s).
       Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.374
       Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.572
       Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.402
       Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.211
       Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.423
       Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.489
       Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.311
       Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.516
       Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.566
       Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.378
       Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.625
       Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.722
      Results saved to runs/val/exp
      

      3. Train

      Close the active learning loop by sampling images from your inference conditions with the roboflow pip package


      Train a YOLOv5s model on the COCO128 dataset with --data coco128.yaml, starting from pretrained --weights yolov5s.pt, or from randomly initialized --weights '' --cfg yolov5s.yaml.

      A Mosaic Dataloader is used for training which combines 4 images into 1 mosaic.

      Label a dataset on Roboflow (optional)

      Roboflow enables you to easily organize, label, and prepare a high quality dataset with your own custom data. Roboflow also makes it easy to establish an active learning pipeline, collaborate with your team on dataset improvement, and integrate directly into your model building workflow with the roboflow pip package.

      #@title Select YOLOv5 🚀 logger {run: 'auto'}
      logger = 'Comet' #@param ['Comet', 'ClearML', 'TensorBoard']
      
      if logger == 'Comet':
        %pip install -q comet_ml
        import comet_ml; comet_ml.init()
      elif logger == 'ClearML':
        %pip install -q clearml
        import clearml; clearml.browser_login()
      elif logger == 'TensorBoard':
        %load_ext tensorboard
        %tensorboard --logdir runs/train
      
      # Train YOLOv5s on COCO128 for 3 epochs
      !python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt --cache
      
      2023-04-09 14:11:38.063605: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
      To enable the following instructions: AVX2 AVX512F FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
      2023-04-09 14:11:39.026661: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
      train: weights=yolov5s.pt, cfg=, data=coco128.yaml, hyp=data/hyps/hyp.scratch-low.yaml, epochs=3, batch_size=16, imgsz=640, rect=False, resume=False, nosave=False, noval=False, noautoanchor=False, noplots=False, evolve=None, bucket=, cache=ram, image_weights=False, device=, multi_scale=False, single_cls=False, optimizer=SGD, sync_bn=False, workers=8, project=runs/train, name=exp, exist_ok=False, quad=False, cos_lr=False, label_smoothing=0.0, patience=100, freeze=[0], save_period=-1, seed=0, local_rank=-1, entity=None, upload_dataset=False, bbox_interval=-1, artifact_alias=latest
      github: up to date with https://github.com/ultralytics/yolov5 ✅
      YOLOv5 🚀 v7.0-136-g71244ae Python-3.9.16 torch-2.0.0+cu118 CUDA:0 (Tesla T4, 15102MiB)
      
      hyperparameters: lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=0.05, cls=0.5, cls_pw=1.0, obj=1.0, obj_pw=1.0, iou_t=0.2, anchor_t=4.0, fl_gamma=0.0, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, mosaic=1.0, mixup=0.0, copy_paste=0.0
      ClearML: run 'pip install clearml' to automatically track, visualize and remotely train YOLOv5 🚀 in ClearML
      Comet: run 'pip install comet_ml' to automatically track and visualize YOLOv5 🚀 runs in Comet
      TensorBoard: Start with 'tensorboard --logdir runs/train', view at http://localhost:6006/
      
      Dataset not found ⚠️, missing paths ['/content/datasets/coco128/images/train2017']
      Downloading https://ultralytics.com/assets/coco128.zip to coco128.zip...
      100% 6.66M/6.66M [00:00<00:00, 75.6MB/s]
      Dataset download success ✅ (0.6s), saved to /content/datasets
      
                       from  n    params  module                                  arguments                     
        0                -1  1      3520  models.common.Conv                      [3, 32, 6, 2, 2]              
        1                -1  1     18560  models.common.Conv                      [32, 64, 3, 2]                
        2                -1  1     18816  models.common.C3                        [64, 64, 1]                   
        3                -1  1     73984  models.common.Conv                      [64, 128, 3, 2]               
        4                -1  2    115712  models.common.C3                        [128, 128, 2]                 
        5                -1  1    295424  models.common.Conv                      [128, 256, 3, 2]              
        6                -1  3    625152  models.common.C3                        [256, 256, 3]                 
        7                -1  1   1180672  models.common.Conv                      [256, 512, 3, 2]              
        8                -1  1   1182720  models.common.C3                        [512, 512, 1]                 
        9                -1  1    656896  models.common.SPPF                      [512, 512, 5]                 
       10                -1  1    131584  models.common.Conv                      [512, 256, 1, 1]              
       11                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
       12           [-1, 6]  1         0  models.common.Concat                    [1]                           
       13                -1  1    361984  models.common.C3                        [512, 256, 1, False]          
       14                -1  1     33024  models.common.Conv                      [256, 128, 1, 1]              
       15                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
       16           [-1, 4]  1         0  models.common.Concat                    [1]                           
       17                -1  1     90880  models.common.C3                        [256, 128, 1, False]          
       18                -1  1    147712  models.common.Conv                      [128, 128, 3, 2]              
       19          [-1, 14]  1         0  models.common.Concat                    [1]                           
       20                -1  1    296448  models.common.C3                        [256, 256, 1, False]          
       21                -1  1    590336  models.common.Conv                      [256, 256, 3, 2]              
       22          [-1, 10]  1         0  models.common.Concat                    [1]                           
       23                -1  1   1182720  models.common.C3                        [512, 512, 1, False]          
       24      [17, 20, 23]  1    229245  models.yolo.Detect                      [80, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]
      Model summary: 214 layers, 7235389 parameters, 7235389 gradients, 16.6 GFLOPs
      
      Transferred 349/349 items from yolov5s.pt
      AMP: checks passed ✅
      optimizer: SGD(lr=0.01) with parameter groups 57 weight(decay=0.0), 60 weight(decay=0.0005), 60 bias
      albumentations: Blur(p=0.01, blur_limit=(3, 7)), MedianBlur(p=0.01, blur_limit=(3, 7)), ToGray(p=0.01), CLAHE(p=0.01, clip_limit=(1, 4.0), tile_grid_size=(8, 8))
      train: Scanning /content/datasets/coco128/labels/train2017... 126 images, 2 backgrounds, 0 corrupt: 100% 128/128 [00:00<00:00, 1709.36it/s]
      train: New cache created: /content/datasets/coco128/labels/train2017.cache
      train: Caching images (0.1GB ram): 100% 128/128 [00:00<00:00, 264.35it/s]
      val: Scanning /content/datasets/coco128/labels/train2017.cache... 126 images, 2 backgrounds, 0 corrupt: 100% 128/128 [00:00<?, ?it/s]
      val: Caching images (0.1GB ram): 100% 128/128 [00:01<00:00, 107.05it/s]
      
      AutoAnchor: 4.27 anchors/target, 0.994 Best Possible Recall (BPR). Current anchors are a good fit to dataset ✅
      Plotting labels to runs/train/exp/labels.jpg... 
      Image sizes 640 train, 640 val
      Using 2 dataloader workers
      Logging results to runs/train/exp
      Starting training for 3 epochs...
      
            Epoch    GPU_mem   box_loss   obj_loss   cls_loss  Instances       Size
              0/2      3.91G    0.04618    0.07209    0.01703        232        640: 100% 8/8 [00:09<00:00,  1.17s/it]
                       Class     Images  Instances          P          R      mAP50   mAP50-95: 100% 4/4 [00:01<00:00,  2.01it/s]
                         all        128        929      0.667      0.602       0.68       0.45
      
            Epoch    GPU_mem   box_loss   obj_loss   cls_loss  Instances       Size
              1/2      4.76G    0.04622    0.06891    0.01817        201        640: 100% 8/8 [00:02<00:00,  3.78it/s]
                       Class     Images  Instances          P          R      mAP50   mAP50-95: 100% 4/4 [00:01<00:00,  2.16it/s]
                         all        128        929      0.709      0.645      0.722      0.478
      
            Epoch    GPU_mem   box_loss   obj_loss   cls_loss  Instances       Size
              2/2      4.76G     0.0436     0.0647    0.01698        227        640: 100% 8/8 [00:01<00:00,  4.19it/s]
                       Class     Images  Instances          P          R      mAP50   mAP50-95: 100% 4/4 [00:01<00:00,  2.95it/s]
                         all        128        929      0.761      0.647      0.735       0.49
      
      3 epochs completed in 0.006 hours.
      Optimizer stripped from runs/train/exp/weights/last.pt, 14.8MB
      Optimizer stripped from runs/train/exp/weights/best.pt, 14.8MB
      
      Validating runs/train/exp/weights/best.pt...
      Fusing layers... 
      Model summary: 157 layers, 7225885 parameters, 0 gradients, 16.4 GFLOPs
                       Class     Images  Instances          P          R      mAP50   mAP50-95: 100% 4/4 [00:06<00:00,  1.56s/it]
                         all        128        929      0.759      0.646      0.734       0.49
                      person        128        254      0.857      0.706      0.805      0.525
                     bicycle        128          6      0.773      0.577      0.725      0.414
                         car        128         46      0.664      0.435      0.551       0.24
                  motorcycle        128          5      0.587        0.8      0.837      0.635
                    airplane        128          6          1      0.989      0.995      0.715
                         bus        128          7      0.635      0.714      0.753      0.651
                       train        128          3      0.686      0.333       0.72      0.504
                       truck        128         12      0.604      0.333      0.472      0.259
                        boat        128          6      0.938      0.333      0.449      0.177
               traffic light        128         14      0.778      0.255      0.401      0.217
                   stop sign        128          2      0.826          1      0.995      0.895
                       bench        128          9      0.711      0.556      0.661      0.313
                        bird        128         16      0.962          1      0.995      0.642
                         cat        128          4      0.868          1      0.995      0.754
                         dog        128          9          1      0.652      0.899      0.651
                       horse        128          2      0.853          1      0.995      0.622
                    elephant        128         17      0.909      0.882      0.934      0.698
                        bear        128          1      0.696          1      0.995      0.995
                       zebra        128          4      0.855          1      0.995      0.905
                     giraffe        128          9      0.788      0.828      0.912      0.701
                    backpack        128          6      0.835        0.5      0.738      0.311
                    umbrella        128         18      0.785      0.814      0.859       0.48
                     handbag        128         19      0.759      0.263      0.366      0.205
                         tie        128          7      0.983      0.714       0.77      0.492
                    suitcase        128          4      0.656          1      0.945      0.631
                     frisbee        128          5      0.721        0.8      0.759      0.724
                        skis        128          1      0.737          1      0.995        0.3
                   snowboard        128          7      0.829      0.696       0.83      0.537
                 sports ball        128          6      0.637      0.667      0.602      0.311
                        kite        128         10      0.636        0.6      0.599      0.226
                baseball bat        128          4      0.501       0.25      0.468      0.205
              baseball glove        128          7      0.483      0.429      0.465      0.292
                  skateboard        128          5      0.932        0.6      0.687      0.493
               tennis racket        128          7       0.77      0.429      0.547      0.332
                      bottle        128         18      0.577      0.379      0.554      0.276
                  wine glass        128         16      0.704      0.875       0.89       0.51
                         cup        128         36      0.841      0.667      0.837      0.533
                        fork        128          6      0.992      0.333       0.45      0.315
                       knife        128         16      0.768      0.688      0.695      0.403
                       spoon        128         22      0.838       0.47      0.639      0.384
                        bowl        128         28      0.764       0.58      0.716      0.513
                      banana        128          1      0.902          1      0.995      0.301
                    sandwich        128          2          1          0      0.359      0.326
                      orange        128          4      0.722       0.75      0.912      0.581
                    broccoli        128         11      0.547      0.364      0.432      0.317
                      carrot        128         24      0.619      0.625      0.724      0.495
                     hot dog        128          2      0.409          1      0.828      0.762
                       pizza        128          5      0.833      0.995      0.962      0.727
                       donut        128         14      0.631          1       0.96      0.839
                        cake        128          4       0.87          1      0.995       0.83
                       chair        128         35      0.583        0.6      0.608      0.317
                       couch        128          6      0.907      0.667      0.815      0.544
                potted plant        128         14      0.739      0.786      0.823       0.48
                         bed        128          3      0.985      0.333       0.83      0.441
                dining table        128         13      0.821      0.357      0.578      0.342
                      toilet        128          2          1      0.988      0.995      0.846
                          tv        128          2       0.57          1      0.995      0.796
                      laptop        128          3          1          0      0.593      0.312
                       mouse        128          2          1          0      0.089     0.0445
                      remote        128          8          1      0.624      0.634      0.538
                  cell phone        128          8      0.622      0.417      0.421      0.187
                   microwave        128          3      0.711          1      0.995      0.766
                        oven        128          5      0.329        0.4       0.43      0.282
                        sink        128          6      0.437      0.333      0.338      0.265
                refrigerator        128          5      0.567        0.8      0.799      0.536
                        book        128         29      0.597      0.257      0.349      0.154
                       clock        128          9      0.765      0.889      0.932      0.736
                        vase        128          2       0.33          1      0.995      0.895
                    scissors        128          1          1          0      0.497     0.0498
                  teddy bear        128         21      0.856      0.569      0.841      0.547
                  toothbrush        128          5        0.8          1      0.928      0.574
      Results saved to runs/train/exp
      

      4. Visualize

      Comet Logging and Visualization 🌟 NEW

      Comet is now fully integrated with YOLOv5. Track and visualize model metrics in real time, save your hyperparameters, datasets, and model checkpoints, and visualize your model predictions with Comet Custom Panels! Comet makes sure you never lose track of your work and makes it easy to share results and collaborate across teams of all sizes!

      Getting started is easy:

      pip install comet_ml  # 1. install
      export COMET_API_KEY=<Your API Key>  # 2. paste API key
      python train.py --img 640 --epochs 3 --data coco128.yaml --weights yolov5s.pt  # 3. train
      

      To learn more about all of the supported Comet features for this integration, check out the Comet Tutorial. If you'd like to learn more about Comet, head over to our documentation. Get started by trying out the Comet Colab Notebook:
      Open In Colab


      Comet Dashboard

      ClearML Logging and Automation 🌟 NEW

      ClearML is completely integrated into YOLOv5 to track your experimentation, manage dataset versions and even remotely execute training runs. To enable ClearML (check cells above):

      You'll get all the great expected features from an experiment manager: live updates, model upload, experiment comparison etc. but ClearML also tracks uncommitted changes and installed packages for example. Thanks to that ClearML Tasks (which is what we call experiments) are also reproducible on different machines! With only 1 extra line, we can schedule a YOLOv5 training task on a queue to be executed by any number of ClearML Agents (workers).

      You can use ClearML Data to version your dataset and then pass it to YOLOv5 simply using its unique ID. This will help you keep track of your data without adding extra hassle. Explore the ClearML Tutorial for details!


      ClearML Experiment Management UI

      Local Logging

      Training results are automatically logged with Tensorboard and CSV loggers to runs/train, with a new experiment directory created for each new training as runs/train/exp2, runs/train/exp3, etc.

      This directory contains train and val statistics, mosaics, labels, predictions and augmentated mosaics, as well as metrics and charts including precision-recall (PR) curves and confusion matrices.

      Local logging results

      Environments

      YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

      Status

      YOLOv5 CI

      If this badge is green, all YOLOv5 GitHub Actions Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training (train.py), testing (val.py), inference (detect.py) and export (export.py) on macOS, Windows, and Ubuntu every 24 hours and on every commit.

      Appendix

      Additional content below.

      # YOLOv5 PyTorch HUB Inference (DetectionModels only)
      import torch
      
      model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True, trust_repo=True)  # or yolov5n - yolov5x6 or custom
      im = 'https://ultralytics.com/images/zidane.jpg'  # file, Path, PIL.Image, OpenCV, nparray, list
      results = model(im)  # inference
      results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
      
    • 📄 val.py 20.2KB
      # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
      """
      Validate a trained YOLOv5 detection model on a detection dataset.
      
      Usage:
          $ python val.py --weights yolov5s.pt --data coco128.yaml --img 640
      
      Usage - formats:
          $ python val.py --weights yolov5s.pt                 # PyTorch
                                    yolov5s.torchscript        # TorchScript
                                    yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                    yolov5s_openvino_model     # OpenVINO
                                    yolov5s.engine             # TensorRT
                                    yolov5s.mlmodel            # CoreML (macOS-only)
                                    yolov5s_saved_model        # TensorFlow SavedModel
                                    yolov5s.pb                 # TensorFlow GraphDef
                                    yolov5s.tflite             # TensorFlow Lite
                                    yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
                                    yolov5s_paddle_model       # PaddlePaddle
      """
      
      import argparse
      import json
      import os
      import subprocess
      import sys
      from pathlib import Path
      
      import numpy as np
      import torch
      from tqdm import tqdm
      
      FILE = Path(__file__).resolve()
      ROOT = FILE.parents[0]  # YOLOv5 root directory
      if str(ROOT) not in sys.path:
          sys.path.append(str(ROOT))  # add ROOT to PATH
      ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
      
      from models.common import DetectMultiBackend
      from utils.callbacks import Callbacks
      from utils.dataloaders import create_dataloader
      from utils.general import (
          LOGGER,
          TQDM_BAR_FORMAT,
          Profile,
          check_dataset,
          check_img_size,
          check_requirements,
          check_yaml,
          coco80_to_coco91_class,
          colorstr,
          increment_path,
          non_max_suppression,
          print_args,
          scale_boxes,
          xywh2xyxy,
          xyxy2xywh,
      )
      from utils.metrics import ConfusionMatrix, ap_per_class, box_iou
      from utils.plots import output_to_target, plot_images, plot_val_study
      from utils.torch_utils import select_device, smart_inference_mode
      
      
      def save_one_txt(predn, save_conf, shape, file):
          """Saves one detection result to a txt file in normalized xywh format, optionally including confidence."""
          gn = torch.tensor(shape)[[1, 0, 1, 0]]  # normalization gain whwh
          for *xyxy, conf, cls in predn.tolist():
              xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
              line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
              with open(file, "a") as f:
                  f.write(("%g " * len(line)).rstrip() % line + "\n")
      
      
      def save_one_json(predn, jdict, path, class_map):
          """
          Saves one JSON detection result with image ID, category ID, bounding box, and score.
      
          Example: {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}
          """
          image_id = int(path.stem) if path.stem.isnumeric() else path.stem
          box = xyxy2xywh(predn[:, :4])  # xywh
          box[:, :2] -= box[:, 2:] / 2  # xy center to top-left corner
          for p, b in zip(predn.tolist(), box.tolist()):
              jdict.append(
                  {
                      "image_id": image_id,
                      "category_id": class_map[int(p[5])],
                      "bbox": [round(x, 3) for x in b],
                      "score": round(p[4], 5),
                  }
              )
      
      
      def process_batch(detections, labels, iouv):
          """
          Return correct prediction matrix.
      
          Arguments:
              detections (array[N, 6]), x1, y1, x2, y2, conf, class
              labels (array[M, 5]), class, x1, y1, x2, y2
          Returns:
              correct (array[N, 10]), for 10 IoU levels
          """
          correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
          iou = box_iou(labels[:, 1:], detections[:, :4])
          correct_class = labels[:, 0:1] == detections[:, 5]
          for i in range(len(iouv)):
              x = torch.where((iou >= iouv[i]) & correct_class)  # IoU > threshold and classes match
              if x[0].shape[0]:
                  matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()  # [label, detect, iou]
                  if x[0].shape[0] > 1:
                      matches = matches[matches[:, 2].argsort()[::-1]]
                      matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
                      # matches = matches[matches[:, 2].argsort()[::-1]]
                      matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
                  correct[matches[:, 1].astype(int), i] = True
          return torch.tensor(correct, dtype=torch.bool, device=iouv.device)
      
      
      @smart_inference_mode()
      def run(
          data,
          weights=None,  # model.pt path(s)
          batch_size=32,  # batch size
          imgsz=640,  # inference size (pixels)
          conf_thres=0.001,  # confidence threshold
          iou_thres=0.6,  # NMS IoU threshold
          max_det=300,  # maximum detections per image
          task="val",  # train, val, test, speed or study
          device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
          workers=8,  # max dataloader workers (per RANK in DDP mode)
          single_cls=False,  # treat as single-class dataset
          augment=False,  # augmented inference
          verbose=False,  # verbose output
          save_txt=False,  # save results to *.txt
          save_hybrid=False,  # save label+prediction hybrid results to *.txt
          save_conf=False,  # save confidences in --save-txt labels
          save_json=False,  # save a COCO-JSON results file
          project=ROOT / "runs/val",  # save to project/name
          name="exp",  # save to project/name
          exist_ok=False,  # existing project/name ok, do not increment
          half=True,  # use FP16 half-precision inference
          dnn=False,  # use OpenCV DNN for ONNX inference
          model=None,
          dataloader=None,
          save_dir=Path(""),
          plots=True,
          callbacks=Callbacks(),
          compute_loss=None,
      ):
          # Initialize/load model and set device
          training = model is not None
          if training:  # called by train.py
              device, pt, jit, engine = next(model.parameters()).device, True, False, False  # get model device, PyTorch model
              half &= device.type != "cpu"  # half precision only supported on CUDA
              model.half() if half else model.float()
          else:  # called directly
              device = select_device(device, batch_size=batch_size)
      
              # Directories
              save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
              (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
      
              # Load model
              model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
              stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
              imgsz = check_img_size(imgsz, s=stride)  # check image size
              half = model.fp16  # FP16 supported on limited backends with CUDA
              if engine:
                  batch_size = model.batch_size
              else:
                  device = model.device
                  if not (pt or jit):
                      batch_size = 1  # export.py models default to batch-size 1
                      LOGGER.info(f"Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models")
      
              # Data
              data = check_dataset(data)  # check
      
          # Configure
          model.eval()
          cuda = device.type != "cpu"
          is_coco = isinstance(data.get("val"), str) and data["val"].endswith(f"coco{os.sep}val2017.txt")  # COCO dataset
          nc = 1 if single_cls else int(data["nc"])  # number of classes
          iouv = torch.linspace(0.5, 0.95, 10, device=device)  # iou vector for mAP@0.5:0.95
          niou = iouv.numel()
      
          # Dataloader
          if not training:
              if pt and not single_cls:  # check --weights are trained on --data
                  ncm = model.model.nc
                  assert ncm == nc, (
                      f"{weights} ({ncm} classes) trained on different --data than what you passed ({nc} "
                      f"classes). Pass correct combination of --weights and --data that are trained together."
                  )
              model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz))  # warmup
              pad, rect = (0.0, False) if task == "speed" else (0.5, pt)  # square inference for benchmarks
              task = task if task in ("train", "val", "test") else "val"  # path to train/val/test images
              dataloader = create_dataloader(
                  data[task],
                  imgsz,
                  batch_size,
                  stride,
                  single_cls,
                  pad=pad,
                  rect=rect,
                  workers=workers,
                  prefix=colorstr(f"{task}: "),
              )[0]
      
          seen = 0
          confusion_matrix = ConfusionMatrix(nc=nc)
          names = model.names if hasattr(model, "names") else model.module.names  # get class names
          if isinstance(names, (list, tuple)):  # old format
              names = dict(enumerate(names))
          class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
          s = ("%22s" + "%11s" * 6) % ("Class", "Images", "Instances", "P", "R", "mAP50", "mAP50-95")
          tp, fp, p, r, f1, mp, mr, map50, ap50, map = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
          dt = Profile(device=device), Profile(device=device), Profile(device=device)  # profiling times
          loss = torch.zeros(3, device=device)
          jdict, stats, ap, ap_class = [], [], [], []
          callbacks.run("on_val_start")
          pbar = tqdm(dataloader, desc=s, bar_format=TQDM_BAR_FORMAT)  # progress bar
          for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
              callbacks.run("on_val_batch_start")
              with dt[0]:
                  if cuda:
                      im = im.to(device, non_blocking=True)
                      targets = targets.to(device)
                  im = im.half() if half else im.float()  # uint8 to fp16/32
                  im /= 255  # 0 - 255 to 0.0 - 1.0
                  nb, _, height, width = im.shape  # batch size, channels, height, width
      
              # Inference
              with dt[1]:
                  preds, train_out = model(im) if compute_loss else (model(im, augment=augment), None)
      
              # Loss
              if compute_loss:
                  loss += compute_loss(train_out, targets)[1]  # box, obj, cls
      
              # NMS
              targets[:, 2:] *= torch.tensor((width, height, width, height), device=device)  # to pixels
              lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else []  # for autolabelling
              with dt[2]:
                  preds = non_max_suppression(
                      preds, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls, max_det=max_det
                  )
      
              # Metrics
              for si, pred in enumerate(preds):
                  labels = targets[targets[:, 0] == si, 1:]
                  nl, npr = labels.shape[0], pred.shape[0]  # number of labels, predictions
                  path, shape = Path(paths[si]), shapes[si][0]
                  correct = torch.zeros(npr, niou, dtype=torch.bool, device=device)  # init
                  seen += 1
      
                  if npr == 0:
                      if nl:
                          stats.append((correct, *torch.zeros((2, 0), device=device), labels[:, 0]))
                          if plots:
                              confusion_matrix.process_batch(detections=None, labels=labels[:, 0])
                      continue
      
                  # Predictions
                  if single_cls:
                      pred[:, 5] = 0
                  predn = pred.clone()
                  scale_boxes(im[si].shape[1:], predn[:, :4], shape, shapes[si][1])  # native-space pred
      
                  # Evaluate
                  if nl:
                      tbox = xywh2xyxy(labels[:, 1:5])  # target boxes
                      scale_boxes(im[si].shape[1:], tbox, shape, shapes[si][1])  # native-space labels
                      labelsn = torch.cat((labels[:, 0:1], tbox), 1)  # native-space labels
                      correct = process_batch(predn, labelsn, iouv)
                      if plots:
                          confusion_matrix.process_batch(predn, labelsn)
                  stats.append((correct, pred[:, 4], pred[:, 5], labels[:, 0]))  # (correct, conf, pcls, tcls)
      
                  # Save/log
                  if save_txt:
                      (save_dir / "labels").mkdir(parents=True, exist_ok=True)
                      save_one_txt(predn, save_conf, shape, file=save_dir / "labels" / f"{path.stem}.txt")
                  if save_json:
                      save_one_json(predn, jdict, path, class_map)  # append to COCO-JSON dictionary
                  callbacks.run("on_val_image_end", pred, predn, path, names, im[si])
      
              # Plot images
              if plots and batch_i < 3:
                  plot_images(im, targets, paths, save_dir / f"val_batch{batch_i}_labels.jpg", names)  # labels
                  plot_images(im, output_to_target(preds), paths, save_dir / f"val_batch{batch_i}_pred.jpg", names)  # pred
      
              callbacks.run("on_val_batch_end", batch_i, im, targets, paths, shapes, preds)
      
          # Compute metrics
          stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)]  # to numpy
          if len(stats) and stats[0].any():
              tp, fp, p, r, f1, ap, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
              ap50, ap = ap[:, 0], ap.mean(1)  # AP@0.5, AP@0.5:0.95
              mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
          nt = np.bincount(stats[3].astype(int), minlength=nc)  # number of targets per class
      
          # Print results
          pf = "%22s" + "%11i" * 2 + "%11.3g" * 4  # print format
          LOGGER.info(pf % ("all", seen, nt.sum(), mp, mr, map50, map))
          if nt.sum() == 0:
              LOGGER.warning(f"WARNING ⚠️ no labels found in {task} set, can not compute metrics without labels")
      
          # Print results per class
          if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
              for i, c in enumerate(ap_class):
                  LOGGER.info(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
      
          # Print speeds
          t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
          if not training:
              shape = (batch_size, 3, imgsz, imgsz)
              LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}" % t)
      
          # Plots
          if plots:
              confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
              callbacks.run("on_val_end", nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix)
      
          # Save JSON
          if save_json and len(jdict):
              w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else ""  # weights
              anno_json = str(Path("../datasets/coco/annotations/instances_val2017.json"))  # annotations
              if not os.path.exists(anno_json):
                  anno_json = os.path.join(data["path"], "annotations", "instances_val2017.json")
              pred_json = str(save_dir / f"{w}_predictions.json")  # predictions
              LOGGER.info(f"\nEvaluating pycocotools mAP... saving {pred_json}...")
              with open(pred_json, "w") as f:
                  json.dump(jdict, f)
      
              try:  # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
                  check_requirements("pycocotools>=2.0.6")
                  from pycocotools.coco import COCO
                  from pycocotools.cocoeval import COCOeval
      
                  anno = COCO(anno_json)  # init annotations api
                  pred = anno.loadRes(pred_json)  # init predictions api
                  eval = COCOeval(anno, pred, "bbox")
                  if is_coco:
                      eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.im_files]  # image IDs to evaluate
                  eval.evaluate()
                  eval.accumulate()
                  eval.summarize()
                  map, map50 = eval.stats[:2]  # update results (mAP@0.5:0.95, mAP@0.5)
              except Exception as e:
                  LOGGER.info(f"pycocotools unable to run: {e}")
      
          # Return results
          model.float()  # for training
          if not training:
              s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
              LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
          maps = np.zeros(nc) + map
          for i, c in enumerate(ap_class):
              maps[c] = ap[i]
          return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
      
      
      def parse_opt():
          """Parses command-line options for YOLOv5 model inference configuration."""
          parser = argparse.ArgumentParser()
          parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
          parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model path(s)")
          parser.add_argument("--batch-size", type=int, default=32, help="batch size")
          parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="inference size (pixels)")
          parser.add_argument("--conf-thres", type=float, default=0.001, help="confidence threshold")
          parser.add_argument("--iou-thres", type=float, default=0.6, help="NMS IoU threshold")
          parser.add_argument("--max-det", type=int, default=300, help="maximum detections per image")
          parser.add_argument("--task", default="val", help="train, val, test, speed or study")
          parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
          parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
          parser.add_argument("--single-cls", action="store_true", help="treat as single-class dataset")
          parser.add_argument("--augment", action="store_true", help="augmented inference")
          parser.add_argument("--verbose", action="store_true", help="report mAP by class")
          parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
          parser.add_argument("--save-hybrid", action="store_true", help="save label+prediction hybrid results to *.txt")
          parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels")
          parser.add_argument("--save-json", action="store_true", help="save a COCO-JSON results file")
          parser.add_argument("--project", default=ROOT / "runs/val", help="save to project/name")
          parser.add_argument("--name", default="exp", help="save to project/name")
          parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
          parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
          parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
          opt = parser.parse_args()
          opt.data = check_yaml(opt.data)  # check YAML
          opt.save_json |= opt.data.endswith("coco.yaml")
          opt.save_txt |= opt.save_hybrid
          print_args(vars(opt))
          return opt
      
      
      def main(opt):
          """Executes YOLOv5 tasks like training, validation, testing, speed, and study benchmarks based on provided
          options.
          """
          check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
      
          if opt.task in ("train", "val", "test"):  # run normally
              if opt.conf_thres > 0.001:  # https://github.com/ultralytics/yolov5/issues/1466
                  LOGGER.info(f"WARNING ⚠️ confidence threshold {opt.conf_thres} > 0.001 produces invalid results")
              if opt.save_hybrid:
                  LOGGER.info("WARNING ⚠️ --save-hybrid will return high mAP from hybrid labels, not from predictions alone")
              run(**vars(opt))
      
          else:
              weights = opt.weights if isinstance(opt.weights, list) else [opt.weights]
              opt.half = torch.cuda.is_available() and opt.device != "cpu"  # FP16 for fastest results
              if opt.task == "speed":  # speed benchmarks
                  # python val.py --task speed --data coco.yaml --batch 1 --weights yolov5n.pt yolov5s.pt...
                  opt.conf_thres, opt.iou_thres, opt.save_json = 0.25, 0.45, False
                  for opt.weights in weights:
                      run(**vars(opt), plots=False)
      
              elif opt.task == "study":  # speed vs mAP benchmarks
                  # python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n.pt yolov5s.pt...
                  for opt.weights in weights:
                      f = f"study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt"  # filename to save to
                      x, y = list(range(256, 1536 + 128, 128)), []  # x axis (image sizes), y axis
                      for opt.imgsz in x:  # img-size
                          LOGGER.info(f"\nRunning {f} --imgsz {opt.imgsz}...")
                          r, _, t = run(**vars(opt), plots=False)
                          y.append(r + t)  # results and times
                      np.savetxt(f, y, fmt="%10.4g")  # save
                  subprocess.run(["zip", "-r", "study.zip", "study_*.txt"])
                  plot_val_study(x=x)  # plot
              else:
                  raise NotImplementedError(f'--task {opt.task} not in ("train", "val", "test", "speed", "study")')
      
      
      if __name__ == "__main__":
          opt = parse_opt()
          main(opt)
      
  • 📄 01.iou.py 1011B
    import numpy as np
    
    
    # 1.定义方法计算IOU
    def Iou(box1, box2):
        # 1.1 使用极坐标形式表示:直接获取两个bbox的坐标,获取两个框的左上角坐标和右下角坐标
        xmin1, ymin1, xmax1, ymax1 = box1
        xmin2, ymin2, xmax2, ymax2 = box2
        # 1.2 获取矩形框交集对应的左上角和右下角的坐标(intersection)
        xx1 = np.max([xmin1, xmin2])
        yy1 = np.max([ymin1, ymin2])
        xx2 = np.min([xmax1, xmax2])
        yy2 = np.min([ymax1, ymax2])
        # 1.3 计算两个矩形框面积
        area1 = (xmax1 - xmin1) * (ymax1 - ymin1)
        area2 = (xmax2 - xmin2) * (ymax2 - ymin2)
        # 1.4 计算交集面积
        # inter_area = (np.max([0, xx2 - xx1])) * (np.max([0, yy2 - yy1]))
        inter_area = (xx2 - xx1) * (yy2 - yy1)
        # 1.5 计算交并比
        iou = inter_area / (area1 + area2 - inter_area)
        return iou
    
    
    if __name__ == "__main__":
        True_bbox, predict_bbox = [100, 35, 398, 400], [40, 150, 355, 398]
        iou = Iou(True_bbox, predict_bbox)
        print(iou)
    
  • 📄 02.nms.py 2.3KB
    import numpy as np
    
    
    def nms(bboxes, confidence_score, threshold):
        """非极大抑制过程
        :param bboxes: 同类别候选框坐标
        :param confidence: 同类别候选框分数
        :param threshold: iou阈值
        :return:
        """
        # 1.强转数组,并取出每个框左上角坐标和右下角坐标
        # 1.1 强转数组
        bboxes = np.array(bboxes)
        score = np.array(confidence_score)
    
        # 1.2 取出n个的极坐标点(左上角坐标和右下角坐标)
        x1 = bboxes[:, 0]
        y1 = bboxes[:, 1]
        x2 = bboxes[:, 2]
        y2 = bboxes[:, 3]
        # 1.3 计算每个框的面积
        areas = (x2 - x1) * (y2 - y1)
    
        # 2.对候选框进行NMS筛选
        # 2.1 返回的框坐标和分数
        picked_boxes = []
        picked_score = []
        # 2.2 对置信度进行排序, 获取排序后的下标序号, argsort默认从小到大排序
        order = np.argsort(score)
        while order.size > 0:
            # 2.3 将当前置信度最大的框加入返回值列表中
            index = order[-1]
            # 2.4 保留该类剩余box中得分最高的一个
            picked_boxes.append(bboxes[index])
            picked_score.append(confidence_score[index])
    
            # 2.5 计算交并比
            # 2.5.1 获取当前置信度最大的候选框与其他任意候选框的相交面积
            x11 = np.maximum(x1[index], x1[order[:-1]])
            y11 = np.maximum(y1[index], y1[order[:-1]])
            x22 = np.minimum(x2[index], x2[order[:-1]])
            y22 = np.minimum(y2[index], y2[order[:-1]])
            # 2.5.2 计算交集的面积,不重叠时面积为0
            w = np.maximum(0.0, x22 - x11)
            h = np.maximum(0.0, y22 - y11)
            intersection = w * h
    
            # 2.5.3 利用相交的面积和两个框自身的面积计算框的交并比
            ratio = intersection / (areas[index] + areas[order[:-1]] - intersection)
            # 2.5.4 保留IoU小于阈值的box
            keep_boxes_indics = np.where(ratio < threshold)
            # 2.5.5 保留剩余的框
            order = order[keep_boxes_indics]
        # 2.6 返回NMS后的框及分类结果
        return picked_boxes, picked_score
    
    
    if __name__ == "__main__":
        boxes = [(187, 82, 337, 317), (150, 67, 305, 282), (246, 121, 368, 304)]
        confidence_score = [0.9, 0.65, 0.8]
        threshold = 0.8
        picked_boxes, picked_score = nms(boxes, confidence_score, threshold)
        print(picked_boxes)
        print(picked_score)
    
  • 📄 03.fasterRCNN.ipynb 1.5MB
    # pytorh相关的工具包
    import torch
    import torchvision
    import torchvision.transforms as transforms
    # from torchvision.io import read_image
    # 读取图像的工具包
    from PIL import Image
    # 绘图
    import matplotlib.pyplot as plt
    

    数据加载

    img = Image.open('img2.jpg')
    
    img.size
    
    (994, 734)
    plt.imshow(img)
    
    <matplotlib.image.AxesImage at 0x180689d6688><Figure size 432x288 with 1 Axes>
    transform = transforms.Compose([transforms.ToTensor(),transforms.Resize((800,800))])
    
    img = transform(img)
    
    img.shape
    
    torch.Size([3, 800, 800])
    plt.imshow(img.permute(1,2,0))
    
    <matplotlib.image.AxesImage at 0x7f97b5ce4278><Figure size 432x288 with 1 Axes>

    模型加载

    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
    
    Downloading: "https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth" to C:\Users\14767/.cache\torch\hub\checkpoints\fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
    HBox(children=(FloatProgress(value=0.0, max=167502836.0), HTML(value='')))
    
    model.eval()
    
    FasterRCNN(
      (transform): GeneralizedRCNNTransform(
          Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
          Resize(min_size=(800,), max_size=1333, mode='bilinear')
      )
      (backbone): BackboneWithFPN(
        (body): IntermediateLayerGetter(
          (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
          (bn1): FrozenBatchNorm2d(64, eps=0.0)
          (relu): ReLU(inplace=True)
          (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
          (layer1): Sequential(
            (0): Bottleneck(
              (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(64, eps=0.0)
              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(64, eps=0.0)
              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(256, eps=0.0)
              (relu): ReLU(inplace=True)
              (downsample): Sequential(
                (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
                (1): FrozenBatchNorm2d(256, eps=0.0)
              )
            )
            (1): Bottleneck(
              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(64, eps=0.0)
              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(64, eps=0.0)
              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(256, eps=0.0)
              (relu): ReLU(inplace=True)
            )
            (2): Bottleneck(
              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(64, eps=0.0)
              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(64, eps=0.0)
              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(256, eps=0.0)
              (relu): ReLU(inplace=True)
            )
          )
          (layer2): Sequential(
            (0): Bottleneck(
              (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(128, eps=0.0)
              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(128, eps=0.0)
              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(512, eps=0.0)
              (relu): ReLU(inplace=True)
              (downsample): Sequential(
                (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
                (1): FrozenBatchNorm2d(512, eps=0.0)
              )
            )
            (1): Bottleneck(
              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(128, eps=0.0)
              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(128, eps=0.0)
              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(512, eps=0.0)
              (relu): ReLU(inplace=True)
            )
            (2): Bottleneck(
              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(128, eps=0.0)
              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(128, eps=0.0)
              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(512, eps=0.0)
              (relu): ReLU(inplace=True)
            )
            (3): Bottleneck(
              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(128, eps=0.0)
              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(128, eps=0.0)
              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(512, eps=0.0)
              (relu): ReLU(inplace=True)
            )
          )
          (layer3): Sequential(
            (0): Bottleneck(
              (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=0.0)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=0.0)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=0.0)
              (relu): ReLU(inplace=True)
              (downsample): Sequential(
                (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)
                (1): FrozenBatchNorm2d(1024, eps=0.0)
              )
            )
            (1): Bottleneck(
              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=0.0)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=0.0)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=0.0)
              (relu): ReLU(inplace=True)
            )
            (2): Bottleneck(
              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=0.0)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=0.0)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=0.0)
              (relu): ReLU(inplace=True)
            )
            (3): Bottleneck(
              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=0.0)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=0.0)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=0.0)
              (relu): ReLU(inplace=True)
            )
            (4): Bottleneck(
              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=0.0)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=0.0)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=0.0)
              (relu): ReLU(inplace=True)
            )
            (5): Bottleneck(
              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=0.0)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=0.0)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=0.0)
              (relu): ReLU(inplace=True)
            )
          )
          (layer4): Sequential(
            (0): Bottleneck(
              (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(512, eps=0.0)
              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(512, eps=0.0)
              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(2048, eps=0.0)
              (relu): ReLU(inplace=True)
              (downsample): Sequential(
                (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)
                (1): FrozenBatchNorm2d(2048, eps=0.0)
              )
            )
            (1): Bottleneck(
              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(512, eps=0.0)
              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(512, eps=0.0)
              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(2048, eps=0.0)
              (relu): ReLU(inplace=True)
            )
            (2): Bottleneck(
              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(512, eps=0.0)
              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(512, eps=0.0)
              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(2048, eps=0.0)
              (relu): ReLU(inplace=True)
            )
          )
        )
        (fpn): FeaturePyramidNetwork(
          (inner_blocks): ModuleList(
            (0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))
            (1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))
            (2): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))
            (3): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))
          )
          (layer_blocks): ModuleList(
            (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          )
          (extra_blocks): LastLevelMaxPool()
        )
      )
      (rpn): RegionProposalNetwork(
        (anchor_generator): AnchorGenerator()
        (head): RPNHead(
          (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (cls_logits): Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))
          (bbox_pred): Conv2d(256, 12, kernel_size=(1, 1), stride=(1, 1))
        )
      )
      (roi_heads): RoIHeads(
        (box_roi_pool): MultiScaleRoIAlign(featmap_names=['0', '1', '2', '3'], output_size=(7, 7), sampling_ratio=2)
        (box_head): TwoMLPHead(
          (fc6): Linear(in_features=12544, out_features=1024, bias=True)
          (fc7): Linear(in_features=1024, out_features=1024, bias=True)
        )
        (box_predictor): FastRCNNPredictor(
          (cls_score): Linear(in_features=1024, out_features=91, bias=True)
          (bbox_pred): Linear(in_features=1024, out_features=364, bias=True)
        )
      )
    )

    模型预测

    # coco数据集的class,共90个类别:人,自行车,火车,。。。
    coco_names = [
        '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
        'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
        'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
        'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
        'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
        'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
        'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
        'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
        'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
        'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
    ]
    
    pred = model([img])
    
    /Users/mac/opt/anaconda3/envs/dltorch/lib/python3.7/site-packages/torch/functional.py:445: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2157.)
      return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
    
    pred
    
    [{'boxes': tensor([[ 31.7399, 328.4528, 460.8738, 793.4400],
              [  5.6514,  88.6435, 216.9703, 635.1252],
              [367.4497, 225.4498, 762.0249, 789.4284],
              [607.8991, 152.3864, 795.5750, 372.0394],
              [183.5409,  34.5758, 396.0948, 661.7242],
              [364.9885,  43.9209, 525.7601, 586.4974],
              [ 13.4759, 211.1637, 290.4755, 748.9048],
              [307.3405,   4.9721, 439.2374, 111.8531],
              [383.3587, 180.4190, 578.4739, 772.2529],
              [407.8741, 139.6278, 681.9370, 533.0029],
              [711.5732,   5.5169, 758.7169, 120.2260],
              [479.1790, 144.8082, 689.6121, 400.8835],
              [244.1306,  36.7394, 487.5280, 744.1151],
              [546.8978, 474.4509, 568.8634, 552.9698],
              [742.9993, 314.1062, 798.9455, 607.5062],
              [742.6909, 317.4082, 796.8760, 619.5782],
              [524.8398, 475.1719, 576.5832, 769.7741],
              [509.4732, 142.5110, 623.7419, 301.2690],
              [469.8165, 142.4818, 800.0000, 396.9042],
              [178.5640, 253.6984, 659.5800, 785.6712],
              [365.3784, 474.1786, 470.2338, 769.1154],
              [160.4500, 265.7687, 387.0724, 610.2078],
              [481.0141, 153.0264, 720.3257, 699.9622],
              [724.8939, 343.9777, 798.2111, 726.9193],
              [439.7558, 568.3492, 517.4477, 787.6546],
              [730.4789, 277.5797, 798.6329, 788.0072],
              [409.8385, 474.7333, 540.6846, 793.6932],
              [720.9894, 114.6369, 753.7390, 180.1143],
              [690.4725, 317.8017, 798.0597, 600.9917],
              [548.9393,  76.3451, 613.9292, 172.6772],
              [720.4436, 118.2152, 754.2703, 179.3526],
              [556.6137, 479.1154, 576.5483, 542.2100],
              [514.3457, 148.3877, 619.4744, 501.8499],
              [197.1474, 368.7211, 431.1885, 764.1635],
              [333.9739, 126.6342, 411.8326, 199.1039],
              [659.6864, 131.1151, 716.2268, 168.4116],
              [318.3166, 531.7626, 407.8087, 717.9441],
              [348.7296,   1.2086, 435.9537,  75.7988]], grad_fn=<StackBackward0>),
      'labels': tensor([ 1,  1,  1, 72,  1,  1,  1, 72,  1,  1, 86,  1,  1, 32, 63, 62, 32,  1,
              72,  1, 32,  1,  1, 63, 32, 62, 32, 44, 63, 72, 86, 32,  1,  1, 79, 85,
              32, 72]),
      'scores': tensor([0.9966, 0.9899, 0.9894, 0.9877, 0.9701, 0.9598, 0.8676, 0.7946, 0.6888,
              0.6239, 0.5967, 0.5822, 0.5759, 0.5339, 0.5014, 0.4489, 0.3447, 0.3405,
              0.2513, 0.2067, 0.1989, 0.1979, 0.1844, 0.1780, 0.1268, 0.0884, 0.0864,
              0.0813, 0.0766, 0.0748, 0.0731, 0.0704, 0.0678, 0.0637, 0.0592, 0.0554,
              0.0535, 0.0532], grad_fn=<IndexBackward0>)}]
    list(pred[0]['labels'].numpy())
    
    [1,
     1,
     1,
     72,
     1,
     1,
     1,
     72,
     1,
     1,
     86,
     1,
     1,
     32,
     63,
     62,
     32,
     1,
     72,
     1,
     32,
     1,
     1,
     63,
     32,
     62,
     32,
     44,
     63,
     72,
     86,
     32,
     1,
     1,
     79,
     85,
     32,
     72]
    pred_class = [coco_names[i] for i in list(pred[0]['labels'].numpy())]
    
    pred_class
    
    ['person',
     'person',
     'person',
     'tv',
     'person',
     'person',
     'person',
     'tv',
     'person',
     'person',
     'vase',
     'person',
     'person',
     'tie',
     'couch',
     'chair',
     'tie',
     'person',
     'tv',
     'person',
     'tie',
     'person',
     'person',
     'couch',
     'tie',
     'chair',
     'tie',
     'bottle',
     'couch',
     'tv',
     'vase',
     'tie',
     'person',
     'person',
     'oven',
     'clock',
     'tie',
     'tv']
    pred[0]['boxes'].detach().numpy()
    
    array([[ 31.739914 , 328.45282  , 460.87378  , 793.44     ],
           [  5.65139  ,  88.64346  , 216.97028  , 635.12524  ],
           [367.4497   , 225.44983  , 762.0249   , 789.4284   ],
           [607.8991   , 152.38641  , 795.575    , 372.03937  ],
           [183.54091  ,  34.575836 , 396.09485  , 661.72424  ],
           [364.9885   ,  43.92093  , 525.76013  , 586.49744  ],
           [ 13.475891 , 211.1637   , 290.47552  , 748.90485  ],
           [307.3405   ,   4.9721336, 439.23743  , 111.85315  ],
           [383.3587   , 180.41904  , 578.47394  , 772.2529   ],
           [407.8741   , 139.62776  , 681.937    , 533.0029   ],
           [711.57324  ,   5.516918 , 758.7169   , 120.22604  ],
           [479.17905  , 144.80818  , 689.61206  , 400.88348  ],
           [244.13058  ,  36.73944  , 487.52798  , 744.1151   ],
           [546.8978   , 474.45093  , 568.8634   , 552.96985  ],
           [742.9993   , 314.10623  , 798.9455   , 607.5062   ],
           [742.6909   , 317.40823  , 796.876    , 619.57825  ],
           [524.8398   , 475.17194  , 576.5832   , 769.7741   ],
           [509.47318  , 142.51099  , 623.74194  , 301.26895  ],
           [469.81653  , 142.48177  , 800.       , 396.90424  ],
           [178.56403  , 253.69843  , 659.58     , 785.6712   ],
           [365.37836  , 474.17865  , 470.23376  , 769.1154   ],
           [160.45001  , 265.76874  , 387.07245  , 610.2078   ],
           [481.0141   , 153.0264   , 720.32574  , 699.96216  ],
           [724.8939   , 343.97772  , 798.21106  , 726.91925  ],
           [439.75577  , 568.3492   , 517.4477   , 787.6546   ],
           [730.47894  , 277.57974  , 798.6329   , 788.0072   ],
           [409.8385   , 474.73334  , 540.68463  , 793.6932   ],
           [720.98944  , 114.63695  , 753.73895  , 180.11427  ],
           [690.47253  , 317.80173  , 798.0597   , 600.9917   ],
           [548.93933  ,  76.3451   , 613.9292   , 172.67722  ],
           [720.4436   , 118.215164 , 754.27026  , 179.35258  ],
           [556.6137   , 479.11536  , 576.5483   , 542.21     ],
           [514.3457   , 148.38771  , 619.47437  , 501.8499   ],
           [197.14737  , 368.72107  , 431.1885   , 764.16345  ],
           [333.97394  , 126.63425  , 411.83258  , 199.10388  ],
           [659.6864   , 131.11514  , 716.2268   , 168.41156  ],
           [318.3166   , 531.7626   , 407.80872  , 717.9441   ],
           [348.72955  ,   1.2085876, 435.95374  ,  75.7988   ]],
          dtype=float32)
    pred_box = [[(box[0],box[1]),(box[2],box[3])] for box in list(pred[0]['boxes'].detach().numpy())]
    
    pred_box
    
    [[(31.739914, 328.45282), (460.87378, 793.44)],
     [(5.65139, 88.64346), (216.97028, 635.12524)],
     [(367.4497, 225.44983), (762.0249, 789.4284)],
     [(607.8991, 152.38641), (795.575, 372.03937)],
     [(183.54091, 34.575836), (396.09485, 661.72424)],
     [(364.9885, 43.92093), (525.76013, 586.49744)],
     [(13.475891, 211.1637), (290.47552, 748.90485)],
     [(307.3405, 4.9721336), (439.23743, 111.85315)],
     [(383.3587, 180.41904), (578.47394, 772.2529)],
     [(407.8741, 139.62776), (681.937, 533.0029)],
     [(711.57324, 5.516918), (758.7169, 120.22604)],
     [(479.17905, 144.80818), (689.61206, 400.88348)],
     [(244.13058, 36.73944), (487.52798, 744.1151)],
     [(546.8978, 474.45093), (568.8634, 552.96985)],
     [(742.9993, 314.10623), (798.9455, 607.5062)],
     [(742.6909, 317.40823), (796.876, 619.57825)],
     [(524.8398, 475.17194), (576.5832, 769.7741)],
     [(509.47318, 142.51099), (623.74194, 301.26895)],
     [(469.81653, 142.48177), (800.0, 396.90424)],
     [(178.56403, 253.69843), (659.58, 785.6712)],
     [(365.37836, 474.17865), (470.23376, 769.1154)],
     [(160.45001, 265.76874), (387.07245, 610.2078)],
     [(481.0141, 153.0264), (720.32574, 699.96216)],
     [(724.8939, 343.97772), (798.21106, 726.91925)],
     [(439.75577, 568.3492), (517.4477, 787.6546)],
     [(730.47894, 277.57974), (798.6329, 788.0072)],
     [(409.8385, 474.73334), (540.68463, 793.6932)],
     [(720.98944, 114.63695), (753.73895, 180.11427)],
     [(690.47253, 317.80173), (798.0597, 600.9917)],
     [(548.93933, 76.3451), (613.9292, 172.67722)],
     [(720.4436, 118.215164), (754.27026, 179.35258)],
     [(556.6137, 479.11536), (576.5483, 542.21)],
     [(514.3457, 148.38771), (619.47437, 501.8499)],
     [(197.14737, 368.72107), (431.1885, 764.16345)],
     [(333.97394, 126.63425), (411.83258, 199.10388)],
     [(659.6864, 131.11514), (716.2268, 168.41156)],
     [(318.3166, 531.7626), (407.80872, 717.9441)],
     [(348.72955, 1.2085876), (435.95374, 75.7988)]]
    # 图像读取
    img = Image.open('img2.jpg').resize((800,800))
    # 图像展示
    fig = plt.imshow(img)
    # 将框绘制在图像上
    for index, boxes in enumerate(pred_box):
        fig.axes.add_patch(plt.Rectangle(
        xy=(boxes[0]), width=boxes[1][0]-boxes[0][0], height=boxes[1][1]-boxes[0][1],
        fill=False, edgecolor="blue", linewidth=1))
    plt.show()
    
    <Figure size 432x288 with 1 Axes>
    the = 0.7
    
    pred_score = list(pred[0]['scores'].detach().numpy())
    
    pred_score
    
    [0.99662447,
     0.98991066,
     0.9894002,
     0.9877265,
     0.9700843,
     0.95981795,
     0.86762965,
     0.79457456,
     0.6888022,
     0.6238965,
     0.5967382,
     0.5821962,
     0.5758556,
     0.53390706,
     0.50142986,
     0.4489189,
     0.3446656,
     0.34053934,
     0.25132728,
     0.20674044,
     0.19888856,
     0.19792202,
     0.18443534,
     0.17802832,
     0.12675783,
     0.08840103,
     0.08642756,
     0.08125935,
     0.07662633,
     0.0747837,
     0.07308833,
     0.07040265,
     0.06783704,
     0.063683495,
     0.0592302,
     0.055396155,
     0.053475436,
     0.053160064]
    pred_t = [pred_score.index(score) for score in pred_score if score > the]
    
    index = pred_t[-1]
    
    pred_box = pred_box[:index+1]
    
    pred_box
    
    [[(31.739914, 328.45282), (460.87378, 793.44)],
     [(5.65139, 88.64346), (216.97028, 635.12524)],
     [(367.4497, 225.44983), (762.0249, 789.4284)],
     [(607.8991, 152.38641), (795.575, 372.03937)],
     [(183.54091, 34.575836), (396.09485, 661.72424)],
     [(364.9885, 43.92093), (525.76013, 586.49744)],
     [(13.475891, 211.1637), (290.47552, 748.90485)],
     [(307.3405, 4.9721336), (439.23743, 111.85315)]]
    # 图像读取
    img = Image.open('img2.jpg').resize((800,800))
    # 图像展示
    fig = plt.imshow(img)
    # 将框绘制在图像上
    for index, boxes in enumerate(pred_box):
        fig.axes.add_patch(plt.Rectangle(
        xy=(boxes[0]), width=boxes[1][0]-boxes[0][0], height=boxes[1][1]-boxes[0][1],
        fill=False, edgecolor="blue", linewidth=1))
    plt.show()
    
    <Figure size 432x288 with 1 Axes>
    pred_class[:index+1]
    
    ['person', 'person', 'person', 'tv', 'person', 'person', 'person', 'tv']

    模型结构详解

    backbone

    transform = transforms.Compose([transforms.ToTensor(),transforms.Resize((800,800))])
    img = transform(img)
    img.shape
    
    torch.Size([3, 800, 800])
    imgs,targets= model.transform([img],targets=None)
    
    imgs
    
    <torchvision.models.detection.image_list.ImageList at 0x7f97769e3908>
    feature_body= model.backbone.body(imgs.tensors)
    
    C2,C3,C4,C5 = feature_body.values()
    
    C2.size()
    
    torch.Size([1, 256, 200, 200])
    C3.size()
    
    torch.Size([1, 512, 100, 100])
    C4.size()
    
    torch.Size([1, 1024, 50, 50])
    C5.size()
    
    torch.Size([1, 2048, 25, 25])
    feature_fpn=model.backbone.fpn(feature_body)
    
    P2,P3,P4,P5,P6=feature_fpn.values()
    
    P2.size()
    
    torch.Size([1, 256, 200, 200])
    P3.size()
    
    torch.Size([1, 256, 100, 100])
    P4.size()
    
    torch.Size([1, 256, 50, 50])
    P5.size()
    
    torch.Size([1, 256, 25, 25])
    P6.size()
    
    torch.Size([1, 256, 13, 13])

    RPN

    anchors = model.rpn.anchor_generator(imgs,[P2,P3,P4,P5,P6])
    
    anchors
    
    [tensor([[ -23.,  -11.,   23.,   11.],
             [ -16.,  -16.,   16.,   16.],
             [ -11.,  -23.,   11.,   23.],
             ...,
             [ 370.,  551., 1094.,  913.],
             [ 476.,  476.,  988.,  988.],
             [ 551.,  370.,  913., 1094.]])]
    anchors[0].size()
    
    torch.Size([159882, 4])
    # 将生成的anchor绘制在图像上
    img_anchor = Image.open('img2.jpg').resize((800,800))
    # 绘制图像
    fig = plt.imshow(img_anchor)
    # 将框绘制在图像上
    for index, boxes in enumerate(anchors[0]):
        if index <300:
            fig.axes.add_patch(plt.Rectangle(
                xy=(boxes[0],boxes[1]), width=boxes[2]-boxes[0], height=boxes[3]-boxes[1],
                fill=False, edgecolor="red", linewidth=1))
    plt.show()
    
    <Figure size 432x288 with 1 Axes>
    cls_logits = model.rpn.head.cls_logits(P3)
    
    cls_logits.size()
    
    torch.Size([1, 3, 100, 100])
    bbox_pred = model.rpn.head.bbox_pred(P4)
    
    bbox_pred.size()
    
    torch.Size([1, 12, 50, 50])
    proposal = model.rpn(imgs,feature_fpn)
    
    proposal
    
    ([tensor([[162.3166,  82.6454, 390.0235, 638.2616],
              [  0.0000,  45.9522, 237.6339, 632.3745],
              [170.6081,  36.5363, 371.3253, 548.1898],
              ...,
              [373.4832, 426.8497, 396.0713, 497.3481],
              [317.5831,   7.1501, 417.4209,  70.2656],
              [395.1561, 506.0393, 427.8982, 546.8923]])],
     {})
    proposal[0][0].size()
    
    torch.Size([1000, 4])
    # 候选区域的绘制
    img = Image.open('img2.jpg').resize((800,800))
    # 显示图像
    fig = plt.imshow(img)
    # 将候选区域显示在图像上
    for index, boxes in enumerate(proposal[0][0]):
        if index <3:
            fig.axes.add_patch(plt.Rectangle(
            xy=(boxes[0],boxes[1]), width=boxes[2]-boxes[0], height=boxes[3]-boxes[1],
            fill=False, edgecolor="blue", linewidth=0.5))
    plt.show()
    
    <Figure size 432x288 with 1 Axes>
    pool_region_list= model.roi_heads.box_roi_pool(feature_fpn,proposal[0],imgs.image_sizes)
    
    pool_region_list.size()
    
    torch.Size([1000, 256, 7, 7])
    detection = model.roi_heads(feature_fpn,proposal[0],imgs.image_sizes)
    
    detection_boxes = detection[0][0]['boxes'].detach().numpy()
    img = Image.open('img2.jpg').resize((800,800))
    # 显示图像
    fig = plt.imshow(img)
    # 将检测结果绘制在图像上
    for index, boxes in enumerate(detection_boxes):
        fig.axes.add_patch(plt.Rectangle(
        xy=(boxes[0],boxes[1]), width=boxes[2]-boxes[0], height=boxes[3]-boxes[1],
        fill=False, edgecolor="purple", linewidth=1))
    plt.show()
    
    <Figure size 432x288 with 1 Axes>

    模型训练

    数据加载

    from voc_data_util import VocDataset,VOC_BBOX_LABEL_NAMES
    
    dataset = VocDataset(
        root='datasets/VOCdevkit/VOC2007/',
        transforms=VocDataset.get_transform())
    
    img,targets = dataset[3]
    
    img.size()
    
    torch.Size([3, 800, 800])
    import numpy as np
    img = img.permute(1,2,0).numpy().astype(np.uint8)
    
    plt.imshow(img)
    
    <matplotlib.image.AxesImage at 0x7f9753d7b518><Figure size 432x288 with 1 Axes>
    targets
    
    {'boxes': tensor([[238.4000, 233.0330, 550.4000, 648.6486]], dtype=torch.float64),
     'labels': tensor([7]),
     'image_id': tensor([3]),
     'area': tensor([33735.]),
     'iscrowd': tensor([0])}
    fig = plt.imshow(img)
    # 将标注信息绘制在图像上
    for index, boxes in enumerate(targets['boxes'].numpy()):
        fig.axes.add_patch(
            plt.Rectangle(xy=(boxes[0], boxes[1]),
                          width=boxes[2] - boxes[0],
                          height=boxes[3] - boxes[1],
                          fill=False,
                          edgecolor="green",
                          linewidth=1))
    plt.show()
    
    <Figure size 432x288 with 1 Axes>

    模型加载

    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
        pretrained=False, num_classes=21, pretrained_backbone=True)
    
    model.train()
    
    FasterRCNN(
      (transform): GeneralizedRCNNTransform(
          Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
          Resize(min_size=(800,), max_size=1333, mode='bilinear')
      )
      (backbone): BackboneWithFPN(
        (body): IntermediateLayerGetter(
          (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
          (bn1): FrozenBatchNorm2d(64, eps=1e-05)
          (relu): ReLU(inplace=True)
          (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
          (layer1): Sequential(
            (0): Bottleneck(
              (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(64, eps=1e-05)
              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(64, eps=1e-05)
              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(256, eps=1e-05)
              (relu): ReLU(inplace=True)
              (downsample): Sequential(
                (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
                (1): FrozenBatchNorm2d(256, eps=1e-05)
              )
            )
            (1): Bottleneck(
              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(64, eps=1e-05)
              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(64, eps=1e-05)
              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(256, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
            (2): Bottleneck(
              (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(64, eps=1e-05)
              (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(64, eps=1e-05)
              (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(256, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
          )
          (layer2): Sequential(
            (0): Bottleneck(
              (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(128, eps=1e-05)
              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(128, eps=1e-05)
              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(512, eps=1e-05)
              (relu): ReLU(inplace=True)
              (downsample): Sequential(
                (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
                (1): FrozenBatchNorm2d(512, eps=1e-05)
              )
            )
            (1): Bottleneck(
              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(128, eps=1e-05)
              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(128, eps=1e-05)
              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(512, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
            (2): Bottleneck(
              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(128, eps=1e-05)
              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(128, eps=1e-05)
              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(512, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
            (3): Bottleneck(
              (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(128, eps=1e-05)
              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(128, eps=1e-05)
              (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(512, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
          )
          (layer3): Sequential(
            (0): Bottleneck(
              (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=1e-05)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=1e-05)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=1e-05)
              (relu): ReLU(inplace=True)
              (downsample): Sequential(
                (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)
                (1): FrozenBatchNorm2d(1024, eps=1e-05)
              )
            )
            (1): Bottleneck(
              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=1e-05)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=1e-05)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
            (2): Bottleneck(
              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=1e-05)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=1e-05)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
            (3): Bottleneck(
              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=1e-05)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=1e-05)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
            (4): Bottleneck(
              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=1e-05)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=1e-05)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
            (5): Bottleneck(
              (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(256, eps=1e-05)
              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(256, eps=1e-05)
              (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(1024, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
          )
          (layer4): Sequential(
            (0): Bottleneck(
              (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(512, eps=1e-05)
              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(512, eps=1e-05)
              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(2048, eps=1e-05)
              (relu): ReLU(inplace=True)
              (downsample): Sequential(
                (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)
                (1): FrozenBatchNorm2d(2048, eps=1e-05)
              )
            )
            (1): Bottleneck(
              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(512, eps=1e-05)
              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(512, eps=1e-05)
              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(2048, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
            (2): Bottleneck(
              (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn1): FrozenBatchNorm2d(512, eps=1e-05)
              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
              (bn2): FrozenBatchNorm2d(512, eps=1e-05)
              (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
              (bn3): FrozenBatchNorm2d(2048, eps=1e-05)
              (relu): ReLU(inplace=True)
            )
          )
        )
        (fpn): FeaturePyramidNetwork(
          (inner_blocks): ModuleList(
            (0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))
            (1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))
            (2): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))
            (3): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))
          )
          (layer_blocks): ModuleList(
            (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          )
          (extra_blocks): LastLevelMaxPool()
        )
      )
      (rpn): RegionProposalNetwork(
        (anchor_generator): AnchorGenerator()
        (head): RPNHead(
          (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (cls_logits): Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))
          (bbox_pred): Conv2d(256, 12, kernel_size=(1, 1), stride=(1, 1))
        )
      )
      (roi_heads): RoIHeads(
        (box_roi_pool): MultiScaleRoIAlign(featmap_names=['0', '1', '2', '3'], output_size=(7, 7), sampling_ratio=2)
        (box_head): TwoMLPHead(
          (fc6): Linear(in_features=12544, out_features=1024, bias=True)
          (fc7): Linear(in_features=1024, out_features=1024, bias=True)
        )
        (box_predictor): FastRCNNPredictor(
          (cls_score): Linear(in_features=1024, out_features=21, bias=True)
          (bbox_pred): Linear(in_features=1024, out_features=84, bias=True)
        )
      )
    )

    RPN正负样本

    rpn_label_matchs,rpn_target_deltas= model.rpn.assign_targets_to_anchors(anchors,[targets])
    
    rpn_label_matchs
    
    [tensor([0., 0., 0.,  ..., 0., 0., 0.])]
    rpn_target_deltas
    
    [tensor([[238.4000, 233.0330, 550.4000, 648.6486],
             [238.4000, 233.0330, 550.4000, 648.6486],
             [238.4000, 233.0330, 550.4000, 648.6486],
             ...,
             [238.4000, 233.0330, 550.4000, 648.6486],
             [238.4000, 233.0330, 550.4000, 648.6486],
             [238.4000, 233.0330, 550.4000, 648.6486]], dtype=torch.float64)]
    positive_anchor_id= torch.where(torch.eq(rpn_label_matchs[0],1))[0].numpy()
    
    positive_anchor_id.shape
    
    (8,)
    anchors[0].size()
    
    torch.Size([159882, 4])
    fig = plt.imshow(img)
    # 遍历所有的正样本的anchor,绘制在图像上
    for idx in positive_anchor_id:
        print(idx)
        if idx == 158513:
            boxes = anchors[0][idx]
            fig.axes.add_patch(plt.Rectangle(
            xy=(boxes[0],boxes[1]), width=boxes[2]-boxes[0], height=boxes[3]-boxes[1],
            fill=False, edgecolor="red", linewidth=1))
    plt.show()
    
    158510
    158513
    158516
    158519
    158585
    158588
    158591
    158594
    <Figure size 432x288 with 1 Axes>
    negative_anchor_id= torch.where(torch.eq(rpn_label_matchs[0],-1))[0].numpy()
    
    negative_anchor_id.shape
    
    (249,)
    fig = plt.imshow(img)
    # 遍历所有的正样本的anchor,绘制在图像上
    for idx in negative_anchor_id:
       # print(idx)
        # 负样本的anchor
        boxes = anchors[0][idx]
        # 取余20为0的负样本进行展示
        fig.axes.add_patch(
            plt.Rectangle(xy=(boxes[0], boxes[1]),
                          width=boxes[2] - boxes[0],
                          height=boxes[3] - boxes[1],
                          fill=False,
                          edgecolor="white",
                          linewidth=1))
    plt.show()
    
    <Figure size 432x288 with 1 Axes>
    regression_targets = model.rpn.box_coder.encode(rpn_target_deltas,anchors)
    
    regression_targets[0].size()
    
    torch.Size([159882, 4])
    proposals,matched_idx,labels,regression_targets= model.roi_heads.select_training_samples(proposal[0],[targets])
    
    proposals[0].size()
    
    torch.Size([512, 4])
    regression_targets[0].size()
    
    torch.Size([512, 4])
    labels
    
    [tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 7, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 7])]
    positive_proposal_id=torch.where(torch.gt(labels[0],0))[0].numpy()
    
    # 绘制部分候选区域的负样本
    fig = plt.imshow(img)
    # 遍历所有的正样本的候选区域,绘制在图像上
    for idx in positive_proposal_id:
        # 正样本的proposals
        boxes = proposals[0][idx]
        # 取余10为0的负样本进行展示
        fig.axes.add_patch(
            plt.Rectangle(xy=(boxes[0], boxes[1]),
                          width=boxes[2] - boxes[0],
                          height=boxes[3] - boxes[1],
                          fill=False,
                          edgecolor="white",
                          linewidth=1))
    plt.show()
    
    <Figure size 432x288 with 1 Axes>
    negitivate_proposal_id = torch.where(torch.eq(labels[0],0))[0].numpy()
    
    # 绘制部分候选区域的负样本
    fig = plt.imshow(img)
    # 遍历所有的负样本的候选区域,绘制在图像上
    for idx in negitivate_proposal_id:
        # 负样本的proposals
        boxes = proposals[0][idx]
        # 取余10为0的负样本进行展示
        if idx % 10 == 0:
            fig.axes.add_patch(
                plt.Rectangle(xy=(boxes[0], boxes[1]),
                              width=boxes[2] - boxes[0],
                              height=boxes[3] - boxes[1],
                              fill=False,
                              edgecolor="white",
                              linewidth=1))
    plt.show()
    
    <Figure size 432x288 with 1 Axes>

    模型训练

    from torch.utils.data import DataLoader
    
    def collate_fn(batch):
        return tuple(zip(*batch))
    
    train_data_loader = DataLoader(dataset,batch_size=2,shuffle = True,collate_fn = collate_fn)
    
    for i,data in enumerate(train_data_loader):
        image,targets = data
        print(len(image))
        print(targets)
        break
    
    2
    ({'boxes': tensor([[224.0000, 140.8000, 425.6000, 522.6666]], dtype=torch.float64), 'labels': tensor([8]), 'image_id': tensor([3183]), 'area': tensor([22554.]), 'iscrowd': tensor([0])}, {'boxes': tensor([[230.4000, 256.0000, 696.0000, 605.8667]], dtype=torch.float64), 'labels': tensor([3]), 'image_id': tensor([2603]), 'area': tensor([47724.]), 'iscrowd': tensor([0])})
    
    # 优化器
    params = [p for p in model.parameters() if p.requires_grad]
    
    optimizer = torch.optim.SGD(params,lr= 0.0001,momentum=0.9,weight_decay=0.0001)
    
    # 遍历每个轮次进行训练
    
    # 1.参数的初始化
    iters = 0
    loss_value = 0
    total_train_loss = []
    # 2.遍历每个轮次进行训练
    for epoch in range(2):
        train_loss = []
        # 3.每个轮次中遍历每个batch
        for images, targets in train_data_loader:
            # 4.把batch数据送入网络中进行预测:损失
            loss_dict = model(images, targets)
            print(loss_dict)
            loss = sum(loss for loss in loss_dict.values())
            loss_value = loss.item()
            train_loss.append(loss_value)
            # 5.进行反向传播:
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print(f"Epoch:{epoch},Batch:{iters},loss:{loss_value}")
            iters+=1
            break
        # 6.保存一此列的参数
        epoch_loss = np.mean(train_loss)
        total_train_loss.append(epoch_loss)
        break
    
    {'loss_classifier': tensor(110.5280, grad_fn=<NllLossBackward0>), 'loss_box_reg': tensor(22.8108, grad_fn=<DivBackward0>), 'loss_objectness': tensor(7.9993, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'loss_rpn_box_reg': tensor(7.4245, dtype=torch.float64, grad_fn=<DivBackward0>)}
    Epoch:0,Batch:0,loss:148.76265928275987
    
    
    
  • 📄 03.fasterRCNN.py 3.4KB
    # 1.pytorh相关的工具包
    import torch
    import torchvision
    import torchvision.transforms as transforms
    from torchvision.io import read_image
    from PIL import Image
    import matplotlib.pyplot as plt
    import warnings
    
    warnings.filterwarnings("ignore")
    # 2.获取图片
    img = Image.open('img3.jpg')
    # plt.imshow(img)
    # plt.show()
    # print(img.size)
    
    # 3.将图片格式转换为tensor形式,大小转换为800x800的大小
    transform = transforms.Compose([transforms.ToTensor(), transforms.Resize((800, 800))])
    img = transform(img)
    
    # 3.1 进行通道的调整并展示图片
    plt.imshow(img.permute(1, 2, 0))
    # print(img.shape)
    plt.show()
    
    
    # 4.设置coco数据集的class,共90个类别:人,自行车,火车,。。。
    coco_names = [
        '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
        'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
        'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
        'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
        'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
        'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
        'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
        'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
        'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
        'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
    ]
    
    # 5.实例化模型
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
    
    # 6.将模型设置为eval模式
    model.eval()
    # print(model)
    
    # 7.直接进行预测
    pred = model([img])
    print(pred)
    
    # 8.获取类别名称,框及对应的置信度
    pred_class = [coco_names[i] for i in list(pred[0]['labels'].numpy())]
    pred_boxes = [[(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]['boxes'].detach().numpy())]
    pred_score = list(pred[0]['scores'].detach().numpy())
    
    print(pred_class)
    print(pred_score)
    
    # # 9.展示绘图结果
    # # 9.1 图像读取
    # img = Image.open('img3.jpg').resize((800, 800))
    # # 9.2 图像展示
    # fig = plt.imshow(img)
    # # # 9.3 将框绘制在图像上
    # # for index, boxes in enumerate(pred_boxes):
    # #     fig.axes.add_patch(plt.Rectangle(
    # #     xy=(boxes[0]), width=boxes[1][0]-boxes[0][0], height=boxes[1][1]-boxes[0][1],
    # #     fill=False, edgecolor="blue", linewidth=1))
    # # plt.show()
    #
    #
    # # 10.过滤分数较低的预测
    # threshold = 0.5
    # # 因pred_score是从大大小进行排列的,只要获取最后一个索引即可
    # pred_t = [pred_score.index(x) for x in pred_score if x > threshold][-1]
    # pred_boxes = pred_boxes[:pred_t + 1]
    # pred_class = pred_class[:pred_t + 1]
    #
    # # 11.将过滤后的结构绘制在图像上
    # # 11.1 图像读取
    # img = Image.open('img3.jpg').resize((800, 800))
    # # 11.2 将框绘制在图像上
    # fig = plt.imshow(img)
    # for index, boxes in enumerate(pred_boxes):
    #     fig.axes.add_patch(plt.Rectangle(
    #         xy=(boxes[0]), width=boxes[1][0] - boxes[0][0], height=boxes[1][1] - boxes[0][1],
    #         fill=False, edgecolor="g", linewidth=1))
    # plt.show()
    
  • 📄 04.fasterRCNN结构介绍.py 3.2KB
    # 1.pytorh相关的工具包
    import torch
    import torchvision
    import torchvision.transforms as transforms
    from torchvision.io import read_image
    from PIL import Image
    import matplotlib.pyplot as plt
    import warnings
    
    warnings.filterwarnings("ignore")
    # 2.图像读取并转换
    img = Image.open('img3.jpg')
    
    transform = transforms.Compose([transforms.ToTensor(), transforms.Resize((800, 800))])
    img_transform = transform(img)
    
    # 3.实例化模型
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
    model.eval()
    
    # 4.送入网络中的数据是【N,C,H,W】
    images = [img_transform]
    images, targets = model.transform(images, targets=None)
    
    # 5.使用backbone获取特征图
    features_body = model.backbone.body(images.tensors)
    C2, C3, C4, C5 = features_body.values()
    # print(C2.shape)
    # print(C3.shape)
    # print(C4.shape)
    # print(C5.shape)
    
    
    # 6.FPN网络融合:C2,C3,C4,C5是resnet提取的特征结果
    features_fpn = model.backbone.fpn(features_body)
    P2, P3, P4, P5, P6 = features_fpn.values()
    # print(P2.shape)
    # print(P3.shape)
    # print(P4.shape)
    # print(P5.shape)
    # print(P6.shape)
    
    # 7.RPN网络
    # 7.1 生成anchors,送入图片信息及相应的特征图
    anchors = model.rpn.anchor_generator(images, [P2, P3, P4, P5, P6])
    # print(anchors)
    # print(anchors[0].shape)
    
    # 7.1.2 将生成的anchor绘制在图像上
    # 将框绘制在图像上
    fig = plt.imshow(img)
    # for index, boxes in enumerate(anchors[0]):
    #     if index < 1000:
    #         fig.axes.add_patch(plt.Rectangle(
    #             xy=(boxes[0], boxes[1]), width=boxes[2] - boxes[0], height=boxes[3] - boxes[1],
    #             fill=False, edgecolor="red", linewidth=1))
    # plt.show()
    
    # 7.2 RPN分类和RPN回归
    # 7.2.1 RPN网络预测分类,返回:logits送入softmax之前的分数
    cls_logits = model.rpn.head.cls_logits(P3)
    # 7.2.2 RPN网络预测目标框
    box_pred = model.rpn.head.bbox_pred(P4)
    # print(cls_logits.size())
    # print(box_pred.size())
    # print(cls_logits)
    
    # 7.3 Proposal层(修正目标框)
    # 7.3.1 RPN网络生成proposal
    proposal = model.rpn(images, features_fpn)
    # print(proposal[0][0].size())
    # print(proposal)
    
    
    # 7.3.2 候选区域的绘制
    # 将候选区域显示在图像上
    fig = plt.imshow(img)
    # for index, boxes in enumerate(proposal[0][0]):
    #     fig.axes.add_patch(plt.Rectangle(
    #         xy=(boxes[0], boxes[1]), width=boxes[2] - boxes[0], height=boxes[3] - boxes[1],
    #         fill=False, edgecolor="blue", linewidth=0.5))
    # plt.show()
    
    # 8.ROI Pooling层实现
    pool_region_list = model.roi_heads.box_roi_pool(features_fpn, proposal[0], images.image_sizes)
    print(pool_region_list.size())
    
    # 9.目标分类与回归
    # 9.1 模型最终输出的结果
    detection = model.roi_heads(features_fpn, proposal[0], images.image_sizes)
    print(detection)
    
    # 9.2 检测结果的绘制
    detection_boxes = detection[0][0]['boxes'].detach().numpy()
    # 9.2.1 将检测结果绘制在图像上
    fig = plt.imshow(img)
    for index, boxes in enumerate(detection_boxes):
        fig.axes.add_patch(plt.Rectangle(
            xy=(boxes[0], boxes[1]), width=boxes[2] - boxes[0], height=boxes[3] - boxes[1],
            fill=False, edgecolor="purple", linewidth=1))
    plt.show()
       #  [x1,y1,x2,y2]
    
  • 📄 05.fasterRCNN训练.py 3.7KB
    # 1.相关依赖包导入
    import torch
    import torchvision
    import torchvision.transforms as transforms
    from torchvision.io import read_image
    from PIL import Image
    import matplotlib.pyplot as plt
    import numpy as np
    # 加载数据集
    from voc_data_util import VocDataset, VOC_BBOX_LABEL_NAMES
    from torch.utils.data import DataLoader
    
    import warnings
    warnings.filterwarnings("ignore")
    
    # 2.数据集加载
    dataset = VocDataset(root='datasets/VOCdevkit/VOC2007', transforms=VocDataset.get_transform())
    
    # # 数据集绘制
    # # 获取图像及对应的标注信息
    img, targets = dataset[141]
    img = img.permute(1, 2, 0).numpy().astype(np.uint8)
    fig = plt.imshow(img)
    # 将标注信息绘制在图像上
    for index, boxes in enumerate(targets['boxes'].numpy()):
        fig.axes.add_patch(
            plt.Rectangle(xy=(boxes[0], boxes[1]),
                          width=boxes[2] - boxes[0],
                          height=boxes[3] - boxes[1],
                          fill=False,
                          edgecolor="green",
                          linewidth=1))
    # plt.show()
    
    # 3.模型加载及实例化
    # N+1包含背景
    
    # 3.模型实例化及参数设置
    num_classes = 21
    # 3.1 模型实例化:不使用预训练模型,修正类别个数,backbone使用预训练结果
    # model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
    #     pretrained=False,
    #     num_classes=num_classes,
    #     pretrained_backbone=True)
    
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False,
                                                                 pretrained_backbone=True,
                                                                 num_classes=num_classes)
    
    
    def collate_fn(batch):
        return tuple(zip(*batch))
    
    
    # 3.2 参数设置
    # 3.2.1 获取batch的数据,送入网络中进行训练
    train_data_loader = DataLoader(
        dataset,
        batch_size=2,
        shuffle=True,
        collate_fn=collate_fn
    )
    
    # 遍历数据,获取第一个batch的数据进行展示
    # for i, data in enumerate(train_data_loader):
    #     inputs, labels = data
    #     # batch的大小
    #     print('图片个数:', len(inputs))
    #     print('目标值:', labels)
    #     break
    
    
    # 获取所有要进行训练的参数,设置优化器
    # 3.2.2 参数冻结
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params, lr=0.00005, momentum=0.9, weight_decay=0.0005)
    # 3.2.3 迭代次数
    itr = 1
    # 3.2.4 存放训练损失
    total_train_loss = []
    # 3.2.5 损失值
    losses_value = 0
    
    # 4.模型训练
    # 4.1 遍历每个轮次进行训练
    for epoch in range(2):
        # 4.2 训练模式
        model.train()
        train_loss = []
        # 4.3 遍历数据,获取图像,目标值和图像id
        for images, targets in train_data_loader:
            # 4.4 将数据送入网络中进行训练,获取损失值(RPN的损失+FastRCNN的损失),将信息打印出来分析
            loss_dict = model(images, targets)
            print(loss_dict)
            # 4.5 将损失求和
            losses = sum(loss for loss in loss_dict.values())
            # 4.6 获取loss值
            losses_value = losses.item()
            train_loss.append(losses_value)
            # 4.7 进行反向传播,更新参数
            optimizer.zero_grad()
            losses.backward()
            optimizer.step()
            # 4.8 日志信息
            print(f"Epoch: {epoch + 1}, Batch: {itr}, Loss: {losses_value}")
            # 4.9 迭代次数增1
            itr += 1
        # 4.10 获取当前轮次的损失
        epoch_train_loss = np.mean(train_loss)
        # 4.11 轮次损失写入到列表中
        total_train_loss.append(epoch_train_loss)
    
    # 5.绘制损失函数变化的曲线
    plt.plot(range(len(total_train_loss)), [loss.numpy() for loss in total_train_loss])
    plt.grid()
    
  • ⚙️ img.jpg 245.9KB
  • ⚙️ img1.jpg 161.0KB
  • ⚙️ img2.jpg 182.6KB
  • ⚙️ img3.jpg 154.7KB
  • 📄 voc_data_util.py 3.8KB
    import torch
    import numpy as np
    import cv2
    import matplotlib.pyplot as plt
    from torchvision import datasets, transforms
    from PIL import Image
    from xml.dom.minidom import parse
    import os
    import albumentations as A
    from albumentations.pytorch.transforms import ToTensorV2
    
    VOC_BBOX_LABEL_NAMES = (
        '_backbone_',
        'aeroplane',
        'bicycle',
        'bird',
        'boat',
        'bottle',
        'bus',
        'car',
        'cat',
        'chair',
        'cow',
        'diningtable',
        'dog',
        'horse',
        'motorbike',
        'person',
        'pottedplant',
        'sheep',
        'sofa',
        'train',
        'tvmonitor')
    
    
    class VocDataset(torch.utils.data.Dataset):
        def __init__(self, root, transforms=None):
            self.root = root
            self.transforms = transforms
            # load all image files, sorting them to ensure that they are aligned
            self.imgs = list(sorted(os.listdir(os.path.join(root, "JPEGImages"))))
            self.bbox_xml = list(sorted(os.listdir(os.path.join(root, "Annotations"))))
    
        def __getitem__(self, idx):
            # load images and bbox
            img_path = os.path.join(self.root, "JPEGImages", self.imgs[idx])
            bbox_xml_path = os.path.join(self.root, "Annotations", self.bbox_xml[idx])
            # img = Image.open(img_path).convert("RGB")
            img = plt.imread(img_path).astype(np.float32)
            if img_path.find('000282.jpg') == 1:
                print(img_path)
                print("索引为:{}".format(idx))
    
            # 读取文件,VOC格式的数据集的标注是xml格式的文件
            dom = parse(bbox_xml_path)
            # 获取文档元素对象
            data = dom.documentElement
            # 获取 objects
            objects = data.getElementsByTagName('object')
            # get bounding box coordinates
            boxes = []
            labels = []
            for object_ in objects:
                # 获取标签中内容
                name = object_.getElementsByTagName('name')[0].childNodes[0].nodeValue  # 就是label,mark_type_1或mark_type_2
                labels.append(VOC_BBOX_LABEL_NAMES.index(name))
                bndbox = object_.getElementsByTagName('bndbox')[0]
                xmin = np.float64(bndbox.getElementsByTagName('xmin')[0].childNodes[0].nodeValue)
                ymin = np.float64(bndbox.getElementsByTagName('ymin')[0].childNodes[0].nodeValue)
                xmax = np.float64(bndbox.getElementsByTagName('xmax')[0].childNodes[0].nodeValue)
                ymax = np.float64(bndbox.getElementsByTagName('ymax')[0].childNodes[0].nodeValue)
                boxes.append([xmin, ymin, xmax, ymax])
    
            boxes = torch.as_tensor(boxes, dtype=torch.float32)
            # there is only one class
            labels = torch.as_tensor(labels, dtype=torch.int64)
    
            image_id = torch.tensor([idx])
            area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
            # suppose all instances are not crowd
            iscrowd = torch.zeros((len(objects),), dtype=torch.int64)
    
            target = {}
            target["boxes"] = boxes
            target["labels"] = labels
            # 由于训练的是目标检测网络,因此没有教程中的target["masks"] = masks
            target["image_id"] = image_id
            target["area"] = area
            target["iscrowd"] = iscrowd
            if self.transforms != None:
                sample = {
                    'image': img,
                    'bboxes': target['boxes'],
                    'labels': labels,
                }
                sample = self.transforms(**sample)
                img = sample['image']
                target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)
    
            return img, target
    
        def __len__(self):
            return len(self.imgs)
    
        @staticmethod
        def get_transform():
            # 图像增强进行翻转和类型的转换
            return A.Compose([
                A.HorizontalFlip(p=0.5),
                A.Resize(height=800, width=800),
                ToTensorV2()
            ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})
    
  • 📄 yolo_detect.py 937B
    """
    YOLOv10 目标检测 Demo(官方 ultralytics >=8.3.0)
    首次运行会自动下载 yolov10n.pt 到 ~/.cache/ultralytics
    """
    import cv2
    from ultralytics import YOLO
    
    def main():
        # 用统一的 YOLO 类加载 YOLOv10 预训练权重
        model = YOLO("yolov10n.pt")  # 自动下载;也可用 yolov10s/m/l/x.pt
    
        # 1) 对单张图片做推理
        img_path = "img.jpg"
        results = model(img_path, imgsz=640, conf=0.25)
        r = results[0]
    
        print("检测到", len(r.boxes), "个目标")
        print("类别:", [model.names[int(c)] for c in r.boxes.cls.tolist()])
    
        # 画框并保存
        annotated = r.plot()
        cv2.imwrite("result.jpg", annotated)
        print("结果已保存为 result.jpg")
    
        # 2)(可选)训练示例 —— 取消注释即可在 COCO8 小数据集上试跑
        # model.train(data="coco8.yaml", epochs=3, imgsz=640, device="cpu")
    
    if __name__ == "__main__":
        main()
    

第4节 · 多模态与 AIGC

🎯 了解 AIGC 背景、图像生成方法与 StableDiffusion 实践(PPT 课件,附文字提取版)。
📎 资料

多模态课件(5 讲,附文字提取版)

  • 📎 01-AIGC 背景.pptx4.8MB
  • 📎 02-图像生成方法.pptx原文 10.0MB 过大未收录 · 保留文字提取版
  • 📎 03-StableDiffusion详解.pptx9.9MB
  • 📎 04-StableDiffusion实践.pptx原文 20.3MB 过大未收录 · 保留文字提取版
  • 📎 05-腾讯云AI绘画.pptx原文 23.1MB 过大未收录 · 保留文字提取版
📄 01-AIGC 背景.pptx 文字提取版

以下为 PPT 文字提取版(图表无法提取,原文见附件)。

第 1 页

AIGC介绍

第 2 页

AIGC是什么?

AIGC全称为AI-Generated Content,人工智能内容生成

1

...................................................................

AIGC的应用场景

AI解决实际问题,提高效率、降低成本

2

...................................................................

AIGC的产品形态

基础层,中间层,应用层

3

...................................................................

第 3 页

AIGC是什么?

AIGC全称为AI-Generated Content,人工智能内容生成。

1

........................................................................

第 4 页

从技术上划分

..............................

1、内容孪生

2、内容编辑

3、内容生成

内容孪生是数字内容映射到另一个模态的技术,包括智能增强和智能转译,应用场景包括图像超分、语音转字幕,文字转语音等

内容编辑通过理解内容及属性控制实现修改,应用场景:场景剪辑、虚拟试衣、人声分离等

内容生成通过学习海量数据中的抽象概念,并生成全新的内容,应用广泛,包括图像生成(AI绘画)、文本生成(AI写作),视频生成、多模态生成等

第 5 页

从生成内容上划分

............................. ...... ...... .

1、文本生成

基于NLP的文本内容生成分为非交互式和交互式,包括摘要/标题生成、聊天机器人,图生文等,ChatGPT、文心一言等是代表性产品

2、图像生成

图像生成分为编辑修改和自主生成,应用广泛,包括超分、修复、人脸替换,文生图等,代表性产品有Midjourney、DALL-E,文心一格等

3、音频生成

音频生成技术可用于C端产品,如语音克隆和场景化语音生成,代表性产品有DeepMusic、WaveNet等

4、视频生成

视频生成分为编辑和自主生成,包括超分、修复、剪辑等,代表性产品有Deepfake、videoGPT等

5、多模态生成

组合搭配,进行模态间的转换生成,如文本生成图像、文本生成音频、文本生成视频、图像生成文本、图像生成视频,代表性产品有DALL-E、MidJourney、Stable Diffusion等

第 6 页

腾讯云AI绘画

............................. ...... ...... .

腾讯云AI绘画是一款基于腾讯优图实验室自研的 AI 图像生成与编辑技术产品,目前提供 API和SaaS两种产品形态。 可以结合输入的文本或图片智能创作出与输入相关的图像内容; 具有更强大的中文理解能力、更多样化的风格选择; 更好支持中文场景下的建筑风景生成、古诗词理解、水墨剪纸等中国元素风格生成; 各种动漫、游戏风格的高精度图像生成和风格转换; 为高质量的内容创作、内容运营提供技术支持。

第 7 页

腾讯混元大模型

............................. ...... ...... .

由腾讯全链路自研的实用级大语言模型,具备强大的中文创作能力,复杂语境下的逻辑推理能力,以及可靠的任务执行能力。用户可以通过腾讯云直接调用API接口,或者将腾讯混元作为基底模型,在公有云上进行精调,为不同产业场景构建专属应用。

第 8 页

腾讯混元大模型

............................. ...... ...... .

多轮对话:具备上下文理解和长文记忆能力,流畅完成各专业领域的多轮问答

内容创作:支持文学创作、文本概要、角色扮演能力,流畅、规范、中立、客观

逻辑推理:准确理解用户意图,基于输入数据或信息进行推理、分析

知识增强:有效解决事实性、时效性问题,提升内容生成效果

第 9 页

AIGC的应用场景

AI解决实际问题,提高效率、降低成本

2

........................................................................

第 10 页

AIGC在影视行业应用

............................................................................

01

在前期创作阶段,AIGC可以通过学习海量剧本,按照预定风格生成剧本,为创作者提供更多灵感和选择。 创作者可以对生成的剧本进行二次筛选与加工,不断完善作品,缩短创作周期。 这种智能辅助创作工具可以带来更多创新和想法,为创作者提供更多灵感和可能性

AIGC学习生成剧本, 提高创作效率

02

在中期拍摄阶段,人工智能的应用能够合成虚拟场景,以此生成无法实拍或成本过高的场景,提升视听体验。 在拍摄前,团队会收集大量场景素材并建模制作虚拟场景。 演员在绿棚中拍摄,实时人员识别与抠图技术会将人物嵌入虚拟场景中进行融合,最终生成视频

中期拍摄可合成虚拟场景

03

在视频后期制作阶段,AIGC技术可应用于画质增强,以提升视频质量。 对于涉及敏感人员如“劣迹艺人”的情况,可使用“人脸替换”和“人声替换”技术对视频进行编辑,以确保内容符合社会价值观。 同时,AI技术可自动对视频片段进行剪辑,极大节省了预告片和片段集锦的制作时间和成本

AIGC增强视频画质 AI剪辑视频片段

第 11 页

AIGC在电商行业应用

............................................................................

AIGC生成3D模型,提升线上购物体验

在商品展示环节,AIGC生成3D模型用于商品展示和虚拟适用,提升线上购物体验。通过这种方式,消费者可以更直观地了解商品的外形、尺寸和效果,提高购买决策的准确性和满意度

在主播打造环节:打造虚拟主播,赋能直播带货;

在主播打造环节,我们注重打造虚拟主播,以创新形式赋能直播带货。通过引入人工智能技术,虚拟主播能够实现自动化直播,提高效率并降低成本

虚拟商城构建,智能聊天机器人赋能线上线下秀场加速演变,提供全新购物场景

在交易场景环节,通过虚拟商城构建和智能聊天机器人,赋能线上和线下秀场加速演变,为消费者提供全新的购物场景,提升购物体验

第 12 页

AIGC在娱乐行业应用

............................................................................

01

人脸美妆、图像融合、黑白上色、风格转换、属性变换

02

社交互动:虚拟主播、网红、聊天机器人、互动游戏

全民娱乐,享受科技带来的美妙体验!通过图像内容生成应用,探索无限创意!尝试人脸美妆,展现个人独特风格;融合照片,创作出更多精彩作品!

社交互动形式日新月异,随着技术的不断发展,虚拟主播、虚拟网红、聊天机器人和聊天互动游戏等新型社交互动形式逐渐成为人们生活中的重要组成部分。这些新型社交互动形式不仅丰富了人们的社交体验,还为社交互动市场注入了新的活力

第 13 页

AIGC在其他行业应用

............................................................................

AIGC为教育工作者提供丰富的教学素材,提升课堂互动和提供更广泛的授课思路

AIGC为教育工作者提供丰富的教学素材,如数字人生成技术生成历史人物对话,提升课堂互动,ChatGPT生成创意性教学方案,提供广泛的授课思路

AIGC技术结合CAD软件,通过文本提示生成特定机构模型,用于工业设计参考

在工业行业,AIGC技术可与工业设计软件CAD融合,通过文本输入提示语生成特定样式的机构模型,供设计者参考。例如,输入“设计一款卫星太阳能电池板可伸缩折翼机构”,AIGC模型即可生成相应的3D设计机构

AIGC突破内容生成行业,提高各行业工作效率与质量,加速企业数字化智能化进程

AIGC在内容生成行业的突破,将提升内容创作者、设计师、工程师、教育工作者等各行业人员的工作效率与质量,同时加速企业数字化与智能化进程,为企业带来更多商业机会和发展空间

第 14 页

AIGC产品形态

基础层,中间层,应用层

3

........................................................................

第 15 页

AIGC产品形态

............................................................................

📄 02-图像生成方法.pptx 文字提取版

以下为 PPT 文字提取版(图表无法提取,原文见附件)。

第 1 页

图像生成方法

第 2 页

图像生成算法

1

...................................................................

基于扩散模型的算法

2

...................................................................

第 3 页

图像生成算法

1

........................................................................

第 4 页

算法分类

............................................................................

2014 变分自编码(VAE)

2014 生成对抗网络(GAN)

2015 扩散模型

扩散是受到非平衡热力学的启发,定义一个扩散步骤的马尔科夫链,并逐渐向数据中添加噪声,然后学习逆扩散过程,从噪声中构建出所需的样本。扩散模型的最初设计是用于去除图像中的噪声。随着降噪系统的训练时间越来越长且越来越好,可以把纯噪声作为输入,生成逼真的图片。

变分自编码器由Kingma等人在2014年提出,与传统的自编码器通过数值方式描述潜空间不同,它以概率方式对潜在空间进行观察,在数据生成方面应用价值较高。

2014年 Ian GoodFellow提出了生成对抗网络,成为最著名的生成模型。GAN使用零和博弈策略学习,在图像生成中应用广泛。以GAN为基础产生了多种变体,如DCGAN,StytleGAN,CycleGAN等。

第 5 页

变分自编码(Variational Autoencoder,VAE)

............................................................................

VAE是一种深度生成模型,描述潜在空间概率,有高应用价值

VAE是深度生成模型

图像经解码器得多种特征,包括微笑、肤色、性别等。自编码器传统生成具体数值潜在特征,目的是生成更近似输入的图像特征。概率分布采样特征值范围,生成图像的潜在特征表示解码器生成的图像

图像解码器生成特征表示

VAE包括编码器和解码器,编码器将高维数据转换为潜在空间概率分布,解码器从采样数据重建生成新数据

VAE编码器解码器, 潜在空间采样

第 6 页

生成对抗网络(Generative Adversarial Networks,GAN)

............................................................................

以生成对抗网络为基础,应用于人脸替换、卡通头像生成,实现图像的超分辨率重建和风格迁移等

判别器:训练阶段,用真实与生成数据训练二分类模型,输出概率值,越接近1,真实数据可能性越大

生成器:随机生成蓝色网络结构

GAN介绍

生成器和判别器

2014年Ian GoodFellow提出生成对抗网络,是最早的生成模型之一, 广泛应用于图像生成,并衍生出多种变体如DCGAN,StytleGAN,CycleGAN等

生成器学习生成合理数据,如图像生成器给定一个向量会生成一张图片,其生成的数据作为判别器的负样本。 判别器判别输入是生成数据还是真实数据,网络输出越接近于0,生成数据可能性越大;反之,真实数据可能性越大

第 7 页

扩散模型(Diffusion Model,里程碑式模型)

............................................................................

第 8 页

扩散模型(Diffusion Model,里程碑式模型)

............................................................................

向原始图像中添加噪声,使原始图像成为随机噪声

从噪声中恢复的原始图像的变种图像。

第 9 页

扩散模型(Diffusion Model,里程碑式模型)

............................................................................

应用: 图像超分、图像上色、文本生成图片、全景图像生成等。 如右图,中间图像作为输入,基于扩散模型,生成左右视角两张图,输入图像与生成图像共同拼接程一张全景图像。

在扩散模型的基础上,各大公司与研究机构研发出了自己的代表产品,例如: DALL-E1~E3( OpenAI ) Imagen( Google ) Stable Diffusion( Stablitity AI ) FLUX( BlackForestLabs )

第 10 页

CLIP

............................................................................

01

CLIP发布

2021年美国OpenAI公司发布了跨模态预训练大模型CLIP,采用双塔模型与比对学习训练方式进行训练,从互联网收集4亿对图文对。CLIP不能用来进行图像生成,却是图像生成的基础

02

模型结构

CLIP将图片与图片描述一起训练,达到的目的:给定一句文本,匹配到与文本内容相符的图片;给定一张图片,匹配到与图片相符的文本。

03

CLIP训练过程

04

CLIP模型应用

跨模态检索:如搜索中文本搜索图片、文本搜索视频,图片搜索文本等。 跨模态内容生成: 文本生成图片(DALL-E3,Stable Diffusion等) 图片生成标题 图片生成描述等。

使用TE和IE进行文本和图像特征提取,其中TE采用TT模型,IE采用CNN或ViT。 进行对比学习,将N个文本特征和N个图像特征两两组合,预测N2个可能的文本-图像对的相似度,使用CS计算相似度。 训练目标是最大化N个正样本的相似度,同时最小化N2-N个负样本的相似度。 训练模型,优化目标函数。 完成训练后,输入文本可预测匹配图片,输入图片可预测匹配文本。

第 11 页

基于扩散模型的算法

2

........................................................................

第 12 页

DALL-E

............................................................................

DALL-E 由美国OpenAI公司发布,在OpenAI网站向公众开放,提供数量有限的免费图像和额外的购买图像服务。

如下图,为DALL-E所生成的各种图像。

第 13 页

Imagen

............................................................................

Imagen是2022年5月谷歌发布的文本到图像的扩散模型,该模型目前不对外开放。用户可通过输入描述性文本,生成图文匹配的图像。

具体的流程如下所示:

第 14 页

Stable Diffusion

............................................................................

AI作画质量高,内容契合度高

AI作画根据prompt提示语“郊区街区一栋房子的照片,灯光明亮的超现实主义艺术,高度细致8K”生成图像,风格与内容契合度高,作画质量高

Stable Diffusion是开源的AI图像生成器

Stable Diffusion训练昂贵,影响巨大

"Stable Diffusion最初用价值不菲的4000台A100显卡训练,体现其财力、理念和贡献精神。其价值和对业界影响堪称AI革命。"

Stable Diffusion是由Stability AI、CompVis团队和RunwayML共同开源的强大文本到图像生成器,免费且开源,是AI图像生成发展的里程碑

📄 03-StableDiffusion详解.pptx 文字提取版

以下为 PPT 文字提取版(图表无法提取,原文见附件)。

第 1 页

Stable Diffusion详解

第 2 页

Stable Diffusion 的简介

概念、背景、特点

1

...................................................................

Stable Diffusion 原理

算法原理 模型构成

2

...................................................................

Stable Diffusion 处理流程

文生图的流程

3

...................................................................

Stable Diffusion 应用场景

文生图 图生图 图像修复

4

...................................................................

第 3 页

Stable Diffusion 简介

概念、背景、特点

1

........................................................................

第 4 页

Stable Diffusion相关概念

......................................................................

第 5 页

Stable Diffusion 是如何诞生的

............................................................................

01

02

04

05

06

03

Stable Diffusion 于 2022 年 8 月 22 日正式开源,融资1亿美元,估值10亿美元

训练模型需要高昂的计算成本和资源要求 Stability 投资资源,三个团队共同开发

Latent Diffusion 模型:慕尼黑大学的机器学习研究小组 CompVis 和纽约的 Runway 团队合作研发,生成速度快、消耗需求小

Stable Diffusion是由Stability AI、CompVis和Runway团队合作开发的

2022年7月, Stable Diffusion 正式亮相 Stable Diffusion 生成模型会更加准确,且支持的图像分辨率也更高,比 Latent Diffusion 模型更加强大。

2023年6月,福布斯报道:Stability 创始人夸大 Stable Diffusion 成就的新闻:https://www.forbes.com/sites

第 6 页

Stable diffusion的特点-可拓展性强

............................................................................

Stable Diffusion拓展性强,可免费学习创作

Stable Diffusion拥有强大的扩展性,任何人都可以通过官方免费公开的论文和模型代码进行学习与创作,解锁AI艺术更多可能性

Stable Diffusion集成了110+功能

Stable Diffusion是一款功能强大的图像处理软件,集成了110多个扩展插件的功能介绍,包括局部重绘、人物姿势控制、图像高清修复和线稿提取等,用户可以根据需要选择不同的插件来扩展软件的功能

精准控制图像分辨率

行业内无法精准控图和图像分辨率低等问题都是其普遍痛点,但这些在 Stable Diffusion 中通过 ControlNet 和 UpScale 等插件都能完美解决。

第 7 页

Stable diffusion的特点-出图快

..................................................-c..........................

使用过 Midjourney 的朋友应该都有体验过排队等了几分钟才出一张图的情况,而要得到一张合适的目标图我们往往需要试验多次才能得到。 由于商业级的 AI 应用都是通过访问云服务器进行操作,除了排队等候服务器响应外,网络延迟和资源过载都会导致绘图效率的降低。

而部署在本地的 Stable Diffusion 绘图效率完全由硬件设备决定,只要显卡算力跟得上,甚至能够实现 3s 一张的超高出图效率。

第 8 页

Stable diffusion的特点-数据安全

............................................................................

因为部署在本地,Stable Diffusion 在数据安全性上是其他在线绘画工具无法比拟的。商业化的绘画应用为了运营和丰富数据,会主动将用户的绘图结果作为案例进行传播。但 Stable Diffusion 所有的绘图过程完全在本地进行,基本不用担心信息泄漏的情况。

第 9 页

Stable Diffusion 运行原理

算法原理 模型构成

2

........................................................................

第 10 页

Stable Diffusion的工作原理

............................................................................

Diffusion模型广泛应用于图像生成领域,其他常见的生成式模型还有GAN、VAE、Flow based Model等。这些模型都基于深度学习进行训练

扩散模型之所以用 Diffusion 来命名,因为它的运作过程就是向训练图像不断地添加噪声,直到变为一张无意义的纯噪声图,再逐步恢复的过程。

第 11 页

加噪过程类似于在干净的水中滴墨汁,颜料会逐渐扩散导致整个水体变浑,AI绘画过程是从模糊的噪声图逐渐变为清晰的图像,即逆向降噪的过程 标准扩散模型在像素空间中进行,需要处理大量数据和高性能显卡 潜在扩散模型先缩小图像再运算,减少数据处理量和硬件要求 潜在扩散模型运算空间为潜空间,即低维空间

Stable Diffusion的工作原理

............................................................................

第 12 页

Diffusion Model的整体思路如下图所示: 正向过程:对于一张原始图片X0,我们给X0添加一个高斯噪声,图片由X0变成变成X1。接着我们会在X1的基础上再添加高斯噪声得到X2。重复上述添加高斯噪声步骤,直到图片变成Xn,由于添加了足够多的高斯噪声,现在的X近似服从高斯分布。 反向过程:随机生成一个服从高斯分布的噪声图片,然后一步一步的减少噪声直到生成想要的图片。 我们希望将t时刻的高斯噪声变成0时刻的图像,是很难一步到位的,和正向过程一样,先考虑t时刻图像和t−1时刻的关系,然后一步步向前推导得出结果。

Diffusion Model的原理

...................... ......................................................

第 13 页

Diffusion Model的实现流程:

Diffusion Model的原理

...................... ......................................................

第 14 页

图像空间是巨大的,具有三个颜色通道的 512×512 图像是一个 786432 维的空间! Stable Diffusion旨在解决速度问题。Stable Diffusion利用潜在空间扩散(Latent diffusion)的模型。它不是在高维图像空间中操作,而是首先将图像压缩到潜空间(Latent space)中,潜空间(Latent space)相对于原始图像空间小了很多倍,所以速度会有较大的提升。

Diffusion Model的原理

...................... ......................................................

第 15 页

Stable Diffusion的架构

............................................................................

Stable Diffusion 模型并不是单一的一个模型,而是多个模型组成的运作系统,其中的技术可以拆解为 3 个结构来看:

第 16 页

Stable Diffusion的架构-ClipText 文本编码器

...........................................................................

CLIP 是由 OpenAI 提出的一种多模态神经网络,它能有效地借助自然语言的监督来学习视觉的概念。典型的视觉数据集是劳动密集型的,创建成本很高,且使用它们训练得到的标准视觉模型只擅长一项任务,导致适应新的任务并不容易。而 CLIP 在各种各样的图像上进行训练,同时依赖于互联网上大量自然语言的监督。

在 SD 模型运作过程中,它可以提取提示词文本部分的特征传递给图像生成器,让模型理解我们输入的提示词内容,从而达到文本控制图像生成的目的。

第 17 页

Stable Diffusion的架构-ClipText 文本编码器

...........................................................................

首先训练一个处理图像的 CNN 和一个处理文本的 Transformer 模型,来预测图像的 Caption 对比学习阶段:给定一个 Batch 的 N 个 (图片,文本) 对,图片输入给 Image Encoder 得到表征 I1,I2,....,In,文本输入给 Text Encoder 得到表征 T1,T2,....,Tn,其中 (Ij,Tj)属于是正样本,(Ij,Tk)属于负样本。最大化 N 个正样本的 Cosine 相似度,最小化N2−N个负样本的 Cosine 相似度。

第 18 页

Stable Diffusion的架构-ClipText 文本编码器

...........................................................................

Zero-Shot Transfer:这个阶段是使用 CLIP 的预训练好的 Image Encoder 和 Text Encoder 来做 Zero-Shot Transfer。比如来一张 ImageNet-1K 验证集的图片,可以使用CLIP 预训练好的模型能完成这个分类的任务 模型的参数是冻结的,在图像生成过程中参数是不会发生变化的

第 19 页

Stable Diffusion的架构-Diffusion 扩散模型

............................................................................

U-Net 原本是用于生物医学图像分割的神经网络模型,因为工作结构像一个 U 型字母,因此被称为 U 型神经网络训练模型。在扩散模型中,U-Net 可以辅助提取并解构训练图像的特征,有了它就能在较少训练样本的情况下获得更加准确多样的数据信息,从而使模型在出图结果上更加精确。

在训练扩散模型过程中,可以采用不同的算法来添加噪音,而 Scheduler 就是用来定义使用哪种算法来运行程序,它可以定义降噪的步骤、是否具备随机性、查找去噪后样本的算法等,因此它又被称为采样算法。在 WebUI 中,它是可调节的一项,我们可以根据图像类型和使用的模型来选择不同的采样器,从而达到更佳的出图效果。

第 20 页

Stable Diffusion的架构-VAE 模型

............................................................................

最后的解码器 VAE,全称是 Variational Auto Encoder 变分自动编码器。简单来说,它的作用就是将高维数据(像素空间)映射到低维空间(潜空间),从而实现数据的压缩和降维。它由编码器(Encoder)和解码器(Decoder)两部分组成,编码器用于将图像信息降维并传入潜空间中,解码器将潜在数据表示转换回原始图像,而在潜在扩散模型的推理生成过程中我们只需用到 VAE 的解码器部分。

整个过程只运行一次即可生成最终的像素图像

第 21 页

Stable Diffusion 处理流程

文生图的流程

3

........................................................................

第 22 页

Stable Diffusion的处理流程

............................................................................

Stable Diffusion从文本生成图像的流程如下所示: Stable Diffusion在潜空间中生成随机张量。可以通过设置随机数生成器的来控制此张量。如果将种子设置为某个值,获得相同的随机张量。 网络 U-Net 将潜在噪声图像和文本提示作为输入,并预测噪声,也在潜在空间(4x64x64 张量)中。 从潜在图像中减去潜在噪声。这将成为新潜在图像 步骤 2 和 3 重复一定数量的采样步骤,例如 20 次 最后,VAE的解码器将潜在图像转换回像素空间。这是运行Stable Diffusion后获得的图像。

第 23 页

Stable Diffusion应用场景

文生图 图生图 图像修复

4

........................................................................

第 24 页

Stable diffusion的应用场景

............................................................................

文生图

Stable Diffusion将输入的文字转化为图像效果 ,使文本更加生动

图生图

图生图就是依赖图片和提示词进行二次创作。从而减少图片生成的随机性,更好的满足我们的需求。 图生图也可生成带有深度信息的图像。

图像修复

图像修复通过替换&填充图像中指定区域的纹理/像素,使其与周围纹理/像素融为一体,达到自然和谐的修复效果。

📄 04-StableDiffusion实践.pptx 文字提取版

以下为 PPT 文字提取版(图表无法提取,原文见附件)。

第 1 页

Stable Diffusion的训练与部署

第 2 页

高性能应用服务HAI

平台介绍 服务优势 应用场景

1

...................................................................

创建Stable Diffusion实例

申请服务 创建服务 WebUI Jupyter_Lab 服务销毁

2

...................................................................

Stable Diffusion的训练

训练方法 Dreambooth训练 LoRA训练

3

...................................................................

Stable Diffusion的预测

Dreambooth预测 LoRA预测

4

...................................................................

第 3 页

高性能应用服务HAI

平台介绍 服务优势 应用场景

1

........................................................................

第 4 页

高性能应用服务HAI

............................................................................

推荐使用搭载N卡的Windows系统,MacOS在默认配置下仅支持CPU运算,图形性能较差,插件支持有限,且训练模型困难。 推荐使用云端部署Stable Diffusion,将其部署在云服务器上使用。

Stable Diffusion模型的配置要求:

第 5 页

高性能应用服务HAI

..........................................................................

澎湃算力,即开即用。以应用为中心,匹配GPU云算力资源,助力中小企业及开发者快速部署LLM、AI作画、数据科学等高性能应用

无需复杂配置,便可享受即开即用的GPU云服务体验。在 HAI 中,根据应用智能匹配并推选出最适合的GPU算力资源,以确保在数据科学、LLM、AI作画等高性能应用中获得最佳性价比。

第 6 页

服务优势

............................................................................

智能选型 根据应用匹配推选GPU算力资源,实现最高性价比。同时,打通必备云服务组件,大幅简化云服务配置流程。

一键部署 分钟级自动构建LLM、AI作画等应用环境。提供多种预装模型环境,包含如StableDiffusion、ChatGLM等热门模型。

可视化界面 提供开发者友好的图形界面,支持JupyterLab、WebUI等多种算力连接方式,AI研究调试超低门槛。

基于腾讯云GPU云服务器底层算力,提供即插即用的高性能云服务。

第 7 页

服务优势

............................................................................

大幅降低GPU云服务器使用门槛,多角度优化产品使用体验,开箱即用

第 8 页

应用场景-AI作画

............................................................................

基于Stable Diffusion开源模型进行AI绘画

场景介绍 AI绘画是一种利用深度学习算法进行创作的绘图方式。广泛应用于数字媒体、游戏、动画、电影、广告等领域。 业务痛点 GPU卡型多样,算力、显存差异大,选型困难 环境配置复杂、模型安装和调试门槛高 各类插件迭代频繁,难以在进行环境管理 产品优势 智能匹配算力,多种算力套餐满足不同需求的绘图性能。 预置主流AI作画模型及常用插件,无需手动部署,支持即开即用。 动态更新模型版本,确保模型版本与时俱进,无需频繁操作。

第 9 页

应用场景-AI对话/写作

............................................................................

场景介绍 大语言模型在广泛的文本数据上进行训练,可以执行广泛的任务,包括文本总结、翻译、情感分析等等。 业务痛点 环境配置复杂,部署难度大。 模型效果调试难,无可视化界面。 模型、数据集储存流程繁琐,难以快捷保存 产品优势 预置国内外主流LLM大语言模型。 支持可视化界面一键登录,方便调优 打通云上存储组件,支持模型、数据集快捷存储

基于开源大语言模型,创作属于自己的Agent、企业知识库

第 10 页

应用场景-AI开发/测试

............................................................................

算法研发:学术研究、论文

场景介绍 面向高校、研究所等大量科研场景,需针对深度学习、机器学习等前沿算法进行开发探索。 业务痛点 需要大量的计算资源进行验证和优化。 涉及多种框架、模型及开发库,环境配置复杂。  产品优势 提供多种高性能GPU云服务器,满足算法验证和测试的需求。 提供公共模型、数据集文件存储桶,优化资源拉取效率。

第 11 页

创建stable diffusion实例

申请服务 创建服务 WebUI Jupyter_lab 服务销毁

2

........................................................................

第 12 页

Stable Diffusion实例创建

............................................................................

使用高性能应用服务HAI快速搭建并使用AI模型 Stable Diffusion ,实现过程如下:

申请高性能应用服务HAI

第 13 页

服务申请

............................................................................

点击链接进入 高性能应用服务 HAI 中【申请资格】

等待审核通过后,点击【前往体检Hai】即可进入高性能应用平台Hai

第 14 页

创建服务

............................................................................

登录 高性能应用平台Hai的【控制台】,如下所示:

点击【新建】即可进入创建实例的页面。

第 15 页

创建服务

..........................................................................

服务创建时,要配置服务器信息,如下所示:

第 16 页

创建服务

..........................................................................

服务创建完成,可查看相关状态,如下所示:

第 17 页

创建服务

............................................................................

也可查看服务的详细信息,如下所示:

第 18 页

webui

............................................................................

要进入 StableDiffusionWebUI ,可选择以下两种方式:

接下来即可进入 Stable Diffusion WebUI 中,进行推理处理。

第 19 页

webui

............................................................................

WebUI 生成图像的界面如下所示:

第 20 页

JupyterLab

............................................................................

要进入 JupyterLab ,可选择以下两种方式:

接下来即可进入 JupyterLab中。

第 21 页

JupterLab

..........................................................................

JupyterLab界面如下所示:

JupyterLab的使用方式与云服务器的使用是一样的。

第 22 页

JupyterLab

JJ..........................................................................

建议使用Stable Diffusion前先安装一些主要的前置应用,如Python、Git和CUDA,以减少报错的可能性并提高下载速度。

第 23 页

Stable Diffusion的训练

训练方法 Dreambooth训练 LoRA训练

3

........................................................................

第 24 页

常见的训练方法

............................................................................

Stable Diffusion有四种训练模型的方法,如下所示:

第 25 页

DreamBooth训练

............................................................................

DreamBooth 主要是对文本生成图片模型进行微调,它使用户可以用几张个性化的图片(比如个人照片,自家宠物照片),以很少的训练成本就可以训练一个专属的模型。

对于这一组个性化的图片,首先需要定义一个专属的 prompt,比如 a [V] dog, [V] 是一个罕见没有什么含义的符号组合。在训练的过程中,首先用文本编码器提取这个专属 prompt 的文本特征,然后用图像编码器提取图片的特征,以文本特征,图像特征,作为扩散模型的输入,进行模型训练。预测输入关键词即可生成对应风格的图片。

第 26 页

DreamBooth训练-实验过程

............................................................................

本实验基于开源代码diffusers,显存需要24G以上,在HAI平台的JupyterLab中实现。代码结构如下所示:

第 27 页

DreamBooth训练-数据准备

............................................................................

收集一些格子衫的图片,将图像统一处理成 512*512 格式的图片:

该数据集中包含103 张格子衫图片,为了使 DreamBooth 模型更加多样化,尽量使用不同的环境、灯光、发型、表情、姿势、角度等,将这些数据放入文件夹train_images中。

第 28 页

DreamBooth训练-开始训练

............................................................................

DreamBooth训练的代码在脚本dreambooth-gezishan.ipynb中,首先是全局变量的设置,如下所示 :

第 29 页

DreamBooth训练-开始训练

............................................................................

接下来实现源模型到diffusers模型的转换,如下所示 :

第 30 页

DreamBooth训练-开始训练

............................................................................

格子衫任务训练的prompt是:Plaidshirtprogrammer,将其配置如下所示:

INSTANCE_PROMPT 中填入 Plaidshirtprogrammer。 可以指定任意关键词,但是需要找一个“不存在的词”。 CLASS_PROMPT 是让AI自动生成 image用的prompt,不要有专属关键词即可。 预览图设置需要修改成你想生成的图片的prompt

第 31 页

DreamBooth训练-开始训练

............................................................................

加速器的配置如下所示:

该配置文件包含了一些与 Accelerated Inference 相关的默认配置,如计算环境、分布式训练设置、混合精度等。这样的配置文件通常用于定义加速推理的运行环境和选项:

第 32 页

DreamBooth训练-开始训练

............................................................................

格子衫任务训练常用参数配置,如下所示:

第 33 页

DreamBooth训练-开始训练

............................................................................

格子衫任务训练参数配置,如下所示: (这些参数一般不需要修改)

第 34 页

DreamBooth训练-开始训练

............................................................................

执行以下代码即可进行模型训练:

第 35 页

DreamBooth训练-开始训练

............................................................................

模型训练的输出如下所示:

第 36 页

DreamBooth训练-开始训练

............................................................................

模型训练完成后,会保存训练权重:

这种模型权重无法直接在WebUI中使用,我们需要将其转换为ckpt的形式,转换代码如下所示:

第 37 页

LoRA训练

....................................................................

LoRA的全称是 Low-Rank Adaptation of Large Language Models,即大语言模型的低阶适应。

第 38 页

Lora训练-实验过程

....................................................................

LoRA训练基于开源代码kohya_ss来实现,显存需要8G以上,在HAI平台的JupyterLab中实现。具体结果如下所示:

第 39 页

Lora训练-准备数据

....................................................................

使用LoRA方法进行训练给我们的任务带上眼镜,我们收集了一系列的数据,如下所示:

该数据集包含20张图片,每张图片对应着一个txt文件,文本文件中记录着对应图片的文本描述。数据放入文件夹train中,存放方式如下所示:

第 40 页

LoRA训练-模型训练

....................................................................

首先配置LoRA训练的预训练模型路径、模型类型、参数化、训练数据集路径以及指定网络结构和超参数。如下所示:

这里需要将前面训练的格子衫模型放入指定文件夹中,并放置好训练图片的位置即可,其他参数按上述进行配置。

第 41 页

LoRA训练-模型训练

....................................................................

配置LoRA训练相关的参数设置。主要配置了图像分辨率、批次大小、最大训练轮数、模型保存频率以及训练模式(是否仅训练U-Net或文本编码器),如下所示:

第 42 页

LoRA训练-模型训练

....................................................................

设置LoRA学习率,其中包括整体模型学习率 (lr)、U-Net 部分的学习率 (unet_lr) 以及文本编码器的学习率 (text_encoder_lr)。选择了余弦退火重启模式作为学习率调度器 (lr_scheduler),并设置了相应的参数,如预热步数 (lr_warmup_steps) 和余弦退火的重启周期数 (lr_restart_cycles)。还设置了优化器的类型 (optimizer_type),默认选择了 "AdamW8bit",如下所示:

第 43 页

LoRA训练-模型训练

....................................................................

最后配置输出模型的名称和保存格式,是否保存训练状态,以及其他与数据处理相关的设置,如下所示:

第 44 页

LoRA训练-模型训练

....................................................................

配置好相应的参数,启动分布式训练的命令,进行模型训练,代码如下所示:

第 45 页

Lora训练-模型训练

....................................................................

训练结果如下所示:

第 46 页

LoRA训练-模型训练

....................................................................

训练模型存储结果如下所示:

接下来我们就可以使用DreamBooth 和LoRA训练出的模型进行预测。

第 47 页

模型训练

............................................................................

Hugging Face(抱脸网):专注于构建、训练和部署先进开源机器学习模型的网站,包含的内容非常广泛,不仅仅包括 AI 绘画,还包括很多其他 AI 领域的东西。 网站地址: https://huggingface.co/models

Civitai(C 站):专业的 AI 绘画模型分享平台,无需注册即可访问和下载模型,包含模型的详细介绍、使用教程、参考图等丰富内容 网站地址: https://civitai.com/

哩布哩布 AI:国内的 AI 绘画原创模型网站,包含大量模型和绘图作品,还支持在线 Stable Diffusion 绘图。 网站地址: https://www.liblib.art/

预训练模型的获取方式:

第 48 页

Stable Diffusion的预测

DreamBooth预测 LoRA预测

4

........................................................................

第 49 页

模型预测

....................................................................

模型预测在webUI中完成:

第 50 页

DreamBooth模型预测

....................................................................

首先我们将格子衫模型添加到stable-diffusion-webui/models/Stable-diffusion文件夹中

第 51 页

DreamBooth模型预测

....................................................................

接下来我们在UI端进行设置,注意需要在正向prompt中添加dreambooth的prompt:

点击【生成】后在右下角即可查看结果。

第 52 页

LoRA模型预测

...............................................................

首先我们将LoRA训练的模型添加到stable-diffusion-webui/models/Lora中:

第 53 页

LoRA模型预测

....................................................................

LoRA预测的方法与DreamBooth的方式是类似的,整个操作流程如下所示:

第 54 页

LoRA模型预测

....................................................................

接下来返回文生图的界面:

点击【生成】后在右下角即可查看结果。

第 55 页

高性能应用服务Hai

平台介绍 服务优势 应用场景

1

...................................................................

创建Stable Diffusion实例

申请服务 创建服务 WebUI Jupyter_Lab 服务销毁

2

...................................................................

Stable Diffusion的训练

训练方法 Dreambooth训练 LoRA训练

3

...................................................................

Stable Diffusion的预测

Dreambooth预测 LoRA预测

4

...................................................................

📄 05-腾讯云AI绘画.pptx 文字提取版

以下为 PPT 文字提取版(图表无法提取,原文见附件)。

第 1 页

腾讯云AI绘画

第 2 页

腾讯云AI绘画介绍

服务形态 产品优势 应用场景

1

...................................................................

API的使用

文生图API 图生图API

2

...................................................................

SaaS应用

文生图 图生图

3

...................................................................

最佳实践

API接口调用 任务管理 小程序开发

4

...................................................................

第 3 页

腾讯云AI绘画介绍

服务形态 产品优势 应用场景

1

........................................................................

第 4 页

腾讯云AI绘画

............................................................................

腾讯云AI绘画是一款基于腾讯优图实验室自研的 AI 图像生成与编辑技术产品,能够结合输入的文本或图片智能创作与输入相关的图像内容。

第 5 页

腾讯云AI绘画

............................................................................

中文理解能力强,可以根据输入的图片或文本创作与之相关的图像内容。

风格选择多样化,支持东方审美的绘画创作能力。

支持中文场景下的建筑风景生成、古诗词理解、水墨剪纸等中国元素风格生成。

支持各种动漫、游戏风格的高精度图像生成和风格转换。

为高质量的内容创作、内容运营提供技术支持。

第 6 页

服务形态

............................................................................

API AI 绘画提供 API 技术服务,需要调用 API 接口来使用,适用于有代码编写基础、对 HTTP 请求和 API 调用有一定了解的开发者。

SaaS 腾讯云智能图像创作平台(Tencent Cloud TI-Image Platform)是一款基于AI算法,通过输入文字或图片即 可生成高质量图片的AI绘画在线创作工具。可以直接通过网页端输入文本或图像,调用腾讯云AI绘画完成图像生成任务,不需编程基础。

腾讯云AI绘画提供 API和SaaS两种服务形态:

第 7 页

产品优势

.........................................................................

自研算法

标准化接口封装,调用简单,兼容性强。

第 8 页

应用场景-文化传播

............................................................................

策划者可以将活动形式与潮流的 AIGC 创意玩法相结合,打造生动有趣的艺术文化主题活动,人人皆可参与互动,自由尝试艺术创作,体验文化多样魅力,提升互动氛围和活动影响力。 典型场景:新年活动、诗意作画、万物皆可剪纸等。

第 9 页

应用场景-风格变换

............................................................................

为图片切换各种创意风格,令图片具备独特的艺术特征,呈现出丰富多彩的艺术灵感。 典型场景:头像风格化、场景风格化。

第 10 页

应用场景-文学配图

............................................................................

对描述型的文稿进行素材配图的创作,输入的文稿可以引发读者的想象力和共情,更适合网络文学、非严肃报告、非科学研究说明文等的文章配图。 典型场景:网络小说封面、公众号推文插图等。

第 11 页

应用场景-设计辅助

............................................................................

游戏、漫画、影视等场景的初始化设计和草图创作,前期借助 AI 能力进行大量的实验作为初步的风格、内容、构图尝试积累,从中选择出相对较好的效果参考进行进一步的高精度内容创作,帮助制作人从大量繁琐冗余的细节工作中抽身。 典型场景:游戏原画设计、漫画草图创作、影视分镜场景构建、海报等物料设计等。

第 12 页

API使用

文生图 图生图

2

........................................................................

第 13 页

API使用-文生图

............................................................................

智能文生图接口将根据输入的描述文本,智能生成与之相关的结果图。 输入:256个字符以内的描述性文本,推荐使用中文。 输出:对应风格及分辨率的 AI 生成图。

第 14 页

API使用-文生图示例

.........................................................................

输入示例:

输出示例:

第 15 页

API使用-文生图示例

.........................................................................

腾讯云AI绘画的API的调用方式可查看连接

第 16 页

API使用-文生图示例

...........................................................................

生成的图片结果如下所示:

第 17 页

API使用-图生图

............................................................................

智能图生图接口将根据输入的图片及辅助描述文本,智能生成与之相关的结果图。 输入:单边分辨率小于2000、转成 Base64 字符串后小于 5MB 的图片,建议同时输入描述文本。 输出:对应风格及分辨率的 AI 生成图。

第 18 页

API使用-图生图示例

...........................................................................

输入示例:

输出示例:

第 19 页

AI绘画图生图示例

.........................................................................

腾讯云AI绘画图生图的API的调用方式可查看连接:

第 20 页

API使用-图生图示例

...........................................................................

第 21 页

SaaS应用

3

........................................................................

腾讯云智能图像创作平台

第 22 页

SaaS应用

.........................................................................

腾讯云智能图像创作平台(Tencent Cloud TI-Image Platform)是一款基于AI算法,通过输入文字或图片即 可生成高质量图片的AI绘画在线创作工具。 核心功能:在线创作,提供文生图/图生图能力 产品链接:https://ti-image.cloud.tencent.com/

第 23 页

SaaS应用

.........................................................................

文生图:支持25种风格 图生图:支持5种风格

第 24 页

SaaS应用

.........................................................................

图像生成之后,满意的作品,可以下载也可手动【存至我的作品】,【存至我的作品】的画作,可在【我的作品】页面管理

第 25 页

SaaS应用

.........................................................................

在【我的作品】中,可进行单张下载或删除,也可进行批量下载或删除。

第 26 页

最佳实践

4

........................................................................

小程序端AI绘画实践

第 27 页

实现思路

.........................................................................

用腾讯云AI绘画实现小程序的AI绘画,主要思路如下所示:

第 28 页

代码结构

.........................................................................

腾讯云AI绘画的小程序实践的代码结构:

第 29 页

小程序的实现方式

.........................................................................

腾讯云AI绘画文生图API调用方式:

第 30 页

小程序的实现方式

.........................................................................

使用Golang调用腾讯云AI绘画接口,在draw.go文件中实现,具体如下所示:

第 31 页

小程序的实现方式

.........................................................................

腾讯云接口调用之后 一、用户把任务提交到cos上,绘画服务通过cos的内容来执行AI画画任务。 二、通过执行shell命令,并将生成好的图片上传到cos,具体实现如下所示:

第 32 页

小程序的实现方式

.........................................................................

通过cos来实现任务管理,实现生成随机字符串、生成任务ID、推送任务、拉取任务、查询任务结果等功能,用于任务的异步处理和结果的存储与查询,该部分代码在job.go文件中,具体实现如下所示: :

第 33 页

小程序的实现方式

.........................................................................

在main.go文件中,小程序要通过中转服务来异步处理消息,该服务主要包括: 一、转发请求到AI绘画。 二、查询AI绘画的结果,具体实现如下所示:

第 34 页

小程序的实现方式

.........................................................................

在main.go文件中,main方法中实现了程序的主要逻辑,具体如下所示:

第 35 页

小程序的实现方式

.........................................................................

小程序实现AI绘画,主要在index.js文件中,主要包括以下内容:

第 36 页

小程序的实现方式

.........................................................................

在index.wxml文件中实现了一个简单的图像生成应用界面。用户可以输入创意提示,点击按钮后,后台会生成图像。用户可以查看任务状态、完成数量以及处理和剩余时间。生成的图像会在页面上展示。整个应用通过与后台服务器的通信实现,具体实现如下所示:

第 37 页

执行方式

.........................................................................

1、切换当前工作目录 "drawserver"中 cd drawserver 2、在进入 "drawer" 目录之后,执行以下命令: go build . Go 编译器在当前目录中构建 Go 程序。点号(.)表示当前目录。go build 命令会编译当前目录中的 Go 源代码文件,并生成一个可执行的二进制文件。 3、运行 "drawserver" 的可执行文件,启动小程序服务端和AI画画服务 ./drawserver -w=true

在终端执行以下过程:

第 38 页

执行结果

.........................................................................

小程序调用腾讯云AI绘画的效果:

ℹ️ 离线版不含附件原文件,需要请到在线网站对应课时下载。