Initial commit

This commit is contained in:
zhengsl 2026-07-27 17:51:49 +08:00
commit f9dff0891f
750 changed files with 103063 additions and 0 deletions

24
.env.example Normal file
View File

@ -0,0 +1,24 @@
DJANGO_DEBUG=true
DJANGO_SECRET_KEY=change-me
DJANGO_ALLOWED_HOSTS=127.0.0.1,localhost,192.168.1.174
DJANGO_CORS_ALLOWED_ORIGINS=http://localhost:3000,ws://localhost:3000
DB_ENGINE=django.db.backends.sqlite3
DB_NAME=
DB_USER=
DB_PASSWORD=
DB_HOST=127.0.0.1
DB_PORT=5432
REDIS_URL=redis://127.0.0.1:6379/1
SM4_PRIVATE_KEY=
FILESPACE_ROOT_PATH=D:\AI-TrainPrediction\AI-TrainPrediction\project_space
TRAIN_LOG_DIR=D:\AI-TrainPrediction\AI-TrainPrediction\project_space
MEDIA_URL=/image/
MEDIA_ROOT=D:\AI-TrainPrediction\AI-TrainPrediction\project_space
IMAGE_BASE_URL=http://localhost:8000/image/
TASK_CALL_BACK_URL=http://localhost:8000/server/
JWT_ACCESS_MINUTES=60
JWT_REFRESH_DAYS=1

View File

