第1节 · AI 编程规范(Vibe Coding)
第二天笔记
第二天
1、今日目标
2、昨日回顾
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模型能力,过滤筛选高潜客户,进行外呼,实现高转化,降低运营成本。
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。先停下让我确认,不要开始写代码。
没有服务器
项目路径: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, 进入浏览器页面后, 点击右上角开发者后台
- 事件配置(默认即可,已经添加了很多事件, 如有需要可以自行添加其他事件)
- 回调配置(默认即可, 已开启长连接接收回调)
- 权限配置(默认即可)
- 版本发布与审核(默认即可, 不用管)
- 验证: 创建完飞书机器人以后, 你就可以在飞书客户端的-->开发者小助手-->打开应用, 找到你创建的机器人了
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
3.3.3.3 openclaw配置
openclaw onboard 进行交互式配置
第一个配置了模型
第二个配置了飞书
第三个配置搜索
3.3.3.3.1 配置大模型
#通过以下命令,进行一个交互设置:
openclaw onboard --classic
- 选quickstart
- 选一个大模型
- 输入完api key 完成配置
3.3.3.3.2 配置飞书
继续, 配置飞书
- 输入飞书的app id, app secret, 然后选open
##### 3.3.3.3 配置联网搜索
作用是给你的 AI 助手(Agent)开通“联网搜索”功能。
配置了这个,你的机器人就能上网查实时信息;不配置,它就只能靠之前的训练数据来回答问题。
- 可以从中选择一个, duckduckGo是免费的
3.3.3.3.4 skills安装
openclaw自动安装一些技能, 可以ctrl + C跳过, 不影响这次飞书的测试. 如果打算安装, 则耐心等待15分钟左右
工具权限的相关命令
# ==================== 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:调整图形 =》布局=》
指定版本安装
采用以下方式安装指定版本,任何系统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、今日目标
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
2.6 claude安装
2.6.1 下载安装git
https://git-scm.com/install/windows
2.6.2 下载安装nodejs
下载安装nodejs https://nodejs.org/en/download/
检查是否安装成功
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
2.6.2 本地公钥复制到github的SSH key中
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.comWhat is your preferred protocol for Git operations?→ HTTPS(或 SSH,你已配好 SSH 可选 SSH)Authenticate Git with your GitHub credentials?→ YesHow 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
# 生成密钥对,-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
- 找到secret and variables-->Actions
github仓库secret 配置
在指定仓库的---->settings--->Actions secrets and variables 中, 配置这三个变量
SSH_PRIVATE_KEY(服务器生成的私钥)
SSH_HOST(服务器的公网IP)
SSH_USER(root)
2.7 standards开发规范
开发规范:提交规范、测试规范、代码规范、进度管理规范等=》给AI看
可读性优先:代码写给人看,其次才是机器执行;拒绝炫技、晦涩写法。
单一职责:函数 / 类 / 组件只做一件事。
最小重复:公共逻辑抽离,禁止复制粘贴代码。
防御性编程:参数校验、空值处理、异常捕获。
可测试:业务逻辑与 UI/IO 解耦,便于单元测试。
兼容性 & 性能前置思考。
约定大于配置:统一项目约定,减少个性化写法。
2.7.1 helloworld跑通验证cicd
第一步:standards存放地址,位于项目目录下
第二步: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按钮即可合并
第三:查看结果
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节 · 计算机视觉:图像分类
01-计算机视觉基础知识
计算机视觉基础CV笔记
一、今日课程内容
- 计算机视觉的定义(了解)
- 计算机视觉常见任务(了解)
- 计算机视觉的应用场景(了解)
二、计算机视觉介绍(了解)
1、定义
- 用【摄像机和电脑】及其他相关设备,对【生物视觉】的一种模拟,任务:理解图片和视频
2、主要任务
- 让【计算机】理解【图片或者视频】中的内容,就像人类和许多其他生物每天所做的那样。
3、任务目标
- 让计算机【理解】图片中的场景(办公室,客厅,咖啡厅等)
- 让计算机【识别】场景中包含的物体(宠物,交通工具,人等)
- 让计算机【定位】物体在图像中的位置(物体的大小,边界等)
- 让计算机理解物体之间的【关系或行为】(是在对话,比赛或吵架等),以及图像表达的意义(喜庆的,悲伤的等)
三、计算机视觉常见任务(了解)
1、三大经典任务
-
图像分类(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年
- 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、网络结构
2、网络特点
(1)网络层数:【8】
(2)卷积层和池化层参数:【卷积核:11,5,3 池化:3X3 ,步长2】
(3)激活函数:【ReLU】
(4)模型复杂度调节:【Dropout】
(5)缓解过拟合操作:【引入了大量的图像增强,如翻转、裁剪和颜色变化,扩大数据集】
3、网络层参数计算
(1)计算公式
(2)AlexNet网络层计算
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)运行结果
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 块
- 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节 · 计算机视觉:目标检测
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、算法流程
- 候选区域生成:使用选择性搜索(Selective Search)的方法找出图片中可能存在目标的侯选区域
- 候选区域搜索依据:使用语义分割的方法,将颜色、边界、纹理等信息作为合并条件,采用多尺度的综合方法 - CNN网络提取特征:选取预训练卷积神经网网络(AlexNet或VGG)用于进行特征提取。
- 采用预训练模型(AlexNet或VGG)在生成的候选区域上进行特征提取
- 特征提取流程:- (1)候选区域需进行裁剪或变形为固定的尺寸,再送入CNN网络中进行特征提取(全连接层输入数据的尺寸是固定的)
- (2)全连接层由1000修改为N+1(预训练模型在ImageNet数据集上获得)
- (3)利用微调后的CNN网络,提取每一个候选区域的特征,获取一个4096维的特征,一幅图像就是2000x4096维特征存储到磁盘中。
- 目标分类:训练支持向量机(SVM)来辨别目标物体和背景,对每个类别,都要训练一个二元SVM。
- 对于N个类别的检测任务,需要训练N(目标类别数目)个SVM分类器
- 对候选区域的特征向量(4096维)进行二分类,判断其是某一类别的目标,还是背景来完成目标分类。 - 目标定位:训练一个线性回归模型,为每个辨识到的物体生成更精确的边界框。
- 训练一个线性回归模型在给定的候选区域的结果上去预测一个新的检测窗口,能够获得更精确的位置 - 模型预测
- 使用选择性搜索的方法从一张图片中提取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个类别中的每一个检测框回归偏移
4、模型预测
(1)输入图像
(2)图像被送入到卷积网络进行特征提取,将通过选择性搜索获取的候选区域映射到特征图中
(3)在特征图上Rol中应用RoIPooling,获取尺寸相同的特征向量
(4)将这些区域传递到全连接的网络中进行分类和回归,得到目标检测的结果。
5、模型总结
Fast R-CNN是对R-CNN模型的一种改进:
- CNN网络不再对每个候选区域进行特征提取,而是直接对整张图像进行出路,这样减少了很多重复计算。
- 用ROI pooling进行特征的尺寸变换,来满足FC全连接层对输入数据尺度的要求。
- 将目标的回归和分类统一在一个网络中,使用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网络的损失函数
-
(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
课堂笔记
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))
```
#### 课堂笔记
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、损失函数
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,更适用于多标签分类任务。
-
相比之前的算法,尤其是针对小目标,精度有显著提升。
- 网络结构
- 整个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相似
-
-
基本组件:
- 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
- 输出(output)
四、模型训练(了解)
1、基本内容
- 正负样本的设置方式,损失函数的构成及训练过程与YOLOv3一致
- 唯一不同是在计算回归值得损失时,V4中采用了IOU的损失
2、IOU损失
- GIOU损失
- DIOU损失
- CIOU损失
- .........
五、模型预测
- 将图片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、模型预测
目标检测代码(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
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.pyruns YOLOv5 Classification inference on a variety of sources, downloading models automatically from the latest YOLOv5 release, and saving results toruns/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)
[34m[1mclassify/predict: [0mweights=['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 [1mruns/predict-cls/exp[0m

2. Validate
Validate a model's accuracy on the Imagenet dataset's
valortestsplits. Models are downloaded automatically from the latest YOLOv5 release. To show results by class use the--verboseflag.# 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
[34m[1mclassify/val: [0mdata=../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.943. Train
Close the active learning loop by sampling images from your inference conditions with the
roboflowpip 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/exp3etc.
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
roboflowpip package.- Custom Training Example: https://blog.roboflow.com/train-yolov5-classification-custom-data/
- Custom Training Notebook:
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
[34m[1mclassify/train: [0mmodel=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 [34m[1mgithub: [0mup 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) [34m[1mTensorBoard: [0mStart 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 [1m/content/datasets/imagenette160[0m [34m[1malbumentations: [0mRandomResizedCrop(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 [34m[1moptimizer:[0m 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 [1mruns/train-cls/exp[0m 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 [1mruns/train-cls/exp[0m 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.app4. 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. trainTo 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:
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):
pip install clearml- run
clearml-initto connect to a ClearML server (deploy your own open-source server, or use our free hosted server)
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!
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 asruns/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.

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):
- Notebooks with free GPU:
- Google Cloud Deep Learning VM. See GCP Quickstart Guide
- Amazon Deep Learning AMI. See AWS Quickstart Guide
- Docker Image. See Docker Quickstart Guide
Status
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.
- Pretrained Models are downloaded
📄 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
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.pyruns YOLOv5 instance segmentation inference on a variety of sources, downloading models automatically from the latest YOLOv5 release, and saving results toruns/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)
[34m[1msegment/predict: [0mweights=['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 [1mruns/predict-seg/exp[0m

2. Validate
Validate a model's accuracy on the COCO dataset's
valortestsplits. Models are downloaded automatically from the latest YOLOv5 release. To show results by class use the--verboseflag.# 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
[34m[1msegment/val: [0mdata=/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 [34m[1mval: [0mScanning /content/datasets/coco/val2017... 4952 images, 48 backgrounds, 0 corrupt: 100% 5000/5000 [00:03<00:00, 1361.31it/s] [34m[1mval: [0mNew 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 [1mruns/val-seg/exp[0m3. Train
Close the active learning loop by sampling images from your inference conditions with the
roboflowpip 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.- Pretrained Models are downloaded
automatically from the latest YOLOv5 release - Datasets available for autodownload include: COCO, COCO128, VOC, Argoverse, VisDrone, GlobalWheat, xView, Objects365, SKU-110K.
- Training Results are saved to
runs/train-seg/with incrementing run directories, i.e.runs/train-seg/exp2,runs/train-seg/exp3etc.
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
roboflowpip package.- Custom Training Example: https://blog.roboflow.com/train-yolov5-instance-segmentation-custom-dataset/
- Custom Training Notebook:
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
[34m[1msegment/train: [0mweights=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 [34m[1mgithub: [0mup 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) [34m[1mhyperparameters: [0mlr0=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 [34m[1mTensorBoard: [0mStart 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 [1m/content/datasets[0m 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 [34m[1mAMP: [0mchecks passed ✅ [34m[1moptimizer:[0m SGD(lr=0.01) with parameter groups 60 weight(decay=0.0), 63 weight(decay=0.0005), 63 bias [34m[1malbumentations: [0mBlur(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)) [34m[1mtrain: [0mScanning /content/datasets/coco128-seg/labels/train2017... 126 images, 2 backgrounds, 0 corrupt: 100% 128/128 [00:00<00:00, 1389.59it/s] [34m[1mtrain: [0mNew cache created: /content/datasets/coco128-seg/labels/train2017.cache [34m[1mtrain: [0mCaching images (0.1GB ram): 100% 128/128 [00:00<00:00, 238.86it/s] [34m[1mval: [0mScanning /content/datasets/coco128-seg/labels/train2017.cache... 126 images, 2 backgrounds, 0 corrupt: 100% 128/128 [00:00<?, ?it/s] [34m[1mval: [0mCaching images (0.1GB ram): 100% 128/128 [00:01<00:00, 98.90it/s] [34m[1mAutoAnchor: [0m4.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 [1mruns/train-seg/exp[0m 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 [1mruns/train-seg/exp[0m4. 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. trainTo 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:
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):
pip install clearml- run
clearml-initto connect to a ClearML server (deploy your own open-source server, or use our free hosted server)
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!
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 asruns/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.

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):
- Notebooks with free GPU:
- Google Cloud Deep Learning VM. See GCP Quickstart Guide
- Amazon Deep Learning AMI. See AWS Quickstart Guide
- Docker Image. See Docker Quickstart Guide
Status
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.
- Pretrained Models are downloaded
📄 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: [](https://badge.fury.io/py/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 <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) ``` |  |  | | ---------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------- | ### 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 <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)了解详细信息,并开始使用: [](https://badge.fury.io/py/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>分割模型使用示例 <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: 推理暂未支持) ``` |  |  | | ---------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------- | ### 模型导出 将 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>分类训练示例 <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/)
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.pyruns YOLOv5 inference on a variety of sources, downloading models automatically from the latest YOLOv5 release, and saving results toruns/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)
[34m[1mdetect: [0mweights=['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 [1mruns/detect/exp[0m

2. Validate
Validate a model's accuracy on the COCO dataset's
valortestsplits. Models are downloaded automatically from the latest YOLOv5 release. To show results by class use the--verboseflag.# 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
[34m[1mval: [0mdata=/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 [34m[1mval: [0mScanning /content/datasets/coco/val2017... 4952 images, 48 backgrounds, 0 corrupt: 100% 5000/5000 [00:02<00:00, 2024.59it/s] [34m[1mval: [0mNew 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 [1mruns/val/exp[0m3. Train
Close the active learning loop by sampling images from your inference conditions with the
roboflowpip 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.- Pretrained Models are downloaded
automatically from the latest YOLOv5 release - Datasets available for autodownload include: COCO, COCO128, VOC, Argoverse, VisDrone, GlobalWheat, xView, Objects365, SKU-110K.
- Training Results are saved to
runs/train/with incrementing run directories, i.e.runs/train/exp2,runs/train/exp3etc.
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
roboflowpip 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 [34m[1mtrain: [0mweights=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 [34m[1mgithub: [0mup 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) [34m[1mhyperparameters: [0mlr0=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 [34m[1mClearML: [0mrun 'pip install clearml' to automatically track, visualize and remotely train YOLOv5 🚀 in ClearML [34m[1mComet: [0mrun 'pip install comet_ml' to automatically track and visualize YOLOv5 🚀 runs in Comet [34m[1mTensorBoard: [0mStart 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 [1m/content/datasets[0m 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 [34m[1mAMP: [0mchecks passed ✅ [34m[1moptimizer:[0m SGD(lr=0.01) with parameter groups 57 weight(decay=0.0), 60 weight(decay=0.0005), 60 bias [34m[1malbumentations: [0mBlur(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)) [34m[1mtrain: [0mScanning /content/datasets/coco128/labels/train2017... 126 images, 2 backgrounds, 0 corrupt: 100% 128/128 [00:00<00:00, 1709.36it/s] [34m[1mtrain: [0mNew cache created: /content/datasets/coco128/labels/train2017.cache [34m[1mtrain: [0mCaching images (0.1GB ram): 100% 128/128 [00:00<00:00, 264.35it/s] [34m[1mval: [0mScanning /content/datasets/coco128/labels/train2017.cache... 126 images, 2 backgrounds, 0 corrupt: 100% 128/128 [00:00<?, ?it/s] [34m[1mval: [0mCaching images (0.1GB ram): 100% 128/128 [00:01<00:00, 107.05it/s] [34m[1mAutoAnchor: [0m4.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 [1mruns/train/exp[0m 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 [1mruns/train/exp[0m4. 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. trainTo 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:
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):
pip install clearml- run
clearml-initto connect to a ClearML server (deploy your own open-source server, or use our free hosted server)
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!
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 asruns/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.

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):
- Notebooks with free GPU:
- Google Cloud Deep Learning VM. See GCP Quickstart Guide
- Amazon Deep Learning AMI. See AWS Quickstart Guide
- Docker Image. See Docker Quickstart Guide
Status
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.
- Pretrained Models are downloaded
📄 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
多模态课件(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绘画的效果:
ℹ️ 离线版不含附件原文件,需要请到在线网站对应课时下载。