@ -0,0 +1,709 @@
# MinIO 文件存储接口文档
## 一、概述
本项目已集成 MinIO 对象存储,并在 Django 后端提供统一的文件存储 API。
### MinIO 服务信息
| 项目 | 值 |
|------|-----|
| API 地址 | `http://127.0.0.1:9000` |
| Web 控制台 | `http://127.0.0.1:9001` |
| 用户名 | `admin` |
| 密码 | `zhengsl2026` |
| 默认 Bucket | `ai-trainprediction` |
### Django 配置
```python
MINIO_ENDPOINT = os.environ.get("MINIO_ENDPOINT", "127.0.0.1:9000")
MINIO_ACCESS_KEY = os.environ.get("MINIO_ACCESS_KEY", "admin")
MINIO_SECRET_KEY = os.environ.get("MINIO_SECRET_KEY", "zhengsl2026")
MINIO_SECURE = _env_bool("MINIO_SECURE", default=False)
MINIO_DEFAULT_BUCKET = os.environ.get("MINIO_DEFAULT_BUCKET", "ai-trainprediction")
MINIO_PRESIGN_EXPIRES = int(os.environ.get("MINIO_PRESIGN_EXPIRES", "3600"))
```
### 认证要求
所有 MinIO 接口均要求 JWT 认证:
```http
Authorization: Bearer {access_token}
```
### 接口前缀
```text
/server/minio/
```
---
## 二、接口总览
| 方法 | 路径 | 功能 |
|------|------|------|
| POST | `/server/minio/bucket/create/` | 创建 Bucket |
| GET | `/server/minio/bucket/list/` | 列出 Bucket |
| DELETE | `/server/minio/bucket/delete/` | 删除 Bucket |
| POST | `/server/minio/file/upload/` | 单文件上传 |
| POST | `/server/minio/files/upload/` | 批量文件上传 |
| POST | `/server/minio/file/upload-zip/` | ZIP 解压后批量入库 |
| POST | `/server/minio/file/upload-video-frames/` | 视频抽帧后图片入库 |
| POST | `/server/minio/file/upload-base64/` | Base64 文件上传 |
| GET | `/server/minio/file/download/` | 流式下载文件 |
| GET | `/server/minio/file/presign/` | 获取预签名下载 URL |
| GET | `/server/minio/file/presign-upload/` | 获取预签名上传 URL |
| GET | `/server/minio/file/preview/` | 流式预览文件 |
| GET | `/server/minio/file/info/` | 获取文件元信息 |
| GET | `/server/minio/files/` | 列出对象 |
| DELETE | `/server/minio/file/delete/` | 删除单个文件 |
| DELETE | `/server/minio/files/delete/` | 批量删除文件 |
| POST | `/server/minio/file/copy/` | 复制文件 |
---
## 三、Bucket 管理
### 3.1 创建 Bucket
```http
POST /server/minio/bucket/create/
Content-Type: application/json
```
请求参数:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `bucket_name` | string | 否 | Bucket 名称,默认 `ai-trainprediction` |
成功响应:
```json
{
"code": 200,
"data": {
"bucket_name": "my-bucket",
"created": true
},
"msg": "Bucket 创建成功"
}
```
说明:
- 若 Bucket 已存在,返回 `created: false`
- MinIO 返回 `BucketAlreadyExists` / `BucketAlreadyOwnedByYou` 时会转换为 `409`
### 3.2 列出 Bucket
```http
GET /server/minio/bucket/list/
```
成功响应:
```json
{
"code": 200,
"data": ["ai-trainprediction", "my-bucket"],
"msg": "获取成功"
}
```
### 3.3 删除 Bucket
```http
DELETE /server/minio/bucket/delete/
Content-Type: application/json
```
请求参数:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `bucket_name` | string | 是 | Bucket 名称 |
成功响应:
```json
{
"code": 200,
"data": {
"bucket_name": "my-bucket"
},
"msg": "Bucket 已删除"
}
```
错误说明:
- `404`Bucket 不存在
- `409`Bucket 非空,无法删除
---
## 四、文件上传
### 4.1 单文件上传
```http
POST /server/minio/file/upload/
Content-Type: multipart/form-data
```
请求参数:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `file` | File | 是 | 上传文件 |
| `bucket_name` | string | 否 | 目标 Bucket默认 `ai-trainprediction` |
| `prefix` | string | 否 | 对象前缀,如 `images/datasets` |
| `object_name` | string | 否 | 指定对象名;不传则自动生成 |
自动命名规则:
```text
{prefix}/YYYY-MM-DD/{uuid12}_{original_filename}
```
成功响应:
```json
{
"code": 200,
"data": {
"bucket_name": "ai-trainprediction",
"object_name": "images/2026-07-01/a1b2c3d4e5f6_photo.jpg",
"original_name": "photo.jpg",
"size": 102400,
"content_type": "image/jpeg",
"presigned_url": "http://127.0.0.1:9000/..."
},
"msg": "上传成功"
}
```
说明:
- 该接口已改为流式上传,不会把整文件一次性读入内存
- 若 `object_name` 明确指定且对象已存在,则会被覆盖
### 4.2 批量文件上传
```http
POST /server/minio/files/upload/
Content-Type: multipart/form-data
```
请求参数:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `files` | File[] | 是 | 多个文件 |
| `bucket_name` | string | 否 | 目标 Bucket |
| `prefix` | string | 否 | 对象前缀 |
成功响应:
```json
{
"code": 200,
"data": [
{
"object_name": "images/2026-07-01/a1b2_photo1.jpg",
"original_name": "photo1.jpg",
"size": 102400,
"content_type": "image/jpeg",
"presigned_url": "http://127.0.0.1:9000/..."
}
],
"msg": "全部上传成功,共 1 个文件"
}
```
### 4.3 Base64 文件上传
```http
POST /server/minio/file/upload-base64/
Content-Type: application/json
```
请求体:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `file_base64` | string | 是 | Base64 编码内容 |
| `filename` | string | 是 | 原始文件名 |
| `bucket_name` | string | 否 | 目标 Bucket |
| `prefix` | string | 否 | 对象前缀 |
请求示例:
```json
{
"file_base64": "/9j/4AAQSkZJRgABAQ...",
"filename": "avatar.png",
"bucket_name": "ai-trainprediction",
"prefix": "avatars"
}
```
成功响应:
```json
{
"code": 200,
"data": {
"bucket_name": "ai-trainprediction",
"object_name": "avatars/2026-07-01/a1b2c3d4e5f6_avatar.png",
"original_name": "avatar.png",
"size": 51200,
"content_type": "image/png",
"presigned_url": "http://127.0.0.1:9000/..."
},
"msg": "上传成功"
}
```
说明:
- 此接口当前只支持 `application/json`
- Base64 解码失败时返回 `400`
### 4.4 ZIP 解压批量入库
```http
POST /server/minio/file/upload-zip/
Content-Type: multipart/form-data
```
请求参数:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `file` | File | 是 | ZIP 压缩包 |
| `bucket_name` | string | 否 | 目标 Bucket |
| `prefix` | string | 否 | MinIO 对象前缀 |
| `images_only` | bool | 否 | 是否仅入库图片文件,默认 `false` |
说明:
- 接口会自动解压 ZIP 内文件并逐个上传到 MinIO
- 默认保留 ZIP 内相对目录结构
- 会自动过滤目录项、隐藏文件和非法路径片段(如 `..`
- `images_only=true` 时,仅上传图片扩展名文件:`jpg/jpeg/png/bmp/gif/webp/tif/tiff`
成功响应:
```json
{
"code": 200,
"data": {
"bucket_name": "ai-trainprediction",
"zip_name": "dataset.zip",
"uploaded_count": 2,
"uploaded": [
{
"source_name": "images/cat.jpg",
"object_name": "datasets/images/cat.jpg",
"size": 102400,
"content_type": "image/jpeg"
}
],
"skipped": ["docs/readme.txt"]
},
"msg": "ZIP 解压入库成功,共上传 2 个文件"
}
```
### 4.5 视频抽帧图片入库
```http
POST /server/minio/file/upload-video-frames/
Content-Type: multipart/form-data
```
请求参数:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `file` | File | 是 | 视频文件 |
| `bucket_name` | string | 否 | 目标 Bucket |
| `prefix` | string | 否 | MinIO 对象前缀 |
| `interval_seconds` | int | 否 | 抽帧时间间隔,单位秒,默认 `1` |
| `start_second` | int | 否 | 开始抽帧时间,默认 `0` |
| `end_second` | int | 否 | 结束抽帧时间,不传则到视频末尾 |
| `max_frames` | int | 否 | 最多生成图片数,默认 `100` |
| `image_format` | string | 否 | 输出图片格式,支持 `jpg/jpeg/png`,默认 `jpg` |
说明:
- 当前支持的视频扩展名:`mp4/avi/mov/mkv/wmv/flv/mpeg/mpg/webm`
- 需要服务端安装 `opencv-python`
- 抽出的图片会存入 `{prefix}/{视频名}_frames/` 目录
- 文件名规则:`frame_00001_1200ms.jpg`
成功响应:
```json
{
"code": 200,
"data": {
"bucket_name": "ai-trainprediction",
"video_name": "demo.mp4",
"fps": 25.0,
"duration_seconds": 12.4,
"interval_seconds": 1,
"start_second": 0,
"end_second": 10,
"max_frames": 20,
"extracted_count": 10,
"extracted": [
{
"object_name": "videos/demo_frames/frame_00001_0ms.jpg",
"frame_index": 0,
"time_second": 0.0
}
]
},
"msg": "视频抽帧入库成功,共生成 10 张图片"
}
```
### 4.6 获取预签名上传 URL
```http
GET /server/minio/file/presign-upload/?bucket_name={bucket}&object_name={object}&expires={seconds}
```
请求参数:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `bucket_name` | string | 否 | 目标 Bucket |
| `object_name` | string | 是 | 目标对象名 |
| `expires` | int | 否 | 有效期秒数,默认 `3600`,最大 `604800` |
成功响应:
```json
{
"code": 200,
"data": {
"presigned_url": "http://127.0.0.1:9000/ai-trainprediction/path/file.jpg?X-Amz-Algorithm=AWS4-HMAC-SHA256&...",
"expires_in": 3600
},
"msg": "获取成功"
}
```
---
## 五、文件下载与预览
### 5.1 流式下载文件
```http
GET /server/minio/file/download/?bucket_name={bucket}&object_name={object}
```
请求参数:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `bucket_name` | string | 否 | Bucket 名称 |
| `object_name` | string | 是 | 对象名 |
说明:
- 返回流式响应,适合大文件下载
- `Content-Disposition``attachment`
### 5.2 获取预签名下载 URL
```http
GET /server/minio/file/presign/?bucket_name={bucket}&object_name={object}&expires={seconds}
```
请求参数:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `bucket_name` | string | 否 | Bucket 名称 |
| `object_name` | string | 是 | 对象名 |
| `expires` | int | 否 | 有效期秒数,默认 `3600`,最大 `604800` |
成功响应:
```json
{
"code": 200,
"data": {
"presigned_url": "http://127.0.0.1:9000/ai-trainprediction/path/file.jpg?X-Amz-Algorithm=AWS4-HMAC-SHA256&...",
"expires_in": 3600
},
"msg": "获取成功"
}
```
说明:
- `expires <= 0` 时返回 `400`
- 超过 7 天会被自动截断为 `604800`
### 5.3 流式预览文件
```http
GET /server/minio/file/preview/?bucket_name={bucket}&object_name={object}
```
说明:
- 返回流式响应
- `Content-Disposition``inline`
- 适合图片、PDF、文本等浏览器可直接预览的文件
---
## 六、文件管理
### 6.1 获取文件元信息
```http
GET /server/minio/file/info/?bucket_name={bucket}&object_name={object}
```
成功响应:
```json
{
"code": 200,
"data": {
"name": "images/2026-07-01/a1b2_photo.jpg",
"size": 102400,
"etag": "d41d8cd98f00b204e9800998ecf8427e",
"content_type": "image/jpeg",
"last_modified": "2026-07-01T12:00:00+00:00",
"presigned_url": "http://127.0.0.1:9000/..."
},
"msg": "获取成功"
}
```
### 6.2 列出对象
```http
GET /server/minio/files/?bucket_name={bucket}&prefix={prefix}
```
请求参数:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `bucket_name` | string | 否 | Bucket 名称 |
| `prefix` | string | 否 | 仅列出指定前缀下的对象 |
成功响应:
```json
{
"code": 200,
"data": [
{
"name": "images/2026-07-01/a1b2_photo.jpg",
"size": 102400,
"last_modified": "2026-07-01T12:00:00+00:00",
"etag": "d41d8cd98f00b204e9800998ecf8427e",
"content_type": "image/jpeg",
"is_dir": false
}
],
"msg": "共 1 个对象"
}
```
### 6.3 删除单个文件
```http
DELETE /server/minio/file/delete/
Content-Type: application/json
```
请求体:
```json
{
"bucket_name": "ai-trainprediction",
"object_name": "images/2026-07-01/a1b2_photo.jpg"
}
```
成功响应:
```json
{
"code": 200,
"data": {
"object_name": "images/2026-07-01/a1b2_photo.jpg"
},
"msg": "删除成功"
}
```
### 6.4 批量删除文件
```http
DELETE /server/minio/files/delete/
Content-Type: application/json
```
请求体:
```json
{
"bucket_name": "ai-trainprediction",
"object_names": [
"images/photo1.jpg",
"images/photo2.png"
]
}
```
成功响应:
```json
{
"code": 200,
"data": {
"deleted": 2
},
"msg": "全部删除成功"
}
```
部分成功响应:
```json
{
"code": 207,
"data": {
"deleted": 1,
"errors": ["images/photo2.png"]
},
"msg": "部分删除1 个失败"
}
```
### 6.5 复制文件
```http
POST /server/minio/file/copy/
Content-Type: application/json
```
请求体:
```json
{
"src_bucket": "ai-trainprediction",
"src_object": "images/photo.jpg",
"dst_bucket": "ai-trainprediction",
"dst_object": "backup/photo.jpg"
}
```
请求参数:
| 参数 | 类型 | 必填 | 说明 |
|------|------|------|------|
| `src_bucket` | string | 否 | 源 Bucket默认 `ai-trainprediction` |
| `src_object` | string | 是 | 源对象名 |
| `dst_bucket` | string | 否 | 目标 Bucket默认 `ai-trainprediction` |
| `dst_object` | string | 是 | 目标对象名 |
成功响应:
```json
{
"code": 200,
"data": {
"object_name": "backup/photo.jpg"
},
"msg": "复制成功"
}
```
说明:
- 若目标 Bucket 不存在,接口会自动创建
---
## 七、错误码说明
| HTTP 状态码 | 含义 |
|-------------|------|
| `200` | 成功 |
| `207` | 批量删除部分成功 |
| `400` | 参数错误、`expires` 不合法、Base64 解码失败等 |
| `401` | 未携带或携带了无效 JWT Token |
| `403` | MinIO 拒绝访问 |
| `404` | Bucket 或对象不存在 |
| `409` | Bucket 已存在、Bucket 非空等冲突 |
| `500` | 其他 MinIO / 服务端错误 |
错误响应格式:
```json
{
"code": 409,
"msg": "Bucket 非空,无法删除",
"detail": "S3 operation failed; code: BucketNotEmpty, ..."
}
```
---
## 八、代码结构
```text
backend/
├── config/settings/base.py
├── config/api_urls.py
└── apps/common/
├── minio_client.py
└── views/minio_storage.py
```
### `minio_client.py`
封装的核心方法包括:
| 方法 | 功能 |
|------|------|
| `bucket_exists()` | 检查 Bucket 是否存在 |
| `create_bucket()` | 创建 Bucket |
| `list_buckets()` | 列出 Bucket |
| `remove_bucket()` | 删除空 Bucket |
| `upload_file()` | 上传本地文件 |
| `upload_bytes()` | 上传字节数据 |
| `upload_stream()` | 流式上传文件对象 |
| `download_file()` | 下载到本地 |
| `download_bytes()` | 下载为 bytes |
| `get_object()` | 获取对象流响应 |
| `get_object_info()` | 获取对象元信息 |
| `object_exists()` | 判断对象是否存在 |
| `list_objects()` | 列出对象 |
| `delete_object()` | 删除单个对象 |
| `delete_objects()` | 批量删除对象 |
| `presign_get()` | 获取预签名下载 URL |
| `presign_upload()` | 获取预签名上传 URL |
| `generate_object_name()` | 生成唯一对象名 |
| `copy_object()` | 复制对象 |
---
## 九、注意事项
1. 所有接口均要求 JWT 认证
2. 下载和预览接口已经改为流式响应,适合较大文件
3. `upload-base64` 当前只支持 `application/json`
4. ZIP 入库默认保留压缩包中的相对目录结构,并会过滤非法路径
5. 视频抽帧依赖 `opencv-python`
6. `expires` 必须大于 0最大有效期为 7 天
7. 指定 `object_name` 时可能覆盖已有对象,请谨慎使用
8. 删除 Bucket 前必须先清空其中所有对象

136
SD_AIBOX_B_H.sql Normal file
View File

@ -0,0 +1,136 @@
/*
Navicat Premium Data Transfer
Source Server : 190
Source Server Type : Oracle
Source Server Version : 110200
Source Host : 172.16.31.190:1521
Source Schema : QGC_REFA
Target Server Type : Oracle
Target Server Version : 110200
File Encoding : 65001
Date: 08/07/2026 18:30:27
*/
-- ----------------------------
-- Table structure for SD_AIBOX_B_H
-- ----------------------------
DROP TABLE "QGC_REFA"."SD_AIBOX_B_H";
CREATE TABLE "QGC_REFA"."SD_AIBOX_B_H" (
"STCD" VARCHAR2(36 BYTE) NOT NULL ,
"STNM" VARCHAR2(255 BYTE) NOT NULL ,
"STTP" VARCHAR2(64 BYTE) NOT NULL ,
"TM" DATE ,
"LGTD" NUMBER(10,6) ,
"LTTD" NUMBER(10,6) ,
"ELEV" NUMBER(10,6) ,
"STLC" VARCHAR2(200 BYTE) ,
"JCDT" DATE ,
"WDDT" DATE ,
"BLDSTT_CODE" NUMBER(1) ,
"USFL" NUMBER(1) ,
"DTIN" NUMBER(1) ,
"DTIN_TM" DATE ,
"MWAY" NUMBER(1) ,
"STINDX" VARCHAR2(1000 BYTE) ,
"ORDER_INDEX" NUMBER(8) ,
"RSTCD" VARCHAR2(64 BYTE) ,
"BASE_ID" VARCHAR2(64 BYTE) ,
"HBRVCD" VARCHAR2(64 BYTE) ,
"RVCD" VARCHAR2(64 BYTE) ,
"INTRODUCE" VARCHAR2(3000 BYTE) ,
"LOGO" VARCHAR2(255 BYTE) ,
"INFFILE" VARCHAR2(1000 BYTE) ,
"FWDX" VARCHAR2(1000 BYTE) ,
"GDFS" VARCHAR2(512 BYTE) ,
"AIPZ" VARCHAR2(1000 BYTE) ,
"MXBB" VARCHAR2(512 BYTE) ,
"IPADDR" VARCHAR2(512 BYTE) ,
"SIMINFO" VARCHAR2(512 BYTE) ,
"LXR" VARCHAR2(512 BYTE) ,
"PURPOSE" VARCHAR2(512 BYTE) ,
"DTFRQCY" NUMBER(7) ,
"REMARK" VARCHAR2(512 BYTE) ,
"VLSR" VARCHAR2(512 BYTE) ,
"VLSR_TM" DATE ,
"RECORD_USER" VARCHAR2(36 BYTE) ,
"RECORD_TIME" DATE DEFAULT SYSDATE NOT NULL ,
"MODIFY_USER" VARCHAR2(36 BYTE) ,
"MODIFY_TIME" DATE ,
"IS_DELETED" NUMBER(1) DEFAULT 0 NOT NULL ,
"DELETE_USER" VARCHAR2(36 BYTE) ,
"DELETE_TIME" DATE
)
TABLESPACE "QGC_REFA_DATA"
LOGGING
NOCOMPRESS
PCTFREE 10
INITRANS 1
STORAGE (
BUFFER_POOL DEFAULT
)
PARALLEL 1
NOCACHE
DISABLE ROW MOVEMENT
;
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."STCD" IS '站码';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."STNM" IS '站名';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."STTP" IS '站类';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."TM" IS '数据时间';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."LGTD" IS '经度(单位:°)';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."LTTD" IS '纬度(单位:°)';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."ELEV" IS '高程单位m';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."STLC" IS '站址/位置';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."JCDT" IS '建成日期';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."WDDT" IS '退役/拆除日期';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."BLDSTT_CODE" IS '建设状态分类';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."USFL" IS '是否启用';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."DTIN" IS '数据是否接入';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."DTIN_TM" IS '数据接入时间';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."MWAY" IS '监测方式';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."STINDX" IS '监测指标';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."ORDER_INDEX" IS '排序';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."RSTCD" IS '所属电站编码';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."BASE_ID" IS '所属水电基地编码';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."HBRVCD" IS '所属水电基地流域编码';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."RVCD" IS '所属流域编码';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."INTRODUCE" IS '简介';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."LOGO" IS 'LOGO图片';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."INFFILE" IS '介绍弹窗图片';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."FWDX" IS '服务对象';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."GDFS" IS '供电方式';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."AIPZ" IS 'AI盒子配置';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."MXBB" IS '模型版本';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."IPADDR" IS 'IP地址';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."SIMINFO" IS 'SIM卡信息';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."LXR" IS '联系人及联系电话';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."PURPOSE" IS '用途';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."DTFRQCY" IS '数据监测频次单位min';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."REMARK" IS '备注';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."VLSR" IS '数据来源';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."VLSR_TM" IS '数据来源时间';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."RECORD_USER" IS '创建人关联SYS_USER.ID';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."RECORD_TIME" IS '创建时间';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."MODIFY_USER" IS '更新人关联SYS_USER.ID';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."MODIFY_TIME" IS '更新时间';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."IS_DELETED" IS '是否已删除0=未删除 1=已删除存储字典值字典项为comm.is_deleted';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."DELETE_USER" IS '删除人关联SYS_USER.ID';
COMMENT ON COLUMN "QGC_REFA"."SD_AIBOX_B_H"."DELETE_TIME" IS '删除时间';
COMMENT ON TABLE "QGC_REFA"."SD_AIBOX_B_H" IS 'AI边缘计算盒子基本属性表';
-- ----------------------------
-- Primary Key structure for table SD_AIBOX_B_H
-- ----------------------------
ALTER TABLE "QGC_REFA"."SD_AIBOX_B_H" ADD CONSTRAINT "SD_AIBOX_B_H_PK" PRIMARY KEY ("STCD");
-- ----------------------------
-- Checks structure for table SD_AIBOX_B_H
-- ----------------------------
ALTER TABLE "QGC_REFA"."SD_AIBOX_B_H" ADD CONSTRAINT "SYS_C0057611" CHECK ("STCD" IS NOT NULL) NOT DEFERRABLE INITIALLY IMMEDIATE NORELY VALIDATE;
ALTER TABLE "QGC_REFA"."SD_AIBOX_B_H" ADD CONSTRAINT "SYS_C0057612" CHECK ("STNM" IS NOT NULL) NOT DEFERRABLE INITIALLY IMMEDIATE NORELY VALIDATE;
ALTER TABLE "QGC_REFA"."SD_AIBOX_B_H" ADD CONSTRAINT "SYS_C0057613" CHECK ("STTP" IS NOT NULL) NOT DEFERRABLE INITIALLY IMMEDIATE NORELY VALIDATE;
ALTER TABLE "QGC_REFA"."SD_AIBOX_B_H" ADD CONSTRAINT "SYS_C0057615" CHECK ("RECORD_TIME" IS NOT NULL) NOT DEFERRABLE INITIALLY IMMEDIATE NORELY VALIDATE;
ALTER TABLE "QGC_REFA"."SD_AIBOX_B_H" ADD CONSTRAINT "SYS_C0057616" CHECK ("IS_DELETED" IS NOT NULL) NOT DEFERRABLE INITIALLY IMMEDIATE NORELY VALIDATE;

1
backend/apps/__init__.py Normal file
View File

@ -0,0 +1 @@

View File

@ -0,0 +1 @@

View File

@ -0,0 +1,68 @@
# Copyright The Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
default_language_version:
python: python3
ci:
autofix_prs: true
autoupdate_commit_msg: "[pre-commit.ci] pre-commit suggestions"
autoupdate_schedule: quarterly
# submodules: true
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v5.0.0
hooks:
- id: end-of-file-fixer
- id: trailing-whitespace
# - id: check-json # skip for incompatibility with .devcontainer/devcontainer.json
- id: check-yaml
- id: check-toml
- id: check-docstring-first
- id: check-executables-have-shebangs
- id: check-case-conflict
# - id: check-added-large-files
# args: ["--maxkb=100", "--enforce-all"]
- id: detect-private-key
# - repo: https://github.com/PyCQA/docformatter
# rev: v1.7.5
# hooks:
# - id: docformatter
# additional_dependencies: [tomli]
# args: ["--in-place"]
- repo: https://github.com/executablebooks/mdformat
rev: 0.7.17
hooks:
- id: mdformat
exclude: '^.*\.md$'
args: ["--number"]
additional_dependencies:
- mdformat-gfm
- mdformat-black
- mdformat_frontmatter
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.6.9
hooks:
# try to fix what is possible
- id: ruff
args: ["--fix", "--ignore", "E501,F401,F403,F841,E741"]
# # perform formatting updates
# - id: ruff-format
# validate if all is fine with preview mode
- id: ruff
args: ["--ignore", "E501,F401,F403,F841,E741"]

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.5 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 162 KiB

View File

@ -0,0 +1,132 @@
import base64
import json
from datetime import datetime
import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as T
from PIL import Image, ImageDraw
class DfineDetector:
def __init__(self, config_path: str, model_path: str, label_map: str, device="cpu", score_thresh=0.3):
# 初始化配置和模型
from src.core import YAMLConfig
cfg = YAMLConfig(config_path, resume=model_path)
if "HGNetv2" in cfg.yaml_cfg:
cfg.yaml_cfg["HGNetv2"]["pretrained"] = False
checkpoint = torch.load(model_path, map_location="cpu")
if "ema" in checkpoint:
state = checkpoint["ema"]["module"]
else:
state = checkpoint["model"]
cfg.model.load_state_dict(state)
class Model(nn.Module):
def __init__(self):
super().__init__()
self.model = cfg.model.deploy()
self.postprocessor = cfg.postprocessor.deploy()
def forward(self, images, orig_target_sizes):
outputs = self.model(images)
outputs = self.postprocessor(outputs, orig_target_sizes)
return outputs
self.device = device
self.model = Model().to(device)
self.model.eval()
self.label_map = json.loads(label_map) # list[dict{"label":"类别英文","name":"中文"}]
self.score_thresh = score_thresh
def detect(self, image_path: str):
try:
im_pil = Image.open(image_path).convert("RGB")
w, h = im_pil.size
orig_size = torch.tensor([[w, h]]).to(self.device)
transforms = T.Compose([
T.Resize((640, 640)),
T.ToTensor(),
])
im_data = transforms(im_pil).unsqueeze(0).to(self.device)
with torch.no_grad():
output = self.model(im_data, orig_size)
labels, boxes, scores = output
idx = scores[0] > self.score_thresh
result_labels = labels[0][idx].detach().cpu().numpy()
result_boxes = boxes[0][idx].detach().cpu().numpy()
result_scores = scores[0][idx].detach().cpu().numpy()
obj_list = []
for l, b, s in zip(result_labels, result_boxes, result_scores):
code = int(l)
# 默认 label_map 结构: [{"label":"person","name":"人"}]
label_str = str(code)
code_name = None
for item in self.label_map:
# 支持数字 or 英文名的映射
if str(item.get("label")) == str(code):
label_str = item["label"]
code_name = item.get("name", "")
break
one_obj = {
"code": label_str, # YOLO 是英文类别,这里建议保持和 label_map 结构统一
"name": code_name,
"score": round(float(s), 2),
"x1": int(b[0]),
"y1": int(b[1]),
"x2": int(b[2]),
"y2": int(b[3])
}
obj_list.append(one_obj)
detection_image = self.draw_and_encode(im_pil, obj_list)
result_json = {
"image": detection_image,
"obj_list": obj_list,
"time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
}
return json.dumps(result_json, ensure_ascii=False)
except Exception as e:
print(f"Error during detection: {e}")
return json.dumps({'error': str(e)}, ensure_ascii=False)
def draw_and_encode(self, im_pil, obj_list):
# 转 opencv 画框
img = np.array(im_pil)
if img.shape[-1] == 3:
pass
else:
img = img[:, :, :3]
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
preset_colors = [
(0, 0, 255), (20, 20, 255), (20, 40, 255), (0, 60, 255), (0, 80, 255),
(0, 100, 255), (0, 120, 255), (0, 140, 255), (0, 160, 255), (0, 180, 255),
(0, 200, 255), (0, 220, 255), (0, 255, 255),
]
category_colors = {}
for i, obj in enumerate(obj_list):
category = obj['code']
if category not in category_colors:
category_colors[category] = preset_colors[i % len(preset_colors)]
x1, y1, x2, y2 = obj['x1'], obj['y1'], obj['x2'], obj['y2']
color = category_colors[category]
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
font = cv2.FONT_HERSHEY_SIMPLEX
text = f"{category}-{obj['score']}"
(text_width, _), _ = cv2.getTextSize(text, font, 0.5, 1)
cv2.rectangle(img, (x1, y1 - 20), (x1 + text_width, y1), color, -1)
cv2.putText(img, text, (x1, y1 - 5), font, 0.5, (255, 255, 255), 1)
_, img_encoded = cv2.imencode('.jpg', img)
img_base64 = base64.b64encode(img_encoded.tobytes()).decode('utf-8')
return img_base64

View File

@ -0,0 +1,48 @@
FROM registry.cn-hangzhou.aliyuncs.com/peterande/dfine:v1
# FULL BUILDING INFO:
# docker login --username=xxx registry.cn-hangzhou.aliyuncs.com
# cd [PATH_2_Dockerfile]
# docker build -t xxx:v1 .
# docker tag xxx:v1 registry.cn-hangzhou.aliyuncs.com/xxx/xxx:v1
# docker push registry.cn-hangzhou.aliyuncs.com/xxx/xxx:v1
# FROM dockerpull.com/nvidia/cuda:12.0.1-cudnn8-devel-ubuntu18.04
# ARG DEBIAN_FRONTEND=noninteractive
# ENV PATH="/root/miniconda3/bin:${PATH}"
# ARG PATH="/root/miniconda3/bin:${PATH}"
# RUN sed -i "s/archive.ubuntu./mirrors.aliyun./g" /etc/apt/sources.list
# RUN sed -i "s/deb.debian.org/mirrors.aliyun.com/g" /etc/apt/sources.list
# RUN sed -i "s/security.debian.org/mirrors.aliyun.com\/debian-security/g" /etc/apt/sources.list
# RUN sed -i 's/archive.ubuntu.com/mirrors.ustc.edu.cn/g' /etc/apt/sources.list
# RUN apt-get update && apt-get install -y --no-install-recommends apt-utils && \
# apt-get upgrade -y && \
# apt-get install -y vim git libgl1-mesa-glx libglib2.0-0 libsm6 && \
# apt-get install -y libxrender1 libxext6 tmux wget htop && \
# apt-get install -y build-essential gcc g++ gdb binutils pciutils net-tools iputils-ping iproute2 git vim wget curl make openssh-server openssh-client tmux tree man unzip unrar
# ENV PYTHONIOENCODING=UTF-8
# RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
# mkdir /root/.conda && \
# bash Miniconda3-latest-Linux-x86_64.sh -b && \
# rm -f Miniconda3-latest-Linux-x86_64.sh && \
# conda init bash
# RUN conda config --set show_channel_urls yes \
# && echo "channels:" > ~/.condarc \
# && echo " - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/" >> ~/.condarc \
# && echo " - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/" >> ~/.condarc \
# && echo "show_channel_urls: true" \
# && cat ~/.condarc \
# && pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple \
# && cat ~/.config/pip/pip.conf
# RUN python3 -m pip install --upgrade pip && \
# python3 -m pip install --upgrade setuptools
# RUN python3 -m pip install jupyterlab pycocotools PyYAML tensorboard scipy
# RUN python3 -m pip --default-timeout=10000 install torch torchvision

View File

@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

View File

@ -0,0 +1,700 @@
<!--# [D-FINE: Redefine Regression Task of DETRs as Fine-grained Distribution Refinement](https://arxiv.org/abs/xxxxxx) -->
English | [简体中文](README_cn.md) | [日本語](README_ja.md) | [English Blog](src/zoo/dfine/blog.md) | [中文博客](src/zoo/dfine/blog_cn.md)
<h2 align="center">
D-FINE: Redefine Regression Task of DETRs as Fine&#8209;grained&nbsp;Distribution&nbsp;Refinement
</h2>
<p align="center">
<a href="https://huggingface.co/spaces/developer0hye/D-FINE">
<img alt="hf" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue">
</a>
<a href="https://github.com/Peterande/D-FINE/blob/master/LICENSE">
<img alt="license" src="https://img.shields.io/badge/LICENSE-Apache%202.0-blue">
</a>
<a href="https://github.com/Peterande/D-FINE/pulls">
<img alt="prs" src="https://img.shields.io/github/issues-pr/Peterande/D-FINE">
</a>
<a href="https://github.com/Peterande/D-FINE/issues">
<img alt="issues" src="https://img.shields.io/github/issues/Peterande/D-FINE?color=olive">
</a>
<a href="https://arxiv.org/abs/2410.13842">
<img alt="arXiv" src="https://img.shields.io/badge/arXiv-2410.13842-red">
</a>
<!-- <a href="mailto: pengyansong@mail.ustc.edu.cn">
<img alt="email" src="https://img.shields.io/badge/contact_me-email-yellow">
</a> -->
<a href="https://results.pre-commit.ci/latest/github/Peterande/D-FINE/master">
<img alt="pre-commit.ci status" src="https://results.pre-commit.ci/badge/github/Peterande/D-FINE/master.svg">
</a>
<a href="https://github.com/Peterande/D-FINE">
<img alt="stars" src="https://img.shields.io/github/stars/Peterande/D-FINE">
</a>
</p>
<p align="center">
📄 This is the official implementation of the paper:
<br>
<a href="https://arxiv.org/abs/2410.13842">D-FINE: Redefine Regression Task of DETRs as Fine-grained Distribution Refinement</a>
</p>
<p align="center">
Yansong Peng, Hebei Li, Peixi Wu, Yueyi Zhang, Xiaoyan Sun, and Feng Wu
</p>
<p align="center">
University of Science and Technology of China
</p>
<p align="center">
<a href="https://paperswithcode.com/sota/real-time-object-detection-on-coco?p=d-fine-redefine-regression-task-in-detrs-as">
<img alt="sota" src="https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/d-fine-redefine-regression-task-in-detrs-as/real-time-object-detection-on-coco">
</a>
</p>
<!-- <table><tr>
<td><img src=https://github.com/Peterande/storage/blob/master/latency.png border=0 width=333></td>
<td><img src=https://github.com/Peterande/storage/blob/master/params.png border=0 width=333></td>
<td><img src=https://github.com/Peterande/storage/blob/master/flops.png border=0 width=333></td>
</tr></table> -->
<p align="center">
<strong>If you like D-FINE, please give us a ⭐! Your support motivates us to keep improving!</strong>
</p>
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/stats_padded.png" width="1000">
</p>
D-FINE is a powerful real-time object detector that redefines the bounding box regression task in DETRs as Fine-grained Distribution Refinement (FDR) and introduces Global Optimal Localization Self-Distillation (GO-LSD), achieving outstanding performance without introducing additional inference and training costs.
<details open>
<summary> Video </summary>
We conduct object detection using D-FINE and YOLO11 on a complex street scene video from [YouTube](https://www.youtube.com/watch?v=CfhEWj9sd9A). Despite challenging conditions such as backlighting, motion blur, and dense crowds, D-FINE-X successfully detects nearly all targets, including subtle small objects like backpacks, bicycles, and traffic lights. Its confidence scores and the localization precision for blurred edges are significantly higher than those of YOLO11.
<!-- We use D-FINE and YOLO11 on a street scene video from [YouTube](https://www.youtube.com/watch?v=CfhEWj9sd9A). Despite challenges like backlighting, motion blur, and dense crowds, D-FINE-X outperforms YOLO11x, detecting more objects with higher confidence and better precision. -->
https://github.com/user-attachments/assets/e5933d8e-3c8a-400e-870b-4e452f5321d9
</details>
## 🚀 Updates
- [x] **\[2024.10.18\]** Release D-FINE series.
- [x] **\[2024.10.25\]** Add custom dataset finetuning configs ([#7](https://github.com/Peterande/D-FINE/issues/7)).
- [x] **\[2024.10.30\]** Update D-FINE-L (E25) pretrained model, with performance improved by 2.0%.
- [x] **\[2024.11.07\]** Release **D-FINE-N**, achiving 42.8% AP<sup>val</sup> on COCO @ 472 FPS<sup>T4</sup>!
## Model Zoo
### COCO
| Model | Dataset | AP<sup>val</sup> | #Params | Latency | GFLOPs | config | checkpoint | logs |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
**D&#8209;FINE&#8209;N** | COCO | **42.8** | 4M | 2.12ms | 7 | [yml](./configs/dfine/dfine_hgnetv2_n_coco.yml) | [42.8](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_n_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_n_coco_log.txt)
**D&#8209;FINE&#8209;S** | COCO | **48.5** | 10M | 3.49ms | 25 | [yml](./configs/dfine/dfine_hgnetv2_s_coco.yml) | [48.5](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_s_coco_log.txt)
**D&#8209;FINE&#8209;M** | COCO | **52.3** | 19M | 5.62ms | 57 | [yml](./configs/dfine/dfine_hgnetv2_m_coco.yml) | [52.3](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_m_coco_log.txt)
**D&#8209;FINE&#8209;L** | COCO | **54.0** | 31M | 8.07ms | 91 | [yml](./configs/dfine/dfine_hgnetv2_l_coco.yml) | [54.0](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_l_coco_log.txt)
**D&#8209;FINE&#8209;X** | COCO | **55.8** | 62M | 12.89ms | 202 | [yml](./configs/dfine/dfine_hgnetv2_x_coco.yml) | [55.8](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_x_coco_log.txt)
### Objects365+COCO
| Model | Dataset | AP<sup>val</sup> | #Params | Latency | GFLOPs | config | checkpoint | logs |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
**D&#8209;FINE&#8209;S** | Objects365+COCO | **50.7** | 10M | 3.49ms | 25 | [yml](./configs/dfine/objects365/dfine_hgnetv2_s_obj2coco.yml) | [50.7](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_obj2coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_s_obj2coco_log.txt)
**D&#8209;FINE&#8209;M** | Objects365+COCO | **55.1** | 19M | 5.62ms | 57 | [yml](./configs/dfine/objects365/dfine_hgnetv2_m_obj2coco.yml) | [55.1](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_obj2coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_m_obj2coco_log.txt)
**D&#8209;FINE&#8209;L** | Objects365+COCO | **57.3** | 31M | 8.07ms | 91 | [yml](./configs/dfine/objects365/dfine_hgnetv2_l_obj2coco.yml) | [57.3](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj2coco_e25.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_l_obj2coco_log_e25.txt)
**D&#8209;FINE&#8209;X** | Objects365+COCO | **59.3** | 62M | 12.89ms | 202 | [yml](./configs/dfine/objects365/dfine_hgnetv2_x_obj2coco.yml) | [59.3](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_obj2coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_x_obj2coco_log.txt)
**We highly recommend that you use the Objects365 pre-trained model for fine-tuning:**
⚠️ **Important**: Please note that this is generally beneficial for complex scene understanding. If your categories are very simple, it might lead to overfitting and suboptimal performance.
<details>
<summary><strong> 🔥 Pretrained Models on Objects365 (Best generalization) </strong></summary>
| Model | Dataset | AP<sup>val</sup> | AP<sup>5000</sup> | #Params | Latency | GFLOPs | config | checkpoint | logs |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
**D&#8209;FINE&#8209;S** | Objects365 | **31.0** | **30.5** | 10M | 3.49ms | 25 | [yml](./configs/dfine/objects365/dfine_hgnetv2_s_obj365.yml) | [30.5](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_s_obj365_log.txt)
**D&#8209;FINE&#8209;M** | Objects365 | **38.6** | **37.4** | 19M | 5.62ms | 57 | [yml](./configs/dfine/objects365/dfine_hgnetv2_m_obj365.yml) | [37.4](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_m_obj365_log.txt)
**D&#8209;FINE&#8209;L** | Objects365 | - | **40.6** | 31M | 8.07ms | 91 | [yml](./configs/dfine/objects365/dfine_hgnetv2_l_obj365.yml) | [40.6](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_l_obj365_log.txt)
**D&#8209;FINE&#8209;L (E25)** | Objects365 | **44.7** | **42.6** | 31M | 8.07ms | 91 | [yml](./configs/dfine/objects365/dfine_hgnetv2_l_obj365.yml) | [42.6](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj365_e25.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_l_obj365_log_e25.txt)
**D&#8209;FINE&#8209;X** | Objects365 | **49.5** | **46.5** | 62M | 12.89ms | 202 | [yml](./configs/dfine/objects365/dfine_hgnetv2_x_obj365.yml) | [46.5](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_x_obj365_log.txt)
- **E25**: Re-trained and extended the pretraining to 25 epochs.
- **AP<sup>val</sup>** is evaluated on *Objects365* full validation set.
- **AP<sup>5000</sup>** is evaluated on the first 5000 samples of the *Objects365* validation set.
</details>
**Notes:**
- **AP<sup>val</sup>** is evaluated on *MSCOCO val2017* dataset.
- **Latency** is evaluated on a single T4 GPU with $batch\\_size = 1$, $fp16$, and $TensorRT==10.4.0$.
- **Objects365+COCO** means finetuned model on *COCO* using pretrained weights trained on *Objects365*.
## Quick start
### Setup
```shell
conda create -n dfine python=3.11.9
conda activate dfine
pip install -r requirements.txt
```
### Data Preparation
<details>
<summary> COCO2017 Dataset </summary>
1. Download COCO2017 from [OpenDataLab](https://opendatalab.com/OpenDataLab/COCO_2017) or [COCO](https://cocodataset.org/#download).
1. Modify paths in [coco_detection.yml](./configs/dataset/coco_detection.yml)
```yaml
train_dataloader:
img_folder: /data/COCO2017/train2017/
ann_file: /data/COCO2017/annotations/instances_train2017.json
val_dataloader:
img_folder: /data/COCO2017/val2017/
ann_file: /data/COCO2017/annotations/instances_val2017.json
```
</details>
<details>
<summary> Objects365 Dataset </summary>
1. Download Objects365 from [OpenDataLab](https://opendatalab.com/OpenDataLab/Objects365).
2. Set the Base Directory:
```shell
export BASE_DIR=/data/Objects365/data
```
3. Extract and organize the downloaded files, resulting directory structure:
```shell
${BASE_DIR}/train
├── images
│ ├── v1
│ │ ├── patch0
│ │ │ ├── 000000000.jpg
│ │ │ ├── 000000001.jpg
│ │ │ └── ... (more images)
│ ├── v2
│ │ ├── patchx
│ │ │ ├── 000000000.jpg
│ │ │ ├── 000000001.jpg
│ │ │ └── ... (more images)
├── zhiyuan_objv2_train.json
```
```shell
${BASE_DIR}/val
├── images
│ ├── v1
│ │ ├── patch0
│ │ │ ├── 000000000.jpg
│ │ │ └── ... (more images)
│ ├── v2
│ │ ├── patchx
│ │ │ ├── 000000000.jpg
│ │ │ └── ... (more images)
├── zhiyuan_objv2_val.json
```
4. Create a New Directory to Store Images from the Validation Set:
```shell
mkdir -p ${BASE_DIR}/train/images_from_val
```
5. Copy the v1 and v2 folders from the val directory into the train/images_from_val directory
```shell
cp -r ${BASE_DIR}/val/images/v1 ${BASE_DIR}/train/images_from_val/
cp -r ${BASE_DIR}/val/images/v2 ${BASE_DIR}/train/images_from_val/
```
6. Run remap_obj365.py to merge a subset of the validation set into the training set. Specifically, this script moves samples with indices between 5000 and 800000 from the validation set to the training set.
```shell
python tools/remap_obj365.py --base_dir ${BASE_DIR}
```
7. Run the resize_obj365.py script to resize any images in the dataset where the maximum edge length exceeds 640 pixels. Use the updated JSON file generated in Step 5 to process the sample data. Ensure that you resize images in both the train and val datasets to maintain consistency.
```shell
python tools/resize_obj365.py --base_dir ${BASE_DIR}
```
8. Modify paths in [obj365_detection.yml](./configs/dataset/obj365_detection.yml)
```yaml
train_dataloader:
img_folder: /data/Objects365/data/train
ann_file: /data/Objects365/data/train/new_zhiyuan_objv2_train_resized.json
val_dataloader:
img_folder: /data/Objects365/data/val/
ann_file: /data/Objects365/data/val/new_zhiyuan_objv2_val_resized.json
```
</details>
<details>
<summary>CrowdHuman</summary>
Download COCO format dataset here: [url](https://aistudio.baidu.com/datasetdetail/231455)
</details>
<details>
<summary>Custom Dataset</summary>
To train on your custom dataset, you need to organize it in the COCO format. Follow the steps below to prepare your dataset:
1. **Set `remap_mscoco_category` to `False`:**
This prevents the automatic remapping of category IDs to match the MSCOCO categories.
```yaml
remap_mscoco_category: False
```
2. **Organize Images:**
Structure your dataset directories as follows:
```shell
dataset/
├── images/
│ ├── train/
│ │ ├── image1.jpg
│ │ ├── image2.jpg
│ │ └── ...
│ ├── val/
│ │ ├── image1.jpg
│ │ ├── image2.jpg
│ │ └── ...
└── annotations/
├── instances_train.json
├── instances_val.json
└── ...
```
- **`images/train/`**: Contains all training images.
- **`images/val/`**: Contains all validation images.
- **`annotations/`**: Contains COCO-formatted annotation files.
3. **Convert Annotations to COCO Format:**
If your annotations are not already in COCO format, you'll need to convert them. You can use the following Python script as a reference or utilize existing tools:
```python
import json
def convert_to_coco(input_annotations, output_annotations):
# Implement conversion logic here
pass
if __name__ == "__main__":
convert_to_coco('path/to/your_annotations.json', 'dataset/annotations/instances_train.json')
```
4. **Update Configuration Files:**
Modify your [custom_detection.yml](./configs/dataset/custom_detection.yml).
```yaml
task: detection
evaluator:
type: CocoEvaluator
iou_types: ['bbox', ]
num_classes: 777 # your dataset classes
remap_mscoco_category: False
train_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/yourdataset/train
ann_file: /data/yourdataset/train/train.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: True
num_workers: 4
drop_last: True
collate_fn:
type: BatchImageCollateFunction
val_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/yourdataset/val
ann_file: /data/yourdataset/val/ann.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: False
num_workers: 4
drop_last: False
collate_fn:
type: BatchImageCollateFunction
```
</details>
## Usage
<details open>
<summary> COCO2017 </summary>
<!-- <summary>1. Training </summary> -->
1. Set Model
```shell
export model=l # n s m l x
```
2. Training
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --use-amp --seed=0
```
<!-- <summary>2. Testing </summary> -->
3. Testing
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --test-only -r model.pth
```
<!-- <summary>3. Tuning </summary> -->
4. Tuning
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --use-amp --seed=0 -t model.pth
```
</details>
<details>
<summary> Objects365 to COCO2017 </summary>
1. Set Model
```shell
export model=l # n s m l x
```
2. Training on Objects365
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/objects365/dfine_hgnetv2_${model}_obj365.yml --use-amp --seed=0
```
3. Tuning on COCO2017
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/objects365/dfine_hgnetv2_${model}_obj2coco.yml --use-amp --seed=0 -t model.pth
```
<!-- <summary>2. Testing </summary> -->
4. Testing
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --test-only -r model.pth
```
</details>
<details>
<summary> Custom Dataset </summary>
1. Set Model
```shell
export model=l # n s m l x
```
2. Training on Custom Dataset
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/dfine_hgnetv2_${model}_custom.yml --use-amp --seed=0
```
<!-- <summary>2. Testing </summary> -->
3. Testing
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/dfine_hgnetv2_${model}_custom.yml --test-only -r model.pth
```
4. Tuning on Custom Dataset
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/objects365/dfine_hgnetv2_${model}_obj2custom.yml --use-amp --seed=0 -t model.pth
```
5. **[Optional]** Modify Class Mappings:
When using the Objects365 pre-trained weights to train on your custom dataset, the example assumes that your dataset only contains the classes `'Person'` and `'Car'`. For faster convergence, you can modify `self.obj365_ids` in `src/solver/_solver.py` as follows:
```python
self.obj365_ids = [0, 5] # Person, Cars
```
You can replace these with any corresponding classes from your dataset. The list of Objects365 classes with their corresponding IDs:
https://github.com/Peterande/D-FINE/blob/352a94ece291e26e1957df81277bef00fe88a8e3/src/solver/_solver.py#L330
New training command:
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/dfine_hgnetv2_${model}_custom.yml --use-amp --seed=0 -t model.pth
```
However, if you don't wish to modify the class mappings, the pre-trained Objects365 weights will still work without any changes. Modifying the class mappings is optional and can potentially accelerate convergence for specific tasks.
</details>
<details>
<summary> Customizing Batch Size </summary>
For example, if you want to double the total batch size when training D-FINE-L on COCO2017, here are the steps you should follow:
1. **Modify your [dataloader.yml](./configs/dfine/include/dataloader.yml)** to increase the `total_batch_size`:
```yaml
train_dataloader:
total_batch_size: 64 # Previously it was 32, now doubled
```
2. **Modify your [dfine_hgnetv2_l_coco.yml](./configs/dfine/dfine_hgnetv2_l_coco.yml)**. Heres how the key parameters should be adjusted:
```yaml
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.000025 # doubled, linear scaling law
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.0005 # doubled, linear scaling law
betas: [0.9, 0.999]
weight_decay: 0.0001 # need a grid search
ema: # added EMA settings
decay: 0.9998 # adjusted by 1 - (1 - decay) * 2
warmups: 500 # halved
lr_warmup_scheduler:
warmup_duration: 250 # halved
```
</details>
<details>
<summary> Customizing Input Size </summary>
If you'd like to train **D-FINE-L** on COCO2017 with an input size of 320x320, follow these steps:
1. **Modify your [dataloader.yml](./configs/dfine/include/dataloader.yml)**:
```yaml
train_dataloader:
dataset:
transforms:
ops:
- {type: Resize, size: [320, 320], }
collate_fn:
base_size: 320
dataset:
transforms:
ops:
- {type: Resize, size: [320, 320], }
```
2. **Modify your [dfine_hgnetv2.yml](./configs/dfine/include/dfine_hgnetv2.yml)**:
```yaml
eval_spatial_size: [320, 320]
```
</details>
## Tools
<details>
<summary> Deployment </summary>
<!-- <summary>4. Export onnx </summary> -->
1. Setup
```shell
pip install onnx onnxsim
export model=l # n s m l x
```
2. Export onnx
```shell
python tools/deployment/export_onnx.py --check -c configs/dfine/dfine_hgnetv2_${model}_coco.yml -r model.pth
```
3. Export [tensorrt](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html)
```shell
trtexec --onnx="model.onnx" --saveEngine="model.engine" --fp16
```
</details>
<details>
<summary> Inference (Visualization) </summary>
1. Setup
```shell
pip install -r tools/inference/requirements.txt
export model=l # n s m l x
```
<!-- <summary>5. Inference </summary> -->
2. Inference (onnxruntime / tensorrt / torch)
Inference on images and videos is now supported.
```shell
python tools/inference/onnx_inf.py --onnx model.onnx --input image.jpg # video.mp4
python tools/inference/trt_inf.py --trt model.engine --input image.jpg
python tools/inference/torch_inf.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml -r model.pth --input image.jpg --device cuda:0
```
</details>
<details>
<summary> Benchmark </summary>
1. Setup
```shell
pip install -r tools/benchmark/requirements.txt
export model=l # n s m l x
```
<!-- <summary>6. Benchmark </summary> -->
2. Model FLOPs, MACs, and Params
```shell
python tools/benchmark/get_info.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml
```
2. TensorRT Latency
```shell
python tools/benchmark/trt_benchmark.py --COCO_dir path/to/COCO2017 --engine_dir model.engine
```
</details>
<details>
<summary> Fiftyone Visualization </summary>
1. Setup
```shell
pip install fiftyone
export model=l # n s m l x
```
4. Voxel51 Fiftyone Visualization ([fiftyone](https://github.com/voxel51/fiftyone))
```shell
python tools/visualization/fiftyone_vis.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml -r model.pth
```
</details>
<details>
<summary> Others </summary>
1. Auto Resume Training
```shell
bash reference/safe_training.sh
```
2. Converting Model Weights
```shell
python reference/convert_weight.py model.pth
```
</details>
## Figures and Visualizations
<details>
<summary> FDR and GO-LSD </summary>
1. Overview of D-FINE with FDR. The probability distributions that act as a more fine-
grained intermediate representation are iteratively refined by the decoder layers in a residual manner.
Non-uniform weighting functions are applied to allow for finer localization.
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/fdr-1.jpg" alt="Fine-grained Distribution Refinement Process" width="1000">
</p>
2. Overview of GO-LSD process. Localization knowledge from the final layers refined
distributions is distilled into earlier layers through DDF loss with decoupled weighting strategies.
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/go_lsd-1.jpg" alt="GO-LSD Process" width="1000">
</p>
</details>
<details open>
<summary> Distributions </summary>
Visualizations of FDR across detection scenarios with initial and refined bounding boxes, along with unweighted and weighted distributions.
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/merged_image.jpg" width="1000">
</p>
</details>
<details>
<summary> Hard Cases </summary>
The following visualization demonstrates D-FINE's predictions in various complex detection scenarios. These include cases with occlusion, low-light conditions, motion blur, depth of field effects, and densely populated scenes. Despite these challenges, D-FINE consistently produces accurate localization results.
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/hard_case-1.jpg" alt="D-FINE Predictions in Challenging Scenarios" width="1000">
</p>
</details>
<!-- <div style="display: flex; flex-wrap: wrap; justify-content: center; margin: 0; padding: 0;">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/merged_image.jpg" style="width:99.96%; margin: 0; padding: 0;" />
</div>
<table><tr>
<td><img src=https://raw.githubusercontent.com/Peterande/storage/master/figs/merged_image.jpg border=0 width=1000></td>
</tr></table> -->
## Citation
If you use `D-FINE` or its methods in your work, please cite the following BibTeX entries:
<details open>
<summary> bibtex </summary>
```latex
@misc{peng2024dfine,
title={D-FINE: Redefine Regression Task in DETRs as Fine-grained Distribution Refinement},
author={Yansong Peng and Hebei Li and Peixi Wu and Yueyi Zhang and Xiaoyan Sun and Feng Wu},
year={2024},
eprint={2410.13842},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
```
</details>
## Acknowledgement
Our work is built upon [RT-DETR](https://github.com/lyuwenyu/RT-DETR).
Thanks to the inspirations from [RT-DETR](https://github.com/lyuwenyu/RT-DETR), [GFocal](https://github.com/implus/GFocal), [LD](https://github.com/HikariTJU/LD), and [YOLOv9](https://github.com/WongKinYiu/yolov9).
✨ Feel free to contribute and reach out if you have any questions! ✨

View File

@ -0,0 +1,673 @@
<!--# [D-FINE: Redefine Regression Task of DETRs as Fine-grained Distribution Refinement](https://arxiv.org/abs/xxxxxx) -->
[English](README.md) | 简体中文 | [日本語](README_ja.md) | [English Blog](src/zoo/dfine/blog.md) | [中文博客](src/zoo/dfine/blog_cn.md)
<h2 align="center">
D-FINE: Redefine Regression Task of DETRs as Fine&#8209;grained&nbsp;Distribution&nbsp;Refinement
</h2>
<p align="center">
<!-- <a href="https://paperswithcode.com/sota/real-time-object-detection-on-coco?p=d-fine-redefine-regression-task-in-detrs-as">
<img alt="sota" src="https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/d-fine-redefine-regression-task-in-detrs-as/real-time-object-detection-on-coco">
</a> -->
<a href="https://github.com/Peterande/D-FINE/blob/master/LICENSE">
<img alt="license" src="https://img.shields.io/badge/LICENSE-Apache%202.0-blue">
</a>
<a href="https://github.com/Peterande/D-FINE/pulls">
<img alt="prs" src="https://img.shields.io/github/issues-pr/Peterande/D-FINE">
</a>
<a href="https://github.com/Peterande/D-FINE/issues">
<img alt="issues" src="https://img.shields.io/github/issues/Peterande/D-FINE?color=olive">
</a>
<a href="https://arxiv.org/abs/2410.13842">
<img alt="arXiv" src="https://img.shields.io/badge/arXiv-2410.13842-red">
</a>
<!-- <a href="mailto: pengyansong@mail.ustc.edu.cn">
<img alt="email" src="https://img.shields.io/badge/contact_me-email-yellow">
</a> -->
<a href="https://results.pre-commit.ci/latest/github/Peterande/D-FINE/master">
<img alt="pre-commit.ci status" src="https://results.pre-commit.ci/badge/github/Peterande/D-FINE/master.svg">
</a>
<a href="https://github.com/Peterande/D-FINE">
<img alt="stars" src="https://img.shields.io/github/stars/Peterande/D-FINE">
</a>
</p>
<p align="center">
📄 这是该文章的官方实现:
<br>
<a href="https://arxiv.org/abs/2410.13842">D-FINE: Redefine Regression Task of DETRs as Fine-grained Distribution Refinement</a>
</p>
<p align="center">
彭岩松,李和倍,吴沛熹,张越一,孙晓艳,吴枫
</p>
<p align="center">
中国科学技术大学
</p>
<p align="center">
<a href="https://paperswithcode.com/sota/real-time-object-detection-on-coco?p=d-fine-redefine-regression-task-in-detrs-as">
<img alt="sota" src="https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/d-fine-redefine-regression-task-in-detrs-as/real-time-object-detection-on-coco">
</a>
</p>
<p align="center"> <strong>如果你喜欢 D-FINE请给我们一个 ⭐!你的支持激励我们不断前进!</strong> </p>
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/stats_padded.png" width="1000">
</p>
D-FINE 是一个强大的实时目标检测器,将 DETR 中的边界框回归任务重新定义为了细粒度的分布优化FDR并引入全局最优的定位自蒸馏GO-LSD在不增加额外推理和训练成本的情况下实现了卓越的性能。
<details open>
<summary> 视频 </summary>
我们分别使用 D-FINE 和 YOLO11 对 [YouTube](https://www.youtube.com/watch?v=CfhEWj9sd9A) 上的一段复杂街景视频进行了目标检测。尽管存在逆光、虚化模糊和密集遮挡等不利因素D-FINE-X 依然成功检测出几乎所有目标,包括背包、自行车和信号灯等难以察觉的小目标,其置信度、以及模糊边缘的定位准确度明显高于 YOLO11x。
https://github.com/user-attachments/assets/e5933d8e-3c8a-400e-870b-4e452f5321d9
</details>
## 🚀 Updates
- [x] **\[2024.10.18\]** 发布 D-FINE 系列。
- [x] **\[2024.10.25\]** 添加了自定义数据集微调配置文件 ([#7](https://github.com/Peterande/D-FINE/issues/7))。
- [x] **\[2024.10.30\]** 更新 D-FINE-L (E25) 预训练模型,性能提升了 2.0%。
- [x] **\[2024.11.07\]** 发布 **D-FINE-N**, 在 COCO 上达到 42.8% AP<sup>val</sup> @ 472 FPS<sup>T4</sup>!
## 模型库
### COCO
| 模型 | 数据集 | AP<sup>val</sup> | 参数量 | 时延 (ms) | GFLOPs | 配置 | 权重 | 日志 |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
**D&#8209;FINE&#8209;N** | COCO | **42.8** | 4M | 2.12ms | 7 | [yml](./configs/dfine/dfine_hgnetv2_n_coco.yml) | [42.8](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_n_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_n_coco_log.txt)
**D&#8209;FINE&#8209;S** | COCO | **48.5** | 10M | 3.49ms | 25 | [yml](./configs/dfine/dfine_hgnetv2_s_coco.yml) | [48.5](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_s_coco_log.txt)
**D&#8209;FINE&#8209;M** | COCO | **52.3** | 19M | 5.62ms | 57 | [yml](./configs/dfine/dfine_hgnetv2_m_coco.yml) | [52.3](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_m_coco_log.txt)
**D&#8209;FINE&#8209;L** | COCO | **54.0** | 31M | 8.07ms | 91 | [yml](./configs/dfine/dfine_hgnetv2_l_coco.yml) | [54.0](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_l_coco_log.txt)
**D&#8209;FINE&#8209;X** | COCO | **55.8** | 62M | 12.89ms | 202 | [yml](./configs/dfine/dfine_hgnetv2_x_coco.yml) | [55.8](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_x_coco_log.txt)
### Objects365+COCO
| 模型 | 数据集 | AP<sup>val</sup> | 参数量 | 时延 (ms) | GFLOPs | 配置 | 权重 | 日志 |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
**D&#8209;FINE&#8209;S** | Objects365+COCO | **50.7** | 10M | 3.49ms | 25 | [yml](./configs/dfine/objects365/dfine_hgnetv2_s_obj2coco.yml) | [50.7](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_obj2coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_s_obj2coco_log.txt)
**D&#8209;FINE&#8209;M** | Objects365+COCO | **55.1** | 19M | 5.62ms | 57 | [yml](./configs/dfine/objects365/dfine_hgnetv2_m_obj2coco.yml) | [55.1](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_obj2coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_m_obj2coco_log.txt)
**D&#8209;FINE&#8209;L** | Objects365+COCO | **57.3** | 31M | 8.07ms | 91 | [yml](./configs/dfine/objects365/dfine_hgnetv2_l_obj2coco.yml) | [57.3](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj2coco_e25.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_l_obj2coco_log_e25.txt)
**D&#8209;FINE&#8209;X** | Objects365+COCO | **59.3** | 62M | 12.89ms | 202 | [yml](./configs/dfine/objects365/dfine_hgnetv2_x_obj2coco.yml) | [59.3](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_obj2coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_x_obj2coco_log.txt)
**我们强烈推荐您使用 Objects365 预训练模型进行微调:**
⚠️ 重要提醒:通常这种预训练模型对复杂场景的理解非常有用。如果您的类别非常简单,请注意,这可能会导致过拟合和次优性能。
<details> <summary><strong> 🔥 Objects365 预训练模型(泛化性最好)</strong></summary>
| 模型 | 数据集 | AP<sup>val</sup> | AP<sup>5000</sup> | 参数量 | 时延 (ms) | GFLOPs | 配置 | 权重 | 日志 |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
**D&#8209;FINE&#8209;S** | Objects365 | **31.0** | **30.5** | 10M | 3.49ms | 25 | [yml](./configs/dfine/objects365/dfine_hgnetv2_s_obj365.yml) | [30.5](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_s_obj365_log.txt)
**D&#8209;FINE&#8209;M** | Objects365 | **38.6** | **37.4** | 19M | 5.62ms | 57 | [yml](./configs/dfine/objects365/dfine_hgnetv2_m_obj365.yml) | [37.4](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_m_obj365_log.txt)
**D&#8209;FINE&#8209;L** | Objects365 | - | **40.6** | 31M | 8.07ms | 91 | [yml](./configs/dfine/objects365/dfine_hgnetv2_l_obj365.yml) | [40.6](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_l_obj365_log.txt)
**D&#8209;FINE&#8209;L (E25)** | Objects365 | **44.7** | **42.6** | 31M | 8.07ms | 91 | [yml](./configs/dfine/objects365/dfine_hgnetv2_l_obj365.yml) | [42.6](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj365_e25.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_l_obj365_log_e25.txt)
**D&#8209;FINE&#8209;X** | Objects365 | **49.5** | **46.5** | 62M | 12.89ms | 202 | [yml](./configs/dfine/objects365/dfine_hgnetv2_x_obj365.yml) | [46.5](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_x_obj365_log.txt)
- **E25**: 重新训练,并将训练延长至 25 个 epoch。
- **AP<sup>val</sup>** 是在 *Objects365* 完整的验证集上进行评估的。
- **AP<sup>5000</sup>** 是在 *Objects365* 验证集的前5000个样本上评估的。
</details>
**注意:**
- **AP<sup>val</sup>** 是在 *MSCOCO val2017* 数据集上评估的。
- **时延** 是在单张 T4 GPU 上以 $batch\\_size = 1$, $fp16$, 和 $TensorRT==10.4.0$ 评估的。
- **Objects365+COCO** 表示使用在 *Objects365* 上预训练的权重在 *COCO* 上微调的模型。
## 快速开始
### 设置
```shell
conda create -n dfine python=3.11.9
conda activate dfine
pip install -r requirements.txt
```
</details>
### 数据集准备
<details>
<summary> COCO2017 数据集 </summary>
1. 从 [OpenDataLab](https://opendatalab.com/OpenDataLab/COCO_2017) 或者 [COCO](https://cocodataset.org/#download) 下载 COCO2017。
1.修改 [coco_detection.yml](./configs/dataset/coco_detection.yml) 中的路径。
```yaml
train_dataloader:
img_folder: /data/COCO2017/train2017/
ann_file: /data/COCO2017/annotations/instances_train2017.json
val_dataloader:
img_folder: /data/COCO2017/val2017/
ann_file: /data/COCO2017/annotations/instances_val2017.json
```
</details>
<details>
<summary> Objects365 数据集 </summary>
1. 从 [OpenDataLab](https://opendatalab.com/OpenDataLab/Objects365) 下载 Objects365。
2. 设置数据集的基础目录:
```shell
export BASE_DIR=/data/Objects365/data
```
3. 解压并整理目录结构如下:
```shell
${BASE_DIR}/train
├── images
│ ├── v1
│ │ ├── patch0
│ │ │ ├── 000000000.jpg
│ │ │ ├── 000000001.jpg
│ │ │ └── ... (more images)
│ ├── v2
│ │ ├── patchx
│ │ │ ├── 000000000.jpg
│ │ │ ├── 000000001.jpg
│ │ │ └── ... (more images)
├── zhiyuan_objv2_train.json
```
```shell
${BASE_DIR}/val
├── images
│ ├── v1
│ │ ├── patch0
│ │ │ ├── 000000000.jpg
│ │ │ └── ... (more images)
│ ├── v2
│ │ ├── patchx
│ │ │ ├── 000000000.jpg
│ │ │ └── ... (more images)
├── zhiyuan_objv2_val.json
```
4. 创建一个新目录来存储验证集中的图像:
```shell
mkdir -p ${BASE_DIR}/train/images_from_val
```
5. 将 val 目录中的 v1 和 v2 文件夹复制到 train/images_from_val 目录中
```shell
cp -r ${BASE_DIR}/val/images/v1 ${BASE_DIR}/train/images_from_val/
cp -r ${BASE_DIR}/val/images/v2 ${BASE_DIR}/train/images_from_val/
```
6. 运行 remap_obj365.py 将验证集中的部分样本合并到训练集中。具体来说,该脚本将索引在 5000 到 800000 之间的样本从验证集移动到训练集。
```shell
python tools/remap_obj365.py --base_dir ${BASE_DIR}
```
7. 运行 resize_obj365.py 脚本,将数据集中任何最大边长超过 640 像素的图像进行大小调整。使用步骤 5 中生成的更新后的 JSON 文件处理样本数据。
```shell
python tools/resize_obj365.py --base_dir ${BASE_DIR}
```
8. 修改 [obj365_detection.yml](./configs/dataset/obj365_detection.yml) 中的路径。
```yaml
train_dataloader:
img_folder: /data/Objects365/data/train
ann_file: /data/Objects365/data/train/new_zhiyuan_objv2_train_resized.json
val_dataloader:
img_folder: /data/Objects365/data/val/
ann_file: /data/Objects365/data/val/new_zhiyuan_objv2_val_resized.json
```
</details>
<details>
<summary>CrowdHuman</summary>
在此下载 COCO 格式的数据集:[链接](https://aistudio.baidu.com/datasetdetail/231455)
</details>
<details>
<summary>自定义数据集</summary>
要在你的自定义数据集上训练,你需要将其组织为 COCO 格式。请按照以下步骤准备你的数据集:
1. **将 `remap_mscoco_category` 设置为 `False`:**
这可以防止类别 ID 自动映射以匹配 MSCOCO 类别。
```yaml
remap_mscoco_category: False
```
2. **组织图像:**
按以下结构组织你的数据集目录:
```shell
dataset/
├── images/
│ ├── train/
│ │ ├── image1.jpg
│ │ ├── image2.jpg
│ │ └── ...
│ ├── val/
│ │ ├── image1.jpg
│ │ ├── image2.jpg
│ │ └── ...
└── annotations/
├── instances_train.json
├── instances_val.json
└── ...
```
- **`images/train/`**: 包含所有训练图像。
- **`images/val/`**: 包含所有验证图像。
- **`annotations/`**: 包含 COCO 格式的注释文件。
3. **将注释转换为 COCO 格式:**
如果你的注释尚未为 COCO 格式,你需要进行转换。你可以参考以下 Python 脚本或使用现有工具:
```python
import json
def convert_to_coco(input_annotations, output_annotations):
# Implement conversion logic here
pass
if __name__ == "__main__":
convert_to_coco('path/to/your_annotations.json', 'dataset/annotations/instances_train.json')
```
4. **更新配置文件:**
修改你的 [custom_detection.yml](./configs/dataset/custom_detection.yml)。
```yaml
task: detection
evaluator:
type: CocoEvaluator
iou_types: ['bbox', ]
num_classes: 777 # your dataset classes
remap_mscoco_category: False
train_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/yourdataset/train
ann_file: /data/yourdataset/train/train.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: True
num_workers: 4
drop_last: True
collate_fn:
type: BatchImageCollateFunction
val_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/yourdataset/val
ann_file: /data/yourdataset/val/ann.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: False
num_workers: 4
drop_last: False
collate_fn:
type: BatchImageCollateFunction
```
</details>
## 使用方法
<details open>
<summary> COCO2017 </summary>
<!-- <summary>1. Training </summary> -->
1. 设置模型
```shell
export model=l # n s m l x
```
2. 训练
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --use-amp --seed=0
```
<!-- <summary>2. Testing </summary> -->
3. 测试
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --test-only -r model.pth
```
<!-- <summary>3. Tuning </summary> -->
4. 微调
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --use-amp --seed=0 -t model.pth
```
</details>
<details>
<summary> 在 Objects365 上训练在COCO2017上微调 </summary>
1. 设置模型
```shell
export model=l # n s m l x
```
2. 在 Objects365 上训练
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/objects365/dfine_hgnetv2_${model}_obj365.yml --use-amp --seed=0
```
3. 在 COCO2017 上微调
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/objects365/dfine_hgnetv2_${model}_obj2coco.yml --use-amp --seed=0 -t model.pth
```
<!-- <summary>2. Testing </summary> -->
4. 测试
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --test-only -r model.pth
```
</details>
<details>
<summary> 自定义数据集 </summary>
1. 设置模型
```shell
export model=l # n s m l x
```
2. 在自定义数据集上训练
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/dfine_hgnetv2_${model}_custom.yml --use-amp --seed=0
```
<!-- <summary>2. Testing </summary> -->
3. 测试
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/dfine_hgnetv2_${model}_custom.yml --test-only -r model.pth
```
4. 在自定义数据集上微调
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/objects365/dfine_hgnetv2_${model}_obj2custom.yml --use-amp --seed=0 -t model.pth
```
5. **[可选项]** 修改类映射:
在使用 Objects365 预训练权重训练自定义数据集时,示例中假设自定义数据集仅有 `'Person'``'Car'` 类,您可以将其替换为数据集中对应的任何类别。为了加快收敛,可以在 `src/solver/_solver.py` 中修改 `self.obj365_ids`,如下所示:
```python
self.obj365_ids = [0, 5] # Person, Cars
```
Objects365 类及其对应 ID 的完整列表:
https://github.com/Peterande/D-FINE/blob/352a94ece291e26e1957df81277bef00fe88a8e3/src/solver/_solver.py#L330
新的训练启动命令:
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/dfine_hgnetv2_${model}_custom.yml --use-amp --seed=0 -t model.pth
```
如果您不想修改类映射,预训练的 Objects365 权重依然可以不做任何更改直接使用。修改类映射是可选的,但针对特定任务可能会加快收敛速度。
</details>
<details>
<summary> 自定义批次大小 </summary>
例如,如果你想在训练 D-FINE-L 时将 COCO2017 的总批次大小增加一倍,请按照以下步骤操作:
1. **修改你的 [dataloader.yml](./configs/dfine/include/dataloader.yml)**,增加 `total_batch_size`
```yaml
train_dataloader:
total_batch_size: 64 # 原来是 32现在增加了一倍
```
2. **修改你的 [dfine_hgnetv2_l_coco.yml](./configs/dfine/dfine_hgnetv2_l_coco.yml)**
```yaml
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.000025 # 翻倍,线性缩放原则
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.0005 # 翻倍,线性缩放原则
betas: [0.9, 0.999]
weight_decay: 0.0001 # 需要网格搜索找到最优值
ema: # 添加 EMA 设置
decay: 0.9998 # 根据 1 - (1 - decay) * 2 调整
warmups: 500 # 减半
lr_warmup_scheduler:
warmup_duration: 250 # 减半
```
</details>
<details>
<summary> 自定义输入尺寸 </summary>
如果你想在 COCO2017 上使用 **D-FINE-L** 进行 320x320 尺寸的图片训练,按照以下步骤操作:
1. **修改你的 [dataloader.yml](./configs/dfine/include/dataloader.yml)**
```yaml
train_dataloader:
dataset:
transforms:
ops:
- {type: Resize, size: [320, 320], }
collate_fn:
base_size: 320
dataset:
transforms:
ops:
- {type: Resize, size: [320, 320], }
```
2. **修改你的 [dfine_hgnetv2.yml](./configs/dfine/include/dfine_hgnetv2.yml)**
```yaml
eval_spatial_size: [320, 320]
```
</details>
## 工具
<details>
<summary> 部署 </summary>
<!-- <summary>4. Export onnx </summary> -->
1. 设置
```shell
pip install onnx onnxsim onnxruntime
export model=l # n s m l x
```
2. 导出 onnx
```shell
python tools/export_onnx.py --check -c configs/dfine/dfine_hgnetv2_${model}_coco.yml -r model.pth
```
3. 导出 [tensorrt](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html)
```shell
trtexec --onnx="model.onnx" --saveEngine="model.engine" --fp16
```
</details>
<details>
<summary> 推理(可视化) </summary>
1. 设置
```shell
pip install -r tools/inference/requirements.txt
export model=l # n s m l x
```
<!-- <summary>5. Inference </summary> -->
2. 推理 (onnxruntime / tensorrt / torch)
目前支持对图像和视频的推理。
```shell
python tools/inference/onnx_inf.py --onnx model.onnx --input image.jpg # video.mp4
python tools/inference/trt_inf.py --trt model.engine --input image.jpg
python tools/inference/torch_inf.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml -r model.pth --input image.jpg --device cuda:0
```
</details>
<details>
<summary> 基准测试 </summary>
1. 设置
```shell
pip install -r tools/benchmark/requirements.txt
export model=l # n s m l x
```
<!-- <summary>6. Benchmark </summary> -->
2. 模型 FLOPs、MACs、参数量
```shell
python tools/benchmark/get_info.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml
```
2. TensorRT 延迟
```shell
python tools/benchmark/trt_benchmark.py --COCO_dir path/to/COCO2017 --engine_dir model.engine
```
</details>
<details>
<summary> Voxel51 Fiftyone 可视化 </summary>
1. 设置
```shell
pip install fiftyone
export model=l # n s m l x
```
4. Voxel51 Fiftyone 可视化 ([fiftyone](https://github.com/voxel51/fiftyone))
```shell
python tools/visualization/fiftyone_vis.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml -r model.pth
```
</details>
<details>
<summary> 其他 </summary>
1. 自动恢复Auto Resume训练
```shell
bash reference/safe_training.sh
```
2. 模型权重转换
```shell
python reference/convert_weight.py model.pth
```
</details>
## 图表与可视化
<details>
<summary> FDR 和 GO-LSD </summary>
D-FINE与FDR概览。概率分布作为更细粒度的中间表征通过解码器层以残差方式进行迭代优化。应用非均匀加权函数以实现更精细的定位。
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/fdr-1.jpg" alt="细粒度分布优化过程" width="1000"> </p>
GO-LSD流程概览。通过DDF损失函数和解耦加权策略将最终层分布中的定位知识蒸馏到前面的层中。
<p align="center"> <img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/go_lsd-1.jpg" alt="GO-LSD流程" width="1000"> </p>
</details>
<details open>
<summary> 分布可视化 </summary>
FDR在检测场景中的可视化包括初始和优化后的边界框以及未加权和加权的分布图。
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/merged_image.jpg" width="1000">
</p>
</details>
<details>
<summary> 困难场景 </summary>
以下可视化展示了D-FINE在各种复杂检测场景中的预测结果。这些场景包括遮挡、低光条件、运动模糊、景深效果和密集场景。尽管面临这些挑战D-FINE依然能够生成准确的定位结果。
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/hard_case-1.jpg" alt="D-FINE在挑战性场景中的预测" width="1000">
</p>
</details>
<!-- <table><tr>
<td><img src=https://raw.githubusercontent.com/Peterande/storage/master/figs/merged_image.jpg border=0 width=1000></td>
</tr></table> -->
## 引用
如果你在工作中使用了 `D-FINE` 或其方法,请引用以下 BibTeX 条目:
<details open>
<summary> bibtex </summary>
```latex
@misc{peng2024dfine,
title={D-FINE: Redefine Regression Task in DETRs as Fine-grained Distribution Refinement},
author={Yansong Peng and Hebei Li and Peixi Wu and Yueyi Zhang and Xiaoyan Sun and Feng Wu},
year={2024},
eprint={2410.13842},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
```
</details>
## 致谢
我们的工作基于 [RT-DETR](https://github.com/lyuwenyu/RT-DETR)。
感谢 [RT-DETR](https://github.com/lyuwenyu/RT-DETR), [GFocal](https://github.com/implus/GFocal), [LD](https://github.com/HikariTJU/LD), 和 [YOLOv9](https://github.com/WongKinYiu/yolov9) 的启发。
✨ 欢迎贡献并在有任何问题时联系我! ✨

View File

@ -0,0 +1,698 @@
<!--# [D-FINE: Redefine Regression Task in DETRs as Fine-grained Distribution Refinement](https://arxiv.org/abs/xxxxxx) -->
[English](README.md) | [简体中文](README_cn.md) | 日本語 | [English Blog](src/zoo/dfine/blog.md) | [中文博客](src/zoo/dfine/blog_cn.md)
<h2 align="center">
D-FINE: Redefine Regression Task of DETRs as Fine&#8209;grained&nbsp;Distribution&nbsp;Refinement
</h2>
<p align="center">
<a href="https://github.com/Peterande/D-FINE/blob/master/LICENSE">
<img alt="license" src="https://img.shields.io/badge/LICENSE-Apache%202.0-blue">
</a>
<a href="https://github.com/Peterande/D-FINE/pulls">
<img alt="prs" src="https://img.shields.io/github/issues-pr/Peterande/D-FINE">
</a>
<a href="https://github.com/Peterande/D-FINE/issues">
<img alt="issues" src="https://img.shields.io/github/issues/Peterande/D-FINE?color=olive">
</a>
<a href="https://arxiv.org/abs/2410.13842">
<img alt="arXiv" src="https://img.shields.io/badge/arXiv-2410.13842-red">
</a>
<!-- <a href="mailto: pengyansong@mail.ustc.edu.cn">
<img alt="email" src="https://img.shields.io/badge/contact_me-email-yellow">
</a> -->
<a href="https://results.pre-commit.ci/latest/github/Peterande/D-FINE/master">
<img alt="pre-commit.ci status" src="https://results.pre-commit.ci/badge/github/Peterande/D-FINE/master.svg">
</a>
<a href="https://github.com/Peterande/D-FINE">
<img alt="stars" src="https://img.shields.io/github/stars/Peterande/D-FINE">
</a>
</p>
<p align="center">
📄 これは論文の公式実装です:
<br>
<a href="https://arxiv.org/abs/2410.13842">D-FINE: Redefine Regression Task of DETRs as Fine-grained Distribution Refinement</a>
</p>
<p align="center">
D-FINE: DETRの回帰タスクを細粒度分布最適化として再定義
</p>
<p align="center">
Yansong Peng, Hebei Li, Peixi Wu, Yueyi Zhang, Xiaoyan Sun, and Feng Wu
</p>
<p align="center">
中国科学技術大学
</p>
<p align="center">
<a href="https://paperswithcode.com/sota/real-time-object-detection-on-coco?p=d-fine-redefine-regression-task-in-detrs-as">
<img alt="sota" src="https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/d-fine-redefine-regression-task-in-detrs-as/real-time-object-detection-on-coco">
</a>
</p>
<!-- <table><tr>
<td><img src=https://github.com/Peterande/storage/blob/master/latency.png border=0 width=333></td>
<td><img src=https://github.com/Peterande/storage/blob/master/params.png border=0 width=333></td>
<td><img src=https://github.com/Peterande/storage/blob/master/flops.png border=0 width=333></td>
</tr></table> -->
<p align="center">
<strong>もしD-FINEが気に入ったら、ぜひ⭐をくださいあなたのサポートが私たちのモチベーションになります</strong>
</p>
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/stats_padded.png" width="1000">
</p>
D-FINEは、DETRの境界ボックス回帰タスクを細粒度分布最適化FDRとして再定義し、グローバル最適な位置特定自己蒸留GO-LSDを導入することで、追加の推論およびトレーニングコストを増やすことなく、優れたパフォーマンスを実現する強力なリアルタイムオブジェクト検出器です。
<details open>
<summary> ビデオ </summary>
D-FINEとYOLO11を使用して、[YouTube](https://www.youtube.com/watch?v=CfhEWj9sd9A)の複雑な街並みのビデオでオブジェクト検出を行いました。逆光、モーションブラー、密集した群衆などの厳しい条件にもかかわらず、D-FINE-Xはほぼすべてのターゲットを検出し、バックパック、自転車、信号機などの微妙な小さなオブジェクトも含まれます。その信頼スコアとぼやけたエッジの位置特定精度はYOLO11よりもはるかに高いです。
<!-- We use D-FINE and YOLO11 on a street scene video from [YouTube](https://www.youtube.com/watch?v=CfhEWj9sd9A). Despite challenges like backlighting, motion blur, and dense crowds, D-FINE-X outperforms YOLO11x, detecting more objects with higher confidence and better precision. -->
https://github.com/user-attachments/assets/e5933d8e-3c8a-400e-870b-4e452f5321d9
</details>
## 🚀 更新情報
- [x] **\[2024.10.18\]** D-FINEシリーズをリリース。
- [x] **\[2024.10.25\]** カスタムデータセットの微調整設定を追加 ([#7](https://github.com/Peterande/D-FINE/issues/7))。
- [x] **\[2024.10.30\]** D-FINE-L (E25) 事前トレーニングモデルを更新し、パフォーマンスが2.0%向上。
- [x] **\[2024.11.07\]** **D-FINE-N** をリリース, COCO で 42.8% の AP<sup>val</sup> を達成 @ 472 FPS<sup>T4</sup>!
## モデルズー
### COCO
| モデル | データセット | AP<sup>val</sup> | パラメータ数 | レイテンシ | GFLOPs | config | checkpoint | logs |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
**D&#8209;FINE&#8209;N** | COCO | **42.8** | 4M | 2.12ms | 7 | [yml](./configs/dfine/dfine_hgnetv2_n_coco.yml) | [42.8](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_n_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_n_coco_log.txt)
**D&#8209;FINE&#8209;S** | COCO | **48.5** | 10M | 3.49ms | 25 | [yml](./configs/dfine/dfine_hgnetv2_s_coco.yml) | [48.5](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_s_coco_log.txt)
**D&#8209;FINE&#8209;M** | COCO | **52.3** | 19M | 5.62ms | 57 | [yml](./configs/dfine/dfine_hgnetv2_m_coco.yml) | [52.3](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_m_coco_log.txt)
**D&#8209;FINE&#8209;L** | COCO | **54.0** | 31M | 8.07ms | 91 | [yml](./configs/dfine/dfine_hgnetv2_l_coco.yml) | [54.0](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_l_coco_log.txt)
**D&#8209;FINE&#8209;X** | COCO | **55.8** | 62M | 12.89ms | 202 | [yml](./configs/dfine/dfine_hgnetv2_x_coco.yml) | [55.8](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/coco/dfine_x_coco_log.txt)
### Objects365+COCO
| モデル | データセット | AP<sup>val</sup> | パラメータ数 | レイテンシ | GFLOPs | config | checkpoint | logs |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
**D&#8209;FINE&#8209;S** | Objects365+COCO | **50.7** | 10M | 3.49ms | 25 | [yml](./configs/dfine/objects365/dfine_hgnetv2_s_obj2coco.yml) | [50.7](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_obj2coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_s_obj2coco_log.txt)
**D&#8209;FINE&#8209;M** | Objects365+COCO | **55.1** | 19M | 5.62ms | 57 | [yml](./configs/dfine/objects365/dfine_hgnetv2_m_obj2coco.yml) | [55.1](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_obj2coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_m_obj2coco_log.txt)
**D&#8209;FINE&#8209;L** | Objects365+COCO | **57.3** | 31M | 8.07ms | 91 | [yml](./configs/dfine/objects365/dfine_hgnetv2_l_obj2coco.yml) | [57.3](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj2coco_e25.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_l_obj2coco_log_e25.txt)
**D&#8209;FINE&#8209;X** | Objects365+COCO | **59.3** | 62M | 12.89ms | 202 | [yml](./configs/dfine/objects365/dfine_hgnetv2_x_obj2coco.yml) | [59.3](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_obj2coco.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj2coco/dfine_x_obj2coco_log.txt)
**微調整のために Objects365 の事前学習モデルを使用することを強くお勧めします:**
⚠️ 重要なお知らせ:このプリトレインモデルは複雑なシーンの理解に有益ですが、カテゴリが非常に単純な場合、過学習や最適ではない性能につながる可能性がありますので、ご注意ください。
<details> <summary><strong> 🔥 Objects365で事前トレーニングされたモデル最良の汎化性能</strong></summary>
| モデル | データセット | AP<sup>val</sup> | AP<sup>5000</sup> | パラメータ数 | レイテンシ | GFLOPs | config | checkpoint | logs |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
**D&#8209;FINE&#8209;S** | Objects365 | **31.0** | **30.5** | 10M | 3.49ms | 25 | [yml](./configs/dfine/objects365/dfine_hgnetv2_s_obj365.yml) | [30.5](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_s_obj365_log.txt)
**D&#8209;FINE&#8209;M** | Objects365 | **38.6** | **37.4** | 19M | 5.62ms | 57 | [yml](./configs/dfine/objects365/dfine_hgnetv2_m_obj365.yml) | [37.4](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_m_obj365_log.txt)
**D&#8209;FINE&#8209;L** | Objects365 | - | **40.6** | 31M | 8.07ms | 91 | [yml](./configs/dfine/objects365/dfine_hgnetv2_l_obj365.yml) | [40.6](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_l_obj365_log.txt)
**D&#8209;FINE&#8209;L (E25)** | Objects365 | **44.7** | **42.6** | 31M | 8.07ms | 91 | [yml](./configs/dfine/objects365/dfine_hgnetv2_l_obj365.yml) | [42.6](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj365_e25.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_l_obj365_log_e25.txt)
**D&#8209;FINE&#8209;X** | Objects365 | **49.5** | **46.5** | 62M | 12.89ms | 202 | [yml](./configs/dfine/objects365/dfine_hgnetv2_x_obj365.yml) | [46.5](https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_obj365.pth) | [url](https://raw.githubusercontent.com/Peterande/storage/refs/heads/master/logs/obj365/dfine_x_obj365_log.txt)
- **E25**: 再トレーニングし、事前トレーニングを25エポックに延長。
- **AP<sup>val</sup>***Objects365* のフルバリデーションセットで評価されます。
- **AP<sup>5000</sup>***Objects365* 検証セットの最初の5000サンプルで評価されます。
</details>
**注意事項:**
- **AP<sup>val</sup>***MSCOCO val2017* データセットで評価されます。
- **レイテンシ** は単一のT4 GPUで $batch\\_size = 1$, $fp16$, および $TensorRT==10.4.0$ で評価されます。
- **Objects365+COCO***Objects365* で事前トレーニングされた重みを使用して *COCO* で微調整されたモデルを意味します。
## クイックスタート
### セットアップ
```shell
conda create -n dfine python=3.11.9
conda activate dfine
pip install -r requirements.txt
```
### データ準備
<details>
<summary> COCO2017 データセット </summary>
1. [OpenDataLab](https://opendatalab.com/OpenDataLab/COCO_2017) または [COCO](https://cocodataset.org/#download) からCOCO2017をダウンロードします。
1. [coco_detection.yml](./configs/dataset/coco_detection.yml) のパスを修正します。
```yaml
train_dataloader:
img_folder: /data/COCO2017/train2017/
ann_file: /data/COCO2017/annotations/instances_train2017.json
val_dataloader:
img_folder: /data/COCO2017/val2017/
ann_file: /data/COCO2017/annotations/instances_val2017.json
```
</details>
<details>
<summary> Objects365 データセット </summary>
1. [OpenDataLab](https://opendatalab.com/OpenDataLab/Objects365) からObjects365をダウンロードします。
2. ベースディレクトリを設定します:
```shell
export BASE_DIR=/data/Objects365/data
```
3. ダウンロードしたファイルを解凍し、以下のディレクトリ構造に整理します:
```shell
${BASE_DIR}/train
├── images
│ ├── v1
│ │ ├── patch0
│ │ │ ├── 000000000.jpg
│ │ │ ├── 000000001.jpg
│ │ │ └── ... (more images)
│ ├── v2
│ │ ├── patchx
│ │ │ ├── 000000000.jpg
│ │ │ ├── 000000001.jpg
│ │ │ └── ... (more images)
├── zhiyuan_objv2_train.json
```
```shell
${BASE_DIR}/val
├── images
│ ├── v1
│ │ ├── patch0
│ │ │ ├── 000000000.jpg
│ │ │ └── ... (more images)
│ ├── v2
│ │ ├── patchx
│ │ │ ├── 000000000.jpg
│ │ │ └── ... (more images)
├── zhiyuan_objv2_val.json
```
4. 検証セットの画像を保存する新しいディレクトリを作成します:
```shell
mkdir -p ${BASE_DIR}/train/images_from_val
```
5. valディレクトリのv1およびv2フォルダをtrain/images_from_valディレクトリにコピーします
```shell
cp -r ${BASE_DIR}/val/images/v1 ${BASE_DIR}/train/images_from_val/
cp -r ${BASE_DIR}/val/images/v2 ${BASE_DIR}/train/images_from_val/
```
6. remap_obj365.pyを実行して、検証セットの一部をトレーニングセットにマージします。具体的には、このスクリプトはインデックスが5000から800000のサンプルを検証セットからトレーニングセットに移動します。
```shell
python tools/remap_obj365.py --base_dir ${BASE_DIR}
```
7. resize_obj365.pyスクリプトを実行して、データセット内の最大エッジ長が640ピクセルを超える画像をリサイズします。ステップ5で生成された更新されたJSONファイルを使用してサンプルデータを処理します。トレーニングセットと検証セットの両方の画像をリサイズして、一貫性を保ちます。
```shell
python tools/resize_obj365.py --base_dir ${BASE_DIR}
```
8. [obj365_detection.yml](./configs/dataset/obj365_detection.yml) のパスを修正します。
```yaml
train_dataloader:
img_folder: /data/Objects365/data/train
ann_file: /data/Objects365/data/train/new_zhiyuan_objv2_train_resized.json
val_dataloader:
img_folder: /data/Objects365/data/val/
ann_file: /data/Objects365/data/val/new_zhiyuan_objv2_val_resized.json
```
</details>
<details>
<summary>CrowdHuman</summary>
こちらからCOCOフォーマットのデータセットをダウンロードしてください[リンク](https://aistudio.baidu.com/datasetdetail/231455)
</details>
<details>
<summary>カスタムデータセット</summary>
カスタムデータセットでトレーニングするには、COCO形式で整理する必要があります。以下の手順に従ってデータセットを準備してください
1. **`remap_mscoco_category``False` に設定します**
これにより、カテゴリIDがMSCOCOカテゴリに自動的にマッピングされるのを防ぎます。
```yaml
remap_mscoco_category: False
```
2. **画像を整理します**
データセットディレクトリを以下のように構造化します:
```shell
dataset/
├── images/
│ ├── train/
│ │ ├── image1.jpg
│ │ ├── image2.jpg
│ │ └── ...
│ ├── val/
│ │ ├── image1.jpg
│ │ ├── image2.jpg
│ │ └── ...
└── annotations/
├── instances_train.json
├── instances_val.json
└── ...
```
- **`images/train/`**: すべてのトレーニング画像を含みます。
- **`images/val/`**: すべての検証画像を含みます。
- **`annotations/`**: COCO形式の注釈ファイルを含みます。
3. **注釈をCOCO形式に変換します**
注釈がまだCOCO形式でない場合は、変換する必要があります。以下のPythonスクリプトを参考にするか、既存のツールを利用してください
```python
import json
def convert_to_coco(input_annotations, output_annotations):
# 変換ロジックをここに実装します
pass
if __name__ == "__main__":
convert_to_coco('path/to/your_annotations.json', 'dataset/annotations/instances_train.json')
```
4. **設定ファイルを更新します**
[custom_detection.yml](./configs/dataset/custom_detection.yml) を修正します。
```yaml
task: detection
evaluator:
type: CocoEvaluator
iou_types: ['bbox', ]
num_classes: 777 # データセットのクラス数
remap_mscoco_category: False
train_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/yourdataset/train
ann_file: /data/yourdataset/train/train.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: True
num_workers: 4
drop_last: True
collate_fn:
type: BatchImageCollateFunction
val_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/yourdataset/val
ann_file: /data/yourdataset/val/ann.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: False
num_workers: 4
drop_last: False
collate_fn:
type: BatchImageCollateFunction
```
</details>
## 使用方法
<details open>
<summary> COCO2017 </summary>
<!-- <summary>1. トレーニング </summary> -->
1. モデルを設定します
```shell
export model=l # n s m l x
```
2. トレーニング
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --use-amp --seed=0
```
<!-- <summary>2. テスト </summary> -->
3. テスト
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --test-only -r model.pth
```
<!-- <summary>3. 微調整 </summary> -->
4. 微調整
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --use-amp --seed=0 -t model.pth
```
</details>
<details>
<summary> Objects365からCOCO2017へ </summary>
1. モデルを設定します
```shell
export model=l # n s m l x
```
2. Objects365でトレーニング
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/objects365/dfine_hgnetv2_${model}_obj365.yml --use-amp --seed=0
```
3. COCO2017で微調整
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/objects365/dfine_hgnetv2_${model}_obj2coco.yml --use-amp --seed=0 -t model.pth
```
<!-- <summary>2. テスト </summary> -->
4. テスト
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml --test-only -r model.pth
```
</details>
<details>
<summary> カスタムデータセット </summary>
1. モデルを設定します
```shell
export model=l # n s m l x
```
2. カスタムデータセットでトレーニング
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/dfine_hgnetv2_${model}_custom.yml --use-amp --seed=0
```
<!-- <summary>2. テスト </summary> -->
3. テスト
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/dfine_hgnetv2_${model}_custom.yml --test-only -r model.pth
```
4. カスタムデータセットで微調整
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/objects365/dfine_hgnetv2_${model}_obj2custom.yml --use-amp --seed=0 -t model.pth
```
5. **[オプション]** クラスマッピングを変更します:
Objects365の事前トレーニング済みの重みを使用してカスタムデータセットでトレーニングする場合、例ではデータセットに `'Person'``'Car'` クラスのみが含まれていると仮定しています。特定のタスクに対して収束を早めるために、`src/solver/_solver.py` の `self.obj365_ids` を以下のように変更できます:
```python
self.obj365_ids = [0, 5] # Person, Cars
```
これらをデータセットの対応するクラスに置き換えることができます。Objects365クラスとその対応IDのリスト
https://github.com/Peterande/D-FINE/blob/352a94ece291e26e1957df81277bef00fe88a8e3/src/solver/_solver.py#L330
新しいトレーニングコマンド:
```shell
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/dfine/custom/dfine_hgnetv2_${model}_custom.yml --use-amp --seed=0 -t model.pth
```
ただし、クラスマッピングを変更したくない場合、事前トレーニング済みのObjects365の重みは変更なしでそのまま使用できます。クラスマッピングの変更はオプションであり、特定のタスクに対して収束を早める可能性があります。
</details>
<details>
<summary> バッチサイズのカスタマイズ </summary>
例えば、COCO2017でD-FINE-Lをトレーニングする際にバッチサイズを2倍にしたい場合、以下の手順に従ってください
1. **[dataloader.yml](./configs/dfine/include/dataloader.yml) を修正して `total_batch_size` を増やします**
```yaml
train_dataloader:
total_batch_size: 64 # 以前は32、今は2倍
```
2. **[dfine_hgnetv2_l_coco.yml](./configs/dfine/dfine_hgnetv2_l_coco.yml) を修正します**。以下のように主要なパラメータを調整します:
```yaml
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.000025 # 2倍、線形スケーリング法則
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.0005 # 2倍、線形スケーリング法則
betas: [0.9, 0.999]
weight_decay: 0.0001 # グリッドサーチが必要です
ema: # EMA設定を追加
decay: 0.9998 # 1 - (1 - decay) * 2 によって調整
warmups: 500 # 半分
lr_warmup_scheduler:
warmup_duration: 250 # 半分
```
</details>
<details>
<summary> 入力サイズのカスタマイズ </summary>
COCO2017で **D-FINE-L** を320x320の入力サイズでトレーニングしたい場合、以下の手順に従ってください
1. **[dataloader.yml](./configs/dfine/include/dataloader.yml) を修正します**
```yaml
train_dataloader:
dataset:
transforms:
ops:
- {type: Resize, size: [320, 320], }
collate_fn:
base_size: 320
dataset:
transforms:
ops:
- {type: Resize, size: [320, 320], }
```
2. **[dfine_hgnetv2.yml](./configs/dfine/include/dfine_hgnetv2.yml) を修正します**
```yaml
eval_spatial_size: [320, 320]
```
</details>
## ツール
<details>
<summary> デプロイ </summary>
<!-- <summary>4. onnxのエクスポート </summary> -->
1. セットアップ
```shell
pip install onnx onnxsim
export model=l # n s m l x
```
2. onnxのエクスポート
```shell
python tools/deployment/export_onnx.py --check -c configs/dfine/dfine_hgnetv2_${model}_coco.yml -r model.pth
```
3. [tensorrt](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html) のエクスポート
```shell
trtexec --onnx="model.onnx" --saveEngine="model.engine" --fp16
```
</details>
<details>
<summary> 推論(可視化) </summary>
1. セットアップ
```shell
pip install -r tools/inference/requirements.txt
export model=l # n s m l x
```
<!-- <summary>5. 推論 </summary> -->
2. 推論 (onnxruntime / tensorrt / torch)
現在、画像とビデオの推論がサポートされています。
```shell
python tools/inference/onnx_inf.py --onnx model.onnx --input image.jpg # video.mp4
python tools/inference/trt_inf.py --trt model.engine --input image.jpg
python tools/inference/torch_inf.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml -r model.pth --input image.jpg --device cuda:0
```
</details>
<details>
<summary> ベンチマーク </summary>
1. セットアップ
```shell
pip install -r tools/benchmark/requirements.txt
export model=l # n s m l x
```
<!-- <summary>6. ベンチマーク </summary> -->
2. モデルのFLOPs、MACs、およびパラメータ数
```shell
python tools/benchmark/get_info.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml
```
2. TensorRTのレイテンシ
```shell
python tools/benchmark/trt_benchmark.py --COCO_dir path/to/COCO2017 --engine_dir model.engine
```
</details>
<details>
<summary> Fiftyoneの可視化 </summary>
1. セットアップ
```shell
pip install fiftyone
export model=l # n s m l x
```
4. Voxel51 Fiftyoneの可視化 ([fiftyone](https://github.com/voxel51/fiftyone))
```shell
python tools/visualization/fiftyone_vis.py -c configs/dfine/dfine_hgnetv2_${model}_coco.yml -r model.pth
```
</details>
<details>
<summary> その他 </summary>
1. 自動再開トレーニング
```shell
bash reference/safe_training.sh
```
2. モデルの重みの変換
```shell
python reference/convert_weight.py model.pth
```
</details>
## 図と可視化
<details>
<summary> FDRとGO-LSD </summary>
1. FDRを搭載したD-FINEの概要。より細粒度の中間表現として機能する確率分布は、残差的にデコーダ層によって逐次最適化されます。
不均一な重み付け関数が適用され、より細かい位置特定が可能になります。
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/fdr-1.jpg" alt="細粒度分布最適化プロセス" width="1000">
</p>
2. GO-LSDプロセスの概要。最終層の最適化された分布からの位置特定知識は、デカップリングされた重み付け戦略を使用してDDF損失を通じて前の層に蒸留されます。
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/go_lsd-1.jpg" alt="GO-LSDプロセス" width="1000">
</p>
</details>
<details open>
<summary> 分布 </summary>
初期および最適化された境界ボックスと、未重み付けおよび重み付けされた分布とともに、さまざまな検出シナリオにおけるFDRの可視化。
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/merged_image.jpg" width="1000">
</p>
</details>
<details>
<summary> 難しいケース </summary>
以下の可視化は、さまざまな複雑な検出シナリオにおけるD-FINEの予測を示しています。これらのシナリオには、遮蔽、低光条件、モーションブラー、被写界深度効果、および密集したシーンが含まれます。これらの課題にもかかわらず、D-FINEは一貫して正確な位置特定結果を生成します。
<p align="center">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/hard_case-1.jpg" alt="複雑なシナリオにおけるD-FINEの予測" width="1000">
</p>
</details>
<!-- <div style="display: flex; flex-wrap: wrap; justify-content: center; margin: 0; padding: 0;">
<img src="https://raw.githubusercontent.com/Peterande/storage/master/figs/merged_image.jpg" style="width:99.96%; margin: 0; padding: 0;" />
</div>
<table><tr>
<td><img src=https://raw.githubusercontent.com/Peterande/storage/master/figs/merged_image.jpg border=0 width=1000></td>
</tr></table> -->
## 引用
もし`D-FINE`やその方法をあなたの仕事で使用する場合、以下のBibTeXエントリを引用してください
<details open>
<summary> bibtex </summary>
```latex
@misc{peng2024dfine,
title={D-FINE: Redefine Regression Task in DETRs as Fine-grained Distribution Refinement},
author={Yansong Peng and Hebei Li and Peixi Wu and Yueyi Zhang and Xiaoyan Sun and Feng Wu},
year={2024},
eprint={2410.13842},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
```
</details>
## 謝辞
私たちの仕事は [RT-DETR](https://github.com/lyuwenyu/RT-DETR) に基づいています。
[RT-DETR](https://github.com/lyuwenyu/RT-DETR), [GFocal](https://github.com/implus/GFocal), [LD](https://github.com/HikariTJU/LD), および [YOLOv9](https://github.com/WongKinYiu/yolov9) からのインスピレーションに感謝します。
✨ 貢献を歓迎し、質問があればお気軽にお問い合わせください! ✨

View File

@ -0,0 +1,91 @@
import os
import json
import xml.etree.ElementTree as ET
from tqdm import tqdm
def collect_train_dataset_voctococo():
XML_DIR = r"D:\project_space\01\006\datasets\002\images" # VOC格式xml文件夹路径
IMG_DIR = r"D:\project_space\01\006\datasets\002\images" # 图片文件夹路径
SAVE_PATH = r"D:\project_space\01\006\datasets\002\images\coco.json" # COCO格式json保存路径
CLASSES = ['0'] # 改成你自己的类别
def get_category_id(name):
return CLASSES.index(name) + 1
# 3. 开始转换
image_id = 1
annotation_id = 1
images = []
annotations = []
categories = []
# 4. 构建categories
for idx, name in enumerate(CLASSES):
categories.append({
"id": idx + 1,
"name": name,
"supercategory": "none"
})
xml_files = [f for f in os.listdir(XML_DIR) if f.endswith('.xml')]
for xml_file in tqdm(xml_files):
xml_path = os.path.join(XML_DIR, xml_file)
tree = ET.parse(xml_path)
root = tree.getroot()
filename = root.find('filename').text
img_path = os.path.join(IMG_DIR, filename)
# 获取图片尺寸
size = root.find('size')
width = int(size.find('width').text)
height = int(size.find('height').text)
# images部分
images.append({
"file_name": filename,
"height": height,
"width": width,
"id": image_id
})
# annotations部分
for obj in root.findall('object'):
name = obj.find('name').text
if name not in CLASSES:
continue
category_id = get_category_id(name)
bndbox = obj.find('bndbox')
xmin = int(float(bndbox.find('xmin').text))
ymin = int(float(bndbox.find('ymin').text))
xmax = int(float(bndbox.find('xmax').text))
ymax = int(float(bndbox.find('ymax').text))
w = xmax - xmin
h = ymax - ymin
# coco标注
annotations.append({
"id": annotation_id,
"image_id": image_id,
"category_id": category_id,
"bbox": [xmin, ymin, w, h],
"area": w * h,
"iscrowd": 0,
"segmentation": []
})
annotation_id += 1
image_id += 1
# 5. 整合并保存
coco_format = {
"images": images,
"type": "instances",
"annotations": annotations,
"categories": categories
}
with open(SAVE_PATH, 'w', encoding='utf-8') as f:
json.dump(coco_format, f, ensure_ascii=False, indent=4)
print(f"转换完成,共有 {len(images)} 张图片,{len(annotations)} 个标注。输出: {SAVE_PATH}")
if __name__ == '__main__':
collect_train_dataset_voctococo()

View File

@ -0,0 +1,41 @@
task: detection
evaluator:
type: CocoEvaluator
iou_types: ['bbox', ]
num_classes: 1
remap_mscoco_category: True
train_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: D:\project_space\01\006\datasets\002\images
ann_file: D:\project_space\01\006\datasets\002\images\coco.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: True
num_workers: 2
drop_last: True
collate_fn:
type: BatchImageCollateFunction
val_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: D:\project_space\01\006\datasets\002\images
ann_file: D:\project_space\01\006\datasets\002\images\coco.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: False
num_workers: 2
drop_last: False
collate_fn:
type: BatchImageCollateFunction

View File

@ -0,0 +1,41 @@
task: detection
evaluator:
type: CocoEvaluator
iou_types: ['bbox', ]
num_classes: 1 # your dataset classes
remap_mscoco_category: False
train_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/CrowdHuman/coco/CrowdHuman_train
ann_file: /data/CrowdHuman/coco/Chuman-train.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: True
num_workers: 4
drop_last: True
collate_fn:
type: BatchImageCollateFunction
val_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/CrowdHuman/coco/CrowdHuman_val
ann_file: /data/CrowdHuman/coco/Chuman-val.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: False
num_workers: 4
drop_last: False
collate_fn:
type: BatchImageCollateFunction

View File

@ -0,0 +1,41 @@
task: detection
evaluator:
type: CocoEvaluator
iou_types: ['bbox', ]
num_classes: 777 # your dataset classes
remap_mscoco_category: False
train_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/yourdataset/train
ann_file: /data/yourdataset/train/train.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: True
num_workers: 4
drop_last: True
collate_fn:
type: BatchImageCollateFunction
val_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/yourdataset/val
ann_file: /data/yourdataset/val/val.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: False
num_workers: 4
drop_last: False
collate_fn:
type: BatchImageCollateFunction

View File

@ -0,0 +1,41 @@
task: detection
evaluator:
type: CocoEvaluator
iou_types: ['bbox', ]
num_classes: 366
remap_mscoco_category: False
train_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/Objects365/data/train
ann_file: /data/Objects365/data/train/new_zhiyuan_objv2_train_resized.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: True
num_workers: 4
drop_last: True
collate_fn:
type: BatchImageCollateFunction
val_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: /data/Objects365/data/val/
ann_file: /data/Objects365/data/val/new_zhiyuan_objv2_val_resized.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: False
num_workers: 4
drop_last: False
collate_fn:
type: BatchImageCollateFunction

View File

@ -0,0 +1,40 @@
task: detection
evaluator:
type: CocoEvaluator
iou_types: ['bbox', ]
num_classes: 20
train_dataloader:
type: DataLoader
dataset:
type: VOCDetection
root: ./dataset/voc/
ann_file: trainval.txt
label_file: label_list.txt
transforms:
type: Compose
ops: ~
shuffle: True
num_workers: 4
drop_last: True
collate_fn:
type: BatchImageCollateFunction
val_dataloader:
type: DataLoader
dataset:
type: VOCDetection
root: ./dataset/voc/
ann_file: test.txt
label_file: label_list.txt
transforms:
type: Compose
ops: ~
shuffle: False
num_workers: 4
drop_last: False
collate_fn:
type: BatchImageCollateFunction

View File

@ -0,0 +1,44 @@
__include__: [
'../../dataset/crowdhuman_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_l_crowdhuman
HGNetv2:
name: 'B4'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000125
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# Increase to search for the optimal ema
epochs: 140
train_dataloader:
dataset:
transforms:
policy:
epoch: 120
collate_fn:
stop_epoch: 120
ema_restart_decay: 0.9999
base_size_repeat: 4

View File

@ -0,0 +1,60 @@
__include__: [
'../../dataset/crowdhuman_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_m_crowdhuman
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B2'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
DFINETransformer:
num_layers: 4 # 5 6
eval_idx: -1 # -2 -3
HybridEncoder:
in_channels: [384, 768, 1536]
hidden_dim: 256
depth_mult: 0.67
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.000025
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.000025
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# Increase to search for the optimal ema
epochs: 220
train_dataloader:
dataset:
transforms:
policy:
epoch: 200
collate_fn:
stop_epoch: 200
ema_restart_decay: 0.9999
base_size_repeat: 6

View File

@ -0,0 +1,82 @@
__include__: [
'../../dataset/crowdhuman_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_n_crowdhuman
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B0'
return_idx: [2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
HybridEncoder:
in_channels: [512, 1024]
feat_strides: [16, 32]
# intra
hidden_dim: 128
use_encoder_idx: [1]
dim_feedforward: 512
# cross
expansion: 0.34
depth_mult: 0.5
DFINETransformer:
feat_channels: [128, 128]
feat_strides: [16, 32]
hidden_dim: 128
dim_feedforward: 512
num_levels: 2
num_layers: 3
eval_idx: -1
num_points: [6, 6]
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0004
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.0004
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.0008
betas: [0.9, 0.999]
weight_decay: 0.0001
# Increase to search for the optimal ema
epochs: 220
train_dataloader:
total_batch_size: 128
dataset:
transforms:
policy:
epoch: 200
collate_fn:
stop_epoch: 200
ema_restart_decay: 0.9999
base_size_repeat: ~
val_dataloader:
total_batch_size: 256

View File

@ -0,0 +1,65 @@
__include__: [
'../../dataset/crowdhuman_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_s_crowdhuman
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B0'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
DFINETransformer:
num_layers: 3 # 4 5 6
eval_idx: -1 # -2 -3 -4
HybridEncoder:
in_channels: [256, 512, 1024]
hidden_dim: 256
depth_mult: 0.34
expansion: 0.5
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0002
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.0002
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.0004
betas: [0.9, 0.999]
weight_decay: 0.0001
# Increase to search for the optimal ema
epochs: 220
train_dataloader:
total_batch_size: 64
dataset:
transforms:
policy:
epoch: 200
collate_fn:
stop_epoch: 200
ema_restart_decay: 0.9999
base_size_repeat: 20
val_dataloader:
total_batch_size: 128

View File

@ -0,0 +1,55 @@
__include__: [
'../../dataset/crowdhuman_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_x_crowdhuman
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B5'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
HybridEncoder:
hidden_dim: 384
dim_feedforward: 2048
DFINETransformer:
feat_channels: [384, 384, 384]
reg_scale: 8
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000025
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# Increase to search for the optimal ema
epochs: 140
train_dataloader:
dataset:
transforms:
policy:
epoch: 120
collate_fn:
stop_epoch: 120
ema_restart_decay: 0.9998
base_size_repeat: 3

View File

@ -0,0 +1,44 @@
__include__: [
'../../dataset/custom_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_l_custom
HGNetv2:
name: 'B4'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000125
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# Increase to search for the optimal ema
epochs: 80 # 72 + 2n
train_dataloader:
dataset:
transforms:
policy:
epoch: 72
collate_fn:
stop_epoch: 72
ema_restart_decay: 0.9999
base_size_repeat: 4

View File

@ -0,0 +1,60 @@
__include__: [
'../../dataset/custom_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_m_custom
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B2'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
DFINETransformer:
num_layers: 4 # 5 6
eval_idx: -1 # -2 -3
HybridEncoder:
in_channels: [384, 768, 1536]
hidden_dim: 256
depth_mult: 0.67
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.000025
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.000025
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# Increase to search for the optimal ema
epochs: 132 # 120 + 4n
train_dataloader:
dataset:
transforms:
policy:
epoch: 120
collate_fn:
stop_epoch: 120
ema_restart_decay: 0.9999
base_size_repeat: 6

View File

@ -0,0 +1,82 @@
__include__: [
'../../dataset/custom_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_n_custom
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B0'
return_idx: [2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
HybridEncoder:
in_channels: [512, 1024]
feat_strides: [16, 32]
# intra
hidden_dim: 128
use_encoder_idx: [1]
dim_feedforward: 512
# cross
expansion: 0.34
depth_mult: 0.5
DFINETransformer:
feat_channels: [128, 128]
feat_strides: [16, 32]
hidden_dim: 128
dim_feedforward: 512
num_levels: 2
num_layers: 3
eval_idx: -1
num_points: [6, 6]
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0004
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.0004
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.0008
betas: [0.9, 0.999]
weight_decay: 0.0001
# Increase to search for the optimal ema
epochs: 220
train_dataloader:
total_batch_size: 128
dataset:
transforms:
policy:
epoch: 200
collate_fn:
stop_epoch: 200
ema_restart_decay: 0.9999
base_size_repeat: ~
val_dataloader:
total_batch_size: 256

View File

@ -0,0 +1,65 @@
__include__: [
'../../dataset/custom_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_s_custom
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B0'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
DFINETransformer:
num_layers: 3 # 4 5 6
eval_idx: -1 # -2 -3 -4
HybridEncoder:
in_channels: [256, 512, 1024]
hidden_dim: 256
depth_mult: 0.34
expansion: 0.5
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0002
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.0002
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.0004
betas: [0.9, 0.999]
weight_decay: 0.0001
# Increase to search for the optimal ema
epochs: 220
train_dataloader:
total_batch_size: 64
dataset:
transforms:
policy:
epoch: 200
collate_fn:
stop_epoch: 200
ema_restart_decay: 0.9999
base_size_repeat: 20
val_dataloader:
total_batch_size: 128

View File

@ -0,0 +1,55 @@
__include__: [
'../../dataset/custom_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_x_custom
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B5'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
HybridEncoder:
hidden_dim: 384
dim_feedforward: 2048
DFINETransformer:
feat_channels: [384, 384, 384]
reg_scale: 8
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000025
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# Increase to search for the optimal ema
epochs: 80 # 72 + 2n
train_dataloader:
dataset:
transforms:
policy:
epoch: 72
collate_fn:
stop_epoch: 72
ema_restart_decay: 0.9998
base_size_repeat: 3

View File

@ -0,0 +1,53 @@
__include__: [
'../../../dataset/custom_detection.yml',
'../../../runtime.yml',
'../../include/dataloader.yml',
'../../include/optimizer.yml',
'../../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_l_obj2custom
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B4'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
pretrained: False
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000125
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
epochs: 36 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 30
collate_fn:
stop_epoch: 30
ema_restart_decay: 0.9999
base_size_repeat: 4
ema:
warmups: 0
lr_warmup_scheduler:
warmup_duration: 0

View File

@ -0,0 +1,66 @@
__include__: [
'../../../dataset/custom_detection.yml',
'../../../runtime.yml',
'../../include/dataloader.yml',
'../../include/optimizer.yml',
'../../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_m_obj2custom
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B2'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
pretrained: False
DFINETransformer:
num_layers: 4 # 5 6
eval_idx: -1 # -2 -3
HybridEncoder:
in_channels: [384, 768, 1536]
hidden_dim: 256
depth_mult: 0.67
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.000025
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.000025
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
epochs: 56 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 48
collate_fn:
stop_epoch: 48
ema_restart_decay: 0.9999
base_size_repeat: 6
ema:
warmups: 0
lr_warmup_scheduler:
warmup_duration: 0

View File

@ -0,0 +1,67 @@
__include__: [
'../../../dataset/custom_detection.yml',
'../../../runtime.yml',
'../../include/dataloader.yml',
'../../include/optimizer.yml',
'../../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_s_obj2custom
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B0'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
pretrained: False
DFINETransformer:
num_layers: 3 # 4 5 6
eval_idx: -1 # -2 -3 -4
HybridEncoder:
in_channels: [256, 512, 1024]
hidden_dim: 256
depth_mult: 0.34
expansion: 0.5
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.000125
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.000125
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
epochs: 64 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 56
collate_fn:
stop_epoch: 56
ema_restart_decay: 0.9999
base_size_repeat: 10
ema:
warmups: 0
lr_warmup_scheduler:
warmup_duration: 0

View File

@ -0,0 +1,62 @@
__include__: [
'../../../dataset/custom_detection.yml',
'../../../runtime.yml',
'../../include/dataloader.yml',
'../../include/optimizer.yml',
'../../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_x_obj2custom
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B5'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
pretrained: False
HybridEncoder:
# intra
hidden_dim: 384
dim_feedforward: 2048
DFINETransformer:
feat_channels: [384, 384, 384]
reg_scale: 8
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000025
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
epochs: 36 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 30
collate_fn:
stop_epoch: 30
ema_restart_decay: 0.9999
base_size_repeat: 3
ema:
warmups: 0
lr_warmup_scheduler:
warmup_duration: 0

View File

@ -0,0 +1,44 @@
__include__: [
'../dataset/coco_detection.yml',
'../runtime.yml',
'./include/dataloader.yml',
'./include/optimizer.yml',
'./include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_l_coco
HGNetv2:
name: 'B4'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000125
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# Increase to search for the optimal ema
epochs: 8 # 72 + 2n
train_dataloader:
dataset:
transforms:
policy:
epoch: 7
collate_fn:
stop_epoch: 7
ema_restart_decay: 0.9999
base_size_repeat: 4

View File

@ -0,0 +1,60 @@
__include__: [
'../dataset/coco_detection.yml',
'../runtime.yml',
'./include/dataloader.yml',
'./include/optimizer.yml',
'./include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_m_coco
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B2'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
DFINETransformer:
num_layers: 4 # 5 6
eval_idx: -1 # -2 -3
HybridEncoder:
in_channels: [384, 768, 1536]
hidden_dim: 256
depth_mult: 0.67
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.00002
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.00002
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.0002
betas: [0.9, 0.999]
weight_decay: 0.0001
# Increase to search for the optimal ema
epochs: 132 # 120 + 4n
train_dataloader:
dataset:
transforms:
policy:
epoch: 120
collate_fn:
stop_epoch: 120
ema_restart_decay: 0.9999
base_size_repeat: 6

View File

@ -0,0 +1,82 @@
__include__: [
'../dataset/coco_detection.yml',
'../runtime.yml',
'./include/dataloader.yml',
'./include/optimizer.yml',
'./include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_n_coco
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B0'
return_idx: [2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
HybridEncoder:
in_channels: [512, 1024]
feat_strides: [16, 32]
# intra
hidden_dim: 128
use_encoder_idx: [1]
dim_feedforward: 512
# cross
expansion: 0.34
depth_mult: 0.5
DFINETransformer:
feat_channels: [128, 128]
feat_strides: [16, 32]
hidden_dim: 128
dim_feedforward: 512
num_levels: 2
num_layers: 3
eval_idx: -1
num_points: [6, 6]
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0004
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.0004
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.0008
betas: [0.9, 0.999]
weight_decay: 0.0001
# Increase to search for the optimal ema
epochs: 160 # 148 + 4n
train_dataloader:
total_batch_size: 128
dataset:
transforms:
policy:
epoch: 148
collate_fn:
stop_epoch: 148
ema_restart_decay: 0.9999
base_size_repeat: ~
val_dataloader:
total_batch_size: 256

View File

@ -0,0 +1,61 @@
__include__: [
'../dataset/coco_detection.yml',
'../runtime.yml',
'./include/dataloader.yml',
'./include/optimizer.yml',
'./include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_s_coco
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B0'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
DFINETransformer:
num_layers: 3 # 4 5 6
eval_idx: -1 # -2 -3 -4
HybridEncoder:
in_channels: [256, 512, 1024]
hidden_dim: 256
depth_mult: 0.34
expansion: 0.5
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0001
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.0001
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.0002
betas: [0.9, 0.999]
weight_decay: 0.0001
# Increase to search for the optimal ema
epochs: 132 # 120 + 4n
train_dataloader:
dataset:
transforms:
policy:
epoch: 120
collate_fn:
stop_epoch: 120
ema_restart_decay: 0.9999
base_size_repeat: 20

View File

@ -0,0 +1,56 @@
__include__: [
'../dataset/coco_detection.yml',
'../runtime.yml',
'./include/dataloader.yml',
'./include/optimizer.yml',
'./include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_x_coco
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B5'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
HybridEncoder:
# intra
hidden_dim: 384
dim_feedforward: 2048
DFINETransformer:
feat_channels: [384, 384, 384]
reg_scale: 8
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000025
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# Increase to search for the optimal ema
epochs: 8 # 72 + 2n
train_dataloader:
dataset:
transforms:
policy:
epoch: 7
collate_fn:
stop_epoch: 7
ema_restart_decay: 0.9998
base_size_repeat: 3

View File

@ -0,0 +1,39 @@
train_dataloader:
dataset:
transforms:
ops:
- {type: RandomPhotometricDistort, p: 0.5}
- {type: RandomZoomOut, fill: 0}
- {type: RandomIoUCrop, p: 0.8}
- {type: SanitizeBoundingBoxes, min_size: 1}
- {type: RandomHorizontalFlip}
- {type: Resize, size: [640, 640], }
- {type: SanitizeBoundingBoxes, min_size: 1}
- {type: ConvertPILImage, dtype: 'float32', scale: True}
- {type: ConvertBoxes, fmt: 'cxcywh', normalize: True}
policy:
name: stop_epoch
epoch: 72 # epoch in [71, ~) stop `ops`
ops: ['RandomPhotometricDistort', 'RandomZoomOut', 'RandomIoUCrop']
collate_fn:
type: BatchImageCollateFunction
base_size: 640
base_size_repeat: 3
stop_epoch: 72 # epoch in [72, ~) stop `multiscales`
shuffle: True
total_batch_size: 32 # total batch size equals to 32 (4 * 8)
num_workers: 4
val_dataloader:
dataset:
transforms:
ops:
- {type: Resize, size: [640, 640], }
- {type: ConvertPILImage, dtype: 'float32', scale: True}
shuffle: False
total_batch_size: 64
num_workers: 4

View File

@ -0,0 +1,82 @@
task: detection
model: DFINE
criterion: DFINECriterion
postprocessor: DFINEPostProcessor
use_focal_loss: True
eval_spatial_size: [640, 640] # h w
DFINE:
backbone: HGNetv2
encoder: HybridEncoder
decoder: DFINETransformer
HGNetv2:
pretrained: True
local_model_dir: weight/hgnetv2/
HybridEncoder:
in_channels: [512, 1024, 2048]
feat_strides: [8, 16, 32]
# intra
hidden_dim: 256
use_encoder_idx: [2]
num_encoder_layers: 1
nhead: 8
dim_feedforward: 1024
dropout: 0.
enc_act: 'gelu'
# cross
expansion: 1.0
depth_mult: 1
act: 'silu'
DFINETransformer:
feat_channels: [256, 256, 256]
feat_strides: [8, 16, 32]
hidden_dim: 256
num_levels: 3
num_layers: 6
eval_idx: -1
num_queries: 300
num_denoising: 100
label_noise_ratio: 0.5
box_noise_scale: 1.0
# NEW
reg_max: 32
reg_scale: 4
# Auxiliary decoder layers dimension scaling
# "eg. If num_layers: 6 eval_idx: -4,
# then layer 3, 4, 5 are auxiliary decoder layers."
layer_scale: 1 # 2
num_points: [3, 6, 3] # [4, 4, 4] [3, 6, 3]
cross_attn_method: default # default, discrete
query_select_method: default # default, agnostic
DFINEPostProcessor:
num_top_queries: 300
DFINECriterion:
weight_dict: {loss_vfl: 1, loss_bbox: 5, loss_giou: 2, loss_fgl: 0.15, loss_ddf: 1.5}
losses: ['vfl', 'boxes', 'local']
alpha: 0.75
gamma: 2.0
reg_max: 32
matcher:
type: HungarianMatcher
weight_dict: {cost_class: 2, cost_bbox: 5, cost_giou: 2}
alpha: 0.25
gamma: 2.0

View File

@ -0,0 +1,36 @@
use_amp: True
use_ema: True
ema:
type: ModelEMA
decay: 0.9999
warmups: 1000
start: 0
epochs: 72
clip_max_norm: 0.1
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm).*$'
lr: 0.0000125
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
lr_scheduler:
type: MultiStepLR
milestones: [500]
gamma: 0.1
lr_warmup_scheduler:
type: LinearWarmup
warmup_duration: 500

View File

@ -0,0 +1,52 @@
__include__: [
'../../dataset/coco_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_l_obj2coco
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B4'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000125
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
epochs: 36 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 30
collate_fn:
stop_epoch: 30
ema_restart_decay: 0.9999
base_size_repeat: 4
ema:
warmups: 0
lr_warmup_scheduler:
warmup_duration: 0

View File

@ -0,0 +1,49 @@
__include__: [
'../../dataset/obj365_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_l_obj365
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B4'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000125
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# weight_decay: 0.00005 # Faster convergence (optional)
epochs: 24 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 500
collate_fn:
stop_epoch: 500
base_size_repeat: 4
checkpoint_freq: 1
print_freq: 1000

View File

@ -0,0 +1,65 @@
__include__: [
'../../dataset/coco_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_m_obj2coco
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B2'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
DFINETransformer:
num_layers: 4 # 5 6
eval_idx: -1 # -2 -3
HybridEncoder:
in_channels: [384, 768, 1536]
hidden_dim: 256
depth_mult: 0.67
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.000025
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.000025
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
epochs: 56 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 48
collate_fn:
stop_epoch: 48
ema_restart_decay: 0.9999
base_size_repeat: 6
ema:
warmups: 0
lr_warmup_scheduler:
warmup_duration: 0

View File

@ -0,0 +1,62 @@
__include__: [
'../../dataset/obj365_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: .output/dfine_hgnetv2_s_obj365
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B2'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
DFINETransformer:
num_layers: 4 # 5 6
eval_idx: -1 # -2 -3
HybridEncoder:
in_channels: [384, 768, 1536]
hidden_dim: 256
depth_mult: 0.67
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.000025
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.000025
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# weight_decay: 0.00005 # Faster convergence (optional)
epochs: 36 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 500
collate_fn:
stop_epoch: 500
base_size_repeat: 6
checkpoint_freq: 1
print_freq: 1000

View File

@ -0,0 +1,88 @@
__include__: [
'../../dataset/coco_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_n_obj2coco
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B0'
return_idx: [2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
HybridEncoder:
in_channels: [512, 1024]
feat_strides: [16, 32]
# intra
hidden_dim: 128
use_encoder_idx: [1]
dim_feedforward: 512
# cross
expansion: 0.34
depth_mult: 0.5
DFINETransformer:
feat_channels: [128, 128]
feat_strides: [16, 32]
hidden_dim: 128
dim_feedforward: 512
num_levels: 2
num_layers: 3
eval_idx: -1
num_points: [6, 6]
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0004
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.0004
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.0008
betas: [0.9, 0.999]
weight_decay: 0.0001
epochs: 64 # Early stop
train_dataloader:
total_batch_size: 128
dataset:
transforms:
policy:
epoch: 56
collate_fn:
stop_epoch: 56
ema_restart_decay: 0.9999
base_size_repeat: ~
ema:
warmups: 0
lr_warmup_scheduler:
warmup_duration: 0
val_dataloader:
total_batch_size: 256

View File

@ -0,0 +1,84 @@
__include__: [
'../../dataset/obj365_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_n_obj365
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B0'
return_idx: [2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
HybridEncoder:
in_channels: [512, 1024]
feat_strides: [16, 32]
# intra
hidden_dim: 128
use_encoder_idx: [1]
dim_feedforward: 512
# cross
expansion: 0.34
depth_mult: 0.5
DFINETransformer:
feat_channels: [128, 128]
feat_strides: [16, 32]
hidden_dim: 128
dim_feedforward: 512
num_levels: 2
num_layers: 3
eval_idx: -1
num_points: [6, 6]
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0004
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.0004
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.0008
betas: [0.9, 0.999]
weight_decay: 0.0001
epochs: 48 # Early stop
train_dataloader:
total_batch_size: 128
dataset:
transforms:
policy:
epoch: 500
collate_fn:
stop_epoch: 500
base_size_repeat: ~
checkpoint_freq: 1
print_freq: 500
val_dataloader:
total_batch_size: 256

View File

@ -0,0 +1,66 @@
__include__: [
'../../dataset/coco_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_s_obj2coco
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B0'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
DFINETransformer:
num_layers: 3 # 4 5 6
eval_idx: -1 # -2 -3 -4
HybridEncoder:
in_channels: [256, 512, 1024]
hidden_dim: 256
depth_mult: 0.34
expansion: 0.5
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.000125
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.000125
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
epochs: 64 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 56
collate_fn:
stop_epoch: 56
ema_restart_decay: 0.9999
base_size_repeat: 10
ema:
warmups: 0
lr_warmup_scheduler:
warmup_duration: 0

View File

@ -0,0 +1,63 @@
__include__: [
'../../dataset/obj365_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_s_obj365
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B0'
return_idx: [1, 2, 3]
freeze_at: -1
freeze_norm: False
use_lab: True
DFINETransformer:
num_layers: 3 # 4 5 6
eval_idx: -1 # -2 -3 -4
HybridEncoder:
in_channels: [256, 512, 1024]
hidden_dim: 256
depth_mult: 0.34
expansion: 0.5
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.000125
-
params: '^(?=.*backbone)(?=.*norm|bn).*$'
lr: 0.000125
weight_decay: 0.
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# weight_decay: 0.00005 # Faster convergence (optional)
epochs: 36 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 500
collate_fn:
stop_epoch: 500
base_size_repeat: 20
checkpoint_freq: 1
print_freq: 1000

View File

@ -0,0 +1,61 @@
__include__: [
'../../dataset/coco_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_x_obj2coco
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B5'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
HybridEncoder:
# intra
hidden_dim: 384
dim_feedforward: 2048
DFINETransformer:
feat_channels: [384, 384, 384]
reg_scale: 8
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000025
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
epochs: 36 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 30
collate_fn:
stop_epoch: 30
ema_restart_decay: 0.9999
base_size_repeat: 3
ema:
warmups: 0
lr_warmup_scheduler:
warmup_duration: 0

View File

@ -0,0 +1,58 @@
__include__: [
'../../dataset/obj365_detection.yml',
'../../runtime.yml',
'../include/dataloader.yml',
'../include/optimizer.yml',
'../include/dfine_hgnetv2.yml',
]
output_dir: ./output/dfine_hgnetv2_x_obj365
DFINE:
backbone: HGNetv2
HGNetv2:
name: 'B5'
return_idx: [1, 2, 3]
freeze_stem_only: True
freeze_at: 0
freeze_norm: True
HybridEncoder:
# intra
hidden_dim: 384
dim_feedforward: 2048
DFINETransformer:
feat_channels: [384, 384, 384]
reg_scale: 8
optimizer:
type: AdamW
params:
-
params: '^(?=.*backbone)(?!.*norm|bn).*$'
lr: 0.0000025
-
params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
weight_decay: 0.
lr: 0.00025
betas: [0.9, 0.999]
weight_decay: 0.000125
# weight_decay: 0.00005 # Faster convergence (optional)
epochs: 24 # Early stop
train_dataloader:
dataset:
transforms:
policy:
epoch: 500
collate_fn:
stop_epoch: 500
base_size_repeat: 3
checkpoint_freq: 1
print_freq: 1000

View File

@ -0,0 +1,24 @@
print_freq: 100
output_dir: './logs'
checkpoint_freq: 12
sync_bn: True
find_unused_parameters: False
use_amp: False
scaler:
type: GradScaler
enabled: True
use_ema: False
ema:
type: ModelEMA
decay: 0.9999
warmups: 1000
use_wandb: False
project_name: D-FINE # for wandb
exp_name: baseline # wandb experiment name

View File

@ -0,0 +1,30 @@
import argparse
import os
import torch
def save_only_ema_weights(checkpoint_file):
"""Extract and save only the EMA weights."""
checkpoint = torch.load(checkpoint_file, map_location="cpu")
weights = {}
if "ema" in checkpoint:
weights["model"] = checkpoint["ema"]["module"]
else:
raise ValueError("The checkpoint does not contain 'ema'.")
dir_name, base_name = os.path.split(checkpoint_file)
name, ext = os.path.splitext(base_name)
output_file = os.path.join(dir_name, f"{name}_converted{ext}")
torch.save(weights, output_file)
print(f"EMA weights saved to {output_file}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Extract and save only EMA weights.")
parser.add_argument("checkpoint_file", type=str, help="Path to the input checkpoint file.")
args = parser.parse_args()
save_only_ema_weights(args.checkpoint_file)

View File

@ -0,0 +1,97 @@
#!/bin/bash
# Function to display the menu for selecting model size
select_model_size() {
echo "Select model size:"
select size in s m l x; do
case $size in
s|m|l|x)
echo "You selected model size: $size"
MODEL_SIZE=$size
break
;;
*)
echo "Invalid selection. Please try again."
;;
esac
done
}
# Function to display the menu for selecting task
select_task() {
echo "Select task:"
select task in obj365 obj2coco coco; do
case $task in
obj365|obj2coco|coco)
echo "You selected task: $task"
TASK=$task
break
;;
*)
echo "Invalid selection. Please try again."
;;
esac
done
}
# Function to ask if the user wants to save logs to a txt file
ask_save_logs() {
while true; do
read -p "Do you want to save logs to a txt file? (y/n): " yn
case $yn in
[Yy]* )
SAVE_LOGS=true
break
;;
[Nn]* )
SAVE_LOGS=false
break
;;
* ) echo "Please answer yes or no.";;
esac
done
}
# Call the functions to let the user select
select_model_size
select_task
ask_save_logs
# Set config file and output directory based on selection
if [ "$TASK" = "coco" ]; then
CONFIG_FILE="configs/dfine/dfine_hgnetv2_${MODEL_SIZE}_${TASK}.yml"
else
CONFIG_FILE="configs/dfine/objects365/dfine_hgnetv2_${MODEL_SIZE}_${TASK}.yml"
fi
OUTPUT_DIR="output/${MODEL_SIZE}_${TASK}"
# Construct the training command
TRAIN_CMD="CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c $CONFIG_FILE --use-amp --seed=0 --output-dir $OUTPUT_DIR"
# Append log redirection if SAVE_LOGS is true
if [ "$SAVE_LOGS" = true ]; then
LOG_FILE="${MODEL_SIZE}_${TASK}.txt"
TRAIN_CMD="$TRAIN_CMD &> \"$LOG_FILE\" 2>&1 &"
else
TRAIN_CMD="$TRAIN_CMD &"
fi
# Run the training command
eval $TRAIN_CMD
if [ $? -ne 0 ]; then
echo "First training failed, restarting with resume option..."
while true; do
RESUME_CMD="CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c $CONFIG_FILE --use-amp --seed=0 --output-dir $OUTPUT_DIR -r ${OUTPUT_DIR}/last.pth"
if [ "$SAVE_LOGS" = true ]; then
LOG_FILE="${MODEL_SIZE}_${TASK}_2.txt"
RESUME_CMD="$RESUME_CMD &> \"$LOG_FILE\" 2>&1 &"
else
RESUME_CMD="$RESUME_CMD &"
fi
eval $RESUME_CMD
if [ $? -eq 0 ]; then
break
fi
done
fi

View File

@ -0,0 +1,9 @@
torch>=2.0.1
torchvision>=0.15.2
faster-coco-eval>=1.6.6
PyYAML
tensorboard
scipy
calflops
transformers
loguru

View File

@ -0,0 +1,6 @@
"""
Copyright (c) 2024 The D-FINE Authors. All Rights Reserved.
"""
# for register purpose
from . import data, nn, optim, zoo

View File

@ -0,0 +1,9 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from ._config import BaseConfig
from .workspace import GLOBAL_CONFIG, create, register
from .yaml_config import YAMLConfig
from .yaml_utils import *

View File

@ -0,0 +1,299 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from pathlib import Path
from typing import Callable, Dict, List
import torch
import torch.nn as nn
from torch.cuda.amp.grad_scaler import GradScaler
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LRScheduler
from torch.utils.data import DataLoader, Dataset
from torch.utils.tensorboard import SummaryWriter
__all__ = [
"BaseConfig",
]
class BaseConfig(object):
# TODO property
def __init__(self) -> None:
super().__init__()
self.task: str = None
# instance / function
self._model: nn.Module = None
self._postprocessor: nn.Module = None
self._criterion: nn.Module = None
self._optimizer: Optimizer = None
self._lr_scheduler: LRScheduler = None
self._lr_warmup_scheduler: LRScheduler = None
self._train_dataloader: DataLoader = None
self._val_dataloader: DataLoader = None
self._ema: nn.Module = None
self._scaler: GradScaler = None
self._train_dataset: Dataset = None
self._val_dataset: Dataset = None
self._collate_fn: Callable = None
self._evaluator: Callable[[nn.Module, DataLoader, str],] = None
self._writer: SummaryWriter = None
# dataset
self.num_workers: int = 0
self.batch_size: int = None
self._train_batch_size: int = None
self._val_batch_size: int = None
self._train_shuffle: bool = None
self._val_shuffle: bool = None
# runtime
self.resume: str = None
self.tuning: str = None
self.epochs: int = None
self.last_epoch: int = -1
self.use_amp: bool = False
self.use_ema: bool = False
self.ema_decay: float = 0.9999
self.ema_warmups: int = 2000
self.sync_bn: bool = False
self.clip_max_norm: float = 0.0
self.find_unused_parameters: bool = None
self.seed: int = None
self.print_freq: int = None
self.checkpoint_freq: int = 1
self.output_dir: str = None
self.summary_dir: str = None
self.device: str = ""
@property
def model(self) -> nn.Module:
return self._model
@model.setter
def model(self, m):
assert isinstance(m, nn.Module), f"{type(m)} != nn.Module, please check your model class"
self._model = m
@property
def postprocessor(self) -> nn.Module:
return self._postprocessor
@postprocessor.setter
def postprocessor(self, m):
assert isinstance(m, nn.Module), f"{type(m)} != nn.Module, please check your model class"
self._postprocessor = m
@property
def criterion(self) -> nn.Module:
return self._criterion
@criterion.setter
def criterion(self, m):
assert isinstance(m, nn.Module), f"{type(m)} != nn.Module, please check your model class"
self._criterion = m
@property
def optimizer(self) -> Optimizer:
return self._optimizer
@optimizer.setter
def optimizer(self, m):
assert isinstance(
m, Optimizer
), f"{type(m)} != optim.Optimizer, please check your model class"
self._optimizer = m
@property
def lr_scheduler(self) -> LRScheduler:
return self._lr_scheduler
@lr_scheduler.setter
def lr_scheduler(self, m):
assert isinstance(
m, LRScheduler
), f"{type(m)} != LRScheduler, please check your model class"
self._lr_scheduler = m
@property
def lr_warmup_scheduler(self) -> LRScheduler:
return self._lr_warmup_scheduler
@lr_warmup_scheduler.setter
def lr_warmup_scheduler(self, m):
self._lr_warmup_scheduler = m
@property
def train_dataloader(self) -> DataLoader:
if self._train_dataloader is None and self.train_dataset is not None:
loader = DataLoader(
self.train_dataset,
batch_size=self.train_batch_size,
num_workers=self.num_workers,
collate_fn=self.collate_fn,
shuffle=self.train_shuffle,
)
loader.shuffle = self.train_shuffle
self._train_dataloader = loader
return self._train_dataloader
@train_dataloader.setter
def train_dataloader(self, loader):
self._train_dataloader = loader
@property
def val_dataloader(self) -> DataLoader:
if self._val_dataloader is None and self.val_dataset is not None:
loader = DataLoader(
self.val_dataset,
batch_size=self.val_batch_size,
num_workers=self.num_workers,
drop_last=False,
collate_fn=self.collate_fn,
shuffle=self.val_shuffle,
persistent_workers=True,
)
loader.shuffle = self.val_shuffle
self._val_dataloader = loader
return self._val_dataloader
@val_dataloader.setter
def val_dataloader(self, loader):
self._val_dataloader = loader
@property
def ema(self) -> nn.Module:
if self._ema is None and self.use_ema and self.model is not None:
from ..optim import ModelEMA
self._ema = ModelEMA(self.model, self.ema_decay, self.ema_warmups)
return self._ema
@ema.setter
def ema(self, obj):
self._ema = obj
@property
def scaler(self) -> GradScaler:
if self._scaler is None and self.use_amp and torch.cuda.is_available():
self._scaler = GradScaler()
return self._scaler
@scaler.setter
def scaler(self, obj: GradScaler):
self._scaler = obj
@property
def val_shuffle(self) -> bool:
if self._val_shuffle is None:
print("warning: set default val_shuffle=False")
return False
return self._val_shuffle
@val_shuffle.setter
def val_shuffle(self, shuffle):
assert isinstance(shuffle, bool), "shuffle must be bool"
self._val_shuffle = shuffle
@property
def train_shuffle(self) -> bool:
if self._train_shuffle is None:
print("warning: set default train_shuffle=True")
return True
return self._train_shuffle
@train_shuffle.setter
def train_shuffle(self, shuffle):
assert isinstance(shuffle, bool), "shuffle must be bool"
self._train_shuffle = shuffle
@property
def train_batch_size(self) -> int:
if self._train_batch_size is None and isinstance(self.batch_size, int):
print(f"warning: set train_batch_size=batch_size={self.batch_size}")
return self.batch_size
return self._train_batch_size
@train_batch_size.setter
def train_batch_size(self, batch_size):
assert isinstance(batch_size, int), "batch_size must be int"
self._train_batch_size = batch_size
@property
def val_batch_size(self) -> int:
if self._val_batch_size is None:
print(f"warning: set val_batch_size=batch_size={self.batch_size}")
return self.batch_size
return self._val_batch_size
@val_batch_size.setter
def val_batch_size(self, batch_size):
assert isinstance(batch_size, int), "batch_size must be int"
self._val_batch_size = batch_size
@property
def train_dataset(self) -> Dataset:
return self._train_dataset
@train_dataset.setter
def train_dataset(self, dataset):
assert isinstance(dataset, Dataset), f"{type(dataset)} must be Dataset"
self._train_dataset = dataset
@property
def val_dataset(self) -> Dataset:
return self._val_dataset
@val_dataset.setter
def val_dataset(self, dataset):
assert isinstance(dataset, Dataset), f"{type(dataset)} must be Dataset"
self._val_dataset = dataset
@property
def collate_fn(self) -> Callable:
return self._collate_fn
@collate_fn.setter
def collate_fn(self, fn):
assert isinstance(fn, Callable), f"{type(fn)} must be Callable"
self._collate_fn = fn
@property
def evaluator(self) -> Callable:
return self._evaluator
@evaluator.setter
def evaluator(self, fn):
assert isinstance(fn, Callable), f"{type(fn)} must be Callable"
self._evaluator = fn
@property
def writer(self) -> SummaryWriter:
if self._writer is None:
if self.summary_dir:
self._writer = SummaryWriter(self.summary_dir)
elif self.output_dir:
self._writer = SummaryWriter(Path(self.output_dir) / "summary")
return self._writer
@writer.setter
def writer(self, m):
assert isinstance(m, SummaryWriter), f"{type(m)} must be SummaryWriter"
self._writer = m
def __repr__(self):
s = ""
for k, v in self.__dict__.items():
if not k.startswith("_"):
s += f"{k}: {v}\n"
return s

View File

@ -0,0 +1,178 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import functools
import importlib
import inspect
from collections import defaultdict
from typing import Any, Dict, List, Optional
GLOBAL_CONFIG = defaultdict(dict)
def register(dct: Any = GLOBAL_CONFIG, name=None, force=False):
"""
dct:
if dct is Dict, register foo into dct as key-value pair
if dct is Clas, register as modules attibute
force
whether force register.
"""
def decorator(foo):
register_name = foo.__name__ if name is None else name
if not force:
if inspect.isclass(dct):
assert not hasattr(dct, foo.__name__), f"module {dct.__name__} has {foo.__name__}"
else:
assert foo.__name__ not in dct, f"{foo.__name__} has been already registered"
if inspect.isfunction(foo):
@functools.wraps(foo)
def wrap_func(*args, **kwargs):
return foo(*args, **kwargs)
if isinstance(dct, dict):
dct[foo.__name__] = wrap_func
elif inspect.isclass(dct):
setattr(dct, foo.__name__, wrap_func)
else:
raise AttributeError("")
return wrap_func
elif inspect.isclass(foo):
dct[register_name] = extract_schema(foo)
else:
raise ValueError(f"Do not support {type(foo)} register")
return foo
return decorator
def extract_schema(module: type):
"""
Args:
module (type),
Return:
Dict,
"""
argspec = inspect.getfullargspec(module.__init__)
arg_names = [arg for arg in argspec.args if arg != "self"]
num_defualts = len(argspec.defaults) if argspec.defaults is not None else 0
num_requires = len(arg_names) - num_defualts
schame = dict()
schame["_name"] = module.__name__
schame["_pymodule"] = importlib.import_module(module.__module__)
schame["_inject"] = getattr(module, "__inject__", [])
schame["_share"] = getattr(module, "__share__", [])
schame["_kwargs"] = {}
for i, name in enumerate(arg_names):
if name in schame["_share"]:
assert i >= num_requires, "share config must have default value."
value = argspec.defaults[i - num_requires]
elif i >= num_requires:
value = argspec.defaults[i - num_requires]
else:
value = None
schame[name] = value
schame["_kwargs"][name] = value
return schame
def create(type_or_name, global_cfg=GLOBAL_CONFIG, **kwargs):
""" """
assert type(type_or_name) in (type, str), "create should be modules or name."
name = type_or_name if isinstance(type_or_name, str) else type_or_name.__name__
if name in global_cfg:
if hasattr(global_cfg[name], "__dict__"):
return global_cfg[name]
else:
raise ValueError("The module {} is not registered".format(name))
cfg = global_cfg[name]
if isinstance(cfg, dict) and "type" in cfg:
_cfg: dict = global_cfg[cfg["type"]]
# clean args
_keys = [k for k in _cfg.keys() if not k.startswith("_")]
for _arg in _keys:
del _cfg[_arg]
_cfg.update(_cfg["_kwargs"]) # restore default args
_cfg.update(cfg) # load config args
_cfg.update(kwargs) # TODO recive extra kwargs
name = _cfg.pop("type") # pop extra key `type` (from cfg)
return create(name, global_cfg)
module = getattr(cfg["_pymodule"], name)
module_kwargs = {}
module_kwargs.update(cfg)
# shared var
for k in cfg["_share"]:
if k in global_cfg:
module_kwargs[k] = global_cfg[k]
else:
module_kwargs[k] = cfg[k]
# inject
for k in cfg["_inject"]:
_k = cfg[k]
if _k is None:
continue
if isinstance(_k, str):
if _k not in global_cfg:
raise ValueError(f"Missing inject config of {_k}.")
_cfg = global_cfg[_k]
if isinstance(_cfg, dict):
module_kwargs[k] = create(_cfg["_name"], global_cfg)
else:
module_kwargs[k] = _cfg
elif isinstance(_k, dict):
if "type" not in _k.keys():
raise ValueError("Missing inject for `type` style.")
_type = str(_k["type"])
if _type not in global_cfg:
raise ValueError(f"Missing {_type} in inspect stage.")
# TODO
_cfg: dict = global_cfg[_type]
# clean args
_keys = [k for k in _cfg.keys() if not k.startswith("_")]
for _arg in _keys:
del _cfg[_arg]
_cfg.update(_cfg["_kwargs"]) # restore default values
_cfg.update(_k) # load config args
name = _cfg.pop("type") # pop extra key (`type` from _k)
module_kwargs[k] = create(name, global_cfg)
else:
raise ValueError(f"Inject does not support {_k}")
# TODO hard code
module_kwargs = {k: v for k, v in module_kwargs.items() if not k.startswith("_")}
# TODO for **kwargs
# extra_args = set(module_kwargs.keys()) - set(arg_names)
# if len(extra_args) > 0:
# raise RuntimeError(f'Error: unknown args {extra_args} for {module}')
return module(**module_kwargs)

View File

@ -0,0 +1,187 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import copy
import re
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from ._config import BaseConfig
from .workspace import create
from .yaml_utils import load_config, merge_config, merge_dict
class YAMLConfig(BaseConfig):
def __init__(self, cfg_path: str, **kwargs) -> None:
super().__init__()
cfg = load_config(cfg_path)
cfg = merge_dict(cfg, kwargs)
self.yaml_cfg = copy.deepcopy(cfg)
for k in super().__dict__:
if not k.startswith("_") and k in cfg:
self.__dict__[k] = cfg[k]
@property
def global_cfg(self):
return merge_config(self.yaml_cfg, inplace=False, overwrite=False)
@property
def model(self) -> torch.nn.Module:
if self._model is None and "model" in self.yaml_cfg:
self._model = create(self.yaml_cfg["model"], self.global_cfg)
return super().model
@property
def postprocessor(self) -> torch.nn.Module:
if self._postprocessor is None and "postprocessor" in self.yaml_cfg:
self._postprocessor = create(self.yaml_cfg["postprocessor"], self.global_cfg)
return super().postprocessor
@property
def criterion(self) -> torch.nn.Module:
if self._criterion is None and "criterion" in self.yaml_cfg:
self._criterion = create(self.yaml_cfg["criterion"], self.global_cfg)
return super().criterion
@property
def optimizer(self) -> optim.Optimizer:
if self._optimizer is None and "optimizer" in self.yaml_cfg:
params = self.get_optim_params(self.yaml_cfg["optimizer"], self.model)
self._optimizer = create("optimizer", self.global_cfg, params=params)
return super().optimizer
@property
def lr_scheduler(self) -> optim.lr_scheduler.LRScheduler:
if self._lr_scheduler is None and "lr_scheduler" in self.yaml_cfg:
self._lr_scheduler = create("lr_scheduler", self.global_cfg, optimizer=self.optimizer)
print(f"Initial lr: {self._lr_scheduler.get_last_lr()}")
return super().lr_scheduler
@property
def lr_warmup_scheduler(self) -> optim.lr_scheduler.LRScheduler:
if self._lr_warmup_scheduler is None and "lr_warmup_scheduler" in self.yaml_cfg:
self._lr_warmup_scheduler = create(
"lr_warmup_scheduler", self.global_cfg, lr_scheduler=self.lr_scheduler
)
return super().lr_warmup_scheduler
@property
def train_dataloader(self) -> DataLoader:
if self._train_dataloader is None and "train_dataloader" in self.yaml_cfg:
self._train_dataloader = self.build_dataloader("train_dataloader")
return super().train_dataloader
@property
def val_dataloader(self) -> DataLoader:
if self._val_dataloader is None and "val_dataloader" in self.yaml_cfg:
self._val_dataloader = self.build_dataloader("val_dataloader")
return super().val_dataloader
@property
def ema(self) -> torch.nn.Module:
if self._ema is None and self.yaml_cfg.get("use_ema", False):
self._ema = create("ema", self.global_cfg, model=self.model)
return super().ema
@property
def scaler(self):
if self._scaler is None and self.yaml_cfg.get("use_amp", False):
self._scaler = create("scaler", self.global_cfg)
return super().scaler
@property
def evaluator(self):
if self._evaluator is None and "evaluator" in self.yaml_cfg:
if self.yaml_cfg["evaluator"]["type"] == "CocoEvaluator":
from ..data import get_coco_api_from_dataset
base_ds = get_coco_api_from_dataset(self.val_dataloader.dataset)
self._evaluator = create("evaluator", self.global_cfg, coco_gt=base_ds)
else:
raise NotImplementedError(f"{self.yaml_cfg['evaluator']['type']}")
return super().evaluator
@property
def use_wandb(self) -> bool:
return self.yaml_cfg.get("use_wandb", False)
@staticmethod
def get_optim_params(cfg: dict, model: nn.Module):
"""
E.g.:
^(?=.*a)(?=.*b).*$ means including a and b
^(?=.*(?:a|b)).*$ means including a or b
^(?=.*a)(?!.*b).*$ means including a, but not b
"""
assert "type" in cfg, ""
cfg = copy.deepcopy(cfg)
if "params" not in cfg:
return model.parameters()
assert isinstance(cfg["params"], list), ""
param_groups = []
visited = []
for pg in cfg["params"]:
pattern = pg["params"]
params = {
k: v
for k, v in model.named_parameters()
if v.requires_grad and len(re.findall(pattern, k)) > 0
}
pg["params"] = params.values()
param_groups.append(pg)
visited.extend(list(params.keys()))
# print(params.keys())
names = [k for k, v in model.named_parameters() if v.requires_grad]
if len(visited) < len(names):
unseen = set(names) - set(visited)
params = {k: v for k, v in model.named_parameters() if v.requires_grad and k in unseen}
param_groups.append({"params": params.values()})
visited.extend(list(params.keys()))
# print(params.keys())
assert len(visited) == len(names), ""
return param_groups
@staticmethod
def get_rank_batch_size(cfg):
"""compute batch size for per rank if total_batch_size is provided."""
assert ("total_batch_size" in cfg or "batch_size" in cfg) and not (
"total_batch_size" in cfg and "batch_size" in cfg
), "`batch_size` or `total_batch_size` should be choosed one"
total_batch_size = cfg.get("total_batch_size", None)
if total_batch_size is None:
bs = cfg.get("batch_size")
else:
from ..misc import dist_utils
assert (
total_batch_size % dist_utils.get_world_size() == 0
), "total_batch_size should be divisible by world size"
bs = total_batch_size // dist_utils.get_world_size()
return bs
def build_dataloader(self, name: str):
bs = self.get_rank_batch_size(self.yaml_cfg[name])
global_cfg = self.global_cfg
if "total_batch_size" in global_cfg[name]:
# pop unexpected key for dataloader init
_ = global_cfg[name].pop("total_batch_size")
print(f"building {name} with batch_size={bs}...")
loader = create(name, global_cfg, batch_size=bs)
loader.shuffle = self.yaml_cfg[name].get("shuffle", False)
return loader

View File

@ -0,0 +1,126 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import copy
import os
from typing import Any, Dict, List, Optional
import yaml
from .workspace import GLOBAL_CONFIG
__all__ = [
"load_config",
"merge_config",
"merge_dict",
"parse_cli",
]
INCLUDE_KEY = "__include__"
def load_config(file_path, cfg=dict()):
"""load config"""
_, ext = os.path.splitext(file_path)
assert ext in [".yml", ".yaml"], "only support yaml files"
with open(file_path) as f:
file_cfg = yaml.load(f, Loader=yaml.Loader)
if file_cfg is None:
return {}
if INCLUDE_KEY in file_cfg:
base_yamls = list(file_cfg[INCLUDE_KEY])
for base_yaml in base_yamls:
if base_yaml.startswith("~"):
base_yaml = os.path.expanduser(base_yaml)
if not base_yaml.startswith("/"):
base_yaml = os.path.join(os.path.dirname(file_path), base_yaml)
with open(base_yaml) as f:
base_cfg = load_config(base_yaml, cfg)
merge_dict(cfg, base_cfg)
return merge_dict(cfg, file_cfg)
def merge_dict(dct, another_dct, inplace=True) -> Dict:
"""merge another_dct into dct"""
def _merge(dct, another) -> Dict:
for k in another:
if k in dct and isinstance(dct[k], dict) and isinstance(another[k], dict):
_merge(dct[k], another[k])
else:
dct[k] = another[k]
return dct
if not inplace:
dct = copy.deepcopy(dct)
return _merge(dct, another_dct)
def dictify(s: str, v: Any) -> Dict:
if "." not in s:
return {s: v}
key, rest = s.split(".", 1)
return {key: dictify(rest, v)}
def parse_cli(nargs: List[str]) -> Dict:
"""
parse command-line arguments
convert `a.c=3 b=10` to `{'a': {'c': 3}, 'b': 10}`
"""
cfg = {}
if nargs is None or len(nargs) == 0:
return cfg
for s in nargs:
s = s.strip()
k, v = s.split("=", 1)
d = dictify(k, yaml.load(v, Loader=yaml.Loader))
cfg = merge_dict(cfg, d)
return cfg
def merge_config(cfg, another_cfg=GLOBAL_CONFIG, inplace: bool = False, overwrite: bool = False):
"""
Merge another_cfg into cfg, return the merged config
Example:
cfg1 = load_config('./dfine_r18vd_6x_coco.yml')
cfg1 = merge_config(cfg, inplace=True)
cfg2 = load_config('./dfine_r50vd_6x_coco.yml')
cfg2 = merge_config(cfg2, inplace=True)
model1 = create(cfg1['model'], cfg1)
model2 = create(cfg2['model'], cfg2)
"""
def _merge(dct, another):
for k in another:
if k not in dct:
dct[k] = another[k]
elif isinstance(dct[k], dict) and isinstance(another[k], dict):
_merge(dct[k], another[k])
elif overwrite:
dct[k] = another[k]
return cfg
if not inplace:
cfg = copy.deepcopy(cfg)
return _merge(cfg, another_cfg)

View File

@ -0,0 +1,20 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from ._misc import convert_to_tv_tensor
from .dataloader import *
from .dataset import *
from .transforms import *
# def set_epoch(self, epoch) -> None:
# self.epoch = epoch
# def _set_epoch_func(datasets):
# """Add `set_epoch` for datasets
# """
# from ..core import register
# for ds in datasets:
# register(ds)(set_epoch)
# _set_epoch_func([CIFAR10, VOCDetection, CocoDetection])

View File

@ -0,0 +1,62 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import importlib.metadata
from torch import Tensor
if "0.15.2" in importlib.metadata.version("torchvision"):
import torchvision
torchvision.disable_beta_transforms_warning()
from torchvision.datapoints import BoundingBox as BoundingBoxes
from torchvision.datapoints import BoundingBoxFormat, Image, Mask, Video
from torchvision.transforms.v2 import SanitizeBoundingBox as SanitizeBoundingBoxes
_boxes_keys = ["format", "spatial_size"]
elif "0.17" > importlib.metadata.version("torchvision") >= "0.16":
import torchvision
torchvision.disable_beta_transforms_warning()
from torchvision.transforms.v2 import SanitizeBoundingBoxes
from torchvision.tv_tensors import BoundingBoxes, BoundingBoxFormat, Image, Mask, Video
_boxes_keys = ["format", "canvas_size"]
elif importlib.metadata.version("torchvision") >= "0.17":
import torchvision
from torchvision.transforms.v2 import SanitizeBoundingBoxes
from torchvision.tv_tensors import BoundingBoxes, BoundingBoxFormat, Image, Mask, Video
_boxes_keys = ["format", "canvas_size"]
else:
raise RuntimeError("Please make sure torchvision version >= 0.15.2")
def convert_to_tv_tensor(tensor: Tensor, key: str, box_format="xyxy", spatial_size=None) -> Tensor:
"""
Args:
tensor (Tensor): input tensor
key (str): transform to key
Return:
Dict[str, TV_Tensor]
"""
assert key in (
"boxes",
"masks",
), "Only support 'boxes' and 'masks'"
if key == "boxes":
box_format = getattr(BoundingBoxFormat, box_format.upper())
_kwargs = dict(zip(_boxes_keys, [box_format, spatial_size]))
return BoundingBoxes(tensor, **_kwargs)
if key == "masks":
return Mask(tensor)

View File

@ -0,0 +1,122 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import random
from functools import partial
import torch
import torch.nn.functional as F
import torch.utils.data as data
import torchvision
import torchvision.transforms.v2 as VT
from torch.utils.data import default_collate
from torchvision.transforms.v2 import InterpolationMode
from torchvision.transforms.v2 import functional as VF
from ..core import register
torchvision.disable_beta_transforms_warning()
__all__ = [
"DataLoader",
"BaseCollateFunction",
"BatchImageCollateFunction",
"batch_image_collate_fn",
]
@register()
class DataLoader(data.DataLoader):
__inject__ = ["dataset", "collate_fn"]
def __repr__(self) -> str:
format_string = self.__class__.__name__ + "("
for n in ["dataset", "batch_size", "num_workers", "drop_last", "collate_fn"]:
format_string += "\n"
format_string += " {0}: {1}".format(n, getattr(self, n))
format_string += "\n)"
return format_string
def set_epoch(self, epoch):
self._epoch = epoch
self.dataset.set_epoch(epoch)
self.collate_fn.set_epoch(epoch)
@property
def epoch(self):
return self._epoch if hasattr(self, "_epoch") else -1
@property
def shuffle(self):
return self._shuffle
@shuffle.setter
def shuffle(self, shuffle):
assert isinstance(shuffle, bool), "shuffle must be a boolean"
self._shuffle = shuffle
@register()
def batch_image_collate_fn(items):
"""only batch image"""
return torch.cat([x[0][None] for x in items], dim=0), [x[1] for x in items]
class BaseCollateFunction(object):
def set_epoch(self, epoch):
self._epoch = epoch
@property
def epoch(self):
return self._epoch if hasattr(self, "_epoch") else -1
def __call__(self, items):
raise NotImplementedError("")
def generate_scales(base_size, base_size_repeat):
scale_repeat = (base_size - int(base_size * 0.75 / 32) * 32) // 32
scales = [int(base_size * 0.75 / 32) * 32 + i * 32 for i in range(scale_repeat)]
scales += [base_size] * base_size_repeat
scales += [int(base_size * 1.25 / 32) * 32 - i * 32 for i in range(scale_repeat)]
return scales
@register()
class BatchImageCollateFunction(BaseCollateFunction):
def __init__(
self,
stop_epoch=None,
ema_restart_decay=0.9999,
base_size=640,
base_size_repeat=None,
) -> None:
super().__init__()
self.base_size = base_size
self.scales = (
generate_scales(base_size, base_size_repeat) if base_size_repeat is not None else None
)
self.stop_epoch = stop_epoch if stop_epoch is not None else 100000000
self.ema_restart_decay = ema_restart_decay
# self.interpolation = interpolation
def __call__(self, items):
images = torch.cat([x[0][None] for x in items], dim=0)
targets = [x[1] for x in items]
if self.scales is not None and self.epoch < self.stop_epoch:
# sz = random.choice(self.scales)
# sz = [sz] if isinstance(sz, int) else list(sz)
# VF.resize(inpt, sz, interpolation=self.interpolation)
sz = random.choice(self.scales)
images = F.interpolate(images, size=sz)
if "masks" in targets[0]:
for tg in targets:
tg["masks"] = F.interpolate(tg["masks"], size=sz, mode="nearest")
raise NotImplementedError("")
return images, targets

View File

@ -0,0 +1,17 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
# from ._dataset import DetDataset
from .cifar_dataset import CIFAR10
from .coco_dataset import (
CocoDetection,
mscoco_category2label,
mscoco_category2name,
mscoco_label2category,
)
from .coco_eval import CocoEvaluator
from .coco_utils import get_coco_api_from_dataset
from .voc_detection import VOCDetection
from .voc_eval import VOCEvaluator

View File

@ -0,0 +1,27 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import torch
import torch.utils.data as data
class DetDataset(data.Dataset):
def __getitem__(self, index):
img, target = self.load_item(index)
if self.transforms is not None:
img, target, _ = self.transforms(img, target, self)
return img, target
def load_item(self, index):
raise NotImplementedError(
"Please implement this function to return item before `transforms`."
)
def set_epoch(self, epoch) -> None:
self._epoch = epoch
@property
def epoch(self):
return self._epoch if hasattr(self, "_epoch") else -1

View File

@ -0,0 +1,25 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from typing import Callable, Optional
import torchvision
from ...core import register
@register()
class CIFAR10(torchvision.datasets.CIFAR10):
__inject__ = ["transform", "target_transform"]
def __init__(
self,
root: str,
train: bool = True,
transform: Optional[Callable] = None,
target_transform: Optional[Callable] = None,
download: bool = False,
) -> None:
super().__init__(root, train, transform, target_transform, download)

View File

@ -0,0 +1,280 @@
"""
Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
Mostly copy-paste from https://github.com/pytorch/vision/blob/13b35ff/references/detection/coco_utils.py
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import faster_coco_eval.core.mask as coco_mask
from faster_coco_eval.utils.pytorch import FasterCocoDetection
import torch
import torchvision
import os
from PIL import Image
from ...core import register
from .._misc import convert_to_tv_tensor
from ._dataset import DetDataset
torchvision.disable_beta_transforms_warning()
Image.MAX_IMAGE_PIXELS = None
__all__ = ["CocoDetection"]
@register()
class CocoDetection(FasterCocoDetection, DetDataset):
__inject__ = [
"transforms",
]
__share__ = ["remap_mscoco_category"]
def __init__(
self, img_folder, ann_file, transforms, return_masks=False, remap_mscoco_category=False
):
super(FasterCocoDetection, self).__init__(img_folder, ann_file)
self._transforms = transforms
self.prepare = ConvertCocoPolysToMask(return_masks)
self.img_folder = img_folder
self.ann_file = ann_file
self.return_masks = return_masks
self.remap_mscoco_category = remap_mscoco_category
def __getitem__(self, idx):
img, target = self.load_item(idx)
if self._transforms is not None:
img, target, _ = self._transforms(img, target, self)
return img, target
def load_item(self, idx):
image, target = super(FasterCocoDetection, self).__getitem__(idx)
image_id = self.ids[idx]
image_path = os.path.join(self.img_folder, self.coco.loadImgs(image_id)[0]["file_name"])
target = {"image_id": image_id, "image_path": image_path, "annotations": target}
if self.remap_mscoco_category:
image, target = self.prepare(image, target, category2label=mscoco_category2label)
else:
image, target = self.prepare(image, target)
target["idx"] = torch.tensor([idx])
if "boxes" in target:
target["boxes"] = convert_to_tv_tensor(
target["boxes"], key="boxes", spatial_size=image.size[::-1]
)
if "masks" in target:
target["masks"] = convert_to_tv_tensor(target["masks"], key="masks")
return image, target
def extra_repr(self) -> str:
s = f" img_folder: {self.img_folder}\n ann_file: {self.ann_file}\n"
s += f" return_masks: {self.return_masks}\n"
if hasattr(self, "_transforms") and self._transforms is not None:
s += f" transforms:\n {repr(self._transforms)}"
if hasattr(self, "_preset") and self._preset is not None:
s += f" preset:\n {repr(self._preset)}"
return s
@property
def categories(
self,
):
return self.coco.dataset["categories"]
@property
def category2name(
self,
):
return {cat["id"]: cat["name"] for cat in self.categories}
@property
def category2label(
self,
):
return {cat["id"]: i for i, cat in enumerate(self.categories)}
@property
def label2category(
self,
):
return {i: cat["id"] for i, cat in enumerate(self.categories)}
def convert_coco_poly_to_mask(segmentations, height, width):
masks = []
for polygons in segmentations:
rles = coco_mask.frPyObjects(polygons, height, width)
mask = coco_mask.decode(rles)
if len(mask.shape) < 3:
mask = mask[..., None]
mask = torch.as_tensor(mask, dtype=torch.uint8)
mask = mask.any(dim=2)
masks.append(mask)
if masks:
masks = torch.stack(masks, dim=0)
else:
masks = torch.zeros((0, height, width), dtype=torch.uint8)
return masks
class ConvertCocoPolysToMask(object):
def __init__(self, return_masks=False):
self.return_masks = return_masks
def __call__(self, image: Image.Image, target, **kwargs):
w, h = image.size
image_id = target["image_id"]
image_id = torch.tensor([image_id])
image_path = target["image_path"]
anno = target["annotations"]
anno = [obj for obj in anno if "iscrowd" not in obj or obj["iscrowd"] == 0]
boxes = [obj["bbox"] for obj in anno]
# guard against no boxes via resizing
boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
boxes[:, 2:] += boxes[:, :2]
boxes[:, 0::2].clamp_(min=0, max=w)
boxes[:, 1::2].clamp_(min=0, max=h)
category2label = kwargs.get("category2label", None)
if category2label is not None:
labels = [category2label[obj["category_id"]] for obj in anno]
else:
labels = [obj["category_id"] for obj in anno]
labels = torch.tensor(labels, dtype=torch.int64)
if self.return_masks:
segmentations = [obj["segmentation"] for obj in anno]
masks = convert_coco_poly_to_mask(segmentations, h, w)
keypoints = None
if anno and "keypoints" in anno[0]:
keypoints = [obj["keypoints"] for obj in anno]
keypoints = torch.as_tensor(keypoints, dtype=torch.float32)
num_keypoints = keypoints.shape[0]
if num_keypoints:
keypoints = keypoints.view(num_keypoints, -1, 3)
keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
boxes = boxes[keep]
labels = labels[keep]
if self.return_masks:
masks = masks[keep]
if keypoints is not None:
keypoints = keypoints[keep]
target = {}
target["boxes"] = boxes
target["labels"] = labels
if self.return_masks:
target["masks"] = masks
target["image_id"] = image_id
target["image_path"] = image_path
if keypoints is not None:
target["keypoints"] = keypoints
# for conversion to coco api
area = torch.tensor([obj["area"] for obj in anno])
iscrowd = torch.tensor([obj["iscrowd"] if "iscrowd" in obj else 0 for obj in anno])
target["area"] = area[keep]
target["iscrowd"] = iscrowd[keep]
target["orig_size"] = torch.as_tensor([int(w), int(h)])
# target["size"] = torch.as_tensor([int(w), int(h)])
return image, target
mscoco_category2name = {
1: "person",
2: "bicycle",
3: "car",
4: "motorcycle",
5: "airplane",
6: "bus",
7: "train",
8: "truck",
9: "boat",
10: "traffic light",
11: "fire hydrant",
13: "stop sign",
14: "parking meter",
15: "bench",
16: "bird",
17: "cat",
18: "dog",
19: "horse",
20: "sheep",
21: "cow",
22: "elephant",
23: "bear",
24: "zebra",
25: "giraffe",
27: "backpack",
28: "umbrella",
31: "handbag",
32: "tie",
33: "suitcase",
34: "frisbee",
35: "skis",
36: "snowboard",
37: "sports ball",
38: "kite",
39: "baseball bat",
40: "baseball glove",
41: "skateboard",
42: "surfboard",
43: "tennis racket",
44: "bottle",
46: "wine glass",
47: "cup",
48: "fork",
49: "knife",
50: "spoon",
51: "bowl",
52: "banana",
53: "apple",
54: "sandwich",
55: "orange",
56: "broccoli",
57: "carrot",
58: "hot dog",
59: "pizza",
60: "donut",
61: "cake",
62: "chair",
63: "couch",
64: "potted plant",
65: "bed",
67: "dining table",
70: "toilet",
72: "tv",
73: "laptop",
74: "mouse",
75: "remote",
76: "keyboard",
77: "cell phone",
78: "microwave",
79: "oven",
80: "toaster",
81: "sink",
82: "refrigerator",
84: "book",
85: "clock",
86: "vase",
87: "scissors",
88: "teddy bear",
89: "hair drier",
90: "toothbrush",
}
mscoco_category2label = {k: i for i, k in enumerate(mscoco_category2name.keys())}
mscoco_label2category = {v: k for k, v in mscoco_category2label.items()}

View File

@ -0,0 +1,22 @@
"""
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
COCO evaluator that works in distributed mode.
Mostly copy-paste from https://github.com/pytorch/vision/blob/edfd5a7/references/detection/coco_eval.py
The difference is that there is less copy-pasting from pycocotools
in the end of the file, as python3 can suppress prints with contextlib
# MiXaiLL76 replacing pycocotools with faster-coco-eval for better performance and support.
"""
from faster_coco_eval.utils.pytorch import FasterCocoEvaluator
from ...core import register
__all__ = [
"CocoEvaluator",
]
@register()
class CocoEvaluator(FasterCocoEvaluator):
pass

View File

@ -0,0 +1,191 @@
"""
copy and modified https://github.com/pytorch/vision/blob/main/references/detection/coco_utils.py
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import faster_coco_eval.core.mask as coco_mask
import torch
import torch.utils.data
import torchvision
import torchvision.transforms.functional as TVF
from faster_coco_eval import COCO
def convert_coco_poly_to_mask(segmentations, height, width):
masks = []
for polygons in segmentations:
rles = coco_mask.frPyObjects(polygons, height, width)
mask = coco_mask.decode(rles)
if len(mask.shape) < 3:
mask = mask[..., None]
mask = torch.as_tensor(mask, dtype=torch.uint8)
mask = mask.any(dim=2)
masks.append(mask)
if masks:
masks = torch.stack(masks, dim=0)
else:
masks = torch.zeros((0, height, width), dtype=torch.uint8)
return masks
class ConvertCocoPolysToMask:
def __call__(self, image, target):
w, h = image.size
image_id = target["image_id"]
anno = target["annotations"]
anno = [obj for obj in anno if obj["iscrowd"] == 0]
boxes = [obj["bbox"] for obj in anno]
# guard against no boxes via resizing
boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
boxes[:, 2:] += boxes[:, :2]
boxes[:, 0::2].clamp_(min=0, max=w)
boxes[:, 1::2].clamp_(min=0, max=h)
classes = [obj["category_id"] for obj in anno]
classes = torch.tensor(classes, dtype=torch.int64)
segmentations = [obj["segmentation"] for obj in anno]
masks = convert_coco_poly_to_mask(segmentations, h, w)
keypoints = None
if anno and "keypoints" in anno[0]:
keypoints = [obj["keypoints"] for obj in anno]
keypoints = torch.as_tensor(keypoints, dtype=torch.float32)
num_keypoints = keypoints.shape[0]
if num_keypoints:
keypoints = keypoints.view(num_keypoints, -1, 3)
keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
boxes = boxes[keep]
classes = classes[keep]
masks = masks[keep]
if keypoints is not None:
keypoints = keypoints[keep]
target = {}
target["boxes"] = boxes
target["labels"] = classes
target["masks"] = masks
target["image_id"] = image_id
if keypoints is not None:
target["keypoints"] = keypoints
# for conversion to coco api
area = torch.tensor([obj["area"] for obj in anno])
iscrowd = torch.tensor([obj["iscrowd"] for obj in anno])
target["area"] = area
target["iscrowd"] = iscrowd
return image, target
def _coco_remove_images_without_annotations(dataset, cat_list=None):
def _has_only_empty_bbox(anno):
return all(any(o <= 1 for o in obj["bbox"][2:]) for obj in anno)
def _count_visible_keypoints(anno):
return sum(sum(1 for v in ann["keypoints"][2::3] if v > 0) for ann in anno)
min_keypoints_per_image = 10
def _has_valid_annotation(anno):
# if it's empty, there is no annotation
if len(anno) == 0:
return False
# if all boxes have close to zero area, there is no annotation
if _has_only_empty_bbox(anno):
return False
# keypoints task have a slight different criteria for considering
# if an annotation is valid
if "keypoints" not in anno[0]:
return True
# for keypoint detection tasks, only consider valid images those
# containing at least min_keypoints_per_image
if _count_visible_keypoints(anno) >= min_keypoints_per_image:
return True
return False
ids = []
for ds_idx, img_id in enumerate(dataset.ids):
ann_ids = dataset.coco.getAnnIds(imgIds=img_id, iscrowd=None)
anno = dataset.coco.loadAnns(ann_ids)
if cat_list:
anno = [obj for obj in anno if obj["category_id"] in cat_list]
if _has_valid_annotation(anno):
ids.append(ds_idx)
dataset = torch.utils.data.Subset(dataset, ids)
return dataset
def convert_to_coco_api(ds):
coco_ds = COCO()
# annotation IDs need to start at 1, not 0, see torchvision issue #1530
ann_id = 1
dataset = {"images": [], "categories": [], "annotations": []}
categories = set()
for img_idx in range(len(ds)):
# find better way to get target
# targets = ds.get_annotations(img_idx)
# img, targets = ds[img_idx]
img, targets = ds.load_item(img_idx)
width, height = img.size
image_id = targets["image_id"].item()
img_dict = {}
img_dict["id"] = image_id
img_dict["width"] = width
img_dict["height"] = height
dataset["images"].append(img_dict)
bboxes = targets["boxes"].clone()
bboxes[:, 2:] -= bboxes[:, :2] # xyxy -> xywh
bboxes = bboxes.tolist()
labels = targets["labels"].tolist()
areas = targets["area"].tolist()
iscrowd = targets["iscrowd"].tolist()
if "masks" in targets:
masks = targets["masks"]
# make masks Fortran contiguous for coco_mask
masks = masks.permute(0, 2, 1).contiguous().permute(0, 2, 1)
if "keypoints" in targets:
keypoints = targets["keypoints"]
keypoints = keypoints.reshape(keypoints.shape[0], -1).tolist()
num_objs = len(bboxes)
for i in range(num_objs):
ann = {}
ann["image_id"] = image_id
ann["bbox"] = bboxes[i]
ann["category_id"] = labels[i]
categories.add(labels[i])
ann["area"] = areas[i]
ann["iscrowd"] = iscrowd[i]
ann["id"] = ann_id
if "masks" in targets:
ann["segmentation"] = coco_mask.encode(masks[i].numpy())
if "keypoints" in targets:
ann["keypoints"] = keypoints[i]
ann["num_keypoints"] = sum(k != 0 for k in keypoints[i][2::3])
dataset["annotations"].append(ann)
ann_id += 1
dataset["categories"] = [{"id": i} for i in sorted(categories)]
coco_ds.dataset = dataset
coco_ds.createIndex()
return coco_ds
def get_coco_api_from_dataset(dataset):
# FIXME: This is... awful?
for _ in range(10):
if isinstance(dataset, torchvision.datasets.CocoDetection):
break
if isinstance(dataset, torch.utils.data.Subset):
dataset = dataset.dataset
if isinstance(dataset, torchvision.datasets.CocoDetection):
return dataset.coco
return convert_to_coco_api(dataset)

View File

@ -0,0 +1,86 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import os
from typing import Callable, Optional
import torch
import torchvision
import torchvision.transforms.functional as TVF
from PIL import Image
from sympy import im
try:
from defusedxml.ElementTree import parse as ET_parse
except ImportError:
from xml.etree.ElementTree import parse as ET_parse
from ...core import register
from .._misc import convert_to_tv_tensor
from ._dataset import DetDataset
@register()
class VOCDetection(torchvision.datasets.VOCDetection, DetDataset):
__inject__ = [
"transforms",
]
def __init__(
self,
root: str,
ann_file: str = "trainval.txt",
label_file: str = "label_list.txt",
transforms: Optional[Callable] = None,
):
with open(os.path.join(root, ann_file), "r") as f:
lines = [x.strip() for x in f.readlines()]
lines = [x.split(" ") for x in lines]
self.images = [os.path.join(root, lin[0]) for lin in lines]
self.targets = [os.path.join(root, lin[1]) for lin in lines]
assert len(self.images) == len(self.targets)
with open(os.path.join(root + label_file), "r") as f:
labels = f.readlines()
labels = [lab.strip() for lab in labels]
self.transforms = transforms
self.labels_map = {lab: i for i, lab in enumerate(labels)}
def __getitem__(self, index: int):
image, target = self.load_item(index)
if self.transforms is not None:
image, target, _ = self.transforms(image, target, self)
# target["orig_size"] = torch.tensor(TVF.get_image_size(image))
return image, target
def load_item(self, index: int):
image = Image.open(self.images[index]).convert("RGB")
target = self.parse_voc_xml(ET_parse(self.annotations[index]).getroot())
output = {}
output["image_id"] = torch.tensor([index])
for k in ["area", "boxes", "labels", "iscrowd"]:
output[k] = []
for blob in target["annotation"]["object"]:
box = [float(v) for v in blob["bndbox"].values()]
output["boxes"].append(box)
output["labels"].append(blob["name"])
output["area"].append((box[2] - box[0]) * (box[3] - box[1]))
output["iscrowd"].append(0)
w, h = image.size
boxes = torch.tensor(output["boxes"]) if len(output["boxes"]) > 0 else torch.zeros(0, 4)
output["boxes"] = convert_to_tv_tensor(
boxes, "boxes", box_format="xyxy", spatial_size=[h, w]
)
output["labels"] = torch.tensor([self.labels_map[lab] for lab in output["labels"]])
output["area"] = torch.tensor(output["area"])
output["iscrowd"] = torch.tensor(output["iscrowd"])
output["orig_size"] = torch.tensor([w, h])
return image, output

View File

@ -0,0 +1,12 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import torch
import torchvision
class VOCEvaluator(object):
def __init__(self) -> None:
pass

View File

@ -0,0 +1,21 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from ._transforms import (
ConvertBoxes,
ConvertPILImage,
EmptyTransform,
Normalize,
PadToSize,
RandomCrop,
RandomHorizontalFlip,
RandomIoUCrop,
RandomPhotometricDistort,
RandomZoomOut,
Resize,
SanitizeBoundingBoxes,
)
from .container import Compose
from .mosaic import Mosaic

View File

@ -0,0 +1,161 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from typing import Any, Dict, List, Optional
import PIL
import PIL.Image
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms.v2 as T
import torchvision.transforms.v2.functional as F
from ...core import register
from .._misc import (
BoundingBoxes,
Image,
Mask,
SanitizeBoundingBoxes,
Video,
_boxes_keys,
convert_to_tv_tensor,
)
torchvision.disable_beta_transforms_warning()
RandomPhotometricDistort = register()(T.RandomPhotometricDistort)
RandomZoomOut = register()(T.RandomZoomOut)
RandomHorizontalFlip = register()(T.RandomHorizontalFlip)
Resize = register()(T.Resize)
# ToImageTensor = register()(T.ToImageTensor)
# ConvertDtype = register()(T.ConvertDtype)
# PILToTensor = register()(T.PILToTensor)
SanitizeBoundingBoxes = register(name="SanitizeBoundingBoxes")(SanitizeBoundingBoxes)
RandomCrop = register()(T.RandomCrop)
Normalize = register()(T.Normalize)
@register()
class EmptyTransform(T.Transform):
def __init__(
self,
) -> None:
super().__init__()
def forward(self, *inputs):
inputs = inputs if len(inputs) > 1 else inputs[0]
return inputs
@register()
class PadToSize(T.Pad):
_transformed_types = (
PIL.Image.Image,
Image,
Video,
Mask,
BoundingBoxes,
)
def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
sp = F.get_spatial_size(flat_inputs[0])
h, w = self.size[1] - sp[0], self.size[0] - sp[1]
self.padding = [0, 0, w, h]
return dict(padding=self.padding)
def __init__(self, size, fill=0, padding_mode="constant") -> None:
if isinstance(size, int):
size = (size, size)
self.size = size
super().__init__(0, fill, padding_mode)
def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
fill = self._fill[type(inpt)]
padding = params["padding"]
return F.pad(inpt, padding=padding, fill=fill, padding_mode=self.padding_mode) # type: ignore[arg-type]
def __call__(self, *inputs: Any) -> Any:
outputs = super().forward(*inputs)
if len(outputs) > 1 and isinstance(outputs[1], dict):
outputs[1]["padding"] = torch.tensor(self.padding)
return outputs
@register()
class RandomIoUCrop(T.RandomIoUCrop):
def __init__(
self,
min_scale: float = 0.3,
max_scale: float = 1,
min_aspect_ratio: float = 0.5,
max_aspect_ratio: float = 2,
sampler_options: Optional[List[float]] = None,
trials: int = 40,
p: float = 1.0,
):
super().__init__(
min_scale, max_scale, min_aspect_ratio, max_aspect_ratio, sampler_options, trials
)
self.p = p
def __call__(self, *inputs: Any) -> Any:
if torch.rand(1) >= self.p:
return inputs if len(inputs) > 1 else inputs[0]
return super().forward(*inputs)
@register()
class ConvertBoxes(T.Transform):
_transformed_types = (BoundingBoxes,)
def __init__(self, fmt="", normalize=False) -> None:
super().__init__()
self.fmt = fmt
self.normalize = normalize
def transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
return self._transform(inpt, params)
def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
spatial_size = getattr(inpt, _boxes_keys[1])
if self.fmt:
in_fmt = inpt.format.value.lower()
inpt = torchvision.ops.box_convert(inpt, in_fmt=in_fmt, out_fmt=self.fmt.lower())
inpt = convert_to_tv_tensor(
inpt, key="boxes", box_format=self.fmt.upper(), spatial_size=spatial_size
)
if self.normalize:
inpt = inpt / torch.tensor(spatial_size[::-1]).tile(2)[None]
return inpt
@register()
class ConvertPILImage(T.Transform):
_transformed_types = (PIL.Image.Image,)
def __init__(self, dtype="float32", scale=True) -> None:
super().__init__()
self.dtype = dtype
self.scale = scale
def transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
return self._transform(inpt, params)
def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
inpt = F.pil_to_tensor(inpt)
if self.dtype == "float32":
inpt = inpt.float()
if self.scale:
inpt = inpt / 255.0
inpt = Image(inpt)
return inpt

View File

@ -0,0 +1,99 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from typing import Any, Dict, List, Optional
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms.v2 as T
from ...core import GLOBAL_CONFIG, register
from ._transforms import EmptyTransform
torchvision.disable_beta_transforms_warning()
@register()
class Compose(T.Compose):
def __init__(self, ops, policy=None) -> None:
transforms = []
if ops is not None:
for op in ops:
if isinstance(op, dict):
name = op.pop("type")
transform = getattr(
GLOBAL_CONFIG[name]["_pymodule"], GLOBAL_CONFIG[name]["_name"]
)(**op)
transforms.append(transform)
op["type"] = name
elif isinstance(op, nn.Module):
transforms.append(op)
else:
raise ValueError("")
else:
transforms = [
EmptyTransform(),
]
super().__init__(transforms=transforms)
if policy is None:
policy = {"name": "default"}
self.policy = policy
self.global_samples = 0
def forward(self, *inputs: Any) -> Any:
return self.get_forward(self.policy["name"])(*inputs)
def get_forward(self, name):
forwards = {
"default": self.default_forward,
"stop_epoch": self.stop_epoch_forward,
"stop_sample": self.stop_sample_forward,
}
return forwards[name]
def default_forward(self, *inputs: Any) -> Any:
sample = inputs if len(inputs) > 1 else inputs[0]
for transform in self.transforms:
sample = transform(sample)
return sample
def stop_epoch_forward(self, *inputs: Any):
sample = inputs if len(inputs) > 1 else inputs[0]
dataset = sample[-1]
cur_epoch = dataset.epoch
policy_ops = self.policy["ops"]
policy_epoch = self.policy["epoch"]
for transform in self.transforms:
if type(transform).__name__ in policy_ops and cur_epoch >= policy_epoch:
pass
else:
sample = transform(sample)
return sample
def stop_sample_forward(self, *inputs: Any):
sample = inputs if len(inputs) > 1 else inputs[0]
dataset = sample[-1]
cur_epoch = dataset.epoch
policy_ops = self.policy["ops"]
policy_sample = self.policy["sample"]
for transform in self.transforms:
if type(transform).__name__ in policy_ops and self.global_samples >= policy_sample:
pass
else:
sample = transform(sample)
self.global_samples += 1
return sample

View File

@ -0,0 +1,172 @@
from typing import List, Optional
import torch
# needed due to empty tensor bug in pytorch and torchvision 0.5
import torchvision
import torchvision.transforms.functional as F
from packaging import version
from torch import Tensor
if version.parse(torchvision.__version__) < version.parse("0.7"):
from torchvision.ops import _new_empty_tensor
from torchvision.ops.misc import _output_size
def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None):
# type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor
"""
Equivalent to nn.functional.interpolate, but with support for empty batch sizes.
This will eventually be supported natively by PyTorch, and this
class can go away.
"""
if version.parse(torchvision.__version__) < version.parse("0.7"):
if input.numel() > 0:
return torch.nn.functional.interpolate(input, size, scale_factor, mode, align_corners)
output_shape = _output_size(2, input, size, scale_factor)
output_shape = list(input.shape[:-2]) + list(output_shape)
return _new_empty_tensor(input, output_shape)
else:
return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners)
def crop(image, target, region):
cropped_image = F.crop(image, *region)
target = target.copy()
i, j, h, w = region
# should we do something wrt the original size?
target["size"] = torch.tensor([h, w])
fields = ["labels", "area", "iscrowd"]
if "boxes" in target:
boxes = target["boxes"]
max_size = torch.as_tensor([w, h], dtype=torch.float32)
cropped_boxes = boxes - torch.as_tensor([j, i, j, i])
cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)
cropped_boxes = cropped_boxes.clamp(min=0)
area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)
target["boxes"] = cropped_boxes.reshape(-1, 4)
target["area"] = area
fields.append("boxes")
if "masks" in target:
# FIXME should we update the area here if there are no boxes?
target["masks"] = target["masks"][:, i : i + h, j : j + w]
fields.append("masks")
# remove elements for which the boxes or masks that have zero area
if "boxes" in target or "masks" in target:
# favor boxes selection when defining which elements to keep
# this is compatible with previous implementation
if "boxes" in target:
cropped_boxes = target["boxes"].reshape(-1, 2, 2)
keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1)
else:
keep = target["masks"].flatten(1).any(1)
for field in fields:
target[field] = target[field][keep]
return cropped_image, target
def hflip(image, target):
flipped_image = F.hflip(image)
w, h = image.size
target = target.copy()
if "boxes" in target:
boxes = target["boxes"]
boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1]) + torch.as_tensor(
[w, 0, w, 0]
)
target["boxes"] = boxes
if "masks" in target:
target["masks"] = target["masks"].flip(-1)
return flipped_image, target
def resize(image, target, size, max_size=None):
# size can be min_size (scalar) or (w, h) tuple
def get_size_with_aspect_ratio(image_size, size, max_size=None):
w, h = image_size
if max_size is not None:
min_original_size = float(min((w, h)))
max_original_size = float(max((w, h)))
if max_original_size / min_original_size * size > max_size:
size = int(round(max_size * min_original_size / max_original_size))
if (w <= h and w == size) or (h <= w and h == size):
return (h, w)
if w < h:
ow = size
oh = int(size * h / w)
else:
oh = size
ow = int(size * w / h)
# r = min(size / min(h, w), max_size / max(h, w))
# ow = int(w * r)
# oh = int(h * r)
return (oh, ow)
def get_size(image_size, size, max_size=None):
if isinstance(size, (list, tuple)):
return size[::-1]
else:
return get_size_with_aspect_ratio(image_size, size, max_size)
size = get_size(image.size, size, max_size)
rescaled_image = F.resize(image, size)
if target is None:
return rescaled_image, None
ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size))
ratio_width, ratio_height = ratios
target = target.copy()
if "boxes" in target:
boxes = target["boxes"]
scaled_boxes = boxes * torch.as_tensor(
[ratio_width, ratio_height, ratio_width, ratio_height]
)
target["boxes"] = scaled_boxes
if "area" in target:
area = target["area"]
scaled_area = area * (ratio_width * ratio_height)
target["area"] = scaled_area
h, w = size
target["size"] = torch.tensor([h, w])
if "masks" in target:
target["masks"] = (
interpolate(target["masks"][:, None].float(), size, mode="nearest")[:, 0] > 0.5
)
return rescaled_image, target
def pad(image, target, padding):
# assumes that we only pad on the bottom right corners
padded_image = F.pad(image, (0, 0, padding[0], padding[1]))
if target is None:
return padded_image, None
target = target.copy()
# should we do something wrt the original size?
target["size"] = torch.tensor(padded_image.size[::-1])
if "masks" in target:
target["masks"] = torch.nn.functional.pad(target["masks"], (0, padding[0], 0, padding[1]))
return padded_image, target

View File

@ -0,0 +1,83 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import random
import torch
import torchvision
import torchvision.transforms.v2 as T
import torchvision.transforms.v2.functional as F
from PIL import Image
from ...core import register
from .._misc import convert_to_tv_tensor
torchvision.disable_beta_transforms_warning()
@register()
class Mosaic(T.Transform):
def __init__(
self,
size,
max_size=None,
) -> None:
super().__init__()
self.resize = T.Resize(size=size, max_size=max_size)
self.crop = T.RandomCrop(size=max_size if max_size else size)
# TODO add arg `output_size` for affine`
# self.random_perspective = T.RandomPerspective(distortion_scale=0.5, p=1., )
self.random_affine = T.RandomAffine(
degrees=0, translate=(0.1, 0.1), scale=(0.5, 1.5), fill=114
)
def forward(self, *inputs):
inputs = inputs if len(inputs) > 1 else inputs[0]
image, target, dataset = inputs
images = []
targets = []
indices = random.choices(range(len(dataset)), k=3)
for i in indices:
image, target = dataset.load_item(i)
image, target = self.resize(image, target)
images.append(image)
targets.append(target)
h, w = F.get_spatial_size(images[0])
offset = [[0, 0], [w, 0], [0, h], [w, h]]
image = Image.new(mode=images[0].mode, size=(w * 2, h * 2), color=0)
for i, im in enumerate(images):
image.paste(im, offset[i])
offset = torch.tensor([[0, 0], [w, 0], [0, h], [w, h]]).repeat(1, 2)
target = {}
for k in targets[0]:
if k == "boxes":
v = [t[k] + offset[i] for i, t in enumerate(targets)]
else:
v = [t[k] for t in targets]
if isinstance(v[0], torch.Tensor):
v = torch.cat(v, dim=0)
target[k] = v
if "boxes" in target:
# target['boxes'] = target['boxes'].clamp(0, 640 * 2 - 1)
w, h = image.size
target["boxes"] = convert_to_tv_tensor(
target["boxes"], "boxes", box_format="xyxy", spatial_size=[h, w]
)
if "masks" in target:
target["masks"] = convert_to_tv_tensor(target["masks"], "masks")
image, target = self.random_affine(image, target)
# image, target = self.resize(image, target)
image, target = self.crop(image, target)
return image, target, dataset

View File

@ -0,0 +1,4 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""

View File

@ -0,0 +1,9 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from .dist_utils import setup_print, setup_seed
from .logger import *
from .profiler_utils import stats
from .visualizer import *

View File

@ -0,0 +1,106 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from typing import List, Tuple
import torch
import torchvision
from torch import Tensor
def generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
return torchvision.ops.generalized_box_iou(boxes1, boxes2)
# elementwise
def elementwise_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
"""
Args:
boxes1, [N, 4]
boxes2, [N, 4]
Returns:
iou, [N, ]
union, [N, ]
"""
area1 = torchvision.ops.box_area(boxes1) # [N, ]
area2 = torchvision.ops.box_area(boxes2) # [N, ]
lt = torch.max(boxes1[:, :2], boxes2[:, :2]) # [N, 2]
rb = torch.min(boxes1[:, 2:], boxes2[:, 2:]) # [N, 2]
wh = (rb - lt).clamp(min=0) # [N, 2]
inter = wh[:, 0] * wh[:, 1] # [N, ]
union = area1 + area2 - inter
iou = inter / union
return iou, union
def elementwise_generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
"""
Args:
boxes1, [N, 4] with [x1, y1, x2, y2]
boxes2, [N, 4] with [x1, y1, x2, y2]
Returns:
giou, [N, ]
"""
assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
iou, union = elementwise_box_iou(boxes1, boxes2)
lt = torch.min(boxes1[:, :2], boxes2[:, :2]) # [N, 2]
rb = torch.max(boxes1[:, 2:], boxes2[:, 2:]) # [N, 2]
wh = (rb - lt).clamp(min=0) # [N, 2]
area = wh[:, 0] * wh[:, 1]
return iou - (area - union) / area
def check_point_inside_box(points: Tensor, boxes: Tensor, eps=1e-9) -> Tensor:
"""
Args:
points, [K, 2], (x, y)
boxes, [N, 4], (x1, y1, y2, y2)
Returns:
Tensor (bool), [K, N]
"""
x, y = [p.unsqueeze(-1) for p in points.unbind(-1)]
x1, y1, x2, y2 = [x.unsqueeze(0) for x in boxes.unbind(-1)]
l = x - x1
t = y - y1
r = x2 - x
b = y2 - y
ltrb = torch.stack([l, t, r, b], dim=-1)
mask = ltrb.min(dim=-1).values > eps
return mask
def point_box_distance(points: Tensor, boxes: Tensor) -> Tensor:
"""
Args:
boxes, [N, 4], (x1, y1, x2, y2)
points, [N, 2], (x, y)
Returns:
Tensor (N, 4), (l, t, r, b)
"""
x1y1, x2y2 = torch.split(boxes, 2, dim=-1)
lt = points - x1y1
rb = x2y2 - points
return torch.concat([lt, rb], dim=-1)
def point_distance_box(points: Tensor, distances: Tensor) -> Tensor:
"""
Args:
points (Tensor), [N, 2], (x, y)
distances (Tensor), [N, 4], (l, t, r, b)
Returns:
boxes (Tensor), (N, 4), (x1, y1, x2, y2)
"""
lt, rb = torch.split(distances, 2, dim=-1)
x1y1 = -lt + points
x2y2 = rb + points
boxes = torch.concat([x1y1, x2y2], dim=-1)
return boxes

View File

@ -0,0 +1,281 @@
"""
reference
- https://github.com/pytorch/vision/blob/main/references/detection/utils.py
- https://github.com/facebookresearch/detr/blob/master/util/misc.py#L406
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import atexit
import os
import random
import time
import numpy as np
import torch
import torch.backends.cudnn
import torch.distributed
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.nn.parallel import DataParallel as DP
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DistributedSampler
# from torch.utils.data.dataloader import DataLoader
from ..data import DataLoader
def setup_distributed(
print_rank: int = 0,
print_method: str = "builtin",
seed: int = None,
):
"""
env setup
args:
print_rank,
print_method, (builtin, rich)
seed,
"""
try:
# https://pytorch.org/docs/stable/elastic/run.html
RANK = int(os.getenv("RANK", -1))
LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))
WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
# torch.distributed.init_process_group(backend=backend, init_method='env://')
torch.distributed.init_process_group(init_method="env://")
torch.distributed.barrier()
rank = torch.distributed.get_rank()
torch.cuda.set_device(rank)
torch.cuda.empty_cache()
enabled_dist = True
if get_rank() == print_rank:
print("Initialized distributed mode...")
except Exception:
enabled_dist = False
print("Not init distributed mode.")
setup_print(get_rank() == print_rank, method=print_method)
if seed is not None:
setup_seed(seed)
return enabled_dist
def setup_print(is_main, method="builtin"):
"""This function disables printing when not in master process"""
import builtins as __builtin__
if method == "builtin":
builtin_print = __builtin__.print
elif method == "rich":
import rich
builtin_print = rich.print
else:
raise AttributeError("")
def print(*args, **kwargs):
force = kwargs.pop("force", False)
if is_main or force:
builtin_print(*args, **kwargs)
__builtin__.print = print
def is_dist_available_and_initialized():
if not torch.distributed.is_available():
return False
if not torch.distributed.is_initialized():
return False
return True
@atexit.register
def cleanup():
"""cleanup distributed environment"""
if is_dist_available_and_initialized():
torch.distributed.barrier()
torch.distributed.destroy_process_group()
def get_rank():
if not is_dist_available_and_initialized():
return 0
return torch.distributed.get_rank()
def get_world_size():
if not is_dist_available_and_initialized():
return 1
return torch.distributed.get_world_size()
def is_main_process():
return get_rank() == 0
def save_on_master(*args, **kwargs):
if is_main_process():
torch.save(*args, **kwargs)
def warp_model(
model: torch.nn.Module,
sync_bn: bool = False,
dist_mode: str = "ddp",
find_unused_parameters: bool = False,
compile: bool = False,
compile_mode: str = "reduce-overhead",
**kwargs,
):
if is_dist_available_and_initialized():
rank = get_rank()
model = nn.SyncBatchNorm.convert_sync_batchnorm(model) if sync_bn else model
if dist_mode == "dp":
model = DP(model, device_ids=[rank], output_device=rank)
elif dist_mode == "ddp":
model = DDP(
model,
device_ids=[rank],
output_device=rank,
find_unused_parameters=find_unused_parameters,
)
else:
raise AttributeError("")
if compile:
model = torch.compile(model, mode=compile_mode)
return model
def de_model(model):
return de_parallel(de_complie(model))
def warp_loader(loader, shuffle=False):
if is_dist_available_and_initialized():
sampler = DistributedSampler(loader.dataset, shuffle=shuffle)
loader = DataLoader(
loader.dataset,
loader.batch_size,
sampler=sampler,
drop_last=loader.drop_last,
collate_fn=loader.collate_fn,
pin_memory=loader.pin_memory,
num_workers=loader.num_workers,
)
return loader
def is_parallel(model) -> bool:
# Returns True if model is of type DP or DDP
return type(model) in (
torch.nn.parallel.DataParallel,
torch.nn.parallel.DistributedDataParallel,
)
def de_parallel(model) -> nn.Module:
# De-parallelize a model: returns single-GPU model if model is of type DP or DDP
return model.module if is_parallel(model) else model
def reduce_dict(data, avg=True):
"""
Args
data dict: input, {k: v, ...}
avg bool: true
"""
world_size = get_world_size()
if world_size < 2:
return data
with torch.no_grad():
keys, values = [], []
for k in sorted(data.keys()):
keys.append(k)
values.append(data[k])
values = torch.stack(values, dim=0)
torch.distributed.all_reduce(values)
if avg is True:
values /= world_size
return {k: v for k, v in zip(keys, values)}
def all_gather(data):
"""
Run all_gather on arbitrary picklable data (not necessarily tensors)
Args:
data: any picklable object
Returns:
list[data]: list of data gathered from each rank
"""
world_size = get_world_size()
if world_size == 1:
return [data]
data_list = [None] * world_size
torch.distributed.all_gather_object(data_list, data)
return data_list
def sync_time():
"""sync_time"""
if torch.cuda.is_available():
torch.cuda.synchronize()
return time.time()
def setup_seed(seed: int, deterministic=False):
"""setup_seed for reproducibility
torch.manual_seed(3407) is all you need. https://arxiv.org/abs/2109.08203
"""
seed = seed + get_rank()
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
# memory will be large when setting deterministic to True
if torch.backends.cudnn.is_available() and deterministic:
torch.backends.cudnn.deterministic = True
# for torch.compile
def check_compile():
import warnings
import torch
gpu_ok = False
if torch.cuda.is_available():
device_cap = torch.cuda.get_device_capability()
if device_cap in ((7, 0), (8, 0), (9, 0)):
gpu_ok = True
if not gpu_ok:
warnings.warn(
"GPU is not NVIDIA V100, A100, or H100. Speedup numbers may be lower " "than expected."
)
return gpu_ok
def is_compile(model):
import torch._dynamo
return type(model) in (torch._dynamo.OptimizedModule,)
def de_complie(model):
return model._orig_mod if is_compile(model) else model

View File

@ -0,0 +1,70 @@
"""
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/util/lazy_loader.py
"""
import importlib
import types
class LazyLoader(types.ModuleType):
"""Lazily import a module, mainly to avoid pulling in large dependencies.
`paddle`, and `ffmpeg` are examples of modules that are large and not always
needed, and this allows them to only be loaded when they are used.
"""
# The lint error here is incorrect.
def __init__(self, local_name, parent_module_globals, name, warning=None):
self._local_name = local_name
self._parent_module_globals = parent_module_globals
self._warning = warning
# These members allows doctest correctly process this module member without
# triggering self._load(). self._load() mutates parant_module_globals and
# triggers a dict mutated during iteration error from doctest.py.
# - for from_module()
self.__module__ = name.rsplit(".", 1)[0]
# - for is_routine()
self.__wrapped__ = None
super(LazyLoader, self).__init__(name)
def _load(self):
"""Load the module and insert it into the parent's globals."""
# Import the target module and insert it into the parent's namespace
module = importlib.import_module(self.__name__)
self._parent_module_globals[self._local_name] = module
# Emit a warning if one was specified
if self._warning:
# logging.warning(self._warning)
# Make sure to only warn once.
self._warning = None
# Update this object's dict so that if someone keeps a reference to the
# LazyLoader, lookups are efficient (__getattr__ is only called on lookups
# that fail).
self.__dict__.update(module.__dict__)
return module
def __getattr__(self, item):
module = self._load()
return getattr(module, item)
def __repr__(self):
# Carefully to not trigger _load, since repr may be called in very
# sensitive places.
return f"<LazyLoader {self.__name__} as {self._local_name}>"
def __dir__(self):
module = self._load()
return dir(module)
# import paddle.nn as nn
# nn = LazyLoader("nn", globals(), "paddle.nn")
# class M(nn.Layer):
# def __init__(self) -> None:
# super().__init__()

View File

@ -0,0 +1,255 @@
"""
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
https://github.com/facebookresearch/detr/blob/main/util/misc.py
Mostly copy-paste from torchvision references.
"""
import datetime
import pickle
import time
from collections import defaultdict, deque
from typing import Dict
import torch
import torch.distributed as tdist
from .dist_utils import get_world_size, is_dist_available_and_initialized
class SmoothedValue(object):
"""Track a series of values and provide access to smoothed values over a
window or the global series average.
"""
def __init__(self, window_size=20, fmt=None):
if fmt is None:
fmt = "{median:.4f} ({global_avg:.4f})"
self.deque = deque(maxlen=window_size)
self.total = 0.0
self.count = 0
self.fmt = fmt
def update(self, value, n=1):
self.deque.append(value)
self.count += n
self.total += value * n
def synchronize_between_processes(self):
"""
Warning: does not synchronize the deque!
"""
if not is_dist_available_and_initialized():
return
t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
tdist.barrier()
tdist.all_reduce(t)
t = t.tolist()
self.count = int(t[0])
self.total = t[1]
@property
def median(self):
d = torch.tensor(list(self.deque))
return d.median().item()
@property
def avg(self):
d = torch.tensor(list(self.deque), dtype=torch.float32)
return d.mean().item()
@property
def global_avg(self):
return self.total / self.count
@property
def max(self):
return max(self.deque)
@property
def value(self):
return self.deque[-1]
def __str__(self):
return self.fmt.format(
median=self.median,
avg=self.avg,
global_avg=self.global_avg,
max=self.max,
value=self.value,
)
def all_gather(data):
"""
Run all_gather on arbitrary picklable data (not necessarily tensors)
Args:
data: any picklable object
Returns:
list[data]: list of data gathered from each rank
"""
world_size = get_world_size()
if world_size == 1:
return [data]
# serialized to a Tensor
buffer = pickle.dumps(data)
storage = torch.ByteStorage.from_buffer(buffer)
tensor = torch.ByteTensor(storage).to("cuda")
# obtain Tensor size of each rank
local_size = torch.tensor([tensor.numel()], device="cuda")
size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)]
tdist.all_gather(size_list, local_size)
size_list = [int(size.item()) for size in size_list]
max_size = max(size_list)
# receiving Tensor from all ranks
# we pad the tensor because torch all_gather does not support
# gathering tensors of different shapes
tensor_list = []
for _ in size_list:
tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda"))
if local_size != max_size:
padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda")
tensor = torch.cat((tensor, padding), dim=0)
tdist.all_gather(tensor_list, tensor)
data_list = []
for size, tensor in zip(size_list, tensor_list):
buffer = tensor.cpu().numpy().tobytes()[:size]
data_list.append(pickle.loads(buffer))
return data_list
def reduce_dict(input_dict, average=True) -> Dict[str, torch.Tensor]:
"""
Args:
input_dict (dict): all the values will be reduced
average (bool): whether to do average or sum
Reduce the values in the dictionary from all processes so that all processes
have the averaged results. Returns a dict with the same fields as
input_dict, after reduction.
"""
world_size = get_world_size()
if world_size < 2:
return input_dict
with torch.no_grad():
names = []
values = []
# sort the keys so that they are consistent across processes
for k in sorted(input_dict.keys()):
names.append(k)
values.append(input_dict[k])
values = torch.stack(values, dim=0)
tdist.all_reduce(values)
if average:
values /= world_size
reduced_dict = {k: v for k, v in zip(names, values)}
return reduced_dict
class MetricLogger(object):
def __init__(self, delimiter="\t"):
self.meters = defaultdict(SmoothedValue)
self.delimiter = delimiter
def update(self, **kwargs):
for k, v in kwargs.items():
if isinstance(v, torch.Tensor):
v = v.item()
assert isinstance(v, (float, int))
self.meters[k].update(v)
def __getattr__(self, attr):
if attr in self.meters:
return self.meters[attr]
if attr in self.__dict__:
return self.__dict__[attr]
raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr))
def __str__(self):
loss_str = []
for name, meter in self.meters.items():
loss_str.append("{}: {}".format(name, str(meter)))
return self.delimiter.join(loss_str)
def synchronize_between_processes(self):
for meter in self.meters.values():
meter.synchronize_between_processes()
def add_meter(self, name, meter):
self.meters[name] = meter
def log_every(self, iterable, print_freq, header=None):
i = 0
if not header:
header = ""
start_time = time.time()
end = time.time()
iter_time = SmoothedValue(fmt="{avg:.4f}")
data_time = SmoothedValue(fmt="{avg:.4f}")
space_fmt = ":" + str(len(str(len(iterable)))) + "d"
if torch.cuda.is_available():
log_msg = self.delimiter.join(
[
header,
"[{0" + space_fmt + "}/{1}]",
"eta: {eta}",
"{meters}",
"time: {time}",
"data: {data}",
"max mem: {memory:.0f}",
]
)
else:
log_msg = self.delimiter.join(
[
header,
"[{0" + space_fmt + "}/{1}]",
"eta: {eta}",
"{meters}",
"time: {time}",
"data: {data}",
]
)
MB = 1024.0 * 1024.0
for obj in iterable:
data_time.update(time.time() - end)
yield obj
iter_time.update(time.time() - end)
if i % print_freq == 0 or i == len(iterable) - 1:
eta_seconds = iter_time.global_avg * (len(iterable) - i)
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
if torch.cuda.is_available():
print(
log_msg.format(
i,
len(iterable),
eta=eta_string,
meters=str(self),
time=str(iter_time),
data=str(data_time),
memory=torch.cuda.max_memory_allocated() / MB,
)
)
else:
print(
log_msg.format(
i,
len(iterable),
eta=eta_string,
meters=str(self),
time=str(iter_time),
data=str(data_time),
)
)
i += 1
end = time.time()
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print(
"{} Total time: {} ({:.4f} s / it)".format(
header, total_time_str, total_time / len(iterable)
)
)

View File

@ -0,0 +1,30 @@
"""
Copyright (c) 2024 The D-FINE Authors. All Rights Reserved.
"""
import copy
from typing import Tuple
from calflops import calculate_flops
def stats(
cfg,
input_shape: Tuple = (1, 3, 640, 640),
) -> Tuple[int, dict]:
base_size = cfg.train_dataloader.collate_fn.base_size
input_shape = (1, 3, base_size, base_size)
model_for_info = copy.deepcopy(cfg.model).deploy()
flops, macs, _ = calculate_flops(
model=model_for_info,
input_shape=input_shape,
output_as_string=True,
output_precision=4,
print_detailed=False,
)
params = sum(p.numel() for p in model_for_info.parameters())
del model_for_info
return params, {"Model FLOPs:%s MACs:%s Params:%s" % (flops, macs, params)}

View File

@ -0,0 +1,121 @@
""" "
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import PIL
import numpy as np
import torch
import torch.utils.data
import torchvision
from typing import List, Dict
torchvision.disable_beta_transforms_warning()
__all__ = ["show_sample", "save_samples"]
def save_samples(samples: torch.Tensor, targets: List[Dict], output_dir: str, split: str, normalized: bool, box_fmt: str):
'''
normalized: whether the boxes are normalized to [0, 1]
box_fmt: 'xyxy', 'xywh', 'cxcywh', D-FINE uses 'cxcywh' for training, 'xyxy' for validation
'''
from torchvision.transforms.functional import to_pil_image
from torchvision.ops import box_convert
from pathlib import Path
from PIL import ImageDraw, ImageFont
import os
os.makedirs(Path(output_dir) / Path(f"{split}_samples"), exist_ok=True)
# Predefined colors (standard color names recognized by PIL)
BOX_COLORS = [
"red", "blue", "green", "orange", "purple",
"cyan", "magenta", "yellow", "lime", "pink",
"teal", "lavender", "brown", "beige", "maroon",
"navy", "olive", "coral", "turquoise", "gold"
]
LABEL_TEXT_COLOR = "white"
font = ImageFont.load_default()
font.size = 32
for i, (sample, target) in enumerate(zip(samples, targets)):
sample_visualization = sample.clone().cpu()
target_boxes = target["boxes"].clone().cpu()
target_labels = target["labels"].clone().cpu()
target_image_id = target["image_id"].item()
target_image_path = target["image_path"]
target_image_path_stem = Path(target_image_path).stem
sample_visualization = to_pil_image(sample_visualization)
sample_visualization_w, sample_visualization_h = sample_visualization.size
# normalized to pixel space
if normalized:
target_boxes[:, 0] = target_boxes[:, 0] * sample_visualization_w
target_boxes[:, 2] = target_boxes[:, 2] * sample_visualization_w
target_boxes[:, 1] = target_boxes[:, 1] * sample_visualization_h
target_boxes[:, 3] = target_boxes[:, 3] * sample_visualization_h
# any box format -> xyxy
target_boxes = box_convert(target_boxes, in_fmt=box_fmt, out_fmt="xyxy")
# clip to image size
target_boxes[:, 0] = torch.clamp(target_boxes[:, 0], 0, sample_visualization_w)
target_boxes[:, 1] = torch.clamp(target_boxes[:, 1], 0, sample_visualization_h)
target_boxes[:, 2] = torch.clamp(target_boxes[:, 2], 0, sample_visualization_w)
target_boxes[:, 3] = torch.clamp(target_boxes[:, 3], 0, sample_visualization_h)
target_boxes = target_boxes.numpy().astype(np.int32)
target_labels = target_labels.numpy().astype(np.int32)
draw = ImageDraw.Draw(sample_visualization)
# draw target boxes
for box, label in zip(target_boxes, target_labels):
x1, y1, x2, y2 = box
# Select color based on class ID
box_color = BOX_COLORS[int(label) % len(BOX_COLORS)]
# Draw box (thick)
draw.rectangle([x1, y1, x2, y2], outline=box_color, width=3)
label_text = f"{label}"
# Measure text size
text_width, text_height = draw.textbbox((0, 0), label_text, font=font)[2:4]
# Draw text background
padding = 2
draw.rectangle(
[x1, y1 - text_height - padding * 2, x1 + text_width + padding * 2, y1],
fill=box_color
)
# Draw text (LABEL_TEXT_COLOR)
draw.text((x1 + padding, y1 - text_height - padding), label_text,
fill=LABEL_TEXT_COLOR, font=font)
save_path = Path(output_dir) / f"{split}_samples" / f"{target_image_id}_{target_image_path_stem}.webp"
sample_visualization.save(save_path)
def show_sample(sample):
"""for coco dataset/dataloader"""
import matplotlib.pyplot as plt
from torchvision.transforms.v2 import functional as F
from torchvision.utils import draw_bounding_boxes
image, target = sample
if isinstance(image, PIL.Image.Image):
image = F.to_image_tensor(image)
image = F.convert_dtype(image, torch.uint8)
annotated_image = draw_bounding_boxes(image, target["boxes"], colors="yellow", width=3)
fig, ax = plt.subplots()
ax.imshow(annotated_image.permute(1, 2, 0).numpy())
ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
fig.tight_layout()
fig.show()
plt.show()

View File

@ -0,0 +1,16 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from .arch import *
#
from .backbone import *
from .backbone import (
FrozenBatchNorm2d,
freeze_batch_norm2d,
get_activation,
)
from .criterion import *
from .postprocessor import *

View File

@ -0,0 +1,7 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from .classification import ClassHead, Classification
from .yolo import YOLO

View File

@ -0,0 +1,45 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import torch
import torch.nn as nn
from ...core import register
__all__ = ["Classification", "ClassHead"]
@register()
class Classification(torch.nn.Module):
__inject__ = ["backbone", "head"]
def __init__(self, backbone: nn.Module, head: nn.Module = None):
super().__init__()
self.backbone = backbone
self.head = head
def forward(self, x):
x = self.backbone(x)
if self.head is not None:
x = self.head(x)
return x
@register()
class ClassHead(nn.Module):
def __init__(self, hidden_dim, num_classes):
super().__init__()
self.pool = nn.AdaptiveAvgPool2d(1)
self.proj = nn.Linear(hidden_dim, num_classes)
def forward(self, x):
x = x[0] if isinstance(x, (list, tuple)) else x
x = self.pool(x)
x = x.reshape(x.shape[0], -1)
x = self.proj(x)
return x

View File

@ -0,0 +1,42 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import torch
from ...core import register
__all__ = [
"YOLO",
]
@register()
class YOLO(torch.nn.Module):
__inject__ = [
"backbone",
"neck",
"head",
]
def __init__(self, backbone: torch.nn.Module, neck, head):
super().__init__()
self.backbone = backbone
self.neck = neck
self.head = head
def forward(self, x, **kwargs):
x = self.backbone(x)
x = self.neck(x)
x = self.head(x)
return x
def deploy(
self,
):
self.eval()
for m in self.modules():
if m is not self and hasattr(m, "deploy"):
m.deploy()
return self

View File

@ -0,0 +1,17 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from .common import (
FrozenBatchNorm2d,
freeze_batch_norm2d,
get_activation,
)
from .csp_darknet import CSPPAN, CSPDarkNet
from .csp_resnet import CSPResNet
from .hgnetv2 import HGNetv2
from .presnet import PResNet
from .test_resnet import MResNet
from .timm_model import TimmModel
from .torchvision_model import TorchVisionModel

View File

@ -0,0 +1,117 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import torch
import torch.nn as nn
class ConvNormLayer(nn.Module):
def __init__(self, ch_in, ch_out, kernel_size, stride, padding=None, bias=False, act=None):
super().__init__()
self.conv = nn.Conv2d(
ch_in,
ch_out,
kernel_size,
stride,
padding=(kernel_size - 1) // 2 if padding is None else padding,
bias=bias,
)
self.norm = nn.BatchNorm2d(ch_out)
self.act = nn.Identity() if act is None else get_activation(act)
def forward(self, x):
return self.act(self.norm(self.conv(x)))
class FrozenBatchNorm2d(nn.Module):
"""copy and modified from https://github.com/facebookresearch/detr/blob/master/models/backbone.py
BatchNorm2d where the batch statistics and the affine parameters are fixed.
Copy-paste from torchvision.misc.ops with added eps before rqsrt,
without which any other models than torchvision.models.resnet[18,34,50,101]
produce nans.
"""
def __init__(self, num_features, eps=1e-5):
super(FrozenBatchNorm2d, self).__init__()
n = num_features
self.register_buffer("weight", torch.ones(n))
self.register_buffer("bias", torch.zeros(n))
self.register_buffer("running_mean", torch.zeros(n))
self.register_buffer("running_var", torch.ones(n))
self.eps = eps
self.num_features = n
def _load_from_state_dict(
self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
):
num_batches_tracked_key = prefix + "num_batches_tracked"
if num_batches_tracked_key in state_dict:
del state_dict[num_batches_tracked_key]
super(FrozenBatchNorm2d, self)._load_from_state_dict(
state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
)
def forward(self, x):
# move reshapes to the beginning
# to make it fuser-friendly
w = self.weight.reshape(1, -1, 1, 1)
b = self.bias.reshape(1, -1, 1, 1)
rv = self.running_var.reshape(1, -1, 1, 1)
rm = self.running_mean.reshape(1, -1, 1, 1)
scale = w * (rv + self.eps).rsqrt()
bias = b - rm * scale
return x * scale + bias
def extra_repr(self):
return "{num_features}, eps={eps}".format(**self.__dict__)
def freeze_batch_norm2d(module: nn.Module) -> nn.Module:
if isinstance(module, nn.BatchNorm2d):
module = FrozenBatchNorm2d(module.num_features)
else:
for name, child in module.named_children():
_child = freeze_batch_norm2d(child)
if _child is not child:
setattr(module, name, _child)
return module
def get_activation(act: str, inplace: bool = True):
"""get activation"""
if act is None:
return nn.Identity()
elif isinstance(act, nn.Module):
return act
act = act.lower()
if act == "silu" or act == "swish":
m = nn.SiLU()
elif act == "relu":
m = nn.ReLU()
elif act == "leaky_relu":
m = nn.LeakyReLU()
elif act == "silu":
m = nn.SiLU()
elif act == "gelu":
m = nn.GELU()
elif act == "hardsigmoid":
m = nn.Hardsigmoid()
else:
raise RuntimeError("")
if hasattr(m, "inplace"):
m.inplace = inplace
return m

View File

@ -0,0 +1,203 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import math
import warnings
import torch
import torch.nn as nn
import torch.nn.functional as F
from ...core import register
from .common import get_activation
def autopad(k, p=None):
if p is None:
p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
return p
def make_divisible(c, d):
return math.ceil(c / d) * d
class Conv(nn.Module):
def __init__(self, cin, cout, k=1, s=1, p=None, g=1, act="silu") -> None:
super().__init__()
self.conv = nn.Conv2d(cin, cout, k, s, autopad(k, p), groups=g, bias=False)
self.bn = nn.BatchNorm2d(cout)
self.act = get_activation(act, inplace=True)
def forward(self, x):
return self.act(self.bn(self.conv(x)))
class Bottleneck(nn.Module):
# Standard bottleneck
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5, act="silu"):
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1, act=act)
self.cv2 = Conv(c_, c2, 3, 1, g=g, act=act)
self.add = shortcut and c1 == c2
def forward(self, x):
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, act="silu"
): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1, act=act)
self.cv2 = Conv(c1, c_, 1, 1, act=act)
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0, act=act) for _ in range(n)))
self.cv3 = Conv(2 * c_, c2, 1, act=act)
def forward(self, x):
return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
class SPPF(nn.Module):
# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
def __init__(self, c1, c2, k=5, act="silu"): # equivalent to SPP(k=(5, 9, 13))
super().__init__()
c_ = c1 // 2 # hidden channels
self.cv1 = Conv(c1, c_, 1, 1, act=act)
self.cv2 = Conv(c_ * 4, c2, 1, 1, act=act)
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
def forward(self, x):
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))
@register()
class CSPDarkNet(nn.Module):
__share__ = ["depth_multi", "width_multi"]
def __init__(
self,
in_channels=3,
width_multi=1.0,
depth_multi=1.0,
return_idx=[2, 3, -1],
act="silu",
) -> None:
super().__init__()
channels = [64, 128, 256, 512, 1024]
channels = [make_divisible(c * width_multi, 8) for c in channels]
depths = [3, 6, 9, 3]
depths = [max(round(d * depth_multi), 1) for d in depths]
self.layers = nn.ModuleList([Conv(in_channels, channels[0], 6, 2, 2, act=act)])
for i, (c, d) in enumerate(zip(channels, depths), 1):
layer = nn.Sequential(
*[Conv(c, channels[i], 3, 2, act=act), C3(channels[i], channels[i], n=d, act=act)]
)
self.layers.append(layer)
self.layers.append(SPPF(channels[-1], channels[-1], k=5, act=act))
self.return_idx = return_idx
self.out_channels = [channels[i] for i in self.return_idx]
self.strides = [[2, 4, 8, 16, 32][i] for i in self.return_idx]
self.depths = depths
self.act = act
def forward(self, x):
outputs = []
for _, m in enumerate(self.layers):
x = m(x)
outputs.append(x)
return [outputs[i] for i in self.return_idx]
@register()
class CSPPAN(nn.Module):
"""
P5 ---> 1x1 ---------------------------------> concat --> c3 --> det
| up | conv /2
P4 ---> concat ---> c3 ---> 1x1 --> concat ---> c3 -----------> det
| up | conv /2
P3 -----------------------> concat ---> c3 ---------------------> det
"""
__share__ = [
"depth_multi",
]
def __init__(self, in_channels=[256, 512, 1024], depth_multi=1.0, act="silu") -> None:
super().__init__()
depth = max(round(3 * depth_multi), 1)
self.out_channels = in_channels
self.fpn_stems = nn.ModuleList(
[
Conv(cin, cout, 1, 1, act=act)
for cin, cout in zip(in_channels[::-1], in_channels[::-1][1:])
]
)
self.fpn_csps = nn.ModuleList(
[
C3(cin, cout, depth, False, act=act)
for cin, cout in zip(in_channels[::-1], in_channels[::-1][1:])
]
)
self.pan_stems = nn.ModuleList([Conv(c, c, 3, 2, act=act) for c in in_channels[:-1]])
self.pan_csps = nn.ModuleList([C3(c, c, depth, False, act=act) for c in in_channels[1:]])
def forward(self, feats):
fpn_feats = []
for i, feat in enumerate(feats[::-1]):
if i == 0:
feat = self.fpn_stems[i](feat)
fpn_feats.append(feat)
else:
_feat = F.interpolate(fpn_feats[-1], scale_factor=2, mode="nearest")
feat = torch.concat([_feat, feat], dim=1)
feat = self.fpn_csps[i - 1](feat)
if i < len(self.fpn_stems):
feat = self.fpn_stems[i](feat)
fpn_feats.append(feat)
pan_feats = []
for i, feat in enumerate(fpn_feats[::-1]):
if i == 0:
pan_feats.append(feat)
else:
_feat = self.pan_stems[i - 1](pan_feats[-1])
feat = torch.concat([_feat, feat], dim=1)
feat = self.pan_csps[i - 1](feat)
pan_feats.append(feat)
return pan_feats
if __name__ == "__main__":
data = torch.rand(1, 3, 320, 640)
width_multi = 0.75
depth_multi = 0.33
m = CSPDarkNet(3, width_multi=width_multi, depth_multi=depth_multi, act="silu")
outputs = m(data)
print([o.shape for o in outputs])
m = CSPPAN(in_channels=m.out_channels, depth_multi=depth_multi, act="silu")
outputs = m(outputs)
print([o.shape for o in outputs])

View File

@ -0,0 +1,302 @@
"""
https://github.com/PaddlePaddle/PaddleDetection/blob/release/2.6/ppdet/modeling/backbones/cspresnet.py
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from collections import OrderedDict
import torch
import torch.nn as nn
import torch.nn.functional as F
from ...core import register
from .common import get_activation
__all__ = ["CSPResNet"]
donwload_url = {
"s": "https://github.com/lyuwenyu/storage/releases/download/v0.1/CSPResNetb_s_pretrained_from_paddle.pth",
"m": "https://github.com/lyuwenyu/storage/releases/download/v0.1/CSPResNetb_m_pretrained_from_paddle.pth",
"l": "https://github.com/lyuwenyu/storage/releases/download/v0.1/CSPResNetb_l_pretrained_from_paddle.pth",
"x": "https://github.com/lyuwenyu/storage/releases/download/v0.1/CSPResNetb_x_pretrained_from_paddle.pth",
}
class ConvBNLayer(nn.Module):
def __init__(self, ch_in, ch_out, filter_size=3, stride=1, groups=1, padding=0, act=None):
super().__init__()
self.conv = nn.Conv2d(
ch_in, ch_out, filter_size, stride, padding, groups=groups, bias=False
)
self.bn = nn.BatchNorm2d(ch_out)
self.act = get_activation(act)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.conv(x)
x = self.bn(x)
x = self.act(x)
return x
class RepVggBlock(nn.Module):
def __init__(self, ch_in, ch_out, act="relu", alpha: bool = False):
super().__init__()
self.ch_in = ch_in
self.ch_out = ch_out
self.conv1 = ConvBNLayer(ch_in, ch_out, 3, stride=1, padding=1, act=None)
self.conv2 = ConvBNLayer(ch_in, ch_out, 1, stride=1, padding=0, act=None)
self.act = get_activation(act)
if alpha:
self.alpha = nn.Parameter(
torch.ones(
1,
)
)
else:
self.alpha = None
def forward(self, x):
if hasattr(self, "conv"):
y = self.conv(x)
else:
if self.alpha:
y = self.conv1(x) + self.alpha * self.conv2(x)
else:
y = self.conv1(x) + self.conv2(x)
y = self.act(y)
return y
def convert_to_deploy(self):
if not hasattr(self, "conv"):
self.conv = nn.Conv2d(self.ch_in, self.ch_out, 3, 1, padding=1)
kernel, bias = self.get_equivalent_kernel_bias()
self.conv.weight.data = kernel
self.conv.bias.data = bias
def get_equivalent_kernel_bias(self):
kernel3x3, bias3x3 = self._fuse_bn_tensor(self.conv1)
kernel1x1, bias1x1 = self._fuse_bn_tensor(self.conv2)
if self.alpha:
return kernel3x3 + self.alpha * self._pad_1x1_to_3x3_tensor(
kernel1x1
), bias3x3 + self.alpha * bias1x1
else:
return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1), bias3x3 + bias1x1
def _pad_1x1_to_3x3_tensor(self, kernel1x1):
if kernel1x1 is None:
return 0
else:
return F.pad(kernel1x1, [1, 1, 1, 1])
def _fuse_bn_tensor(self, branch: ConvBNLayer):
if branch is None:
return 0, 0
kernel = branch.conv.weight
running_mean = branch.norm.running_mean
running_var = branch.norm.running_var
gamma = branch.norm.weight
beta = branch.norm.bias
eps = branch.norm.eps
std = (running_var + eps).sqrt()
t = (gamma / std).reshape(-1, 1, 1, 1)
return kernel * t, beta - running_mean * gamma / std
class BasicBlock(nn.Module):
def __init__(self, ch_in, ch_out, act="relu", shortcut=True, use_alpha=False):
super().__init__()
assert ch_in == ch_out
self.conv1 = ConvBNLayer(ch_in, ch_out, 3, stride=1, padding=1, act=act)
self.conv2 = RepVggBlock(ch_out, ch_out, act=act, alpha=use_alpha)
self.shortcut = shortcut
def forward(self, x):
y = self.conv1(x)
y = self.conv2(y)
if self.shortcut:
return x + y
else:
return y
class EffectiveSELayer(nn.Module):
"""Effective Squeeze-Excitation
From `CenterMask : Real-Time Anchor-Free Instance Segmentation` - https://arxiv.org/abs/1911.06667
"""
def __init__(self, channels, act="hardsigmoid"):
super(EffectiveSELayer, self).__init__()
self.fc = nn.Conv2d(channels, channels, kernel_size=1, padding=0)
self.act = get_activation(act)
def forward(self, x: torch.Tensor):
x_se = x.mean((2, 3), keepdim=True)
x_se = self.fc(x_se)
x_se = self.act(x_se)
return x * x_se
class CSPResStage(nn.Module):
def __init__(self, block_fn, ch_in, ch_out, n, stride, act="relu", attn="eca", use_alpha=False):
super().__init__()
ch_mid = (ch_in + ch_out) // 2
if stride == 2:
self.conv_down = ConvBNLayer(ch_in, ch_mid, 3, stride=2, padding=1, act=act)
else:
self.conv_down = None
self.conv1 = ConvBNLayer(ch_mid, ch_mid // 2, 1, act=act)
self.conv2 = ConvBNLayer(ch_mid, ch_mid // 2, 1, act=act)
self.blocks = nn.Sequential(
*[
block_fn(ch_mid // 2, ch_mid // 2, act=act, shortcut=True, use_alpha=use_alpha)
for i in range(n)
]
)
if attn:
self.attn = EffectiveSELayer(ch_mid, act="hardsigmoid")
else:
self.attn = None
self.conv3 = ConvBNLayer(ch_mid, ch_out, 1, act=act)
def forward(self, x):
if self.conv_down is not None:
x = self.conv_down(x)
y1 = self.conv1(x)
y2 = self.blocks(self.conv2(x))
y = torch.concat([y1, y2], dim=1)
if self.attn is not None:
y = self.attn(y)
y = self.conv3(y)
return y
@register()
class CSPResNet(nn.Module):
layers = [3, 6, 6, 3]
channels = [64, 128, 256, 512, 1024]
model_cfg = {
"s": {
"depth_mult": 0.33,
"width_mult": 0.50,
},
"m": {
"depth_mult": 0.67,
"width_mult": 0.75,
},
"l": {
"depth_mult": 1.00,
"width_mult": 1.00,
},
"x": {
"depth_mult": 1.33,
"width_mult": 1.25,
},
}
def __init__(
self,
name: str,
act="silu",
return_idx=[1, 2, 3],
use_large_stem=True,
use_alpha=False,
pretrained=False,
):
super().__init__()
depth_mult = self.model_cfg[name]["depth_mult"]
width_mult = self.model_cfg[name]["width_mult"]
channels = [max(round(c * width_mult), 1) for c in self.channels]
layers = [max(round(l * depth_mult), 1) for l in self.layers]
act = get_activation(act)
if use_large_stem:
self.stem = nn.Sequential(
OrderedDict(
[
(
"conv1",
ConvBNLayer(3, channels[0] // 2, 3, stride=2, padding=1, act=act),
),
(
"conv2",
ConvBNLayer(
channels[0] // 2, channels[0] // 2, 3, stride=1, padding=1, act=act
),
),
(
"conv3",
ConvBNLayer(
channels[0] // 2, channels[0], 3, stride=1, padding=1, act=act
),
),
]
)
)
else:
self.stem = nn.Sequential(
OrderedDict(
[
(
"conv1",
ConvBNLayer(3, channels[0] // 2, 3, stride=2, padding=1, act=act),
),
(
"conv2",
ConvBNLayer(
channels[0] // 2, channels[0], 3, stride=1, padding=1, act=act
),
),
]
)
)
n = len(channels) - 1
self.stages = nn.Sequential(
OrderedDict(
[
(
str(i),
CSPResStage(
BasicBlock,
channels[i],
channels[i + 1],
layers[i],
2,
act=act,
use_alpha=use_alpha,
),
)
for i in range(n)
]
)
)
self._out_channels = channels[1:]
self._out_strides = [4 * 2**i for i in range(n)]
self.return_idx = return_idx
if pretrained:
if isinstance(pretrained, bool) or "http" in pretrained:
state = torch.hub.load_state_dict_from_url(donwload_url[name], map_location="cpu")
else:
state = torch.load(pretrained, map_location="cpu")
self.load_state_dict(state)
print(f"Load CSPResNet_{name} state_dict")
def forward(self, x):
x = self.stem(x)
outs = []
for idx, stage in enumerate(self.stages):
x = stage(x)
if idx in self.return_idx:
outs.append(x)
return outs

View File

@ -0,0 +1,579 @@
"""
reference
- https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
Copyright (c) 2024 The D-FINE Authors. All Rights Reserved.
"""
import logging
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from ...core import register
from .common import FrozenBatchNorm2d
# Constants for initialization
kaiming_normal_ = nn.init.kaiming_normal_
zeros_ = nn.init.zeros_
ones_ = nn.init.ones_
__all__ = ["HGNetv2"]
def safe_barrier():
if torch.distributed.is_available() and torch.distributed.is_initialized():
torch.distributed.barrier()
else:
pass
def safe_get_rank():
if torch.distributed.is_available() and torch.distributed.is_initialized():
return torch.distributed.get_rank()
else:
return 0
class LearnableAffineBlock(nn.Module):
def __init__(self, scale_value=1.0, bias_value=0.0):
super().__init__()
self.scale = nn.Parameter(torch.tensor([scale_value]), requires_grad=True)
self.bias = nn.Parameter(torch.tensor([bias_value]), requires_grad=True)
def forward(self, x):
return self.scale * x + self.bias
class ConvBNAct(nn.Module):
def __init__(
self,
in_chs,
out_chs,
kernel_size,
stride=1,
groups=1,
padding="",
use_act=True,
use_lab=False,
):
super().__init__()
self.use_act = use_act
self.use_lab = use_lab
if padding == "same":
self.conv = nn.Sequential(
nn.ZeroPad2d([0, 1, 0, 1]),
nn.Conv2d(in_chs, out_chs, kernel_size, stride, groups=groups, bias=False),
)
else:
self.conv = nn.Conv2d(
in_chs,
out_chs,
kernel_size,
stride,
padding=(kernel_size - 1) // 2,
groups=groups,
bias=False,
)
self.bn = nn.BatchNorm2d(out_chs)
if self.use_act:
self.act = nn.ReLU()
else:
self.act = nn.Identity()
if self.use_act and self.use_lab:
self.lab = LearnableAffineBlock()
else:
self.lab = nn.Identity()
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
x = self.act(x)
x = self.lab(x)
return x
class LightConvBNAct(nn.Module):
def __init__(
self,
in_chs,
out_chs,
kernel_size,
groups=1,
use_lab=False,
):
super().__init__()
self.conv1 = ConvBNAct(
in_chs,
out_chs,
kernel_size=1,
use_act=False,
use_lab=use_lab,
)
self.conv2 = ConvBNAct(
out_chs,
out_chs,
kernel_size=kernel_size,
groups=out_chs,
use_act=True,
use_lab=use_lab,
)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
return x
class StemBlock(nn.Module):
# for HGNetv2
def __init__(self, in_chs, mid_chs, out_chs, use_lab=False):
super().__init__()
self.stem1 = ConvBNAct(
in_chs,
mid_chs,
kernel_size=3,
stride=2,
use_lab=use_lab,
)
self.stem2a = ConvBNAct(
mid_chs,
mid_chs // 2,
kernel_size=2,
stride=1,
use_lab=use_lab,
)
self.stem2b = ConvBNAct(
mid_chs // 2,
mid_chs,
kernel_size=2,
stride=1,
use_lab=use_lab,
)
self.stem3 = ConvBNAct(
mid_chs * 2,
mid_chs,
kernel_size=3,
stride=2,
use_lab=use_lab,
)
self.stem4 = ConvBNAct(
mid_chs,
out_chs,
kernel_size=1,
stride=1,
use_lab=use_lab,
)
self.pool = nn.MaxPool2d(kernel_size=2, stride=1, ceil_mode=True)
def forward(self, x):
x = self.stem1(x)
x = F.pad(x, (0, 1, 0, 1))
x2 = self.stem2a(x)
x2 = F.pad(x2, (0, 1, 0, 1))
x2 = self.stem2b(x2)
x1 = self.pool(x)
x = torch.cat([x1, x2], dim=1)
x = self.stem3(x)
x = self.stem4(x)
return x
class EseModule(nn.Module):
def __init__(self, chs):
super().__init__()
self.conv = nn.Conv2d(
chs,
chs,
kernel_size=1,
stride=1,
padding=0,
)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
identity = x
x = x.mean((2, 3), keepdim=True)
x = self.conv(x)
x = self.sigmoid(x)
return torch.mul(identity, x)
class HG_Block(nn.Module):
def __init__(
self,
in_chs,
mid_chs,
out_chs,
layer_num,
kernel_size=3,
residual=False,
light_block=False,
use_lab=False,
agg="ese",
drop_path=0.0,
):
super().__init__()
self.residual = residual
self.layers = nn.ModuleList()
for i in range(layer_num):
if light_block:
self.layers.append(
LightConvBNAct(
in_chs if i == 0 else mid_chs,
mid_chs,
kernel_size=kernel_size,
use_lab=use_lab,
)
)
else:
self.layers.append(
ConvBNAct(
in_chs if i == 0 else mid_chs,
mid_chs,
kernel_size=kernel_size,
stride=1,
use_lab=use_lab,
)
)
# feature aggregation
total_chs = in_chs + layer_num * mid_chs
if agg == "se":
aggregation_squeeze_conv = ConvBNAct(
total_chs,
out_chs // 2,
kernel_size=1,
stride=1,
use_lab=use_lab,
)
aggregation_excitation_conv = ConvBNAct(
out_chs // 2,
out_chs,
kernel_size=1,
stride=1,
use_lab=use_lab,
)
self.aggregation = nn.Sequential(
aggregation_squeeze_conv,
aggregation_excitation_conv,
)
else:
aggregation_conv = ConvBNAct(
total_chs,
out_chs,
kernel_size=1,
stride=1,
use_lab=use_lab,
)
att = EseModule(out_chs)
self.aggregation = nn.Sequential(
aggregation_conv,
att,
)
self.drop_path = nn.Dropout(drop_path) if drop_path else nn.Identity()
def forward(self, x):
identity = x
output = [x]
for layer in self.layers:
x = layer(x)
output.append(x)
x = torch.cat(output, dim=1)
x = self.aggregation(x)
if self.residual:
x = self.drop_path(x) + identity
return x
class HG_Stage(nn.Module):
def __init__(
self,
in_chs,
mid_chs,
out_chs,
block_num,
layer_num,
downsample=True,
light_block=False,
kernel_size=3,
use_lab=False,
agg="se",
drop_path=0.0,
):
super().__init__()
self.downsample = downsample
if downsample:
self.downsample = ConvBNAct(
in_chs,
in_chs,
kernel_size=3,
stride=2,
groups=in_chs,
use_act=False,
use_lab=use_lab,
)
else:
self.downsample = nn.Identity()
blocks_list = []
for i in range(block_num):
blocks_list.append(
HG_Block(
in_chs if i == 0 else out_chs,
mid_chs,
out_chs,
layer_num,
residual=False if i == 0 else True,
kernel_size=kernel_size,
light_block=light_block,
use_lab=use_lab,
agg=agg,
drop_path=drop_path[i] if isinstance(drop_path, (list, tuple)) else drop_path,
)
)
self.blocks = nn.Sequential(*blocks_list)
def forward(self, x):
x = self.downsample(x)
x = self.blocks(x)
return x
@register()
class HGNetv2(nn.Module):
"""
HGNetV2
Args:
stem_channels: list. Number of channels for the stem block.
stage_type: str. The stage configuration of HGNet. such as the number of channels, stride, etc.
use_lab: boolean. Whether to use LearnableAffineBlock in network.
lr_mult_list: list. Control the learning rate of different stages.
Returns:
model: nn.Layer. Specific HGNetV2 model depends on args.
"""
arch_configs = {
"B0": {
"stem_channels": [3, 16, 16],
"stage_config": {
# in_channels, mid_channels, out_channels, num_blocks, downsample, light_block, kernel_size, layer_num
"stage1": [16, 16, 64, 1, False, False, 3, 3],
"stage2": [64, 32, 256, 1, True, False, 3, 3],
"stage3": [256, 64, 512, 2, True, True, 5, 3],
"stage4": [512, 128, 1024, 1, True, True, 5, 3],
},
"url": "https://github.com/Peterande/storage/releases/download/dfinev1.0/PPHGNetV2_B0_stage1.pth",
},
"B1": {
"stem_channels": [3, 24, 32],
"stage_config": {
# in_channels, mid_channels, out_channels, num_blocks, downsample, light_block, kernel_size, layer_num
"stage1": [32, 32, 64, 1, False, False, 3, 3],
"stage2": [64, 48, 256, 1, True, False, 3, 3],
"stage3": [256, 96, 512, 2, True, True, 5, 3],
"stage4": [512, 192, 1024, 1, True, True, 5, 3],
},
"url": "https://github.com/Peterande/storage/releases/download/dfinev1.0/PPHGNetV2_B1_stage1.pth",
},
"B2": {
"stem_channels": [3, 24, 32],
"stage_config": {
# in_channels, mid_channels, out_channels, num_blocks, downsample, light_block, kernel_size, layer_num
"stage1": [32, 32, 96, 1, False, False, 3, 4],
"stage2": [96, 64, 384, 1, True, False, 3, 4],
"stage3": [384, 128, 768, 3, True, True, 5, 4],
"stage4": [768, 256, 1536, 1, True, True, 5, 4],
},
"url": "https://github.com/Peterande/storage/releases/download/dfinev1.0/PPHGNetV2_B2_stage1.pth",
},
"B3": {
"stem_channels": [3, 24, 32],
"stage_config": {
# in_channels, mid_channels, out_channels, num_blocks, downsample, light_block, kernel_size, layer_num
"stage1": [32, 32, 128, 1, False, False, 3, 5],
"stage2": [128, 64, 512, 1, True, False, 3, 5],
"stage3": [512, 128, 1024, 3, True, True, 5, 5],
"stage4": [1024, 256, 2048, 1, True, True, 5, 5],
},
"url": "https://github.com/Peterande/storage/releases/download/dfinev1.0/PPHGNetV2_B3_stage1.pth",
},
"B4": {
"stem_channels": [3, 32, 48],
"stage_config": {
# in_channels, mid_channels, out_channels, num_blocks, downsample, light_block, kernel_size, layer_num
"stage1": [48, 48, 128, 1, False, False, 3, 6],
"stage2": [128, 96, 512, 1, True, False, 3, 6],
"stage3": [512, 192, 1024, 3, True, True, 5, 6],
"stage4": [1024, 384, 2048, 1, True, True, 5, 6],
},
"url": "https://github.com/Peterande/storage/releases/download/dfinev1.0/PPHGNetV2_B4_stage1.pth",
},
"B5": {
"stem_channels": [3, 32, 64],
"stage_config": {
# in_channels, mid_channels, out_channels, num_blocks, downsample, light_block, kernel_size, layer_num
"stage1": [64, 64, 128, 1, False, False, 3, 6],
"stage2": [128, 128, 512, 2, True, False, 3, 6],
"stage3": [512, 256, 1024, 5, True, True, 5, 6],
"stage4": [1024, 512, 2048, 2, True, True, 5, 6],
},
"url": "https://github.com/Peterande/storage/releases/download/dfinev1.0/PPHGNetV2_B5_stage1.pth",
},
"B6": {
"stem_channels": [3, 48, 96],
"stage_config": {
# in_channels, mid_channels, out_channels, num_blocks, downsample, light_block, kernel_size, layer_num
"stage1": [96, 96, 192, 2, False, False, 3, 6],
"stage2": [192, 192, 512, 3, True, False, 3, 6],
"stage3": [512, 384, 1024, 6, True, True, 5, 6],
"stage4": [1024, 768, 2048, 3, True, True, 5, 6],
},
"url": "https://github.com/Peterande/storage/releases/download/dfinev1.0/PPHGNetV2_B6_stage1.pth",
},
}
def __init__(
self,
name,
use_lab=False,
return_idx=[1, 2, 3],
freeze_stem_only=True,
freeze_at=0,
freeze_norm=True,
pretrained=True,
local_model_dir="weight/hgnetv2/",
):
super().__init__()
self.use_lab = use_lab
self.return_idx = return_idx
stem_channels = self.arch_configs[name]["stem_channels"]
stage_config = self.arch_configs[name]["stage_config"]
download_url = self.arch_configs[name]["url"]
self._out_strides = [4, 8, 16, 32]
self._out_channels = [stage_config[k][2] for k in stage_config]
# stem
self.stem = StemBlock(
in_chs=stem_channels[0],
mid_chs=stem_channels[1],
out_chs=stem_channels[2],
use_lab=use_lab,
)
# stages
self.stages = nn.ModuleList()
for i, k in enumerate(stage_config):
(
in_channels,
mid_channels,
out_channels,
block_num,
downsample,
light_block,
kernel_size,
layer_num,
) = stage_config[k]
self.stages.append(
HG_Stage(
in_channels,
mid_channels,
out_channels,
block_num,
layer_num,
downsample,
light_block,
kernel_size,
use_lab,
)
)
if freeze_at >= 0:
self._freeze_parameters(self.stem)
if not freeze_stem_only:
for i in range(min(freeze_at + 1, len(self.stages))):
self._freeze_parameters(self.stages[i])
if freeze_norm:
self._freeze_norm(self)
if pretrained:
RED, GREEN, RESET = "\033[91m", "\033[92m", "\033[0m"
try:
# If the file doesn't exist locally, download from the URL
if safe_get_rank() == 0:
print(
GREEN
+ "If the pretrained HGNetV2 can't be downloaded automatically. Please check your network connection."
+ RESET
)
print(
GREEN
+ "Please check your network connection. Or download the model manually from "
+ RESET
+ f"{download_url}"
+ GREEN
+ " to "
+ RESET
+ f"{local_model_dir}."
+ RESET
)
state = torch.hub.load_state_dict_from_url(
download_url, map_location="cpu", model_dir=local_model_dir
)
print(f"Loaded stage1 {name} HGNetV2 from URL.")
# Wait for rank 0 to download the model
safe_barrier()
# All processes load the downloaded model
model_path = local_model_dir + "PPHGNetV2_" + name + "_stage1.pth"
state = torch.load(model_path, map_location="cpu")
self.load_state_dict(state)
print(f"Loaded stage1 {name} HGNetV2 from URL.")
except (Exception, KeyboardInterrupt) as e:
if safe_get_rank() == 0:
print(f"{str(e)}")
logging.error(
RED + "CRITICAL WARNING: Failed to load pretrained HGNetV2 model" + RESET
)
logging.error(
GREEN
+ "Please check your network connection. Or download the model manually from "
+ RESET
+ f"{download_url}"
+ GREEN
+ " to "
+ RESET
+ f"{local_model_dir}."
+ RESET
)
exit()
def _freeze_norm(self, m: nn.Module):
if isinstance(m, nn.BatchNorm2d):
m = FrozenBatchNorm2d(m.num_features)
else:
for name, child in m.named_children():
_child = self._freeze_norm(child)
if _child is not child:
setattr(m, name, _child)
return m
def _freeze_parameters(self, m: nn.Module):
for p in m.parameters():
p.requires_grad = False
def forward(self, x):
x = self.stem(x)
outs = []
for idx, stage in enumerate(self.stages):
x = stage(x)
if idx in self.return_idx:
outs.append(x)
return outs

View File

@ -0,0 +1,263 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from collections import OrderedDict
import torch
import torch.nn as nn
import torch.nn.functional as F
from ...core import register
from .common import FrozenBatchNorm2d, get_activation
__all__ = ["PResNet"]
ResNet_cfg = {
18: [2, 2, 2, 2],
34: [3, 4, 6, 3],
50: [3, 4, 6, 3],
101: [3, 4, 23, 3],
# 152: [3, 8, 36, 3],
}
donwload_url = {
18: "https://github.com/lyuwenyu/storage/releases/download/v0.1/ResNet18_vd_pretrained_from_paddle.pth",
34: "https://github.com/lyuwenyu/storage/releases/download/v0.1/ResNet34_vd_pretrained_from_paddle.pth",
50: "https://github.com/lyuwenyu/storage/releases/download/v0.1/ResNet50_vd_ssld_v2_pretrained_from_paddle.pth",
101: "https://github.com/lyuwenyu/storage/releases/download/v0.1/ResNet101_vd_ssld_pretrained_from_paddle.pth",
}
class ConvNormLayer(nn.Module):
def __init__(self, ch_in, ch_out, kernel_size, stride, padding=None, bias=False, act=None):
super().__init__()
self.conv = nn.Conv2d(
ch_in,
ch_out,
kernel_size,
stride,
padding=(kernel_size - 1) // 2 if padding is None else padding,
bias=bias,
)
self.norm = nn.BatchNorm2d(ch_out)
self.act = get_activation(act)
def forward(self, x):
return self.act(self.norm(self.conv(x)))
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, ch_in, ch_out, stride, shortcut, act="relu", variant="b"):
super().__init__()
self.shortcut = shortcut
if not shortcut:
if variant == "d" and stride == 2:
self.short = nn.Sequential(
OrderedDict(
[
("pool", nn.AvgPool2d(2, 2, 0, ceil_mode=True)),
("conv", ConvNormLayer(ch_in, ch_out, 1, 1)),
]
)
)
else:
self.short = ConvNormLayer(ch_in, ch_out, 1, stride)
self.branch2a = ConvNormLayer(ch_in, ch_out, 3, stride, act=act)
self.branch2b = ConvNormLayer(ch_out, ch_out, 3, 1, act=None)
self.act = nn.Identity() if act is None else get_activation(act)
def forward(self, x):
out = self.branch2a(x)
out = self.branch2b(out)
if self.shortcut:
short = x
else:
short = self.short(x)
out = out + short
out = self.act(out)
return out
class BottleNeck(nn.Module):
expansion = 4
def __init__(self, ch_in, ch_out, stride, shortcut, act="relu", variant="b"):
super().__init__()
if variant == "a":
stride1, stride2 = stride, 1
else:
stride1, stride2 = 1, stride
width = ch_out
self.branch2a = ConvNormLayer(ch_in, width, 1, stride1, act=act)
self.branch2b = ConvNormLayer(width, width, 3, stride2, act=act)
self.branch2c = ConvNormLayer(width, ch_out * self.expansion, 1, 1)
self.shortcut = shortcut
if not shortcut:
if variant == "d" and stride == 2:
self.short = nn.Sequential(
OrderedDict(
[
("pool", nn.AvgPool2d(2, 2, 0, ceil_mode=True)),
("conv", ConvNormLayer(ch_in, ch_out * self.expansion, 1, 1)),
]
)
)
else:
self.short = ConvNormLayer(ch_in, ch_out * self.expansion, 1, stride)
self.act = nn.Identity() if act is None else get_activation(act)
def forward(self, x):
out = self.branch2a(x)
out = self.branch2b(out)
out = self.branch2c(out)
if self.shortcut:
short = x
else:
short = self.short(x)
out = out + short
out = self.act(out)
return out
class Blocks(nn.Module):
def __init__(self, block, ch_in, ch_out, count, stage_num, act="relu", variant="b"):
super().__init__()
self.blocks = nn.ModuleList()
for i in range(count):
self.blocks.append(
block(
ch_in,
ch_out,
stride=2 if i == 0 and stage_num != 2 else 1,
shortcut=False if i == 0 else True,
variant=variant,
act=act,
)
)
if i == 0:
ch_in = ch_out * block.expansion
def forward(self, x):
out = x
for block in self.blocks:
out = block(out)
return out
@register()
class PResNet(nn.Module):
def __init__(
self,
depth,
variant="d",
num_stages=4,
return_idx=[0, 1, 2, 3],
act="relu",
freeze_at=-1,
freeze_norm=True,
pretrained=False,
):
super().__init__()
block_nums = ResNet_cfg[depth]
ch_in = 64
if variant in ["c", "d"]:
conv_def = [
[3, ch_in // 2, 3, 2, "conv1_1"],
[ch_in // 2, ch_in // 2, 3, 1, "conv1_2"],
[ch_in // 2, ch_in, 3, 1, "conv1_3"],
]
else:
conv_def = [[3, ch_in, 7, 2, "conv1_1"]]
self.conv1 = nn.Sequential(
OrderedDict(
[
(name, ConvNormLayer(cin, cout, k, s, act=act))
for cin, cout, k, s, name in conv_def
]
)
)
ch_out_list = [64, 128, 256, 512]
block = BottleNeck if depth >= 50 else BasicBlock
_out_channels = [block.expansion * v for v in ch_out_list]
_out_strides = [4, 8, 16, 32]
self.res_layers = nn.ModuleList()
for i in range(num_stages):
stage_num = i + 2
self.res_layers.append(
Blocks(
block, ch_in, ch_out_list[i], block_nums[i], stage_num, act=act, variant=variant
)
)
ch_in = _out_channels[i]
self.return_idx = return_idx
self.out_channels = [_out_channels[_i] for _i in return_idx]
self.out_strides = [_out_strides[_i] for _i in return_idx]
if freeze_at >= 0:
self._freeze_parameters(self.conv1)
for i in range(min(freeze_at, num_stages)):
self._freeze_parameters(self.res_layers[i])
if freeze_norm:
self._freeze_norm(self)
if pretrained:
if isinstance(pretrained, bool) or "http" in pretrained:
state = torch.hub.load_state_dict_from_url(
donwload_url[depth], map_location="cpu", model_dir="weight"
)
else:
state = torch.load(pretrained, map_location="cpu")
self.load_state_dict(state)
print(f"Load PResNet{depth} state_dict")
def _freeze_parameters(self, m: nn.Module):
for p in m.parameters():
p.requires_grad = False
def _freeze_norm(self, m: nn.Module):
if isinstance(m, nn.BatchNorm2d):
m = FrozenBatchNorm2d(m.num_features)
else:
for name, child in m.named_children():
_child = self._freeze_norm(child)
if _child is not child:
setattr(m, name, _child)
return m
def forward(self, x):
conv1 = self.conv1(x)
x = F.max_pool2d(conv1, kernel_size=3, stride=2, padding=1)
outs = []
for idx, stage in enumerate(self.res_layers):
x = stage(x)
if idx in self.return_idx:
outs.append(x)
return outs

View File

@ -0,0 +1,83 @@
from collections import OrderedDict
import torch
import torch.nn as nn
import torch.nn.functional as F
from ...core import register
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, in_planes, planes, stride=1):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(
in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.shortcut = nn.Sequential()
if stride != 1 or in_planes != self.expansion * planes:
self.shortcut = nn.Sequential(
nn.Conv2d(
in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False
),
nn.BatchNorm2d(self.expansion * planes),
)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
out += self.shortcut(x)
out = F.relu(out)
return out
class _ResNet(nn.Module):
def __init__(self, block, num_blocks, num_classes=10):
super().__init__()
self.in_planes = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
self.linear = nn.Linear(512 * block.expansion, num_classes)
def _make_layer(self, block, planes, num_blocks, stride):
strides = [stride] + [1] * (num_blocks - 1)
layers = []
for stride in strides:
layers.append(block(self.in_planes, planes, stride))
self.in_planes = planes * block.expansion
return nn.Sequential(*layers)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = F.avg_pool2d(out, 4)
out = out.view(out.size(0), -1)
out = self.linear(out)
return out
@register()
class MResNet(nn.Module):
def __init__(self, num_classes=10, num_blocks=[2, 2, 2, 2]) -> None:
super().__init__()
self.model = _ResNet(BasicBlock, num_blocks, num_classes)
def forward(self, x):
return self.model(x)

View File

@ -0,0 +1,66 @@
"""Copyright(c) 2023 lyuwenyu. All Rights Reserved.
https://towardsdatascience.com/getting-started-with-pytorch-image-models-timm-a-practitioners-guide-4e77b4bf9055#0583
"""
import torch
from torchvision.models.feature_extraction import create_feature_extractor, get_graph_node_names
from ...core import register
from .utils import IntermediateLayerGetter
@register()
class TimmModel(torch.nn.Module):
def __init__(
self, name, return_layers, pretrained=False, exportable=True, features_only=True, **kwargs
) -> None:
super().__init__()
import timm
model = timm.create_model(
name,
pretrained=pretrained,
exportable=exportable,
features_only=features_only,
**kwargs,
)
# nodes, _ = get_graph_node_names(model)
# print(nodes)
# features = {'': ''}
# model = create_feature_extractor(model, return_nodes=features)
assert set(return_layers).issubset(
model.feature_info.module_name()
), f"return_layers should be a subset of {model.feature_info.module_name()}"
# self.model = model
self.model = IntermediateLayerGetter(model, return_layers)
return_idx = [model.feature_info.module_name().index(name) for name in return_layers]
self.strides = [model.feature_info.reduction()[i] for i in return_idx]
self.channels = [model.feature_info.channels()[i] for i in return_idx]
self.return_idx = return_idx
self.return_layers = return_layers
def forward(self, x: torch.Tensor):
outputs = self.model(x)
# outputs = [outputs[i] for i in self.return_idx]
return outputs
if __name__ == "__main__":
model = TimmModel(name="resnet34", return_layers=["layer2", "layer3"])
data = torch.rand(1, 3, 640, 640)
outputs = model(data)
for output in outputs:
print(output.shape)
"""
model:
type: TimmModel
name: resnet34
return_layers: ['layer2', 'layer4']
"""

View File

@ -0,0 +1,50 @@
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import torch
import torchvision
from ...core import register
from .utils import IntermediateLayerGetter
__all__ = ["TorchVisionModel"]
@register()
class TorchVisionModel(torch.nn.Module):
def __init__(self, name, return_layers, weights=None, **kwargs) -> None:
super().__init__()
if weights is not None:
weights = getattr(torchvision.models.get_model_weights(name), weights)
model = torchvision.models.get_model(name, weights=weights, **kwargs)
# TODO hard code.
if hasattr(model, "features"):
model = IntermediateLayerGetter(model.features, return_layers)
else:
model = IntermediateLayerGetter(model, return_layers)
self.model = model
def forward(self, x):
return self.model(x)
# TorchVisionModel('swin_t', return_layers=['5', '7'])
# TorchVisionModel('resnet34', return_layers=['layer2','layer3', 'layer4'])
# TorchVisionModel:
# name: swin_t
# return_layers: ['5', '7']
# weights: DEFAULT
# model:
# type: TorchVisionModel
# name: resnet34
# return_layers: ['layer2','layer3', 'layer4']
# weights: DEFAULT

View File

@ -0,0 +1,56 @@
"""
https://github.com/pytorch/vision/blob/main/torchvision/models/_utils.py
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from collections import OrderedDict
from typing import Dict, List
import torch.nn as nn
class IntermediateLayerGetter(nn.ModuleDict):
"""
Module wrapper that returns intermediate layers from a model
It has a strong assumption that the modules have been registered
into the model in the same order as they are used.
This means that one should **not** reuse the same nn.Module
twice in the forward if you want this to work.
Additionally, it is only able to query submodules that are directly
assigned to the model. So if `model` is passed, `model.feature1` can
be returned, but not `model.feature1.layer2`.
"""
_version = 3
def __init__(self, model: nn.Module, return_layers: List[str]) -> None:
if not set(return_layers).issubset([name for name, _ in model.named_children()]):
raise ValueError(
"return_layers are not present in model. {}".format(
[name for name, _ in model.named_children()]
)
)
orig_return_layers = return_layers
return_layers = {str(k): str(k) for k in return_layers}
layers = OrderedDict()
for name, module in model.named_children():
layers[name] = module
if name in return_layers:
del return_layers[name]
if not return_layers:
break
super().__init__(layers)
self.return_layers = orig_return_layers
def forward(self, x):
outputs = []
for name, module in self.items():
x = module(x)
if name in self.return_layers:
outputs.append(x)
return outputs

Some files were not shown because too many files have changed in this diff Show More