diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/.gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..56c1a921a59416a36501c0eee98935e2e58ec20b --- /dev/null +++ b/.gitignore @@ -0,0 +1,30 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +env/ +venv/ +ENV/ +.venv + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS +.DS_Store +Thumbs.db + +# Development documentation (not needed for Spaces) +INFERENCE_REQUIRED.md +REQUIRED_FILES.md +FILES_TO_COPY.md + +# Old config files (if not needed) +# 10deimv2_dinov3_s_coco.yml + diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b06686fa2e8856da58eacd62b286a7282492b088 --- /dev/null +++ b/README.md @@ -0,0 +1,99 @@ +--- +title: DEIMv2 Floorplan Symbol Detection +emoji: 🏗️ +colorFrom: blue +colorTo: green +sdk: gradio +sdk_version: 5.0.0 +app_file: app.py +pinned: false +license: AItech +--- + +# DEIMv2 図面記号検出デモ + +DEIMv2 (Detection Enhanced by Interaction Module v2) を使用した図面記号検出デモアプリケーションです。 + +## 機能 + +- PNG形式の図面画像からの記号検出（16クラス対応） +- タイル推論による大サイズ画像の処理（640×640タイル、128pxオーバーラップ） +- 検出結果の可視化（バウンディングボックス + ラベル + スコア） +- 記号ごとの個数集計表示 +- スコア閾値の調整（デフォルト: 0.9） +- クラスフィルタリング機能（検出するクラスを選択可能） +- NMS（Non-Maximum Suppression）による重複検出の自動マージ + +## 使用方法 + +1. 左側に **PNG形式の図面** をアップロード +2. **詳細設定**を開いて以下を調整（必要に応じて）: + - **スコア閾値**: 検出の信頼度閾値（デフォルト: 0.9） + - **検出するクラス**: チェックボックスで検出したいクラスを選択（デフォルト: door1のみ） +3. 「検出を実行」ボタンをクリック +4. 中央に **検出結果付き図面**、右側に **記号名称と個数** が表示されます + +## モデル + +- **DEIMv2**: DINOv3STAsバックボーン（ViT-Tiny）を使用した物体検出モデル +- **検出対象**: 16クラスの図面記号 + - `kanki`, `kanki_shikaku`, `kanki_regisuta` + - `window1`, `window2` + - `door1`, `door2` + - `bathtub1`, `konro1`, `sink1`, `toilet1` + - `kasaikeihou1`, `kasaikeihou2` + - `houi1`, `houi2`, `houi3` + +## ハードウェア要件 + +- **本番環境**: Hugging Face Spaces ZeroGPUを使用 +- **ローカル開発**: CPU Basic（2 vCPU、16GBメモリ）以上 +- **GPU**: 推論速度向上のため、NVIDIA T4以上のGPUを推奨（オプション） +- **ディスク**: モデルファイル（.pt）はGit LFSで管理されます + +## ファイル構成 + +- `app.py`: Gradio UI + 推論パイプライン +- `detection.py`: DEIMv2推論ラッパー（タイル推論、NMS統合） +- `configs/deimv2_floorplan.yaml`: モデル設定ファイル +- `models/best_stg2.pth`: モデル重みファイル（Git LFS） +- `engine/`: DEIMv2エンジンモジュール + - `core/`: YAMLConfig関連モジュール + - `yaml_config.py`: YAMLConfigクラス + - `_config.py`: BaseConfigクラス + - `workspace.py`: オブジェクト作成ユーティリティ + - `yaml_utils.py`: YAML読み込みユーティリティ + - `backbone/`: バックボーンモジュール（DINOv3STAs等） + - `deim/`: DEIMv2検出モジュール + - `data/`: データローダー・データセット + - `optim/`: オプティマイザー・学習率スケジューラー + - `solver/`: 学習・推論エンジン +- `requirements.txt`: Python依存関係 + +## 技術詳細 + +### タイル推論 +大きな画像を640×640ピクセルのタイルに分割して推論します。タイル間は128ピクセルのオーバーラップを持ち、境界付近の記号も確実に検出できます。 + +### NMS（Non-Maximum Suppression） +タイル推論により生じる重複検出を、クラスごとにIoU閾値0.4でNMSを適用して統合します。 + +### ZeroGPU対応 +本番環境ではHugging Face SpacesのZeroGPUを使用しています。`@spaces.GPU`デコレータにより、GPUが利用可能な場合に自動的にGPUを使用し、利用できない場合はCPUにフォールバックします。デバイスの決定は推論実行時に動的に行われます（`detection.py`の`_get_device()`関数）。 + +## セットアップ注意事項 + +⚠️ **重要**: `engine/`ディレクトリ内のモジュールはDEIMv2リポジトリから必要な実装をコピーしています。 +本番環境で動作させるには、すべての依存モジュールが正しくインポート可能であることを確認してください。 + +## トラブルシューティング + +- **モデルファイルが見つからない**: Git LFSが正しく設定されているか確認してください +- **メモリ不足**: CPU UpgradeまたはGPUオプションの使用を検討してください +- **推論エラー**: 画像形式がPNG形式であることを確認してください +- **検出結果が表示されない**: スコア閾値を下げる（例: 0.5）か、検出するクラスをすべて選択してください +- **デバッグモード**: 環境変数 `DEBUG_DEIMV2=1` を設定すると、詳細なデバッグ情報が出力されます + +## ライセンス + +MIT diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..4fb477b682e2260a9288b6ba94f1e59d450fa4a7 --- /dev/null +++ b/app.py @@ -0,0 +1,260 @@ +# app.py +from collections import Counter +from functools import lru_cache +from typing import Tuple, Dict, Any, List +import os +import yaml + +import gradio as gr +import numpy as np +import spaces +from PIL import Image, ImageDraw, ImageFont + +from detection import run_inference, Detection + +# Gradio 5.0.x と JSON コンポーネントの組み合わせで +# /info 生成時に json_schema_to_python_type が bool を dict とみなして落ちる +# 既知バグがあるため、bool を安全に処理するようにパッチを当てる。 +try: + from gradio_client import utils as grc_utils + + _orig_json_schema_to_python_type = grc_utils._json_schema_to_python_type # type: ignore[attr-defined] + _orig_json_schema_to_python_type_public = grc_utils.json_schema_to_python_type + + def _json_schema_to_python_type_safe(schema, defs=None): # type: ignore[override] + if isinstance(schema, bool): + return "Any" + return _orig_json_schema_to_python_type(schema, defs) + + grc_utils._json_schema_to_python_type = _json_schema_to_python_type_safe # type: ignore[attr-defined] + + def _json_schema_to_python_type_safe_public(schema): + if isinstance(schema, bool): + return "Any" + return _orig_json_schema_to_python_type_public(schema) + + grc_utils.json_schema_to_python_type = _json_schema_to_python_type_safe_public +except Exception: + # パッチが失敗してもアプリ起動は継続する + pass + + +@lru_cache(maxsize=1) +def load_class_names() -> List[str]: + """ + 設定ファイルからクラスリストを読み込む + """ + config_path = "configs/deimv2_floorplan.yaml" + try: + with open(config_path, 'r', encoding='utf-8') as f: + config = yaml.safe_load(f) + # Modelセクションからclass_namesを取得 + if 'Model' in config and 'class_names' in config['Model']: + return config['Model']['class_names'] + else: + # フォールバック: デフォルトのクラスリスト + return ["kanki", "kanki_shikaku", "kanki_regisuta", "window1", "window2", + "door1", "door2", "bathtub1", "konro1", "sink1", "toilet1", + "kasaikeihou1", "kasaikeihou2", "houi1", "houi2", "houi3"] + except Exception as e: + # エラー時はデフォルトのクラスリストを返す + print(f"Warning: Failed to load class names from config: {e}") + return ["kanki", "kanki_shikaku", "kanki_regisuta", "window1", "window2", + "door1", "door2", "bathtub1", "konro1", "sink1", "toilet1", + "kasaikeihou1", "kasaikeihou2", "houi1", "houi2", "houi3"] + + +def pil_to_np(img: Image.Image) -> np.ndarray: + return np.array(img.convert("RGB")) + + +def draw_detections( + image_pil: Image.Image, + detections: List[Detection], +) -> Image.Image: + """検出結果を図面に重ねて描画""" + draw = ImageDraw.Draw(image_pil) + + # クラスごとの色マッピング + color_map = { + "kanki": (255, 0, 0), # 赤 + "door1": (0, 0, 255), # 青 + "door2": (255, 255, 0), # 黄 + } + default_color = (0, 255, 0) # デフォルト色（緑） + + try: + font = ImageFont.truetype("DejaVuSans.ttf", 24) + except Exception: + font = ImageFont.load_default() + + for (x1, y1, x2, y2, label, score) in detections: + # ラベルに応じた色を取得 + color = color_map.get(label, default_color) + + # bbox + draw.rectangle([(x1, y1), (x2, y2)], outline=color, width=3) + + # ラベル＋スコア + text = f"{label} {score:.2f}" + # textsizeは非推奨のため、textbboxを使用（互換性のためフォールバックあり） + try: + bbox = draw.textbbox((0, 0), text, font=font) + tw = bbox[2] - bbox[0] + th = bbox[3] - bbox[1] + except AttributeError: + # 古いPillowバージョン用のフォールバック + tw, th = draw.textsize(text, font=font) + draw.rectangle( + [(x1, y1 - th - 2), (x1 + tw + 2, y1)], + fill=color, + ) + draw.text((x1 + 1, y1 - th - 2), text, fill=(0, 0, 0), font=font) + + return image_pil + + +def summarize_detections(detections: List[Detection]) -> List[List[Any]]: + """ラベルごとの個数を集計してDataframe用のデータ形式にする""" + labels = [d[4] for d in detections] + counter = Counter(labels) + + # ヘッダー行 + data = [["記号名称", "個数"]] + + # データ行（個数の降順、ラベル名の昇順でソート） + for label, count in sorted(counter.items(), key=lambda x: (-x[1], x[0])): + data.append([label, count]) + + return data + + +def inference_pipeline( + image: Image.Image, + score_thresh: float = 0.8, + selected_classes: List[str] = None, +) -> Tuple[Image.Image, List[List[Any]]]: + """Gradio から呼ばれるメイン処理""" + if image is None: + raise gr.Error("PNG形式の図面をアップロードしてください。") + + try: + # 画像が文字列（ファイルパス）の場合はPIL Imageに変換 + if isinstance(image, str): + try: + img_pil = Image.open(image).convert("RGB") + except Exception as e: + raise gr.Error(f"画像ファイルの読み込みに失敗しました: {str(e)}") + else: + # 既にPIL Imageオブジェクトの場合 + img_pil = image.convert("RGB") + + img_np = pil_to_np(img_pil) + + # DEIMv2 推論 + detections = run_inference(img_np, score_thresh=score_thresh) + + # クラスフィルタリング: 選択されたクラスのみを残す + if selected_classes is not None and len(selected_classes) > 0: + # 選択されたクラスリストに含まれる検出結果のみをフィルタリング + filtered_detections = [ + det for det in detections + if det[4] in selected_classes # det[4]はlabel_name + ] + detections = filtered_detections + + # 描画 + vis_pil = draw_detections(img_pil.copy(), detections) + + # 集計 + summary = summarize_detections(detections) + + return vis_pil, summary + except Exception as e: + error_msg = f"エラーが発生しました: {str(e)}" + raise gr.Error(error_msg) + + +@spaces.GPU +def gpu_inference( + image: Image.Image, + score_thresh: float = 0.9, # UIのデフォルト値と統一 + selected_classes: List[str] = None, +): + """Spaces ZeroGPU が検出できるようにデコレータ付きの推論関数を用意""" + return inference_pipeline(image, score_thresh, selected_classes) + + +# ========================= +# Gradio UI +# ========================= +with gr.Blocks(title="DEIMv2 Floorplan Symbol Detection") as demo: + gr.Markdown( + """ +# 図面記号検出デモ（by AItech） + +1. 左側に **PNG図面** をアップロード +2. 「検出を実行」を押す +3. 中央に **検出結果付き図面**、右側に **記号名称＋個数** が表示されます。 +""" + ) + + # クラスリストを読み込む + class_names = load_class_names() + + with gr.Row(): + # 左: 入力 + with gr.Column(scale=1): + input_image = gr.Image( + label="入力図面 (PNG)", + type="pil", + image_mode="RGB", + ) + + # 詳細設定タブ（デフォルトは閉じた状態） + with gr.Accordion("詳細設定", open=False): + score_thresh = gr.Slider( + minimum=0.0, + maximum=1.0, + value=0.9, + step=0.05, + label="スコア閾値", + ) + selected_classes = gr.CheckboxGroup( + choices=class_names, + value=["door1"], # デフォルトでdoor1のみ選択 + label="検出するクラス", + info="選択したクラスの検出結果のみが表示されます", + ) + + run_button = gr.Button("検出を実行", variant="primary") + + # 中央: 出力画像 + with gr.Column(scale=2): + output_image = gr.Image( + label="検出結果付き図面", + type="pil", + ) + + # 右: サイドバー（記号名称＋個数） + with gr.Column(scale=1): + summary_dataframe = gr.Dataframe( + label="検出サマリ (記号名称と個数)", + headers=["記号名称", "個数"], + interactive=False, + ) + + # ボタンの動作 + run_button.click( + fn=gpu_inference, + inputs=[input_image, score_thresh, selected_classes], + outputs=[output_image, summary_dataframe], + ) + +# Gradio 5では、Spaces上ではdemoオブジェクトを直接エクスポートするだけで動作します +# ローカルテスト時のみdemo.launch()を呼び出します +if __name__ == "__main__": + demo.launch(server_name="0.0.0.0", server_port=7860) + +# Spaces上では、demoオブジェクトを直接エクスポートします +# Gradio 5は自動的にdemoオブジェクトを検出して起動します diff --git a/configs/base/dataloader.yml b/configs/base/dataloader.yml new file mode 100644 index 0000000000000000000000000000000000000000..22de3aa4645e9620c0297396da56c06c8b47e8a4 --- /dev/null +++ b/configs/base/dataloader.yml @@ -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'] # Mosaicを除外 + + 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 diff --git a/configs/base/deim.yml b/configs/base/deim.yml new file mode 100644 index 0000000000000000000000000000000000000000..3aa63588df77c2063c00632a90dfd45dbdee0ef5 --- /dev/null +++ b/configs/base/deim.yml @@ -0,0 +1,48 @@ +# Dense O2O +train_dataloader: + dataset: + transforms: + ops: + - {type: Mosaic, output_size: 320, rotation_range: 10, translation_range: [0.1, 0.1], scaling_range: [0.5, 1.5], + probability: 1.0, fill_value: 0, use_cache: True, max_cached_images: 50, random_pop: True} + - {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: + epoch: [4, 29, 50] # list + ops: ['Mosaic', 'RandomPhotometricDistort', 'RandomZoomOut', 'RandomIoUCrop'] + mosaic_prob: 0.5 + + collate_fn: + mixup_prob: 0.5 + mixup_epochs: [4, 29] + stop_epoch: 50 # epoch in [72, ~) stop `multiscales` + +# Unfreezing BN +HGNetv2: + freeze_at: -1 # 0 default + freeze_norm: False # True default + +# Activation +DFINETransformer: + activation: silu + mlp_act: silu + +## Our LR-Scheduler +lrsheduler: flatcosine +lr_gamma: 0.5 +warmup_iter: 2000 +flat_epoch: 29 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 8 + +## Our Loss +DEIMCriterion: + weight_dict: {loss_mal: 1, loss_bbox: 5, loss_giou: 2, loss_fgl: 0.15, loss_ddf: 1.5} + losses: ['mal', 'boxes', 'local'] + gamma: 1.5 \ No newline at end of file diff --git a/configs/base/deimv2.yml b/configs/base/deimv2.yml new file mode 100644 index 0000000000000000000000000000000000000000..7428898f8f734a7cb48cadcd4c7a9bd6f719d9a2 --- /dev/null +++ b/configs/base/deimv2.yml @@ -0,0 +1,144 @@ +task: detection + +model: DEIM +criterion: DEIMCriterion +postprocessor: PostProcessor + +use_focal_loss: True +eval_spatial_size: [640, 640] # h w +checkpoint_freq: 5 # save freq + +DEIM: + backbone: HGNetv2 + encoder: HybridEncoder + decoder: DEIMTransformer + +HGNetv2: + name: 'B4' + return_idx: [1, 2, 3] + freeze_at: -1 # 0 default + freeze_stem_only: True + freeze_norm: False # True default + 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' + + # New + version: deim + csp_type: csp2 + fuse_op: sum + +DEIMTransformer: + 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 + + reg_max: 32 + reg_scale: 4 + 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 + + # Act + activation: silu + mlp_act: silu + + # FFN + dim_feedforward: 2048 + +PostProcessor: + num_top_queries: 300 + + +## DEIM LR-Scheduler +epoches: 58 # 72 + 2n # Increase to search for the optimal ema + +lrsheduler: flatcosine +lr_gamma: 0.5 +warmup_iter: 2000 +flat_epoch: 29 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 8 + +## Dense O2O: Mosaic + Mixup + CopyBlend +train_dataloader: + dataset: + transforms: + ops: + - {type: Mosaic, output_size: 320, rotation_range: 10, translation_range: [0.1, 0.1], scaling_range: [0.5, 1.5], + probability: 1.0, fill_value: 0, use_cache: True, max_cached_images: 50, random_pop: True} + - {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} + # Mosaic options + policy: + epoch: [4, 29, 50] # list + ops: ['Mosaic', 'RandomPhotometricDistort', 'RandomZoomOut', 'RandomIoUCrop'] + mosaic_prob: 0.5 + + collate_fn: + # Mixup options + mixup_prob: 0.5 + mixup_epochs: [4, 29] + stop_epoch: 50 # epoch in [72, ~) stop `multiscales` + # CopyBlend options + copyblend_prob: 0.5 + copyblend_epochs: [4, 50] + area_threshold: 100 + num_objects: 3 + with_expand: True + expand_ratios: [0.1, 0.25] + + ema_restart_decay: 0.9999 + base_size_repeat: 4 + +## DEIM Loss +DEIMCriterion: + weight_dict: {loss_mal: 1, loss_bbox: 5, loss_giou: 2, loss_fgl: 0.15, loss_ddf: 1.5} + losses: ['mal', 'boxes', 'local'] + gamma: 1.5 + alpha: 0.75 + reg_max: 32 + + matcher: + type: HungarianMatcher + weight_dict: {cost_class: 2, cost_bbox: 5, cost_giou: 2} + alpha: 0.25 + gamma: 2.0 + # change matcher + change_matcher: True + iou_order_alpha: 4.0 + matcher_change_epoch: 45 \ No newline at end of file diff --git a/configs/base/dfine_hgnetv2.yml b/configs/base/dfine_hgnetv2.yml new file mode 100644 index 0000000000000000000000000000000000000000..e9de6d1b10341f1e5f3d525213aacb48a19e5aaf --- /dev/null +++ b/configs/base/dfine_hgnetv2.yml @@ -0,0 +1,90 @@ +task: detection + +model: DEIM +criterion: DEIMCriterion +postprocessor: PostProcessor + +use_focal_loss: True +eval_spatial_size: [640, 640] # h w +checkpoint_freq: 4 # save freq + +DEIM: + backbone: HGNetv2 + encoder: HybridEncoder + decoder: DFINETransformer + +# Add, default for step lr scheduler +lrsheduler: flatcosine +lr_gamma: 1 +warmup_iter: 500 +flat_epoch: 4000000 +no_aug_epoch: 0 + +HGNetv2: + pretrained: True + local_model_dir: ../RT-DETR-main/D-FINE/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 + + +PostProcessor: + num_top_queries: 300 + + +DEIMCriterion: + 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 \ No newline at end of file diff --git a/configs/base/optimizer.yml b/configs/base/optimizer.yml new file mode 100644 index 0000000000000000000000000000000000000000..db490088f0220b72309f1d7a4ab1ca6aafb45322 --- /dev/null +++ b/configs/base/optimizer.yml @@ -0,0 +1,35 @@ +use_amp: True +use_ema: True +ema: + type: ModelEMA + decay: 0.9999 + warmups: 1000 + start: 0 + +epoches: 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 diff --git a/configs/base/rt_deim.yml b/configs/base/rt_deim.yml new file mode 100644 index 0000000000000000000000000000000000000000..d195ce9ea09165289e4c72301ea14b7ac45971ce --- /dev/null +++ b/configs/base/rt_deim.yml @@ -0,0 +1,49 @@ +# Dense O2O +train_dataloader: + dataset: + transforms: + ops: + - {type: Mosaic, output_size: 320, rotation_range: 10, translation_range: [0.1, 0.1], scaling_range: [0.5, 1.5], + probability: 1.0, fill_value: 0, use_cache: False, max_cached_images: 50, random_pop: True} + - {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: + epoch: [4, 29, 50] # list + ops: ['Mosaic', 'RandomPhotometricDistort', 'RandomZoomOut', 'RandomIoUCrop'] + mosaic_prob: 0.5 + + collate_fn: + mixup_prob: 0.5 + mixup_epochs: [4, 29] + stop_epoch: 50 # epoch in [72, ~) stop `multiscales` + +# Unfreezing BN +PResNet: + freeze_at: -1 # default 0 + freeze_norm: False # default True + +# Activation +RTDETRTransformerv2: + query_pos_method: as_reg + activation: silu + mlp_act: silu + +## Our LR-Scheduler +lrsheduler: flatcosine +lr_gamma: 0.5 +warmup_iter: 2000 +flat_epoch: 29 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 8 + +## Our Loss +DEIMCriterion: + weight_dict: {loss_mal: 1, loss_bbox: 5, loss_giou: 2} + losses: ['mal', 'boxes', ] + gamma: 1.5 \ No newline at end of file diff --git a/configs/base/rt_optimizer.yml b/configs/base/rt_optimizer.yml new file mode 100644 index 0000000000000000000000000000000000000000..0dbada062e6325083f3e79991b1965d2c0fd7901 --- /dev/null +++ b/configs/base/rt_optimizer.yml @@ -0,0 +1,37 @@ +use_amp: True +use_ema: True +ema: + type: ModelEMA + decay: 0.9999 + warmups: 2000 + start: 0 + +epoches: 72 +clip_max_norm: 0.1 + +train_dataloader: + total_batch_size: 16 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm).*$' + lr: 0.00001 + - + params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0001 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +lr_scheduler: + type: MultiStepLR + milestones: [1000] + gamma: 0.1 + + +lr_warmup_scheduler: + type: LinearWarmup + warmup_duration: 2000 diff --git a/configs/base/rtdetrv2_r50vd.yml b/configs/base/rtdetrv2_r50vd.yml new file mode 100644 index 0000000000000000000000000000000000000000..e00de349bc6ef8d1a09a19634dc6e9d0c6b4ac41 --- /dev/null +++ b/configs/base/rtdetrv2_r50vd.yml @@ -0,0 +1,90 @@ +task: detection + +model: DEIM +criterion: DEIMCriterion +postprocessor: PostProcessor + +use_focal_loss: True +eval_spatial_size: [640, 640] # h w +checkpoint_freq: 4 # save freq + +DEIM: + backbone: PResNet + encoder: HybridEncoder + decoder: RTDETRTransformerv2 + + +# Add, default for step lr scheduler +lrsheduler: flatcosine +lr_gamma: 1 +warmup_iter: 2000 +flat_epoch: 4000000 +no_aug_epoch: 0 + +PResNet: + depth: 50 + variant: d + freeze_at: 0 + return_idx: [1, 2, 3] + num_stages: 4 + freeze_norm: True + pretrained: True + local_model_dir: ../RT-DETR-main/rtdetrv2_pytorch/INK1k/ + + +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' + version: rt_detrv2 # pay attention to this + + +RTDETRTransformerv2: + feat_channels: [256, 256, 256] + feat_strides: [8, 16, 32] + hidden_dim: 256 + num_levels: 3 + + num_layers: 6 + num_queries: 300 + + num_denoising: 100 + label_noise_ratio: 0.5 + box_noise_scale: 1.0 # 1.0 0.4 + + eval_idx: -1 + + # NEW, can be chosen + num_points: [4, 4, 4] # [3,3,3] [2,2,2] + cross_attn_method: default # default, discrete + query_select_method: default # default, agnostic + + +PostProcessor: + num_top_queries: 300 + +DEIMCriterion: + weight_dict: {loss_vfl: 1, loss_bbox: 5, loss_giou: 2,} + losses: ['vfl', 'boxes', ] + alpha: 0.75 + gamma: 2.0 + use_uni_set: False + + matcher: + type: HungarianMatcher + weight_dict: {cost_class: 2, cost_bbox: 5, cost_giou: 2} + alpha: 0.25 + gamma: 2.0 \ No newline at end of file diff --git a/configs/dataset/coco_detection.yml b/configs/dataset/coco_detection.yml new file mode 100644 index 0000000000000000000000000000000000000000..51851f4764e939a19b94316e3769d61692717eac --- /dev/null +++ b/configs/dataset/coco_detection.yml @@ -0,0 +1,40 @@ +task: detection + +evaluator: + type: CocoEvaluator + iou_types: ['bbox'] + +num_classes: 16 +remap_mscoco_category: False # カテゴリIDが0-15に変更済みのためFalse + +train_dataloader: + type: DataLoader + dataset: + type: CocoDetection + img_folder: /content/DEIMv2/dataset/train/images + ann_file: /content/DEIMv2/dataset/train/annotations/train_annotations.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: /content/DEIMv2/dataset/validation/images + ann_file: /content/DEIMv2/dataset/validation/annotations/validation_annotations.json + return_masks: False + transforms: + type: Compose + ops: ~ + shuffle: False + num_workers: 2 + drop_last: False + collate_fn: + type: BatchImageCollateFunction diff --git a/configs/dataset/crowdhuman_detection.yml b/configs/dataset/crowdhuman_detection.yml new file mode 100644 index 0000000000000000000000000000000000000000..f4dc707db67f651533671034dea92144083f74f9 --- /dev/null +++ b/configs/dataset/crowdhuman_detection.yml @@ -0,0 +1,41 @@ +task: detection + +evaluator: + type: CocoEvaluator + iou_types: ['bbox', ] + +num_classes: 2 # your dataset classes +remap_mscoco_category: False + +train_dataloader: + type: DataLoader + dataset: + type: CocoDetection + img_folder: /datassd/coco/crowd_human_coco/CrowdHuman_train + ann_file: /datassd/coco/crowd_human_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: /datassd/coco/crowd_human_coco/CrowdHuman_val + ann_file: /datassd/coco/crowd_human_coco/Chuman-val.json + return_masks: False + transforms: + type: Compose + ops: ~ + shuffle: False + num_workers: 4 + drop_last: False + collate_fn: + type: BatchImageCollateFunction diff --git a/configs/dataset/custom_detection.yml b/configs/dataset/custom_detection.yml new file mode 100644 index 0000000000000000000000000000000000000000..35435ad68e29d99d8f9f69100cd56a2c403fe710 --- /dev/null +++ b/configs/dataset/custom_detection.yml @@ -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 diff --git a/configs/dataset/obj365_detection.yml b/configs/dataset/obj365_detection.yml new file mode 100644 index 0000000000000000000000000000000000000000..e843e85bf2d53de3e61fdd109cf51ab9fc9957e3 --- /dev/null +++ b/configs/dataset/obj365_detection.yml @@ -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: /home/Dataset/objects365/train + ann_file: /home/Dataset/objects365/train/new_zhiyuan_objv2_train_resized640.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: /home/Dataset/objects365/val + ann_file: /home/Dataset/objects365/val/new_zhiyuan_objv2_val_resized640.json + return_masks: False + transforms: + type: Compose + ops: ~ + shuffle: False + num_workers: 4 + drop_last: False + collate_fn: + type: BatchImageCollateFunction diff --git a/configs/dataset/voc_detection.yml b/configs/dataset/voc_detection.yml new file mode 100644 index 0000000000000000000000000000000000000000..1f9ceeb8881653d496ac5fd02c465aea5306d72f --- /dev/null +++ b/configs/dataset/voc_detection.yml @@ -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 diff --git a/configs/deim_dfine/deim_hgnetv2_l_coco.yml b/configs/deim_dfine/deim_hgnetv2_l_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..6b35a78e453d52d15291fd91dd04c9a55cfec8af --- /dev/null +++ b/configs/deim_dfine/deim_hgnetv2_l_coco.yml @@ -0,0 +1,37 @@ +__include__: [ + './dfine_hgnetv2_l_coco.yml', + '../base/deim.yml' +] + +output_dir: ./outputs/deim_hgnetv2_l_coco + +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.000125 + +# Increase to search for the optimal ema +epoches: 58 # 72 + 2n + +## Our LR-Scheduler +flat_epoch: 29 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 8 + +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 29, 50] # list + + collate_fn: + mixup_epochs: [4, 29] + stop_epoch: 50 \ No newline at end of file diff --git a/configs/deim_dfine/deim_hgnetv2_m_coco.yml b/configs/deim_dfine/deim_hgnetv2_m_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..9fa5167620c57f6fdb14892dd1cf9a00839fde92 --- /dev/null +++ b/configs/deim_dfine/deim_hgnetv2_m_coco.yml @@ -0,0 +1,39 @@ +__include__: [ + './dfine_hgnetv2_m_coco.yml', + '../base/deim.yml' +] + +output_dir: ./outputs/deim_hgnetv2_m_coco + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*bn).*$' + lr: 0.00004 + - + params: '^(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0004 + betas: [0.9, 0.999] + weight_decay: 0.0001 + + +# Increase to search for the optimal ema +epoches: 102 # 120 + 4n + +## Our LR-Scheduler +flat_epoch: 49 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 12 + +## Our DataAug +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 49, 90] # list + + collate_fn: + mixup_epochs: [4, 49] + stop_epoch: 90 \ No newline at end of file diff --git a/configs/deim_dfine/deim_hgnetv2_n_coco.yml b/configs/deim_dfine/deim_hgnetv2_n_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..62db245d96e4cb402e774efe9cb26fc1f401e40e --- /dev/null +++ b/configs/deim_dfine/deim_hgnetv2_n_coco.yml @@ -0,0 +1,44 @@ +__include__: [ + './dfine_hgnetv2_n_coco.yml', + '../base/deim.yml' +] + +output_dir: ./deim_outputs/deim_hgnetv2_n_coco + +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 +epoches: 160 # 148 + 12 + +## Our LR-Scheduler +flat_epoch: 7800 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 12 +lr_gamma: 1.0 + +## Our DataAug +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 78, 148] # list + + collate_fn: + mixup_epochs: [4, 78] + stop_epoch: 148 + base_size_repeat: ~ \ No newline at end of file diff --git a/configs/deim_dfine/deim_hgnetv2_s_coco.yml b/configs/deim_dfine/deim_hgnetv2_s_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..68ea99aae9926cefb33b3380d37a9cbd70ed28eb --- /dev/null +++ b/configs/deim_dfine/deim_hgnetv2_s_coco.yml @@ -0,0 +1,39 @@ +__include__: [ + './dfine_hgnetv2_s_coco.yml', + '../base/deim.yml' +] + +output_dir: ./outputs/deim_hgnetv2_s_coco + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*bn).*$' + lr: 0.0002 + - + params: '^(?=.*(?:norm|bn)).*$' # except bias + weight_decay: 0. + + lr: 0.0004 + betas: [0.9, 0.999] + weight_decay: 0.0001 + + +# Increase to search for the optimal ema +epoches: 132 # 120 + 4n + +## Our LR-Scheduler +flat_epoch: 64 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 12 + +## Our DataAug +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 64, 120] # list + + collate_fn: + mixup_epochs: [4, 64] + stop_epoch: 120 \ No newline at end of file diff --git a/configs/deim_dfine/deim_hgnetv2_x_coco.yml b/configs/deim_dfine/deim_hgnetv2_x_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..8ec7f1b611089c6b8aa3a974c980f862cf821679 --- /dev/null +++ b/configs/deim_dfine/deim_hgnetv2_x_coco.yml @@ -0,0 +1,37 @@ +__include__: [ + './dfine_hgnetv2_x_coco.yml', + '../base/deim.yml' +] + +output_dir: ./outputs/deim_hgnetv2_x_coco + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm|bn).*$' + lr: 0.000005 + - + params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0005 + betas: [0.9, 0.999] + weight_decay: 0.000125 + +# Increase to search for the optimal ema +epoches: 58 # 72 + 2n + +## Our LR-Scheduler +flat_epoch: 29 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 8 + +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 29, 50] # list + + collate_fn: + mixup_epochs: [4, 29] + stop_epoch: 50 \ No newline at end of file diff --git a/configs/deim_dfine/dfine_hgnetv2_l_coco.yml b/configs/deim_dfine/dfine_hgnetv2_l_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..16c6002434659b5918db2f08955f054ca9c83d81 --- /dev/null +++ b/configs/deim_dfine/dfine_hgnetv2_l_coco.yml @@ -0,0 +1,44 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/dfine_hgnetv2.yml', +] + +output_dir: ./outputs/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 +epoches: 80 # 72 + 2n +train_dataloader: + dataset: + transforms: + policy: + epoch: 72 + collate_fn: + stop_epoch: 72 + ema_restart_decay: 0.9999 + base_size_repeat: 4 diff --git a/configs/deim_dfine/dfine_hgnetv2_m_coco.yml b/configs/deim_dfine/dfine_hgnetv2_m_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..6b5f917bd5723b931bb09e0389eb49e2e15af8c8 --- /dev/null +++ b/configs/deim_dfine/dfine_hgnetv2_m_coco.yml @@ -0,0 +1,60 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/dfine_hgnetv2.yml', +] + +output_dir: ./output/dfine_hgnetv2_m_coco + + +DEIM: + 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 +epoches: 132 # 120 + 4n +train_dataloader: + dataset: + transforms: + policy: + epoch: 120 + collate_fn: + stop_epoch: 120 + ema_restart_decay: 0.9999 + base_size_repeat: 6 diff --git a/configs/deim_dfine/dfine_hgnetv2_n_coco.yml b/configs/deim_dfine/dfine_hgnetv2_n_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..c45e65357cbfe547851e8b7385be67ca7226f6cd --- /dev/null +++ b/configs/deim_dfine/dfine_hgnetv2_n_coco.yml @@ -0,0 +1,82 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/dfine_hgnetv2.yml', +] + +output_dir: ./output/dfine_hgnetv2_n_coco + + +DEIM: + 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 +epoches: 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 diff --git a/configs/deim_dfine/dfine_hgnetv2_s_coco.yml b/configs/deim_dfine/dfine_hgnetv2_s_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..33857bc47fe0311586a28b16f7c2d58af909ffbc --- /dev/null +++ b/configs/deim_dfine/dfine_hgnetv2_s_coco.yml @@ -0,0 +1,61 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/dfine_hgnetv2.yml', +] + +output_dir: ./output/dfine_hgnetv2_s_coco + + +DEIM: + 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 +epoches: 132 # 120 + 4n +train_dataloader: + dataset: + transforms: + policy: + epoch: 120 + collate_fn: + stop_epoch: 120 + ema_restart_decay: 0.9999 + base_size_repeat: 20 diff --git a/configs/deim_dfine/dfine_hgnetv2_x_coco.yml b/configs/deim_dfine/dfine_hgnetv2_x_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..46ec15753906b37f9d0f61bf8e0637c77c295251 --- /dev/null +++ b/configs/deim_dfine/dfine_hgnetv2_x_coco.yml @@ -0,0 +1,56 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/dfine_hgnetv2.yml', +] + +output_dir: ./output/dfine_hgnetv2_x_coco + + +DEIM: + 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 +epoches: 80 # 72 + 2n +train_dataloader: + dataset: + transforms: + policy: + epoch: 72 + collate_fn: + stop_epoch: 72 + ema_restart_decay: 0.9998 + base_size_repeat: 3 diff --git a/configs/deim_dfine/object365/deim_hgnetv2_x_obj2coco_24e.yml b/configs/deim_dfine/object365/deim_hgnetv2_x_obj2coco_24e.yml new file mode 100644 index 0000000000000000000000000000000000000000..28fcd4c12f2ab91bffb111e33b18b612a9bda12d --- /dev/null +++ b/configs/deim_dfine/object365/deim_hgnetv2_x_obj2coco_24e.yml @@ -0,0 +1,50 @@ +__include__: [ + './dfine_hgnetv2_x_obj2coco.yml', + '../../base/deim.yml' +] + +output_dir: ./deim_outputs/deim_hgnetv2_x_obj2coco_24e + +HGNetv2: + freeze_at: 0 # 0 default + freeze_norm: True # True default + +# Activation +DFINETransformer: + activation: relu + mlp_act: relu + +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 +epoches: 24 # 72 + 2n + +## Our LR-Scheduler +lrsheduler: flatcosine +lr_gamma: 1 +warmup_iter: 0 # 0 +flat_epoch: 12000 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 4 + +## Our DataAug +train_dataloader: + dataset: + transforms: + policy: + epoch: [2, 12, 20] # list + + collate_fn: + mixup_epochs: [2, 12] + stop_epoch: 20 \ No newline at end of file diff --git a/configs/deim_dfine/object365/dfine_hgnetv2_x_obj2coco.yml b/configs/deim_dfine/object365/dfine_hgnetv2_x_obj2coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..c9c711b678e953d7589ad106c2aef064bd906ae6 --- /dev/null +++ b/configs/deim_dfine/object365/dfine_hgnetv2_x_obj2coco.yml @@ -0,0 +1,57 @@ +__include__: [ + '../../dataset/coco_detection.yml', + '../../runtime.yml', + '../../base/dataloader.yml', + '../../base/optimizer.yml', + '../../base/dfine_hgnetv2.yml', +] + +output_dir: ./outputs/dfine_hgnetv2_x_obj2coco + +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 + + +epoches: 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 diff --git a/configs/deim_rtdetrv2/deim_r101vd_60e_coco.yml b/configs/deim_rtdetrv2/deim_r101vd_60e_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..538a9afb8b9e1187d3ebf2cccd963d6eec35fd70 --- /dev/null +++ b/configs/deim_rtdetrv2/deim_r101vd_60e_coco.yml @@ -0,0 +1,36 @@ +__include__: [ + './rtdetrv2_r101vd_6x_coco.yml', + '../base/rt_deim.yml', +] + +output_dir: ./outputs/deim_rtdetrv2_r101vd_60e_coco + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm).*$' + lr: 0.000002 + - + params: '^(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0002 + betas: [0.9, 0.999] + weight_decay: 0.0001 + + +# change part +epoches: 60 +flat_epoch: 34 # 4 + 60 / 2 +no_aug_epoch: 2 + +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 34, 58] # list + + collate_fn: + mixup_epochs: [4, 34] + stop_epoch: 58 diff --git a/configs/deim_rtdetrv2/deim_r18vd_120e_coco.yml b/configs/deim_rtdetrv2/deim_r18vd_120e_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..2b2069fa727453a3f34f9a67c5b9f59ffce7773d --- /dev/null +++ b/configs/deim_rtdetrv2/deim_r18vd_120e_coco.yml @@ -0,0 +1,32 @@ +__include__: [ + './rtdetrv2_r18vd_120e_coco.yml', + '../base/rt_deim.yml', +] + +output_dir: ./output/deim_rtdetrv2_r18vd_120e_coco + +optimizer: + type: AdamW + params: + - + params: '^(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0002 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +# change part +epoches: 120 +flat_epoch: 64 # 4 + 120 / 2 +no_aug_epoch: 3 + +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 64, 117] # list + + collate_fn: + mixup_epochs: [4, 64] + stop_epoch: 117 \ No newline at end of file diff --git a/configs/deim_rtdetrv2/deim_r34vd_120e_coco.yml b/configs/deim_rtdetrv2/deim_r34vd_120e_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..fb9d23f72a49df713c5d0c4e91ed5e1b95d57e73 --- /dev/null +++ b/configs/deim_rtdetrv2/deim_r34vd_120e_coco.yml @@ -0,0 +1,36 @@ +__include__: [ + './rtdetrv2_r34vd_120e_coco.yml', + '../base/rt_deim.yml', +] + +output_dir: ./outputs/deim_rtdetrv2_r34vd_120e_coco + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm).*$' + lr: 0.0001 + - + params: '^(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0002 + betas: [0.9, 0.999] + weight_decay: 0.0001 + + +# change part +epoches: 120 +flat_epoch: 64 +no_aug_epoch: 3 + +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 64, 117] # list + + collate_fn: + mixup_epochs: [4, 64] + stop_epoch: 117 \ No newline at end of file diff --git a/configs/deim_rtdetrv2/deim_r50vd_60e_coco.yml b/configs/deim_rtdetrv2/deim_r50vd_60e_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..7427c57248f0f2740512fa65b6e5642d2da99709 --- /dev/null +++ b/configs/deim_rtdetrv2/deim_r50vd_60e_coco.yml @@ -0,0 +1,35 @@ +__include__: [ + './rtdetrv2_r50vd_6x_coco.yml', + '../base/rt_deim.yml', +] + +output_dir: ./outputs/deim_rtdetrv2_r50vd_60e_coco + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm).*$' + lr: 0.00002 + - + params: '^(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0002 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +# change part +epoches: 60 +flat_epoch: 34 # 4 + 60 / 2 +no_aug_epoch: 2 + +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 34, 58] # list + + collate_fn: + mixup_epochs: [4, 34] + stop_epoch: 58 \ No newline at end of file diff --git a/configs/deim_rtdetrv2/deim_r50vd_m_60e_coco.yml b/configs/deim_rtdetrv2/deim_r50vd_m_60e_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..b08bdc3713cb32fc23fc95ea9fa37f94fc7feb43 --- /dev/null +++ b/configs/deim_rtdetrv2/deim_r50vd_m_60e_coco.yml @@ -0,0 +1,39 @@ +__include__: [ + './rtdetrv2_r50vd_m_7x_coco.yml', + '../base/rt_deim.yml', +] + +output_dir: ./outputs/deim_rtdetrv2_r50vd_m_60e_coco + +RTDETRTransformerv2: + eval_idx: 2 # use 3th decoder layer to eval + num_layers: 3 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm).*$' + lr: 0.00002 + - + params: '^(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0002 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +# change part +epoches: 60 +flat_epoch: 34 # 4 + 60 / 2 +no_aug_epoch: 2 + +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 34, 58] # list + + collate_fn: + mixup_epochs: [4, 34] + stop_epoch: 58 \ No newline at end of file diff --git a/configs/deim_rtdetrv2/rtdetrv2_r101vd_6x_coco.yml b/configs/deim_rtdetrv2/rtdetrv2_r101vd_6x_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..256a089b2886fc1830923e957732b7a07bc4273b --- /dev/null +++ b/configs/deim_rtdetrv2/rtdetrv2_r101vd_6x_coco.yml @@ -0,0 +1,40 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/rt_optimizer.yml', + '../base/rtdetrv2_r50vd.yml', +] + + +output_dir: ./outputs/rtdetrv2_r101vd_6x_coco + + +PResNet: + depth: 101 + + +HybridEncoder: + # intra + hidden_dim: 384 + dim_feedforward: 2048 + + +RTDETRTransformerv2: + feat_channels: [384, 384, 384] + + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm|bn).*$' + lr: 0.000001 + - + params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$' # only encoder + decoder norm + weight_decay: 0. + + lr: 0.0001 + betas: [0.9, 0.999] + weight_decay: 0.0001 + diff --git a/configs/deim_rtdetrv2/rtdetrv2_r18vd_120e_coco.yml b/configs/deim_rtdetrv2/rtdetrv2_r18vd_120e_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..04d4d533da9a0bcd6f64e9eb2ddbcdd76ddf8edf --- /dev/null +++ b/configs/deim_rtdetrv2/rtdetrv2_r18vd_120e_coco.yml @@ -0,0 +1,44 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/rt_optimizer.yml', + '../base/rtdetrv2_r50vd.yml', +] + + +output_dir: ./output/rtdetrv2_r18vd_120e_coco + + +PResNet: + depth: 18 + freeze_at: -1 + freeze_norm: False + pretrained: True + +HybridEncoder: + in_channels: [128, 256, 512] + hidden_dim: 256 + expansion: 0.5 + +RTDETRTransformerv2: + num_layers: 3 + + +epoches: 120 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + +train_dataloader: + dataset: + transforms: + policy: + epoch: 117 + collate_fn: + scales: ~ \ No newline at end of file diff --git a/configs/deim_rtdetrv2/rtdetrv2_r34vd_120e_coco.yml b/configs/deim_rtdetrv2/rtdetrv2_r34vd_120e_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..9cfb522a1870818e2411d863b40f5f35b2289b12 --- /dev/null +++ b/configs/deim_rtdetrv2/rtdetrv2_r34vd_120e_coco.yml @@ -0,0 +1,57 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/rt_optimizer.yml', + '../base/rtdetrv2_r50vd.yml', +] + + +output_dir: ./outputs/rtdetrv2_r34vd_120e_coco + + +PResNet: + depth: 34 + freeze_at: -1 + freeze_norm: False + pretrained: True + + +HybridEncoder: + in_channels: [128, 256, 512] + hidden_dim: 256 + expansion: 0.5 + + +RTDETRTransformerv2: + num_layers: 4 + + +epoches: 120 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm|bn).*$' + lr: 0.00005 + - + params: '^(?=.*backbone)(?=.*norm|bn).*$' + lr: 0.00005 + weight_decay: 0. + - + params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$' + weight_decay: 0. + + lr: 0.0001 + betas: [0.9, 0.999] + weight_decay: 0.0001 + + +train_dataloader: + dataset: + transforms: + policy: + epoch: 117 + collate_fn: + stop_epoch: 117 diff --git a/configs/deim_rtdetrv2/rtdetrv2_r50vd_6x_coco.yml b/configs/deim_rtdetrv2/rtdetrv2_r50vd_6x_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..3ffe8505be0e15b5f22f088d9215da86fdabc970 --- /dev/null +++ b/configs/deim_rtdetrv2/rtdetrv2_r50vd_6x_coco.yml @@ -0,0 +1,25 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/rt_optimizer.yml', + '../base/rtdetrv2_r50vd.yml', +] + + +output_dir: ./outputs/rtdetrv2_r50vd_6x_coco + + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm).*$' + lr: 0.00001 + - + params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0001 + betas: [0.9, 0.999] + weight_decay: 0.0001 \ No newline at end of file diff --git a/configs/deim_rtdetrv2/rtdetrv2_r50vd_m_7x_coco.yml b/configs/deim_rtdetrv2/rtdetrv2_r50vd_m_7x_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..760a3866760fd487a4788a9d9efa249e5b65d6a4 --- /dev/null +++ b/configs/deim_rtdetrv2/rtdetrv2_r50vd_m_7x_coco.yml @@ -0,0 +1,43 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/rt_optimizer.yml', + '../base/rtdetrv2_r50vd.yml', +] + +output_dir: ./outputs/rtdetrv2_r50vd_m_6x_coco + + +HybridEncoder: + expansion: 0.5 + + +RTDETRTransformerv2: + eval_idx: 2 # use 3th decoder layer to eval + + +epoches: 84 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm).*$' + lr: 0.00001 + - + params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0001 + betas: [0.9, 0.999] + weight_decay: 0.0001 + + +train_dataloader: + dataset: + transforms: + policy: + epoch: 81 + collate_fn: + stop_epoch: 81 \ No newline at end of file diff --git a/configs/deimv2/deimv2_dinov3_l_coco.yml b/configs/deimv2/deimv2_dinov3_l_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..7bb97cb7bbcad41c59ae8d9f01994862296d79b5 --- /dev/null +++ b/configs/deimv2/deimv2_dinov3_l_coco.yml @@ -0,0 +1,104 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml', +] + + +output_dir: ./outputs/deimv2_dinov3_l_coco + +DEIM: + backbone: DINOv3STAs + +DINOv3STAs: + name: dinov3_vits16 + weights_path: ./ckpts/dinov3_vits16_pretrain_lvd1689m-08c60483.pth + interaction_indexes: [5,8,11] # only need the [1/8, 1/16, 1/32] + finetune: True + conv_inplane: 32 + hidden_dim: 224 + +HybridEncoder: + in_channels: [224, 224, 224] + hidden_dim: 224 + dim_feedforward: 896 + +DEIMTransformer: + feat_channels: [224, 224, 224] + hidden_dim: 224 + num_layers: 4 + eval_idx: -1 + dim_feedforward: 1792 + +## DEIM LR-Scheduler +epoches: 68 # 72 + 2n # Increase to search for the optimal ema + +lrsheduler: flatcosine +lr_gamma: 0.5 +warmup_iter: 2000 +flat_epoch: 34 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 8 + +## Optimizer +optimizer: + type: AdamW + params: + - + # except norm/bn/bias in self.dinov3 + params: '^(?=.*.dinov3)(?!.*(?:norm|bn|bias)).*$' + lr: 0.0000125 + - + # including norm/bn/bias in self.dinov3 + params: '^(?=.*.dinov3)(?=.*(?:norm|bn|bias)).*$' + lr: 0.0000125 + weight_decay: 0. + - + # including norm/bn/bias except for the self.dinov3 + params: '^(?=.*(?:sta|encoder|decoder))(?=.*(?:norm|bn|bias)).*$' + weight_decay: 0. + + lr: 0.0005 + betas: [0.9, 0.999] + weight_decay: 0.000125 + + +## Dense O2O: Mosaic + Mixup + CopyBlend +train_dataloader: + dataset: + transforms: + ops: + - {type: Mosaic, output_size: 320, rotation_range: 10, translation_range: [0.1, 0.1], scaling_range: [0.5, 1.5], + probability: 1.0, fill_value: 0, use_cache: True, max_cached_images: 50, random_pop: True} + - {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: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} + - {type: ConvertBoxes, fmt: 'cxcywh', normalize: True} + policy: + epoch: [4, 34, 60] # list + + collate_fn: + mixup_epochs: [4, 34] + stop_epoch: 60 + copyblend_epochs: [4, 60] + base_size_repeat: 3 + +val_dataloader: + dataset: + transforms: + ops: + - {type: Resize, size: [640, 640], } + - {type: ConvertPILImage, dtype: 'float32', scale: True} + - {type: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} + +## DEIM Loss +DEIMCriterion: + matcher: + matcher_change_epoch: 50 \ No newline at end of file diff --git a/configs/deimv2/deimv2_dinov3_m_coco.yml b/configs/deimv2/deimv2_dinov3_m_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..c4b17334138bad9e8bd431e858b50abc6c2514b3 --- /dev/null +++ b/configs/deimv2/deimv2_dinov3_m_coco.yml @@ -0,0 +1,107 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml', +] + +output_dir: ./outputs/deimv2_dinov3_m_coco + +DEIM: + backbone: DINOv3STAs + +DINOv3STAs: + name: vit_tinyplus + embed_dim: 256 + weights_path: ./ckpts/vittplus_distill.pt + interaction_indexes: [3, 7, 11] # only need the [1/8, 1/16, 1/32] + num_heads: 4 + +HybridEncoder: + in_channels: [256, 256, 256] + depth_mult: 1 + expansion: 0.67 + hidden_dim: 256 + dim_feedforward: 512 + + +DEIMTransformer: + feat_channels: [256, 256, 256] + hidden_dim: 256 + dim_feedforward: 512 + num_layers: 4 # 4 5 6 + eval_idx: -1 # -2 -3 -4 + +optimizer: + type: AdamW + + params: + - + # except norm/bn/bias in self.dinov3 + params: '^(?=.*.dinov3)(?!.*(?:norm|bn|bias)).*$' + lr: 0.000025 + - + # including norm/bn/bias in self.dinov3 + params: '^(?=.*.dinov3)(?=.*(?:norm|bn|bias)).*$' + lr: 0.000025 + weight_decay: 0. + - + # including norm/bn/bias except for the self.dinov3 + params: '^(?=.*(?:sta|encoder|decoder))(?=.*(?:norm|bn|bias)).*$' + weight_decay: 0. + + lr: 0.0005 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +epoches: 102 # 120 + 4n + +## Our LR-Scheduler +flat_epoch: 49 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 12 + + +## Our DataAug +train_dataloader: + dataset: + transforms: + ops: + - {type: Mosaic, output_size: 320, rotation_range: 10, translation_range: [0.1, 0.1], scaling_range: [0.5, 1.5], + probability: 1.0, fill_value: 0, use_cache: True, max_cached_images: 50, random_pop: True} + - {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: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} + - {type: ConvertBoxes, fmt: 'cxcywh', normalize: True} + policy: + epoch: [4, 49, 90] # list + + collate_fn: + mixup_prob: 0.5 + ema_restart_decay: 0.9999 + base_size_repeat: 6 + mixup_epochs: [4, 49] + stop_epoch: 90 + copyblend_epochs: [4, 90] + + +val_dataloader: + dataset: + transforms: + ops: + - {type: Resize, size: [640, 640], } + - {type: ConvertPILImage, dtype: 'float32', scale: True} + - {type: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} + +DEIMCriterion: + matcher: + # new matcher + change_matcher: True + iou_order_alpha: 4.0 + matcher_change_epoch: 80 diff --git a/configs/deimv2/deimv2_dinov3_s_coco.yml b/configs/deimv2/deimv2_dinov3_s_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..c8ec7ea3b95902fcd88592cbd5ce478a09b788c0 --- /dev/null +++ b/configs/deimv2/deimv2_dinov3_s_coco.yml @@ -0,0 +1,108 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml', +] + +output_dir: ./outputs/deimv2_dinov3_s_coco + +DEIM: + backbone: DINOv3STAs + +DINOv3STAs: + name: vit_tiny + embed_dim: 192 + weights_path: ./ckpts/vitt_distill.pt + interaction_indexes: [3, 7, 11] # only need the [1/8, 1/16, 1/32] + num_heads: 3 + +HybridEncoder: + in_channels: [192, 192, 192] + depth_mult: 0.67 + expansion: 0.34 + hidden_dim: 192 + dim_feedforward: 512 + +DEIMTransformer: + feat_channels: [192, 192, 192] + hidden_dim: 192 + dim_feedforward: 512 + num_layers: 4 # 4 5 6 + eval_idx: -1 # -2 -3 -4 + + +## Optimizer +optimizer: + type: AdamW + + params: + - + # except norm/bn/bias in self.dinov3 + params: '^(?=.*.dinov3)(?!.*(?:norm|bn|bias)).*$' + lr: 0.000025 + - + # including all norm/bn/bias in self.dinov3 + params: '^(?=.*.dinov3)(?=.*(?:norm|bn|bias)).*$' + lr: 0.000025 + weight_decay: 0. + - + # including all norm/bn/bias except for the self.dinov3 + params: '^(?=.*(?:sta|encoder|decoder))(?=.*(?:norm|bn|bias)).*$' + weight_decay: 0. + + lr: 0.0005 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +# Increase to search for the optimal ema +epoches: 132 # 120 + 4n + +## Our LR-Scheduler +flat_epoch: 64 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 12 + +## Our DataAug +train_dataloader: + dataset: + transforms: + ops: + - {type: Mosaic, output_size: 320, rotation_range: 10, translation_range: [0.1, 0.1], scaling_range: [0.5, 1.5], + probability: 1.0, fill_value: 0, use_cache: True, max_cached_images: 50, random_pop: True} + - {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: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} + - {type: ConvertBoxes, fmt: 'cxcywh', normalize: True} + policy: + epoch: [4, 64, 120] # list + + collate_fn: + base_size: 640 + mixup_prob: 0.5 + ema_restart_decay: 0.9999 + base_size_repeat: 20 + mixup_epochs: [4, 64] + stop_epoch: 120 + copyblend_epochs: [4, 120] + +val_dataloader: + dataset: + transforms: + ops: + - {type: Resize, size: [640, 640], } + - {type: ConvertPILImage, dtype: 'float32', scale: True} + - {type: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} + +DEIMCriterion: + matcher: + # change matcher + change_matcher: True + iou_order_alpha: 4.0 + matcher_change_epoch: 100 diff --git a/configs/deimv2/deimv2_dinov3_x_coco.yml b/configs/deimv2/deimv2_dinov3_x_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..f85120b4905dc25318f564d749fd2ba789811fed --- /dev/null +++ b/configs/deimv2/deimv2_dinov3_x_coco.yml @@ -0,0 +1,94 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml', +] + + +output_dir: ./outputs/deimv2_dinov3_x_coco + +DEIM: + backbone: DINOv3STAs + +DINOv3STAs: + name: dinov3_vits16plus + weights_path: ./ckpts/dinov3_vits16plus_pretrain_lvd1689m-4057cbaa.pth + interaction_indexes: [5,8,11] # only need the [1/8, 1/16, 1/32] + finetune: True + conv_inplane: 64 + hidden_dim: 256 + +HybridEncoder: + in_channels: [256, 256, 256] + # intra + hidden_dim: 256 + dim_feedforward: 1024 + + # cross + expansion: 1.25 + depth_mult: 1.37 + +DEIMTransformer: + num_layers: 6 + eval_idx: -1 + feat_channels: [256, 256, 256] + # reg_scale: 8 + hidden_dim: 256 + dim_feedforward: 2048 + +optimizer: + type: AdamW + params: + - + # except norm/bn/bias in self.dinov3 + params: '^(?=.*.dinov3)(?!.*(?:norm|bn|bias)).*$' + lr: 0.00001 + - + # including norm/bn/bias in self.dinov3 + params: '^(?=.*.dinov3)(?=.*(?:norm|bn|bias)).*$' + lr: 0.00001 + weight_decay: 0. + - + # including norm/bn/bias except for the self.dinov3 + params: '^(?=.*(?:sta|encoder|decoder))(?=.*(?:norm|bn|bias)).*$' + weight_decay: 0. + + lr: 0.0005 + betas: [0.9, 0.999] + weight_decay: 0.000125 + +## Dense O2O: Mosaic + Mixup + CopyBlend +train_dataloader: + dataset: + transforms: + ops: + - {type: Mosaic, output_size: 320, rotation_range: 10, translation_range: [0.1, 0.1], scaling_range: [0.5, 1.5], + probability: 1.0, fill_value: 0, use_cache: True, max_cached_images: 50, random_pop: True} + - {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: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} + - {type: ConvertBoxes, fmt: 'cxcywh', normalize: True} + policy: + epoch: [4, 29, 50] # list + + collate_fn: + mixup_epochs: [4, 29] + stop_epoch: 50 + copyblend_epochs: [4, 50] + base_size_repeat: 3 + +val_dataloader: + dataset: + transforms: + ops: + - {type: Resize, size: [640, 640], } + - {type: ConvertPILImage, dtype: 'float32', scale: True} + - {type: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} \ No newline at end of file diff --git a/configs/deimv2/deimv2_hgnetv2_atto_coco.yml b/configs/deimv2/deimv2_hgnetv2_atto_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..4d770494fd5b99bb50199ae809920c57a97d98a7 --- /dev/null +++ b/configs/deimv2/deimv2_hgnetv2_atto_coco.yml @@ -0,0 +1,123 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml', +] + +output_dir: ./outputs/deimv2_hgnetv2_atto_coco + +DEIM: + encoder: LiteEncoder + +HGNetv2: + name: 'Atto' + return_idx: [2] + freeze_at: -1 + freeze_norm: False + use_lab: True + +LiteEncoder: + in_channels: [256] + feat_strides: [16] + # intra + hidden_dim: 64 + + # cross + expansion: 0.34 + depth_mult: 0.5 + act: 'silu' + + +DEIMTransformer: + feat_channels: [64, 64] + feat_strides: [16, 32] + hidden_dim: 64 + num_levels: 2 + num_points: [4, 2] + + num_layers: 3 + eval_idx: -1 + num_queries: 100 + + # FFN + dim_feedforward: 160 + + # New options for DEIMv2 + share_bbox_head: True + use_gateway: False + +# Increase to search for the optimal ema +epoches: 500 # 468 + 32 + +## Our LR-Scheduler +warmup_iter: 4000 +flat_epoch: 250 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 32 +lr_gamma: 0.5 + +optimizer: + type: AdamW + params: + - params: '^(?=.*backbone)(?!.*norm|bn).*$' + lr: 0.001 + - params: '^(?=.*backbone)(?=.*norm|bn).*$' + lr: 0.001 + weight_decay: 0. + - params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$' # except bias + weight_decay: 0. + + lr: 0.002 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +eval_spatial_size: [320, 320] +train_dataloader: + total_batch_size: 128 + dataset: + transforms: + ops: + - {type: Mosaic, output_size: 160, rotation_range: 10, translation_range: [0.1, 0.1], scaling_range: [0.5, 1.5], + probability: 1.0, fill_value: 0, use_cache: True, max_cached_images: 50, random_pop: True} + - {type: RandomPhotometricDistort, p: 0.5} + - {type: RandomZoomOut, fill: 0} + - {type: RandomIoUCrop, p: 0.8} + - {type: SanitizeBoundingBoxes, min_size: 12} + - {type: RandomHorizontalFlip} + - {type: Resize, size: [320, 320], } + - {type: SanitizeBoundingBoxes, min_size: 12} + - {type: ConvertPILImage, dtype: 'float32', scale: True} + - {type: ConvertBoxes, fmt: 'cxcywh', normalize: True} + policy: + epoch: [4, 250, 400] # list + mosaic_prob: 0.3 + + collate_fn: + mixup_prob: 0.0 + mixup_epochs: [40000, 15000] + copyblend_prob: 0.0 + copyblend_epochs: [40000, 15000] + + stop_epoch: 468 # 468 + 32 + ema_restart_decay: 0.9999 + base_size: 320 + base_size_repeat: ~ + +val_dataloader: + total_batch_size: 256 + dataset: + transforms: + ops: + - {type: Resize, size: [320, 320], } + - {type: ConvertPILImage, dtype: 'float32', scale: True} + shuffle: False + num_workers: 16 + + +DEIMCriterion: + losses: ['mal', 'boxes'] # , 'local' + use_uni_set: False + + matcher: + matcher_change_epoch: 450 # FIX This \ No newline at end of file diff --git a/configs/deimv2/deimv2_hgnetv2_femto_coco.yml b/configs/deimv2/deimv2_hgnetv2_femto_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..7a9a2952da71aa0d01a5fd5cc28b57b346f49837 --- /dev/null +++ b/configs/deimv2/deimv2_hgnetv2_femto_coco.yml @@ -0,0 +1,128 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml', +] + +output_dir: ./outputs/deimv2_hgnetv2_femto_coco + +DEIM: + encoder: LiteEncoder + +HGNetv2: + name: 'Femto' + return_idx: [2] + freeze_at: -1 + freeze_norm: False + use_lab: True + +LiteEncoder: + in_channels: [512] + feat_strides: [16] + + # intra + hidden_dim: 96 + + # cross + expansion: 0.34 + depth_mult: 0.5 + act: 'silu' + + +DEIMTransformer: + feat_channels: [96, 96] + feat_strides: [16, 32] + hidden_dim: 96 + num_levels: 2 + num_points: [4, 2] + + num_layers: 3 + eval_idx: -1 + num_queries: 150 + + # FFN + dim_feedforward: 256 + + # New options for DEIMv2 + share_bbox_head: True + use_gateway: False + +# Increase to search for the optimal ema +epoches: 500 # 468 + 32 + +## Our LR-Scheduler +warmup_iter: 4000 +flat_epoch: 250 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 32 +lr_gamma: 0.5 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm|bn).*$' + lr: 0.0008 + - + params: '^(?=.*backbone)(?=.*norm|bn).*$' + lr: 0.0008 + weight_decay: 0. + - # not opt + params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$' + weight_decay: 0. + + lr: 0.0016 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +eval_spatial_size: [416, 416] +train_dataloader: + total_batch_size: 128 + dataset: + transforms: + ops: + - {type: Mosaic, output_size: 208, rotation_range: 10, translation_range: [0.1, 0.1], scaling_range: [0.5, 1.5], + probability: 1.0, fill_value: 0, use_cache: True, max_cached_images: 50, random_pop: True} + - {type: RandomPhotometricDistort, p: 0.5} + - {type: RandomZoomOut, fill: 0} + - {type: RandomIoUCrop, p: 0.8} + - {type: SanitizeBoundingBoxes, min_size: 10} + - {type: RandomHorizontalFlip} + - {type: Resize, size: [416, 416], } + - {type: SanitizeBoundingBoxes, min_size: 10} + - {type: ConvertPILImage, dtype: 'float32', scale: True} + - {type: ConvertBoxes, fmt: 'cxcywh', normalize: True} + policy: + epoch: [4, 250, 400] # list + ops: ['Mosaic', 'RandomPhotometricDistort', 'RandomZoomOut', 'RandomIoUCrop'] + mosaic_prob: 0.5 + + collate_fn: + mixup_prob: 0.0 + mixup_epochs: [40000, 15000] + copyblend_prob: 0.0 + copyblend_epochs: [40000, 15000] + + stop_epoch: 468 # 468 + 32 + ema_restart_decay: 0.9999 + base_size: 416 + base_size_repeat: ~ + +val_dataloader: + total_batch_size: 256 + dataset: + transforms: + ops: + - {type: Resize, size: [416, 416], } + - {type: ConvertPILImage, dtype: 'float32', scale: True} + shuffle: False + num_workers: 16 + + +DEIMCriterion: + losses: ['mal', 'boxes'] # , 'local' + use_uni_set: False + + matcher: + matcher_change_epoch: 450 # FIX This \ No newline at end of file diff --git a/configs/deimv2/deimv2_hgnetv2_l_coco.yml b/configs/deimv2/deimv2_hgnetv2_l_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..0d94babe4b29ed306c1e1f5b1ab5352379242799 --- /dev/null +++ b/configs/deimv2/deimv2_hgnetv2_l_coco.yml @@ -0,0 +1,24 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml' +] + +output_dir: ./outputs/deimv2_hgnetv2_l_coco + + +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.000125 diff --git a/configs/deimv2/deimv2_hgnetv2_m_coco.yml b/configs/deimv2/deimv2_hgnetv2_m_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..d95fcf3bb45c04ab8ca1cd199a7ce47b2a95f683 --- /dev/null +++ b/configs/deimv2/deimv2_hgnetv2_m_coco.yml @@ -0,0 +1,72 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml' +] + +output_dir: ./outputs/deimv2_hgnetv2_m_coco + +HGNetv2: + name: 'B2' + return_idx: [1, 2, 3] + freeze_at: -1 + freeze_norm: False + use_lab: True + +HybridEncoder: + in_channels: [384, 768, 1536] + hidden_dim: 256 + depth_mult: 0.67 + +DEIMTransformer: + num_layers: 4 # 5 6 + eval_idx: -1 # -2 -3 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*bn).*$' + lr: 0.00004 + - + params: '^(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0004 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +# Increase to search for the optimal ema +epoches: 102 # 120 + 4n + +## Our LR-Scheduler +flat_epoch: 49 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 12 + +## Our DataAug +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 49, 90] # list + + collate_fn: + ema_restart_decay: 0.9999 + base_size_repeat: 6 + mixup_epochs: [4, 49] + stop_epoch: 90 + copyblend_prob: 0.5 + copyblend_epochs: [4, 90] + area_threshold: 100 + num_objects: 3 + with_expand: True + expand_ratios: [0.1, 0.25] + +DEIMCriterion: + matcher: + # new matcher + change_matcher: True + iou_order_alpha: 4.0 + matcher_change_epoch: 80 \ No newline at end of file diff --git a/configs/deimv2/deimv2_hgnetv2_n_coco.yml b/configs/deimv2/deimv2_hgnetv2_n_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..00ceea488b1eeec50204fb3464c4c725fafcc9fe --- /dev/null +++ b/configs/deimv2/deimv2_hgnetv2_n_coco.yml @@ -0,0 +1,96 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml' +] + +output_dir: ./outputs/deimv2_hgnetv2_n_coco + +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 + + version: 'dfine' + +DEIMTransformer: + feat_channels: [128, 128] + feat_strides: [16, 32] + hidden_dim: 128 + num_levels: 2 + num_points: [6, 6] + + num_layers: 3 + eval_idx: -1 + + # FFN + dim_feedforward: 512 + +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 +epoches: 160 # 148 + 12 + +## Our LR-Scheduler +flat_epoch: 7800 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 12 +lr_gamma: 1.0 + +## Our DataAug +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 78, 148] # list + + collate_fn: + ema_restart_decay: 0.9999 + base_size_repeat: ~ + mixup_epochs: [4, 78] + stop_epoch: 148 + copyblend_prob: 0.4 + copyblend_epochs: [4, 78] # CP half + area_threshold: 100 + num_objects: 3 + with_expand: True + expand_ratios: [0.1, 0.25] + +DEIMCriterion: + matcher: + # new matcher + change_matcher: True + iou_order_alpha: 4.0 + matcher_change_epoch: 136 \ No newline at end of file diff --git a/configs/deimv2/deimv2_hgnetv2_pico_coco.yml b/configs/deimv2/deimv2_hgnetv2_pico_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..71c29f65afd4db99b57d291de5789ad7d7b63240 --- /dev/null +++ b/configs/deimv2/deimv2_hgnetv2_pico_coco.yml @@ -0,0 +1,128 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml', +] + +output_dir: ./outputs/deimv2_hgnetv2_pico_coco + +DEIM: + encoder: LiteEncoder + decoder: DEIMTransformer + +HGNetv2: + name: 'Pico' + return_idx: [2] + freeze_at: -1 + freeze_norm: False + use_lab: True + +LiteEncoder: + in_channels: [512] + feat_strides: [16] + + # intra + hidden_dim: 112 + + # cross + expansion: 0.34 + depth_mult: 0.5 + act: 'silu' + + +DEIMTransformer: + feat_channels: [112, 112] + feat_strides: [16, 32] + hidden_dim: 112 + num_levels: 2 + num_points: [4, 2] + + num_layers: 3 + eval_idx: -1 + num_queries: 200 + + # FFN + dim_feedforward: 320 + + # New options for DEIMv2 + share_bbox_head: True + use_gateway: False + +# Increase to search for the optimal ema +epoches: 500 # 468 + 32 + +## Our LR-Scheduler +warmup_iter: 4000 +flat_epoch: 250 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 32 +lr_gamma: 0.5 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm|bn).*$' + lr: 0.0008 + - + params: '^(?=.*backbone)(?=.*norm|bn).*$' + lr: 0.0008 + weight_decay: 0. + - # not opt + params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$' + weight_decay: 0. + + lr: 0.0016 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +eval_spatial_size: [640, 640] +train_dataloader: + total_batch_size: 128 + dataset: + transforms: + ops: + - {type: Mosaic, output_size: 320, rotation_range: 10, translation_range: [0.1, 0.1], scaling_range: [0.5, 1.5], + probability: 1.0, fill_value: 0, use_cache: True, max_cached_images: 50, random_pop: True} + - {type: RandomPhotometricDistort, p: 0.5} + - {type: RandomZoomOut, fill: 0} + - {type: RandomIoUCrop, p: 0.8} + - {type: SanitizeBoundingBoxes, min_size: 8} + - {type: RandomHorizontalFlip} + - {type: Resize, size: [640, 640], } + - {type: SanitizeBoundingBoxes, min_size: 8} + - {type: ConvertPILImage, dtype: 'float32', scale: True} + - {type: ConvertBoxes, fmt: 'cxcywh', normalize: True} + policy: + epoch: [4, 250, 400] # list + ops: ['Mosaic', 'RandomPhotometricDistort', 'RandomZoomOut', 'RandomIoUCrop'] + mosaic_prob: 0.5 + + collate_fn: + mixup_prob: 0.0 + mixup_epochs: [40000, 15000] + copyblend_prob: 0.0 + copyblend_epochs: [40000, 15000] + stop_epoch: 468 # 468 + 32 + ema_restart_decay: 0.9999 + base_size: 640 + base_size_repeat: ~ + +val_dataloader: + total_batch_size: 256 + dataset: + transforms: + ops: + - {type: Resize, size: [640, 640], } + - {type: ConvertPILImage, dtype: 'float32', scale: True} + shuffle: False + num_workers: 16 + + +DEIMCriterion: + losses: ['mal', 'boxes'] # , 'local' + use_uni_set: False + + matcher: + matcher_change_epoch: 450 # FIX This \ No newline at end of file diff --git a/configs/deimv2/deimv2_hgnetv2_s_coco.yml b/configs/deimv2/deimv2_hgnetv2_s_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..b543760f9f759823f81a7abbb6f93858e584aa87 --- /dev/null +++ b/configs/deimv2/deimv2_hgnetv2_s_coco.yml @@ -0,0 +1,76 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml' +] + +output_dir: ./outputs/deimv2_hgnetv2_s_coco + +HGNetv2: + name: 'B0' + return_idx: [1, 2, 3] + freeze_at: -1 + freeze_norm: False + use_lab: True + +HybridEncoder: + in_channels: [256, 512, 1024] + hidden_dim: 256 + depth_mult: 0.34 + expansion: 0.5 + + version: 'dfine' + +DEIMTransformer: + num_layers: 3 # 4 5 6 + eval_idx: -1 # -2 -3 -4 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*bn).*$' + lr: 0.0002 + - + params: '^(?=.*(?:norm|bn)).*$' # except bias + weight_decay: 0. + + lr: 0.0004 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +# Increase to search for the optimal ema +epoches: 132 # 120 + 4n + +## Our LR-Scheduler +flat_epoch: 64 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 12 + +## Our DataAug +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 64, 120] # list + + collate_fn: + ema_restart_decay: 0.9999 + base_size_repeat: 20 + mixup_epochs: [4, 64] + stop_epoch: 120 + copyblend_prob: 0.5 + # copyblend_epochs: [4, 64] # from v11 to v12: copy-paste continues only half epochs + copyblend_epochs: [4, 120] + area_threshold: 100 + num_objects: 3 + with_expand: True + expand_ratios: [0.1, 0.25] + +DEIMCriterion: + matcher: + # new matcher + change_matcher: True + iou_order_alpha: 4.0 + matcher_change_epoch: 100 \ No newline at end of file diff --git a/configs/deimv2/deimv2_hgnetv2_x_coco.yml b/configs/deimv2/deimv2_hgnetv2_x_coco.yml new file mode 100644 index 0000000000000000000000000000000000000000..0355d6e314a4ef127b1c3bb25d2978a8cbedb4a5 --- /dev/null +++ b/configs/deimv2/deimv2_hgnetv2_x_coco.yml @@ -0,0 +1,60 @@ +__include__: [ + '../dataset/coco_detection.yml', + '../runtime.yml', + '../base/dataloader.yml', + '../base/optimizer.yml', + '../base/deimv2.yml' +] + +output_dir: ./outputs/deimv2_hgnetv2_x_coco + + +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 + +DEIMTransformer: + feat_channels: [384, 384, 384] # [256, 256, 256] + reg_scale: 8 # 4 + + # FFN + dim_feedforward: 2048 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm|bn).*$' + lr: 0.000005 + - + params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.0005 + betas: [0.9, 0.999] + weight_decay: 0.000125 + +# Increase to search for the optimal ema +epoches: 58 # 72 + 2n + +## Our LR-Scheduler +flat_epoch: 29 # 4 + epoch // 2, e.g., 40 = 4 + 72 / 2 +no_aug_epoch: 8 + +train_dataloader: + dataset: + transforms: + policy: + epoch: [4, 29, 50] # list + + collate_fn: + ema_restart_decay: 0.9998 + base_size_repeat: 3 diff --git a/configs/deimv2_floorplan.yaml b/configs/deimv2_floorplan.yaml new file mode 100644 index 0000000000000000000000000000000000000000..40371132d89c14bd027fbb434291e6dd37edfca9 --- /dev/null +++ b/configs/deimv2_floorplan.yaml @@ -0,0 +1,170 @@ +__include__: [ + '../configs/dataset/coco_detection.yml', # 10newdataset/coco_detection.yml (16クラス用) + '../configs/runtime.yml', + '../configs/base/dataloader.yml', + # '../configs/base/optimizer.yml', + '../configs/base/deimv2.yml', +] + +output_dir: ./outputs/9sdeimv2_dinov3_s_coco + +DEIM: + backbone: DINOv3STAs + +Model: + num_classes: 16 + class_names: ["kanki", "kanki_shikaku", "kanki_regisuta", "window1", "window2", "door1", "door2", "bathtub1", "konro1", "sink1", "toilet1", "kasaikeihou1", "kasaikeihou2", "houi1", "houi2", "houi3"] + +DINOv3STAs: + name: vit_tiny + embed_dim: 192 + weights_path: ./ckpts/vitt_distill.pt # 事前学習を使わないなら行ごと削除 + interaction_indexes: [3, 7, 11] + num_heads: 3 + +HybridEncoder: + in_channels: [192, 192, 192] + depth_mult: 0.67 + expansion: 0.34 + hidden_dim: 192 + dim_feedforward: 512 + +DEIMTransformer: + feat_channels: [192, 192, 192] + hidden_dim: 192 + dim_feedforward: 512 + num_layers: 4 # 4 5 6 + eval_idx: -1 # -2 -3 -4 + + +## Optimizer +optimizer: + type: AdamW + + params: + - + # except norm/bn/bias in self.dinov3 + params: '^(?=.*.dinov3)(?!.*(?:norm|bn|bias)).*$' + lr: 0.000025 + - + # including all norm/bn/bias in self.dinov3 + params: '^(?=.*.dinov3)(?=.*(?:norm|bn|bias)).*$' + lr: 0.000025 + weight_decay: 0. + - + # including all norm/bn/bias except for the self.dinov3 + params: '^(?=.*(?:sta|encoder|decoder))(?=.*(?:norm|bn|bias)).*$' + weight_decay: 0. + + lr: 0.0005 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +epoches: 400 +flat_epoch: 196 +no_aug_epoch: 46 + +# optimizer.ymlから必要な設定を手動で追加 +use_amp: True +use_ema: True +ema: + type: ModelEMA + decay: 0.9999 + warmups: 1000 + start: 0 + +clip_max_norm: 0.1 +sync_bn: True +find_unused_parameters: True + +# 学習率スケジューリング設定 +# CosineAnnealingLR専用設定（パラメータを最小限に） +lr_scheduler: + type: CosineAnnealingLR + T_max: 400 + eta_min: 0.0000001 + +# オプション2: MultiStepLR (段階的減少) +# lr_scheduler: +# type: MultiStepLR +# milestones: [120, 160] +# gamma: 0.1 + +# オプション3: OneCycleLR (1サイクル学習) +# lr_scheduler: +# type: OneCycleLR +# max_lr: 0.0005 +# total_steps: 200 +# pct_start: 0.3 +# anneal_strategy: 'cos' + +lr_warmup_scheduler: + type: LinearWarmup + warmup_duration: 1000 + +# 既存のflatcosineスケジューラーを無効化 +lrsheduler: null + +# deimv2.ymlのflatcosineスケジューラーも無効化 +lr_gamma: null +warmup_iter: null +flat_epoch: null +no_aug_epoch: null + + +# ---- Data Aug / Loader（図面＋640px＋OOM対策）---- +train_dataloader: + dataset: + transforms: + ops: + # 640でのピーク抑制のためMosaicは確率低め/スケール幅絞り + - {type: Mosaic, output_size: 640, rotation_range: 8, translation_range: [0.1, 0.1], + scaling_range: [0.9, 1.1], probability: 0.2, fill_value: 0, use_cache: True, + max_cached_images: 20, random_pop: True} + - {type: RandomPhotometricDistort, p: 0.2} + - {type: RandomZoomOut, fill: 0} + - {type: RandomIoUCrop, p: 0.6} + - {type: SanitizeBoundingBoxes, min_size: 1} + - {type: RandomHorizontalFlip} + - {type: RandomRotation, degrees: [90, 180, 270, 360], p: 0.5} # 修正版で有効化 + - {type: Resize, size: [640, 640]} # ★ 640固定 + - {type: SanitizeBoundingBoxes, min_size: 1} + - {type: ConvertPILImage, dtype: 'float32', scale: True} + - {type: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} + - {type: ConvertBoxes, fmt: 'cxcywh', normalize: True} + policy: + epoch: [8, 192, 352] # 400epochに合わせて調整 + + collate_fn: # 線画での崩れ防止＆メモリ抑制 + ema_restart_decay: 0.9999 + base_size_repeat: 1 # ★ 1にして実質マルチスケールOFF + stop_epoch: 352 # 400epochの90%程度で停止 + copyblend_epochs: [8, 352] # 400epochに合わせて調整 + + # 実装が読む場合のみ有効。読まない場合は base/dataloader.yml や起動引数で制御 + total_batch_size: 4 # ★ まずは 4 に落として安定化 + +val_dataloader: + dataset: + transforms: + ops: + - {type: Resize, size: [640, 640]} + - {type: ConvertPILImage, dtype: 'float32', scale: True} + - {type: Normalize, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225]} + total_batch_size: 6 # 評価も同程度に + +DEIMCriterion: + matcher: + change_matcher: True + iou_order_alpha: 4.0 + matcher_change_epoch: 300 + gamma: 1.5 + alpha: 0.75 + weight_dict: {loss_mal: 1, loss_bbox: 5, loss_giou: 2, loss_fgl: 0.15, loss_ddf: 1.5} + losses: [mal, boxes, local] + +# 出力設定 - 最後のエポック必ず保存 +output: + save_last: true + save_interval: 5 # チェックポイント保存間隔 + checkpoint_freq: 5 # 学習ループでの保存頻度 diff --git a/configs/runtime.yml b/configs/runtime.yml new file mode 100644 index 0000000000000000000000000000000000000000..8397ce1ff91246825e39f0530b544daaa0f891fe --- /dev/null +++ b/configs/runtime.yml @@ -0,0 +1,20 @@ +print_freq: 500 +output_dir: './logs' +checkpoint_freq: 12 + + +sync_bn: True +find_unused_parameters: True + + +use_amp: False +scaler: + type: GradScaler + enabled: True + + +use_ema: False +ema: + type: ModelEMA + decay: 0.9999 + warmups: 1000 diff --git a/detection.py b/detection.py new file mode 100644 index 0000000000000000000000000000000000000000..95472e340574b9424ae3aba4bf7fe2e77ed932bb --- /dev/null +++ b/detection.py @@ -0,0 +1,672 @@ +# detection.py +from functools import lru_cache +from typing import List, Tuple +import os +import numpy as np +import torch +import torch.nn as nn +from PIL import Image +import torchvision.transforms as T + +# デバッグ出力制御フラグ（環境変数で制御） +DEBUG_DEIMV2 = os.getenv("DEBUG_DEIMV2", "0") == "1" + +# YAMLConfigをインポート（engineパッケージ経由でレジストリをロード） +# モジュール登録のために、すべての必要なモジュールを明示的にインポート +# 重要: モジュールファイルを直接インポートすることで、@register()デコレータを確実に実行する + +# まず、engineパッケージ全体をインポート +try: + import engine + from engine import YAMLConfig +except ImportError: + from engine.core.yaml_config import YAMLConfig + +# 次に、すべての必要なモジュールファイルを直接インポート +# これにより、@register()デコレータが実行され、GLOBAL_CONFIGに登録される +try: + # Backboneモジュール + import engine.backbone.dinov3_adapter # DINOv3STAs + # DEIMモジュール（すべての重要なクラスを含む） + import engine.deim.hybrid_encoder # HybridEncoder - 必須 + import engine.deim.deim_decoder # DEIMTransformer - 必須 + import engine.deim.deim # DEIM - 必須 + import engine.deim.postprocessor # PostProcessor - 必須 + import engine.deim.matcher # HungarianMatcher + import engine.deim.deim_criterion # DEIMCriterion + # その他のモジュールもインポート（念のため） + import engine.deim + import engine.backbone + import engine.data + import engine.optim +except ImportError as e: + # インポートエラーは警告として出力（デバッグ用） + import warnings + warnings.warn(f"Some engine modules could not be imported: {e}") + +# (x1, y1, x2, y2, label_name, score) +Detection = Tuple[float, float, float, float, str, float] + +# ★ここを自分のファイル名に合わせる +MODEL_CONFIG_PATH = "configs/deimv2_floorplan.yaml" +MODEL_WEIGHTS_PATH = "models/best_stg2.pth" + + +def _get_device(): + """ + ZeroGPU対応: デバイスを遅延決定する。 + ZeroGPU環境では、import時点ではGPUが利用できないため、 + この関数を呼び出した時点でデバイスを決定する。 + """ + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + +# クラスID→記号名マッピング +# クラス名リスト: ["kanki", "kanki_shikaku", "kanki_regisuta", "window1", "window2", "door1", "door2", "bathtub1", "konro1", "sink1", "toilet1", "kasaikeihou1", "kasaikeihou2", "houi1", "houi2", "houi3"] +label_map = { + 0: "kanki", # kanki + 5: "door1", + 6: "door2", +} + + +@lru_cache(maxsize=1) +def load_deimv2_model(): + """ + HF Spaces 起動時に一度だけ呼ばれて、cfg + model + transform をキャッシュする。 + もとの main(args) でやっていた処理をここに移植。 + """ + import os + + # モジュール登録の確認と強制インポート + # トップレベルでインポート済みのはずだが、念のため確認して再インポート + from engine.core.workspace import GLOBAL_CONFIG + + # 必要なモジュールが登録されているか確認 + required_modules = { + 'HybridEncoder': 'engine.deim.hybrid_encoder', + 'DEIMTransformer': 'engine.deim.deim_decoder', + 'PostProcessor': 'engine.deim.postprocessor', + 'DINOv3STAs': 'engine.backbone.dinov3_adapter', + 'DEIM': 'engine.deim.deim', + } + + missing_modules = {name: module_path for name, module_path in required_modules.items() + if name not in GLOBAL_CONFIG} + + if missing_modules: + # まだ登録されていない場合は、強制的にインポート + import importlib + for name, module_path in missing_modules.items(): + try: + importlib.import_module(module_path) + except ModuleNotFoundError as e: + # 依存関係の問題を明確に示す + missing_dep = str(e).split("'")[1] if "'" in str(e) else str(e) + raise RuntimeError( + f"Failed to import module {name} from {module_path} due to missing dependency: {missing_dep}. " + f"Please install it with: pip install {missing_dep}. " + f"Available registered modules: {list(GLOBAL_CONFIG.keys())[:20]}" + ) + except Exception as e: + raise RuntimeError( + f"Failed to import and register module {name} from {module_path}: {e}. " + f"Available registered modules: {list(GLOBAL_CONFIG.keys())[:20]}" + ) + + # 再確認 + still_missing = [name for name in required_modules.keys() if name not in GLOBAL_CONFIG] + if still_missing: + raise RuntimeError( + f"Modules still not registered after import: {still_missing}. " + f"Available registered modules: {list(GLOBAL_CONFIG.keys())}" + ) + + # パスの確認（絶対パスに変換） + config_path = os.path.abspath(MODEL_CONFIG_PATH) + weights_path = os.path.abspath(MODEL_WEIGHTS_PATH) + + if not os.path.exists(config_path): + raise FileNotFoundError(f"設定ファイルが見つかりません: {config_path}") + if not os.path.exists(weights_path): + raise FileNotFoundError(f"モデルファイルが見つかりません: {weights_path}") + + try: + cfg = YAMLConfig(config_path, resume=weights_path) + except Exception as e: + raise RuntimeError(f"YAMLConfigの初期化に失敗しました: {e}") + + # もとのスクリプトと同じ処理 + if 'HGNetv2' in cfg.yaml_cfg: + cfg.yaml_cfg['HGNetv2']['pretrained'] = False + + # ZeroGPU対応: デバイスを遅延決定 + device = _get_device() + print(f"[DEBUG] 使用デバイス: {device}") + + try: + checkpoint = torch.load(weights_path, map_location=device) + if 'ema' in checkpoint: + state = checkpoint['ema']['module'] + else: + state = checkpoint['model'] + + # 訓練時と設定のズレがある場合でも動かせるように緩めにロード + model_state = cfg.model.state_dict() + + # デバッグ情報: チェックポイント内のデコーダー関連キーのリストアップ + decoder_keys_in_checkpoint = [k for k in state.keys() if 'decoder' in k.lower()] + print(f"[DEBUG] チェックポイント内のデコーダー関連キー数: {len(decoder_keys_in_checkpoint)}") + if decoder_keys_in_checkpoint: + print(f"[DEBUG] チェックポイント内のデコーダー関連キー（最初の20件）: {decoder_keys_in_checkpoint[:20]}") + + # デバッグ情報: モデル内のデコーダー関連キーのリストアップ + decoder_keys_in_model = [k for k in model_state.keys() if 'decoder' in k.lower()] + print(f"[DEBUG] モデル内のデコーダー関連キー数: {len(decoder_keys_in_model)}") + if decoder_keys_in_model: + print(f"[DEBUG] モデル内のデコーダー関連キー（最初の20件）: {decoder_keys_in_model[:20]}") + + # デバッグ情報: デコーダーパラメータの値の確認（読み込み前） + if decoder_keys_in_checkpoint: + first_decoder_key = decoder_keys_in_checkpoint[0] + if first_decoder_key in state: + print(f"[DEBUG] チェックポイント内のデコーダーパラメータ '{first_decoder_key}' の値の範囲: min={state[first_decoder_key].min():.6f}, max={state[first_decoder_key].max():.6f}, mean={state[first_decoder_key].mean():.6f}") + + compatible_state = {} + skipped = [] + for k, v in state.items(): + if k in model_state and model_state[k].shape == v.shape: + compatible_state[k] = v + else: + skipped.append(k) + + load_result = cfg.model.load_state_dict(compatible_state, strict=False) + + # デバッグ情報: キーのマッチング確認 + matched_decoder_keys = [k for k in decoder_keys_in_checkpoint if k in model_state and k in compatible_state] + unmatched_decoder_keys = [k for k in decoder_keys_in_checkpoint if k not in model_state or k not in compatible_state] + print(f"[DEBUG] マッチしたデコーダー関連キー数: {len(matched_decoder_keys)}") + print(f"[DEBUG] マッチしなかったデコーダー関連キー数: {len(unmatched_decoder_keys)}") + if unmatched_decoder_keys: + print(f"[DEBUG] マッチしなかったデコーダー関連キー（最初の20件）: {unmatched_decoder_keys[:20]}") + + # デバッグ情報: デコーダーパラメータの値の確認（読み込み後） + decoder_params_after = {k: v for k, v in cfg.model.named_parameters() if 'decoder' in k.lower()} + if decoder_params_after: + first_decoder_param_after = next(iter(decoder_params_after.values())) + print(f"[DEBUG] 読み込み後のデコーダーパラメータの値の範囲: min={first_decoder_param_after.min():.6f}, max={first_decoder_param_after.max():.6f}, mean={first_decoder_param_after.mean():.6f}") + # すべてのデコーダーパラメータの値の範囲も確認 + all_decoder_values = torch.cat([p.flatten() for p in decoder_params_after.values()]) + print(f"[DEBUG] 読み込み後の全デコーダーパラメータの値の範囲: min={all_decoder_values.min():.6f}, max={all_decoder_values.max():.6f}, mean={all_decoder_values.mean():.6f}") + + # デバッグ情報: 読み込み統計 + print(f"[DEBUG] チェックポイント読み込み統計:") + print(f" - チェックポイント内のキー数: {len(state)}") + print(f" - モデル内のキー数: {len(model_state)}") + print(f" - 読み込んだキー数: {len(compatible_state)}") + print(f" - 形状不一致でスキップ: {len(skipped)}") + print(f" - 読み込み後の欠落キー: {len(load_result.missing_keys) if load_result.missing_keys else 0}") + print(f" - 読み込み後の予期しないキー: {len(load_result.unexpected_keys) if load_result.unexpected_keys else 0}") + + if skipped or load_result.missing_keys or load_result.unexpected_keys: + print("Warning: partial checkpoint load.") + if skipped: + print(f" shape-mismatched skipped keys: {len(skipped)}") + # 重要なキー（decoder, head関連）を優先表示 + important_skipped = [k for k in skipped if any(x in k for x in ['decoder', 'head', 'class', 'bbox', 'query'])] + if important_skipped: + print(f" [重要] スキップされたキー（decoder/head関連）: {important_skipped[:10]}") + if len(skipped) <= 20: + print(f" すべてのスキップされたキー: {skipped}") + else: + print(f" スキップされたキー（最初の20件）: {skipped[:20]}") + if load_result.missing_keys: + print(f" missing keys after load: {len(load_result.missing_keys)}") + # 重要なキー（decoder, head関連）を優先表示 + important_missing = [k for k in load_result.missing_keys if any(x in k for x in ['decoder', 'head', 'class', 'bbox', 'query'])] + if important_missing: + print(f" [重要] 欠落キー（decoder/head関連）: {important_missing[:20]}") + if len(load_result.missing_keys) <= 30: + print(f" すべての欠落キー: {load_result.missing_keys}") + else: + print(f" 欠落キー（最初の30件）: {list(load_result.missing_keys)[:30]}") + if load_result.unexpected_keys: + print(f" unexpected keys after load: {len(load_result.unexpected_keys)}") + if len(load_result.unexpected_keys) <= 20: + print(f" 予期しないキー: {load_result.unexpected_keys}") + else: + print(f" 予期しないキー（最初の20件）: {load_result.unexpected_keys[:20]}") + else: + print(f"[DEBUG] モデル重みの読み込み: 成功 (読み込んだキー数: {len(compatible_state)})") + + # デバッグ情報: モデルのパラメータ統計（読み込み後） + total_params = sum(p.numel() for p in cfg.model.parameters()) + trainable_params = sum(p.numel() for p in cfg.model.parameters() if p.requires_grad) + print(f"[DEBUG] モデルパラメータ統計:") + print(f" - 総パラメータ数: {total_params:,}") + print(f" - 学習可能パラメータ数: {trainable_params:,}") + + # デバッグ情報: デコーダーとヘッドのパラメータが初期化されているか確認 + decoder_params = {k: v for k, v in cfg.model.named_parameters() if 'decoder' in k} + head_params = {k: v for k, v in cfg.model.named_parameters() if any(x in k for x in ['head', 'class', 'bbox'])} + print(f"[DEBUG] デコーダー/ヘッドパラメータ:") + print(f" - デコーダーパラメータ数: {len(decoder_params)}") + print(f" - ヘッドパラメータ数: {len(head_params)}") + if decoder_params: + # 最初のデコーダーパラメータの統計を確認 + first_decoder_param = next(iter(decoder_params.values())) + print(f" - デコーダーパラメータの値の範囲: min={first_decoder_param.min():.6f}, max={first_decoder_param.max():.6f}, mean={first_decoder_param.mean():.6f}") + if head_params: + # 最初のヘッドパラメータの統計を確認 + first_head_param = next(iter(head_params.values())) + print(f" - ヘッドパラメータの値の範囲: min={first_head_param.min():.6f}, max={first_head_param.max():.6f}, mean={first_head_param.mean():.6f}") + except Exception as e: + raise RuntimeError(f"モデルの重みの読み込みに失敗しました: {e}") + + class Model(nn.Module): + def __init__(self, cfg, device): + super().__init__() + self.device = device + self.model = cfg.model.eval().to(device) + self.postprocessor = cfg.postprocessor.eval().to(device) + + # デバッグ情報: モデルとポストプロセッサの設定 + print(f"[DEBUG] モデル構築:") + print(f" - モデルタイプ: {type(self.model).__name__}") + print(f" - ポストプロセッサタイプ: {type(self.postprocessor).__name__}") + if hasattr(self.postprocessor, 'use_focal_loss'): + print(f" - use_focal_loss: {self.postprocessor.use_focal_loss}") + if hasattr(self.postprocessor, 'num_classes'): + print(f" - num_classes: {self.postprocessor.num_classes}") + if hasattr(self.postprocessor, 'num_top_queries'): + print(f" - num_top_queries: {self.postprocessor.num_top_queries}") + + def forward(self, images, orig_target_sizes): + outputs = self.model(images) + outputs = self.postprocessor(outputs, orig_target_sizes) + return outputs + + model = Model(cfg, device) + + # eval_spatial_sizeが設定にない場合は、val_dataloaderのResizeサイズから取得 + # デフォルトは640x640 + if "eval_spatial_size" in cfg.yaml_cfg: + img_size = cfg.yaml_cfg["eval_spatial_size"] + else: + # val_dataloaderのtransformsから取得を試みる + val_transforms = cfg.yaml_cfg.get("val_dataloader", {}).get("dataset", {}).get("transforms", {}).get("ops", []) + img_size = 640 # デフォルト値 + for op in val_transforms: + if isinstance(op, dict) and op.get("type") == "Resize": + size = op.get("size", [640, 640]) + img_size = size[0] if isinstance(size, list) else size + break + + vit_backbone = cfg.yaml_cfg.get('DINOv3STAs', False) + + if vit_backbone: + transforms = T.Compose([ + T.Resize(img_size), + T.ToTensor(), + T.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225], + ), + ]) + else: + transforms = T.Compose([ + T.Resize(img_size), + T.ToTensor(), + ]) + + return model, transforms + + +def run_inference_single_tile( + model, + transforms, + tile_pil: Image.Image, + tile_x: int, + tile_y: int, + tile_w: int, + tile_h: int, + score_thresh: float = 0.8, +) -> List[Detection]: + """ + 単一タイルに対する推論を実行し、元の画像座標系に変換して返す。 + """ + tile_w_actual, tile_h_actual = tile_pil.size + + # ZeroGPU対応: モデルからデバイスを動的に取得 + device = next(model.model.parameters()).device + + # デバッグ情報: タイルの基本情報（DEBUG_DEIMV2フラグで制御） + if DEBUG_DEIMV2: + is_first_few = tile_x < 2000 and tile_y < 2000 # 最初の数タイルのみ + if is_first_few: + print(f"[DEBUG] タイル処理開始: 座標({tile_x},{tile_y}), サイズ{tile_w}×{tile_h}") + print(f"[DEBUG] PIL画像サイズ: {tile_w_actual}×{tile_h_actual}") + # タイル画像の統計情報 + tile_np = np.array(tile_pil) + print(f"[DEBUG] タイル画像値の範囲: min={tile_np.min()}, max={tile_np.max()}, mean={tile_np.mean():.2f}") + + # タイルサイズをorig_sizeとして設定 + orig_size = torch.tensor([[tile_w_actual, tile_h_actual]], device=device).float() + + # 前処理 + im_tensor = transforms(tile_pil).unsqueeze(0).to(device) + + # デバッグ情報: リサイズ後の確認（DEBUG_DEIMV2フラグで制御） + if DEBUG_DEIMV2: + is_first_few = tile_x < 2000 and tile_y < 2000 + if is_first_few: + print(f"[DEBUG] リサイズ後のテンソル形状: {im_tensor.shape}") + # リサイズ前後のサイズ比較 + if hasattr(transforms, 'transforms'): + for t in transforms.transforms: + if isinstance(t, T.Resize): + print(f"[DEBUG] Resize設定: {t.size}") + break + print(f"[DEBUG] 前処理後のテンソル: shape={im_tensor.shape}, dtype={im_tensor.dtype}") + print(f"[DEBUG] テンソル値の範囲: min={im_tensor.min():.4f}, max={im_tensor.max():.4f}, mean={im_tensor.mean():.4f}") + print(f"[DEBUG] orig_size: {orig_size}") + + with torch.no_grad(): + # デバッグモードの場合のみ、モデルの生の出力を確認（推論を2回実行） + if DEBUG_DEIMV2: + model_outputs = model.model(im_tensor) + is_first_few = tile_x < 2000 and tile_y < 2000 + if is_first_few: + print(f"[DEBUG] モデル生出力: pred_logits.shape={model_outputs['pred_logits'].shape}, pred_boxes.shape={model_outputs['pred_boxes'].shape}") + print(f"[DEBUG] pred_logits範囲: min={model_outputs['pred_logits'].min():.4f}, max={model_outputs['pred_logits'].max():.4f}, mean={model_outputs['pred_logits'].mean():.4f}") + # クラス別のlogitsの最大値を確認 + logits_max_per_class = model_outputs['pred_logits'].max(dim=1)[0] # [1, 16] + print(f"[DEBUG] クラス別最大logits: {logits_max_per_class[0].cpu().numpy()}") + # sigmoid後のスコアも確認 + scores_raw = torch.sigmoid(model_outputs['pred_logits']) + scores_max_per_class = scores_raw.max(dim=1)[0] # [1, 16] + print(f"[DEBUG] クラス別最大スコア（sigmoid後）: {scores_max_per_class[0].cpu().numpy()}") + print(f"[DEBUG] pred_boxes範囲: min={model_outputs['pred_boxes'].min():.4f}, max={model_outputs['pred_boxes'].max():.4f}, mean={model_outputs['pred_boxes'].mean():.4f}") + print(f"[DEBUG] pred_boxes形状（cxcywh形式）: 最初の5件={model_outputs['pred_boxes'][0, :5, :]}") + + # ポストプロセッサ前の座標変換を確認 + import torchvision.ops + bbox_pred_raw = torchvision.ops.box_convert(model_outputs['pred_boxes'], in_fmt='cxcywh', out_fmt='xyxy') + print(f"[DEBUG] cxcywh→xyxy変換後（正規化座標）: 最初の5件={bbox_pred_raw[0, :5, :]}") + print(f"[DEBUG] xyxy変換後の範囲: min={bbox_pred_raw.min():.4f}, max={bbox_pred_raw.max():.4f}") + + # 本番推論（1回のみ実行） + outputs = model(im_tensor, orig_size) + + if not outputs or len(outputs) == 0: + if DEBUG_DEIMV2: + print(f"[DEBUG] モデル出力が空です") + return [] + + out = outputs[0] + labels = out['labels'].detach().cpu().numpy() + boxes = out['boxes'].detach().cpu().numpy() + scores = out['scores'].detach().cpu().numpy() + + # デバッグ情報: ポストプロセッサ後の出力（DEBUG_DEIMV2フラグで制御） + if DEBUG_DEIMV2: + print(f"[DEBUG] ポストプロセッサ後: labels.shape={labels.shape}, boxes.shape={boxes.shape}, scores.shape={scores.shape}") + print(f"[DEBUG] boxes範囲: x1=[{boxes[:, 0].min():.1f}, {boxes[:, 0].max():.1f}], y1=[{boxes[:, 1].min():.1f}, {boxes[:, 1].max():.1f}], x2=[{boxes[:, 2].min():.1f}, {boxes[:, 2].max():.1f}], y2=[{boxes[:, 3].min():.1f}, {boxes[:, 3].max():.1f}]") + + # デバッグ情報: タイルの検出結果詳細 + print(f"[DEBUG] タイル検出結果（フィルタリング前）: {len(scores)}件") + if len(scores) > 0: + print(f"[DEBUG] スコア範囲: min={scores.min():.4f}, max={scores.max():.4f}, mean={scores.mean():.4f}") + print(f"[DEBUG] スコア分布:") + print(f" - 0.0-0.3: {(scores < 0.3).sum()}件") + print(f" - 0.3-0.5: {((scores >= 0.3) & (scores < 0.5)).sum()}件") + print(f" - 0.5-0.7: {((scores >= 0.5) & (scores < 0.7)).sum()}件") + print(f" - 0.7-0.9: {((scores >= 0.7) & (scores < 0.9)).sum()}件") + print(f" - 0.9-1.0: {(scores >= 0.9).sum()}件") + + # ラベルの分布 + unique_labels, label_counts = np.unique(labels, return_counts=True) + print(f"[DEBUG] ラベル分布:") + for label_id, count in zip(unique_labels, label_counts): + label_name = label_map.get(int(label_id), f"class_{int(label_id)}") + print(f" - クラスID {int(label_id)} ({label_name}): {count}件") + + # スコア閾値以上の検出数 + above_thresh = scores >= score_thresh + print(f"[DEBUG] スコア閾値({score_thresh})以上の検出数: {above_thresh.sum()}件") + + # 上位5件の詳細を表示 + if len(scores) > 0: + top_indices = np.argsort(scores)[::-1][:5] + print(f"[DEBUG] 上位5件の検出結果:") + for rank, idx in enumerate(top_indices, 1): + label_id = int(labels[idx]) + label_name = label_map.get(label_id, f"class_{label_id}") + score = float(scores[idx]) + x1, y1, x2, y2 = boxes[idx] + print(f" [{rank}] {label_name}, スコア={score:.4f}, bbox=({x1:.1f},{y1:.1f},{x2:.1f},{y2:.1f})") + + detections: List[Detection] = [] + filtered_by_thresh = 0 + filtered_by_label = 0 + + for label, box, score in zip(labels, boxes, scores): + score = float(score) + if score < score_thresh: + filtered_by_thresh += 1 + continue + + x1, y1, x2, y2 = [float(v) for v in box.tolist()] + label_id = int(label) + label_name = label_map.get(label_id, f"class_{label_id}") + + # label_mapに存在しないクラスもカウント（デバッグ用） + if label_id not in label_map: + filtered_by_label += 1 + + # タイル座標を元の画像座標に変換 + x1_orig = x1 + tile_x + y1_orig = y1 + tile_y + x2_orig = x2 + tile_x + y2_orig = y2 + tile_y + + detections.append((x1_orig, y1_orig, x2_orig, y2_orig, label_name, score)) + + if DEBUG_DEIMV2 and (filtered_by_thresh > 0 or filtered_by_label > 0): + print(f"[DEBUG] フィルタリング: スコア閾値で{filtered_by_thresh}件、label_mapで{filtered_by_label}件除外") + + return detections + + +def run_inference( + image_np: np.ndarray, + score_thresh: float = 0.8, + tile_size: int = 640, + tile_overlap: int = 128, +) -> List[Detection]: + """ + タイル推論を実行する。 + 大きな画像をタイルに分割して推論し、結果を統合する。 + + Args: + image_np: RGB np.ndarray (H, W, 3) + score_thresh: スコア閾値 + tile_size: タイルサイズ（デフォルト: 640） + tile_overlap: タイル間のオーバーラップ（デフォルト: 128） + + Returns: + [(x1,y1,x2,y2,label_name,score), ...] + """ + try: + model, transforms = load_deimv2_model() + except Exception as e: + raise RuntimeError(f"モデルの読み込みに失敗しました: {e}") + + try: + # numpy → PIL + if image_np.dtype != np.uint8: + if image_np.max() <= 1.0: + image_np = (image_np * 255).astype(np.uint8) + else: + image_np = image_np.astype(np.uint8) + + im_pil = Image.fromarray(image_np).convert("RGB") + img_w, img_h = im_pil.size + + if DEBUG_DEIMV2: + print(f"[DEBUG] ===== タイル推論開始 =====") + print(f"[DEBUG] 入力画像サイズ: {img_w}×{img_h} (ピクセル数: {img_w*img_h:,})") + print(f"[DEBUG] タイルサイズ: {tile_size}×{tile_size}") + print(f"[DEBUG] タイルオーバーラップ: {tile_overlap}px") + + # タイルに分割 + step = tile_size - tile_overlap + tiles = [] + tile_coords = [] + + for y in range(0, img_h, step): + for x in range(0, img_w, step): + # タイルの範囲を計算 + tile_x = x + tile_y = y + tile_x_end = min(x + tile_size, img_w) + tile_y_end = min(y + tile_size, img_h) + tile_w = tile_x_end - tile_x + tile_h = tile_y_end - tile_y + + # タイルを切り出し + tile = im_pil.crop((tile_x, tile_y, tile_x_end, tile_y_end)) + + # デバッグ情報: タイルの確認（DEBUG_DEIMV2フラグで制御） + if DEBUG_DEIMV2 and len(tiles) < 3: + print(f"[DEBUG] タイル {len(tiles)+1} 詳細:") + print(f" - 切り出し範囲: ({tile_x}, {tile_y}) → ({tile_x_end}, {tile_y_end})") + print(f" - タイルサイズ: {tile_w}×{tile_h}") + print(f" - 実際のPIL画像サイズ: {tile.size}") + # タイル画像の統計情報 + tile_np = np.array(tile) + print(f" - 画像値の範囲: min={tile_np.min()}, max={tile_np.max()}, mean={tile_np.mean():.2f}") + # タイル画像を保存（デバッグ用、最初の3タイルのみ） + try: + debug_dir = "debug_tiles" + os.makedirs(debug_dir, exist_ok=True) + tile.save(f"{debug_dir}/tile_{len(tiles)+1}_x{tile_x}_y{tile_y}.png") + print(f" - タイル画像を保存: {debug_dir}/tile_{len(tiles)+1}_x{tile_x}_y{tile_y}.png") + except Exception as e: + print(f" - タイル画像の保存に失敗: {e}") + + tiles.append(tile) + tile_coords.append((tile_x, tile_y, tile_w, tile_h)) + + if DEBUG_DEIMV2: + print(f"[DEBUG] タイル数: {len(tiles)}") + + # 各タイルに対して推論 + all_detections = [] + for i, (tile, (tile_x, tile_y, tile_w, tile_h)) in enumerate(zip(tiles, tile_coords)): + if DEBUG_DEIMV2: + print(f"[DEBUG] タイル {i+1}/{len(tiles)}: 座標({tile_x},{tile_y}), サイズ{tile_w}×{tile_h}") + + tile_detections = run_inference_single_tile( + model, transforms, tile, + tile_x, tile_y, tile_w, tile_h, + score_thresh=score_thresh + ) + + if DEBUG_DEIMV2: + print(f"[DEBUG] 検出数: {len(tile_detections)}件") + all_detections.extend(tile_detections) + + if DEBUG_DEIMV2: + print(f"[DEBUG] 総検出数（重複あり）: {len(all_detections)}件") + + if len(all_detections) == 0: + if DEBUG_DEIMV2: + print(f"[DEBUG] =========================") + return [] + + # 全検出結果の統計情報（DEBUG_DEIMV2フラグで制御） + if DEBUG_DEIMV2 and len(all_detections) > 0: + all_scores = [det[5] for det in all_detections] # scoreは6番目の要素 + all_labels = [det[4] for det in all_detections] # label_nameは5番目の要素 + + print(f"[DEBUG] 全タイル統合後の統計:") + print(f" - スコア範囲: min={min(all_scores):.4f}, max={max(all_scores):.4f}, mean={sum(all_scores)/len(all_scores):.4f}") + + # ラベルごとの集計 + from collections import Counter + label_counter = Counter(all_labels) + print(f" - ラベル別検出数:") + for label_name, count in sorted(label_counter.items(), key=lambda x: -x[1]): + print(f" - {label_name}: {count}件") + + # NMS（Non-Maximum Suppression）で重複検出をマージ + # クラスごとにNMSを適用（異なるクラス間の重複は許可） + from torchvision.ops import nms + + # ZeroGPU対応: モデルからデバイスを動的に取得 + device = next(model.model.parameters()).device + + # クラスごとにグループ化 + detections_by_class = {} + for det in all_detections: + x1, y1, x2, y2, label_name, score = det + if label_name not in detections_by_class: + detections_by_class[label_name] = [] + detections_by_class[label_name].append((x1, y1, x2, y2, score)) + + if DEBUG_DEIMV2: + print(f"[DEBUG] NMS適用前: {len(detections_by_class)}クラス、合計{len(all_detections)}件") + + merged_detections = [] + nms_removed = 0 + for label_name, boxes_scores in detections_by_class.items(): + if len(boxes_scores) == 0: + continue + + before_nms = len(boxes_scores) + + # テンソルに変換 + boxes_tensor = torch.tensor([[x1, y1, x2, y2] for x1, y1, x2, y2, _ in boxes_scores], device=device) + scores_tensor = torch.tensor([score for _, _, _, _, score in boxes_scores], device=device) + + # NMS適用（IoU閾値: 0.4 - より厳しく重複を削除） + keep_indices = nms(boxes_tensor, scores_tensor, iou_threshold=0.4) + + # マージ後の検出を追加 + for idx in keep_indices.cpu().numpy(): + x1, y1, x2, y2, score = boxes_scores[idx] + merged_detections.append((x1, y1, x2, y2, label_name, score)) + + after_nms = len(keep_indices) + removed = before_nms - after_nms + nms_removed += removed + if DEBUG_DEIMV2 and removed > 0: + print(f"[DEBUG] {label_name}: NMSで{removed}件削除 ({before_nms}→{after_nms}件)") + + if DEBUG_DEIMV2: + print(f"[DEBUG] NMS適用後: {len(merged_detections)}件 (合計{nms_removed}件削除)") + + # 最終結果の統計 + if len(merged_detections) > 0: + final_scores = [det[5] for det in merged_detections] + final_labels = [det[4] for det in merged_detections] + final_label_counter = Counter(final_labels) + + print(f"[DEBUG] 最終検出結果:") + print(f" - 総検出数: {len(merged_detections)}件") + print(f" - スコア範囲: min={min(final_scores):.4f}, max={max(final_scores):.4f}, mean={sum(final_scores)/len(final_scores):.4f}") + print(f" - ラベル別:") + for label_name, count in sorted(final_label_counter.items(), key=lambda x: -x[1]): + print(f" - {label_name}: {count}件") + + # 上位10件を表示 + sorted_detections = sorted(merged_detections, key=lambda x: x[5], reverse=True)[:10] + print(f"[DEBUG] 上位10件の検出結果:") + for rank, (x1, y1, x2, y2, label_name, score) in enumerate(sorted_detections, 1): + print(f" [{rank}] {label_name}, スコア={score:.4f}, bbox=({x1:.1f},{y1:.1f},{x2:.1f},{y2:.1f})") + + print(f"[DEBUG] =========================") + + return merged_detections + except Exception as e: + raise RuntimeError(f"推論の実行に失敗しました: {e}") diff --git a/engine/__init__.py b/engine/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..69baa01f55ae4799118a52fb6290ae7a2006d87a --- /dev/null +++ b/engine/__init__.py @@ -0,0 +1,16 @@ +""" +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +""" + +# for register purpose +from . import optim +from . import data +from . import deim + +from .backbone import * + +from .backbone import ( + get_activation, + FrozenBatchNorm2d, + freeze_batch_norm2d, +) \ No newline at end of file diff --git a/engine/backbone/__init__.py b/engine/backbone/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5ab896cb7b4ffe931cd9142e3c175ad3e6354e83 --- /dev/null +++ b/engine/backbone/__init__.py @@ -0,0 +1,22 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +from .common import ( + get_activation, + FrozenBatchNorm2d, + freeze_batch_norm2d, +) +from .presnet import PResNet +from .test_resnet import MResNet + +from .timm_model import TimmModel +from .torchvision_model import TorchVisionModel + +from .csp_resnet import CSPResNet +from .csp_darknet import CSPDarkNet, CSPPAN + +from .hgnetv2 import HGNetv2 + +from .dinov3_adapter import * diff --git a/engine/backbone/common.py b/engine/backbone/common.py new file mode 100644 index 0000000000000000000000000000000000000000..bcbe0313df9f5467ff58ac7833370ddd636e4238 --- /dev/null +++ b/engine/backbone/common.py @@ -0,0 +1,116 @@ +""" +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 diff --git a/engine/backbone/csp_darknet.py b/engine/backbone/csp_darknet.py new file mode 100644 index 0000000000000000000000000000000000000000..a1f949358817b3dcaa932f716ec7ef0383ce89c8 --- /dev/null +++ b/engine/backbone/csp_darknet.py @@ -0,0 +1,179 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import math +import warnings + +from .common import get_activation +from ..core import register + + +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., 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]) diff --git a/engine/backbone/csp_resnet.py b/engine/backbone/csp_resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..2d22b28c232be2a9fccbd73363e92207921380f8 --- /dev/null +++ b/engine/backbone/csp_resnet.py @@ -0,0 +1,277 @@ +""" +https://github.com/PaddlePaddle/PaddleDetection/blob/release/2.6/ppdet/modeling/backbones/cspresnet.py + +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections import OrderedDict + +from .common import get_activation + +from ..core import register + +__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 diff --git a/engine/backbone/dinov3/__init__.py b/engine/backbone/dinov3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cff9e183ee08080f7d20c555c3f49d563244521a --- /dev/null +++ b/engine/backbone/dinov3/__init__.py @@ -0,0 +1 @@ +from .vision_transformer import DinoVisionTransformer \ No newline at end of file diff --git a/engine/backbone/dinov3/layers/__init__.py b/engine/backbone/dinov3/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5b82c261aaba6f0b7b871662f3549062e69928f8 --- /dev/null +++ b/engine/backbone/dinov3/layers/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +from .attention import CausalSelfAttention, LinearKMaskedBias, SelfAttention +from .block import CausalSelfAttentionBlock, SelfAttentionBlock +from .ffn_layers import Mlp, SwiGLUFFN +from .fp8_linear import convert_linears_to_fp8 +from .layer_scale import LayerScale +from .patch_embed import PatchEmbed +from .rms_norm import RMSNorm +from .rope_position_encoding import RopePositionEmbedding diff --git a/engine/backbone/dinov3/layers/attention.py b/engine/backbone/dinov3/layers/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..de66d3e4312550b177638a88aed61b29e413771f --- /dev/null +++ b/engine/backbone/dinov3/layers/attention.py @@ -0,0 +1,164 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import math +from typing import List, Tuple + +import torch +import torch.nn.functional as F +from ..utils import cat_keep_shapes, uncat_with_shapes +from torch import Tensor, nn + + +# RoPE-related functions: +def rope_rotate_half(x: Tensor) -> Tensor: + # x: [ x0 x1 x2 x3 x4 x5] + # out: [-x3 -x4 -x5 x0 x1 x2] + x1, x2 = x.chunk(2, dim=-1) + return torch.cat([-x2, x1], dim=-1) + + +def rope_apply(x: Tensor, sin: Tensor, cos: Tensor) -> Tensor: + # x: [..., D], eg [x0, x1, x2, x3, x4, x5] + # sin: [..., D], eg [sin0, sin1, sin2, sin0, sin1, sin2] + # cos: [..., D], eg [cos0, cos1, cos2, cos0, cos1, cos2] + return (x * cos) + (rope_rotate_half(x) * sin) + + +class LinearKMaskedBias(nn.Linear): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + o = self.out_features + assert o % 3 == 0 + if self.bias is not None: + self.register_buffer("bias_mask", torch.full_like(self.bias, fill_value=math.nan)) + + def forward(self, input: Tensor) -> Tensor: + masked_bias = self.bias * self.bias_mask.to(self.bias.dtype) if self.bias is not None else None + return F.linear(input, self.weight, masked_bias) + + +class SelfAttention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + proj_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + mask_k_bias: bool = False, + device=None, + ) -> None: + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + + linear_class = LinearKMaskedBias if mask_k_bias else nn.Linear + self.qkv = linear_class(dim, dim * 3, bias=qkv_bias, device=device) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim, bias=proj_bias, device=device) + self.proj_drop = nn.Dropout(proj_drop) + + def apply_rope(self, q: Tensor, k: Tensor, rope: Tensor | Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]: + # All operations will use the dtype of rope, the output is cast back to the dtype of q and k + q_dtype = q.dtype + k_dtype = k.dtype + sin, cos = rope + rope_dtype = sin.dtype + q = q.to(dtype=rope_dtype) + k = k.to(dtype=rope_dtype) + N = q.shape[-2] + prefix = N - sin.shape[-2] + assert prefix >= 0 + q_prefix = q[:, :, :prefix, :] + q = rope_apply(q[:, :, prefix:, :], sin, cos) # [B, head, hw, D//head] + q = torch.cat((q_prefix, q), dim=-2) # [B, head, N, D//head] + k_prefix = k[:, :, :prefix, :] + k = rope_apply(k[:, :, prefix:, :], sin, cos) # [B, head, hw, D//head] + k = torch.cat((k_prefix, k), dim=-2) # [B, head, N, D//head] + q = q.to(dtype=q_dtype) + k = k.to(dtype=k_dtype) + return q, k + + def forward(self, x: Tensor, attn_bias=None, rope: Tensor = None) -> Tensor: + qkv = self.qkv(x) + attn_v = self.compute_attention(qkv=qkv, attn_bias=attn_bias, rope=rope) + x = self.proj(attn_v) + x = self.proj_drop(x) + return x + + def forward_list(self, x_list, attn_bias=None, rope_list=None) -> List[Tensor]: + assert len(x_list) == len(rope_list) # should be enforced by the Block + x_flat, shapes, num_tokens = cat_keep_shapes(x_list) + qkv_flat = self.qkv(x_flat) + qkv_list = uncat_with_shapes(qkv_flat, shapes, num_tokens) + att_out = [] + for _, (qkv, _, rope) in enumerate(zip(qkv_list, shapes, rope_list)): + att_out.append(self.compute_attention(qkv, attn_bias=attn_bias, rope=rope)) + x_flat, shapes, num_tokens = cat_keep_shapes(att_out) + x_flat = self.proj(x_flat) + return uncat_with_shapes(x_flat, shapes, num_tokens) + + def compute_attention(self, qkv: Tensor, attn_bias=None, rope=None) -> Tensor: + assert attn_bias is None + B, N, _ = qkv.shape + C = self.qkv.in_features + + qkv = qkv.reshape(B, N, 3, self.num_heads, C // self.num_heads) + q, k, v = torch.unbind(qkv, 2) + q, k, v = [t.transpose(1, 2) for t in [q, k, v]] + if rope is not None: + q, k = self.apply_rope(q, k, rope) + x = torch.nn.functional.scaled_dot_product_attention(q, k, v) + x = x.transpose(1, 2) + return x.reshape([B, N, C]) + + +class CausalSelfAttention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + proj_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ) -> None: + super().__init__() + self.dim = dim + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = attn_drop + self.proj = nn.Linear(dim, dim, bias=proj_bias) + self.proj_drop = nn.Dropout(proj_drop) + + def init_weights( + self, init_attn_std: float | None = None, init_proj_std: float | None = None, factor: float = 1.0 + ) -> None: + init_attn_std = init_attn_std or (self.dim**-0.5) + init_proj_std = init_proj_std or init_attn_std * factor + nn.init.normal_(self.qkv.weight, std=init_attn_std) + nn.init.normal_(self.proj.weight, std=init_proj_std) + if self.qkv.bias is not None: + nn.init.zeros_(self.qkv.bias) + if self.proj.bias is not None: + nn.init.zeros_(self.proj.bias) + + def forward(self, x: Tensor, is_causal: bool = True) -> Tensor: + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) + q, k, v = torch.unbind(qkv, 2) + q, k, v = [t.transpose(1, 2) for t in [q, k, v]] + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=None, dropout_p=self.attn_drop if self.training else 0, is_causal=is_causal + ) + x = x.transpose(1, 2).contiguous().view(B, N, C) + x = self.proj_drop(self.proj(x)) + return x diff --git a/engine/backbone/dinov3/layers/block.py b/engine/backbone/dinov3/layers/block.py new file mode 100644 index 0000000000000000000000000000000000000000..b29f93f0378e23d1542e464184b219871444d28b --- /dev/null +++ b/engine/backbone/dinov3/layers/block.py @@ -0,0 +1,269 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +from typing import Callable, List, Optional + +import torch +from torch import Tensor, nn + +from ..utils import cat_keep_shapes, uncat_with_shapes + +from .attention import CausalSelfAttention, SelfAttention +from .ffn_layers import Mlp +from .layer_scale import LayerScale # , DropPath + +torch._dynamo.config.automatic_dynamic_shapes = False +torch._dynamo.config.accumulated_cache_size_limit = 1024 + + +class SelfAttentionBlock(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + ffn_ratio: float = 4.0, + qkv_bias: bool = False, + proj_bias: bool = True, + ffn_bias: bool = True, + drop: float = 0.0, + attn_drop: float = 0.0, + init_values=None, + drop_path: float = 0.0, + act_layer: Callable[..., nn.Module] = nn.GELU, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + attn_class: Callable[..., nn.Module] = SelfAttention, + ffn_layer: Callable[..., nn.Module] = Mlp, + mask_k_bias: bool = False, + device=None, + ) -> None: + super().__init__() + # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}") + self.norm1 = norm_layer(dim) + self.attn = attn_class( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + attn_drop=attn_drop, + proj_drop=drop, + mask_k_bias=mask_k_bias, + device=device, + ) + self.ls1 = LayerScale(dim, init_values=init_values, device=device) if init_values else nn.Identity() + + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * ffn_ratio) + self.mlp = ffn_layer( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + bias=ffn_bias, + device=device, + ) + self.ls2 = LayerScale(dim, init_values=init_values, device=device) if init_values else nn.Identity() + + self.sample_drop_ratio = drop_path + + @staticmethod + def _maybe_index_rope(rope: tuple[Tensor, Tensor] | None, indices: Tensor) -> tuple[Tensor, Tensor] | None: + if rope is None: + return None + + sin, cos = rope + assert sin.ndim == cos.ndim + if sin.ndim == 4: + # If the rope embedding has a batch dimension (is different for each batch element), index into it + return sin[indices], cos[indices] # [batch, heads, patches, embed_dim] + else: + # No batch dimension, do not index + return sin, cos # [heads, patches, embed_dim] or [patches, embed_dim] + + def _forward(self, x: Tensor, rope=None) -> Tensor: + """ + This is the reference implementation for a single tensor, matching what is done below for a list. + We call the list op on [x] instead of this function. + """ + b, _, _ = x.shape + sample_subset_size = max(int(b * (1 - self.sample_drop_ratio)), 1) + residual_scale_factor = b / sample_subset_size + + if self.training and self.sample_drop_ratio > 0.0: + indices_1 = (torch.randperm(b, device=x.device))[:sample_subset_size] + + x_subset_1 = x[indices_1] + rope_subset = self._maybe_index_rope(rope, indices_1) + residual_1 = self.attn(self.norm1(x_subset_1), rope=rope_subset) + + x_attn = torch.index_add( + x, + dim=0, + source=self.ls1(residual_1), + index=indices_1, + alpha=residual_scale_factor, + ) + + indices_2 = (torch.randperm(b, device=x.device))[:sample_subset_size] + + x_subset_2 = x_attn[indices_2] + residual_2 = self.mlp(self.norm2(x_subset_2)) + + x_ffn = torch.index_add( + x_attn, + dim=0, + source=self.ls2(residual_2), + index=indices_2, + alpha=residual_scale_factor, + ) + else: + x_attn = x + self.ls1(self.attn(self.norm1(x), rope=rope)) + x_ffn = x_attn + self.ls2(self.mlp(self.norm2(x_attn))) + + return x_ffn + + def _forward_list(self, x_list: List[Tensor], rope_list=None) -> List[Tensor]: + """ + This list operator concatenates the tokens from the list of inputs together to save + on the elementwise operations. Torch-compile memory-planning allows hiding the overhead + related to concat ops. + """ + b_list = [x.shape[0] for x in x_list] + sample_subset_sizes = [max(int(b * (1 - self.sample_drop_ratio)), 1) for b in b_list] + residual_scale_factors = [b / sample_subset_size for b, sample_subset_size in zip(b_list, sample_subset_sizes)] + + if self.training and self.sample_drop_ratio > 0.0: + indices_1_list = [ + (torch.randperm(b, device=x.device))[:sample_subset_size] + for x, b, sample_subset_size in zip(x_list, b_list, sample_subset_sizes) + ] + x_subset_1_list = [x[indices_1] for x, indices_1 in zip(x_list, indices_1_list)] + + if rope_list is not None: + rope_subset_list = [ + self._maybe_index_rope(rope, indices_1) for rope, indices_1 in zip(rope_list, indices_1_list) + ] + else: + rope_subset_list = rope_list + + flattened, shapes, num_tokens = cat_keep_shapes(x_subset_1_list) + norm1 = uncat_with_shapes(self.norm1(flattened), shapes, num_tokens) + residual_1_list = self.attn.forward_list(norm1, rope_list=rope_subset_list) + + x_attn_list = [ + torch.index_add( + x, + dim=0, + source=self.ls1(residual_1), + index=indices_1, + alpha=residual_scale_factor, + ) + for x, residual_1, indices_1, residual_scale_factor in zip( + x_list, residual_1_list, indices_1_list, residual_scale_factors + ) + ] + + indices_2_list = [ + (torch.randperm(b, device=x.device))[:sample_subset_size] + for x, b, sample_subset_size in zip(x_list, b_list, sample_subset_sizes) + ] + x_subset_2_list = [x[indices_2] for x, indices_2 in zip(x_attn_list, indices_2_list)] + flattened, shapes, num_tokens = cat_keep_shapes(x_subset_2_list) + norm2_flat = self.norm2(flattened) + norm2_list = uncat_with_shapes(norm2_flat, shapes, num_tokens) + + residual_2_list = self.mlp.forward_list(norm2_list) + + x_ffn = [ + torch.index_add( + x_attn, + dim=0, + source=self.ls2(residual_2), + index=indices_2, + alpha=residual_scale_factor, + ) + for x_attn, residual_2, indices_2, residual_scale_factor in zip( + x_attn_list, residual_2_list, indices_2_list, residual_scale_factors + ) + ] + else: + x_out = [] + for x, rope in zip(x_list, rope_list): + x_attn = x + self.ls1(self.attn(self.norm1(x), rope=rope)) + x_ffn = x_attn + self.ls2(self.mlp(self.norm2(x_attn))) + x_out.append(x_ffn) + x_ffn = x_out + + return x_ffn + + def forward(self, x_or_x_list, rope_or_rope_list=None) -> List[Tensor]: + if isinstance(x_or_x_list, Tensor): + # for reference: + # return self._forward(x_or_x_list, rope=rope_or_rope_list) + # in order to match implementations we call the list op: + return self._forward_list([x_or_x_list], rope_list=[rope_or_rope_list])[0] + elif isinstance(x_or_x_list, list): + if rope_or_rope_list is None: + rope_or_rope_list = [None for x in x_or_x_list] + # return [self._forward(x, rope=rope) for x, rope in zip(x_or_x_list, rope_or_rope_list)] + return self._forward_list(x_or_x_list, rope_list=rope_or_rope_list) + else: + raise AssertionError + + +class CausalSelfAttentionBlock(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + ffn_ratio: float = 4.0, + ls_init_value: Optional[float] = None, + is_causal: bool = True, + act_layer: Callable = nn.GELU, + norm_layer: Callable = nn.LayerNorm, + dropout_prob: float = 0.0, + ): + super().__init__() + + self.dim = dim + self.is_causal = is_causal + self.ls1 = LayerScale(dim, init_values=ls_init_value) if ls_init_value else nn.Identity() + self.attention_norm = norm_layer(dim) + self.attention = CausalSelfAttention(dim, num_heads, attn_drop=dropout_prob, proj_drop=dropout_prob) + + self.ffn_norm = norm_layer(dim) + ffn_hidden_dim = int(dim * ffn_ratio) + self.feed_forward = Mlp( + in_features=dim, + hidden_features=ffn_hidden_dim, + drop=dropout_prob, + act_layer=act_layer, + ) + + self.ls2 = LayerScale(dim, init_values=ls_init_value) if ls_init_value else nn.Identity() + + def init_weights( + self, + init_attn_std: float | None = None, + init_proj_std: float | None = None, + init_fc_std: float | None = None, + factor: float = 1.0, + ) -> None: + init_attn_std = init_attn_std or (self.dim**-0.5) + init_proj_std = init_proj_std or init_attn_std * factor + init_fc_std = init_fc_std or (2 * self.dim) ** -0.5 + self.attention.init_weights(init_attn_std, init_proj_std) + self.attention_norm.reset_parameters() + nn.init.normal_(self.feed_forward.fc1.weight, std=init_fc_std) + nn.init.normal_(self.feed_forward.fc2.weight, std=init_proj_std) + self.ffn_norm.reset_parameters() + + def forward( + self, + x: torch.Tensor, + ): + + x_attn = x + self.ls1(self.attention(self.attention_norm(x), self.is_causal)) + x_ffn = x_attn + self.ls2(self.feed_forward(self.ffn_norm(x_attn))) + return x_ffn diff --git a/engine/backbone/dinov3/layers/dino_head.py b/engine/backbone/dinov3/layers/dino_head.py new file mode 100644 index 0000000000000000000000000000000000000000..bb71f35fc7ecf15e31963eb76d21626ccdea9b90 --- /dev/null +++ b/engine/backbone/dinov3/layers/dino_head.py @@ -0,0 +1,67 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import torch +import torch.nn as nn +from torch.nn.init import trunc_normal_ + + +class DINOHead(nn.Module): + def __init__( + self, + in_dim, + out_dim, + use_bn=False, + nlayers=3, + hidden_dim=2048, + bottleneck_dim=256, + mlp_bias=True, + ): + super().__init__() + nlayers = max(nlayers, 1) + self.mlp = _build_mlp( + nlayers, + in_dim, + bottleneck_dim, + hidden_dim=hidden_dim, + use_bn=use_bn, + bias=mlp_bias, + ) + self.last_layer = nn.Linear(bottleneck_dim, out_dim, bias=False) + + def init_weights(self) -> None: + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x, no_last_layer=False, only_last_layer=False): + if not only_last_layer: + x = self.mlp(x) + eps = 1e-6 if x.dtype == torch.float16 else 1e-12 + x = nn.functional.normalize(x, dim=-1, p=2, eps=eps) + if not no_last_layer: + x = self.last_layer(x) + return x + + +def _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=None, use_bn=False, bias=True): + if nlayers == 1: + return nn.Linear(in_dim, bottleneck_dim, bias=bias) + else: + layers = [nn.Linear(in_dim, hidden_dim, bias=bias)] + if use_bn: + layers.append(nn.BatchNorm1d(hidden_dim)) + layers.append(nn.GELU()) + for _ in range(nlayers - 2): + layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias)) + if use_bn: + layers.append(nn.BatchNorm1d(hidden_dim)) + layers.append(nn.GELU()) + layers.append(nn.Linear(hidden_dim, bottleneck_dim, bias=bias)) + return nn.Sequential(*layers) diff --git a/engine/backbone/dinov3/layers/ffn_layers.py b/engine/backbone/dinov3/layers/ffn_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..cd533049c13b7644cc71bae4c63dc19ee30cbcc7 --- /dev/null +++ b/engine/backbone/dinov3/layers/ffn_layers.py @@ -0,0 +1,77 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +from typing import Callable, List, Optional + +import torch.nn.functional as F +from torch import Tensor, nn + +from ..utils import cat_keep_shapes, uncat_with_shapes + + +class ListForwardMixin(object): + def forward(self, x: Tensor): + raise NotImplementedError + + def forward_list(self, x_list: List[Tensor]) -> List[Tensor]: + x_flat, shapes, num_tokens = cat_keep_shapes(x_list) + x_flat = self.forward(x_flat) + return uncat_with_shapes(x_flat, shapes, num_tokens) + + +class Mlp(nn.Module, ListForwardMixin): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = nn.GELU, + drop: float = 0.0, + bias: bool = True, + device=None, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias, device=device) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias, device=device) + self.drop = nn.Dropout(drop) + + def forward(self, x: Tensor) -> Tensor: + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class SwiGLUFFN(nn.Module, ListForwardMixin): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Optional[Callable[..., nn.Module]] = None, + drop: float = 0.0, + bias: bool = True, + align_to: int = 8, + device=None, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + d = int(hidden_features * 2 / 3) + swiglu_hidden_features = d + (-d % align_to) + self.w1 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device) + self.w2 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device) + self.w3 = nn.Linear(swiglu_hidden_features, out_features, bias=bias, device=device) + + def forward(self, x: Tensor) -> Tensor: + x1 = self.w1(x) + x2 = self.w2(x) + hidden = F.silu(x1) * x2 + return self.w3(hidden) diff --git a/engine/backbone/dinov3/layers/fp8_linear.py b/engine/backbone/dinov3/layers/fp8_linear.py new file mode 100644 index 0000000000000000000000000000000000000000..b54db7aa7f9fd45aec5d064346ec768b52509db9 --- /dev/null +++ b/engine/backbone/dinov3/layers/fp8_linear.py @@ -0,0 +1,141 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import re + +import torch + +from ..layers.attention import LinearKMaskedBias +from ..utils import named_replace + +# avoid division by zero when calculating scale +EPS = 1e-12 + + +def scale(t, amax_t): + max_v = torch.finfo(torch.float8_e4m3fn).max + scale_t = torch.clamp(amax_t.float(), min=EPS) / max_v + t_fp8 = (t / scale_t).to(torch.float8_e4m3fn) + return t_fp8, scale_t + + +def matmul(first, amax_first, second_t, amax_second_t, bias): + first_fp8, scale_first = scale(first, amax_first) + second_t_fp8, scale_second_t = scale(second_t, amax_second_t) + # PyTorch's row-wise scaled matmul kernel is based on CUTLASS and is quite + # slow. Hence we fall back to an "unscaled" matmul, which uses cuBLAS, and + # apply the scale manually afterwards. + output = torch._scaled_mm( + first_fp8, + second_t_fp8.t(), + scale_a=scale_first.new_ones((1, 1)), + scale_b=scale_second_t.t().new_ones((1, 1)), + bias=None, + out_dtype=torch.bfloat16, + use_fast_accum=False, + ) + output = (output * scale_first * scale_second_t.t()).to(torch.bfloat16) + if bias is not None: + output = output + bias + return output + + +@torch.compiler.allow_in_graph +class Fp8LinearFn(torch.autograd.Function): + @staticmethod + def forward(ctx, a, b_t, bias): + amax_a = a.abs().amax(dim=-1, keepdim=True) + amax_b_t = b_t.abs().amax(dim=-1, keepdim=True) + out = matmul(a, amax_a, b_t, amax_b_t, bias) + + ctx.a_requires_grad = a.requires_grad + ctx.b_requires_grad = b_t.requires_grad + ctx.bias_requires_grad = bias.requires_grad if bias is not None else False + + ctx.save_for_backward(a, b_t, amax_b_t.max()) + + return out + + @staticmethod + def backward(ctx, grad_out): + a, b_t, amax_b = ctx.saved_tensors + + if ctx.a_requires_grad: + b = b_t.t().contiguous() + amax_grad_out = grad_out.abs().amax(dim=-1, keepdim=True) + amax_b = amax_b.repeat(b.shape[0], 1) + grad_a = matmul(grad_out, amax_grad_out, b, amax_b, None) + else: + grad_a = None + if ctx.b_requires_grad: + grad_b = grad_out.t() @ a + else: + grad_b = None + if ctx.bias_requires_grad: + grad_bias = grad_out.sum(dim=0) + else: + grad_bias = None + + return grad_a, grad_b, grad_bias + + +class Fp8Linear(torch.nn.Linear): + def forward(self, input: torch.Tensor) -> torch.Tensor: + out = Fp8LinearFn.apply(input.flatten(end_dim=-2), self.weight, self.bias) + out = out.unflatten(0, input.shape[:-1]) + return out + + +class Fp8LinearKMaskedBias(LinearKMaskedBias): + def forward(self, input: torch.Tensor) -> torch.Tensor: + masked_bias = self.bias * self.bias_mask if self.bias is not None else None + out = Fp8LinearFn.apply(input.flatten(end_dim=-2), self.weight, masked_bias) + out = out.unflatten(0, input.shape[:-1]) + return out + + +def convert_linears_to_fp8(root_module: torch.nn.Module, *, filter: str) -> torch.nn.Module: + filter_re = re.compile(filter) + total_count = 0 + + def replace(module: torch.nn.Module, name: str) -> torch.nn.Module: + nonlocal total_count + if not isinstance(module, torch.nn.Linear) or not filter_re.search(name): + return module + if type(module) == torch.nn.Linear: + new_cls = Fp8Linear + elif type(module) == LinearKMaskedBias: + new_cls = Fp8LinearKMaskedBias + else: + assert False, str(type(module)) + if module.in_features % 64 != 0 or module.out_features % 64 != 0: + # This is not a strict requirement, but H100 TensorCores for fp8 + # operate on tiles of 64 elements anyways, and Inductor sometimes + # pads inner dims to become multiples of 64. Also, if one day we + # switch back to cuBLAS, it artificially requires dims to be + # multiples of 16. + raise RuntimeError( + "fp8 requires all dimensions to be multiples of 64 " "(consider using ffn_layer=swiglu64 or higher)" + ) + new_module = new_cls( + in_features=module.in_features, + out_features=module.out_features, + bias=module.bias is not None, + dtype=module.weight.dtype, + device=module.weight.device, + ) + new_module.weight = module.weight + new_module.bias = module.bias + total_count += 1 + return new_module + + out = named_replace(replace, root_module) + assert total_count > 0, "fp8: no layer found to convert" + # Force re-compile everything + torch._dynamo.reset_code_caches() + from torch._inductor.cudagraph_trees import reset_cudagraph_trees + + reset_cudagraph_trees() + return out diff --git a/engine/backbone/dinov3/layers/layer_scale.py b/engine/backbone/dinov3/layers/layer_scale.py new file mode 100644 index 0000000000000000000000000000000000000000..0b72b7c64c9cc38fd4e3db63e9c90f0158caa36c --- /dev/null +++ b/engine/backbone/dinov3/layers/layer_scale.py @@ -0,0 +1,29 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +from typing import Union + +import torch +from torch import Tensor, nn + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: Union[float, Tensor] = 1e-5, + inplace: bool = False, + device=None, + ) -> None: + super().__init__() + self.inplace = inplace + self.gamma = nn.Parameter(torch.empty(dim, device=device)) + self.init_values = init_values + + def reset_parameters(self): + nn.init.constant_(self.gamma, self.init_values) + + def forward(self, x: Tensor) -> Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma diff --git a/engine/backbone/dinov3/layers/patch_embed.py b/engine/backbone/dinov3/layers/patch_embed.py new file mode 100644 index 0000000000000000000000000000000000000000..760343f14cd0c1c2bbb2c70d43f82eb0bb1fddf4 --- /dev/null +++ b/engine/backbone/dinov3/layers/patch_embed.py @@ -0,0 +1,89 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import math +from typing import Callable, Tuple, Union + +from torch import Tensor, nn + + +def make_2tuple(x): + if isinstance(x, tuple): + assert len(x) == 2 + return x + + assert isinstance(x, int) + return (x, x) + + +class PatchEmbed(nn.Module): + """ + 2D image to patch embedding: (B,C,H,W) -> (B,N,D) + + Args: + img_size: Image size. + patch_size: Patch token size. + in_chans: Number of input image channels. + embed_dim: Number of linear projection output channels. + norm_layer: Normalization layer. + """ + + def __init__( + self, + img_size: Union[int, Tuple[int, int]] = 224, + patch_size: Union[int, Tuple[int, int]] = 16, + in_chans: int = 3, + embed_dim: int = 768, + norm_layer: Callable | None = None, + flatten_embedding: bool = True, + ) -> None: + super().__init__() + + image_HW = make_2tuple(img_size) + patch_HW = make_2tuple(patch_size) + patch_grid_size = ( + image_HW[0] // patch_HW[0], + image_HW[1] // patch_HW[1], + ) + + self.img_size = image_HW + self.patch_size = patch_HW + self.patches_resolution = patch_grid_size + self.num_patches = patch_grid_size[0] * patch_grid_size[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.flatten_embedding = flatten_embedding + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x: Tensor) -> Tensor: + _, _, H, W = x.shape + # patch_H, patch_W = self.patch_size + # assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}" + # assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}" + + x = self.proj(x) # B C H W + H, W = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) # B HW C + x = self.norm(x) + if not self.flatten_embedding: + x = x.reshape(-1, H, W, self.embed_dim) # B H W C + return x + + def flops(self) -> float: + Ho, Wo = self.patches_resolution + flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) + if self.norm is not None: + flops += Ho * Wo * self.embed_dim + return flops + + def reset_parameters(self): + k = 1 / (self.in_chans * (self.patch_size[0] ** 2)) + nn.init.uniform_(self.proj.weight, -math.sqrt(k), math.sqrt(k)) + if self.proj.bias is not None: + nn.init.uniform_(self.proj.bias, -math.sqrt(k), math.sqrt(k)) diff --git a/engine/backbone/dinov3/layers/rms_norm.py b/engine/backbone/dinov3/layers/rms_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..1d0a89c47c5e71687cadbf47fef567b2c6a2b3b4 --- /dev/null +++ b/engine/backbone/dinov3/layers/rms_norm.py @@ -0,0 +1,24 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import torch +from torch import Tensor, nn + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def reset_parameters(self) -> None: + nn.init.constant_(self.weight, 1) + + def _norm(self, x: Tensor) -> Tensor: + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x: Tensor) -> Tensor: + output = self._norm(x.float()).type_as(x) + return output * self.weight diff --git a/engine/backbone/dinov3/layers/rope_position_encoding.py b/engine/backbone/dinov3/layers/rope_position_encoding.py new file mode 100644 index 0000000000000000000000000000000000000000..2635d09e7732fb4c146f57d3aef19c2d3e5668ec --- /dev/null +++ b/engine/backbone/dinov3/layers/rope_position_encoding.py @@ -0,0 +1,121 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import math +from typing import Literal + +import numpy as np +import torch +from torch import Tensor, nn + + +# RoPE positional embedding with no mixing of coordinates (axial) and no learnable weights +# Supports two parametrizations of the rope parameters: either using `base` or `min_period` and `max_period`. +class RopePositionEmbedding(nn.Module): + def __init__( + self, + embed_dim: int, + *, + num_heads: int, + base: float | None = 100.0, + min_period: float | None = None, + max_period: float | None = None, + normalize_coords: Literal["min", "max", "separate"] = "separate", + shift_coords: float | None = None, + jitter_coords: float | None = None, + rescale_coords: float | None = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + ): + super().__init__() + assert embed_dim % (4 * num_heads) == 0 + both_periods = min_period is not None and max_period is not None + if (base is None and not both_periods) or (base is not None and both_periods): + raise ValueError("Either `base` or `min_period`+`max_period` must be provided.") + + D_head = embed_dim // num_heads + self.base = base + self.min_period = min_period + self.max_period = max_period + self.D_head = D_head + self.normalize_coords = normalize_coords + self.shift_coords = shift_coords + self.jitter_coords = jitter_coords + self.rescale_coords = rescale_coords + + # Needs persistent=True because we do teacher.load_state_dict(student.state_dict()) to initialize the teacher + self.dtype = dtype # Don't rely on self.periods.dtype + self.register_buffer( + "periods", + torch.empty(D_head // 4, device=device, dtype=dtype), + persistent=True, + ) + self._init_weights() + + def forward(self, *, H: int, W: int) -> tuple[Tensor, Tensor]: + device = self.periods.device + dtype = self.dtype + dd = {"device": device, "dtype": dtype} + + # Prepare coords in range [-1, +1] + if self.normalize_coords == "max": + max_HW = max(H, W) + coords_h = torch.arange(0.5, H, **dd) / max_HW # [H] + coords_w = torch.arange(0.5, W, **dd) / max_HW # [W] + elif self.normalize_coords == "min": + min_HW = min(H, W) + coords_h = torch.arange(0.5, H, **dd) / min_HW # [H] + coords_w = torch.arange(0.5, W, **dd) / min_HW # [W] + elif self.normalize_coords == "separate": + coords_h = torch.arange(0.5, H, **dd) / H # [H] + coords_w = torch.arange(0.5, W, **dd) / W # [W] + else: + raise ValueError(f"Unknown normalize_coords: {self.normalize_coords}") + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1) # [H, W, 2] + coords = coords.flatten(0, 1) # [HW, 2] + coords = 2.0 * coords - 1.0 # Shift range [0, 1] to [-1, +1] + + # Shift coords by adding a uniform value in [-shift, shift] + if self.training and self.shift_coords is not None: + shift_hw = torch.empty(2, **dd).uniform_(-self.shift_coords, self.shift_coords) + coords += shift_hw[None, :] + + # Jitter coords by multiplying the range [-1, 1] by a log-uniform value in [1/jitter, jitter] + if self.training and self.jitter_coords is not None: + jitter_max = np.log(self.jitter_coords) + jitter_min = -jitter_max + jitter_hw = torch.empty(2, **dd).uniform_(jitter_min, jitter_max).exp() + coords *= jitter_hw[None, :] + + # Rescale coords by multiplying the range [-1, 1] by a log-uniform value in [1/rescale, rescale] + if self.training and self.rescale_coords is not None: + rescale_max = np.log(self.rescale_coords) + rescale_min = -rescale_max + rescale_hw = torch.empty(1, **dd).uniform_(rescale_min, rescale_max).exp() + coords *= rescale_hw + + # Prepare angles and sin/cos + angles = 2 * math.pi * coords[:, :, None] / self.periods[None, None, :] # [HW, 2, D//4] + angles = angles.flatten(1, 2) # [HW, D//2] + angles = angles.tile(2) # [HW, D] + cos = torch.cos(angles) # [HW, D] + sin = torch.sin(angles) # [HW, D] + + return (sin, cos) # 2 * [HW, D] + + def _init_weights(self): + device = self.periods.device + dtype = self.dtype + if self.base is not None: + periods = self.base ** ( + 2 * torch.arange(self.D_head // 4, device=device, dtype=dtype) / (self.D_head // 2) + ) # [D//4] + else: + base = self.max_period / self.min_period + exponents = torch.linspace(0, 1, self.D_head // 4, device=device, dtype=dtype) # [D//4] range [0, 1] + periods = base**exponents # range [1, max_period / min_period] + periods = periods / base # range [min_period / max_period, 1] + periods = periods * self.max_period # range [min_period, max_period] + self.periods.data = periods diff --git a/engine/backbone/dinov3/layers/sparse_linear.py b/engine/backbone/dinov3/layers/sparse_linear.py new file mode 100644 index 0000000000000000000000000000000000000000..46edbb6673d56148a1f3317cbbf8d041b0c5908b --- /dev/null +++ b/engine/backbone/dinov3/layers/sparse_linear.py @@ -0,0 +1,90 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import logging +from typing import Callable + +import torch +import torch.nn as nn +import torch.nn.functional as F +import xformers.ops as xops + +from ..utils import named_apply, named_replace + +logger = logging.getLogger("dinov3") + + +class LinearW24(torch.nn.Linear): + ALGO = "largest_abs_values_greedy" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.sparsity_enabled = False + + def forward(self, input: torch.Tensor) -> torch.Tensor: + if not self.sparsity_enabled: + return super().forward(input) + + input_shape = input.shape + input = input.flatten(end_dim=-2) + dim0 = input.shape[0] + if dim0 % 8 != 0: + # NOTE: This should be torch-compiled away + input = F.pad(input, [0, 0, 0, -dim0 % 8]) + w_sparse = xops.sparsify24( + self.weight, + algo=self.ALGO, + gradient="ste", + backend="cusparselt", + ) + return F.linear(input, w_sparse, self.bias,)[ + :dim0 + ].unflatten(dim=0, sizes=input_shape[:-1]) + + +def replace_linears_with_sparse_linear(root_module: nn.Module, *, filter_fn: Callable[[str], bool]) -> nn.Module: + total_count = 0 + + def replace(module: nn.Module, name: str) -> nn.Module: + nonlocal total_count + if not isinstance(module, nn.Linear) or not filter_fn(name): + return module + assert type(module) == nn.Linear, "Subtypes not supported" + new_module = LinearW24( + in_features=module.in_features, + out_features=module.out_features, + bias=module.bias is not None, + dtype=module.weight.dtype, + device=module.weight.device, + ) + new_module.weight = module.weight + new_module.bias = module.bias + total_count += 1 + return new_module + + out = named_replace(replace, root_module) + assert total_count > 0, "2:4 sparsity: no layer found to sparsify" + return out + + +def update_24sparsity(root_module: nn.Module, enabled: bool) -> int: + num_modified = 0 + + def maybe_apply_sparsity(module: nn.Module, name: str) -> nn.Module: + nonlocal num_modified + if not isinstance(module, LinearW24): + return module + num_modified += 1 + module.sparsity_enabled = enabled + logger.info(f"- {'' if module.sparsity_enabled else 'de'}sparsifying {name}") + return module + + named_apply(maybe_apply_sparsity, root_module) + # Force re-compile everything + torch._dynamo.reset_code_caches() + from torch._inductor.cudagraph_trees import reset_cudagraph_trees + + reset_cudagraph_trees() + return num_modified diff --git a/engine/backbone/dinov3/utils/__init__.py b/engine/backbone/dinov3/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2794ac284e9bba24c6cee6a3eb5ecf7722f8734c --- /dev/null +++ b/engine/backbone/dinov3/utils/__init__.py @@ -0,0 +1,16 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +from .dtype import as_torch_dtype +from .utils import ( + cat_keep_shapes, + count_parameters, + fix_random_seeds, + get_conda_env, + get_sha, + named_apply, + named_replace, + uncat_with_shapes, +) diff --git a/engine/backbone/dinov3/utils/cluster.py b/engine/backbone/dinov3/utils/cluster.py new file mode 100644 index 0000000000000000000000000000000000000000..25df3bc556241b1efc22908c7c832d5f3751682f --- /dev/null +++ b/engine/backbone/dinov3/utils/cluster.py @@ -0,0 +1,103 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import os +from enum import Enum +from pathlib import Path +from typing import Any, Dict, Optional + + +class ClusterType(Enum): + CW = "cw" + + +def _guess_cluster_type() -> ClusterType: + return ClusterType.CW + + +def get_cluster_type( + cluster_type: Optional[ClusterType] = None, +) -> Optional[ClusterType]: + if cluster_type is None: + return _guess_cluster_type() + + return cluster_type + + +def get_slurm_account(cluster_type: Optional[ClusterType] = None) -> Optional[str]: + cluster_type = get_cluster_type(cluster_type) + if cluster_type is None: + return None + return { + ClusterType.CW: "fair_amaia_cw_explore", + }[cluster_type] + + +def get_checkpoint_path(cluster_type: Optional[ClusterType] = None) -> Optional[Path]: + cluster_type = get_cluster_type(cluster_type) + if cluster_type is None: + return None + + CHECKPOINT_DIRNAMES = { + ClusterType.CW: "", + } + return Path("/") / CHECKPOINT_DIRNAMES[cluster_type] + + +def get_user_checkpoint_path( + cluster_type: Optional[ClusterType] = None, +) -> Optional[Path]: + checkpoint_path = get_checkpoint_path(cluster_type) + if checkpoint_path is None: + return None + + username = os.environ.get("USER") + assert username is not None + return checkpoint_path / username + + +def get_slurm_qos(cluster_type: Optional[ClusterType] = None) -> Optional[str]: + cluster_type = get_cluster_type(cluster_type) + if cluster_type is None: + return None + + return { + ClusterType.CW: "explore", + }.get(cluster_type) + + +def get_slurm_partition(cluster_type: Optional[ClusterType] = None) -> Optional[str]: + cluster_type = get_cluster_type(cluster_type) + if cluster_type is None: + return None + + SLURM_PARTITIONS = { + ClusterType.CW: "learn", + } + return SLURM_PARTITIONS[cluster_type] + + +def get_slurm_executor_parameters( + nodes: int, + num_gpus_per_node: int, + cluster_type: Optional[ClusterType] = None, + **kwargs, +) -> Dict[str, Any]: + # create default parameters + params = { + "mem_gb": 0, # Requests all memory on a node, see https://slurm.schedmd.com/sbatch.html + "gpus_per_node": num_gpus_per_node, + "tasks_per_node": num_gpus_per_node, # one task per GPU + "cpus_per_task": 10, + "nodes": nodes, + "slurm_partition": get_slurm_partition(cluster_type), + } + # apply cluster-specific adjustments + cluster_type = get_cluster_type(cluster_type) + if cluster_type == ClusterType.CW: + params["cpus_per_task"] = 16 + # set additional parameters / apply overrides + params.update(kwargs) + return params diff --git a/engine/backbone/dinov3/utils/custom_callable.py b/engine/backbone/dinov3/utils/custom_callable.py new file mode 100644 index 0000000000000000000000000000000000000000..cb7c2f762835a6027f94006fad3360cf19ca4be3 --- /dev/null +++ b/engine/backbone/dinov3/utils/custom_callable.py @@ -0,0 +1,47 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import contextlib +import importlib +import inspect +import os +import sys +from pathlib import Path + + +@contextlib.contextmanager +def _load_modules_from_dir(dir_: str): + sys.path.insert(0, dir_) + yield + sys.path.pop(0) + + +def load_custom_callable(module_path: str | Path, callable_name: str): + module_full_path = os.path.realpath(module_path) + assert os.path.exists(module_full_path), f"module {module_full_path} does not exist" + module_dir, module_filename = os.path.split(module_full_path) + module_name, _ = os.path.splitext(module_filename) + + with _load_modules_from_dir(module_dir): + module = importlib.import_module(module_name) + if inspect.getfile(module) != module_full_path: + importlib.reload(module) + callable_ = getattr(module, callable_name) + + return callable_ + + +@contextlib.contextmanager +def change_working_dir_and_pythonpath(new_dir): + old_dir = Path.cwd() + new_dir = Path(new_dir).expanduser().resolve().as_posix() + old_pythonpath = sys.path.copy() + sys.path.insert(0, new_dir) + os.chdir(new_dir) + try: + yield + finally: + os.chdir(old_dir) + sys.path = old_pythonpath diff --git a/engine/backbone/dinov3/utils/dtype.py b/engine/backbone/dinov3/utils/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..b2795ed42b512ce51890c8592db5c364b18a5f4c --- /dev/null +++ b/engine/backbone/dinov3/utils/dtype.py @@ -0,0 +1,35 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +from typing import Dict, Union + +import numpy as np +import torch + +TypeSpec = Union[str, np.dtype, torch.dtype] + + +_NUMPY_TO_TORCH_DTYPE: Dict[np.dtype, torch.dtype] = { + np.dtype("bool"): torch.bool, + np.dtype("uint8"): torch.uint8, + np.dtype("int8"): torch.int8, + np.dtype("int16"): torch.int16, + np.dtype("int32"): torch.int32, + np.dtype("int64"): torch.int64, + np.dtype("float16"): torch.float16, + np.dtype("float32"): torch.float32, + np.dtype("float64"): torch.float64, + np.dtype("complex64"): torch.complex64, + np.dtype("complex128"): torch.complex128, +} + + +def as_torch_dtype(dtype: TypeSpec) -> torch.dtype: + if isinstance(dtype, torch.dtype): + return dtype + if isinstance(dtype, str): + dtype = np.dtype(dtype) + assert isinstance(dtype, np.dtype), f"Expected an instance of nunpy dtype, got {type(dtype)}" + return _NUMPY_TO_TORCH_DTYPE[dtype] diff --git a/engine/backbone/dinov3/utils/utils.py b/engine/backbone/dinov3/utils/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7b267033f49dfa0819406f48152f490b7c17ac94 --- /dev/null +++ b/engine/backbone/dinov3/utils/utils.py @@ -0,0 +1,130 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import logging +import os +import random +import subprocess +from typing import Callable, List, Optional, Tuple + +import numpy as np +import torch +from torch import Tensor, nn + +logger = logging.getLogger("dinov3") + + +def cat_keep_shapes(x_list: List[Tensor]) -> Tuple[Tensor, List[Tuple[int]], List[int]]: + shapes = [x.shape for x in x_list] + num_tokens = [x.select(dim=-1, index=0).numel() for x in x_list] + flattened = torch.cat([x.flatten(0, -2) for x in x_list]) + return flattened, shapes, num_tokens + + +def uncat_with_shapes(flattened: Tensor, shapes: List[Tuple[int]], num_tokens: List[int]) -> List[Tensor]: + outputs_splitted = torch.split_with_sizes(flattened, num_tokens, dim=0) + shapes_adjusted = [shape[:-1] + torch.Size([flattened.shape[-1]]) for shape in shapes] + outputs_reshaped = [o.reshape(shape) for o, shape in zip(outputs_splitted, shapes_adjusted)] + return outputs_reshaped + + +def named_replace( + fn: Callable, + module: nn.Module, + name: str = "", + depth_first: bool = True, + include_root: bool = False, +) -> nn.Module: + if not depth_first and include_root: + module = fn(module=module, name=name) + for child_name_o, child_module in list(module.named_children()): + child_name = ".".join((name, child_name_o)) if name else child_name_o + new_child = named_replace( + fn=fn, + module=child_module, + name=child_name, + depth_first=depth_first, + include_root=True, + ) + setattr(module, child_name_o, new_child) + + if depth_first and include_root: + module = fn(module=module, name=name) + return module + + +def named_apply( + fn: Callable, + module: nn.Module, + name: str = "", + depth_first: bool = True, + include_root: bool = False, +) -> nn.Module: + if not depth_first and include_root: + fn(module=module, name=name) + for child_name, child_module in module.named_children(): + child_name = ".".join((name, child_name)) if name else child_name + named_apply( + fn=fn, + module=child_module, + name=child_name, + depth_first=depth_first, + include_root=True, + ) + if depth_first and include_root: + fn(module=module, name=name) + return module + + +def fix_random_seeds(seed: int = 31): + """ + Fix random seeds. + """ + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed) + random.seed(seed) + + +def get_sha() -> str: + cwd = os.path.dirname(os.path.abspath(__file__)) + + def _run(command): + return subprocess.check_output(command, cwd=cwd).decode("ascii").strip() + + sha = "N/A" + diff = "clean" + branch = "N/A" + try: + sha = _run(["git", "rev-parse", "HEAD"]) + subprocess.check_output(["git", "diff"], cwd=cwd) + diff = _run(["git", "diff-index", "HEAD"]) + diff = "has uncommited changes" if diff else "clean" + branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"]) + except Exception: + pass + message = f"sha: {sha}, status: {diff}, branch: {branch}" + return message + + +def get_conda_env() -> Tuple[Optional[str], Optional[str]]: + conda_env_name = os.environ.get("CONDA_DEFAULT_ENV") + conda_env_path = os.environ.get("CONDA_PREFIX") + return conda_env_name, conda_env_path + + +def count_parameters(module: nn.Module) -> int: + c = 0 + for m in module.parameters(): + c += m.nelement() + return c + + +def has_batchnorms(model: nn.Module) -> bool: + bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm) + for _, module in model.named_modules(): + if isinstance(module, bn_types): + return True + return False diff --git a/engine/backbone/dinov3/vision_transformer.py b/engine/backbone/dinov3/vision_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..e525cedd6d460b9087f64b1695cf9ead5c7efa79 --- /dev/null +++ b/engine/backbone/dinov3/vision_transformer.py @@ -0,0 +1,392 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import os +import logging +from enum import Enum +from functools import partial +from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, Union + +import torch +import torch.nn.init +from torch import Tensor, nn + +from .layers import LayerScale, Mlp, PatchEmbed, RMSNorm, RopePositionEmbedding, SelfAttentionBlock, SwiGLUFFN +from .utils import named_apply + + +class Weights(Enum): + LVD1689M = "LVD1689M" + SAT493M = "SAT493M" + +configs = { + 'dinov3_vits16': { + 'img_size': 224, + 'patch_size': 16, + 'in_chans': 3, + 'pos_embed_rope_base': 100, + 'pos_embed_rope_normalize_coords': "separate", + 'pos_embed_rope_rescale_coords': 2, + 'pos_embed_rope_dtype': "fp32", + 'embed_dim': 384, + 'depth': 12, + 'num_heads': 6, + 'ffn_ratio': 4, + 'qkv_bias': True, + 'drop_path_rate': 0.0, + 'layerscale_init': 1.0e-05, + 'norm_layer': "layernormbf16", + 'ffn_layer': "mlp", + 'ffn_bias': True, + 'proj_bias': True, + 'n_storage_tokens': 4, + 'mask_k_bias': True, + 'pretrained': True, + 'weights': Weights.LVD1689M, + 'compact_arch_name': "vits", + 'check_hash': False, + }, + + 'dinov3_vits16plus': { + "img_size": 224, + "patch_size": 16, + "in_chans": 3, + "pos_embed_rope_base": 100, + "pos_embed_rope_normalize_coords": "separate", + "pos_embed_rope_rescale_coords": 2, + "pos_embed_rope_dtype": "fp32", + "embed_dim": 384, + "depth": 12, + "num_heads": 6, + "ffn_ratio": 6, + "qkv_bias": True, + "drop_path_rate": 0.0, + "layerscale_init": 1.0e-05, + "norm_layer": "layernormbf16", + "ffn_layer": "swiglu", + "ffn_bias": True, + "proj_bias": True, + "n_storage_tokens": 4, + "mask_k_bias": True, + "pretrained": True, + "weights": Weights.LVD1689M, + "compact_arch_name": "vitsplus", + "check_hash": False, + } +} + +logger = logging.getLogger("dinov3") + +ffn_layer_dict = { + "mlp": Mlp, + "swiglu": SwiGLUFFN, + "swiglu32": partial(SwiGLUFFN, align_to=32), + "swiglu64": partial(SwiGLUFFN, align_to=64), + "swiglu128": partial(SwiGLUFFN, align_to=128), +} + +norm_layer_dict = { + "layernorm": partial(nn.LayerNorm, eps=1e-6), + "layernormbf16": partial(nn.LayerNorm, eps=1e-5), + "rmsnorm": RMSNorm, +} + +dtype_dict = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + + +def init_weights_vit(module: nn.Module, name: str = ""): + if isinstance(module, nn.Linear): + torch.nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + if isinstance(module, nn.LayerNorm): + module.reset_parameters() + if isinstance(module, LayerScale): + module.reset_parameters() + if isinstance(module, PatchEmbed): + module.reset_parameters() + if isinstance(module, RMSNorm): + module.reset_parameters() + + +class DinoVisionTransformer(nn.Module): + def __init__( + self, + name, + ): + super().__init__() + + img_size = configs[name]['img_size'] + patch_size = configs[name]['patch_size'] + in_chans = configs[name]['in_chans'] + pos_embed_rope_min_period = None + pos_embed_rope_max_period = None + pos_embed_rope_shift_coords = None + pos_embed_rope_jitter_coords = None + pos_embed_rope_rescale_coords = None + pos_embed_rope_base = configs[name]['pos_embed_rope_base'] + pos_embed_rope_normalize_coords = configs[name]['pos_embed_rope_normalize_coords'] + pos_embed_rope_rescale_coords = configs[name]['pos_embed_rope_rescale_coords'] + pos_embed_rope_dtype = configs[name]['pos_embed_rope_dtype'] + embed_dim = configs[name]['embed_dim'] + depth = configs[name]['depth'] + num_heads = configs[name]['num_heads'] + ffn_ratio = configs[name]['ffn_ratio'] + qkv_bias = configs[name]['qkv_bias'] + drop_path_rate = configs[name]['drop_path_rate'] + layerscale_init = configs[name]['layerscale_init'] + norm_layer = configs[name]['norm_layer'] + ffn_layer = configs[name]['ffn_layer'] + ffn_bias = configs[name]['ffn_bias'] + proj_bias = configs[name]['proj_bias'] + n_storage_tokens = configs[name]['n_storage_tokens'] + mask_k_bias = configs[name]['mask_k_bias'] + pretrained = configs[name]['pretrained'] + weights = configs[name]['weights'] + compact_arch_name = configs[name]['compact_arch_name'] + check_hash = configs[name]['check_hash'] + untie_cls_and_patch_norms = False + untie_global_and_local_cls_norm = False + device = None + + norm_layer_cls = norm_layer_dict[norm_layer] + + self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models + self.n_blocks = depth + self.num_heads = num_heads + self.patch_size = patch_size + + self.patch_embed = PatchEmbed( + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + flatten_embedding=False, + ) + + self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim, device=device)) + self.n_storage_tokens = n_storage_tokens + if self.n_storage_tokens > 0: + self.storage_tokens = nn.Parameter(torch.empty(1, n_storage_tokens, embed_dim, device=device)) + logger.info(f"using base={pos_embed_rope_base} for rope new") + logger.info(f"using min_period={pos_embed_rope_min_period} for rope new") + logger.info(f"using max_period={pos_embed_rope_max_period} for rope new") + logger.info(f"using normalize_coords={pos_embed_rope_normalize_coords} for rope new") + logger.info(f"using shift_coords={pos_embed_rope_shift_coords} for rope new") + logger.info(f"using rescale_coords={pos_embed_rope_rescale_coords} for rope new") + logger.info(f"using jitter_coords={pos_embed_rope_jitter_coords} for rope new") + logger.info(f"using dtype={pos_embed_rope_dtype} for rope new") + self.rope_embed = RopePositionEmbedding( + embed_dim=embed_dim, + num_heads=num_heads, + base=pos_embed_rope_base, + min_period=pos_embed_rope_min_period, + max_period=pos_embed_rope_max_period, + normalize_coords=pos_embed_rope_normalize_coords, + shift_coords=pos_embed_rope_shift_coords, + jitter_coords=pos_embed_rope_jitter_coords, + rescale_coords=pos_embed_rope_rescale_coords, + dtype=dtype_dict[pos_embed_rope_dtype], + device=device, + ) + logger.info(f"using {ffn_layer} layer as FFN") + ffn_layer_cls = ffn_layer_dict[ffn_layer] + ffn_ratio_sequence = [ffn_ratio] * depth + blocks_list = [ + SelfAttentionBlock( + dim=embed_dim, + num_heads=num_heads, + ffn_ratio=ffn_ratio_sequence[i], + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ffn_bias=ffn_bias, + drop_path=drop_path_rate, + norm_layer=norm_layer_cls, + act_layer=nn.GELU, + ffn_layer=ffn_layer_cls, + init_values=layerscale_init, + mask_k_bias=mask_k_bias, + device=device, + ) + for i in range(depth) + ] + + self.chunked_blocks = False + self.blocks = nn.ModuleList(blocks_list) + + # This norm is applied to everything, or when untying, to patch and mask tokens. + self.norm = norm_layer_cls(embed_dim) + + self.untie_cls_and_patch_norms = untie_cls_and_patch_norms + if untie_cls_and_patch_norms: + # When untying, this norm is applied to CLS tokens and registers. + self.cls_norm = norm_layer_cls(embed_dim) + else: + self.cls_norm = None + + self.untie_global_and_local_cls_norm = untie_global_and_local_cls_norm + if untie_global_and_local_cls_norm: + # When untying, this norm is applied to local CLS tokens and registers. + # This norm is never used during eval. + self.local_cls_norm = norm_layer_cls(embed_dim) + else: + self.local_cls_norm = None + self.head = nn.Identity() + self.mask_token = nn.Parameter(torch.empty(1, embed_dim, device=device)) + + self.init_weights() + + def init_weights(self): + self.rope_embed._init_weights() + nn.init.normal_(self.cls_token, std=0.02) + if self.n_storage_tokens > 0: + nn.init.normal_(self.storage_tokens, std=0.02) + nn.init.zeros_(self.mask_token) + named_apply(init_weights_vit, self) + + def prepare_tokens_with_masks(self, x: Tensor, masks=None) -> Tuple[Tensor, Tuple[int]]: + x = self.patch_embed(x) + B, H, W, _ = x.shape + x = x.flatten(1, 2) + + if masks is not None: + x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x) + cls_token = self.cls_token + else: + cls_token = self.cls_token + 0 * self.mask_token + if self.n_storage_tokens > 0: + storage_tokens = self.storage_tokens + else: + storage_tokens = torch.empty( + 1, + 0, + cls_token.shape[-1], + dtype=cls_token.dtype, + device=cls_token.device, + ) + + x = torch.cat( + [ + cls_token.expand(B, -1, -1), + storage_tokens.expand(B, -1, -1), + x, + ], + dim=1, + ) + + return x, (H, W) + + def forward_features_list(self, x_list: List[Tensor], masks_list: List[Tensor]) -> List[Dict[str, Tensor]]: + x = [] + rope = [] + for t_x, t_masks in zip(x_list, masks_list): + t2_x, hw_tuple = self.prepare_tokens_with_masks(t_x, t_masks) + x.append(t2_x) + rope.append(hw_tuple) + for _, blk in enumerate(self.blocks): + if self.rope_embed is not None: + rope_sincos = [self.rope_embed(H=H, W=W) for H, W in rope] + else: + rope_sincos = [None for r in rope] + x = blk(x, rope_sincos) + all_x = x + output = [] + for idx, (x, masks) in enumerate(zip(all_x, masks_list)): + if self.untie_cls_and_patch_norms or self.untie_global_and_local_cls_norm: + if self.untie_global_and_local_cls_norm and self.training and idx == 1: + # Assume second entry of list corresponds to local crops. + # We only ever apply this during training. + x_norm_cls_reg = self.local_cls_norm(x[:, : self.n_storage_tokens + 1]) + elif self.untie_cls_and_patch_norms: + x_norm_cls_reg = self.cls_norm(x[:, : self.n_storage_tokens + 1]) + else: + x_norm_cls_reg = self.norm(x[:, : self.n_storage_tokens + 1]) + x_norm_patch = self.norm(x[:, self.n_storage_tokens + 1 :]) + else: + x_norm = self.norm(x) + x_norm_cls_reg = x_norm[:, : self.n_storage_tokens + 1] + x_norm_patch = x_norm[:, self.n_storage_tokens + 1 :] + output.append( + { + "x_norm_clstoken": x_norm_cls_reg[:, 0], + "x_storage_tokens": x_norm_cls_reg[:, 1:], + "x_norm_patchtokens": x_norm_patch, + "x_prenorm": x, + "masks": masks, + } + ) + return output + + def forward_features(self, x: Tensor | List[Tensor], masks: Optional[Tensor] = None) -> List[Dict[str, Tensor]]: + if isinstance(x, torch.Tensor): + return self.forward_features_list([x], [masks])[0] + else: + return self.forward_features_list(x, masks) + + def _get_intermediate_layers_not_chunked(self, x: Tensor, n: int = 1) -> List[Tensor]: + x, (H, W) = self.prepare_tokens_with_masks(x) + # If n is an int, take the n last blocks. If it's a list, take them + output, total_block_len = [], len(self.blocks) + blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n + for i, blk in enumerate(self.blocks): + if self.rope_embed is not None: + rope_sincos = self.rope_embed(H=H, W=W) + else: + rope_sincos = None + x = blk(x, rope_sincos) + if i in blocks_to_take: + output.append(x) + assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found" + return output + + def get_intermediate_layers( + self, + x: torch.Tensor, + *, + n: Union[int, Sequence] = 1, # Layers or n last layers to take + reshape: bool = False, + return_class_token: bool = False, + return_extra_tokens: bool = False, + norm: bool = True, + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor, ...]]]: + outputs = self._get_intermediate_layers_not_chunked(x, n) + if norm: + outputs_normed = [] + for out in outputs: + if self.untie_cls_and_patch_norms: + x_norm_cls_reg = self.cls_norm(out[:, : self.n_storage_tokens + 1]) + x_norm_patch = self.norm(out[:, self.n_storage_tokens + 1 :]) + outputs_normed.append(torch.cat((x_norm_cls_reg, x_norm_patch), dim=1)) + else: + outputs_normed.append(self.norm(out)) + outputs = outputs_normed + class_tokens = [out[:, 0] for out in outputs] + extra_tokens = [out[:, 1 : self.n_storage_tokens + 1] for out in outputs] + outputs = [out[:, self.n_storage_tokens + 1 :] for out in outputs] + if reshape: + B, _, h, w = x.shape + outputs = [ + out.reshape(B, h // self.patch_size, w // self.patch_size, -1).permute(0, 3, 1, 2).contiguous() + for out in outputs + ] + if not return_class_token and not return_extra_tokens: + return tuple(outputs) + elif return_class_token and not return_extra_tokens: + return tuple(zip(outputs, class_tokens)) + elif not return_class_token and return_extra_tokens: + return tuple(zip(outputs, extra_tokens)) + elif return_class_token and return_extra_tokens: + return tuple(zip(outputs, class_tokens, extra_tokens)) + + def forward(self, *args, is_training: bool = False, **kwargs) -> List[Dict[str, Tensor]] | Tensor: + ret = self.forward_features(*args, **kwargs) + if is_training: + return ret + else: + return self.head(ret["x_norm_clstoken"]) diff --git a/engine/backbone/dinov3_adapter.py b/engine/backbone/dinov3_adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..d81498ad0d8f97fe7e6f48a5ce9c8da2704d1db4 --- /dev/null +++ b/engine/backbone/dinov3_adapter.py @@ -0,0 +1,171 @@ +""" +DEIMv2: Real-Time Object Detection Meets DINOv3 +Copyright (c) 2025 The DEIMv2 Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from DINOv3 (https://github.com/facebookresearch/dinov3) + +Copyright (c) Meta Platforms, Inc. and affiliates. + +This software may be used and distributed in accordance with +the terms of the DINOv3 License Agreement. +""" + +import os + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp + +from functools import partial +from ..core import register +from .vit_tiny import VisionTransformer +from .dinov3 import DinoVisionTransformer + + +class SpatialPriorModulev2(nn.Module): + def __init__(self, inplanes=16): + super().__init__() + + # 1/4 + self.stem = nn.Sequential( + *[ + nn.Conv2d(3, inplanes, kernel_size=3, stride=2, padding=1, bias=False), + nn.SyncBatchNorm(inplanes), + nn.GELU(), + nn.MaxPool2d(kernel_size=3, stride=2, padding=1), + ] + ) + # 1/8 + self.conv2 = nn.Sequential( + *[ + nn.Conv2d(inplanes, 2 * inplanes, kernel_size=3, stride=2, padding=1, bias=False), + nn.SyncBatchNorm(2 * inplanes), + ] + ) + # 1/16 + self.conv3 = nn.Sequential( + *[ + nn.GELU(), + nn.Conv2d(2 * inplanes, 4 * inplanes, kernel_size=3, stride=2, padding=1, bias=False), + nn.SyncBatchNorm(4 * inplanes), + ] + ) + # 1/32 + self.conv4 = nn.Sequential( + *[ + nn.GELU(), + nn.Conv2d(4 * inplanes, 4 * inplanes, kernel_size=3, stride=2, padding=1, bias=False), + nn.SyncBatchNorm(4 * inplanes), + ] + ) + + def forward(self, x): + c1 = self.stem(x) + c2 = self.conv2(c1) # 1/8 + c3 = self.conv3(c2) # 1/16 + c4 = self.conv4(c3) # 1/32 + + return c2, c3, c4 + + +@register() +class DINOv3STAs(nn.Module): + def __init__( + self, + name=None, + weights_path=None, + interaction_indexes=[], + finetune=True, + embed_dim=192, + num_heads=3, + patch_size=16, + use_sta=True, + conv_inplane=16, + hidden_dim=None, + ): + super(DINOv3STAs, self).__init__() + if 'dinov3' in name: + self.dinov3 = DinoVisionTransformer(name=name) + if weights_path is not None and os.path.exists(weights_path): + print(f'Loading ckpt from {weights_path}...') + self.dinov3.load_state_dict(torch.load(weights_path)) + else: + print('Training DINOv3 from scratch...') + else: + self.dinov3 = VisionTransformer(embed_dim=embed_dim, num_heads=num_heads, return_layers=interaction_indexes) + if weights_path is not None and os.path.exists(weights_path): + print(f'Loading ckpt from {weights_path}...') + self.dinov3._model.load_state_dict(torch.load(weights_path)) + else: + print('Training ViT-Tiny from scratch...') + + embed_dim = self.dinov3.embed_dim + self.interaction_indexes = interaction_indexes + self.patch_size = patch_size + + if not finetune: + self.dinov3.eval() + self.dinov3.requires_grad_(False) + + # init the feature pyramid + self.use_sta = use_sta + if use_sta: + print(f"Using Lite Spatial Prior Module with inplanes={conv_inplane}") + self.sta = SpatialPriorModulev2(inplanes=conv_inplane) + else: + conv_inplane = 0 + + # linear projection + hidden_dim = hidden_dim if hidden_dim is not None else embed_dim + self.convs = nn.ModuleList([ + nn.Conv2d(embed_dim + conv_inplane*2, hidden_dim, kernel_size=1, stride=1, padding=0, bias=False), + nn.Conv2d(embed_dim + conv_inplane*4, hidden_dim, kernel_size=1, stride=1, padding=0, bias=False), + nn.Conv2d(embed_dim + conv_inplane*4, hidden_dim, kernel_size=1, stride=1, padding=0, bias=False) + ]) + # norm + self.norms = nn.ModuleList([ + nn.SyncBatchNorm(hidden_dim), + nn.SyncBatchNorm(hidden_dim), + nn.SyncBatchNorm(hidden_dim) + ]) + + def forward(self, x): + # Code for matching with oss + H_c, W_c = x.shape[2] // 16, x.shape[3] // 16 + H_toks, W_toks = x.shape[2] // self.patch_size, x.shape[3] // self.patch_size + bs, C, h, w = x.shape + + if len(self.interaction_indexes) > 0 and not isinstance(self.dinov3, VisionTransformer): + all_layers = self.dinov3.get_intermediate_layers( + x, n=self.interaction_indexes, return_class_token=True + ) + else: + all_layers = self.dinov3(x) + + if len(all_layers) == 1: # repeat the same layer for all the three scales + all_layers = [all_layers[0], all_layers[0], all_layers[0]] + + sem_feats = [] + num_scales = len(all_layers) - 2 + for i, sem_feat in enumerate(all_layers): + feat, _ = sem_feat + sem_feat = feat.transpose(1, 2).view(bs, -1, H_c, W_c).contiguous() # [B, D, H, W] + resize_H, resize_W = int(H_c * 2**(num_scales-i)), int(W_c * 2**(num_scales-i)) + sem_feat = F.interpolate(sem_feat, size=[resize_H, resize_W], mode="bilinear", align_corners=False) + sem_feats.append(sem_feat) + + # fusion + fused_feats = [] + if self.use_sta: + detail_feats = self.sta(x) + for sem_feat, detail_feat in zip(sem_feats, detail_feats): + fused_feats.append(torch.cat([sem_feat, detail_feat], dim=1)) + else: + fused_feats = sem_feats + + c2 = self.norms[0](self.convs[0](fused_feats[0])) + c3 = self.norms[1](self.convs[1](fused_feats[1])) + c4 = self.norms[2](self.convs[2](fused_feats[2])) + + return c2, c3, c4 \ No newline at end of file diff --git a/engine/backbone/hgnetv2.py b/engine/backbone/hgnetv2.py new file mode 100644 index 0000000000000000000000000000000000000000..e70fc6dd3114cc2a496416b53e9437c0c9640740 --- /dev/null +++ b/engine/backbone/hgnetv2.py @@ -0,0 +1,627 @@ +""" +DEIMv2: Real-Time Object Detection Meets DINOv3 +Copyright (c) 2025 The DEIMv2 Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINEr) + +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 torch +import torch.nn as nn +import torch.nn.functional as F +import os +from .common import FrozenBatchNorm2d +from ..core import register +import logging +from .common import get_activation + +# Constants for initialization +kaiming_normal_ = nn.init.kaiming_normal_ +zeros_ = nn.init.zeros_ +ones_ = nn.init.ones_ + +__all__ = ['HGNetv2'] + + +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, + act='relu', + ): + 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() + self.act = get_activation(act) + 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, + act='relu', + ): + super().__init__() + self.conv1 = ConvBNAct( + in_chs, + out_chs, + kernel_size=1, + use_act=False, + use_lab=use_lab, + act=act, + ) + self.conv2 = ConvBNAct( + out_chs, + out_chs, + kernel_size=kernel_size, + groups=out_chs, + use_act=True, + use_lab=use_lab, + act=act, + ) + + 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, act='relu'): + super().__init__() + self.stem1 = ConvBNAct( + in_chs, + mid_chs, + kernel_size=3, + stride=2, + use_lab=use_lab, + act=act, + ) + self.stem2a = ConvBNAct( + mid_chs, + mid_chs // 2, + kernel_size=2, + stride=1, + use_lab=use_lab, + act=act, + ) + self.stem2b = ConvBNAct( + mid_chs // 2, + mid_chs, + kernel_size=2, + stride=1, + use_lab=use_lab, + act=act, + ) + self.stem3 = ConvBNAct( + mid_chs * 2, + mid_chs, + kernel_size=3, + stride=2, + use_lab=use_lab, + act=act, + ) + self.stem4 = ConvBNAct( + mid_chs, + out_chs, + kernel_size=1, + stride=1, + use_lab=use_lab, + act=act, + ) + 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., + act='relu', + ): + 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, + act=act, + ) + ) + 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, + act=act, + ) + ) + + # 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, + act=act, + ) + aggregation_excitation_conv = ConvBNAct( + out_chs // 2, + out_chs, + kernel_size=1, + stride=1, + use_lab=use_lab, + act=act, + ) + 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, + act=act, + ) + 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., + act='relu', + ): + 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, + act=act, + ) + 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, + act=act, + ) + ) + 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 = { + 'Atto': { # only 3 stages + '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, 256, 1, True, True, 3, 3], + }, + 'url': 'https://github.com/Peterande/storage/releases/download/dfinev1.0/PPHGNetV2_B0_stage1.pth' + }, + 'Femto': { # only 3 stages + '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, 1, True, True, 5, 3], + }, + 'url': 'https://github.com/Peterande/storage/releases/download/dfinev1.0/PPHGNetV2_B0_stage1.pth' + }, + 'Pico': { # only 3 stages + '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], + }, + 'url': 'https://github.com/Peterande/storage/releases/download/dfinev1.0/PPHGNetV2_B0_stage1.pth' + }, + '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/', + act='relu', + ): + 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] + print(f" ### Backbone.act: {act} ### ") + print(f" ### Backbone.act: {act} ### ") + + # stem + self.stem = StemBlock( + in_chs=stem_channels[0], + mid_chs=stem_channels[1], + out_chs=stem_channels[2], + use_lab=use_lab, + act=act) + + # 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, + act=act) + ) + + 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 (name in ['Atto', 'Femto', 'Pico']): + model_path = local_model_dir + 'PPHGNetV2_' + 'B0' + '_stage1.pth' + else: + model_path = local_model_dir + 'PPHGNetV2_' + name + '_stage1.pth' + if os.path.exists(model_path): + state = torch.load(model_path, map_location='cpu') + print(f"Loaded stage1 {name} HGNetV2 from local file.") + else: + # If the file doesn't exist locally, download from the URL + if torch.distributed.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) + torch.distributed.barrier() + else: + torch.distributed.barrier() + state = torch.load(local_model_dir) + + print(f"Loaded stage1 {name} HGNetV2 from URL.") + + if ('Atto' == name): + self.load_partial_state_dict(self, state) + elif ('Femto' == name) or ('Pico' == name): + missing_keys, unexpected_keys = self.load_state_dict(state, strict=False) + print("Missing keys:", missing_keys) + print("Unexpected keys:", unexpected_keys) + else: + self.load_state_dict(state) + + except (Exception, KeyboardInterrupt) as e: + if torch.distributed.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() + + @staticmethod + def load_partial_state_dict(model, state_dict): + model_dict = model.state_dict() + # 只保留shape完全一致的参数 + filtered_dict = {k: v for k, v in state_dict.items() + if k in model_dict and v.shape == model_dict[k].shape} + + # 更新模型参数 + model_dict.update(filtered_dict) + model.load_state_dict(model_dict, strict=False) + missing = set(model_dict.keys()) - set(filtered_dict.keys()) + unexpected = set(state_dict.keys()) - set(filtered_dict.keys()) + print("Missing keys:", missing) + print(" #########################################################") + print("Unexpected keys:", unexpected) + + 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 diff --git a/engine/backbone/ms_deform_attn.py b/engine/backbone/ms_deform_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..08ab9317393d3e2d5938559270e02d9fca7454ef --- /dev/null +++ b/engine/backbone/ms_deform_attn.py @@ -0,0 +1,220 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This software may be used and distributed in accordance with +# the terms of the DINOv3 License Agreement. + +import math +import warnings + +import torch +import torch.nn.functional as F +from torch import nn +from torch.autograd import Function +# from torch.amp import custom_fwd, custom_bwd + +from torch.autograd.function import once_differentiable +from torch.nn.init import constant_, xavier_uniform_ + +try: + import MultiScaleDeformableAttention as MSDA +except ImportError: + # if we just care about inference, we don't need + # the compiled extension for multi-scale deformable attention + MSDA = None + + +class MSDeformAttnFunction(Function): + @staticmethod + # @custom_fwd(device_type="cuda", cast_inputs=torch.float32) + def forward( + ctx, value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, im2col_step + ): + ctx.im2col_step = im2col_step + output = ms_deform_attn_core_pytorch( + value, + value_spatial_shapes, + # value_level_start_index, + sampling_locations, + attention_weights, + ) + ctx.save_for_backward( + value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights + ) + return output + + @staticmethod + @once_differentiable + # @custom_bwd(device_type="cuda") + def backward(ctx, grad_output): + if MSDA is None: + raise RuntimeError( + "MultiScaleDeformableAttention is not available, " + "please compile with CUDA if you want to train a " + "segmentation head with deformable attention" + ) + value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights = ctx.saved_tensors + grad_value, grad_sampling_loc, grad_attn_weight = MSDA.ms_deform_attn_backward( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + grad_output, + ctx.im2col_step, + ) + + return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None + + +def ms_deform_attn_core_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights): + # for debug and test only, + # need to use cuda version instead + N_, S_, M_, D_ = value.shape + _, Lq_, M_, L_, P_, _ = sampling_locations.shape + value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1) + sampling_grids = 2 * sampling_locations - 1 + sampling_value_list = [] + for lid_, (H_, W_) in enumerate(value_spatial_shapes): + # N_, H_*W_, M_, D_ -> N_, H_*W_, M_*D_ -> N_, M_*D_, H_*W_ -> N_*M_, D_, H_, W_ + value_l_ = value_list[lid_].flatten(2).transpose(1, 2).reshape(N_ * M_, D_, H_, W_) + # N_, Lq_, M_, P_, 2 -> N_, M_, Lq_, P_, 2 -> N_*M_, Lq_, P_, 2 + sampling_grid_l_ = sampling_grids[:, :, :, lid_].transpose(1, 2).flatten(0, 1) + # N_*M_, D_, Lq_, P_ + sampling_value_l_ = F.grid_sample( + value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False + ) + sampling_value_list.append(sampling_value_l_) + # (N_, Lq_, M_, L_, P_) -> (N_, M_, Lq_, L_, P_) -> (N_, M_, 1, Lq_, L_*P_) + attention_weights = attention_weights.transpose(1, 2).reshape(N_ * M_, 1, Lq_, L_ * P_) + output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(N_, M_ * D_, Lq_) + return output.transpose(1, 2).contiguous() + + +def _is_power_of_2(n): + if (not isinstance(n, int)) or (n < 0): + raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n))) + return (n & (n - 1) == 0) and n != 0 + + +class MSDeformAttn(nn.Module): + def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4, ratio=1.0): + """Multi-Scale Deformable Attention Module. + + :param d_model hidden dimension + :param n_levels number of feature levels + :param n_heads number of attention heads + :param n_points number of sampling points per attention head per feature level + """ + super().__init__() + if d_model % n_heads != 0: + raise ValueError("d_model must be divisible by n_heads, but got {} and {}".format(d_model, n_heads)) + _d_per_head = d_model // n_heads + # you'd better set _d_per_head to a power of 2 + # which is more efficient in our CUDA implementation + if not _is_power_of_2(_d_per_head): + warnings.warn( + "You'd better set d_model in MSDeformAttn to make " + "the dimension of each attention head a power of 2 " + "which is more efficient in our CUDA implementation." + ) + + self.im2col_step = 64 + + self.d_model = d_model + self.n_levels = n_levels + self.n_heads = n_heads + self.n_points = n_points + self.ratio = ratio + self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2) + self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points) + self.value_proj = nn.Linear(d_model, int(d_model * ratio)) + self.output_proj = nn.Linear(int(d_model * ratio), d_model) + + self._reset_parameters() + + self.ms_deformable_attn_core = ms_deform_attn_core_pytorch + def _reset_parameters(self): + constant_(self.sampling_offsets.weight.data, 0.0) + thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads) + grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) + grid_init = ( + (grid_init / grid_init.abs().max(-1, keepdim=True)[0]) + .view(self.n_heads, 1, 1, 2) + .repeat(1, self.n_levels, self.n_points, 1) + ) + for i in range(self.n_points): + grid_init[:, :, i, :] *= i + 1 + + with torch.no_grad(): + self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1)) + constant_(self.attention_weights.weight.data, 0.0) + constant_(self.attention_weights.bias.data, 0.0) + xavier_uniform_(self.value_proj.weight.data) + constant_(self.value_proj.bias.data, 0.0) + xavier_uniform_(self.output_proj.weight.data) + constant_(self.output_proj.bias.data, 0.0) + + def forward( + self, + query, + reference_points, + input_flatten, + input_spatial_shapes, + input_level_start_index, + input_padding_mask=None, + ): + """ + :param query (N, Length_{query}, C) + :param reference_points (N, Length_{query}, n_levels, 2), range in [0, 1], top-left (0,0), bottom-right (1, 1), including padding area + or (N, Length_{query}, n_levels, 4), add additional (w, h) to form reference boxes + :param input_flatten (N, \\sum_{l=0}^{L-1} H_l \\cdot W_l, C) + :param input_spatial_shapes (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})] + :param input_level_start_index (n_levels, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}] + :param input_padding_mask (N, \\sum_{l=0}^{L-1} H_l \\cdot W_l), True for padding elements, False for non-padding elements + + :return output (N, Length_{query}, C) + """ + + N, Len_q, _ = query.shape + N, Len_in, _ = input_flatten.shape + assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in + + value = self.value_proj(input_flatten) + if input_padding_mask is not None: + value = value.masked_fill(input_padding_mask[..., None], float(0)) + + value = value.view(N, Len_in, self.n_heads, int(self.ratio * self.d_model) // self.n_heads) + sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2) + attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points) + attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points) + + if reference_points.shape[-1] == 2: + offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1) + sampling_locations = ( + reference_points[:, :, None, :, None, :] + + sampling_offsets / offset_normalizer[None, None, None, :, None, :] + ) + elif reference_points.shape[-1] == 4: + sampling_locations = ( + reference_points[:, :, None, :, None, :2] + + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5 + ) + else: + raise ValueError( + "Last dim of reference_points must be 2 or 4, but get {} instead.".format(reference_points.shape[-1]) + ) + + # output_cpp = MSDeformAttnFunction.apply( + # value, + # input_spatial_shapes, + # input_level_start_index, + # sampling_locations, + # attention_weights, + # self.im2col_step, + # ) + output = self.ms_deformable_attn_core(value, input_spatial_shapes, sampling_locations, attention_weights) + + # print("C++ version deformable attention", output_cpp.sum().item(), output_cpp.mean().item(), output_cpp.max().item(), output_cpp.min().item()) + # print("PyTorch version deformable attention", output.sum().item(), output.mean().item(), output.max().item(), output.min().item()) + output = self.output_proj(output) + return output diff --git a/engine/backbone/presnet.py b/engine/backbone/presnet.py new file mode 100644 index 0000000000000000000000000000000000000000..bfc486ca098105593e1f946e8a02e1ead98420f5 --- /dev/null +++ b/engine/backbone/presnet.py @@ -0,0 +1,259 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" +import torch +import torch.nn as nn +import torch.nn.functional as F + +from collections import OrderedDict + +from .common import get_activation, FrozenBatchNorm2d + +from ..core import register +import os + + +__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', +} + +local_weights = { + 18: "ResNet18_vd_pretrained_from_paddle.pth", + 34: "ResNet34_vd_pretrained_from_paddle.pth", + 50: "ResNet50_vd_ssld_v2_pretrained_from_paddle.pth", + 101: "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, + local_model_dir='weights/resnets', + ): + 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: + model_path = local_model_dir + local_weights[depth] + if os.path.exists(model_path): + state = torch.load(model_path, map_location='cpu') + print(f"Loaded PResNet{depth} from local file@{model_path}.") + else: + if isinstance(pretrained, bool) or 'http' in pretrained: + state = torch.hub.load_state_dict_from_url(donwload_url[depth], map_location='cpu', model_dir=local_model_dir) + 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 diff --git a/engine/backbone/test_resnet.py b/engine/backbone/test_resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..3804f380a20202845ed7e2930dc612510acc6575 --- /dev/null +++ b/engine/backbone/test_resnet.py @@ -0,0 +1,80 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from collections import OrderedDict + + +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) diff --git a/engine/backbone/timm_model.py b/engine/backbone/timm_model.py new file mode 100644 index 0000000000000000000000000000000000000000..5f5a97cc74a7399f8fdf8216d782fe332996570c --- /dev/null +++ b/engine/backbone/timm_model.py @@ -0,0 +1,69 @@ +"""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 get_graph_node_names, create_feature_extractor + +from .utils import IntermediateLayerGetter +from ..core import register + + +@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'] + """ diff --git a/engine/backbone/torchvision_model.py b/engine/backbone/torchvision_model.py new file mode 100644 index 0000000000000000000000000000000000000000..d8d2010e377dde7dac81a7d1b96485887eb8f805 --- /dev/null +++ b/engine/backbone/torchvision_model.py @@ -0,0 +1,48 @@ +""" +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) + + 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 diff --git a/engine/backbone/utils.py b/engine/backbone/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7c4e0901e30d8a6dbd2e0e69ac13479da6a9e3f7 --- /dev/null +++ b/engine/backbone/utils.py @@ -0,0 +1,54 @@ +""" +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 diff --git a/engine/backbone/vit_tiny.py b/engine/backbone/vit_tiny.py new file mode 100644 index 0000000000000000000000000000000000000000..e00291394d91466f41e9829c1a8b2c5f32e1e862 --- /dev/null +++ b/engine/backbone/vit_tiny.py @@ -0,0 +1,301 @@ +""" +DEIMv2: Real-Time Object Detection Meets DINOv3 +Copyright (c) 2025 The DEIMv2 Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from DINOv3 (https://github.com/facebookresearch/dinov3) +Modified from https://huggingface.co/spaces/Hila/RobustViT/blob/main/ViT/ViT_new.py + +""" +import math +import warnings +from functools import partial +from typing import List, Literal, Tuple + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import nn + + +class RopePositionEmbedding(nn.Module): + def __init__( + self, + embed_dim: int, + *, + num_heads: int, + base: float | None = 100.0, + min_period: float | None = None, + max_period: float | None = None, + normalize_coords: Literal["min", "max", "separate"] = "separate", + shift_coords: float | None = None, + jitter_coords: float | None = None, + rescale_coords: float | None = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + ): + super().__init__() + head_dim = embed_dim // num_heads + assert head_dim % 4 == 0, "Head dimension must be divisible by 4 for 2D RoPE" + both_periods = min_period is not None and max_period is not None + if (base is None and not both_periods) or (base is not None and both_periods): + raise ValueError("Either `base` or `min_period`+`max_period` must be provided.") + + self.base = base + self.min_period = min_period + self.max_period = max_period + self.D_head = head_dim + self.normalize_coords = normalize_coords + self.shift_coords = shift_coords + self.jitter_coords = jitter_coords + self.rescale_coords = rescale_coords + self.dtype = dtype + self.register_buffer( + "periods", + torch.empty(head_dim // 4, device=device, dtype=dtype), + persistent=True, + ) + self._init_weights() + + def forward(self, *, H: int, W: int) -> Tuple[torch.Tensor, torch.Tensor]: + device = self.periods.device + dtype = self.dtype if self.dtype is not None else torch.get_default_dtype() + dd = {"device": device, "dtype": dtype} + + if self.normalize_coords == "max": + max_HW = max(H, W) + coords_h = torch.arange(0.5, H, **dd) / max_HW + coords_w = torch.arange(0.5, W, **dd) / max_HW + elif self.normalize_coords == "separate": + coords_h = torch.arange(0.5, H, **dd) / H + coords_w = torch.arange(0.5, W, **dd) / W + else: # min + min_HW = min(H, W) + coords_h = torch.arange(0.5, H, **dd) / min_HW + coords_w = torch.arange(0.5, W, **dd) / min_HW + + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1) + coords = coords.flatten(0, 1) + coords = 2.0 * coords - 1.0 + + if self.training and self.shift_coords is not None: + coords += torch.empty(2, **dd).uniform_(-self.shift_coords, self.shift_coords)[None, :] + if self.training and self.jitter_coords is not None: + jitter = (torch.empty(2, **dd).uniform_(-np.log(self.jitter_coords), np.log(self.jitter_coords))).exp() + coords *= jitter[None, :] + if self.training and self.rescale_coords is not None: + rescale = (torch.empty(1, **dd).uniform_(-np.log(self.rescale_coords), np.log(self.rescale_coords))).exp() + coords *= rescale + + angles = 2 * math.pi * coords[:, :, None] / self.periods[None, None, :] + angles = angles.flatten(1, 2).repeat(1, 2) + + sin = torch.sin(angles) + cos = torch.cos(angles) + return sin.unsqueeze(0).unsqueeze(0), cos.unsqueeze(0).unsqueeze(0) + + def _init_weights(self): + device = self.periods.device + dtype = self.dtype if self.dtype is not None else torch.get_default_dtype() + if self.base is not None: + periods = self.base ** (2 * torch.arange(self.D_head // 4, device=device, dtype=dtype) / (self.D_head // 2)) + else: + base = self.max_period / self.min_period + exponents = torch.linspace(0, 1, self.D_head // 4, device=device, dtype=dtype) + periods = self.max_period * (base ** (exponents - 1)) + self.periods.data.copy_(periods) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rope(x, sin, cos): + """Applies RoPE to the input tensor.""" + return (x * cos) + (rotate_half(x) * sin) + + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + def forward(self, x): + x = self.fc1(x); x = self.act(x); x = self.drop(x); x = self.fc2(x); x = self.drop(x) + return x + + +class PatchEmbed(nn.Module): + def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): + super().__init__() + img_size = (img_size, img_size) if isinstance(img_size, int) else img_size + patch_size = (patch_size, patch_size) if isinstance(patch_size, int) else patch_size + self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) + self.num_patches = self.grid_size[0] * self.grid_size[1] + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + def forward(self, x): + return self.proj(x).flatten(2).transpose(1, 2) + + +def drop_path(x, drop_prob: float = 0., training: bool = False): + if drop_prob == 0. or not training: return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) + output = x.div(keep_prob) * random_tensor.floor() + return output + + +class DropPath(nn.Module): + def __init__(self, drop_prob=None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) + + +def _no_grad_trunc_normal_(tensor, mean, std, a, b): + def norm_cdf(x): return (1. + math.erf(x / math.sqrt(2.))) / 2. + if (mean < a - 2 * std) or (mean > b + 2 * std): + warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. The distribution of values may be incorrect.", stacklevel=2) + with torch.no_grad(): + l = norm_cdf((a - mean) / std); u = norm_cdf((b - mean) / std) + tensor.uniform_(2 * l - 1, 2 * u - 1); tensor.erfinv_(); tensor.mul_(std * math.sqrt(2.)); tensor.add_(mean); tensor.clamp_(min=a, max=b) + return tensor + + +def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.): + return _no_grad_trunc_normal_(tensor, mean, std, a, b) + + +class Attention(nn.Module): + def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim ** -0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = attn_drop + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x, rope_sincos=None): + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + + if rope_sincos is not None: + sin, cos = rope_sincos + q_cls, q_patch = q[:, :, :1, :], q[:, :, 1:, :] + k_cls, k_patch = k[:, :, :1, :], k[:, :, 1:, :] + + q_patch = apply_rope(q_patch, sin, cos) + k_patch = apply_rope(k_patch, sin, cos) + + q = torch.cat((q_cls, q_patch), dim=2) + k = torch.cat((k_cls, k_patch), dim=2) + + x = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop) + x = x.transpose(1, 2).reshape([B, N, C]) + x = self.proj(x) + x = self.proj_drop(x) + return x + +class Block(nn.Module): + def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + self.mlp = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop) + + def forward(self, x, rope_sincos=None): + attn_output = self.attn(self.norm1(x), rope_sincos=rope_sincos) + x = x + self.drop_path(attn_output) + x = x + self.drop_path(self.mlp(self.norm2(x))) + return x + + +class VisionTransformer(nn.Module): + def __init__( + self, img_size=224, patch_size=16, in_chans=3, embed_dim=192, depth=12, + num_heads=3, mlp_ratio=4., qkv_bias=True, drop_rate=0., attn_drop_rate=0., + drop_path_rate=0., return_layers=[3, 7, 11], embed_layer=PatchEmbed, + norm_layer=None, act_layer=None + ): + super().__init__() + self.num_features = self.embed_dim = embed_dim + self.num_tokens = 1 + self.return_layers = return_layers + norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) + act_layer = act_layer or nn.GELU + + self._model = nn.Module() + self._model.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) + self.patch_size = patch_size + self._model.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] + self._model.blocks = nn.ModuleList([ + Block( + dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], + norm_layer=norm_layer, act_layer=act_layer + ) for i in range(depth) + ]) + + self._model.rope_embed = RopePositionEmbedding( + embed_dim=embed_dim, num_heads=num_heads, base=100.0, + normalize_coords="separate", shift_coords=None, jitter_coords=None, + rescale_coords=None, dtype=None, device=None, + ) + self.init_weights() + + def init_weights(self): + trunc_normal_(self._model.cls_token, std=.02) + self._model.rope_embed._init_weights() + self.apply(self._init_vit_weights) + + def _init_vit_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)): + nn.init.zeros_(m.bias); nn.init.ones_(m.weight) + + @torch.jit.ignore + def no_weight_decay(self): + return {'cls_token'} + + def get_model(self): + return self._model + + def feature_dim(self): + return self.embed_dim + + def forward(self, x): + outs = [] + B, C, H, W = x.shape + + x_embed = self._model.patch_embed(x) + cls_token = self._model.cls_token.expand(x_embed.shape[0], -1, -1) + x = torch.cat((cls_token, x_embed), dim=1) + + patch_grid_h = H // self.patch_size + patch_grid_w = W // self.patch_size + rope_sincos = self._model.rope_embed(H=patch_grid_h, W=patch_grid_w) + + for i, blk in enumerate(self._model.blocks): + x = blk(x, rope_sincos=rope_sincos) + if i in self.return_layers: + outs.append((x[:, 1:], x[:, 0])) + return outs diff --git a/engine/core/__init__.py b/engine/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..61dbfa350a42e8b525225f715dbaf532d0c07051 --- /dev/null +++ b/engine/core/__init__.py @@ -0,0 +1,9 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +from .workspace import GLOBAL_CONFIG, register, create +from .yaml_utils import * +from ._config import BaseConfig +from .yaml_config import YAMLConfig diff --git a/engine/core/_config.py b/engine/core/_config.py new file mode 100644 index 0000000000000000000000000000000000000000..b8b601fe9ad93df6439a6668fe1f7ca4987119ed --- /dev/null +++ b/engine/core/_config.py @@ -0,0 +1,297 @@ +""" +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 torch.utils.data import Dataset, DataLoader +from torch.optim import Optimizer +from torch.optim.lr_scheduler import LRScheduler +from torch.cuda.amp.grad_scaler import GradScaler +from torch.utils.tensorboard import SummaryWriter + +from pathlib import Path +from typing import Callable, List, Dict + + +__all__ = ['BaseConfig', ] + + +class BaseConfig(object): + + 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.epoches :int = None + self.last_epoch :int = -1 + + self.lrsheduler: str = None + self.lr_gamma: float = None + self.no_aug_epoch: int = None + self.warmup_iter: int = None + self.flat_epoch: int = None + + 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. + 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 diff --git a/engine/core/workspace.py b/engine/core/workspace.py new file mode 100644 index 0000000000000000000000000000000000000000..2f9d2a146605e7b8c58c92eedbd6f5eb0981372d --- /dev/null +++ b/engine/core/workspace.py @@ -0,0 +1,177 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +import inspect +import importlib +import functools +from collections import defaultdict +from typing import Any, Dict, Optional, List + + +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) + 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.') + + _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}') + + module_kwargs = {k: v for k, v in module_kwargs.items() if not k.startswith('_')} + + # 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) diff --git a/engine/core/yaml_config.py b/engine/core/yaml_config.py new file mode 100644 index 0000000000000000000000000000000000000000..bdd27b41da8193aafe50d531adb5ceace0c50442 --- /dev/null +++ b/engine/core/yaml_config.py @@ -0,0 +1,174 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader + +import re +import copy + +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 + + @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 diff --git a/engine/core/yaml_utils.py b/engine/core/yaml_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..411e416d4f2ceb3e54c130860be7f765ed0e28af --- /dev/null +++ b/engine/core/yaml_utils.py @@ -0,0 +1,126 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +import os +import copy +import yaml +from typing import Any, Dict, Optional, List + +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) diff --git a/engine/data/__init__.py b/engine/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..72a9cc86a2f51347efd71a27b74583be030f6e40 --- /dev/null +++ b/engine/data/__init__.py @@ -0,0 +1,23 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +from .dataset import * +from .transforms import * +from .dataloader import * + +from ._misc import convert_to_tv_tensor + + + + +# 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]) diff --git a/engine/data/_misc.py b/engine/data/_misc.py new file mode 100644 index 0000000000000000000000000000000000000000..22c333fe9a5292bf29c2356a45641c840fce3002 --- /dev/null +++ b/engine/data/_misc.py @@ -0,0 +1,56 @@ +""" +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, Mask, Image, 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, Mask, Image, 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, Mask, Image, 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) diff --git a/engine/data/dataloader.py b/engine/data/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..388106b080744d2fd7136648a4f7e47b202e6ffb --- /dev/null +++ b/engine/data/dataloader.py @@ -0,0 +1,366 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE) +Copyright (c) 2024 D-FINE authors. All Rights Reserved. +""" + +import torch +import torch.utils.data as data +import torch.nn.functional as F +from torch.utils.data import default_collate + +import torchvision +import torchvision.transforms.v2 as VT +from torchvision.transforms.v2 import functional as VF, InterpolationMode + +import random +from functools import partial + +from ..core import register +torchvision.disable_beta_transforms_warning() +from copy import deepcopy +from PIL import Image, ImageDraw +import os +from collections import defaultdict, deque + + +__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, + mixup_prob=0.0, + mixup_epochs=[0, 0], + copyblend_prob=0.0, + copyblend_epochs=[0, 0], + copyblend_type='blend', + conflict_with_mixup=False, + area_threshold=100, + num_objects=3, + with_expand=False, + expand_ratios=[0.1, 0.25], + random_num_objects=False, + data_vis=False, + vis_save='./vis_dataset/' + ) -> 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.mixup_prob, self.mixup_epochs = mixup_prob, mixup_epochs + + self.copyblend_prob, self.copyblend_epochs, self.copyblend_type = copyblend_prob, copyblend_epochs, copyblend_type + self.area_threshold, self.num_objects = area_threshold, num_objects + self.data_vis, self.vis_save = data_vis, vis_save + self.with_expand, self.expand_ratios, self.random_num_objects = with_expand, expand_ratios, random_num_objects + self.conflict_with_mixup = conflict_with_mixup # 是否冲突 + + if self.mixup_prob > 0 or self.copyblend_prob > 0: + if os.path.isdir(self.vis_save): + for file in os.listdir(self.vis_save): + os.remove('{}/{}'.format(self.vis_save, file)) + os.makedirs(self.vis_save, exist_ok=True) if self.data_vis else None + + if self.mixup_prob > 0: + print(" ### Using MixUp with Prob@{} in {} epochs ### ".format(mixup_prob, mixup_epochs)) + if self.copyblend_prob > 0: + print(" ### Using CopyBlend-{} with Prob@{} in {} epochs ### ".format(copyblend_type, copyblend_prob, copyblend_epochs)) + print(f' ### CopyBlend -- area threshold@{area_threshold} and num of object@{num_objects} ### ') + if self.with_expand: + print(f' ### CopyBlend -- expand@{expand_ratios} ### ') + if self.random_num_objects: + print(f' ### CopyBlend -- random num of objects@{[1, self.num_objects]} ### ') + + if stop_epoch is not None: + print(" ### Multi-scale Training until {} epochs ### ".format(self.stop_epoch)) + print(" ### Multi-scales@ {} ### ".format(self.scales)) + self.print_info_flag = True + self.print_copyblend_flag = True + # self.interpolation = interpolation + + def apply_mixup(self, images, targets): + """ + Applies Mixup augmentation to the batch if conditions are met. + + Args: + images (torch.Tensor): Batch of images. + targets (list[dict]): List of target dictionaries corresponding to images. + + Returns: + tuple: Updated images and targets + """ + # Log when Mixup is permanently disabled + if self.epoch == self.mixup_epochs[-1] and self.print_info_flag: + print(f" ### Attention --- Mixup is closed after epoch@ {self.epoch} ###") + self.print_info_flag = False + + MixUp_flag, CopyBlend_flag = False, False + beta = round(random.uniform(0.45, 0.55), 6) + # Apply Mixup if within specified epoch range and probability threshold + if random.random() < self.mixup_prob and self.mixup_epochs[0] <= self.epoch < self.mixup_epochs[-1]: + # Generate mixup ratio + beta = round(random.uniform(0.45, 0.55), 6) + MixUp_flag = True + + # Mix images + images = images.roll(shifts=1, dims=0).mul_(1.0 - beta).add_(images.mul(beta)) + + # Prepare targets for Mixup + shifted_targets = targets[-1:] + targets[:-1] + updated_targets = deepcopy(targets) + + for i in range(len(targets)): + # Combine boxes, labels, and areas from original and shifted targets + updated_targets[i]['boxes'] = torch.cat([targets[i]['boxes'], shifted_targets[i]['boxes']], dim=0) + updated_targets[i]['labels'] = torch.cat([targets[i]['labels'], shifted_targets[i]['labels']], dim=0) + updated_targets[i]['area'] = torch.cat([targets[i]['area'], shifted_targets[i]['area']], dim=0) + + # Add mixup ratio to targets + updated_targets[i]['mixup'] = torch.tensor( + [beta] * len(targets[i]['labels']) + [1.0 - beta] * len(shifted_targets[i]['labels']), + dtype=torch.float32 + ) + targets = updated_targets + + elif (self.copyblend_epochs[0] <= self.epoch < self.copyblend_epochs[-1] and random.random() < self.copyblend_prob): + if self.epoch == self.copyblend_epochs[-1] and self.print_copyblend_flag: + print(f" ### Attention --- CopyBlend closed after epoch@ {self.epoch} ###") + self.print_copyblend_flag = False + + CopyBlend_flag = True + objects_pool = defaultdict(list) + img_height, img_width = images[0].shape[-2:] + + # get all valid objects in batch + for i in range(len(images)): + source_boxes = targets[i]['boxes'] + source_labels = targets[i]['labels'] + source_areas = targets[i]['area'] + + # filter valid objects + valid_objects = [idx for idx in range(len(source_boxes)) if source_areas[idx] >= self.area_threshold] + for idx in valid_objects: + objects_pool['boxes'].append(source_boxes[idx]) + objects_pool['labels'].append(source_labels[idx]) + objects_pool['areas'].append(source_areas[idx]) + objects_pool['image_idx'].append(i) + objects_pool['image_height'].append(img_height) + objects_pool['image_width'].append(img_width) + + # check if objects_pool is empty + if len(objects_pool['boxes']) == 0: + return images, targets + + # convert list to tensor for convenient operation + for key in ['boxes', 'labels', 'areas']: + objects_pool[key] = torch.stack(objects_pool[key]) if objects_pool[key] else torch.tensor([]) + + # apply CopyBlend + batch_size = len(images) + updated_images = images.clone() + updated_targets = deepcopy(targets) + + for i in range(batch_size): + # randomly decide the number of objects to blend + if self.random_num_objects: + num_objects = random.randint(1, min(self.num_objects, len(objects_pool['boxes']))) + else: + num_objects = min(self.num_objects, len(objects_pool['boxes'])) + + # randomly select objects to blend + selected_indices = random.sample(range(len(objects_pool['boxes'])), num_objects) + + blend_boxes = [] + blend_labels = [] + blend_areas = [] + blend_mixup_ratios = [] + + for idx in selected_indices: + # get source object information + box = objects_pool['boxes'][idx] + label = objects_pool['labels'][idx] + area = objects_pool['areas'][idx] + source_idx = objects_pool['image_idx'][idx] + source_height = objects_pool['image_height'][idx] + source_width = objects_pool['image_width'][idx] + + # calculate source object size and position + cx, cy, w, h = box + x1_src, y1_src = int((cx - w / 2) * source_width), int((cy - h / 2) * source_height) + x2_src, y2_src = int((cx + w / 2) * source_width), int((cy + h / 2) * source_height) + + # check if source object is out of bound + x1_src, y1_src = max(x1_src, 0), max(y1_src, 0) + x2_src, y2_src = min(x2_src, img_width), min(y2_src, img_height) + new_w_px, new_h_px = x2_src - x1_src, y2_src - y1_src + # check if source object is valid + if new_w_px <= 0 or new_h_px <= 0: + continue + + # randomly determine blend position + x1 = random.randint(0, img_width - new_w_px) if new_w_px < img_width else 0 + y1 = random.randint(0, img_height - new_h_px) if new_h_px < img_height else 0 + # after the above limit, [x2, y2] will not be out of bound, so no need to check + x2, y2 = x1 + new_w_px, y1 + new_h_px + + # calculate new normalized coordinates + new_cx, new_cy = (x1 + new_w_px / 2) / img_width, (y1 + new_h_px / 2) / img_height + new_w, new_h = new_w_px / img_width, new_h_px / img_height + + # add to blend list - use original unexpanded box + blend_boxes.append(torch.tensor([new_cx, new_cy, new_w, new_h])) + blend_labels.append(label) + blend_areas.append(area) + # mixup ratio + blend_mixup_ratios.append(1.0 - beta) + + # handle expanded area + if self.with_expand: + alpha = round(random.uniform(self.expand_ratios[0], self.expand_ratios[1]), 6) + expand_w, expand_h = int(new_w_px * alpha), int(new_h_px * alpha) + # check if out of bound: get the best offset in GT image + x1_expand, y1_expand = x1_src - max(x1_src - expand_w, 0), y1_src - max(y1_src - expand_h, 0) + x2_expand, y2_expand = min(x2_src + expand_w, img_width) - x2_src, min(y2_src + expand_h, img_height) - y2_src + # check if out of bound: whether the expanded area is out of bound in blend image + new_x1_expand, new_y1_expand = x1 - max(x1 - x1_expand, 0), y1 - max(y1 - y1_expand, 0) + new_x2_expand, new_y2_expand = min(x2 + x2_expand, img_width) - x2, min(y2 + y2_expand, img_height) - y2 + # update + x1_src, y1_src, x2_src, y2_src = x1_src - new_x1_expand, y1_src - new_y1_expand, x2_src + new_x2_expand, y2_src + new_y2_expand + x1, y1, x2, y2 = x1 - new_x1_expand, y1 - new_y1_expand, x2 + new_x2_expand, y2 + new_y2_expand + + # blend original area first + copy_patch_orig = images[source_idx, :, y1_src:y2_src, x1_src:x2_src] + if self.copyblend_type == 'blend': + blended_patch = updated_images[i, :, y1:y2, x1:x2] * beta + copy_patch_orig * (1 - beta) + updated_images[i, :, y1:y2, x1:x2] = blended_patch + else: + updated_images[i, :, y1:y2, x1:x2] = copy_patch_orig + + # add blended objects to targets + if len(blend_boxes) > 0: + blend_boxes = torch.stack(blend_boxes) + blend_labels = torch.stack(blend_labels) + blend_areas = torch.stack(blend_areas) + + # add mixup ratio + updated_targets[i]['mixup'] = torch.tensor( + [1.0] * len(updated_targets[i]['boxes']) + blend_mixup_ratios, + dtype=torch.float32 + ) + # update targets + updated_targets[i]['boxes'] = torch.cat([updated_targets[i]['boxes'], blend_boxes]) + updated_targets[i]['labels'] = torch.cat([updated_targets[i]['labels'], blend_labels]) + updated_targets[i]['area'] = torch.cat([updated_targets[i]['area'], blend_areas]) + + images, targets = updated_images, updated_targets + + if self.data_vis and CopyBlend_flag: + for i in range(len(updated_targets)): + image_tensor = images[i] + if image_tensor.min() < 0: # use normalization + image_tensor = image_tensor * torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) \ + + torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + image_tensor_uint8 = (image_tensor * 255).type(torch.uint8) + image_numpy = image_tensor_uint8.numpy().transpose((1, 2, 0)) + pilImage = Image.fromarray(image_numpy) + draw = ImageDraw.Draw(pilImage) + print('mix_vis:', i, 'boxes.len=', len(updated_targets[i]['boxes'])) + for box in updated_targets[i]['boxes']: + draw.rectangle([int(box[0]*640 - (box[2]*640)/2), int(box[1]*640 - (box[3]*640)/2), + int(box[0]*640 + (box[2]*640)/2), int(box[1]*640 + (box[3]*640)/2)], outline=(255,255,0)) + pilImage.save(self.vis_save + str(i) + "_"+ str(len(updated_targets[i]['boxes'])) +'_out.jpg') + + return images, targets + + def __call__(self, items): + images = torch.cat([x[0][None] for x in items], dim=0) + targets = [x[1] for x in items] + + # Mixup + images, targets = self.apply_mixup(images, targets) + + 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 diff --git a/engine/data/dataset/__init__.py b/engine/data/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c946e0cdfc02d2e23ef5d34c2ce3f190a007dc70 --- /dev/null +++ b/engine/data/dataset/__init__.py @@ -0,0 +1,16 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +# from ._dataset import DetDataset +from .coco_dataset import CocoDetection +from .coco_dataset import ( + mscoco_category2name, + mscoco_category2label, + 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 diff --git a/engine/data/dataset/_dataset.py b/engine/data/dataset/_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..66784a86d871078a5f497a47635d6cab044b8d4d --- /dev/null +++ b/engine/data/dataset/_dataset.py @@ -0,0 +1,24 @@ +""" +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 diff --git a/engine/data/dataset/coco_dataset.py b/engine/data/dataset/coco_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..83202e74019f07d259ba00989744684d91498e7d --- /dev/null +++ b/engine/data/dataset/coco_dataset.py @@ -0,0 +1,264 @@ +""" +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 torch +import torch.utils.data + +import torchvision + +from PIL import Image +import faster_coco_eval +import faster_coco_eval.core.mask as coco_mask +from ._dataset import DetDataset +from .._misc import convert_to_tv_tensor +from ...core import register + +torchvision.disable_beta_transforms_warning() +faster_coco_eval.init_as_pycocotools() +Image.MAX_IMAGE_PIXELS = None + +__all__ = ['CocoDetection'] + + +@register() +class CocoDetection(torchvision.datasets.CocoDetection, DetDataset): + __inject__ = ['transforms', ] + __share__ = ['remap_mscoco_category'] + + def __init__(self, img_folder, ann_file, transforms, return_masks=False, remap_mscoco_category=False): + super(CocoDetection, 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(CocoDetection, self).__getitem__(idx) + image_id = self.ids[idx] + target = {'image_id': image_id, '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]) + + 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 + 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()} diff --git a/engine/data/dataset/coco_eval.py b/engine/data/dataset/coco_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..75f6bd8ddde3fa164fbc87f499f3f3a2919f4ee5 --- /dev/null +++ b/engine/data/dataset/coco_eval.py @@ -0,0 +1,200 @@ +""" +# 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 +""" +import os +import contextlib +import copy +import numpy as np +import torch + +from faster_coco_eval import COCO, COCOeval_faster +import faster_coco_eval.core.mask as mask_util +from ...core import register +from ...misc import dist_utils +__all__ = ['CocoEvaluator',] + + +@register() +class CocoEvaluator(object): + def __init__(self, coco_gt, iou_types): + assert isinstance(iou_types, (list, tuple)) + coco_gt = copy.deepcopy(coco_gt) + self.coco_gt : COCO = coco_gt + self.iou_types = iou_types + + self.coco_eval = {} + for iou_type in iou_types: + self.coco_eval[iou_type] = COCOeval_faster(coco_gt, iouType=iou_type, print_function=print, separate_eval=True) + + self.img_ids = [] + self.eval_imgs = {k: [] for k in iou_types} + + def cleanup(self): + self.coco_eval = {} + for iou_type in self.iou_types: + self.coco_eval[iou_type] = COCOeval_faster(self.coco_gt, iouType=iou_type, print_function=print, separate_eval=True) + self.img_ids = [] + self.eval_imgs = {k: [] for k in self.iou_types} + + + def update(self, predictions): + img_ids = list(np.unique(list(predictions.keys()))) + self.img_ids.extend(img_ids) + + for iou_type in self.iou_types: + results = self.prepare(predictions, iou_type) + coco_eval = self.coco_eval[iou_type] + + # suppress pycocotools prints + with open(os.devnull, 'w') as devnull: + with contextlib.redirect_stdout(devnull): + coco_dt = self.coco_gt.loadRes(results) if results else COCO() + coco_eval.cocoDt = coco_dt + coco_eval.params.imgIds = list(img_ids) + coco_eval.evaluate() + + self.eval_imgs[iou_type].append(np.array(coco_eval._evalImgs_cpp).reshape(len(coco_eval.params.catIds), len(coco_eval.params.areaRng), len(coco_eval.params.imgIds))) + + def synchronize_between_processes(self): + for iou_type in self.iou_types: + img_ids, eval_imgs = merge(self.img_ids, self.eval_imgs[iou_type]) + + coco_eval = self.coco_eval[iou_type] + coco_eval.params.imgIds = img_ids + coco_eval._paramsEval = copy.deepcopy(coco_eval.params) + coco_eval._evalImgs_cpp = eval_imgs + + def accumulate(self): + for coco_eval in self.coco_eval.values(): + coco_eval.accumulate() + + def summarize(self): + for iou_type, coco_eval in self.coco_eval.items(): + print("IoU metric: {}".format(iou_type)) + coco_eval.summarize() + + def prepare(self, predictions, iou_type): + if iou_type == "bbox": + return self.prepare_for_coco_detection(predictions) + elif iou_type == "segm": + return self.prepare_for_coco_segmentation(predictions) + elif iou_type == "keypoints": + return self.prepare_for_coco_keypoint(predictions) + else: + raise ValueError("Unknown iou type {}".format(iou_type)) + + def prepare_for_coco_detection(self, predictions): + coco_results = [] + for original_id, prediction in predictions.items(): + if len(prediction) == 0: + continue + + boxes = prediction["boxes"] + boxes = convert_to_xywh(boxes).tolist() + scores = prediction["scores"].tolist() + labels = prediction["labels"].tolist() + + coco_results.extend( + [ + { + "image_id": original_id, + "category_id": labels[k], + "bbox": box, + "score": scores[k], + } + for k, box in enumerate(boxes) + ] + ) + return coco_results + + def prepare_for_coco_segmentation(self, predictions): + coco_results = [] + for original_id, prediction in predictions.items(): + if len(prediction) == 0: + continue + + scores = prediction["scores"] + labels = prediction["labels"] + masks = prediction["masks"] + + masks = masks > 0.5 + + scores = prediction["scores"].tolist() + labels = prediction["labels"].tolist() + + rles = [ + mask_util.encode(np.array(mask[0, :, :, np.newaxis], dtype=np.uint8, order="F"))[0] + for mask in masks + ] + for rle in rles: + rle["counts"] = rle["counts"].decode("utf-8") + + coco_results.extend( + [ + { + "image_id": original_id, + "category_id": labels[k], + "segmentation": rle, + "score": scores[k], + } + for k, rle in enumerate(rles) + ] + ) + return coco_results + + def prepare_for_coco_keypoint(self, predictions): + coco_results = [] + for original_id, prediction in predictions.items(): + if len(prediction) == 0: + continue + + boxes = prediction["boxes"] + boxes = convert_to_xywh(boxes).tolist() + scores = prediction["scores"].tolist() + labels = prediction["labels"].tolist() + keypoints = prediction["keypoints"] + keypoints = keypoints.flatten(start_dim=1).tolist() + + coco_results.extend( + [ + { + "image_id": original_id, + "category_id": labels[k], + 'keypoints': keypoint, + "score": scores[k], + } + for k, keypoint in enumerate(keypoints) + ] + ) + return coco_results + + +def convert_to_xywh(boxes): + xmin, ymin, xmax, ymax = boxes.unbind(1) + return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=1) + +def merge(img_ids, eval_imgs): + all_img_ids = dist_utils.all_gather(img_ids) + all_eval_imgs = dist_utils.all_gather(eval_imgs) + + merged_img_ids = [] + for p in all_img_ids: + merged_img_ids.extend(p) + + merged_eval_imgs = [] + for p in all_eval_imgs: + merged_eval_imgs.extend(p) + + + merged_img_ids = np.array(merged_img_ids) + merged_eval_imgs = np.concatenate(merged_eval_imgs, axis=2).ravel() + # merged_eval_imgs = np.array(merged_eval_imgs).T.ravel() + + # keep only unique (and in sorted order) images + merged_img_ids, idx = np.unique(merged_img_ids, return_index=True) + + return merged_img_ids.tolist(), merged_eval_imgs.tolist() diff --git a/engine/data/dataset/coco_utils.py b/engine/data/dataset/coco_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6b81b5ea9524618fbbae575b89d42e780dc5a050 --- /dev/null +++ b/engine/data/dataset/coco_utils.py @@ -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 torch +import torch.utils.data +import torchvision +import torchvision.transforms.functional as TVF +import faster_coco_eval.core.mask as coco_mask +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): + 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) diff --git a/engine/data/dataset/voc_detection.py b/engine/data/dataset/voc_detection.py new file mode 100644 index 0000000000000000000000000000000000000000..1c6dd13d0c83d0cbb69b6388ef76e40731cb5dde --- /dev/null +++ b/engine/data/dataset/voc_detection.py @@ -0,0 +1,76 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +from sympy import im +import torch +import torchvision +import torchvision.transforms.functional as TVF + +import os +from PIL import Image +from typing import Optional, Callable + +try: + from defusedxml.ElementTree import parse as ET_parse +except ImportError: + from xml.etree.ElementTree import parse as ET_parse + +from ._dataset import DetDataset +from .._misc import convert_to_tv_tensor +from ...core import register + +@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 diff --git a/engine/data/dataset/voc_eval.py b/engine/data/dataset/voc_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..0bee50ae4e12f4f688083cd2e5620b9602b55e7e --- /dev/null +++ b/engine/data/dataset/voc_eval.py @@ -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 diff --git a/engine/data/transforms/__init__.py b/engine/data/transforms/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5567d6b8042e93236d54398e9146acaed066676d --- /dev/null +++ b/engine/data/transforms/__init__.py @@ -0,0 +1,22 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + + +from ._transforms import ( + EmptyTransform, + RandomPhotometricDistort, + RandomZoomOut, + RandomIoUCrop, + RandomHorizontalFlip, + Resize, + PadToSize, + SanitizeBoundingBoxes, + RandomCrop, + Normalize, + ConvertBoxes, + ConvertPILImage, +) +from .container import Compose +from .mosaic import Mosaic \ No newline at end of file diff --git a/engine/data/transforms/_transforms.py b/engine/data/transforms/_transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..31588df5203041730da89b7231479b5b4fc92f20 --- /dev/null +++ b/engine/data/transforms/_transforms.py @@ -0,0 +1,137 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +import torch +import torch.nn as nn + +import torchvision +import torchvision.transforms.v2 as T +import torchvision.transforms.v2.functional as F + +import PIL +import PIL.Image + +from typing import Any, Dict, List, Optional + +from .._misc import convert_to_tv_tensor, _boxes_keys +from .._misc import Image, Video, Mask, BoundingBoxes +from .._misc import SanitizeBoundingBoxes + +from ...core import register +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: + 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: + inpt = F.pil_to_tensor(inpt) + if self.dtype == 'float32': + inpt = inpt.float() + + if self.scale: + inpt = inpt / 255. + + inpt = Image(inpt) + + return inpt diff --git a/engine/data/transforms/container.py b/engine/data/transforms/container.py new file mode 100644 index 0000000000000000000000000000000000000000..203cb11933d0e51366fba03571c1491541b8ba3b --- /dev/null +++ b/engine/data/transforms/container.py @@ -0,0 +1,126 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE) +Copyright (c) 2024 D-FINE authors. All Rights Reserved. +""" + +import torch +import torch.nn as nn + +import torchvision +import torchvision.transforms.v2 as T + +from typing import Any, Dict, List, Optional + +from ._transforms import EmptyTransform +from ...core import register, GLOBAL_CONFIG +torchvision.disable_beta_transforms_warning() +import random + + +@register() +class Compose(T.Compose): + def __init__(self, ops, policy=None, mosaic_prob=-0.1) -> 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 + print(" ### Transform @{} ### ".format(type(transform).__name__)) + + elif isinstance(op, nn.Module): + transforms.append(op) + + else: + raise ValueError('') + else: + transforms =[EmptyTransform(), ] + + super().__init__(transforms=transforms) + + self.mosaic_prob = mosaic_prob + if policy is None: + policy = {'name': 'default'} + else: + if self.mosaic_prob > 0: + print(" ### Mosaic with Prob.@{} and ZoomOut/IoUCrop existed ### ".format(self.mosaic_prob)) + print(" ### ImgTransforms Epochs: {} ### ".format(policy['epoch'])) + print(' ### Policy_ops@{} ###'.format(policy['ops'])) + self.global_samples = 0 + self.policy = policy + + 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'] + + if isinstance(policy_epoch, list) and len(policy_epoch) == 3: # 4-stages + if policy_epoch[0] <= cur_epoch < policy_epoch[1]: + with_mosaic = random.random() <= self.mosaic_prob # Probility for Mosaic + else: + with_mosaic = False + for transform in self.transforms: + if (type(transform).__name__ in policy_ops and cur_epoch < policy_epoch[0]): # first stage: NoAug + pass + elif (type(transform).__name__ in policy_ops and cur_epoch >= policy_epoch[-1]): # last stage: NoAug + pass + else: + # Using Mosaic for [policy_epoch[0], policy_epoch[1]] with probability + if (type(transform).__name__ == 'Mosaic' and not with_mosaic): + pass + # Mosaic and Zoomout/IoUCrop can not be co-existed in the same sample + elif (type(transform).__name__ == 'RandomZoomOut' or type(transform).__name__ == 'RandomIoUCrop') and with_mosaic: + pass + else: + sample = transform(sample) + else: # the default data scheduler + 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 diff --git a/engine/data/transforms/functional.py b/engine/data/transforms/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..0106c5233b5ca6a191577c78c8cb058d22c59338 --- /dev/null +++ b/engine/data/transforms/functional.py @@ -0,0 +1,168 @@ +import torch +import torchvision.transforms.functional as F + +from packaging import version +from typing import Optional, List +from torch import Tensor + +# needed due to empty tensor bug in pytorch and torchvision 0.5 +import torchvision +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: + 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 diff --git a/engine/data/transforms/mosaic.py b/engine/data/transforms/mosaic.py new file mode 100644 index 0000000000000000000000000000000000000000..fb065e5abc0c59535fc6f4993273f2c473af1460 --- /dev/null +++ b/engine/data/transforms/mosaic.py @@ -0,0 +1,168 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +""" + +import torch +import torchvision.transforms.v2 as T +import torchvision.transforms.v2.functional as F +import random +from PIL import Image + +from .._misc import convert_to_tv_tensor +from ...core import register + + +@register() +class Mosaic(T.Transform): + """ + Applies Mosaic augmentation to a batch of images. Combines four randomly selected images + into a single composite image with randomized transformations. + """ + + def __init__(self, output_size=320, max_size=None, rotation_range=0, translation_range=(0.1, 0.1), + scaling_range=(0.5, 1.5), probability=1.0, fill_value=114, use_cache=True, max_cached_images=50, + random_pop=True) -> None: + """ + Args: + output_size (int): Target size for resizing individual images. + rotation_range (float): Range of rotation in degrees for affine transformation. + translation_range (tuple): Range of translation for affine transformation. + scaling_range (tuple): Range of scaling factors for affine transformation. + probability (float): Probability of applying the Mosaic augmentation. + fill_value (int): Fill value for padding or affine transformations. + use_cache (bool): Whether to use cache. Defaults to True. + max_cached_images (int): The maximum length of the cache. + random_pop (bool): Whether to randomly pop a result from the cache. + """ + super().__init__() + self.resize = T.Resize(size=output_size, max_size=max_size) + self.probability = probability + self.affine_transform = T.RandomAffine(degrees=rotation_range, translate=translation_range, + scale=scaling_range, fill=fill_value) + self.use_cache = use_cache + self.mosaic_cache = [] + self.max_cached_images = max_cached_images + self.random_pop = random_pop + + def load_samples_from_dataset(self, image, target, dataset): + """Loads and resizes a set of images and their corresponding targets.""" + # Append the main image + get_size_func = F.get_size if hasattr(F, "get_size") else F.get_spatial_size # torchvision >=0.17 is get_size + image, target = self.resize(image, target) + resized_images, resized_targets = [image], [target] + max_height, max_width = get_size_func(resized_images[0]) + + # randomly select 3 images + sample_indices = random.choices(range(len(dataset)), k=3) + for idx in sample_indices: + # image, target = dataset.load_item(idx) + image, target = self.resize(dataset.load_item(idx)) + height, width = get_size_func(image) + max_height, max_width = max(max_height, height), max(max_width, width) + resized_images.append(image) + resized_targets.append(target) + + return resized_images, resized_targets, max_height, max_width + + def load_samples_from_cache(self, image, target, cache): + image, target = self.resize(image, target) + cache.append(dict(img=image, labels=target)) + + if len(cache) > self.max_cached_images: + if self.random_pop: + index = random.randint(0, len(cache) - 2) # do not remove last image + else: + index = 0 + cache.pop(index) + sample_indices = random.choices(range(len(cache)), k=3) + mosaic_samples = [dict(img=cache[idx]["img"].copy(), labels=self._clone(cache[idx]["labels"])) for idx in + sample_indices] # sample 3 images + mosaic_samples = [dict(img=image.copy(), labels=self._clone(target))] + mosaic_samples + + get_size_func = F.get_size if hasattr(F, "get_size") else F.get_spatial_size + sizes = [get_size_func(mosaic_samples[idx]["img"]) for idx in range(4)] + max_height = max(size[0] for size in sizes) + max_width = max(size[1] for size in sizes) + + return mosaic_samples, max_height, max_width + + def create_mosaic_from_cache(self, mosaic_samples, max_height, max_width): + placement_offsets = [[0, 0], [max_width, 0], [0, max_height], [max_width, max_height]] + merged_image = Image.new(mode=mosaic_samples[0]["img"].mode, size=(max_width * 2, max_height * 2), color=0) + offsets = torch.tensor([[0, 0], [max_width, 0], [0, max_height], [max_width, max_height]]).repeat(1, 2) + + mosaic_target = [] + for i, sample in enumerate(mosaic_samples): + img = sample["img"] + target = sample["labels"] + + merged_image.paste(img, placement_offsets[i]) + target['boxes'] = target['boxes'] + offsets[i] + mosaic_target.append(target) + + merged_target = {} + for key in mosaic_target[0]: + merged_target[key] = torch.cat([target[key] for target in mosaic_target]) + + return merged_image, merged_target + + def create_mosaic_from_dataset(self, images, targets, max_height, max_width): + """Creates a mosaic image by combining multiple images.""" + placement_offsets = [[0, 0], [max_width, 0], [0, max_height], [max_width, max_height]] + merged_image = Image.new(mode=images[0].mode, size=(max_width * 2, max_height * 2), color=0) + for i, img in enumerate(images): + merged_image.paste(img, placement_offsets[i]) + + """Merges targets into a single target dictionary for the mosaic.""" + offsets = torch.tensor([[0, 0], [max_width, 0], [0, max_height], [max_width, max_height]]).repeat(1, 2) + merged_target = {} + for key in targets[0]: + if key == 'boxes': + values = [target[key] + offsets[i] for i, target in enumerate(targets)] + else: + values = [target[key] for target in targets] + + merged_target[key] = torch.cat(values, dim=0) if isinstance(values[0], torch.Tensor) else values + + return merged_image, merged_target + + @staticmethod + def _clone(tensor_dict): + return {key: value.clone() for (key, value) in tensor_dict.items()} + + def forward(self, *inputs): + """ + Args: + inputs (tuple): Input tuple containing (image, target, dataset). + + Returns: + tuple: Augmented (image, target, dataset). + """ + if len(inputs) == 1: + inputs = inputs[0] + image, target, dataset = inputs + + # Skip mosaic augmentation with probability 1 - self.probability + if self.probability < 1.0 and random.random() > self.probability: + return image, target, dataset + + # Prepare mosaic components + if self.use_cache: + mosaic_samples, max_height, max_width = self.load_samples_from_cache(image, target, self.mosaic_cache) + mosaic_image, mosaic_target = self.create_mosaic_from_cache(mosaic_samples, max_height, max_width) + else: + resized_images, resized_targets, max_height, max_width = self.load_samples_from_dataset(image, target,dataset) + mosaic_image, mosaic_target = self.create_mosaic_from_dataset(resized_images, resized_targets, max_height, max_width) + + # Clamp boxes and convert target formats + if 'boxes' in mosaic_target: + mosaic_target['boxes'] = convert_to_tv_tensor(mosaic_target['boxes'], 'boxes', box_format='xyxy', + spatial_size=mosaic_image.size[::-1]) + if 'masks' in mosaic_target: + mosaic_target['masks'] = convert_to_tv_tensor(mosaic_target['masks'], 'masks') + + # Apply affine transformations + mosaic_image, mosaic_target = self.affine_transform(mosaic_image, mosaic_target) + + return mosaic_image, mosaic_target, dataset diff --git a/engine/deim/__init__.py b/engine/deim/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6517d1da48a2b12990945665a2cf6f79ec20ef17 --- /dev/null +++ b/engine/deim/__init__.py @@ -0,0 +1,23 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + + +from .deim import DEIM + +from .matcher import HungarianMatcher + +from .hybrid_encoder import HybridEncoder +from .lite_encoder import LiteEncoder + + +from .dfine_decoder import DFINETransformer +from .rtdetrv2_decoder import RTDETRTransformerv2 + +from .postprocessor import PostProcessor +from .deim_criterion import DEIMCriterion +from .deim_decoder import DEIMTransformer \ No newline at end of file diff --git a/engine/deim/box_ops.py b/engine/deim/box_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ede1b32470d71c6f4b0e492fcd655ca2d44ac40b --- /dev/null +++ b/engine/deim/box_ops.py @@ -0,0 +1,90 @@ +""" +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +https://github.com/facebookresearch/detr/blob/main/util/box_ops.py +""" + +import torch +from torch import Tensor +from torchvision.ops.boxes import box_area + + +def box_cxcywh_to_xyxy(x): + x_c, y_c, w, h = x.unbind(-1) + b = [(x_c - 0.5 * w.clamp(min=0.0)), (y_c - 0.5 * h.clamp(min=0.0)), + (x_c + 0.5 * w.clamp(min=0.0)), (y_c + 0.5 * h.clamp(min=0.0))] + return torch.stack(b, dim=-1) + + +def box_xyxy_to_cxcywh(x: Tensor) -> Tensor: + x0, y0, x1, y1 = x.unbind(-1) + b = [(x0 + x1) / 2, (y0 + y1) / 2, + (x1 - x0), (y1 - y0)] + return torch.stack(b, dim=-1) + + +# modified from torchvision to also return the union +def box_iou(boxes1: Tensor, boxes2: Tensor): + area1 = box_area(boxes1) + area2 = box_area(boxes2) + + lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2] + rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2] + + wh = (rb - lt).clamp(min=0) # [N,M,2] + inter = wh[:, :, 0] * wh[:, :, 1] # [N,M] + + union = area1[:, None] + area2 - inter + + iou = inter / union + return iou, union + + +def generalized_box_iou(boxes1, boxes2): + """ + Generalized IoU from https://giou.stanford.edu/ + + The boxes should be in [x0, y0, x1, y1] format + + Returns a [N, M] pairwise matrix, where N = len(boxes1) + and M = len(boxes2) + """ + # degenerate boxes gives inf / nan results + # so do an early check + assert (boxes1[:, 2:] >= boxes1[:, :2]).all() + assert (boxes2[:, 2:] >= boxes2[:, :2]).all() + iou, union = box_iou(boxes1, boxes2) + + lt = torch.min(boxes1[:, None, :2], boxes2[:, :2]) + rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:]) + + wh = (rb - lt).clamp(min=0) # [N,M,2] + area = wh[:, :, 0] * wh[:, :, 1] + + return iou - (area - union) / area + + +def masks_to_boxes(masks): + """Compute the bounding boxes around the provided masks + + The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions. + + Returns a [N, 4] tensors, with the boxes in xyxy format + """ + if masks.numel() == 0: + return torch.zeros((0, 4), device=masks.device) + + h, w = masks.shape[-2:] + + y = torch.arange(0, h, dtype=torch.float) + x = torch.arange(0, w, dtype=torch.float) + y, x = torch.meshgrid(y, x) + + x_mask = (masks * x.unsqueeze(0)) + x_max = x_mask.flatten(1).max(-1)[0] + x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] + + y_mask = (masks * y.unsqueeze(0)) + y_max = y_mask.flatten(1).max(-1)[0] + y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] + + return torch.stack([x_min, y_min, x_max, y_max], 1) \ No newline at end of file diff --git a/engine/deim/deim.py b/engine/deim/deim.py new file mode 100644 index 0000000000000000000000000000000000000000..7170fb5a67d31c91325e4ca391dd81c54cdcb8f6 --- /dev/null +++ b/engine/deim/deim.py @@ -0,0 +1,38 @@ +""" +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +""" + +import torch.nn as nn +from ..core import register + + +__all__ = ['DEIM', ] + + +@register() +class DEIM(nn.Module): + __inject__ = ['backbone', 'encoder', 'decoder', ] + + def __init__(self, \ + backbone: nn.Module, + encoder: nn.Module, + decoder: nn.Module, + ): + super().__init__() + self.backbone = backbone + self.decoder = decoder + self.encoder = encoder + + def forward(self, x, targets=None): + x = self.backbone(x) + x = self.encoder(x) + x = self.decoder(x, targets) + + return x + + def deploy(self, ): + self.eval() + for m in self.modules(): + if hasattr(m, 'convert_to_deploy'): + m.convert_to_deploy() + return self diff --git a/engine/deim/deim_criterion.py b/engine/deim/deim_criterion.py new file mode 100644 index 0000000000000000000000000000000000000000..aef446878f438011a413cc4af1b6851b95580625 --- /dev/null +++ b/engine/deim/deim_criterion.py @@ -0,0 +1,485 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE/) +Copyright (c) 2024 D-FINE Authors. All Rights Reserved. +""" + +import torch +import torch.nn as nn +import torch.distributed +import torch.nn.functional as F +import torchvision + +import copy + +from .dfine_utils import bbox2distance +from .box_ops import box_cxcywh_to_xyxy, box_iou, generalized_box_iou +from ..misc.dist_utils import get_world_size, is_dist_available_and_initialized +from ..core import register + + +@register() +class DEIMCriterion(nn.Module): + """ This class computes the loss for DEIM. + """ + __share__ = ['num_classes', ] + __inject__ = ['matcher', ] + + def __init__(self, \ + matcher, + weight_dict, + losses, + alpha=0.2, + gamma=2.0, + num_classes=80, + reg_max=32, + boxes_weight_format=None, + share_matched_indices=False, + mal_alpha=None, + use_uni_set=True, + ): + """Create the criterion. + Parameters: + matcher: module able to compute a matching between targets and proposals. + weight_dict: dict containing as key the names of the losses and as values their relative weight. + losses: list of all the losses to be applied. See get_loss for list of available losses. + num_classes: number of object categories, omitting the special no-object category. + reg_max (int): Max number of the discrete bins in D-FINE. + boxes_weight_format: format for boxes weight (iou, ). + """ + super().__init__() + self.num_classes = num_classes + self.matcher = matcher + self.weight_dict = weight_dict + self.losses = losses + self.boxes_weight_format = boxes_weight_format + self.share_matched_indices = share_matched_indices + self.alpha = alpha + self.gamma = gamma + self.fgl_targets, self.fgl_targets_dn = None, None + self.own_targets, self.own_targets_dn = None, None + self.reg_max = reg_max + self.num_pos, self.num_neg = None, None + self.mal_alpha = mal_alpha + self.use_uni_set = use_uni_set + + def loss_labels_focal(self, outputs, targets, indices, num_boxes): + assert 'pred_logits' in outputs + src_logits = outputs['pred_logits'] + idx = self._get_src_permutation_idx(indices) + target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)]) + target_classes = torch.full(src_logits.shape[:2], self.num_classes, + dtype=torch.int64, device=src_logits.device) + target_classes[idx] = target_classes_o + target = F.one_hot(target_classes, num_classes=self.num_classes+1)[..., :-1] + loss = torchvision.ops.sigmoid_focal_loss(src_logits, target, self.alpha, self.gamma, reduction='none') + loss = loss.mean(1).sum() * src_logits.shape[1] / num_boxes + + return {'loss_focal': loss} + + def loss_labels_vfl(self, outputs, targets, indices, num_boxes, values=None): + assert 'pred_boxes' in outputs + idx = self._get_src_permutation_idx(indices) + if values is None: + src_boxes = outputs['pred_boxes'][idx] + target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0) + ious, _ = box_iou(box_cxcywh_to_xyxy(src_boxes), box_cxcywh_to_xyxy(target_boxes)) + ious = torch.diag(ious).detach() + else: + ious = values + + src_logits = outputs['pred_logits'] + target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)]) + target_classes = torch.full(src_logits.shape[:2], self.num_classes, + dtype=torch.int64, device=src_logits.device) + target_classes[idx] = target_classes_o + target = F.one_hot(target_classes, num_classes=self.num_classes + 1)[..., :-1] + + target_score_o = torch.zeros_like(target_classes, dtype=src_logits.dtype) + target_score_o[idx] = ious.to(target_score_o.dtype) + target_score = target_score_o.unsqueeze(-1) * target + + pred_score = F.sigmoid(src_logits).detach() + weight = self.alpha * pred_score.pow(self.gamma) * (1 - target) + target_score + + loss = F.binary_cross_entropy_with_logits(src_logits, target_score, weight=weight, reduction='none') + loss = loss.mean(1).sum() * src_logits.shape[1] / num_boxes + return {'loss_vfl': loss} + + def loss_labels_mal(self, outputs, targets, indices, num_boxes, values=None): + assert 'pred_boxes' in outputs + idx = self._get_src_permutation_idx(indices) + if values is None: + src_boxes = outputs['pred_boxes'][idx] + target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0) + ious, _ = box_iou(box_cxcywh_to_xyxy(src_boxes), box_cxcywh_to_xyxy(target_boxes)) + ious = torch.diag(ious).detach() + else: + ious = values + + src_logits = outputs['pred_logits'] + target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)]) + target_classes = torch.full(src_logits.shape[:2], self.num_classes, + dtype=torch.int64, device=src_logits.device) + target_classes[idx] = target_classes_o + target = F.one_hot(target_classes, num_classes=self.num_classes + 1)[..., :-1] + + target_score_o = torch.zeros_like(target_classes, dtype=src_logits.dtype) + target_score_o[idx] = ious.to(target_score_o.dtype) + target_score = target_score_o.unsqueeze(-1) * target + + pred_score = F.sigmoid(src_logits).detach() + target_score = target_score.pow(self.gamma) + if self.mal_alpha != None: + weight = self.mal_alpha * pred_score.pow(self.gamma) * (1 - target) + target + else: + weight = pred_score.pow(self.gamma) * (1 - target) + target + + # print(" ### DEIM-gamma{}-alpha{} ### ".format(self.gamma, self.mal_alpha)) + loss = F.binary_cross_entropy_with_logits(src_logits, target_score, weight=weight, reduction='none') + loss = loss.mean(1).sum() * src_logits.shape[1] / num_boxes + return {'loss_mal': loss} + + def loss_boxes(self, outputs, targets, indices, num_boxes, boxes_weight=None): + """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss + targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4] + The target boxes are expected in format (center_x, center_y, w, h), normalized by the image size. + """ + assert 'pred_boxes' in outputs + idx = self._get_src_permutation_idx(indices) + src_boxes = outputs['pred_boxes'][idx] + target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0) + losses = {} + loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none') + losses['loss_bbox'] = loss_bbox.sum() / num_boxes + + loss_giou = 1 - torch.diag(generalized_box_iou(\ + box_cxcywh_to_xyxy(src_boxes), box_cxcywh_to_xyxy(target_boxes))) + loss_giou = loss_giou if boxes_weight is None else loss_giou * boxes_weight + losses['loss_giou'] = loss_giou.sum() / num_boxes + + return losses + + def loss_local(self, outputs, targets, indices, num_boxes, T=5): + """Compute Fine-Grained Localization (FGL) Loss + and Decoupled Distillation Focal (DDF) Loss. """ + + losses = {} + if 'pred_corners' in outputs: + idx = self._get_src_permutation_idx(indices) + target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0) + + pred_corners = outputs['pred_corners'][idx].reshape(-1, (self.reg_max+1)) + ref_points = outputs['ref_points'][idx].detach() + with torch.no_grad(): + if self.fgl_targets_dn is None and 'is_dn' in outputs: + self.fgl_targets_dn= bbox2distance(ref_points, box_cxcywh_to_xyxy(target_boxes), + self.reg_max, outputs['reg_scale'], outputs['up']) + if self.fgl_targets is None and 'is_dn' not in outputs: + self.fgl_targets = bbox2distance(ref_points, box_cxcywh_to_xyxy(target_boxes), + self.reg_max, outputs['reg_scale'], outputs['up']) + + target_corners, weight_right, weight_left = self.fgl_targets_dn if 'is_dn' in outputs else self.fgl_targets + + ious = torch.diag(box_iou(\ + box_cxcywh_to_xyxy(outputs['pred_boxes'][idx]), box_cxcywh_to_xyxy(target_boxes))[0]) + weight_targets = ious.unsqueeze(-1).repeat(1, 1, 4).reshape(-1).detach() + + losses['loss_fgl'] = self.unimodal_distribution_focal_loss( + pred_corners, target_corners, weight_right, weight_left, weight_targets, avg_factor=num_boxes) + + if 'teacher_corners' in outputs: + pred_corners = outputs['pred_corners'].reshape(-1, (self.reg_max+1)) + target_corners = outputs['teacher_corners'].reshape(-1, (self.reg_max+1)) + if not torch.equal(pred_corners, target_corners): + weight_targets_local = outputs['teacher_logits'].sigmoid().max(dim=-1)[0] + + mask = torch.zeros_like(weight_targets_local, dtype=torch.bool) + mask[idx] = True + mask = mask.unsqueeze(-1).repeat(1, 1, 4).reshape(-1) + + weight_targets_local[idx] = ious.reshape_as(weight_targets_local[idx]).to(weight_targets_local.dtype) + weight_targets_local = weight_targets_local.unsqueeze(-1).repeat(1, 1, 4).reshape(-1).detach() + + loss_match_local = weight_targets_local * (T ** 2) * (nn.KLDivLoss(reduction='none') + (F.log_softmax(pred_corners / T, dim=1), F.softmax(target_corners.detach() / T, dim=1))).sum(-1) + if 'is_dn' not in outputs: + batch_scale = 8 / outputs['pred_boxes'].shape[0] # Avoid the influence of batch size per GPU + self.num_pos, self.num_neg = (mask.sum() * batch_scale) ** 0.5, ((~mask).sum() * batch_scale) ** 0.5 + loss_match_local1 = loss_match_local[mask].mean() if mask.any() else 0 + loss_match_local2 = loss_match_local[~mask].mean() if (~mask).any() else 0 + losses['loss_ddf'] = (loss_match_local1 * self.num_pos + loss_match_local2 * self.num_neg) / (self.num_pos + self.num_neg) + + return losses + + def _get_src_permutation_idx(self, indices): + # permute predictions following indices + batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)]) + src_idx = torch.cat([src for (src, _) in indices]) + return batch_idx, src_idx + + def _get_tgt_permutation_idx(self, indices): + # permute targets following indices + batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)]) + tgt_idx = torch.cat([tgt for (_, tgt) in indices]) + return batch_idx, tgt_idx + + def _get_go_indices(self, indices, indices_aux_list): + """Get a matching union set across all decoder layers. """ + results = [] + for indices_aux in indices_aux_list: + indices = [(torch.cat([idx1[0], idx2[0]]), torch.cat([idx1[1], idx2[1]])) + for idx1, idx2 in zip(indices.copy(), indices_aux.copy())] + + for ind in [torch.cat([idx[0][:, None], idx[1][:, None]], 1) for idx in indices]: + unique, counts = torch.unique(ind, return_counts=True, dim=0) + count_sort_indices = torch.argsort(counts, descending=True) + unique_sorted = unique[count_sort_indices] + column_to_row = {} + for idx in unique_sorted: + row_idx, col_idx = idx[0].item(), idx[1].item() + if row_idx not in column_to_row: + column_to_row[row_idx] = col_idx + final_rows = torch.tensor(list(column_to_row.keys()), device=ind.device) + final_cols = torch.tensor(list(column_to_row.values()), device=ind.device) + results.append((final_rows.long(), final_cols.long())) + return results + + def _clear_cache(self): + self.fgl_targets, self.fgl_targets_dn = None, None + self.own_targets, self.own_targets_dn = None, None + self.num_pos, self.num_neg = None, None + + def get_loss(self, loss, outputs, targets, indices, num_boxes, **kwargs): + loss_map = { + 'boxes': self.loss_boxes, + 'focal': self.loss_labels_focal, + 'vfl': self.loss_labels_vfl, + 'mal': self.loss_labels_mal, + 'local': self.loss_local, + } + assert loss in loss_map, f'do you really want to compute {loss} loss?' + return loss_map[loss](outputs, targets, indices, num_boxes, **kwargs) + + def forward(self, outputs, targets, epoch=0, **kwargs): + """ This performs the loss computation. + Parameters: + outputs: dict of tensors, see the output specification of the model for the format + targets: list of dicts, such that len(targets) == batch_size. + The expected keys in each dict depends on the losses applied, see each loss' doc + """ + outputs_without_aux = {k: v for k, v in outputs.items() if 'aux' not in k} + + # Retrieve the matching between the outputs of the last layer and the targets + indices = self.matcher(outputs_without_aux, targets, epoch=epoch)['indices'] + self._clear_cache() + + # Get the matching union set across all decoder layers. + if 'aux_outputs' in outputs: + indices_aux_list, cached_indices, cached_indices_enc = [], [], [] + aux_outputs_list = outputs['aux_outputs'] + if 'pre_outputs' in outputs: + aux_outputs_list = outputs['aux_outputs'] + [outputs['pre_outputs']] + for i, aux_outputs in enumerate(aux_outputs_list): + indices_aux = self.matcher(aux_outputs, targets, epoch=epoch)['indices'] + cached_indices.append(indices_aux) + indices_aux_list.append(indices_aux) + for i, aux_outputs in enumerate(outputs['enc_aux_outputs']): + indices_enc = self.matcher(aux_outputs, targets, epoch=epoch)['indices'] + cached_indices_enc.append(indices_enc) + indices_aux_list.append(indices_enc) + indices_go = self._get_go_indices(indices, indices_aux_list) + + num_boxes_go = sum(len(x[0]) for x in indices_go) + num_boxes_go = torch.as_tensor([num_boxes_go], dtype=torch.float, device=next(iter(outputs.values())).device) + if is_dist_available_and_initialized(): + torch.distributed.all_reduce(num_boxes_go) + num_boxes_go = torch.clamp(num_boxes_go / get_world_size(), min=1).item() + else: + assert 'aux_outputs' in outputs, '' + + # Compute the average number of target boxes accross all nodes, for normalization purposes + num_boxes = sum(len(t["labels"]) for t in targets) + num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device) + if is_dist_available_and_initialized(): + torch.distributed.all_reduce(num_boxes) + num_boxes = torch.clamp(num_boxes / get_world_size(), min=1).item() + + # Compute all the requested losses, main loss + losses = {} + for loss in self.losses: + use_uni_set = self.use_uni_set and (loss in ['boxes', 'local']) + indices_in = indices_go if use_uni_set else indices + num_boxes_in = num_boxes_go if use_uni_set else num_boxes + meta = self.get_loss_meta_info(loss, outputs, targets, indices_in) + l_dict = self.get_loss(loss, outputs, targets, indices_in, num_boxes_in, **meta) + l_dict = {k: l_dict[k] * self.weight_dict[k] for k in l_dict if k in self.weight_dict} + losses.update(l_dict) + + # In case of auxiliary losses, we repeat this process with the output of each intermediate layer. + if 'aux_outputs' in outputs: + for i, aux_outputs in enumerate(outputs['aux_outputs']): + if 'local' in self.losses: # only work for local loss + aux_outputs['up'], aux_outputs['reg_scale'] = outputs['up'], outputs['reg_scale'] + for loss in self.losses: + use_uni_set = self.use_uni_set and (loss in ['boxes', 'local']) + indices_in = indices_go if use_uni_set else cached_indices[i] + num_boxes_in = num_boxes_go if use_uni_set else num_boxes + meta = self.get_loss_meta_info(loss, aux_outputs, targets, indices_in) + l_dict = self.get_loss(loss, aux_outputs, targets, indices_in, num_boxes_in, **meta) + + l_dict = {k: l_dict[k] * self.weight_dict[k] for k in l_dict if k in self.weight_dict} + l_dict = {k + f'_aux_{i}': v for k, v in l_dict.items()} + losses.update(l_dict) + + # In case of auxiliary traditional head output at first decoder layer. just for dfine + if 'pre_outputs' in outputs: + aux_outputs = outputs['pre_outputs'] + for loss in self.losses: + use_uni_set = self.use_uni_set and (loss in ['boxes', 'local']) + indices_in = indices_go if use_uni_set else cached_indices[-1] + num_boxes_in = num_boxes_go if use_uni_set else num_boxes + meta = self.get_loss_meta_info(loss, aux_outputs, targets, indices_in) + l_dict = self.get_loss(loss, aux_outputs, targets, indices_in, num_boxes_in, **meta) + + l_dict = {k: l_dict[k] * self.weight_dict[k] for k in l_dict if k in self.weight_dict} + l_dict = {k + '_pre': v for k, v in l_dict.items()} + losses.update(l_dict) + + # In case of encoder auxiliary losses. + if 'enc_aux_outputs' in outputs: + assert 'enc_meta' in outputs, '' + class_agnostic = outputs['enc_meta']['class_agnostic'] + if class_agnostic: + orig_num_classes = self.num_classes + self.num_classes = 1 + enc_targets = copy.deepcopy(targets) + for t in enc_targets: + t['labels'] = torch.zeros_like(t["labels"]) + else: + enc_targets = targets + + for i, aux_outputs in enumerate(outputs['enc_aux_outputs']): + for loss in self.losses: + use_uni_set = self.use_uni_set and (loss == 'boxes') + indices_in = indices_go if use_uni_set else cached_indices_enc[i] + num_boxes_in = num_boxes_go if use_uni_set else num_boxes + meta = self.get_loss_meta_info(loss, aux_outputs, enc_targets, indices_in) + l_dict = self.get_loss(loss, aux_outputs, enc_targets, indices_in, num_boxes_in, **meta) + l_dict = {k: l_dict[k] * self.weight_dict[k] for k in l_dict if k in self.weight_dict} + l_dict = {k + f'_enc_{i}': v for k, v in l_dict.items()} + losses.update(l_dict) + + if class_agnostic: + self.num_classes = orig_num_classes + + # In case of cdn auxiliary losses. + if 'dn_outputs' in outputs: + assert 'dn_meta' in outputs, '' + indices_dn = self.get_cdn_matched_indices(outputs['dn_meta'], targets) + dn_num_boxes = num_boxes * outputs['dn_meta']['dn_num_group'] + + for i, aux_outputs in enumerate(outputs['dn_outputs']): + if 'local' in self.losses: # only work for local loss + aux_outputs['is_dn'] = True + aux_outputs['up'], aux_outputs['reg_scale'] = outputs['up'], outputs['reg_scale'] + for loss in self.losses: + meta = self.get_loss_meta_info(loss, aux_outputs, targets, indices_dn) + l_dict = self.get_loss(loss, aux_outputs, targets, indices_dn, dn_num_boxes, **meta) + l_dict = {k: l_dict[k] * self.weight_dict[k] for k in l_dict if k in self.weight_dict} + l_dict = {k + f'_dn_{i}': v for k, v in l_dict.items()} + losses.update(l_dict) + + # In case of auxiliary traditional head output at first decoder layer, just for dfine + if 'dn_pre_outputs' in outputs: + aux_outputs = outputs['dn_pre_outputs'] + for loss in self.losses: + meta = self.get_loss_meta_info(loss, aux_outputs, targets, indices_dn) + l_dict = self.get_loss(loss, aux_outputs, targets, indices_dn, dn_num_boxes, **meta) + l_dict = {k: l_dict[k] * self.weight_dict[k] for k in l_dict if k in self.weight_dict} + l_dict = {k + '_dn_pre': v for k, v in l_dict.items()} + losses.update(l_dict) + + # For debugging Objects365 pre-train. + losses = {k:torch.nan_to_num(v, nan=0.0) for k, v in losses.items()} + return losses + + def get_loss_meta_info(self, loss, outputs, targets, indices): + if self.boxes_weight_format is None: + return {} + + src_boxes = outputs['pred_boxes'][self._get_src_permutation_idx(indices)] + target_boxes = torch.cat([t['boxes'][j] for t, (_, j) in zip(targets, indices)], dim=0) + + if self.boxes_weight_format == 'iou': + iou, _ = box_iou(box_cxcywh_to_xyxy(src_boxes.detach()), box_cxcywh_to_xyxy(target_boxes)) + iou = torch.diag(iou) + elif self.boxes_weight_format == 'giou': + iou = torch.diag(generalized_box_iou(\ + box_cxcywh_to_xyxy(src_boxes.detach()), box_cxcywh_to_xyxy(target_boxes))) + else: + raise AttributeError() + + if loss in ('boxes', ): + meta = {'boxes_weight': iou} + elif loss in ('vfl', 'mal'): + meta = {'values': iou} + else: + meta = {} + + return meta + + @staticmethod + def get_cdn_matched_indices(dn_meta, targets): + """get_cdn_matched_indices + """ + dn_positive_idx, dn_num_group = dn_meta["dn_positive_idx"], dn_meta["dn_num_group"] + num_gts = [len(t['labels']) for t in targets] + device = targets[0]['labels'].device + + dn_match_indices = [] + for i, num_gt in enumerate(num_gts): + if num_gt > 0: + gt_idx = torch.arange(num_gt, dtype=torch.int64, device=device) + gt_idx = gt_idx.tile(dn_num_group) + assert len(dn_positive_idx[i]) == len(gt_idx) + dn_match_indices.append((dn_positive_idx[i], gt_idx)) + else: + dn_match_indices.append((torch.zeros(0, dtype=torch.int64, device=device), \ + torch.zeros(0, dtype=torch.int64, device=device))) + + return dn_match_indices + + + def feature_loss_function(self, fea, target_fea): + loss = (fea - target_fea) ** 2 * ((fea > 0) | (target_fea > 0)).float() + return torch.abs(loss) + + + def unimodal_distribution_focal_loss(self, pred, label, weight_right, weight_left, weight=None, reduction='sum', avg_factor=None): + dis_left = label.long() + dis_right = dis_left + 1 + + loss = F.cross_entropy(pred, dis_left, reduction='none') * weight_left.reshape(-1) \ + + F.cross_entropy(pred, dis_right, reduction='none') * weight_right.reshape(-1) + + if weight is not None: + weight = weight.float() + loss = loss * weight + + if avg_factor is not None: + loss = loss.sum() / avg_factor + elif reduction == 'mean': + loss = loss.mean() + elif reduction == 'sum': + loss = loss.sum() + + return loss + + def get_gradual_steps(self, outputs): + num_layers = len(outputs['aux_outputs']) + 1 if 'aux_outputs' in outputs else 1 + step = .5 / (num_layers - 1) + opt_list = [.5 + step * i for i in range(num_layers)] if num_layers > 1 else [1] + return opt_list diff --git a/engine/deim/deim_decoder.py b/engine/deim/deim_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..93cbbc53617db0043c2f0ae07aa83b762b44312b --- /dev/null +++ b/engine/deim/deim_decoder.py @@ -0,0 +1,611 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE/) +Copyright (c) 2024 D-FINE Authors. All Rights Reserved. +""" + +import math +import copy +import functools +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.nn.init as init +from typing import List + +from ..core import register +from .denoising import get_contrastive_denoising_training_group +from .utils import deformable_attention_core_func_v2, get_activation, inverse_sigmoid, bias_init_with_prob + +from .dfine_decoder import MSDeformableAttention, LQE, Integral +from .dfine_utils import weighting_function, distance2bbox +from .deim_utils import RMSNorm, SwiGLUFFN, Gate, MLP + +__all__ = ['DEIMTransformer'] + + +class TransformerDecoderLayer(nn.Module): + def __init__(self, + d_model=256, + n_head=8, + dim_feedforward=1024, + dropout=0., + activation='relu', + n_levels=4, + n_points=4, + cross_attn_method='default', + layer_scale=None, + use_gateway=False, + ): + super(TransformerDecoderLayer, self).__init__() + + if layer_scale is not None: + print(f" --- Wide Layer@{layer_scale} ---") + dim_feedforward = round(layer_scale * dim_feedforward) + d_model = round(layer_scale * d_model) + + # self attention + self.self_attn = nn.MultiheadAttention(d_model, n_head, dropout=dropout, batch_first=True) + self.dropout1 = nn.Dropout(dropout) + self.norm1 = RMSNorm(d_model) + + # cross attention + self.cross_attn = MSDeformableAttention(d_model, n_head, n_levels, n_points, method=cross_attn_method) + self.dropout2 = nn.Dropout(dropout) + + self.use_gateway = use_gateway + if use_gateway: + self.gateway = Gate(d_model, use_rmsnorm=True) + else: + self.norm2 = RMSNorm(d_model) + + # ffn + self.swish_ffn = SwiGLUFFN(d_model, dim_feedforward // 2, d_model) + self.dropout4 = nn.Dropout(dropout) + self.norm3 = RMSNorm(d_model) + + def with_pos_embed(self, tensor, pos): + return tensor if pos is None else tensor + pos + + def forward(self, + target, + reference_points, + value, + spatial_shapes, + attn_mask=None, + query_pos_embed=None): + + # self attention + q = k = self.with_pos_embed(target, query_pos_embed) + + target2, _ = self.self_attn(q, k, value=target, attn_mask=attn_mask) + target = target + self.dropout1(target2) + target = self.norm1(target) + + # cross attention + target2 = self.cross_attn(\ + self.with_pos_embed(target, query_pos_embed), + reference_points, + value, + spatial_shapes) + + if self.use_gateway: + target = self.gateway(target, self.dropout2(target2)) + else: + target = target + self.dropout2(target2) + target = self.norm2(target) + + # ffn + target2 = self.swish_ffn(target) + target = target + self.dropout4(target2) + target = self.norm3(target.clamp(min=-65504, max=65504)) + + return target + + +class TransformerDecoder(nn.Module): + """ + Transformer Decoder implementing Fine-grained Distribution Refinement (FDR). + + This decoder refines object detection predictions through iterative updates across multiple layers, + utilizing attention mechanisms, location quality estimators, and distribution refinement techniques + to improve bounding box accuracy and robustness. + """ + + def __init__(self, hidden_dim, decoder_layer, decoder_layer_wide, num_layers, num_head, reg_max, reg_scale, up, + eval_idx=-1, layer_scale=2, act='relu'): + super(TransformerDecoder, self).__init__() + self.hidden_dim = hidden_dim + self.num_layers = num_layers + self.layer_scale = layer_scale + self.num_head = num_head + self.eval_idx = eval_idx if eval_idx >= 0 else num_layers + eval_idx + self.up, self.reg_scale, self.reg_max = up, reg_scale, reg_max + self.layers = nn.ModuleList([copy.deepcopy(decoder_layer) for _ in range(self.eval_idx + 1)] \ + + [copy.deepcopy(decoder_layer_wide) for _ in range(num_layers - self.eval_idx - 1)]) + self.lqe_layers = nn.ModuleList([copy.deepcopy(LQE(4, 64, 2, reg_max, act=act)) for _ in range(num_layers)]) + + def value_op(self, memory, value_proj, value_scale, memory_mask, memory_spatial_shapes): + """ + Preprocess values for MSDeformableAttention. + """ + value = value_proj(memory) if value_proj is not None else memory + value = F.interpolate(memory, size=value_scale) if value_scale is not None else value + if memory_mask is not None: + value = value * memory_mask.to(value.dtype).unsqueeze(-1) + value = value.reshape(value.shape[0], value.shape[1], self.num_head, -1) + split_shape = [h * w for h, w in memory_spatial_shapes] + return value.permute(0, 2, 3, 1).split(split_shape, dim=-1) + + def convert_to_deploy(self): + self.project = weighting_function(self.reg_max, self.up, self.reg_scale, deploy=True) + self.layers = self.layers[:self.eval_idx + 1] + self.lqe_layers = nn.ModuleList([nn.Identity()] * (self.eval_idx) + [self.lqe_layers[self.eval_idx]]) + + def forward(self, + target, + ref_points_unact, + memory, + spatial_shapes, + bbox_head, + score_head, + query_pos_head, + pre_bbox_head, + integral, + up, + reg_scale, + attn_mask=None, + memory_mask=None, + dn_meta=None): + output = target + output_detach = pred_corners_undetach = 0 + value = self.value_op(memory, None, None, memory_mask, spatial_shapes) + + dec_out_bboxes = [] + dec_out_logits = [] + dec_out_pred_corners = [] + dec_out_refs = [] + if not hasattr(self, 'project'): + project = weighting_function(self.reg_max, up, reg_scale) + else: + project = self.project + + ref_points_detach = F.sigmoid(ref_points_unact) + query_pos_embed = query_pos_head(ref_points_detach).clamp(min=-10, max=10) + + for i, layer in enumerate(self.layers): + ref_points_input = ref_points_detach.unsqueeze(2) + + if i >= self.eval_idx + 1 and self.layer_scale > 1: + query_pos_embed = F.interpolate(query_pos_embed, scale_factor=self.layer_scale) + value = self.value_op(memory, None, query_pos_embed.shape[-1], memory_mask, spatial_shapes) + output = F.interpolate(output, size=query_pos_embed.shape[-1]) + output_detach = output.detach() + + output = layer(output, ref_points_input, value, spatial_shapes, attn_mask, query_pos_embed) + + if i == 0 : + # Initial bounding box predictions with inverse sigmoid refinement + pre_bboxes = F.sigmoid(pre_bbox_head(output) + inverse_sigmoid(ref_points_detach)) + pre_scores = score_head[0](output) + ref_points_initial = pre_bboxes.detach() + + # Refine bounding box corners using FDR, integrating previous layer's corrections + pred_corners = bbox_head[i](output + output_detach) + pred_corners_undetach + inter_ref_bbox = distance2bbox(ref_points_initial, integral(pred_corners, project), reg_scale) + + if self.training or i == self.eval_idx: + scores = score_head[i](output) + # Lqe does not affect the performance here. + scores = self.lqe_layers[i](scores, pred_corners) + dec_out_logits.append(scores) + dec_out_bboxes.append(inter_ref_bbox) + dec_out_pred_corners.append(pred_corners) + dec_out_refs.append(ref_points_initial) + + if not self.training: + break + + pred_corners_undetach = pred_corners + ref_points_detach = inter_ref_bbox.detach() + output_detach = output.detach() + + return torch.stack(dec_out_bboxes), torch.stack(dec_out_logits), \ + torch.stack(dec_out_pred_corners), torch.stack(dec_out_refs), pre_bboxes, pre_scores + + +@register() +class DEIMTransformer(nn.Module): + __share__ = ['num_classes', 'eval_spatial_size'] + + def __init__(self, + num_classes=80, + hidden_dim=256, + num_queries=300, + feat_channels=[512, 1024, 2048], + feat_strides=[8, 16, 32], + num_levels=3, + num_points=4, + nhead=8, + num_layers=6, + dim_feedforward=1024, + dropout=0., + activation="relu", + num_denoising=100, + label_noise_ratio=0.5, + box_noise_scale=1.0, + learn_query_content=False, + eval_spatial_size=None, + eval_idx=-1, + eps=1e-2, + aux_loss=True, + cross_attn_method='default', + query_select_method='default', + reg_max=32, + reg_scale=4., + layer_scale=1, + mlp_act='relu', + use_gateway=True, + share_bbox_head=False, + share_score_head=False, + ): + super().__init__() + assert len(feat_channels) <= num_levels + assert len(feat_strides) == len(feat_channels) + + for _ in range(num_levels - len(feat_strides)): + feat_strides.append(feat_strides[-1] * 2) + + self.hidden_dim = hidden_dim + scaled_dim = round(layer_scale*hidden_dim) + self.nhead = nhead + self.feat_strides = feat_strides + self.num_levels = num_levels + self.num_classes = num_classes + self.num_queries = num_queries + self.eps = eps + self.num_layers = num_layers + self.eval_spatial_size = eval_spatial_size + self.aux_loss = aux_loss + self.reg_max = reg_max + + assert query_select_method in ('default', 'one2many', 'agnostic'), '' + assert cross_attn_method in ('default', 'discrete'), '' + self.cross_attn_method = cross_attn_method + self.query_select_method = query_select_method + # -- print the parameters + print(f" --- Use Gateway@{use_gateway} ---") + print(f" --- Use Share Bbox Head@{share_bbox_head} ---") + print(f" --- Use Share Score Head@{share_score_head} ---") + + # backbone feature projection + self._build_input_proj_layer(feat_channels) + + # Transformer module + self.up = nn.Parameter(torch.tensor([0.5]), requires_grad=False) + self.reg_scale = nn.Parameter(torch.tensor([reg_scale]), requires_grad=False) + decoder_layer = TransformerDecoderLayer(hidden_dim, nhead, dim_feedforward, dropout, \ + activation, num_levels, num_points, cross_attn_method=cross_attn_method, use_gateway=use_gateway) + decoder_layer_wide = TransformerDecoderLayer(hidden_dim, nhead, dim_feedforward, dropout, \ + activation, num_levels, num_points, cross_attn_method=cross_attn_method, layer_scale=layer_scale, use_gateway=use_gateway) + self.decoder = TransformerDecoder(hidden_dim, decoder_layer, decoder_layer_wide, num_layers, nhead, + reg_max, self.reg_scale, self.up, eval_idx, layer_scale, act=activation) + # denoising + self.num_denoising = num_denoising + self.label_noise_ratio = label_noise_ratio + self.box_noise_scale = box_noise_scale + if num_denoising > 0: + self.denoising_class_embed = nn.Embedding(num_classes+1, hidden_dim, padding_idx=num_classes) + init.normal_(self.denoising_class_embed.weight[:-1]) + + # decoder embedding + self.learn_query_content = learn_query_content + if learn_query_content: + self.tgt_embed = nn.Embedding(num_queries, hidden_dim) + + if query_select_method == 'agnostic': + self.enc_score_head = nn.Linear(hidden_dim, 1) + else: + self.enc_score_head = nn.Linear(hidden_dim, num_classes) + self.enc_bbox_head = MLP(hidden_dim, hidden_dim, 4, 3, act=mlp_act) + + self.query_pos_head = MLP(4, hidden_dim, hidden_dim, 3, act=mlp_act) + + # decoder head + self.pre_bbox_head = MLP(hidden_dim, hidden_dim, 4, 3, act=mlp_act) + self.integral = Integral(self.reg_max) + + self.eval_idx = eval_idx if eval_idx >= 0 else num_layers + eval_idx + dec_score_head = nn.Linear(hidden_dim, num_classes) + self.dec_score_head = nn.ModuleList( + [dec_score_head if share_score_head else copy.deepcopy(dec_score_head) for _ in range(self.eval_idx + 1)] + + [copy.deepcopy(dec_score_head) for _ in range(num_layers - self.eval_idx - 1)]) + + # Share the same bbox head for all layers + dec_bbox_head = MLP(hidden_dim, hidden_dim, 4 * (self.reg_max+1), 3, act=mlp_act) + self.dec_bbox_head = nn.ModuleList( + [dec_bbox_head if share_bbox_head else copy.deepcopy(dec_bbox_head) for _ in range(self.eval_idx + 1)] + + [MLP(scaled_dim, scaled_dim, 4 * (self.reg_max+1), 3, act=mlp_act) for _ in range(num_layers - self.eval_idx - 1)]) + + # init encoder output anchors and valid_mask + if self.eval_spatial_size: + anchors, valid_mask = self._generate_anchors() + self.register_buffer('anchors', anchors) + self.register_buffer('valid_mask', valid_mask) + # init encoder output anchors and valid_mask + if self.eval_spatial_size: + self.anchors, self.valid_mask = self._generate_anchors() + + + self._reset_parameters(feat_channels) + + def convert_to_deploy(self): + self.dec_score_head = nn.ModuleList([nn.Identity()] * (self.eval_idx) + [self.dec_score_head[self.eval_idx]]) + self.dec_bbox_head = nn.ModuleList( + [self.dec_bbox_head[i] if i <= self.eval_idx else nn.Identity() for i in range(len(self.dec_bbox_head))] + ) + + def _reset_parameters(self, feat_channels): + bias = bias_init_with_prob(0.01) + init.constant_(self.enc_score_head.bias, bias) + init.constant_(self.enc_bbox_head.layers[-1].weight, 0) + init.constant_(self.enc_bbox_head.layers[-1].bias, 0) + + init.constant_(self.pre_bbox_head.layers[-1].weight, 0) + init.constant_(self.pre_bbox_head.layers[-1].bias, 0) + + for cls_, reg_ in zip(self.dec_score_head, self.dec_bbox_head): + init.constant_(cls_.bias, bias) + if hasattr(reg_, 'layers'): + init.constant_(reg_.layers[-1].weight, 0) + init.constant_(reg_.layers[-1].bias, 0) + + if self.learn_query_content: + init.xavier_uniform_(self.tgt_embed.weight) + init.xavier_uniform_(self.query_pos_head.layers[0].weight) + init.xavier_uniform_(self.query_pos_head.layers[1].weight) + init.xavier_uniform_(self.query_pos_head.layers[-1].weight) + for m, in_channels in zip(self.input_proj, feat_channels): + if in_channels != self.hidden_dim: + init.xavier_uniform_(m[0].weight) + + def _build_input_proj_layer(self, feat_channels): + self.input_proj = nn.ModuleList() + for in_channels in feat_channels: + if in_channels == self.hidden_dim: + self.input_proj.append(nn.Identity()) + else: + self.input_proj.append( + nn.Sequential(OrderedDict([ + ('conv', nn.Conv2d(in_channels, self.hidden_dim, 1, bias=False)), + ('norm', nn.BatchNorm2d(self.hidden_dim,))]) + ) + ) + + in_channels = feat_channels[-1] + + for _ in range(self.num_levels - len(feat_channels)): + if in_channels == self.hidden_dim: + self.input_proj.append(nn.Identity()) + else: + self.input_proj.append( + nn.Sequential(OrderedDict([ + ('conv', nn.Conv2d(in_channels, self.hidden_dim, 3, 2, padding=1, bias=False)), + ('norm', nn.BatchNorm2d(self.hidden_dim))]) + ) + ) + in_channels = self.hidden_dim + + def _get_encoder_input(self, feats: List[torch.Tensor]): + # get projection features + proj_feats = [self.input_proj[i](feat) for i, feat in enumerate(feats)] + if self.num_levels > len(proj_feats): + len_srcs = len(proj_feats) + for i in range(len_srcs, self.num_levels): + if i == len_srcs: + proj_feats.append(self.input_proj[i](feats[-1])) + else: + proj_feats.append(self.input_proj[i](proj_feats[-1])) + + # get encoder inputs + feat_flatten = [] + spatial_shapes = [] + for i, feat in enumerate(proj_feats): + _, _, h, w = feat.shape + # [b, c, h, w] -> [b, h*w, c] + feat_flatten.append(feat.flatten(2).permute(0, 2, 1)) + # [num_levels, 2] + spatial_shapes.append([h, w]) + + # [b, l, c] + feat_flatten = torch.concat(feat_flatten, 1) + return feat_flatten, spatial_shapes + + def _generate_anchors(self, + spatial_shapes=None, + grid_size=0.05, + dtype=torch.float32, + device='cpu'): + if spatial_shapes is None: + spatial_shapes = [] + eval_h, eval_w = self.eval_spatial_size + for s in self.feat_strides: + spatial_shapes.append([int(eval_h / s), int(eval_w / s)]) + + anchors = [] + for lvl, (h, w) in enumerate(spatial_shapes): + grid_y, grid_x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing='ij') + grid_xy = torch.stack([grid_x, grid_y], dim=-1) + grid_xy = (grid_xy.unsqueeze(0) + 0.5) / torch.tensor([w, h], dtype=dtype) + wh = torch.ones_like(grid_xy) * grid_size * (2.0 ** lvl) + lvl_anchors = torch.concat([grid_xy, wh], dim=-1).reshape(-1, h * w, 4) + anchors.append(lvl_anchors) + + anchors = torch.concat(anchors, dim=1).to(device) + valid_mask = ((anchors > self.eps) * (anchors < 1 - self.eps)).all(-1, keepdim=True) + anchors = torch.log(anchors / (1 - anchors)) + anchors = torch.where(valid_mask, anchors, torch.inf) + + return anchors, valid_mask + + + def _get_decoder_input(self, + memory: torch.Tensor, + spatial_shapes, + denoising_logits=None, + denoising_bbox_unact=None): + + # prepare input for decoder + if self.training or self.eval_spatial_size is None: + anchors, valid_mask = self._generate_anchors(spatial_shapes, device=memory.device) + else: + anchors = self.anchors + valid_mask = self.valid_mask + if memory.shape[0] > 1: + anchors = anchors.repeat(memory.shape[0], 1, 1) + + # memory = torch.where(valid_mask, memory, 0) + memory = valid_mask.to(memory.dtype) * memory + + enc_outputs_logits :torch.Tensor = self.enc_score_head(memory) + + # select topk queries + enc_topk_memory, enc_topk_logits, enc_topk_anchors = \ + self._select_topk(memory, enc_outputs_logits, anchors, self.num_queries) + + enc_topk_bbox_unact :torch.Tensor = self.enc_bbox_head(enc_topk_memory) + enc_topk_anchors + + enc_topk_bboxes_list, enc_topk_logits_list = [], [] + if self.training: + enc_topk_bboxes = F.sigmoid(enc_topk_bbox_unact) + enc_topk_bboxes_list.append(enc_topk_bboxes) + enc_topk_logits_list.append(enc_topk_logits) + + if self.learn_query_content: + content = self.tgt_embed.weight.unsqueeze(0).tile([memory.shape[0], 1, 1]) + else: + content = enc_topk_memory.detach() + + enc_topk_bbox_unact = enc_topk_bbox_unact.detach() + + if denoising_bbox_unact is not None: + enc_topk_bbox_unact = torch.concat([denoising_bbox_unact, enc_topk_bbox_unact], dim=1) + content = torch.concat([denoising_logits, content], dim=1) + + return content, enc_topk_bbox_unact, enc_topk_bboxes_list, enc_topk_logits_list + + def _select_topk(self, memory: torch.Tensor, outputs_logits: torch.Tensor, outputs_anchors_unact: torch.Tensor, topk: int): + if self.query_select_method == 'default': + _, topk_ind = torch.topk(outputs_logits.max(-1).values, topk, dim=-1) + + elif self.query_select_method == 'one2many': + _, topk_ind = torch.topk(outputs_logits.flatten(1), topk, dim=-1) + topk_ind = topk_ind // self.num_classes + + elif self.query_select_method == 'agnostic': + _, topk_ind = torch.topk(outputs_logits.squeeze(-1), topk, dim=-1) + + topk_ind: torch.Tensor + + topk_anchors = outputs_anchors_unact.gather(dim=1, \ + index=topk_ind.unsqueeze(-1).repeat(1, 1, outputs_anchors_unact.shape[-1])) + + topk_logits = outputs_logits.gather(dim=1, \ + index=topk_ind.unsqueeze(-1).repeat(1, 1, outputs_logits.shape[-1])) if self.training else None + + topk_memory = memory.gather(dim=1, \ + index=topk_ind.unsqueeze(-1).repeat(1, 1, memory.shape[-1])) + + return topk_memory, topk_logits, topk_anchors + + def forward(self, feats, targets=None): + # input projection and embedding + memory, spatial_shapes = self._get_encoder_input(feats) + + # prepare denoising training + if self.training and self.num_denoising > 0: + denoising_logits, denoising_bbox_unact, attn_mask, dn_meta = \ + get_contrastive_denoising_training_group(targets, \ + self.num_classes, + self.num_queries, + self.denoising_class_embed, + num_denoising=self.num_denoising, + label_noise_ratio=self.label_noise_ratio, + box_noise_scale=1.0, + ) + else: + denoising_logits, denoising_bbox_unact, attn_mask, dn_meta = None, None, None, None + + init_ref_contents, init_ref_points_unact, enc_topk_bboxes_list, enc_topk_logits_list = \ + self._get_decoder_input(memory, spatial_shapes, denoising_logits, denoising_bbox_unact) + + # decoder + out_bboxes, out_logits, out_corners, out_refs, pre_bboxes, pre_logits = self.decoder( + init_ref_contents, + init_ref_points_unact, + memory, + spatial_shapes, + self.dec_bbox_head, + self.dec_score_head, + self.query_pos_head, + self.pre_bbox_head, + self.integral, + self.up, + self.reg_scale, + attn_mask=attn_mask, + dn_meta=dn_meta) + + if self.training and dn_meta is not None: + # the output from the first decoder layer, only one + dn_pre_logits, pre_logits = torch.split(pre_logits, dn_meta['dn_num_split'], dim=1) + dn_pre_bboxes, pre_bboxes = torch.split(pre_bboxes, dn_meta['dn_num_split'], dim=1) + + dn_out_logits, out_logits = torch.split(out_logits, dn_meta['dn_num_split'], dim=2) + dn_out_bboxes, out_bboxes = torch.split(out_bboxes, dn_meta['dn_num_split'], dim=2) + + dn_out_corners, out_corners = torch.split(out_corners, dn_meta['dn_num_split'], dim=2) + dn_out_refs, out_refs = torch.split(out_refs, dn_meta['dn_num_split'], dim=2) + + if self.training: + out = {'pred_logits': out_logits[-1], 'pred_boxes': out_bboxes[-1], 'pred_corners': out_corners[-1], + 'ref_points': out_refs[-1], 'up': self.up, 'reg_scale': self.reg_scale} + else: + out = {'pred_logits': out_logits[-1], 'pred_boxes': out_bboxes[-1]} + + if self.training and self.aux_loss: + out['aux_outputs'] = self._set_aux_loss2(out_logits[:-1], out_bboxes[:-1], out_corners[:-1], out_refs[:-1], + out_corners[-1], out_logits[-1]) + out['enc_aux_outputs'] = self._set_aux_loss(enc_topk_logits_list, enc_topk_bboxes_list) + out['pre_outputs'] = {'pred_logits': pre_logits, 'pred_boxes': pre_bboxes} + out['enc_meta'] = {'class_agnostic': self.query_select_method == 'agnostic'} + + if dn_meta is not None: + out['dn_outputs'] = self._set_aux_loss2(dn_out_logits, dn_out_bboxes, dn_out_corners, dn_out_refs, + dn_out_corners[-1], dn_out_logits[-1]) + out['dn_pre_outputs'] = {'pred_logits': dn_pre_logits, 'pred_boxes': dn_pre_bboxes} + out['dn_meta'] = dn_meta + + return out + + + @torch.jit.unused + def _set_aux_loss(self, outputs_class, outputs_coord): + # this is a workaround to make torchscript happy, as torchscript + # doesn't support dictionary with non-homogeneous values, such + # as a dict having both a Tensor and a list. + return [{'pred_logits': a, 'pred_boxes': b} for a, b in zip(outputs_class, outputs_coord)] + + + @torch.jit.unused + def _set_aux_loss2(self, outputs_class, outputs_coord, outputs_corners, outputs_ref, + teacher_corners=None, teacher_logits=None): + # this is a workaround to make torchscript happy, as torchscript + # doesn't support dictionary with non-homogeneous values, such + # as a dict having both a Tensor and a list. + return [{'pred_logits': a, 'pred_boxes': b, 'pred_corners': c, 'ref_points': d, + 'teacher_corners': teacher_corners, 'teacher_logits': teacher_logits} + for a, b, c, d in zip(outputs_class, outputs_coord, outputs_corners, outputs_ref)] \ No newline at end of file diff --git a/engine/deim/deim_utils.py b/engine/deim/deim_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..292e7a3a17c7fe7ab6e676217a912326adc1d3ce --- /dev/null +++ b/engine/deim/deim_utils.py @@ -0,0 +1,83 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.nn.init as init + +from .utils import get_activation, bias_init_with_prob + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.scale = nn.Parameter(torch.ones(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x.float()).type_as(x) + output = output * self.scale + return output + + def extra_repr(self) -> str: + return f'dim={self.dim}, eps={self.eps}' + +# default 3-layer MLP +class MLP(nn.Module): + def __init__(self, input_dim, hidden_dim, output_dim, num_layers=3, act='relu'): + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) + self.act = get_activation(act) + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = self.act(layer(x)) if i < self.num_layers - 1 else layer(x) + return x + +# Taken from: https://github.com/facebookresearch/dinov2/blob/main/dinov2/layers/swiglu_ffn.py#L14-L34 +class SwiGLUFFN(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: int, + out_features: int, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + self._reset_parameters() + + def _reset_parameters(self): + init.xavier_uniform_(self.w12.weight) + init.constant_(self.w12.bias, 0) + init.xavier_uniform_(self.w3.weight) + init.constant_(self.w3.bias, 0) + + def forward(self, x): + x12 = self.w12(x) + x1, x2 = x12.chunk(2, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) + + +class Gate(nn.Module): + def __init__(self, d_model, use_rmsnorm=False): + super(Gate, self).__init__() + self.gate = nn.Linear(2 * d_model, 2 * d_model) + bias = bias_init_with_prob(0.5) + init.constant_(self.gate.bias, bias) + init.constant_(self.gate.weight, 0) + self.norm = RMSNorm(d_model) if use_rmsnorm else nn.LayerNorm(d_model) + + def forward(self, x1, x2): + gate_input = torch.cat([x1, x2], dim=-1) + gates = torch.sigmoid(self.gate(gate_input)) + gate1, gate2 = gates.chunk(2, dim=-1) + return self.norm(gate1 * x1 + gate2 * x2) \ No newline at end of file diff --git a/engine/deim/denoising.py b/engine/deim/denoising.py new file mode 100644 index 0000000000000000000000000000000000000000..4601af206ff84d6730d69afb0b84cf38755fc2e2 --- /dev/null +++ b/engine/deim/denoising.py @@ -0,0 +1,106 @@ +"""Copyright(c) 2023 lyuwenyu. All Rights Reserved. +Modifications Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +""" + +import torch + +from .utils import inverse_sigmoid +from .box_ops import box_cxcywh_to_xyxy, box_xyxy_to_cxcywh + + + +def get_contrastive_denoising_training_group(targets, + num_classes, + num_queries, + class_embed, + num_denoising=100, + label_noise_ratio=0.5, + box_noise_scale=1.0,): + """cnd""" + if num_denoising <= 0: + return None, None, None, None + + num_gts = [len(t['labels']) for t in targets] + device = targets[0]['labels'].device + + max_gt_num = max(num_gts) + if max_gt_num == 0: + return None, None, None, None + + num_group = num_denoising // max_gt_num + num_group = 1 if num_group == 0 else num_group + # pad gt to max_num of a batch + bs = len(num_gts) + + input_query_class = torch.full([bs, max_gt_num], num_classes, dtype=torch.int32, device=device) + input_query_bbox = torch.zeros([bs, max_gt_num, 4], device=device) + pad_gt_mask = torch.zeros([bs, max_gt_num], dtype=torch.bool, device=device) + + for i in range(bs): + num_gt = num_gts[i] + if num_gt > 0: + input_query_class[i, :num_gt] = targets[i]['labels'] + input_query_bbox[i, :num_gt] = targets[i]['boxes'] + pad_gt_mask[i, :num_gt] = 1 + # each group has positive and negative queries. + input_query_class = input_query_class.tile([1, 2 * num_group]) + input_query_bbox = input_query_bbox.tile([1, 2 * num_group, 1]) + pad_gt_mask = pad_gt_mask.tile([1, 2 * num_group]) + # positive and negative mask + negative_gt_mask = torch.zeros([bs, max_gt_num * 2, 1], device=device) + negative_gt_mask[:, max_gt_num:] = 1 + negative_gt_mask = negative_gt_mask.tile([1, num_group, 1]) + positive_gt_mask = 1 - negative_gt_mask + # contrastive denoising training positive index + positive_gt_mask = positive_gt_mask.squeeze(-1) * pad_gt_mask + dn_positive_idx = torch.nonzero(positive_gt_mask)[:, 1] + dn_positive_idx = torch.split(dn_positive_idx, [n * num_group for n in num_gts]) + # total denoising queries + num_denoising = int(max_gt_num * 2 * num_group) + + if label_noise_ratio > 0: + mask = torch.rand_like(input_query_class, dtype=torch.float) < (label_noise_ratio * 0.5) + # randomly put a new one here + new_label = torch.randint_like(mask, 0, num_classes, dtype=input_query_class.dtype) + input_query_class = torch.where(mask & pad_gt_mask, new_label, input_query_class) + + if box_noise_scale > 0: + known_bbox = box_cxcywh_to_xyxy(input_query_bbox) + diff = torch.tile(input_query_bbox[..., 2:] * 0.5, [1, 1, 2]) * box_noise_scale + rand_sign = torch.randint_like(input_query_bbox, 0, 2) * 2.0 - 1.0 + rand_part = torch.rand_like(input_query_bbox) + rand_part = (rand_part + 1.0) * negative_gt_mask + rand_part * (1 - negative_gt_mask) + known_bbox += (rand_sign * rand_part * diff) + known_bbox = torch.clip(known_bbox, min=0.0, max=1.0) + input_query_bbox = box_xyxy_to_cxcywh(known_bbox) + input_query_bbox[input_query_bbox < 0] *= -1 + input_query_bbox_unact = inverse_sigmoid(input_query_bbox) + + input_query_logits = class_embed(input_query_class) + + tgt_size = num_denoising + num_queries + attn_mask = torch.full([tgt_size, tgt_size], False, dtype=torch.bool, device=device) + # match query cannot see the reconstruction + attn_mask[num_denoising:, :num_denoising] = True + + # reconstruct cannot see each other + for i in range(num_group): + if i == 0: + attn_mask[max_gt_num * 2 * i: max_gt_num * 2 * (i + 1), max_gt_num * 2 * (i + 1): num_denoising] = True + if i == num_group - 1: + attn_mask[max_gt_num * 2 * i: max_gt_num * 2 * (i + 1), :max_gt_num * i * 2] = True + else: + attn_mask[max_gt_num * 2 * i: max_gt_num * 2 * (i + 1), max_gt_num * 2 * (i + 1): num_denoising] = True + attn_mask[max_gt_num * 2 * i: max_gt_num * 2 * (i + 1), :max_gt_num * 2 * i] = True + + dn_meta = { + "dn_positive_idx": dn_positive_idx, + "dn_num_group": num_group, + "dn_num_split": [num_denoising, num_queries] + } + + # print(input_query_class.shape) # torch.Size([4, 196, 256]) + # print(input_query_bbox.shape) # torch.Size([4, 196, 4]) + # print(attn_mask.shape) # torch.Size([496, 496]) + + return input_query_logits, input_query_bbox_unact, attn_mask, dn_meta diff --git a/engine/deim/dfine_decoder.py b/engine/deim/dfine_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..b929ca6c42f02361a4f06f096aa3dbf08ae03d07 --- /dev/null +++ b/engine/deim/dfine_decoder.py @@ -0,0 +1,790 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE/) +Copyright (c) 2024 D-FINE Authors. All Rights Reserved. +""" + +import math +import copy +import functools +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.nn.init as init +from typing import List + +from .dfine_utils import weighting_function, distance2bbox +from .denoising import get_contrastive_denoising_training_group +from .utils import deformable_attention_core_func_v2, get_activation, inverse_sigmoid +from .utils import bias_init_with_prob +from ..core import register + +__all__ = ['DFINETransformer'] + + +class MLP(nn.Module): + def __init__(self, input_dim, hidden_dim, output_dim, num_layers, act='relu'): + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) + self.act = get_activation(act) + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = self.act(layer(x)) if i < self.num_layers - 1 else layer(x) + return x + + +class MSDeformableAttention(nn.Module): + def __init__( + self, + embed_dim=256, + num_heads=8, + num_levels=4, + num_points=4, + method='default', + offset_scale=0.5, + ): + """Multi-Scale Deformable Attention + """ + super(MSDeformableAttention, self).__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.num_levels = num_levels + self.offset_scale = offset_scale + + if isinstance(num_points, list): + assert len(num_points) == num_levels, '' + num_points_list = num_points + else: + num_points_list = [num_points for _ in range(num_levels)] + + self.num_points_list = num_points_list + + num_points_scale = [1/n for n in num_points_list for _ in range(n)] + self.register_buffer('num_points_scale', torch.tensor(num_points_scale, dtype=torch.float32)) + + self.total_points = num_heads * sum(num_points_list) + self.method = method + + self.head_dim = embed_dim // num_heads + assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + + self.sampling_offsets = nn.Linear(embed_dim, self.total_points * 2) + self.attention_weights = nn.Linear(embed_dim, self.total_points) + + self.ms_deformable_attn_core = functools.partial(deformable_attention_core_func_v2, method=self.method) + + self._reset_parameters() + + if method == 'discrete': + for p in self.sampling_offsets.parameters(): + p.requires_grad = False + + def _reset_parameters(self): + # sampling_offsets + init.constant_(self.sampling_offsets.weight, 0) + thetas = torch.arange(self.num_heads, dtype=torch.float32) * (2.0 * math.pi / self.num_heads) + grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) + grid_init = grid_init / grid_init.abs().max(-1, keepdim=True).values + grid_init = grid_init.reshape(self.num_heads, 1, 2).tile([1, sum(self.num_points_list), 1]) + scaling = torch.concat([torch.arange(1, n + 1) for n in self.num_points_list]).reshape(1, -1, 1) + grid_init *= scaling + self.sampling_offsets.bias.data[...] = grid_init.flatten() + + # attention_weights + init.constant_(self.attention_weights.weight, 0) + init.constant_(self.attention_weights.bias, 0) + + + def forward(self, + query: torch.Tensor, + reference_points: torch.Tensor, + value: torch.Tensor, + value_spatial_shapes: List[int]): + """ + Args: + query (Tensor): [bs, query_length, C] + reference_points (Tensor): [bs, query_length, n_levels, 2], range in [0, 1], top-left (0,0), + bottom-right (1, 1), including padding area + value (Tensor): [bs, value_length, C] + value_spatial_shapes (List): [n_levels, 2], [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})] + + Returns: + output (Tensor): [bs, Length_{query}, C] + """ + bs, Len_q = query.shape[:2] + + sampling_offsets: torch.Tensor = self.sampling_offsets(query) + sampling_offsets = sampling_offsets.reshape(bs, Len_q, self.num_heads, sum(self.num_points_list), 2) + + attention_weights = self.attention_weights(query).reshape(bs, Len_q, self.num_heads, sum(self.num_points_list)) + attention_weights = F.softmax(attention_weights, dim=-1) + + if reference_points.shape[-1] == 2: + offset_normalizer = torch.tensor(value_spatial_shapes) + offset_normalizer = offset_normalizer.flip([1]).reshape(1, 1, 1, self.num_levels, 1, 2) + sampling_locations = reference_points.reshape(bs, Len_q, 1, self.num_levels, 1, 2) + sampling_offsets / offset_normalizer + elif reference_points.shape[-1] == 4: + # reference_points [8, 480, None, 1, 4] + # sampling_offsets [8, 480, 8, 12, 2] + num_points_scale = self.num_points_scale.to(dtype=query.dtype).unsqueeze(-1) + offset = sampling_offsets * num_points_scale * reference_points[:, :, None, :, 2:] * self.offset_scale + sampling_locations = reference_points[:, :, None, :, :2] + offset + else: + raise ValueError( + "Last dim of reference_points must be 2 or 4, but get {} instead.". + format(reference_points.shape[-1])) + + output = self.ms_deformable_attn_core(value, value_spatial_shapes, sampling_locations, attention_weights, self.num_points_list) + + return output + + +class TransformerDecoderLayer(nn.Module): + def __init__(self, + d_model=256, + n_head=8, + dim_feedforward=1024, + dropout=0., + activation='relu', + n_levels=4, + n_points=4, + cross_attn_method='default', + layer_scale=None): + super(TransformerDecoderLayer, self).__init__() + if layer_scale is not None: + dim_feedforward = round(layer_scale * dim_feedforward) + d_model = round(layer_scale * d_model) + + # self attention + self.self_attn = nn.MultiheadAttention(d_model, n_head, dropout=dropout, batch_first=True) + self.dropout1 = nn.Dropout(dropout) + self.norm1 = nn.LayerNorm(d_model) + + # cross attention + self.cross_attn = MSDeformableAttention(d_model, n_head, n_levels, n_points, \ + method=cross_attn_method) + self.dropout2 = nn.Dropout(dropout) + + # gate + self.gateway = Gate(d_model) + + # ffn + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.activation = get_activation(activation) + self.dropout3 = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + self.dropout4 = nn.Dropout(dropout) + self.norm3 = nn.LayerNorm(d_model) + + self._reset_parameters() + + def _reset_parameters(self): + init.xavier_uniform_(self.linear1.weight) + init.xavier_uniform_(self.linear2.weight) + + def with_pos_embed(self, tensor, pos): + return tensor if pos is None else tensor + pos + + def forward_ffn(self, tgt): + return self.linear2(self.dropout3(self.activation(self.linear1(tgt)))) + + def forward(self, + target, + reference_points, + value, + spatial_shapes, + attn_mask=None, + query_pos_embed=None): + + # self attention + q = k = self.with_pos_embed(target, query_pos_embed) + + target2, _ = self.self_attn(q, k, value=target, attn_mask=attn_mask) + target = target + self.dropout1(target2) + target = self.norm1(target) + + # cross attention + target2 = self.cross_attn(\ + self.with_pos_embed(target, query_pos_embed), + reference_points, + value, + spatial_shapes) + + target = self.gateway(target, self.dropout2(target2)) + + # ffn + target2 = self.forward_ffn(target) + target = target + self.dropout4(target2) + target = self.norm3(target.clamp(min=-65504, max=65504)) + + return target + + +class Gate(nn.Module): + def __init__(self, d_model): + super(Gate, self).__init__() + self.gate = nn.Linear(2 * d_model, 2 * d_model) + bias = bias_init_with_prob(0.5) + init.constant_(self.gate.bias, bias) + init.constant_(self.gate.weight, 0) + self.norm = nn.LayerNorm(d_model) + + def forward(self, x1, x2): + gate_input = torch.cat([x1, x2], dim=-1) + gates = torch.sigmoid(self.gate(gate_input)) + gate1, gate2 = gates.chunk(2, dim=-1) + return self.norm(gate1 * x1 + gate2 * x2) + + +class Integral(nn.Module): + """ + A static layer that calculates integral results from a distribution. + + This layer computes the target location using the formula: `sum{Pr(n) * W(n)}`, + where Pr(n) is the softmax probability vector representing the discrete + distribution, and W(n) is the non-uniform Weighting Function. + + Args: + reg_max (int): Max number of the discrete bins. Default is 32. + It can be adjusted based on the dataset or task requirements. + """ + + def __init__(self, reg_max=32): + super(Integral, self).__init__() + self.reg_max = reg_max + + def forward(self, x, project): + shape = x.shape + x = F.softmax(x.reshape(-1, self.reg_max + 1), dim=1) + x = F.linear(x, project.to(x.device)).reshape(-1, 4) + return x.reshape(list(shape[:-1]) + [-1]) + + +class LQE(nn.Module): + def __init__(self, k, hidden_dim, num_layers, reg_max, act='relu'): + super(LQE, self).__init__() + self.k = k + self.reg_max = reg_max + self.reg_conf = MLP(4 * (k + 1), hidden_dim, 1, num_layers, act=act) + init.constant_(self.reg_conf.layers[-1].bias, 0) + init.constant_(self.reg_conf.layers[-1].weight, 0) + + def forward(self, scores, pred_corners): + B, L, _ = pred_corners.size() + prob = F.softmax(pred_corners.reshape(B, L, 4, self.reg_max+1), dim=-1) + prob_topk, _ = prob.topk(self.k, dim=-1) + stat = torch.cat([prob_topk, prob_topk.mean(dim=-1, keepdim=True)], dim=-1) + quality_score = self.reg_conf(stat.reshape(B, L, -1)) + return scores + quality_score + + +class TransformerDecoder(nn.Module): + """ + Transformer Decoder implementing Fine-grained Distribution Refinement (FDR). + + This decoder refines object detection predictions through iterative updates across multiple layers, + utilizing attention mechanisms, location quality estimators, and distribution refinement techniques + to improve bounding box accuracy and robustness. + """ + + def __init__(self, hidden_dim, decoder_layer, decoder_layer_wide, num_layers, num_head, reg_max, reg_scale, up, + eval_idx=-1, layer_scale=2, act='relu'): + super(TransformerDecoder, self).__init__() + self.hidden_dim = hidden_dim + self.num_layers = num_layers + self.layer_scale = layer_scale + self.num_head = num_head + self.eval_idx = eval_idx if eval_idx >= 0 else num_layers + eval_idx + self.up, self.reg_scale, self.reg_max = up, reg_scale, reg_max + self.layers = nn.ModuleList([copy.deepcopy(decoder_layer) for _ in range(self.eval_idx + 1)] \ + + [copy.deepcopy(decoder_layer_wide) for _ in range(num_layers - self.eval_idx - 1)]) + self.lqe_layers = nn.ModuleList([copy.deepcopy(LQE(4, 64, 2, reg_max, act=act)) for _ in range(num_layers)]) + + def value_op(self, memory, value_proj, value_scale, memory_mask, memory_spatial_shapes): + """ + Preprocess values for MSDeformableAttention. + """ + value = value_proj(memory) if value_proj is not None else memory + value = F.interpolate(memory, size=value_scale) if value_scale is not None else value + if memory_mask is not None: + value = value * memory_mask.to(value.dtype).unsqueeze(-1) + value = value.reshape(value.shape[0], value.shape[1], self.num_head, -1) + split_shape = [h * w for h, w in memory_spatial_shapes] + return value.permute(0, 2, 3, 1).split(split_shape, dim=-1) + + def convert_to_deploy(self): + self.project = weighting_function(self.reg_max, self.up, self.reg_scale, deploy=True) + self.layers = self.layers[:self.eval_idx + 1] + self.lqe_layers = nn.ModuleList([nn.Identity()] * (self.eval_idx) + [self.lqe_layers[self.eval_idx]]) + + def forward(self, + target, + ref_points_unact, + memory, + spatial_shapes, + bbox_head, + score_head, + query_pos_head, + pre_bbox_head, + integral, + up, + reg_scale, + attn_mask=None, + memory_mask=None, + dn_meta=None): + output = target + output_detach = pred_corners_undetach = 0 + value = self.value_op(memory, None, None, memory_mask, spatial_shapes) + + dec_out_bboxes = [] + dec_out_logits = [] + dec_out_pred_corners = [] + dec_out_refs = [] + if not hasattr(self, 'project'): + project = weighting_function(self.reg_max, up, reg_scale) + else: + project = self.project + + ref_points_detach = F.sigmoid(ref_points_unact) + + for i, layer in enumerate(self.layers): + ref_points_input = ref_points_detach.unsqueeze(2) + query_pos_embed = query_pos_head(ref_points_detach).clamp(min=-10, max=10) + + if i >= self.eval_idx + 1 and self.layer_scale > 1: + query_pos_embed = F.interpolate(query_pos_embed, scale_factor=self.layer_scale) + value = self.value_op(memory, None, query_pos_embed.shape[-1], memory_mask, spatial_shapes) + output = F.interpolate(output, size=query_pos_embed.shape[-1]) + output_detach = output.detach() + + output = layer(output, ref_points_input, value, spatial_shapes, attn_mask, query_pos_embed) + + if i == 0 : + # Initial bounding box predictions with inverse sigmoid refinement + pre_bboxes = F.sigmoid(pre_bbox_head(output) + inverse_sigmoid(ref_points_detach)) + pre_scores = score_head[0](output) + ref_points_initial = pre_bboxes.detach() + + # Refine bounding box corners using FDR, integrating previous layer's corrections + pred_corners = bbox_head[i](output + output_detach) + pred_corners_undetach + inter_ref_bbox = distance2bbox(ref_points_initial, integral(pred_corners, project), reg_scale) + + if self.training or i == self.eval_idx: + scores = score_head[i](output) + # Lqe does not affect the performance here. + scores = self.lqe_layers[i](scores, pred_corners) + dec_out_logits.append(scores) + dec_out_bboxes.append(inter_ref_bbox) + dec_out_pred_corners.append(pred_corners) + dec_out_refs.append(ref_points_initial) + + if not self.training: + break + + pred_corners_undetach = pred_corners + ref_points_detach = inter_ref_bbox.detach() + output_detach = output.detach() + + return torch.stack(dec_out_bboxes), torch.stack(dec_out_logits), \ + torch.stack(dec_out_pred_corners), torch.stack(dec_out_refs), pre_bboxes, pre_scores + + +@register() +class DFINETransformer(nn.Module): + __share__ = ['num_classes', 'eval_spatial_size'] + + def __init__(self, + num_classes=80, + hidden_dim=256, + num_queries=300, + feat_channels=[512, 1024, 2048], + feat_strides=[8, 16, 32], + num_levels=3, + num_points=4, + nhead=8, + num_layers=6, + dim_feedforward=1024, + dropout=0., + activation="relu", + num_denoising=100, + label_noise_ratio=0.5, + box_noise_scale=1.0, + learn_query_content=False, + eval_spatial_size=None, + eval_idx=-1, + eps=1e-2, + aux_loss=True, + cross_attn_method='default', + query_select_method='default', + reg_max=32, + reg_scale=4., + layer_scale=1, + mlp_act='relu', + ): + super().__init__() + assert len(feat_channels) <= num_levels + assert len(feat_strides) == len(feat_channels) + + for _ in range(num_levels - len(feat_strides)): + feat_strides.append(feat_strides[-1] * 2) + + self.hidden_dim = hidden_dim + scaled_dim = round(layer_scale*hidden_dim) + self.nhead = nhead + self.feat_strides = feat_strides + self.num_levels = num_levels + self.num_classes = num_classes + self.num_queries = num_queries + self.eps = eps + self.num_layers = num_layers + self.eval_spatial_size = eval_spatial_size + self.aux_loss = aux_loss + self.reg_max = reg_max + + assert query_select_method in ('default', 'one2many', 'agnostic'), '' + assert cross_attn_method in ('default', 'discrete'), '' + self.cross_attn_method = cross_attn_method + self.query_select_method = query_select_method + + # backbone feature projection + self._build_input_proj_layer(feat_channels) + + # Transformer module + self.up = nn.Parameter(torch.tensor([0.5]), requires_grad=False) + self.reg_scale = nn.Parameter(torch.tensor([reg_scale]), requires_grad=False) + decoder_layer = TransformerDecoderLayer(hidden_dim, nhead, dim_feedforward, dropout, \ + activation, num_levels, num_points, cross_attn_method=cross_attn_method) + decoder_layer_wide = TransformerDecoderLayer(hidden_dim, nhead, dim_feedforward, dropout, \ + activation, num_levels, num_points, cross_attn_method=cross_attn_method, layer_scale=layer_scale) + self.decoder = TransformerDecoder(hidden_dim, decoder_layer, decoder_layer_wide, num_layers, nhead, + reg_max, self.reg_scale, self.up, eval_idx, layer_scale, act=activation) + # denoising + self.num_denoising = num_denoising + self.label_noise_ratio = label_noise_ratio + self.box_noise_scale = box_noise_scale + if num_denoising > 0: + self.denoising_class_embed = nn.Embedding(num_classes+1, hidden_dim, padding_idx=num_classes) + init.normal_(self.denoising_class_embed.weight[:-1]) + + # decoder embedding + self.learn_query_content = learn_query_content + if learn_query_content: + self.tgt_embed = nn.Embedding(num_queries, hidden_dim) + self.query_pos_head = MLP(4, 2 * hidden_dim, hidden_dim, 2, act=mlp_act) + + # if num_select_queries != self.num_queries: + # layer = TransformerEncoderLayer(hidden_dim, nhead, dim_feedforward, activation='gelu') + # self.encoder = TransformerEncoder(layer, 1) + + self.enc_output = nn.Sequential(OrderedDict([ + ('proj', nn.Linear(hidden_dim, hidden_dim)), + ('norm', nn.LayerNorm(hidden_dim,)), + ])) + + if query_select_method == 'agnostic': + self.enc_score_head = nn.Linear(hidden_dim, 1) + else: + self.enc_score_head = nn.Linear(hidden_dim, num_classes) + + self.enc_bbox_head = MLP(hidden_dim, hidden_dim, 4, 3, act=mlp_act) + + # decoder head + self.eval_idx = eval_idx if eval_idx >= 0 else num_layers + eval_idx + self.dec_score_head = nn.ModuleList( + [nn.Linear(hidden_dim, num_classes) for _ in range(self.eval_idx + 1)] + + [nn.Linear(scaled_dim, num_classes) for _ in range(num_layers - self.eval_idx - 1)]) + self.pre_bbox_head = MLP(hidden_dim, hidden_dim, 4, 3, act=mlp_act) + self.dec_bbox_head = nn.ModuleList( + [MLP(hidden_dim, hidden_dim, 4 * (self.reg_max+1), 3, act=mlp_act) for _ in range(self.eval_idx + 1)] + + [MLP(scaled_dim, scaled_dim, 4 * (self.reg_max+1), 3, act=mlp_act) for _ in range(num_layers - self.eval_idx - 1)]) + self.integral = Integral(self.reg_max) + + # init encoder output anchors and valid_mask + if self.eval_spatial_size: + anchors, valid_mask = self._generate_anchors() + self.register_buffer('anchors', anchors) + self.register_buffer('valid_mask', valid_mask) + # init encoder output anchors and valid_mask + if self.eval_spatial_size: + self.anchors, self.valid_mask = self._generate_anchors() + + + self._reset_parameters(feat_channels) + + def convert_to_deploy(self): + self.dec_score_head = nn.ModuleList([nn.Identity()] * (self.eval_idx) + [self.dec_score_head[self.eval_idx]]) + self.dec_bbox_head = nn.ModuleList( + [self.dec_bbox_head[i] if i <= self.eval_idx else nn.Identity() for i in range(len(self.dec_bbox_head))] + ) + + def _reset_parameters(self, feat_channels): + bias = bias_init_with_prob(0.01) + init.constant_(self.enc_score_head.bias, bias) + init.constant_(self.enc_bbox_head.layers[-1].weight, 0) + init.constant_(self.enc_bbox_head.layers[-1].bias, 0) + + init.constant_(self.pre_bbox_head.layers[-1].weight, 0) + init.constant_(self.pre_bbox_head.layers[-1].bias, 0) + + for cls_, reg_ in zip(self.dec_score_head, self.dec_bbox_head): + init.constant_(cls_.bias, bias) + if hasattr(reg_, 'layers'): + init.constant_(reg_.layers[-1].weight, 0) + init.constant_(reg_.layers[-1].bias, 0) + + init.xavier_uniform_(self.enc_output[0].weight) + if self.learn_query_content: + init.xavier_uniform_(self.tgt_embed.weight) + init.xavier_uniform_(self.query_pos_head.layers[0].weight) + init.xavier_uniform_(self.query_pos_head.layers[1].weight) + for m, in_channels in zip(self.input_proj, feat_channels): + if in_channels != self.hidden_dim: + init.xavier_uniform_(m[0].weight) + + def _build_input_proj_layer(self, feat_channels): + self.input_proj = nn.ModuleList() + for in_channels in feat_channels: + if in_channels == self.hidden_dim: + self.input_proj.append(nn.Identity()) + else: + self.input_proj.append( + nn.Sequential(OrderedDict([ + ('conv', nn.Conv2d(in_channels, self.hidden_dim, 1, bias=False)), + ('norm', nn.BatchNorm2d(self.hidden_dim,))]) + ) + ) + + in_channels = feat_channels[-1] + + for _ in range(self.num_levels - len(feat_channels)): + if in_channels == self.hidden_dim: + self.input_proj.append(nn.Identity()) + else: + self.input_proj.append( + nn.Sequential(OrderedDict([ + ('conv', nn.Conv2d(in_channels, self.hidden_dim, 3, 2, padding=1, bias=False)), + ('norm', nn.BatchNorm2d(self.hidden_dim))]) + ) + ) + in_channels = self.hidden_dim + + def _get_encoder_input(self, feats: List[torch.Tensor]): + # get projection features + proj_feats = [self.input_proj[i](feat) for i, feat in enumerate(feats)] + if self.num_levels > len(proj_feats): + len_srcs = len(proj_feats) + for i in range(len_srcs, self.num_levels): + if i == len_srcs: + proj_feats.append(self.input_proj[i](feats[-1])) + else: + proj_feats.append(self.input_proj[i](proj_feats[-1])) + + # get encoder inputs + feat_flatten = [] + spatial_shapes = [] + for i, feat in enumerate(proj_feats): + _, _, h, w = feat.shape + # [b, c, h, w] -> [b, h*w, c] + feat_flatten.append(feat.flatten(2).permute(0, 2, 1)) + # [num_levels, 2] + spatial_shapes.append([h, w]) + + # [b, l, c] + feat_flatten = torch.concat(feat_flatten, 1) + return feat_flatten, spatial_shapes + + def _generate_anchors(self, + spatial_shapes=None, + grid_size=0.05, + dtype=torch.float32, + device='cpu'): + if spatial_shapes is None: + spatial_shapes = [] + eval_h, eval_w = self.eval_spatial_size + for s in self.feat_strides: + spatial_shapes.append([int(eval_h / s), int(eval_w / s)]) + + anchors = [] + for lvl, (h, w) in enumerate(spatial_shapes): + grid_y, grid_x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing='ij') + grid_xy = torch.stack([grid_x, grid_y], dim=-1) + grid_xy = (grid_xy.unsqueeze(0) + 0.5) / torch.tensor([w, h], dtype=dtype) + wh = torch.ones_like(grid_xy) * grid_size * (2.0 ** lvl) + lvl_anchors = torch.concat([grid_xy, wh], dim=-1).reshape(-1, h * w, 4) + anchors.append(lvl_anchors) + + anchors = torch.concat(anchors, dim=1).to(device) + valid_mask = ((anchors > self.eps) * (anchors < 1 - self.eps)).all(-1, keepdim=True) + anchors = torch.log(anchors / (1 - anchors)) + anchors = torch.where(valid_mask, anchors, torch.inf) + + return anchors, valid_mask + + + def _get_decoder_input(self, + memory: torch.Tensor, + spatial_shapes, + denoising_logits=None, + denoising_bbox_unact=None): + + # prepare input for decoder + if self.training or self.eval_spatial_size is None: + anchors, valid_mask = self._generate_anchors(spatial_shapes, device=memory.device) + else: + anchors = self.anchors + valid_mask = self.valid_mask + if memory.shape[0] > 1: + anchors = anchors.repeat(memory.shape[0], 1, 1) + + # memory = torch.where(valid_mask, memory, 0) + memory = valid_mask.to(memory.dtype) * memory + + output_memory :torch.Tensor = self.enc_output(memory) + enc_outputs_logits :torch.Tensor = self.enc_score_head(output_memory) + + enc_topk_bboxes_list, enc_topk_logits_list = [], [] + enc_topk_memory, enc_topk_logits, enc_topk_anchors = \ + self._select_topk(output_memory, enc_outputs_logits, anchors, self.num_queries) + + enc_topk_bbox_unact :torch.Tensor = self.enc_bbox_head(enc_topk_memory) + enc_topk_anchors + + if self.training: + enc_topk_bboxes = F.sigmoid(enc_topk_bbox_unact) + enc_topk_bboxes_list.append(enc_topk_bboxes) + enc_topk_logits_list.append(enc_topk_logits) + + # if self.num_select_queries != self.num_queries: + # raise NotImplementedError('') + + if self.learn_query_content: + content = self.tgt_embed.weight.unsqueeze(0).tile([memory.shape[0], 1, 1]) + else: + content = enc_topk_memory.detach() + + enc_topk_bbox_unact = enc_topk_bbox_unact.detach() + + if denoising_bbox_unact is not None: + enc_topk_bbox_unact = torch.concat([denoising_bbox_unact, enc_topk_bbox_unact], dim=1) + content = torch.concat([denoising_logits, content], dim=1) + + return content, enc_topk_bbox_unact, enc_topk_bboxes_list, enc_topk_logits_list + + def _select_topk(self, memory: torch.Tensor, outputs_logits: torch.Tensor, outputs_anchors_unact: torch.Tensor, topk: int): + if self.query_select_method == 'default': + _, topk_ind = torch.topk(outputs_logits.max(-1).values, topk, dim=-1) + + elif self.query_select_method == 'one2many': + _, topk_ind = torch.topk(outputs_logits.flatten(1), topk, dim=-1) + topk_ind = topk_ind // self.num_classes + + elif self.query_select_method == 'agnostic': + _, topk_ind = torch.topk(outputs_logits.squeeze(-1), topk, dim=-1) + + topk_ind: torch.Tensor + + topk_anchors = outputs_anchors_unact.gather(dim=1, \ + index=topk_ind.unsqueeze(-1).repeat(1, 1, outputs_anchors_unact.shape[-1])) + + topk_logits = outputs_logits.gather(dim=1, \ + index=topk_ind.unsqueeze(-1).repeat(1, 1, outputs_logits.shape[-1])) if self.training else None + + topk_memory = memory.gather(dim=1, \ + index=topk_ind.unsqueeze(-1).repeat(1, 1, memory.shape[-1])) + + return topk_memory, topk_logits, topk_anchors + + def forward(self, feats, targets=None): + # input projection and embedding + memory, spatial_shapes = self._get_encoder_input(feats) + + # prepare denoising training + if self.training and self.num_denoising > 0: + denoising_logits, denoising_bbox_unact, attn_mask, dn_meta = \ + get_contrastive_denoising_training_group(targets, \ + self.num_classes, + self.num_queries, + self.denoising_class_embed, + num_denoising=self.num_denoising, + label_noise_ratio=self.label_noise_ratio, + box_noise_scale=1.0, + ) + else: + denoising_logits, denoising_bbox_unact, attn_mask, dn_meta = None, None, None, None + + init_ref_contents, init_ref_points_unact, enc_topk_bboxes_list, enc_topk_logits_list = \ + self._get_decoder_input(memory, spatial_shapes, denoising_logits, denoising_bbox_unact) + + # decoder + out_bboxes, out_logits, out_corners, out_refs, pre_bboxes, pre_logits = self.decoder( + init_ref_contents, + init_ref_points_unact, + memory, + spatial_shapes, + self.dec_bbox_head, + self.dec_score_head, + self.query_pos_head, + self.pre_bbox_head, + self.integral, + self.up, + self.reg_scale, + attn_mask=attn_mask, + dn_meta=dn_meta) + + if self.training and dn_meta is not None: + # the output from the first decoder layer, only one + dn_pre_logits, pre_logits = torch.split(pre_logits, dn_meta['dn_num_split'], dim=1) + dn_pre_bboxes, pre_bboxes = torch.split(pre_bboxes, dn_meta['dn_num_split'], dim=1) + + dn_out_logits, out_logits = torch.split(out_logits, dn_meta['dn_num_split'], dim=2) + dn_out_bboxes, out_bboxes = torch.split(out_bboxes, dn_meta['dn_num_split'], dim=2) + + dn_out_corners, out_corners = torch.split(out_corners, dn_meta['dn_num_split'], dim=2) + dn_out_refs, out_refs = torch.split(out_refs, dn_meta['dn_num_split'], dim=2) + + + if self.training: + out = {'pred_logits': out_logits[-1], 'pred_boxes': out_bboxes[-1], 'pred_corners': out_corners[-1], + 'ref_points': out_refs[-1], 'up': self.up, 'reg_scale': self.reg_scale} + else: + out = {'pred_logits': out_logits[-1], 'pred_boxes': out_bboxes[-1]} + + if self.training and self.aux_loss: + out['aux_outputs'] = self._set_aux_loss2(out_logits[:-1], out_bboxes[:-1], out_corners[:-1], out_refs[:-1], + out_corners[-1], out_logits[-1]) + out['enc_aux_outputs'] = self._set_aux_loss(enc_topk_logits_list, enc_topk_bboxes_list) + out['pre_outputs'] = {'pred_logits': pre_logits, 'pred_boxes': pre_bboxes} + out['enc_meta'] = {'class_agnostic': self.query_select_method == 'agnostic'} + + if dn_meta is not None: + out['dn_outputs'] = self._set_aux_loss2(dn_out_logits, dn_out_bboxes, dn_out_corners, dn_out_refs, + dn_out_corners[-1], dn_out_logits[-1]) + out['dn_pre_outputs'] = {'pred_logits': dn_pre_logits, 'pred_boxes': dn_pre_bboxes} + out['dn_meta'] = dn_meta + + return out + + + @torch.jit.unused + def _set_aux_loss(self, outputs_class, outputs_coord): + # this is a workaround to make torchscript happy, as torchscript + # doesn't support dictionary with non-homogeneous values, such + # as a dict having both a Tensor and a list. + return [{'pred_logits': a, 'pred_boxes': b} for a, b in zip(outputs_class, outputs_coord)] + + + @torch.jit.unused + def _set_aux_loss2(self, outputs_class, outputs_coord, outputs_corners, outputs_ref, + teacher_corners=None, teacher_logits=None): + # this is a workaround to make torchscript happy, as torchscript + # doesn't support dictionary with non-homogeneous values, such + # as a dict having both a Tensor and a list. + return [{'pred_logits': a, 'pred_boxes': b, 'pred_corners': c, 'ref_points': d, + 'teacher_corners': teacher_corners, 'teacher_logits': teacher_logits} + for a, b, c, d in zip(outputs_class, outputs_coord, outputs_corners, outputs_ref)] diff --git a/engine/deim/dfine_utils.py b/engine/deim/dfine_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c0864e20f4a660a32caba2526f032f99f918fafe --- /dev/null +++ b/engine/deim/dfine_utils.py @@ -0,0 +1,156 @@ +""" +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +""" + +import torch +from .box_ops import box_xyxy_to_cxcywh + + +def weighting_function(reg_max, up, reg_scale, deploy=False): + """ + Generates the non-uniform Weighting Function W(n) for bounding box regression. + + Args: + reg_max (int): Max number of the discrete bins. + up (Tensor): Controls upper bounds of the sequence, + where maximum offset is ±up * H / W. + reg_scale (float): Controls the curvature of the Weighting Function. + Larger values result in flatter weights near the central axis W(reg_max/2)=0 + and steeper weights at both ends. + deploy (bool): If True, uses deployment mode settings. + + Returns: + Tensor: Sequence of Weighting Function. + """ + if deploy: + upper_bound1 = (abs(up[0]) * abs(reg_scale)).item() + upper_bound2 = (abs(up[0]) * abs(reg_scale) * 2).item() + step = (upper_bound1 + 1) ** (2 / (reg_max - 2)) + left_values = [-(step) ** i + 1 for i in range(reg_max // 2 - 1, 0, -1)] + right_values = [(step) ** i - 1 for i in range(1, reg_max // 2)] + values = [-upper_bound2] + left_values + [torch.zeros_like(up[0][None])] + right_values + [upper_bound2] + return torch.tensor(values, dtype=up.dtype, device=up.device) + else: + upper_bound1 = abs(up[0]) * abs(reg_scale) + upper_bound2 = abs(up[0]) * abs(reg_scale) * 2 + step = (upper_bound1 + 1) ** (2 / (reg_max - 2)) + left_values = [-(step) ** i + 1 for i in range(reg_max // 2 - 1, 0, -1)] + right_values = [(step) ** i - 1 for i in range(1, reg_max // 2)] + values = [-upper_bound2] + left_values + [torch.zeros_like(up[0][None])] + right_values + [upper_bound2] + return torch.cat(values, 0) + + +def translate_gt(gt, reg_max, reg_scale, up): + """ + Decodes bounding box ground truth (GT) values into distribution-based GT representations. + + This function maps continuous GT values into discrete distribution bins, which can be used + for regression tasks in object detection models. It calculates the indices of the closest + bins to each GT value and assigns interpolation weights to these bins based on their proximity + to the GT value. + + Args: + gt (Tensor): Ground truth bounding box values, shape (N, ). + reg_max (int): Maximum number of discrete bins for the distribution. + reg_scale (float): Controls the curvature of the Weighting Function. + up (Tensor): Controls the upper bounds of the Weighting Function. + + Returns: + Tuple[Tensor, Tensor, Tensor]: + - indices (Tensor): Index of the left bin closest to each GT value, shape (N, ). + - weight_right (Tensor): Weight assigned to the right bin, shape (N, ). + - weight_left (Tensor): Weight assigned to the left bin, shape (N, ). + """ + gt = gt.reshape(-1) + function_values = weighting_function(reg_max, up, reg_scale) + + # Find the closest left-side indices for each value + diffs = function_values.unsqueeze(0) - gt.unsqueeze(1) + mask = diffs <= 0 + closest_left_indices = torch.sum(mask, dim=1) - 1 + + # Calculate the weights for the interpolation + indices = closest_left_indices.float() + + weight_right = torch.zeros_like(indices) + weight_left = torch.zeros_like(indices) + + valid_idx_mask = (indices >= 0) & (indices < reg_max) + valid_indices = indices[valid_idx_mask].long() + + # Obtain distances + left_values = function_values[valid_indices] + right_values = function_values[valid_indices + 1] + + left_diffs = torch.abs(gt[valid_idx_mask] - left_values) + right_diffs = torch.abs(right_values - gt[valid_idx_mask]) + + # Valid weights + weight_right[valid_idx_mask] = left_diffs / (left_diffs + right_diffs) + weight_left[valid_idx_mask] = 1.0 - weight_right[valid_idx_mask] + + # Invalid weights (out of range) + invalid_idx_mask_neg = (indices < 0) + weight_right[invalid_idx_mask_neg] = 0.0 + weight_left[invalid_idx_mask_neg] = 1.0 + indices[invalid_idx_mask_neg] = 0.0 + + invalid_idx_mask_pos = (indices >= reg_max) + weight_right[invalid_idx_mask_pos] = 1.0 + weight_left[invalid_idx_mask_pos] = 0.0 + indices[invalid_idx_mask_pos] = reg_max - 0.1 + + return indices, weight_right, weight_left + + +def distance2bbox(points, distance, reg_scale): + """ + Decodes edge-distances into bounding box coordinates. + + Args: + points (Tensor): (B, N, 4) or (N, 4) format, representing [x, y, w, h], + where (x, y) is the center and (w, h) are width and height. + distance (Tensor): (B, N, 4) or (N, 4), representing distances from the + point to the left, top, right, and bottom boundaries. + + reg_scale (float): Controls the curvature of the Weighting Function. + + Returns: + Tensor: Bounding boxes in (N, 4) or (B, N, 4) format [cx, cy, w, h]. + """ + reg_scale = abs(reg_scale) + x1 = points[..., 0] - (0.5 * reg_scale + distance[..., 0]) * (points[..., 2] / reg_scale) + y1 = points[..., 1] - (0.5 * reg_scale + distance[..., 1]) * (points[..., 3] / reg_scale) + x2 = points[..., 0] + (0.5 * reg_scale + distance[..., 2]) * (points[..., 2] / reg_scale) + y2 = points[..., 1] + (0.5 * reg_scale + distance[..., 3]) * (points[..., 3] / reg_scale) + + bboxes = torch.stack([x1, y1, x2, y2], -1) + + return box_xyxy_to_cxcywh(bboxes) + + +def bbox2distance(points, bbox, reg_max, reg_scale, up, eps=0.1): + """ + Converts bounding box coordinates to distances from a reference point. + + Args: + points (Tensor): (n, 4) [x, y, w, h], where (x, y) is the center. + bbox (Tensor): (n, 4) bounding boxes in "xyxy" format. + reg_max (float): Maximum bin value. + reg_scale (float): Controling curvarture of W(n). + up (Tensor): Controling upper bounds of W(n). + eps (float): Small value to ensure target < reg_max. + + Returns: + Tensor: Decoded distances. + """ + reg_scale = abs(reg_scale) + left = (points[:, 0] - bbox[:, 0]) / (points[..., 2] / reg_scale + 1e-16) - 0.5 * reg_scale + top = (points[:, 1] - bbox[:, 1]) / (points[..., 3] / reg_scale + 1e-16) - 0.5 * reg_scale + right = (bbox[:, 2] - points[:, 0]) / (points[..., 2] / reg_scale + 1e-16) - 0.5 * reg_scale + bottom = (bbox[:, 3] - points[:, 1]) / (points[..., 3] / reg_scale + 1e-16) - 0.5 * reg_scale + four_lens = torch.stack([left, top, right, bottom], -1) + four_lens, weight_right, weight_left = translate_gt(four_lens, reg_max, reg_scale, up) + if reg_max is not None: + four_lens = four_lens.clamp(min=0, max=reg_max-eps) + return four_lens.reshape(-1).detach(), weight_right.detach(), weight_left.detach() diff --git a/engine/deim/hybrid_encoder.py b/engine/deim/hybrid_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..77a74725bc966f130837ff7edbaa7a0730497ff6 --- /dev/null +++ b/engine/deim/hybrid_encoder.py @@ -0,0 +1,498 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE/) +Copyright (c) 2024 D-FINE Authors. All Rights Reserved. +""" + +import copy +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .utils import get_activation + +from ..core import register + +__all__ = ['HybridEncoder'] + + +class ConvNormLayer_fuse(nn.Module): + def __init__(self, ch_in, ch_out, kernel_size, stride, g=1, padding=None, bias=False, act=None): + super().__init__() + padding = (kernel_size-1)//2 if padding is None else padding + self.conv = nn.Conv2d( + ch_in, + ch_out, + kernel_size, + stride, + groups=g, + padding=padding, + bias=bias) + self.norm = nn.BatchNorm2d(ch_out) + self.act = nn.Identity() if act is None else get_activation(act) + self.ch_in, self.ch_out, self.kernel_size, self.stride, self.g, self.padding, self.bias = \ + ch_in, ch_out, kernel_size, stride, g, padding, bias + + def forward(self, x): + if hasattr(self, 'conv_bn_fused'): + y = self.conv_bn_fused(x) + else: + y = self.norm(self.conv(x)) + return self.act(y) + + def convert_to_deploy(self): + if not hasattr(self, 'conv_bn_fused'): + self.conv_bn_fused = nn.Conv2d( + self.ch_in, + self.ch_out, + self.kernel_size, + self.stride, + groups=self.g, + padding=self.padding, + bias=True) + + kernel, bias = self.get_equivalent_kernel_bias() + self.conv_bn_fused.weight.data = kernel + self.conv_bn_fused.bias.data = bias + self.__delattr__('conv') + self.__delattr__('norm') + + def get_equivalent_kernel_bias(self): + kernel3x3, bias3x3 = self._fuse_bn_tensor() + + return kernel3x3, bias3x3 + + def _fuse_bn_tensor(self): + kernel = self.conv.weight + running_mean = self.norm.running_mean + running_var = self.norm.running_var + gamma = self.norm.weight + beta = self.norm.bias + eps = self.norm.eps + std = (running_var + eps).sqrt() + t = (gamma / std).reshape(-1, 1, 1, 1) + return kernel * t, beta - running_mean * gamma / std + + +class ConvNormLayer(nn.Module): + def __init__(self, ch_in, ch_out, kernel_size, stride, g=1, padding=None, bias=False, act=None): + super().__init__() + padding = (kernel_size-1)//2 if padding is None else padding + self.conv = nn.Conv2d( + ch_in, + ch_out, + kernel_size, + stride, + groups=g, + padding=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))) + + +# self.cv1 = Conv(c1, c2, 1, 1) +# self.cv2 = Conv(c2, c2, k=k, s=s, g=c2, act=False) +class SCDown(nn.Module): + def __init__(self, c1, c2, k, s, act=None): + super().__init__() + self.cv1 = ConvNormLayer_fuse(c1, c2, 1, 1) + self.cv2 = ConvNormLayer_fuse(c2, c2, k, s, c2) + + def forward(self, x): + return self.cv2(self.cv1(x)) + + +class VGGBlock(nn.Module): + def __init__(self, ch_in, ch_out, act='relu'): + super().__init__() + self.ch_in = ch_in + self.ch_out = ch_out + self.conv1 = ConvNormLayer(ch_in, ch_out, 3, 1, padding=1, act=None) + self.conv2 = ConvNormLayer(ch_in, ch_out, 1, 1, padding=0, act=None) + self.act = nn.Identity() if act is None else get_activation(act) + + def forward(self, x): + if hasattr(self, 'conv'): + y = self.conv(x) + else: + y = self.conv1(x) + self.conv2(x) + + return self.act(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 + self.__delattr__('conv1') + self.__delattr__('conv2') + + def get_equivalent_kernel_bias(self): + kernel3x3, bias3x3 = self._fuse_bn_tensor(self.conv1) + kernel1x1, bias1x1 = self._fuse_bn_tensor(self.conv2) + + 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: ConvNormLayer): + 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 CSPLayer(nn.Module): + def __init__(self, + in_channels, + out_channels, + num_blocks=3, + expansion=1.0, + bias=False, + act="silu", + bottletype=VGGBlock): + super(CSPLayer, self).__init__() + hidden_channels = int(out_channels * expansion) + self.conv1 = ConvNormLayer_fuse(in_channels, hidden_channels, 1, 1, bias=bias, act=act) + self.conv2 = ConvNormLayer_fuse(in_channels, hidden_channels, 1, 1, bias=bias, act=act) + self.bottlenecks = nn.Sequential(*[ + bottletype(hidden_channels, hidden_channels, act=act) for _ in range(num_blocks) + ]) + if hidden_channels != out_channels: + self.conv3 = ConvNormLayer_fuse(hidden_channels, out_channels, 1, 1, bias=bias, act=act) + else: + self.conv3 = nn.Identity() + + def forward(self, x): + x_2 = self.conv2(x) + x_1 = self.conv1(x) + x_1 = self.bottlenecks(x_1) + return self.conv3(x_1 + x_2) + +class RepNCSPELAN4(nn.Module): + # csp-elan + def __init__(self, c1, c2, c3, c4, n=3, + bias=False, + act="silu", + csp_type='csp2', + ): + super().__init__() + self.c = c3//2 + self.cv1 = ConvNormLayer_fuse(c1, c3, 1, 1, bias=bias, act=act) + if csp_type == 'csp2': + CSPLayer = CSPLayer2 + self.cv2 = nn.Sequential(CSPLayer(c3//2, c4, n, 1, bias=bias, act=act, bottletype=VGGBlock), ConvNormLayer_fuse(c4, c4, 3, 1, bias=bias, act=act)) + self.cv3 = nn.Sequential(CSPLayer(c4, c4, n, 1, bias=bias, act=act, bottletype=VGGBlock), ConvNormLayer_fuse(c4, c4, 3, 1, bias=bias, act=act)) + self.cv4 = ConvNormLayer_fuse(c3+(2*c4), c2, 1, 1, bias=bias, act=act) + + def forward_chunk(self, x): + y = list(self.cv1(x).chunk(2, 1)) + y.extend((m(y[-1])) for m in [self.cv2, self.cv3]) + return self.cv4(torch.cat(y, 1)) + + def forward(self, x): + y = list(self.cv1(x).split((self.c, self.c), 1)) + y.extend(m(y[-1]) for m in [self.cv2, self.cv3]) + return self.cv4(torch.cat(y, 1)) + + +# This layer is equivalent to RepC3 in YOLOs repo +class CSPLayer2(nn.Module): + def __init__(self, + in_channels, + out_channels, + num_blocks=3, + expansion=1.0, + bias=False, + act="silu", + bottletype=VGGBlock, + ): + super(CSPLayer2, self).__init__() + hidden_channels = int(out_channels * expansion) + + self.conv1 = ConvNormLayer_fuse(in_channels, hidden_channels * 2, 1, 1, bias=bias, act=act) + self.bottlenecks = nn.Sequential(*[ + bottletype(hidden_channels, hidden_channels, act=act) for _ in range(num_blocks) + ]) + if hidden_channels != out_channels: + self.conv3 = ConvNormLayer_fuse(hidden_channels, out_channels, 1, 1, bias=bias, act=act) + else: + self.conv3 = nn.Identity() + + def forward(self, x): + y = list(self.conv1(x).chunk(2, 1)) + return self.conv3(y[0] + self.bottlenecks(y[1])) + +class RepNCSPELAN5(nn.Module): + # csp-elan + def __init__(self, c1, c2, c3, c4, n=3, bias=False, act="silu"): + super().__init__() + self.c = c3 // 2 + self.cv1 = ConvNormLayer_fuse(c1, c3, 1, 1, bias=bias, act=act) + + self.cv2 = nn.Sequential(CSPLayer2(c3//2, c4, n, 1, bias=bias, act=act, bottletype=VGGBlock)) + self.cv3 = nn.Sequential(CSPLayer2(c4, c4, n, 1, bias=bias, act=act, bottletype=VGGBlock)) + self.cv4 = ConvNormLayer_fuse(c3+(2*c4), c2, 1, 1, bias=bias, act=act) + + def forward_chunk(self, x): + y = list(self.cv1(x).chunk(2, 1)) + y.extend((m(y[-1])) for m in [self.cv2, self.cv3]) + out = self.cv4(torch.cat(y, 1)) + return out + + def forward(self, x): + y = list(self.cv1(x).split((self.c, self.c), 1)) + y.extend(m(y[-1]) for m in [self.cv2, self.cv3]) + out = self.cv4(torch.cat(y, 1)) + return out + +# transformer +class TransformerEncoderLayer(nn.Module): + def __init__(self, + d_model, + nhead, + dim_feedforward=2048, + dropout=0.1, + activation="relu", + normalize_before=False): + super().__init__() + self.normalize_before = normalize_before + + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout, batch_first=True) + + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.activation = get_activation(activation) + + @staticmethod + def with_pos_embed(tensor, pos_embed): + return tensor if pos_embed is None else tensor + pos_embed + + def forward(self, src, src_mask=None, pos_embed=None) -> torch.Tensor: + residual = src + if self.normalize_before: + src = self.norm1(src) + q = k = self.with_pos_embed(src, pos_embed) + src, _ = self.self_attn(q, k, value=src, attn_mask=src_mask) + + src = residual + self.dropout1(src) + if not self.normalize_before: + src = self.norm1(src) + + residual = src + if self.normalize_before: + src = self.norm2(src) + src = self.linear2(self.dropout(self.activation(self.linear1(src)))) + src = residual + self.dropout2(src) + if not self.normalize_before: + src = self.norm2(src) + return src + + +class TransformerEncoder(nn.Module): + def __init__(self, encoder_layer, num_layers, norm=None): + super(TransformerEncoder, self).__init__() + self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(num_layers)]) + self.num_layers = num_layers + self.norm = norm + + def forward(self, src, src_mask=None, pos_embed=None) -> torch.Tensor: + output = src + for layer in self.layers: + output = layer(output, src_mask=src_mask, pos_embed=pos_embed) + + if self.norm is not None: + output = self.norm(output) + + return output + + +@register() +class HybridEncoder(nn.Module): + __share__ = ['eval_spatial_size', ] + + def __init__(self, + in_channels=[512, 1024, 2048], + feat_strides=[8, 16, 32], + hidden_dim=256, + nhead=8, + dim_feedforward = 1024, + dropout=0.0, + enc_act='gelu', + use_encoder_idx=[2], + num_encoder_layers=1, + pe_temperature=10000, + expansion=1.0, + depth_mult=1.0, + act='silu', + eval_spatial_size=None, + version='dfine', + csp_type='csp', + fuse_op='cat', + ): + super().__init__() + self.in_channels = in_channels + self.feat_strides = feat_strides + self.hidden_dim = hidden_dim + self.use_encoder_idx = use_encoder_idx + self.num_encoder_layers = num_encoder_layers + self.pe_temperature = pe_temperature + self.eval_spatial_size = eval_spatial_size + self.out_channels = [hidden_dim for _ in range(len(in_channels))] + self.out_strides = feat_strides + self.fuse_op = fuse_op + + # channel projection + self.input_proj = nn.ModuleList() + for in_channel in in_channels: + if in_channel != hidden_dim: + proj = nn.Sequential(OrderedDict([ + ('conv', nn.Conv2d(in_channel, hidden_dim, kernel_size=1, bias=False)), + ('norm', nn.BatchNorm2d(hidden_dim)) + ])) + else: + proj = nn.Identity() + self.input_proj.append(proj) + + # encoder transformer + encoder_layer = TransformerEncoderLayer( + hidden_dim, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + activation=enc_act + ) + + self.encoder = nn.ModuleList([ + TransformerEncoder(copy.deepcopy(encoder_layer), num_encoder_layers) for _ in range(len(use_encoder_idx)) + ]) + + input_dim = hidden_dim if self.fuse_op == 'sum' else hidden_dim * 2 # deim use sum instead of cat + + Lateral_Conv = ConvNormLayer_fuse(hidden_dim, hidden_dim, 1, 1) + SCDown_Conv = nn.Sequential(SCDown(hidden_dim, hidden_dim, 3, 2)) + + c1, c2, c3, c4, num_blocks = input_dim, hidden_dim, hidden_dim*2, round(expansion * hidden_dim // 2), round(3 * depth_mult) + if version == 'dfine': + Fuse_Block = RepNCSPELAN4(c1=c1, c2=c2, c3=c3, c4=c4, n=num_blocks, act=act, csp_type=csp_type) + elif version == 'deim': + Fuse_Block = RepNCSPELAN5(c1=c1, c2=c2, c3=c3, c4=c4, n=num_blocks, act=act) + else: # RT-DETR + Fuse_Block = CSPLayer(in_channels=c1, out_channels=c2, num_blocks=num_blocks, act=act, \ + expansion=expansion, bottletype=VGGBlock) + Lateral_Conv = ConvNormLayer_fuse(hidden_dim, hidden_dim, 1, 1, act=act) + SCDown_Conv = ConvNormLayer_fuse(hidden_dim, hidden_dim, 3, 2, act=act) + + # top-down fpn + self.lateral_convs = nn.ModuleList() + self.fpn_blocks = nn.ModuleList() + for _ in range(len(in_channels) - 1, 0, -1): + self.lateral_convs.append(copy.deepcopy(Lateral_Conv)) + self.fpn_blocks.append(copy.deepcopy(Fuse_Block)) + + # bottom-up pan + self.downsample_convs = nn.ModuleList() + self.pan_blocks = nn.ModuleList() + for _ in range(len(in_channels) - 1): + self.downsample_convs.append(copy.deepcopy(SCDown_Conv)) + self.pan_blocks.append(copy.deepcopy(Fuse_Block)) + + self._reset_parameters() + + def _reset_parameters(self): + if self.eval_spatial_size: + for idx in self.use_encoder_idx: + stride = self.feat_strides[idx] + pos_embed = self.build_2d_sincos_position_embedding( + self.eval_spatial_size[1] // stride, self.eval_spatial_size[0] // stride, + self.hidden_dim, self.pe_temperature) + setattr(self, f'pos_embed{idx}', pos_embed) + # self.register_buffer(f'pos_embed{idx}', pos_embed) + + @staticmethod + def build_2d_sincos_position_embedding(w, h, embed_dim=256, temperature=10000.): + """ + """ + grid_w = torch.arange(int(w), dtype=torch.float32) + grid_h = torch.arange(int(h), dtype=torch.float32) + grid_w, grid_h = torch.meshgrid(grid_w, grid_h, indexing='ij') + assert embed_dim % 4 == 0, \ + 'Embed dimension must be divisible by 4 for 2D sin-cos position embedding' + pos_dim = embed_dim // 4 + omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim + omega = 1. / (temperature ** omega) + + out_w = grid_w.flatten()[..., None] @ omega[None] + out_h = grid_h.flatten()[..., None] @ omega[None] + + return torch.concat([out_w.sin(), out_w.cos(), out_h.sin(), out_h.cos()], dim=1)[None, :, :] + + def forward(self, feats): + assert len(feats) == len(self.in_channels) + proj_feats = [self.input_proj[i](feat) for i, feat in enumerate(feats)] + + # encoder + if self.num_encoder_layers > 0: + for i, enc_ind in enumerate(self.use_encoder_idx): + h, w = proj_feats[enc_ind].shape[2:] + # flatten [B, C, H, W] to [B, HxW, C] + src_flatten = proj_feats[enc_ind].flatten(2).permute(0, 2, 1) + if self.training or self.eval_spatial_size is None: + pos_embed = self.build_2d_sincos_position_embedding( + w, h, self.hidden_dim, self.pe_temperature).to(src_flatten.device) + else: + pos_embed = getattr(self, f'pos_embed{enc_ind}', None).to(src_flatten.device) + + memory :torch.Tensor = self.encoder[i](src_flatten, pos_embed=pos_embed) + proj_feats[enc_ind] = memory.permute(0, 2, 1).reshape(-1, self.hidden_dim, h, w).contiguous() + + # broadcasting and fusion + inner_outs = [proj_feats[-1]] + for idx in range(len(self.in_channels) - 1, 0, -1): + feat_heigh = inner_outs[0] + feat_low = proj_feats[idx - 1] + feat_heigh = self.lateral_convs[len(self.in_channels) - 1 - idx](feat_heigh) + inner_outs[0] = feat_heigh + upsample_feat = F.interpolate(feat_heigh, scale_factor=2., mode='nearest') + fused_feat = (upsample_feat + feat_low) \ + if self.fuse_op == 'sum' else torch.concat([upsample_feat, feat_low], dim=1) + inner_out = self.fpn_blocks[len(self.in_channels)-1-idx](fused_feat) + inner_outs.insert(0, inner_out) + + outs = [inner_outs[0]] + for idx in range(len(self.in_channels) - 1): + feat_low = outs[-1] + feat_height = inner_outs[idx + 1] + downsample_feat = self.downsample_convs[idx](feat_low) + fused_feat = (downsample_feat + feat_height) \ + if self.fuse_op == 'sum' else torch.concat([downsample_feat, feat_height], dim=1) + out = self.pan_blocks[idx](fused_feat) + outs.append(out) + + return outs diff --git a/engine/deim/lite_encoder.py b/engine/deim/lite_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..b28b88079d56196138c498c28ef71fbc7e50e12b --- /dev/null +++ b/engine/deim/lite_encoder.py @@ -0,0 +1,107 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE/) +Copyright (c) 2024 D-FINE Authors. All Rights Reserved. +""" + +import copy +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +from functools import partial + +from .utils import get_activation + +from ..core import register +from .hybrid_encoder import ConvNormLayer_fuse +from .hybrid_encoder import RepNCSPELAN4 + +__all__ = ['LiteEncoder'] + + +# Copy from https://github.com/meituan/YOLOv6/blob/main/yolov6/layers/common.py#L695 +class GAP_Fusion(nn.Module): + '''BiFusion Block in PAN''' + def __init__(self, in_channels, out_channels, act=None): + super().__init__() + self.cv = ConvNormLayer_fuse(out_channels, out_channels, 1, 1, act=act) + + def forward(self, x): + # global average pooling + gap = F.adaptive_avg_pool2d(x, 1) + x = x + gap + return self.cv(x) + +# Two-scale encoder +@register() +class LiteEncoder(nn.Module): + __share__ = ['eval_spatial_size', ] + + def __init__(self, + in_channels=[512], + feat_strides=[16], + hidden_dim=256, + expansion=1.0, + depth_mult=1.0, + act='silu', + eval_spatial_size=None, + csp_type='csp2', + ): + super().__init__() + self.in_channels = in_channels + self.feat_strides = feat_strides + self.hidden_dim = hidden_dim + self.eval_spatial_size = eval_spatial_size + self.out_channels = [hidden_dim for _ in range(len(in_channels))] + self.out_strides = feat_strides + + # channel projection: unify the channel dimension of the input features + self.input_proj = nn.ModuleList() + for in_channel in in_channels: + proj = nn.Sequential(OrderedDict([ + ('conv', nn.Conv2d(in_channel, hidden_dim, kernel_size=1, bias=False)), + ('norm', nn.BatchNorm2d(hidden_dim)) + ])) + + self.input_proj.append(proj) + + # get the small-scale feature + down_sample = nn.Sequential( # avg pooling + nn.AvgPool2d(kernel_size=3, stride=2, padding=1), + nn.Conv2d(hidden_dim, hidden_dim, 1, 1, bias=False), + nn.BatchNorm2d(hidden_dim), + get_activation(act) + ) + self.down_sample1 = copy.deepcopy(down_sample) + self.down_sample2 = copy.deepcopy(down_sample) + + # Bi-Fusion + self.bi_fusion = GAP_Fusion(hidden_dim, hidden_dim, act=act) + + # fuse block + c1, c2, c3, c4, num_blocks = hidden_dim, hidden_dim, hidden_dim*2, round(expansion * hidden_dim // 2), round(3 * depth_mult) + fuse_block = RepNCSPELAN4(c1=c1, c2=c2, c3=c3, c4=c4, n=num_blocks, act=act, csp_type=csp_type) + self.fpn_block = copy.deepcopy(fuse_block) + self.pan_block = copy.deepcopy(fuse_block) + + def forward(self, feats): + assert len(feats) == len(self.in_channels) + proj_feats = [self.input_proj[i](feat) for i, feat in enumerate(feats)] + proj_feats.append(self.down_sample1(proj_feats[-1])) # get the small-scale feature + + # fuse the global feature and the small-scale feature + proj_feats[-1] = self.bi_fusion(proj_feats[-1]) + + outs = [] + # fpn + fuse_feat = proj_feats[0] + F.interpolate(proj_feats[1], scale_factor=2., mode='nearest') + outs.append(self.fpn_block(fuse_feat)) + + fuse_feat = proj_feats[1] + self.down_sample2(outs[-1]) + outs.append(self.pan_block(fuse_feat)) + + return outs \ No newline at end of file diff --git a/engine/deim/matcher.py b/engine/deim/matcher.py new file mode 100644 index 0000000000000000000000000000000000000000..bf358aca3181b89f66700390b0fcd26ee2d50a26 --- /dev/null +++ b/engine/deim/matcher.py @@ -0,0 +1,151 @@ +""" +Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +Modules to compute the matching cost and solve the corresponding LSAP. + +Copyright (c) 2024 The D-FINE Authors All Rights Reserved. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from scipy.optimize import linear_sum_assignment +from typing import Dict + +from .box_ops import box_cxcywh_to_xyxy, generalized_box_iou, box_iou + +from ..core import register +import numpy as np + + +@register() +class HungarianMatcher(nn.Module): + """This class computes an assignment between the targets and the predictions of the network + + For efficiency reasons, the targets don't include the no_object. Because of this, in general, + there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, + while the others are un-matched (and thus treated as non-objects). + """ + + __share__ = ['use_focal_loss', ] + + def __init__(self, weight_dict, use_focal_loss=False, alpha=0.25, gamma=2.0, + change_matcher=False, iou_order_alpha=1.0, matcher_change_epoch=10000): + """Creates the matcher + + Params: + cost_class: This is the relative weight of the classification error in the matching cost + cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost + cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost + """ + super().__init__() + self.cost_class = weight_dict['cost_class'] + self.cost_bbox = weight_dict['cost_bbox'] + self.cost_giou = weight_dict['cost_giou'] + + self.change_matcher = change_matcher + self.iou_order_alpha = iou_order_alpha + self.matcher_change_epoch = matcher_change_epoch + if self.change_matcher: + print(f"Using the new matching cost with iou_order_alpha = {iou_order_alpha} at epoch {matcher_change_epoch}") + + self.use_focal_loss = use_focal_loss + self.alpha = alpha + self.gamma = gamma + + assert self.cost_class != 0 or self.cost_bbox != 0 or self.cost_giou != 0, "all costs cant be 0" + + @torch.no_grad() + def forward(self, outputs: Dict[str, torch.Tensor], targets, return_topk=False, epoch=0): + """ Performs the matching + + Params: + outputs: This is a dict that contains at least these entries: + "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits + "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates + + targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing: + "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth + objects in the target) containing the class labels + "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates + + Returns: + A list of size batch_size, containing tuples of (index_i, index_j) where: + - index_i is the indices of the selected predictions (in order) + - index_j is the indices of the corresponding selected targets (in order) + For each batch element, it holds: + len(index_i) = len(index_j) = min(num_queries, num_target_boxes) + """ + bs, num_queries = outputs["pred_logits"].shape[:2] + + # We flatten to compute the cost matrices in a batch + if self.use_focal_loss: + out_prob = F.sigmoid(outputs["pred_logits"].flatten(0, 1)) + else: + out_prob = outputs["pred_logits"].flatten(0, 1).softmax(-1) # [batch_size * num_queries, num_classes] + + out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4] + + # Also concat the target labels and boxes + tgt_ids = torch.cat([v["labels"] for v in targets]) + tgt_bbox = torch.cat([v["boxes"] for v in targets]) + + if self.change_matcher and epoch >= self.matcher_change_epoch: + # Compute the class_score + class_score = out_prob[:, tgt_ids] # shape = [batch_size * num_queries, gt num within a batch] + + # # Compute iou + bbox_iou, _ = box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox)) + + # Final cost matrix + C = (-1) * (class_score * torch.pow(bbox_iou, self.iou_order_alpha)) + else: + # Compute the classification cost. Contrary to the loss, we don't use the NLL, + # but approximate it in 1 - proba[target class]. + # The 1 is a constant that doesn't change the matching, it can be ommitted. + if self.use_focal_loss: + out_prob = out_prob[:, tgt_ids] + neg_cost_class = (1 - self.alpha) * (out_prob ** self.gamma) * (-(1 - out_prob + 1e-8).log()) + pos_cost_class = self.alpha * ((1 - out_prob) ** self.gamma) * (-(out_prob + 1e-8).log()) + cost_class = pos_cost_class - neg_cost_class + else: + cost_class = -out_prob[:, tgt_ids] + + # Compute the L1 cost between boxes + cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) + + # Compute the giou cost betwen boxes + cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox)) + + # Final cost matrix 3 * self.cost_bbox + 2 * self.cost_class + self.cost_giou + C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou + + C = C.view(bs, num_queries, -1).cpu() + + sizes = [len(v["boxes"]) for v in targets] + C = torch.nan_to_num(C, nan=1.0) + indices_pre = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))] + indices = [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices_pre] + + # Compute topk indices + if return_topk: + return {'indices_o2m': self.get_top_k_matches(C, sizes=sizes, k=return_topk, initial_indices=indices_pre)} + + return {'indices': indices} # , 'indices_o2m': C.min(-1)[1]} + + def get_top_k_matches(self, C, sizes, k=1, initial_indices=None): + indices_list = [] + # C_original = C.clone() + for i in range(k): + indices_k = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))] if i > 0 else initial_indices + indices_list.append([ + (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) + for i, j in indices_k + ]) + for c, idx_k in zip(C.split(sizes, -1), indices_k): + idx_k = np.stack(idx_k) + c[:, idx_k] = 1e6 + indices_list = [(torch.cat([indices_list[i][j][0] for i in range(k)], dim=0), + torch.cat([indices_list[i][j][1] for i in range(k)], dim=0)) for j in range(len(sizes))] + # C.copy_(C_original) + return indices_list diff --git a/engine/deim/postprocessor.py b/engine/deim/postprocessor.py new file mode 100644 index 0000000000000000000000000000000000000000..8574592a4b06444ecf7577766b11d340e2e25fa5 --- /dev/null +++ b/engine/deim/postprocessor.py @@ -0,0 +1,92 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import torchvision + +from ..core import register + + +__all__ = ['PostProcessor'] + + +def mod(a, b): + out = a - a // b * b + return out + + +@register() +class PostProcessor(nn.Module): + __share__ = [ + 'num_classes', + 'use_focal_loss', + 'num_top_queries', + 'remap_mscoco_category' + ] + + def __init__( + self, + num_classes=80, + use_focal_loss=True, + num_top_queries=300, + remap_mscoco_category=False + ) -> None: + super().__init__() + self.use_focal_loss = use_focal_loss + self.num_top_queries = num_top_queries + self.num_classes = int(num_classes) + self.remap_mscoco_category = remap_mscoco_category + self.deploy_mode = False + + def extra_repr(self) -> str: + return f'use_focal_loss={self.use_focal_loss}, num_classes={self.num_classes}, num_top_queries={self.num_top_queries}' + + # def forward(self, outputs, orig_target_sizes): + def forward(self, outputs, orig_target_sizes: torch.Tensor): + logits, boxes = outputs['pred_logits'], outputs['pred_boxes'] + # orig_target_sizes = torch.stack([t["orig_size"] for t in targets], dim=0) + + bbox_pred = torchvision.ops.box_convert(boxes, in_fmt='cxcywh', out_fmt='xyxy') + bbox_pred *= orig_target_sizes.repeat(1, 2).unsqueeze(1) + + if self.use_focal_loss: + scores = F.sigmoid(logits) + scores, index = torch.topk(scores.flatten(1), self.num_top_queries, dim=-1) + # labels = index % self.num_classes + labels = mod(index, self.num_classes) + index = index // self.num_classes + boxes = bbox_pred.gather(dim=1, index=index.unsqueeze(-1).repeat(1, 1, bbox_pred.shape[-1])) + + else: + scores = F.softmax(logits)[:, :, :-1] + scores, labels = scores.max(dim=-1) + if scores.shape[1] > self.num_top_queries: + scores, index = torch.topk(scores, self.num_top_queries, dim=-1) + labels = torch.gather(labels, dim=1, index=index) + boxes = torch.gather(boxes, dim=1, index=index.unsqueeze(-1).tile(1, 1, boxes.shape[-1])) + + if self.deploy_mode: + return labels, boxes, scores + + if self.remap_mscoco_category: + from ..data.dataset import mscoco_label2category + labels = torch.tensor([mscoco_label2category[int(x.item())] for x in labels.flatten()])\ + .to(boxes.device).reshape(labels.shape) + + results = [] + for lab, box, sco in zip(labels, boxes, scores): + result = dict(labels=lab, boxes=box, scores=sco) + results.append(result) + + return results + + + def deploy(self, ): + self.eval() + self.deploy_mode = True + return self diff --git a/engine/deim/rtdetrv2_decoder.py b/engine/deim/rtdetrv2_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..d7261f7d5523e0b2628d491c2106c38a33ea9561 --- /dev/null +++ b/engine/deim/rtdetrv2_decoder.py @@ -0,0 +1,622 @@ +"""Copyright(c) 2023 lyuwenyu. All Rights Reserved. +Modifications Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +""" + +import math +import copy +import functools +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.nn.init as init +from typing import List + +from .denoising import get_contrastive_denoising_training_group +from .utils import bias_init_with_prob, get_activation, inverse_sigmoid +from .utils import deformable_attention_core_func_v2 + +from ..core import register + +__all__ = ['RTDETRTransformerv2'] + + +class MLP(nn.Module): + def __init__(self, input_dim, hidden_dim, output_dim, num_layers, act='relu'): + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) + self.act = get_activation(act) + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = self.act(layer(x)) if i < self.num_layers - 1 else layer(x) + return x + + +class MSDeformableAttention(nn.Module): + def __init__( + self, + embed_dim=256, + num_heads=8, + num_levels=4, + num_points=4, + method='default', + offset_scale=0.5, + value_shape='default', + ): + """Multi-Scale Deformable Attention + """ + super(MSDeformableAttention, self).__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.num_levels = num_levels + self.offset_scale = offset_scale + + if isinstance(num_points, list): + assert len(num_points) == num_levels, '' + num_points_list = num_points + else: + num_points_list = [num_points for _ in range(num_levels)] + + self.num_points_list = num_points_list + + num_points_scale = [1/n for n in num_points_list for _ in range(n)] + self.register_buffer('num_points_scale', torch.tensor(num_points_scale, dtype=torch.float32)) + + self.total_points = num_heads * sum(num_points_list) + self.method = method + + self.head_dim = embed_dim // num_heads + assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + + self.sampling_offsets = nn.Linear(embed_dim, self.total_points * 2) + self.attention_weights = nn.Linear(embed_dim, self.total_points) + self.value_proj = nn.Linear(embed_dim, embed_dim) + self.output_proj = nn.Linear(embed_dim, embed_dim) + + self.ms_deformable_attn_core = functools.partial(deformable_attention_core_func_v2, + method=self.method, value_shape=value_shape) + + self._reset_parameters() + + if method == 'discrete': + for p in self.sampling_offsets.parameters(): + p.requires_grad = False + + def _reset_parameters(self): + # sampling_offsets + init.constant_(self.sampling_offsets.weight, 0) + thetas = torch.arange(self.num_heads, dtype=torch.float32) * (2.0 * math.pi / self.num_heads) + grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) + grid_init = grid_init / grid_init.abs().max(-1, keepdim=True).values + grid_init = grid_init.reshape(self.num_heads, 1, 2).tile([1, sum(self.num_points_list), 1]) + scaling = torch.concat([torch.arange(1, n + 1) for n in self.num_points_list]).reshape(1, -1, 1) + grid_init *= scaling + self.sampling_offsets.bias.data[...] = grid_init.flatten() + + # attention_weights + init.constant_(self.attention_weights.weight, 0) + init.constant_(self.attention_weights.bias, 0) + + # proj + init.xavier_uniform_(self.value_proj.weight) + init.constant_(self.value_proj.bias, 0) + init.xavier_uniform_(self.output_proj.weight) + init.constant_(self.output_proj.bias, 0) + + + def forward(self, + query: torch.Tensor, + reference_points: torch.Tensor, + value: torch.Tensor, + value_spatial_shapes: List[int], + value_mask: torch.Tensor=None): + """ + Args: + query (Tensor): [bs, query_length, C] + reference_points (Tensor): [bs, query_length, n_levels, 2], range in [0, 1], top-left (0,0), + bottom-right (1, 1), including padding area + value (Tensor): [bs, value_length, C] + value_spatial_shapes (List): [n_levels, 2], [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})] + value_mask (Tensor): [bs, value_length], True for non-padding elements, False for padding elements + + Returns: + output (Tensor): [bs, Length_{query}, C] + """ + bs, Len_q = query.shape[:2] + Len_v = value.shape[1] + + value = self.value_proj(value) + if value_mask is not None: + value = value * value_mask.to(value.dtype).unsqueeze(-1) + + value = value.reshape(bs, Len_v, self.num_heads, self.head_dim) + + sampling_offsets: torch.Tensor = self.sampling_offsets(query) + sampling_offsets = sampling_offsets.reshape(bs, Len_q, self.num_heads, sum(self.num_points_list), 2) + + attention_weights = self.attention_weights(query).reshape(bs, Len_q, self.num_heads, sum(self.num_points_list)) + attention_weights = F.softmax(attention_weights, dim=-1).reshape(bs, Len_q, self.num_heads, sum(self.num_points_list)) + + if reference_points.shape[-1] == 2: + offset_normalizer = torch.tensor(value_spatial_shapes) + offset_normalizer = offset_normalizer.flip([1]).reshape(1, 1, 1, self.num_levels, 1, 2) + sampling_locations = reference_points.reshape(bs, Len_q, 1, self.num_levels, 1, 2) + sampling_offsets / offset_normalizer + elif reference_points.shape[-1] == 4: + # reference_points [8, 480, None, 1, 4] + # sampling_offsets [8, 480, 8, 12, 2] + num_points_scale = self.num_points_scale.to(dtype=query.dtype).unsqueeze(-1) + offset = sampling_offsets * num_points_scale * reference_points[:, :, None, :, 2:] * self.offset_scale + sampling_locations = reference_points[:, :, None, :, :2] + offset + else: + raise ValueError( + "Last dim of reference_points must be 2 or 4, but get {} instead.". + format(reference_points.shape[-1])) + + output = self.ms_deformable_attn_core(value, value_spatial_shapes, sampling_locations, attention_weights, self.num_points_list) + + output = self.output_proj(output) + + return output + + +class TransformerDecoderLayer(nn.Module): + def __init__(self, + d_model=256, + n_head=8, + dim_feedforward=1024, + dropout=0., + activation='relu', + n_levels=4, + n_points=4, + cross_attn_method='default', + value_shape='default', + ): + super(TransformerDecoderLayer, self).__init__() + + # self attention + self.self_attn = nn.MultiheadAttention(d_model, n_head, dropout=dropout, batch_first=True) + self.dropout1 = nn.Dropout(dropout) + self.norm1 = nn.LayerNorm(d_model) + + # cross attention + self.cross_attn = MSDeformableAttention(d_model, n_head, n_levels, n_points, method=cross_attn_method, value_shape=value_shape) + self.dropout2 = nn.Dropout(dropout) + self.norm2 = nn.LayerNorm(d_model) + + # ffn + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.activation = get_activation(activation) + self.dropout3 = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + self.dropout4 = nn.Dropout(dropout) + self.norm3 = nn.LayerNorm(d_model) + + self._reset_parameters() + + def _reset_parameters(self): + init.xavier_uniform_(self.linear1.weight) + init.xavier_uniform_(self.linear2.weight) + + def with_pos_embed(self, tensor, pos): + return tensor if pos is None else tensor + pos + + def forward_ffn(self, tgt): + return self.linear2(self.dropout3(self.activation(self.linear1(tgt)))) + + def forward(self, + target, + reference_points, + memory, + memory_spatial_shapes, + attn_mask=None, + memory_mask=None, + query_pos_embed=None): + # self attention + q = k = self.with_pos_embed(target, query_pos_embed) + + target2, _ = self.self_attn(q, k, value=target, attn_mask=attn_mask) + target = target + self.dropout1(target2) + target = self.norm1(target) + + # cross attention + target2 = self.cross_attn(\ + self.with_pos_embed(target, query_pos_embed), + reference_points, + memory, + memory_spatial_shapes, + memory_mask) + target = target + self.dropout2(target2) + target = self.norm2(target) + + # ffn + target2 = self.forward_ffn(target) + target = target + self.dropout4(target2) + target = self.norm3(target) + + return target + + +class TransformerDecoder(nn.Module): + def __init__(self, hidden_dim, decoder_layer, num_layers, eval_idx=-1): + super(TransformerDecoder, self).__init__() + self.layers = nn.ModuleList([copy.deepcopy(decoder_layer) for _ in range(num_layers)]) + self.hidden_dim = hidden_dim + self.num_layers = num_layers + self.eval_idx = eval_idx if eval_idx >= 0 else num_layers + eval_idx + + def forward(self, + target, + ref_points_unact, + memory, + memory_spatial_shapes, + bbox_head, + score_head, + query_pos_head, + attn_mask=None, + memory_mask=None): + dec_out_bboxes = [] + dec_out_logits = [] + ref_points_detach = F.sigmoid(ref_points_unact) + + output = target + for i, layer in enumerate(self.layers): + ref_points_input = ref_points_detach.unsqueeze(2) + query_pos_embed = query_pos_head(ref_points_detach) + + output = layer(output, ref_points_input, memory, memory_spatial_shapes, attn_mask, memory_mask, query_pos_embed) + + inter_ref_bbox = F.sigmoid(bbox_head[i](output) + inverse_sigmoid(ref_points_detach)) + + if self.training: + dec_out_logits.append(score_head[i](output)) + if i == 0: + dec_out_bboxes.append(inter_ref_bbox) + else: + dec_out_bboxes.append(F.sigmoid(bbox_head[i](output) + inverse_sigmoid(ref_points))) + + elif i == self.eval_idx: + dec_out_logits.append(score_head[i](output)) + dec_out_bboxes.append(inter_ref_bbox) + break + + ref_points = inter_ref_bbox + ref_points_detach = inter_ref_bbox.detach() + + return torch.stack(dec_out_bboxes), torch.stack(dec_out_logits) + + +@register() +class RTDETRTransformerv2(nn.Module): + __share__ = ['num_classes', 'eval_spatial_size'] + + def __init__(self, + num_classes=80, + hidden_dim=256, + num_queries=300, + feat_channels=[512, 1024, 2048], + feat_strides=[8, 16, 32], + num_levels=3, + num_points=4, + nhead=8, + num_layers=6, + dim_feedforward=1024, + dropout=0., + activation="relu", + num_denoising=100, + label_noise_ratio=0.5, + box_noise_scale=1.0, + learn_query_content=False, + eval_spatial_size=None, + eval_idx=-1, + eps=1e-2, + aux_loss=True, + cross_attn_method='default', + query_select_method='default', + value_shape='reshape', + mlp_act='relu', + query_pos_method='default', + ): + super().__init__() + assert len(feat_channels) <= num_levels + assert len(feat_strides) == len(feat_channels) + + for _ in range(num_levels - len(feat_strides)): + feat_strides.append(feat_strides[-1] * 2) + + self.hidden_dim = hidden_dim + self.nhead = nhead + self.feat_strides = feat_strides + self.num_levels = num_levels + self.num_classes = num_classes + self.num_queries = num_queries + self.eps = eps + self.num_layers = num_layers + self.eval_spatial_size = eval_spatial_size + self.aux_loss = aux_loss + + assert query_select_method in ('default', 'one2many', 'agnostic'), '' + assert cross_attn_method in ('default', 'discrete'), '' + self.cross_attn_method = cross_attn_method + self.query_select_method = query_select_method + + # backbone feature projection + self._build_input_proj_layer(feat_channels) + + # Transformer module + decoder_layer = TransformerDecoderLayer(hidden_dim, nhead, dim_feedforward, dropout, \ + activation, num_levels, num_points, cross_attn_method=cross_attn_method, value_shape=value_shape) + self.decoder = TransformerDecoder(hidden_dim, decoder_layer, num_layers, eval_idx) + + # denoising + self.num_denoising = num_denoising + self.label_noise_ratio = label_noise_ratio + self.box_noise_scale = box_noise_scale + if num_denoising > 0: + self.denoising_class_embed = nn.Embedding(num_classes+1, hidden_dim, padding_idx=num_classes) + init.normal_(self.denoising_class_embed.weight[:-1]) + + # decoder embedding + self.learn_query_content = learn_query_content + if learn_query_content: + self.tgt_embed = nn.Embedding(num_queries, hidden_dim) + + if query_pos_method == 'as_reg': + self.query_pos_head = MLP(4, hidden_dim, hidden_dim, 3, act=mlp_act) + print(" ### Query Position Embedding@{} ### ".format(query_pos_method)) + else: + self.query_pos_head = MLP(4, 2 * hidden_dim, hidden_dim, 2, act=mlp_act) + + # if num_select_queries != self.num_queries: + # layer = TransformerEncoderLayer(hidden_dim, nhead, dim_feedforward, activation='gelu') + # self.encoder = TransformerEncoder(layer, 1) + + self.enc_output = nn.Sequential(OrderedDict([ + ('proj', nn.Linear(hidden_dim, hidden_dim)), + ('norm', nn.LayerNorm(hidden_dim,)), + ])) + + if query_select_method == 'agnostic': + self.enc_score_head = nn.Linear(hidden_dim, 1) + else: + self.enc_score_head = nn.Linear(hidden_dim, num_classes) + + self.enc_bbox_head = MLP(hidden_dim, hidden_dim, 4, 3, act=mlp_act) + + # decoder head + self.dec_score_head = nn.ModuleList([ + nn.Linear(hidden_dim, num_classes) for _ in range(num_layers) + ]) + self.dec_bbox_head = nn.ModuleList([ + MLP(hidden_dim, hidden_dim, 4, 3, act=mlp_act) for _ in range(num_layers) + ]) + + # init encoder output anchors and valid_mask + if self.eval_spatial_size: + anchors, valid_mask = self._generate_anchors() + self.register_buffer('anchors', anchors) + self.register_buffer('valid_mask', valid_mask) + + self._reset_parameters() + + def _reset_parameters(self): + bias = bias_init_with_prob(0.01) + init.constant_(self.enc_score_head.bias, bias) + init.constant_(self.enc_bbox_head.layers[-1].weight, 0) + init.constant_(self.enc_bbox_head.layers[-1].bias, 0) + + for _cls, _reg in zip(self.dec_score_head, self.dec_bbox_head): + init.constant_(_cls.bias, bias) + init.constant_(_reg.layers[-1].weight, 0) + init.constant_(_reg.layers[-1].bias, 0) + + init.xavier_uniform_(self.enc_output[0].weight) + if self.learn_query_content: + init.xavier_uniform_(self.tgt_embed.weight) + init.xavier_uniform_(self.query_pos_head.layers[0].weight) + init.xavier_uniform_(self.query_pos_head.layers[1].weight) + for m in self.input_proj: + init.xavier_uniform_(m[0].weight) + + def _build_input_proj_layer(self, feat_channels): + self.input_proj = nn.ModuleList() + for in_channels in feat_channels: + self.input_proj.append( + nn.Sequential(OrderedDict([ + ('conv', nn.Conv2d(in_channels, self.hidden_dim, 1, bias=False)), + ('norm', nn.BatchNorm2d(self.hidden_dim,))]) + ) + ) + + in_channels = feat_channels[-1] + + for _ in range(self.num_levels - len(feat_channels)): + self.input_proj.append( + nn.Sequential(OrderedDict([ + ('conv', nn.Conv2d(in_channels, self.hidden_dim, 3, 2, padding=1, bias=False)), + ('norm', nn.BatchNorm2d(self.hidden_dim))]) + ) + ) + in_channels = self.hidden_dim + + def _get_encoder_input(self, feats: List[torch.Tensor]): + # get projection features + proj_feats = [self.input_proj[i](feat) for i, feat in enumerate(feats)] + if self.num_levels > len(proj_feats): + len_srcs = len(proj_feats) + for i in range(len_srcs, self.num_levels): + if i == len_srcs: + proj_feats.append(self.input_proj[i](feats[-1])) + else: + proj_feats.append(self.input_proj[i](proj_feats[-1])) + + # get encoder inputs + feat_flatten = [] + spatial_shapes = [] + for i, feat in enumerate(proj_feats): + _, _, h, w = feat.shape + # [b, c, h, w] -> [b, h*w, c] + feat_flatten.append(feat.flatten(2).permute(0, 2, 1)) + # [num_levels, 2] + spatial_shapes.append([h, w]) + # [b, l, c] + feat_flatten = torch.concat(feat_flatten, 1) + return feat_flatten, spatial_shapes + + def _generate_anchors(self, + spatial_shapes=None, + grid_size=0.05, + dtype=torch.float32, + device='cpu'): + if spatial_shapes is None: + spatial_shapes = [] + eval_h, eval_w = self.eval_spatial_size + for s in self.feat_strides: + spatial_shapes.append([int(eval_h / s), int(eval_w / s)]) + + anchors = [] + for lvl, (h, w) in enumerate(spatial_shapes): + grid_y, grid_x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing='ij') + grid_xy = torch.stack([grid_x, grid_y], dim=-1) + grid_xy = (grid_xy.unsqueeze(0) + 0.5) / torch.tensor([w, h], dtype=dtype) + wh = torch.ones_like(grid_xy) * grid_size * (2.0 ** lvl) + lvl_anchors = torch.concat([grid_xy, wh], dim=-1).reshape(-1, h * w, 4) + anchors.append(lvl_anchors) + + anchors = torch.concat(anchors, dim=1).to(device) + valid_mask = ((anchors > self.eps) * (anchors < 1 - self.eps)).all(-1, keepdim=True) + anchors = torch.log(anchors / (1 - anchors)) + anchors = torch.where(valid_mask, anchors, torch.inf) + + return anchors, valid_mask + + + def _get_decoder_input(self, + memory: torch.Tensor, + spatial_shapes, + denoising_logits=None, + denoising_bbox_unact=None): + + # prepare input for decoder + if self.training or self.eval_spatial_size is None: + anchors, valid_mask = self._generate_anchors(spatial_shapes, device=memory.device) + else: + anchors = self.anchors + valid_mask = self.valid_mask + + # memory = torch.where(valid_mask, memory, 0) + memory = valid_mask.to(memory.dtype) * memory + + output_memory :torch.Tensor = self.enc_output(memory) + enc_outputs_logits :torch.Tensor = self.enc_score_head(output_memory) + enc_outputs_coord_unact :torch.Tensor = self.enc_bbox_head(output_memory) + anchors + + enc_topk_bboxes_list, enc_topk_logits_list = [], [] + enc_topk_memory, enc_topk_logits, enc_topk_bbox_unact = \ + self._select_topk(output_memory, enc_outputs_logits, enc_outputs_coord_unact, self.num_queries) + + if self.training: + enc_topk_bboxes = F.sigmoid(enc_topk_bbox_unact) + enc_topk_bboxes_list.append(enc_topk_bboxes) + enc_topk_logits_list.append(enc_topk_logits) + + # if self.num_select_queries != self.num_queries: + # raise NotImplementedError('') + + if self.learn_query_content: + content = self.tgt_embed.weight.unsqueeze(0).tile([memory.shape[0], 1, 1]) + else: + content = enc_topk_memory.detach() + + enc_topk_bbox_unact = enc_topk_bbox_unact.detach() + + if denoising_bbox_unact is not None: + enc_topk_bbox_unact = torch.concat([denoising_bbox_unact, enc_topk_bbox_unact], dim=1) + content = torch.concat([denoising_logits, content], dim=1) + + return content, enc_topk_bbox_unact, enc_topk_bboxes_list, enc_topk_logits_list + + def _select_topk(self, memory: torch.Tensor, outputs_logits: torch.Tensor, outputs_coords_unact: torch.Tensor, topk: int): + if self.query_select_method == 'default': + _, topk_ind = torch.topk(outputs_logits.max(-1).values, topk, dim=-1) + + elif self.query_select_method == 'one2many': + _, topk_ind = torch.topk(outputs_logits.flatten(1), topk, dim=-1) + topk_ind = topk_ind // self.num_classes + + elif self.query_select_method == 'agnostic': + _, topk_ind = torch.topk(outputs_logits.squeeze(-1), topk, dim=-1) + + topk_ind: torch.Tensor + + topk_coords = outputs_coords_unact.gather(dim=1, \ + index=topk_ind.unsqueeze(-1).repeat(1, 1, outputs_coords_unact.shape[-1])) + + topk_logits = outputs_logits.gather(dim=1, \ + index=topk_ind.unsqueeze(-1).repeat(1, 1, outputs_logits.shape[-1])) + + topk_memory = memory.gather(dim=1, \ + index=topk_ind.unsqueeze(-1).repeat(1, 1, memory.shape[-1])) + + return topk_memory, topk_logits, topk_coords + + + def forward(self, feats, targets=None): + # input projection and embedding + memory, spatial_shapes = self._get_encoder_input(feats) + + # prepare denoising training + if self.training and self.num_denoising > 0: + denoising_logits, denoising_bbox_unact, attn_mask, dn_meta = \ + get_contrastive_denoising_training_group(targets, \ + self.num_classes, + self.num_queries, + self.denoising_class_embed, + num_denoising=self.num_denoising, + label_noise_ratio=self.label_noise_ratio, + box_noise_scale=self.box_noise_scale, ) + else: + denoising_logits, denoising_bbox_unact, attn_mask, dn_meta = None, None, None, None + + init_ref_contents, init_ref_points_unact, enc_topk_bboxes_list, enc_topk_logits_list = \ + self._get_decoder_input(memory, spatial_shapes, denoising_logits, denoising_bbox_unact) + + # decoder + out_bboxes, out_logits = self.decoder( + init_ref_contents, + init_ref_points_unact, + memory, + spatial_shapes, + self.dec_bbox_head, + self.dec_score_head, + self.query_pos_head, + attn_mask=attn_mask) + + if self.training and dn_meta is not None: + dn_out_bboxes, out_bboxes = torch.split(out_bboxes, dn_meta['dn_num_split'], dim=2) + dn_out_logits, out_logits = torch.split(out_logits, dn_meta['dn_num_split'], dim=2) + + out = {'pred_logits': out_logits[-1], 'pred_boxes': out_bboxes[-1]} + + if self.training and self.aux_loss: + out['aux_outputs'] = self._set_aux_loss(out_logits[:-1], out_bboxes[:-1]) + out['enc_aux_outputs'] = self._set_aux_loss(enc_topk_logits_list, enc_topk_bboxes_list) + out['enc_meta'] = {'class_agnostic': self.query_select_method == 'agnostic'} + + if dn_meta is not None: + out['dn_outputs'] = self._set_aux_loss(dn_out_logits, dn_out_bboxes) + out['dn_meta'] = dn_meta + + return out + + + @torch.jit.unused + def _set_aux_loss(self, outputs_class, outputs_coord): + # this is a workaround to make torchscript happy, as torchscript + # doesn't support dictionary with non-homogeneous values, such + # as a dict having both a Tensor and a list. + return [{'pred_logits': a, 'pred_boxes': b} + for a, b in zip(outputs_class, outputs_coord)] \ No newline at end of file diff --git a/engine/deim/utils.py b/engine/deim/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e8fe1eb42d872d6550083764daf35b2cb3f91fd4 --- /dev/null +++ b/engine/deim/utils.py @@ -0,0 +1,181 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE) +Copyright (c) 2023 . All Rights Reserved. +""" + +import math +from typing import List + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def inverse_sigmoid(x: torch.Tensor, eps: float=1e-5) -> torch.Tensor: + x = x.clip(min=0., max=1.) + return torch.log(x.clip(min=eps) / (1 - x).clip(min=eps)) + + +def bias_init_with_prob(prior_prob=0.01): + """initialize conv/fc bias value according to a given probability value.""" + bias_init = float(-math.log((1 - prior_prob) / prior_prob)) + return bias_init + + +def deformable_attention_core_func(value, value_spatial_shapes, sampling_locations, attention_weights): + """ + Args: + value (Tensor): [bs, value_length, n_head, c] + value_spatial_shapes (Tensor|List): [n_levels, 2] + value_level_start_index (Tensor|List): [n_levels] + sampling_locations (Tensor): [bs, query_length, n_head, n_levels, n_points, 2] + attention_weights (Tensor): [bs, query_length, n_head, n_levels, n_points] + + Returns: + output (Tensor): [bs, Length_{query}, C] + """ + bs, _, n_head, c = value.shape + _, Len_q, _, n_levels, n_points, _ = sampling_locations.shape + + split_shape = [h * w for h, w in value_spatial_shapes] + value_list = value.split(split_shape, dim=1) + sampling_grids = 2 * sampling_locations - 1 + sampling_value_list = [] + for level, (h, w) in enumerate(value_spatial_shapes): + # N_, H_*W_, M_, D_ -> N_, H_*W_, M_*D_ -> N_, M_*D_, H_*W_ -> N_*M_, D_, H_, W_ + value_l_ = value_list[level].flatten(2).permute( + 0, 2, 1).reshape(bs * n_head, c, h, w) + # N_, Lq_, M_, P_, 2 -> N_, M_, Lq_, P_, 2 -> N_*M_, Lq_, P_, 2 + sampling_grid_l_ = sampling_grids[:, :, :, level].permute( + 0, 2, 1, 3, 4).flatten(0, 1) + # N_*M_, D_, Lq_, P_ + sampling_value_l_ = F.grid_sample( + value_l_, + sampling_grid_l_, + mode='bilinear', + padding_mode='zeros', + align_corners=False) + sampling_value_list.append(sampling_value_l_) + # (N_, Lq_, M_, L_, P_) -> (N_, M_, Lq_, L_, P_) -> (N_*M_, 1, Lq_, L_*P_) + attention_weights = attention_weights.permute(0, 2, 1, 3, 4).reshape( + bs * n_head, 1, Len_q, n_levels * n_points) + output = (torch.stack( + sampling_value_list, dim=-2).flatten(-2) * + attention_weights).sum(-1).reshape(bs, n_head * c, Len_q) + + return output.permute(0, 2, 1) + + + +def deformable_attention_core_func_v2(\ + value: torch.Tensor, + value_spatial_shapes, + sampling_locations: torch.Tensor, + attention_weights: torch.Tensor, + num_points_list: List[int], + method='default', + value_shape='default', + ): + """ + Args: + value (Tensor): [bs, value_length, n_head, c] + value_spatial_shapes (Tensor|List): [n_levels, 2] + value_level_start_index (Tensor|List): [n_levels] + sampling_locations (Tensor): [bs, query_length, n_head, n_levels * n_points, 2] + attention_weights (Tensor): [bs, query_length, n_head, n_levels * n_points] + + Returns: + output (Tensor): [bs, Length_{query}, C] + """ + if value_shape == 'default': + bs, n_head, c, _ = value[0].shape + elif value_shape == 'reshape': # reshape following RT-DETR + bs, _, n_head, c = value.shape + split_shape = [h * w for h, w in value_spatial_shapes] + value = value.permute(0, 2, 3, 1).flatten(0, 1).split(split_shape, dim=-1) + _, Len_q, _, _, _ = sampling_locations.shape + + # sampling_offsets [8, 480, 8, 12, 2] + if method == 'default': + sampling_grids = 2 * sampling_locations - 1 + + elif method == 'discrete': + sampling_grids = sampling_locations + + sampling_grids = sampling_grids.permute(0, 2, 1, 3, 4).flatten(0, 1) + sampling_locations_list = sampling_grids.split(num_points_list, dim=-2) + + sampling_value_list = [] + for level, (h, w) in enumerate(value_spatial_shapes): + value_l = value[level].reshape(bs * n_head, c, h, w) + sampling_grid_l: torch.Tensor = sampling_locations_list[level] + + if method == 'default': + sampling_value_l = F.grid_sample( + value_l, + sampling_grid_l, + mode='bilinear', + padding_mode='zeros', + align_corners=False) + + elif method == 'discrete': + # n * m, seq, n, 2 + sampling_coord = (sampling_grid_l * torch.tensor([[w, h]], device=value_l.device) + 0.5).to(torch.int64) + + # FIX ME? for rectangle input + sampling_coord = sampling_coord.clamp(0, h - 1) + sampling_coord = sampling_coord.reshape(bs * n_head, Len_q * num_points_list[level], 2) + + s_idx = torch.arange(sampling_coord.shape[0], device=value_l.device).unsqueeze(-1).repeat(1, sampling_coord.shape[1]) + sampling_value_l: torch.Tensor = value_l[s_idx, :, sampling_coord[..., 1], sampling_coord[..., 0]] # n l c + + sampling_value_l = sampling_value_l.permute(0, 2, 1).reshape(bs * n_head, c, Len_q, num_points_list[level]) + + sampling_value_list.append(sampling_value_l) + + attn_weights = attention_weights.permute(0, 2, 1, 3).reshape(bs * n_head, 1, Len_q, sum(num_points_list)) + weighted_sample_locs = torch.concat(sampling_value_list, dim=-1) * attn_weights + output = weighted_sample_locs.sum(-1).reshape(bs, n_head * c, Len_q) + + return output.permute(0, 2, 1) + + +def get_activation(act: str, inpace: 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 = inpace + + return m diff --git a/engine/misc/__init__.py b/engine/misc/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..acd6b469a789ea7f9e8ced6656e60e7a9c1cfb99 --- /dev/null +++ b/engine/misc/__init__.py @@ -0,0 +1,9 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +from .logger import * +from .visualizer import * +from .dist_utils import setup_seed, setup_print +from .profiler_utils import stats diff --git a/engine/misc/box_ops.py b/engine/misc/box_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..bdaa0cf8a10dcdaa6c15e30a543c30d3c063830c --- /dev/null +++ b/engine/misc/box_ops.py @@ -0,0 +1,105 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +import torch +import torchvision +from torch import Tensor +from typing import List, Tuple + + +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 diff --git a/engine/misc/dist_utils.py b/engine/misc/dist_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..368d4353a2ed213113b42bcd2b9f7bd595b27ede --- /dev/null +++ b/engine/misc/dist_utils.py @@ -0,0 +1,268 @@ +""" +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 os +import time +import random +import numpy as np +import atexit + +import torch +import torch.nn as nn +import torch.distributed +import torch.backends.cudnn + +from torch.nn.parallel import DataParallel as DP +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + +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 torch + import warnings + 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 diff --git a/engine/misc/lazy_loader.py b/engine/misc/lazy_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..e99ce5995f9ab8f4181623605e41eb4174596959 --- /dev/null +++ b/engine/misc/lazy_loader.py @@ -0,0 +1,70 @@ +""" +https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/util/lazy_loader.py +""" + + +import types +import importlib + +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"" + + 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__() diff --git a/engine/misc/logger.py b/engine/misc/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..fd020905fc09e31768c28b8b9fc6b30f9ab8c846 --- /dev/null +++ b/engine/misc/logger.py @@ -0,0 +1,238 @@ +""" +# 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 time +import pickle +import datetime +from collections import defaultdict, deque +from typing import Dict + +import torch +import torch.distributed as tdist + +from .dist_utils import is_dist_available_and_initialized, get_world_size + + +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))) diff --git a/engine/misc/profiler_utils.py b/engine/misc/profiler_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f328cc35836f641ca388c529d1807a4fee224981 --- /dev/null +++ b/engine/misc/profiler_utils.py @@ -0,0 +1,26 @@ +""" +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +""" + +import copy +from calflops import calculate_flops +from typing import Tuple + +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)} diff --git a/engine/misc/visualizer.py b/engine/misc/visualizer.py new file mode 100644 index 0000000000000000000000000000000000000000..4e14eef99c17a688d941cca5ac7baf18ac4f42e3 --- /dev/null +++ b/engine/misc/visualizer.py @@ -0,0 +1,33 @@ +"""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +import PIL +import torch +import torch.utils.data +import torchvision +torchvision.disable_beta_transforms_warning() + +__all__ = ['show_sample'] + +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() diff --git a/engine/optim/__init__.py b/engine/optim/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a6b3bc6a619c7c56fe5b15784bed6cd800c3aeb2 --- /dev/null +++ b/engine/optim/__init__.py @@ -0,0 +1,9 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +from .ema import * +from .optim import * +from .amp import * +from .warmup import * diff --git a/engine/optim/amp.py b/engine/optim/amp.py new file mode 100644 index 0000000000000000000000000000000000000000..6af85e5cd110d932a0534a82b5e93a8a467dbce0 --- /dev/null +++ b/engine/optim/amp.py @@ -0,0 +1,14 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + + +import torch.cuda.amp as amp + +from ..core import register + + +__all__ = ['GradScaler'] + +GradScaler = register()(amp.grad_scaler.GradScaler) diff --git a/engine/optim/ema.py b/engine/optim/ema.py new file mode 100644 index 0000000000000000000000000000000000000000..1c234347d8e640a54fe7925ee23aa22c5155f3fd --- /dev/null +++ b/engine/optim/ema.py @@ -0,0 +1,102 @@ +""" +D-FINE: Redefine Regression Task of DETRs as Fine-grained Distribution Refinement +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright (c) 2023 lyuwenyu. All Rights Reserved. +""" + + +import torch +import torch.nn as nn + +import math +from copy import deepcopy + +from ..core import register +from ..misc import dist_utils + +__all__ = ['ModelEMA'] + + +@register() +class ModelEMA(object): + """ + Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models + Keep a moving average of everything in the model state_dict (parameters and buffers). + This is intended to allow functionality like + https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage + A smoothed version of the weights is necessary for some training schemes to perform well. + This class is sensitive where it is initialized in the sequence of model init, + GPU assignment and distributed training wrappers. + """ + def __init__(self, model: nn.Module, decay: float=0.9999, warmups: int=1000, start: int=0): + super().__init__() + + self.module = deepcopy(dist_utils.de_parallel(model)).eval() + # if next(model.parameters()).device.type != 'cpu': + # self.module.half() # FP16 EMA + + self.decay = decay + self.warmups = warmups + self.before_start = 0 + self.start = start + self.updates = 0 # number of EMA updates + if warmups == 0: + self.decay_fn = lambda x: decay + else: + self.decay_fn = lambda x: decay * (1 - math.exp(-x / warmups)) # decay exponential ramp (to help early epochs) + + for p in self.module.parameters(): + p.requires_grad_(False) + + + def update(self, model: nn.Module): + if self.before_start < self.start: + self.before_start += 1 + return + # Update EMA parameters + with torch.no_grad(): + self.updates += 1 + d = self.decay_fn(self.updates) + msd = dist_utils.de_parallel(model).state_dict() + for k, v in self.module.state_dict().items(): + if v.dtype.is_floating_point: + v *= d + v += (1 - d) * msd[k].detach() + + def to(self, *args, **kwargs): + self.module = self.module.to(*args, **kwargs) + return self + + def state_dict(self, ): + return dict(module=self.module.state_dict(), updates=self.updates) + + def load_state_dict(self, state, strict=True): + self.module.load_state_dict(state['module'], strict=strict) + if 'updates' in state: + self.updates = state['updates'] + + def forwad(self, ): + raise RuntimeError('ema...') + + def extra_repr(self) -> str: + return f'decay={self.decay}, warmups={self.warmups}' + + + +class ExponentialMovingAverage(torch.optim.swa_utils.AveragedModel): + """Maintains moving averages of model parameters using an exponential decay. + ``ema_avg = decay * avg_model_param + (1 - decay) * model_param`` + `torch.optim.swa_utils.AveragedModel `_ + is used to compute the EMA. + """ + def __init__(self, model, decay, device="cpu", use_buffers=True): + + self.decay_fn = lambda x: decay * (1 - math.exp(-x / 2000)) + + def ema_avg(avg_model_param, model_param, num_averaged): + decay = self.decay_fn(num_averaged) + return decay * avg_model_param + (1 - decay) * model_param + + super().__init__(model, device, ema_avg, use_buffers=use_buffers) diff --git a/engine/optim/lr_scheduler.py b/engine/optim/lr_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..b5902b3497311c944a762361e1b199665cbe1184 --- /dev/null +++ b/engine/optim/lr_scheduler.py @@ -0,0 +1,73 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +""" + +import math +from functools import partial + + +def flat_cosine_schedule(total_iter, warmup_iter, flat_iter, no_aug_iter, current_iter, init_lr, min_lr): + """ + Computes the learning rate using a warm-up, flat, and cosine decay schedule. + + Args: + total_iter (int): Total number of iterations. + warmup_iter (int): Number of iterations for warm-up phase. + flat_iter (int): Number of iterations for flat phase. + no_aug_iter (int): Number of iterations for no-augmentation phase. + current_iter (int): Current iteration. + init_lr (float): Initial learning rate. + min_lr (float): Minimum learning rate. + + Returns: + float: Calculated learning rate. + """ + if current_iter <= warmup_iter: + return init_lr * (current_iter / float(warmup_iter)) ** 2 + elif warmup_iter < current_iter <= flat_iter: + return init_lr + elif current_iter >= total_iter - no_aug_iter: + return min_lr + else: + cosine_decay = 0.5 * (1 + math.cos(math.pi * (current_iter - flat_iter) / + (total_iter - flat_iter - no_aug_iter))) + return min_lr + (init_lr - min_lr) * cosine_decay + + +class FlatCosineLRScheduler: + """ + Learning rate scheduler with warm-up, optional flat phase, and cosine decay following RTMDet. + + Args: + optimizer (torch.optim.Optimizer): Optimizer instance. + lr_gamma (float): Scaling factor for the minimum learning rate. + iter_per_epoch (int): Number of iterations per epoch. + total_epochs (int): Total number of training epochs. + warmup_epochs (int): Number of warm-up epochs. + flat_epochs (int): Number of flat epochs (for flat-cosine scheduler). + no_aug_epochs (int): Number of no-augmentation epochs. + """ + def __init__(self, optimizer, lr_gamma, iter_per_epoch, total_epochs, + warmup_iter, flat_epochs, no_aug_epochs, scheduler_type="cosine"): + self.base_lrs = [group["initial_lr"] for group in optimizer.param_groups] + self.min_lrs = [base_lr * lr_gamma for base_lr in self.base_lrs] + + total_iter = int(iter_per_epoch * total_epochs) + no_aug_iter = int(iter_per_epoch * no_aug_epochs) + flat_iter = int(iter_per_epoch * flat_epochs) + + print(self.base_lrs, self.min_lrs, total_iter, warmup_iter, flat_iter, no_aug_iter) + self.lr_func = partial(flat_cosine_schedule, total_iter, warmup_iter, flat_iter, no_aug_iter) + + def step(self, current_iter, optimizer): + """ + Updates the learning rate of the optimizer at the current iteration. + + Args: + current_iter (int): Current iteration. + optimizer (torch.optim.Optimizer): Optimizer instance. + """ + for i, group in enumerate(optimizer.param_groups): + group["lr"] = self.lr_func(current_iter, self.base_lrs[i], self.min_lrs[i]) + return optimizer diff --git a/engine/optim/optim.py b/engine/optim/optim.py new file mode 100644 index 0000000000000000000000000000000000000000..f4830c6606493612724c66094eebc9fd9331618c --- /dev/null +++ b/engine/optim/optim.py @@ -0,0 +1,25 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + + +import torch.optim as optim +import torch.optim.lr_scheduler as lr_scheduler + +from ..core import register + + +__all__ = ['AdamW', 'SGD', 'Adam', 'MultiStepLR', 'CosineAnnealingLR', 'OneCycleLR', 'LambdaLR'] + + + +SGD = register()(optim.SGD) +Adam = register()(optim.Adam) +AdamW = register()(optim.AdamW) + + +MultiStepLR = register()(lr_scheduler.MultiStepLR) +CosineAnnealingLR = register()(lr_scheduler.CosineAnnealingLR) +OneCycleLR = register()(lr_scheduler.OneCycleLR) +LambdaLR = register()(lr_scheduler.LambdaLR) diff --git a/engine/optim/warmup.py b/engine/optim/warmup.py new file mode 100644 index 0000000000000000000000000000000000000000..86e319b59b4c71c9e0c8297c11dcf818a908300d --- /dev/null +++ b/engine/optim/warmup.py @@ -0,0 +1,48 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +from torch.optim.lr_scheduler import LRScheduler + +from ..core import register + + +class Warmup(object): + def __init__(self, lr_scheduler: LRScheduler, warmup_duration: int, last_step: int=-1) -> None: + self.lr_scheduler = lr_scheduler + self.warmup_end_values = [pg['lr'] for pg in lr_scheduler.optimizer.param_groups] + self.last_step = last_step + self.warmup_duration = warmup_duration + self.step() + + def state_dict(self): + return {k: v for k, v in self.__dict__.items() if k != 'lr_scheduler'} + + def load_state_dict(self, state_dict): + self.__dict__.update(state_dict) + + def get_warmup_factor(self, step, **kwargs): + raise NotImplementedError + + def step(self, ): + self.last_step += 1 + if self.last_step >= self.warmup_duration: + return + factor = self.get_warmup_factor(self.last_step) + for i, pg in enumerate(self.lr_scheduler.optimizer.param_groups): + pg['lr'] = factor * self.warmup_end_values[i] + + def finished(self, ): + if self.last_step >= self.warmup_duration: + return True + return False + + +@register() +class LinearWarmup(Warmup): + def __init__(self, lr_scheduler: LRScheduler, warmup_duration: int, last_step: int = -1) -> None: + super().__init__(lr_scheduler, warmup_duration, last_step) + + def get_warmup_factor(self, step): + return min(1.0, (step + 1) / self.warmup_duration) diff --git a/engine/solver/__init__.py b/engine/solver/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a6a56c9f502c236b5cb26d8a2934d941c6a8e2a1 --- /dev/null +++ b/engine/solver/__init__.py @@ -0,0 +1,17 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +from ._solver import BaseSolver +from .clas_solver import ClasSolver +from .det_solver import DetSolver + + + +from typing import Dict + +TASKS :Dict[str, BaseSolver] = { + 'classification': ClasSolver, + 'detection': DetSolver, +} diff --git a/engine/solver/_solver.py b/engine/solver/_solver.py new file mode 100644 index 0000000000000000000000000000000000000000..925e45f0caa9b30e546c3082cae1598f24a7e287 --- /dev/null +++ b/engine/solver/_solver.py @@ -0,0 +1,696 @@ +import torch +import torch.nn as nn + +from datetime import datetime +from pathlib import Path +from typing import Dict +import atexit + +from ..misc import dist_utils +from ..core import BaseConfig + + +def to(m: nn.Module, device: str): + if m is None: + return None + return m.to(device) + + +def remove_module_prefix(state_dict): + new_state_dict = {} + for k, v in state_dict.items(): + if k.startswith('module.'): + new_state_dict[k[7:]] = v + else: + new_state_dict[k] = v + return new_state_dict + + +class BaseSolver(object): + def __init__(self, cfg: BaseConfig) -> None: + self.cfg = cfg + self.obj365_ids = [ + 0, 46, 5, 58, 114, 55, 116, 65, 21, 40, 176, 127, 249, 24, 56, 139, 92, 78, 99, 96, + 144, 295, 178, 180, 38, 39, 13, 43, 120, 219, 148, 173, 165, 154, 137, 113, 145, 146, + 204, 8, 35, 10, 88, 84, 93, 26, 112, 82, 265, 104, 141, 152, 234, 143, 150, 97, 2, + 50, 25, 75, 98, 153, 37, 73, 115, 132, 106, 61, 163, 134, 277, 81, 133, 18, 94, 30, + 169, 70, 328, 226 + ] + def _setup(self): + """Avoid instantiating unnecessary classes""" + cfg = self.cfg + if cfg.device: + device = torch.device(cfg.device) + else: + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + self.model = cfg.model + + # NOTE: Must load_tuning_state before EMA instance building + if self.cfg.tuning: + print(f'Tuning checkpoint from {self.cfg.tuning}') + self.load_tuning_state(self.cfg.tuning) + + self.model = dist_utils.warp_model( + self.model.to(device), sync_bn=cfg.sync_bn, find_unused_parameters=cfg.find_unused_parameters + ) + + self.criterion = self.to(cfg.criterion, device) + self.postprocessor = self.to(cfg.postprocessor, device) + + self.ema = self.to(cfg.ema, device) + self.scaler = cfg.scaler + + self.device = device + self.last_epoch = self.cfg.last_epoch + + self.output_dir = Path(cfg.output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + self.writer = cfg.writer + + if self.writer: + atexit.register(self.writer.close) + if dist_utils.is_main_process(): + self.writer.add_text('config', '{:s}'.format(cfg.__repr__()), 0) + + def cleanup(self): + if self.writer: + atexit.register(self.writer.close) + + def train(self): + self._setup() + self.optimizer = self.cfg.optimizer + self.lr_scheduler = self.cfg.lr_scheduler + self.lr_warmup_scheduler = self.cfg.lr_warmup_scheduler + + self.train_dataloader = dist_utils.warp_loader( + self.cfg.train_dataloader, shuffle=self.cfg.train_dataloader.shuffle + ) + self.val_dataloader = dist_utils.warp_loader( + self.cfg.val_dataloader, shuffle=self.cfg.val_dataloader.shuffle + ) + + self.evaluator = self.cfg.evaluator + + # NOTE: Instantiating order + if self.cfg.resume: + print(f'Resume checkpoint from {self.cfg.resume}') + self.load_resume_state(self.cfg.resume) + + def eval(self): + self._setup() + + self.val_dataloader = dist_utils.warp_loader( + self.cfg.val_dataloader, shuffle=self.cfg.val_dataloader.shuffle + ) + + self.evaluator = self.cfg.evaluator + + if self.cfg.resume: + print(f'Resume checkpoint from {self.cfg.resume}') + self.load_resume_state(self.cfg.resume) + + def to(self, module, device): + return module.to(device) if hasattr(module, 'to') else module + + def state_dict(self): + """State dict, train/eval""" + state = {} + state['date'] = datetime.now().isoformat() + + # For resume + state['last_epoch'] = self.last_epoch + + for k, v in self.__dict__.items(): + if hasattr(v, 'state_dict'): + v = dist_utils.de_parallel(v) + state[k] = v.state_dict() + + return state + + def load_state_dict(self, state): + """Load state dict, train/eval""" + if 'last_epoch' in state: + self.last_epoch = state['last_epoch'] + print('Load last_epoch') + + for k, v in self.__dict__.items(): + if hasattr(v, 'load_state_dict') and k in state: + v = dist_utils.de_parallel(v) + v.load_state_dict(state[k]) + print(f'Load {k}.state_dict') + + if hasattr(v, 'load_state_dict') and k not in state: + if k == 'ema': + model = getattr(self, 'model', None) + if model is not None: + ema = dist_utils.de_parallel(v) + model_state_dict = remove_module_prefix(model.state_dict()) + ema.load_state_dict({'module': model_state_dict}) + print(f'Load {k}.state_dict from model.state_dict') + else: + print(f'Not load {k}.state_dict') + + def load_resume_state(self, path: str): + """Load resume""" + if path.startswith('http'): + state = torch.hub.load_state_dict_from_url(path, map_location='cpu') + else: + state = torch.load(path, map_location='cpu') + + # state['model'] = remove_module_prefix(state['model']) + self.load_state_dict(state) + + def load_tuning_state(self, path: str): + """Load model for tuning and adjust mismatched head parameters""" + if path.startswith('http'): + state = torch.hub.load_state_dict_from_url(path, map_location='cpu') + else: + state = torch.load(path, map_location='cpu') + + module = dist_utils.de_parallel(self.model) + + # Load the appropriate state dict + if 'ema' in state: + pretrain_state_dict = state['ema']['module'] + else: + pretrain_state_dict = state['model'] + + # Adjust head parameters between datasets + try: + adjusted_state_dict = self._adjust_head_parameters(module.state_dict(), pretrain_state_dict) + stat, infos = self._matched_state(module.state_dict(), adjusted_state_dict) + except Exception: + stat, infos = self._matched_state(module.state_dict(), pretrain_state_dict) + + module.load_state_dict(stat, strict=False) + print(f'Load model.state_dict, {infos}') + + @staticmethod + def _matched_state(state: Dict[str, torch.Tensor], params: Dict[str, torch.Tensor]): + missed_list = [] + unmatched_list = [] + matched_state = {} + for k, v in state.items(): + if k in params: + if v.shape == params[k].shape: + matched_state[k] = params[k] + else: + unmatched_list.append(k) + else: + missed_list.append(k) + + return matched_state, {'missed': missed_list, 'unmatched': unmatched_list} + + def _adjust_head_parameters(self, cur_state_dict, pretrain_state_dict): + """Adjust head parameters between datasets.""" + # List of parameters to adjust + if pretrain_state_dict['decoder.denoising_class_embed.weight'].size() != \ + cur_state_dict['decoder.denoising_class_embed.weight'].size(): + del pretrain_state_dict['decoder.denoising_class_embed.weight'] + + head_param_names = [ + 'decoder.enc_score_head.weight', + 'decoder.enc_score_head.bias' + ] + for i in range(8): + head_param_names.append(f'decoder.dec_score_head.{i}.weight') + head_param_names.append(f'decoder.dec_score_head.{i}.bias') + + adjusted_params = [] + + for param_name in head_param_names: + if param_name in cur_state_dict and param_name in pretrain_state_dict: + cur_tensor = cur_state_dict[param_name] + pretrain_tensor = pretrain_state_dict[param_name] + adjusted_tensor = self.map_class_weights(cur_tensor, pretrain_tensor) + if adjusted_tensor is not None: + pretrain_state_dict[param_name] = adjusted_tensor + adjusted_params.append(param_name) + else: + print(f"Cannot adjust parameter '{param_name}' due to size mismatch.") + + return pretrain_state_dict + + def map_class_weights(self, cur_tensor, pretrain_tensor): + """Map class weights from pretrain model to current model based on class IDs.""" + if pretrain_tensor.size() == cur_tensor.size(): + return pretrain_tensor + + adjusted_tensor = cur_tensor.clone() + adjusted_tensor.requires_grad = False + + if pretrain_tensor.size() > cur_tensor.size(): + for coco_id, obj_id in enumerate(self.obj365_ids): + adjusted_tensor[coco_id] = pretrain_tensor[obj_id+1] + else: + for coco_id, obj_id in enumerate(self.obj365_ids): + adjusted_tensor[obj_id+1] = pretrain_tensor[coco_id] + + return adjusted_tensor + + def fit(self): + raise NotImplementedError('') + + def val(self): + raise NotImplementedError('') + +# obj365_classes = [ +# 'Person', 'Sneakers', 'Chair', 'Other Shoes', 'Hat', 'Car', 'Lamp', 'Glasses', +# 'Bottle', 'Desk', 'Cup', 'Street Lights', 'Cabinet/shelf', 'Handbag/Satchel', +# 'Bracelet', 'Plate', 'Picture/Frame', 'Helmet', 'Book', 'Gloves', 'Storage box', +# 'Boat', 'Leather Shoes', 'Flower', 'Bench', 'Potted Plant', 'Bowl/Basin', 'Flag', +# 'Pillow', 'Boots', 'Vase', 'Microphone', 'Necklace', 'Ring', 'SUV', 'Wine Glass', +# 'Belt', 'Moniter/TV', 'Backpack', 'Umbrella', 'Traffic Light', 'Speaker', 'Watch', +# 'Tie', 'Trash bin Can', 'Slippers', 'Bicycle', 'Stool', 'Barrel/bucket', 'Van', +# 'Couch', 'Sandals', 'Bakset', 'Drum', 'Pen/Pencil', 'Bus', 'Wild Bird', 'High Heels', +# 'Motorcycle', 'Guitar', 'Carpet', 'Cell Phone', 'Bread', 'Camera', 'Canned', 'Truck', +# 'Traffic cone', 'Cymbal', 'Lifesaver', 'Towel', 'Stuffed Toy', 'Candle', 'Sailboat', +# 'Laptop', 'Awning', 'Bed', 'Faucet', 'Tent', 'Horse', 'Mirror', 'Power outlet', +# 'Sink', 'Apple', 'Air Conditioner', 'Knife', 'Hockey Stick', 'Paddle', 'Pickup Truck', +# 'Fork', 'Traffic Sign', 'Ballon', 'Tripod', 'Dog', 'Spoon', 'Clock', 'Pot', 'Cow', +# 'Cake', 'Dinning Table', 'Sheep', 'Hanger', 'Blackboard/Whiteboard', 'Napkin', +# 'Other Fish', 'Orange/Tangerine', 'Toiletry', 'Keyboard', 'Tomato', 'Lantern', +# 'Machinery Vehicle', 'Fan', 'Green Vegetables', 'Banana', 'Baseball Glove', +# 'Airplane', 'Mouse', 'Train', 'Pumpkin', 'Soccer', 'Skiboard', 'Luggage', 'Nightstand', +# 'Tea pot', 'Telephone', 'Trolley', 'Head Phone', 'Sports Car', 'Stop Sign', 'Dessert', +# 'Scooter', 'Stroller', 'Crane', 'Remote', 'Refrigerator', 'Oven', 'Lemon', 'Duck', +# 'Baseball Bat', 'Surveillance Camera', 'Cat', 'Jug', 'Broccoli', 'Piano', 'Pizza', +# 'Elephant', 'Skateboard', 'Surfboard', 'Gun', 'Skating and Skiing shoes', 'Gas stove', +# 'Donut', 'Bow Tie', 'Carrot', 'Toilet', 'Kite', 'Strawberry', 'Other Balls', 'Shovel', +# 'Pepper', 'Computer Box', 'Toilet Paper', 'Cleaning Products', 'Chopsticks', 'Microwave', +# 'Pigeon', 'Baseball', 'Cutting/chopping Board', 'Coffee Table', 'Side Table', 'Scissors', +# 'Marker', 'Pie', 'Ladder', 'Snowboard', 'Cookies', 'Radiator', 'Fire Hydrant', 'Basketball', +# 'Zebra', 'Grape', 'Giraffe', 'Potato', 'Sausage', 'Tricycle', 'Violin', 'Egg', +# 'Fire Extinguisher', 'Candy', 'Fire Truck', 'Billards', 'Converter', 'Bathtub', +# 'Wheelchair', 'Golf Club', 'Briefcase', 'Cucumber', 'Cigar/Cigarette ', 'Paint Brush', +# 'Pear', 'Heavy Truck', 'Hamburger', 'Extractor', 'Extention Cord', 'Tong', +# 'Tennis Racket', 'Folder', 'American Football', 'earphone', 'Mask', 'Kettle', +# 'Tennis', 'Ship', 'Swing', 'Coffee Machine', 'Slide', 'Carriage', 'Onion', +# 'Green beans', 'Projector', 'Frisbee', 'Washing Machine/Drying Machine', 'Chicken', +# 'Printer', 'Watermelon', 'Saxophone', 'Tissue', 'Toothbrush', 'Ice cream', +# 'Hotair ballon', 'Cello', 'French Fries', 'Scale', 'Trophy', 'Cabbage', 'Hot dog', +# 'Blender', 'Peach', 'Rice', 'Wallet/Purse', 'Volleyball', 'Deer', 'Goose', 'Tape', +# 'Tablet', 'Cosmetics', 'Trumpet', 'Pineapple', 'Golf Ball', 'Ambulance', 'Parking meter', +# 'Mango', 'Key', 'Hurdle', 'Fishing Rod', 'Medal', 'Flute', 'Brush', 'Penguin', +# 'Megaphone', 'Corn', 'Lettuce', 'Garlic', 'Swan', 'Helicopter', 'Green Onion', +# 'Sandwich', 'Nuts', 'Speed Limit Sign', 'Induction Cooker', 'Broom', 'Trombone', +# 'Plum', 'Rickshaw', 'Goldfish', 'Kiwi fruit', 'Router/modem', 'Poker Card', 'Toaster', +# 'Shrimp', 'Sushi', 'Cheese', 'Notepaper', 'Cherry', 'Pliers', 'CD', 'Pasta', 'Hammer', +# 'Cue', 'Avocado', 'Hamimelon', 'Flask', 'Mushroon', 'Screwdriver', 'Soap', 'Recorder', +# 'Bear', 'Eggplant', 'Board Eraser', 'Coconut', 'Tape Measur/ Ruler', 'Pig', +# 'Showerhead', 'Globe', 'Chips', 'Steak', 'Crosswalk Sign', 'Stapler', 'Campel', +# 'Formula 1 ', 'Pomegranate', 'Dishwasher', 'Crab', 'Hoverboard', 'Meat ball', +# 'Rice Cooker', 'Tuba', 'Calculator', 'Papaya', 'Antelope', 'Parrot', 'Seal', +# 'Buttefly', 'Dumbbell', 'Donkey', 'Lion', 'Urinal', 'Dolphin', 'Electric Drill', +# 'Hair Dryer', 'Egg tart', 'Jellyfish', 'Treadmill', 'Lighter', 'Grapefruit', +# 'Game board', 'Mop', 'Radish', 'Baozi', 'Target', 'French', 'Spring Rolls', 'Monkey', +# 'Rabbit', 'Pencil Case', 'Yak', 'Red Cabbage', 'Binoculars', 'Asparagus', 'Barbell', +# 'Scallop', 'Noddles', 'Comb', 'Dumpling', 'Oyster', 'Table Teniis paddle', +# 'Cosmetics Brush/Eyeliner Pencil', 'Chainsaw', 'Eraser', 'Lobster', 'Durian', 'Okra', +# 'Lipstick', 'Cosmetics Mirror', 'Curling', 'Table Tennis ' +# ] + +# coco_classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', +# 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', +# 'stop sign', 'parking meter', 'bench', 'wild bird', 'cat', 'dog', +# 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', +# 'backpack', 'umbrella', 'handbag/satchel', 'tie', 'luggage', 'frisbee', +# 'skating and skiing shoes', 'snowboard', 'baseball', 'kite', 'baseball bat', +# 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', +# 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl/basin', +# 'banana', 'apple', 'sandwich', 'orange/tangerine', 'broccoli', 'carrot', +# 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', +# 'potted plant', 'bed', 'dinning table', 'toilet', 'moniter/tv', 'laptop', +# 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', +# 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', +# 'vase', 'scissors', 'stuffed toy', 'hair dryer', 'toothbrush'] + + +# obj365_classes = [ +# (0, 'Person'), +# (1, 'Sneakers'), +# (2, 'Chair'), +# (3, 'Other Shoes'), +# (4, 'Hat'), +# (5, 'Car'), +# (6, 'Lamp'), +# (7, 'Glasses'), +# (8, 'Bottle'), +# (9, 'Desk'), +# (10, 'Cup'), +# (11, 'Street Lights'), +# (12, 'Cabinet/shelf'), +# (13, 'Handbag/Satchel'), +# (14, 'Bracelet'), +# (15, 'Plate'), +# (16, 'Picture/Frame'), +# (17, 'Helmet'), +# (18, 'Book'), +# (19, 'Gloves'), +# (20, 'Storage box'), +# (21, 'Boat'), +# (22, 'Leather Shoes'), +# (23, 'Flower'), +# (24, 'Bench'), +# (25, 'Potted Plant'), +# (26, 'Bowl/Basin'), +# (27, 'Flag'), +# (28, 'Pillow'), +# (29, 'Boots'), +# (30, 'Vase'), +# (31, 'Microphone'), +# (32, 'Necklace'), +# (33, 'Ring'), +# (34, 'SUV'), +# (35, 'Wine Glass'), +# (36, 'Belt'), +# (37, 'Monitor/TV'), +# (38, 'Backpack'), +# (39, 'Umbrella'), +# (40, 'Traffic Light'), +# (41, 'Speaker'), +# (42, 'Watch'), +# (43, 'Tie'), +# (44, 'Trash bin Can'), +# (45, 'Slippers'), +# (46, 'Bicycle'), +# (47, 'Stool'), +# (48, 'Barrel/bucket'), +# (49, 'Van'), +# (50, 'Couch'), +# (51, 'Sandals'), +# (52, 'Basket'), +# (53, 'Drum'), +# (54, 'Pen/Pencil'), +# (55, 'Bus'), +# (56, 'Wild Bird'), +# (57, 'High Heels'), +# (58, 'Motorcycle'), +# (59, 'Guitar'), +# (60, 'Carpet'), +# (61, 'Cell Phone'), +# (62, 'Bread'), +# (63, 'Camera'), +# (64, 'Canned'), +# (65, 'Truck'), +# (66, 'Traffic cone'), +# (67, 'Cymbal'), +# (68, 'Lifesaver'), +# (69, 'Towel'), +# (70, 'Stuffed Toy'), +# (71, 'Candle'), +# (72, 'Sailboat'), +# (73, 'Laptop'), +# (74, 'Awning'), +# (75, 'Bed'), +# (76, 'Faucet'), +# (77, 'Tent'), +# (78, 'Horse'), +# (79, 'Mirror'), +# (80, 'Power outlet'), +# (81, 'Sink'), +# (82, 'Apple'), +# (83, 'Air Conditioner'), +# (84, 'Knife'), +# (85, 'Hockey Stick'), +# (86, 'Paddle'), +# (87, 'Pickup Truck'), +# (88, 'Fork'), +# (89, 'Traffic Sign'), +# (90, 'Balloon'), +# (91, 'Tripod'), +# (92, 'Dog'), +# (93, 'Spoon'), +# (94, 'Clock'), +# (95, 'Pot'), +# (96, 'Cow'), +# (97, 'Cake'), +# (98, 'Dining Table'), +# (99, 'Sheep'), +# (100, 'Hanger'), +# (101, 'Blackboard/Whiteboard'), +# (102, 'Napkin'), +# (103, 'Other Fish'), +# (104, 'Orange/Tangerine'), +# (105, 'Toiletry'), +# (106, 'Keyboard'), +# (107, 'Tomato'), +# (108, 'Lantern'), +# (109, 'Machinery Vehicle'), +# (110, 'Fan'), +# (111, 'Green Vegetables'), +# (112, 'Banana'), +# (113, 'Baseball Glove'), +# (114, 'Airplane'), +# (115, 'Mouse'), +# (116, 'Train'), +# (117, 'Pumpkin'), +# (118, 'Soccer'), +# (119, 'Skiboard'), +# (120, 'Luggage'), +# (121, 'Nightstand'), +# (122, 'Tea pot'), +# (123, 'Telephone'), +# (124, 'Trolley'), +# (125, 'Head Phone'), +# (126, 'Sports Car'), +# (127, 'Stop Sign'), +# (128, 'Dessert'), +# (129, 'Scooter'), +# (130, 'Stroller'), +# (131, 'Crane'), +# (132, 'Remote'), +# (133, 'Refrigerator'), +# (134, 'Oven'), +# (135, 'Lemon'), +# (136, 'Duck'), +# (137, 'Baseball Bat'), +# (138, 'Surveillance Camera'), +# (139, 'Cat'), +# (140, 'Jug'), +# (141, 'Broccoli'), +# (142, 'Piano'), +# (143, 'Pizza'), +# (144, 'Elephant'), +# (145, 'Skateboard'), +# (146, 'Surfboard'), +# (147, 'Gun'), +# (148, 'Skating and Skiing Shoes'), +# (149, 'Gas Stove'), +# (150, 'Donut'), +# (151, 'Bow Tie'), +# (152, 'Carrot'), +# (153, 'Toilet'), +# (154, 'Kite'), +# (155, 'Strawberry'), +# (156, 'Other Balls'), +# (157, 'Shovel'), +# (158, 'Pepper'), +# (159, 'Computer Box'), +# (160, 'Toilet Paper'), +# (161, 'Cleaning Products'), +# (162, 'Chopsticks'), +# (163, 'Microwave'), +# (164, 'Pigeon'), +# (165, 'Baseball'), +# (166, 'Cutting/chopping Board'), +# (167, 'Coffee Table'), +# (168, 'Side Table'), +# (169, 'Scissors'), +# (170, 'Marker'), +# (171, 'Pie'), +# (172, 'Ladder'), +# (173, 'Snowboard'), +# (174, 'Cookies'), +# (175, 'Radiator'), +# (176, 'Fire Hydrant'), +# (177, 'Basketball'), +# (178, 'Zebra'), +# (179, 'Grape'), +# (180, 'Giraffe'), +# (181, 'Potato'), +# (182, 'Sausage'), +# (183, 'Tricycle'), +# (184, 'Violin'), +# (185, 'Egg'), +# (186, 'Fire Extinguisher'), +# (187, 'Candy'), +# (188, 'Fire Truck'), +# (189, 'Billiards'), +# (190, 'Converter'), +# (191, 'Bathtub'), +# (192, 'Wheelchair'), +# (193, 'Golf Club'), +# (194, 'Briefcase'), +# (195, 'Cucumber'), +# (196, 'Cigar/Cigarette'), +# (197, 'Paint Brush'), +# (198, 'Pear'), +# (199, 'Heavy Truck'), +# (200, 'Hamburger'), +# (201, 'Extractor'), +# (202, 'Extension Cord'), +# (203, 'Tong'), +# (204, 'Tennis Racket'), +# (205, 'Folder'), +# (206, 'American Football'), +# (207, 'Earphone'), +# (208, 'Mask'), +# (209, 'Kettle'), +# (210, 'Tennis'), +# (211, 'Ship'), +# (212, 'Swing'), +# (213, 'Coffee Machine'), +# (214, 'Slide'), +# (215, 'Carriage'), +# (216, 'Onion'), +# (217, 'Green Beans'), +# (218, 'Projector'), +# (219, 'Frisbee'), +# (220, 'Washing Machine/Drying Machine'), +# (221, 'Chicken'), +# (222, 'Printer'), +# (223, 'Watermelon'), +# (224, 'Saxophone'), +# (225, 'Tissue'), +# (226, 'Toothbrush'), +# (227, 'Ice Cream'), +# (228, 'Hot Air Balloon'), +# (229, 'Cello'), +# (230, 'French Fries'), +# (231, 'Scale'), +# (232, 'Trophy'), +# (233, 'Cabbage'), +# (234, 'Hot Dog'), +# (235, 'Blender'), +# (236, 'Peach'), +# (237, 'Rice'), +# (238, 'Wallet/Purse'), +# (239, 'Volleyball'), +# (240, 'Deer'), +# (241, 'Goose'), +# (242, 'Tape'), +# (243, 'Tablet'), +# (244, 'Cosmetics'), +# (245, 'Trumpet'), +# (246, 'Pineapple'), +# (247, 'Golf Ball'), +# (248, 'Ambulance'), +# (249, 'Parking Meter'), +# (250, 'Mango'), +# (251, 'Key'), +# (252, 'Hurdle'), +# (253, 'Fishing Rod'), +# (254, 'Medal'), +# (255, 'Flute'), +# (256, 'Brush'), +# (257, 'Penguin'), +# (258, 'Megaphone'), +# (259, 'Corn'), +# (260, 'Lettuce'), +# (261, 'Garlic'), +# (262, 'Swan'), +# (263, 'Helicopter'), +# (264, 'Green Onion'), +# (265, 'Sandwich'), +# (266, 'Nuts'), +# (267, 'Speed Limit Sign'), +# (268, 'Induction Cooker'), +# (269, 'Broom'), +# (270, 'Trombone'), +# (271, 'Plum'), +# (272, 'Rickshaw'), +# (273, 'Goldfish'), +# (274, 'Kiwi Fruit'), +# (275, 'Router/Modem'), +# (276, 'Poker Card'), +# (277, 'Toaster'), +# (278, 'Shrimp'), +# (279, 'Sushi'), +# (280, 'Cheese'), +# (281, 'Notepaper'), +# (282, 'Cherry'), +# (283, 'Pliers'), +# (284, 'CD'), +# (285, 'Pasta'), +# (286, 'Hammer'), +# (287, 'Cue'), +# (288, 'Avocado'), +# (289, 'Hami Melon'), +# (290, 'Flask'), +# (291, 'Mushroom'), +# (292, 'Screwdriver'), +# (293, 'Soap'), +# (294, 'Recorder'), +# (295, 'Bear'), +# (296, 'Eggplant'), +# (297, 'Board Eraser'), +# (298, 'Coconut'), +# (299, 'Tape Measure/Ruler'), +# (300, 'Pig'), +# (301, 'Showerhead'), +# (302, 'Globe'), +# (303, 'Chips'), +# (304, 'Steak'), +# (305, 'Crosswalk Sign'), +# (306, 'Stapler'), +# (307, 'Camel'), +# (308, 'Formula 1'), +# (309, 'Pomegranate'), +# (310, 'Dishwasher'), +# (311, 'Crab'), +# (312, 'Hoverboard'), +# (313, 'Meatball'), +# (314, 'Rice Cooker'), +# (315, 'Tuba'), +# (316, 'Calculator'), +# (317, 'Papaya'), +# (318, 'Antelope'), +# (319, 'Parrot'), +# (320, 'Seal'), +# (321, 'Butterfly'), +# (322, 'Dumbbell'), +# (323, 'Donkey'), +# (324, 'Lion'), +# (325, 'Urinal'), +# (326, 'Dolphin'), +# (327, 'Electric Drill'), +# (328, 'Hair Dryer'), +# (329, 'Egg Tart'), +# (330, 'Jellyfish'), +# (331, 'Treadmill'), +# (332, 'Lighter'), +# (333, 'Grapefruit'), +# (334, 'Game Board'), +# (335, 'Mop'), +# (336, 'Radish'), +# (337, 'Baozi'), +# (338, 'Target'), +# (339, 'French'), +# (340, 'Spring Rolls'), +# (341, 'Monkey'), +# (342, 'Rabbit'), +# (343, 'Pencil Case'), +# (344, 'Yak'), +# (345, 'Red Cabbage'), +# (346, 'Binoculars'), +# (347, 'Asparagus'), +# (348, 'Barbell'), +# (349, 'Scallop'), +# (350, 'Noodles'), +# (351, 'Comb'), +# (352, 'Dumpling'), +# (353, 'Oyster'), +# (354, 'Table Tennis Paddle'), +# (355, 'Cosmetics Brush/Eyeliner Pencil'), +# (356, 'Chainsaw'), +# (357, 'Eraser'), +# (358, 'Lobster'), +# (359, 'Durian'), +# (360, 'Okra'), +# (361, 'Lipstick'), +# (362, 'Cosmetics Mirror'), +# (363, 'Curling'), +# (364, 'Table Tennis') +# ] diff --git a/engine/solver/clas_engine.py b/engine/solver/clas_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..dee29b575ef3115af774b8208d97208f3fb1be39 --- /dev/null +++ b/engine/solver/clas_engine.py @@ -0,0 +1,74 @@ +""" +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 ..misc import (MetricLogger, SmoothedValue, reduce_dict) + + +def train_one_epoch(model: nn.Module, criterion: nn.Module, dataloader, optimizer, ema, epoch, device): + """ + """ + model.train() + + metric_logger = MetricLogger(delimiter=" ") + metric_logger.add_meter('lr', SmoothedValue(window_size=1, fmt='{value:.6f}')) + print_freq = 100 + header = 'Epoch: [{}]'.format(epoch) + + for imgs, labels in metric_logger.log_every(dataloader, print_freq, header): + imgs = imgs.to(device) + labels = labels.to(device) + + preds = model(imgs) + loss: torch.Tensor = criterion(preds, labels, epoch) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + if ema is not None: + ema.update(model) + + loss_reduced_values = {k: v.item() for k, v in reduce_dict({'loss': loss}).items()} + metric_logger.update(**loss_reduced_values) + metric_logger.update(lr=optimizer.param_groups[0]["lr"]) + + metric_logger.synchronize_between_processes() + print("Averaged stats:", metric_logger) + + stats = {k: meter.global_avg for k, meter in metric_logger.meters.items()} + return stats + + + +@torch.no_grad() +def evaluate(model, criterion, dataloader, device): + model.eval() + + metric_logger = MetricLogger(delimiter=" ") + # metric_logger.add_meter('acc', SmoothedValue(window_size=1, fmt='{global_avg:.4f}')) + # metric_logger.add_meter('loss', SmoothedValue(window_size=1, fmt='{value:.2f}')) + metric_logger.add_meter('acc', SmoothedValue(window_size=1)) + metric_logger.add_meter('loss', SmoothedValue(window_size=1)) + + header = 'Test:' + for imgs, labels in metric_logger.log_every(dataloader, 10, header): + imgs, labels = imgs.to(device), labels.to(device) + preds = model(imgs) + + acc = (preds.argmax(dim=-1) == labels).sum() / preds.shape[0] + loss = criterion(preds, labels) + + dict_reduced = reduce_dict({'acc': acc, 'loss': loss}) + reduced_values = {k: v.item() for k, v in dict_reduced.items()} + metric_logger.update(**reduced_values) + + metric_logger.synchronize_between_processes() + print("Averaged stats:", metric_logger) + + stats = {k: meter.global_avg for k, meter in metric_logger.meters.items()} + return stats diff --git a/engine/solver/clas_solver.py b/engine/solver/clas_solver.py new file mode 100644 index 0000000000000000000000000000000000000000..e49907191717800c1ec3b943a53ceeeeb842a6e2 --- /dev/null +++ b/engine/solver/clas_solver.py @@ -0,0 +1,71 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +import time +import json +import datetime +from pathlib import Path + +import torch +import torch.nn as nn + +from ..misc import dist_utils +from ._solver import BaseSolver +from .clas_engine import train_one_epoch, evaluate + + +class ClasSolver(BaseSolver): + + def fit(self, ): + print("Start training") + self.train() + args = self.cfg + + n_parameters = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + print('Number of params:', n_parameters) + + output_dir = Path(args.output_dir) + output_dir.mkdir(exist_ok=True) + + start_time = time.time() + start_epoch = self.last_epoch + 1 + for epoch in range(start_epoch, args.epoches): + + if dist_utils.is_dist_available_and_initialized(): + self.train_dataloader.sampler.set_epoch(epoch) + + train_stats = train_one_epoch(self.model, + self.criterion, + self.train_dataloader, + self.optimizer, + self.ema, + epoch=epoch, + device=self.device) + self.lr_scheduler.step() + self.last_epoch += 1 + + if output_dir: + checkpoint_paths = [output_dir / 'checkpoint.pth'] + # extra checkpoint before LR drop and every 100 epochs + if (epoch + 1) % args.checkpoint_freq == 0: + checkpoint_paths.append(output_dir / f'checkpoint{epoch:04}.pth') + for checkpoint_path in checkpoint_paths: + dist_utils.save_on_master(self.state_dict(epoch), checkpoint_path) + + module = self.ema.module if self.ema else self.model + test_stats = evaluate(module, self.criterion, self.val_dataloader, self.device) + + log_stats = {**{f'train_{k}': v for k, v in train_stats.items()}, + **{f'test_{k}': v for k, v in test_stats.items()}, + 'epoch': epoch, + 'n_parameters': n_parameters} + + if output_dir and dist_utils.is_main_process(): + with (output_dir / "log.txt").open("a") as f: + f.write(json.dumps(log_stats) + "\n") + + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print('Training time {}'.format(total_time_str)) diff --git a/engine/solver/det_engine.py b/engine/solver/det_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..5d53dda7385ad66c78b60ac142797abbbea08092 --- /dev/null +++ b/engine/solver/det_engine.py @@ -0,0 +1,177 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from DETR (https://github.com/facebookresearch/detr/blob/main/engine.py) +Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +""" + + +import sys +import math +from typing import Iterable + +import torch +import torch.amp +from torch.utils.tensorboard import SummaryWriter +from torch.cuda.amp.grad_scaler import GradScaler + +from ..optim import ModelEMA, Warmup +from ..data import CocoEvaluator +from ..misc import MetricLogger, SmoothedValue, dist_utils + + +def train_one_epoch(self_lr_scheduler, lr_scheduler, model: torch.nn.Module, criterion: torch.nn.Module, + data_loader: Iterable, optimizer: torch.optim.Optimizer, + device: torch.device, epoch: int, max_norm: float = 0, **kwargs): + model.train() + criterion.train() + metric_logger = MetricLogger(delimiter=" ") + metric_logger.add_meter('lr', SmoothedValue(window_size=1, fmt='{value:.6f}')) + header = 'Epoch: [{}]'.format(epoch) + + print_freq = kwargs.get('print_freq', 10) + writer :SummaryWriter = kwargs.get('writer', None) + + ema :ModelEMA = kwargs.get('ema', None) + scaler :GradScaler = kwargs.get('scaler', None) + lr_warmup_scheduler :Warmup = kwargs.get('lr_warmup_scheduler', None) + + cur_iters = epoch * len(data_loader) + + for i, (samples, targets) in enumerate(metric_logger.log_every(data_loader, print_freq, header)): + samples = samples.to(device) + targets = [{k: v.to(device) for k, v in t.items()} for t in targets] + global_step = epoch * len(data_loader) + i + metas = dict(epoch=epoch, step=i, global_step=global_step, epoch_step=len(data_loader)) + + if scaler is not None: + with torch.autocast(device_type=str(device), cache_enabled=True): + outputs = model(samples, targets=targets) + + if torch.isnan(outputs['pred_boxes']).any() or torch.isinf(outputs['pred_boxes']).any(): + print(outputs['pred_boxes']) + state = model.state_dict() + new_state = {} + for key, value in model.state_dict().items(): + # Replace 'module' with 'model' in each key + new_key = key.replace('module.', '') + # Add the updated key-value pair to the state dictionary + state[new_key] = value + new_state['model'] = state + dist_utils.save_on_master(new_state, "./NaN.pth") + + with torch.autocast(device_type=str(device), enabled=False): + loss_dict = criterion(outputs, targets, **metas) + + loss = sum(loss_dict.values()) + scaler.scale(loss).backward() + + if max_norm > 0: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) + + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad() + + else: + outputs = model(samples, targets=targets) + loss_dict = criterion(outputs, targets, **metas) + + loss : torch.Tensor = sum(loss_dict.values()) + optimizer.zero_grad() + loss.backward() + + if max_norm > 0: + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) + + optimizer.step() + + # ema + if ema is not None: + ema.update(model) + + if self_lr_scheduler: + optimizer = lr_scheduler.step(cur_iters + i, optimizer) + else: + if lr_warmup_scheduler is not None: + lr_warmup_scheduler.step() + + loss_dict_reduced = dist_utils.reduce_dict(loss_dict) + loss_value = sum(loss_dict_reduced.values()) + + if not math.isfinite(loss_value): + print("Loss is {}, stopping training".format(loss_value)) + print(loss_dict_reduced) + sys.exit(1) + + metric_logger.update(loss=loss_value, **loss_dict_reduced) + metric_logger.update(lr=optimizer.param_groups[0]["lr"]) + + if writer and dist_utils.is_main_process() and global_step % 10 == 0: + writer.add_scalar('Loss/total', loss_value.item(), global_step) + for j, pg in enumerate(optimizer.param_groups): + writer.add_scalar(f'Lr/pg_{j}', pg['lr'], global_step) + for k, v in loss_dict_reduced.items(): + writer.add_scalar(f'Loss/{k}', v.item(), global_step) + + # gather the stats from all processes + metric_logger.synchronize_between_processes() + print("Averaged stats:", metric_logger) + return {k: meter.global_avg for k, meter in metric_logger.meters.items()} + + +@torch.no_grad() +def evaluate(model: torch.nn.Module, criterion: torch.nn.Module, postprocessor, data_loader, coco_evaluator: CocoEvaluator, device): + model.eval() + criterion.eval() + coco_evaluator.cleanup() + + metric_logger = MetricLogger(delimiter=" ") + # metric_logger.add_meter('class_error', SmoothedValue(window_size=1, fmt='{value:.2f}')) + header = 'Test:' + + # iou_types = tuple(k for k in ('segm', 'bbox') if k in postprocessor.keys()) + iou_types = coco_evaluator.iou_types + # coco_evaluator = CocoEvaluator(base_ds, iou_types) + # coco_evaluator.coco_eval[iou_types[0]].params.iouThrs = [0, 0.1, 0.5, 0.75] + + for samples, targets in metric_logger.log_every(data_loader, 10, header): + samples = samples.to(device) + targets = [{k: v.to(device) for k, v in t.items()} for t in targets] + + outputs = model(samples) + + orig_target_sizes = torch.stack([t["orig_size"] for t in targets], dim=0) + + results = postprocessor(outputs, orig_target_sizes) + + # if 'segm' in postprocessor.keys(): + # target_sizes = torch.stack([t["size"] for t in targets], dim=0) + # results = postprocessor['segm'](results, outputs, orig_target_sizes, target_sizes) + + res = {target['image_id'].item(): output for target, output in zip(targets, results)} + if coco_evaluator is not None: + coco_evaluator.update(res) + + # gather the stats from all processes + metric_logger.synchronize_between_processes() + print("Averaged stats:", metric_logger) + if coco_evaluator is not None: + coco_evaluator.synchronize_between_processes() + + # accumulate predictions from all images + if coco_evaluator is not None: + coco_evaluator.accumulate() + coco_evaluator.summarize() + + stats = {} + # stats = {k: meter.global_avg for k, meter in metric_logger.meters.items()} + if coco_evaluator is not None: + if 'bbox' in iou_types: + stats['coco_eval_bbox'] = coco_evaluator.coco_eval['bbox'].stats.tolist() + if 'segm' in iou_types: + stats['coco_eval_masks'] = coco_evaluator.coco_eval['segm'].stats.tolist() + + return stats, coco_evaluator diff --git a/engine/solver/det_solver.py b/engine/solver/det_solver.py new file mode 100644 index 0000000000000000000000000000000000000000..9c5d0d9292700fdb946bc4159682566465b1e5e7 --- /dev/null +++ b/engine/solver/det_solver.py @@ -0,0 +1,201 @@ +""" +DEIM: DETR with Improved Matching for Fast Convergence +Copyright (c) 2024 The DEIM Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE) +Copyright (c) 2024 D-FINE authors. All Rights Reserved. +""" + +import time +import json +import datetime + +import torch + +from ..misc import dist_utils, stats + +from ._solver import BaseSolver +from .det_engine import train_one_epoch, evaluate +from ..optim.lr_scheduler import FlatCosineLRScheduler + + +class DetSolver(BaseSolver): + + def fit(self, ): + self.train() + args = self.cfg + + n_parameters, model_stats = stats(self.cfg) + print(model_stats) + print("-"*42 + "Start training" + "-"*43) + + for i, (name, param) in enumerate(self.model.named_parameters()): + if i in [194, 195]: + print(f"Index {i}: {name} - requires_grad: {param.requires_grad}") + + self.self_lr_scheduler = False + if args.lrsheduler is not None: + iter_per_epoch = len(self.train_dataloader) + print(" ## Using Self-defined Scheduler-{} ## ".format(args.lrsheduler)) + self.lr_scheduler = FlatCosineLRScheduler(self.optimizer, args.lr_gamma, iter_per_epoch, total_epochs=args.epoches, + warmup_iter=args.warmup_iter, flat_epochs=args.flat_epoch, no_aug_epochs=args.no_aug_epoch) + self.self_lr_scheduler = True + n_parameters = sum([p.numel() for p in self.model.parameters() if p.requires_grad]) + print(f'number of trainable parameters: {n_parameters}') + + n_parameters = sum([p.numel() for p in self.model.parameters() if not p.requires_grad]) + print(f'number of non-trainable parameters: {n_parameters}') + + top1 = 0 + best_stat = {'epoch': -1, } + # evaluate again before resume training + if self.last_epoch > 0: + module = self.ema.module if self.ema else self.model + test_stats, coco_evaluator = evaluate( + module, + self.criterion, + self.postprocessor, + self.val_dataloader, + self.evaluator, + self.device + ) + for k in test_stats: + best_stat['epoch'] = self.last_epoch + best_stat[k] = test_stats[k][0] + top1 = test_stats[k][0] + print(f'best_stat: {best_stat}') + + best_stat_print = best_stat.copy() + start_time = time.time() + start_epoch = self.last_epoch + 1 + for epoch in range(start_epoch, args.epoches): + + self.train_dataloader.set_epoch(epoch) + # self.train_dataloader.dataset.set_epoch(epoch) + if dist_utils.is_dist_available_and_initialized(): + self.train_dataloader.sampler.set_epoch(epoch) + + if epoch == self.train_dataloader.collate_fn.stop_epoch: + self.load_resume_state(str(self.output_dir / 'best_stg1.pth')) + self.ema.decay = self.train_dataloader.collate_fn.ema_restart_decay + print(f'Refresh EMA at epoch {epoch} with decay {self.ema.decay}') + + train_stats = train_one_epoch( + self.self_lr_scheduler, + self.lr_scheduler, + self.model, + self.criterion, + self.train_dataloader, + self.optimizer, + self.device, + epoch, + max_norm=args.clip_max_norm, + print_freq=args.print_freq, + ema=self.ema, + scaler=self.scaler, + lr_warmup_scheduler=self.lr_warmup_scheduler, + writer=self.writer + ) + + if not self.self_lr_scheduler: # update by epoch + if self.lr_warmup_scheduler is None or self.lr_warmup_scheduler.finished(): + self.lr_scheduler.step() + + self.last_epoch += 1 + + if self.output_dir and epoch < self.train_dataloader.collate_fn.stop_epoch: + checkpoint_paths = [self.output_dir / 'last.pth'] + # extra checkpoint before LR drop and every 100 epochs + if (epoch + 1) % args.checkpoint_freq == 0: + checkpoint_paths.append(self.output_dir / f'checkpoint{epoch:04}.pth') + for checkpoint_path in checkpoint_paths: + dist_utils.save_on_master(self.state_dict(), checkpoint_path) + + module = self.ema.module if self.ema else self.model + test_stats, coco_evaluator = evaluate( + module, + self.criterion, + self.postprocessor, + self.val_dataloader, + self.evaluator, + self.device + ) + + for k in test_stats: + if self.writer and dist_utils.is_main_process(): + for i, v in enumerate(test_stats[k]): + self.writer.add_scalar(f'Test/{k}_{i}'.format(k), v, epoch) + + if k in best_stat: + best_stat['epoch'] = epoch if test_stats[k][0] > best_stat[k] else best_stat['epoch'] + best_stat[k] = max(best_stat[k], test_stats[k][0]) + else: + best_stat['epoch'] = epoch + best_stat[k] = test_stats[k][0] + + if best_stat[k] > top1: + best_stat_print['epoch'] = epoch + top1 = best_stat[k] + if self.output_dir: + if epoch >= self.train_dataloader.collate_fn.stop_epoch: + dist_utils.save_on_master(self.state_dict(), self.output_dir / 'best_stg2.pth') + else: + dist_utils.save_on_master(self.state_dict(), self.output_dir / 'best_stg1.pth') + + best_stat_print[k] = max(best_stat[k], top1) + print(f'best_stat: {best_stat_print}') # global best + + if best_stat['epoch'] == epoch and self.output_dir: + if epoch >= self.train_dataloader.collate_fn.stop_epoch: + if test_stats[k][0] > top1: + top1 = test_stats[k][0] + dist_utils.save_on_master(self.state_dict(), self.output_dir / 'best_stg2.pth') + else: + top1 = max(test_stats[k][0], top1) + dist_utils.save_on_master(self.state_dict(), self.output_dir / 'best_stg1.pth') + + elif epoch >= self.train_dataloader.collate_fn.stop_epoch: + best_stat = {'epoch': -1, } + self.ema.decay -= 0.0001 + self.load_resume_state(str(self.output_dir / 'best_stg1.pth')) + print(f'Refresh EMA at epoch {epoch} with decay {self.ema.decay}') + + + log_stats = { + **{f'train_{k}': v for k, v in train_stats.items()}, + **{f'test_{k}': v for k, v in test_stats.items()}, + 'epoch': epoch, + 'n_parameters': n_parameters + } + + if self.output_dir and dist_utils.is_main_process(): + with (self.output_dir / "log.txt").open("a") as f: + f.write(json.dumps(log_stats) + "\n") + + # for evaluation logs + if coco_evaluator is not None: + (self.output_dir / 'eval').mkdir(exist_ok=True) + if "bbox" in coco_evaluator.coco_eval: + filenames = ['latest.pth'] + if epoch % 50 == 0: + filenames.append(f'{epoch:03}.pth') + for name in filenames: + torch.save(coco_evaluator.coco_eval["bbox"].eval, + self.output_dir / "eval" / name) + + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print('Training time {}'.format(total_time_str)) + + + def val(self, ): + self.eval() + + module = self.ema.module if self.ema else self.model + test_stats, coco_evaluator = evaluate(module, self.criterion, self.postprocessor, + self.val_dataloader, self.evaluator, self.device) + + if self.output_dir: + dist_utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval, self.output_dir / "eval.pth") + + return diff --git a/models/best_stg2.onnx b/models/best_stg2.onnx new file mode 100644 index 0000000000000000000000000000000000000000..f31bb95125374ef76aacc12aca895656284c69e5 --- /dev/null +++ b/models/best_stg2.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9306b85b293433584f3519ef1e4180409e0fdb762fe7ebeb0134b6d480d4ba20 +size 39979898 diff --git a/models/best_stg2.pth b/models/best_stg2.pth new file mode 100644 index 0000000000000000000000000000000000000000..a118f7ef90a037e203e4c210c80efedce7e44446 --- /dev/null +++ b/models/best_stg2.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:044acb9b75b37319f68e8327bbf51cf95ad9e64297b8bd8b33ebaf0b33f6fac4 +size 156612525 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..9f5f5f8f7557f17ad204b1d3ab8eeb8f6faac634 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,18 @@ +gradio>=5.1.0 +spaces +torch==2.5.1 +torchvision==0.20.1 +numpy +Pillow>=9.0.0 +opencv-python-headless +huggingface_hub==0.25.2 +PyYAML>=6.0 +tensorboard +scipy>=1.7.0 +faster-coco-eval>=1.6.7 +calflops +transformers +onnxruntime>=1.16.0 +# DEIMv2関連の依存関係 +timm +omegaconf diff --git a/tools/benchmark/dataset.py b/tools/benchmark/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..76fa6491bac0719b0ad0bffe844d51a900038227 --- /dev/null +++ b/tools/benchmark/dataset.py @@ -0,0 +1,105 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +import os +import glob +from PIL import Image + +import torch +import torch.utils.data as data +import torchvision +import torchvision.transforms as T +import torchvision.transforms.functional as F + +Image.MAX_IMAGE_PIXELS = None + +class ToTensor(T.ToTensor): + def __init__(self) -> None: + super().__init__() + + def __call__(self, pic): + if isinstance(pic, torch.Tensor): + return pic + return super().__call__(pic) + +class PadToSize(T.Pad): + def __init__(self, size, fill=0, padding_mode='constant'): + super().__init__(0, fill, padding_mode) + self.size = size + self.fill = fill + + def __call__(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be padded. + + Returns: + PIL Image or Tensor: Padded image. + """ + w, h = F.get_image_size(img) + padding = (0, 0, self.size[0] - w, self.size[1] - h) + return F.pad(img, padding, self.fill, self.padding_mode) + + +class Dataset(data.Dataset): + def __init__(self, img_dir: str='', preprocess: T.Compose=None, device='cuda:0') -> None: + super().__init__() + + self.device = device + self.size = 640 + + self.im_path_list = list(glob.glob(os.path.join(img_dir, '*.jpg'))) + + if preprocess is None: + self.preprocess = T.Compose([ + T.Resize(size=639, max_size=640), + PadToSize(size=(640, 640), fill=114), + ToTensor(), + T.ConvertImageDtype(torch.float), + ]) + else: + self.preprocess = preprocess + + def __len__(self, ): + return len(self.im_path_list) + + def __getitem__(self, index): + # im = Image.open(self.img_path_list[index]).convert('RGB') + im = torchvision.io.read_file(self.im_path_list[index]) + im = torchvision.io.decode_jpeg(im, mode=torchvision.io.ImageReadMode.RGB, device=self.device) + _, h, w = im.shape # c,h,w + + im = self.preprocess(im) + + blob = { + 'images': im, + 'im_shape': torch.tensor([self.size, self.size]).to(im.device), + 'scale_factor': torch.tensor([self.size / h, self.size / w]).to(im.device), + 'orig_target_sizes': torch.tensor([w, h]).to(im.device), + } + + return blob + + @staticmethod + def post_process(): + pass + + @staticmethod + def collate_fn(): + pass + + +def draw_nms_result(blob, outputs, draw_score_threshold=0.25, name=''): + '''show result + Keys: + 'num_dets', 'det_boxes', 'det_scores', 'det_classes' + ''' + for i in range(blob['image'].shape[0]): + det_scores = outputs['det_scores'][i] + det_boxes = outputs['det_boxes'][i][det_scores > draw_score_threshold] + + im = (blob['image'][i] * 255).to(torch.uint8) + im = torchvision.utils.draw_bounding_boxes(im, boxes=det_boxes, width=2) + Image.fromarray(im.permute(1, 2, 0).cpu().numpy()).save(f'test_{name}_{i}.jpg') diff --git a/tools/benchmark/get_info.py b/tools/benchmark/get_info.py new file mode 100644 index 0000000000000000000000000000000000000000..b72efa35b599f2bd8d4b8440e505ba5c6ec8f2ca --- /dev/null +++ b/tools/benchmark/get_info.py @@ -0,0 +1,50 @@ +""" +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +""" + +import os +import sys +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '../..')) + +import argparse +from calflops import calculate_flops +from engine.core import YAMLConfig + +import torch +import torch.nn as nn + +def custom_repr(self): + return f'{{Tensor:{tuple(self.shape)}}} {original_repr(self)}' +original_repr = torch.Tensor.__repr__ +torch.Tensor.__repr__ = custom_repr + +def main(args, ): + """main + """ + cfg = YAMLConfig(args.config, resume=None) + class Model_for_flops(nn.Module): + def __init__(self, ) -> None: + super().__init__() + self.model = cfg.model.deploy() + + def forward(self, images): + outputs = self.model(images) + return outputs + + model = Model_for_flops().eval() + + flops, macs, _ = calculate_flops(model=model, + input_shape=(1, 3, 640, 640), + output_as_string=True, + output_precision=4) + params = sum(p.numel() for p in model.parameters()) + print("Model FLOPs:%s MACs:%s Params:%s \n" %(flops, macs, params)) + + +if __name__ == '__main__': + + parser = argparse.ArgumentParser() + parser.add_argument('--config', '-c', default= "configs/dfine/dfine_hgnetv2_l_coco.yml", type=str) + args = parser.parse_args() + + main(args) diff --git a/tools/benchmark/requirements.txt b/tools/benchmark/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..55a3c0f0a63f0ad81e1d41a19753d7c550231c27 --- /dev/null +++ b/tools/benchmark/requirements.txt @@ -0,0 +1,6 @@ +onnxruntime +tensorrt +pycuda +calflops +tqdm +# onnx_graphsurgeon # for YOLOs diff --git a/tools/benchmark/trt_benchmark.py b/tools/benchmark/trt_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..a650ac06f62b74034ee39a9ae66edb046486ad7c --- /dev/null +++ b/tools/benchmark/trt_benchmark.py @@ -0,0 +1,207 @@ +""" +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +""" + +import tensorrt as trt +import pycuda.driver as cuda +from utils import TimeProfiler +import numpy as np +import os +import time +import torch + +from collections import namedtuple, OrderedDict +import glob +import argparse +from dataset import Dataset +from tqdm import tqdm + + +def parse_args(): + parser = argparse.ArgumentParser(description='Argument Parser Example') + parser.add_argument('--COCO_dir', + type=str, + default='/data/COCO2017/val2017', + help="Directory for images to perform inference on.") + parser.add_argument("--engine_dir", + type=str, + help="Directory containing model engine files.") + parser.add_argument('--busy', + action='store_true', + help="Flag to indicate that other processes may be running.") + args = parser.parse_args() + return args + +class TRTInference(object): + def __init__(self, engine_path, device='cuda', backend='torch', max_batch_size=32, verbose=False): + self.engine_path = engine_path + self.device = device + self.backend = backend + self.max_batch_size = max_batch_size + + self.logger = trt.Logger(trt.Logger.VERBOSE) if verbose else trt.Logger(trt.Logger.INFO) + self.engine = self.load_engine(engine_path) + self.context = self.engine.create_execution_context() + self.bindings = self.get_bindings(self.engine, self.context, self.max_batch_size, self.device) + self.bindings_addr = OrderedDict((n, v.ptr) for n, v in self.bindings.items()) + self.input_names = self.get_input_names() + self.output_names = self.get_output_names() + + if self.backend == 'cuda': + self.stream = cuda.Stream() + self.time_profile = TimeProfiler() + self.time_profile_dataset = TimeProfiler() + + def init(self): + self.dynamic = False + + def load_engine(self, path): + trt.init_libnvinfer_plugins(self.logger, '') + with open(path, 'rb') as f, trt.Runtime(self.logger) as runtime: + return runtime.deserialize_cuda_engine(f.read()) + + def get_input_names(self): + names = [] + for _, name in enumerate(self.engine): + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + names.append(name) + return names + + def get_output_names(self): + names = [] + for _, name in enumerate(self.engine): + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT: + names.append(name) + return names + + def get_bindings(self, engine, context, max_batch_size=32, device=None): + Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr')) + bindings = OrderedDict() + for i, name in enumerate(engine): + shape = engine.get_tensor_shape(name) + dtype = trt.nptype(engine.get_tensor_dtype(name)) + + if shape[0] == -1: + dynamic = True + shape[0] = max_batch_size + if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + context.set_input_shape(name, shape) + + if self.backend == 'cuda': + if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + data = np.random.randn(*shape).astype(dtype) + ptr = cuda.mem_alloc(data.nbytes) + bindings[name] = Binding(name, dtype, shape, data, ptr) + else: + data = cuda.pagelocked_empty(trt.volume(shape), dtype) + ptr = cuda.mem_alloc(data.nbytes) + bindings[name] = Binding(name, dtype, shape, data, ptr) + else: + data = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device) + bindings[name] = Binding(name, dtype, shape, data, data.data_ptr()) + return bindings + + def run_torch(self, blob): + for n in self.input_names: + if self.bindings[n].shape != blob[n].shape: + self.context.set_input_shape(n, blob[n].shape) + self.bindings[n] = self.bindings[n]._replace(shape=blob[n].shape) + + self.bindings_addr.update({n: blob[n].data_ptr() for n in self.input_names}) + self.context.execute_v2(list(self.bindings_addr.values())) + outputs = {n: self.bindings[n].data for n in self.output_names} + return outputs + + def async_run_cuda(self, blob): + for n in self.input_names: + cuda.memcpy_htod_async(self.bindings_addr[n], blob[n], self.stream) + + bindings_addr = [int(v) for _, v in self.bindings_addr.items()] + self.context.execute_async_v2(bindings=bindings_addr, stream_handle=self.stream.handle) + + outputs = {} + for n in self.output_names: + cuda.memcpy_dtoh_async(self.bindings[n].data, self.bindings[n].ptr, self.stream) + outputs[n] = self.bindings[n].data + + self.stream.synchronize() + + return outputs + + def __call__(self, blob): + if self.backend == 'torch': + return self.run_torch(blob) + elif self.backend == 'cuda': + return self.async_run_cuda(blob) + + def synchronize(self): + if self.backend == 'torch' and torch.cuda.is_available(): + torch.cuda.synchronize() + elif self.backend == 'cuda': + self.stream.synchronize() + + def warmup(self, blob, n): + for _ in range(n): + _ = self(blob) + + def speed(self, blob, n, nonempty_process=False): + times = [] + self.time_profile_dataset.reset() + for i in tqdm(range(n), desc="Running Inference", unit="iteration"): + self.time_profile.reset() + with self.time_profile_dataset: + img = blob[i] + if img['images'] is not None: + img['image'] = img['input'] = img['images'].unsqueeze(0) + else: + img['images'] = img['input'] = img['image'].unsqueeze(0) + with self.time_profile: + _ = self(img) + times.append(self.time_profile.total) + + # end-to-end model only + times = sorted(times) + if len(times) > 100 and nonempty_process: + times = times[:100] + + avg_time = sum(times) / len(times) # Calculate the average of the remaining times + return avg_time + +def main(): + FLAGS = parse_args() + dataset = Dataset(FLAGS.infer_dir) + im = torch.ones(1, 3, 640, 640).cuda() + blob = { + 'image': im, + 'images': im, + 'input': im, + 'im_shape': torch.tensor([640, 640]).to(im.device), + 'scale_factor': torch.tensor([1, 1]).to(im.device), + 'orig_target_sizes': torch.tensor([640, 640]).to(im.device), + } + + engine_files = glob.glob(os.path.join(FLAGS.models_dir, "*.engine")) + results = [] + + for engine_file in engine_files: + print(f"Testing engine: {engine_file}") + model = TRTInference(engine_file, max_batch_size=1, verbose=False) + model.init() + model.warmup(blob, 1000) + t = [] + for _ in range(1): + t.append(model.speed(dataset, 1000, FLAGS.busy)) + avg_latency = 1000 * torch.tensor(t).mean() + results.append((engine_file, avg_latency)) + print(f"Engine: {engine_file}, Latency: {avg_latency:.2f} ms") + + del model + torch.cuda.empty_cache() + time.sleep(1) + + sorted_results = sorted(results, key=lambda x: x[1]) + for engine_file, latency in sorted_results: + print(f"Engine: {engine_file}, Latency: {latency:.2f} ms") + +if __name__ == '__main__': + main() diff --git a/tools/benchmark/utils.py b/tools/benchmark/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..23e1800a2deaf84f3fec46c1a40b6e29f4772719 --- /dev/null +++ b/tools/benchmark/utils.py @@ -0,0 +1,80 @@ +import time +import contextlib +import numpy as np +from PIL import Image +from collections import OrderedDict + +import onnx +import torch +import onnx_graphsurgeon + + +def to_binary_data(path, size=(640, 640), output_name='input_tensor.bin'): + '''--loadInputs='image:input_tensor.bin' + ''' + im = Image.open(path).resize(size) + data = np.asarray(im, dtype=np.float32).transpose(2, 0, 1)[None] / 255. + data.tofile(output_name) + + +def yolo_insert_nms(path, score_threshold=0.01, iou_threshold=0.7, max_output_boxes=300, simplify=False): + ''' + http://www.xavierdupre.fr/app/onnxcustom/helpsphinx/api/onnxops/onnx__EfficientNMS_TRT.html + https://huggingface.co/spaces/muttalib1326/Punjabi_Character_Detection/blob/3dd1e17054c64e5f6b2254278f96cfa2bf418cd4/utils/add_nms.py + ''' + onnx_model = onnx.load(path) + + if simplify: + from onnxsim import simplify + onnx_model, _ = simplify(onnx_model, overwrite_input_shapes={'image': [1, 3, 640, 640]}) + + graph = onnx_graphsurgeon.import_onnx(onnx_model) + graph.toposort() + graph.fold_constants() + graph.cleanup() + + topk = max_output_boxes + attrs = OrderedDict(plugin_version='1', + background_class=-1, + max_output_boxes=topk, + score_threshold=score_threshold, + iou_threshold=iou_threshold, + score_activation=False, + box_coding=0, ) + + outputs = [onnx_graphsurgeon.Variable('num_dets', np.int32, [-1, 1]), + onnx_graphsurgeon.Variable('det_boxes', np.float32, [-1, topk, 4]), + onnx_graphsurgeon.Variable('det_scores', np.float32, [-1, topk]), + onnx_graphsurgeon.Variable('det_classes', np.int32, [-1, topk])] + + graph.layer(op='EfficientNMS_TRT', + name="batched_nms", + inputs=[graph.outputs[0], + graph.outputs[1]], + outputs=outputs, + attrs=attrs, ) + + graph.outputs = outputs + graph.cleanup().toposort() + + onnx.save(onnx_graphsurgeon.export_onnx(graph), 'yolo_w_nms.onnx') + + +class TimeProfiler(contextlib.ContextDecorator): + def __init__(self, ): + self.total = 0 + + def __enter__(self, ): + self.start = self.time() + return self + + def __exit__(self, type, value, traceback): + self.total += self.time() - self.start + + def reset(self, ): + self.total = 0 + + def time(self, ): + if torch.cuda.is_available(): + torch.cuda.synchronize() + return time.time() diff --git a/tools/dataset/remap_obj365.py b/tools/dataset/remap_obj365.py new file mode 100644 index 0000000000000000000000000000000000000000..f76214e7a05b5c158deaecbe4f4994e020bf8226 --- /dev/null +++ b/tools/dataset/remap_obj365.py @@ -0,0 +1,139 @@ +""" +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +""" + +import json +import os +import argparse + + +def update_image_paths(images, new_prefix): + print('Updating image paths with new prefix...') + for img in images: + split = img['file_name'].split('/')[1:] + img['file_name'] = os.path.join(new_prefix, *split) + print('Image paths updated.') + return images + +def create_split_annotations(original_annotations, split_image_ids, new_prefix, output_file): + print(f'Creating split annotations for {output_file}...') + new_images = [img for img in original_annotations['images'] if img['id'] in split_image_ids] + print(f'Number of images selected: {len(new_images)}') + if new_prefix is not None: + new_images = update_image_paths(new_images, new_prefix) + + new_annotations = { + 'images': new_images, + 'annotations': [ann for ann in original_annotations['annotations'] if ann['image_id'] in split_image_ids], + 'categories': original_annotations['categories'] + } + print(f'Number of annotations selected: {len(new_annotations["annotations"])}') + with open(output_file, 'w') as f: + json.dump(new_annotations, f) + print(f'Annotations saved to {output_file}') + +def parse_arguments(): + parser = argparse.ArgumentParser(description='Split and update dataset annotations.') + parser.add_argument( + '--base_dir', + type=str, + default='/datassd/objects365', + help='Base directory of the dataset, e.g., /data/Objects365/data' + ) + parser.add_argument( + '--new_val_size', + type=int, + default=5000, + help='Number of images to include in the new validation set (default: 5000)' + ) + parser.add_argument( + '--output_suffix', + type=str, + default='new', + help='Suffix to add to new annotation files (default: new)' + ) + return parser.parse_args() + +def main(): + args = parse_arguments() + base_dir = args.base_dir + new_val_size = args.new_val_size + output_suffix = args.output_suffix + + # Define paths based on the base directory + original_train_ann_file = os.path.join(base_dir, 'train', 'zhiyuan_objv2_train.json') + original_val_ann_file = os.path.join(base_dir, 'val', 'zhiyuan_objv2_val.json') + + new_val_ann_file = os.path.join(base_dir, 'val', f'{output_suffix}_zhiyuan_objv2_val.json') + new_train_ann_file = os.path.join(base_dir, 'train', f'{output_suffix}_zhiyuan_objv2_train.json') + + # Check if original annotation files exist + if not os.path.isfile(original_train_ann_file): + print(f'Error: Training annotation file not found at {original_train_ann_file}') + return + if not os.path.isfile(original_val_ann_file): + print(f'Error: Validation annotation file not found at {original_val_ann_file}') + return + + # Load the original training and validation annotations + print('Loading original training annotations...') + with open(original_train_ann_file, 'r') as f: + train_annotations = json.load(f) + print('Training annotations loaded.') + + print('Loading original validation annotations...') + with open(original_val_ann_file, 'r') as f: + val_annotations = json.load(f) + print('Validation annotations loaded.') + + # Extract image IDs from the original validation set + print('Extracting image IDs from the validation set...') + val_image_ids = [img['id'] for img in val_annotations['images']] + print(f'Total validation images: {len(val_image_ids)}') + + # Split image IDs for the new training and validation sets + print(f'Splitting validation images into new validation set of size {new_val_size} and training set...') + new_val_image_ids = val_image_ids[:new_val_size] + new_train_image_ids = val_image_ids[new_val_size:] + print(f'New validation set size: {len(new_val_image_ids)}') + print(f'New training set size from validation images: {len(new_train_image_ids)}') + + # Create new validation annotation file + print('Creating new validation annotations...') + create_split_annotations(val_annotations, new_val_image_ids, None, new_val_ann_file) + print('New validation annotations created.') + + # Combine the remaining validation images and annotations with the original training data + print('Preparing new training images and annotations...') + new_train_images = [img for img in val_annotations['images'] if img['id'] in new_train_image_ids] + print(f'Number of images from validation to add to training: {len(new_train_images)}') + new_train_images = update_image_paths(new_train_images, 'images_from_val') + new_train_annotations = [ann for ann in val_annotations['annotations'] if ann['image_id'] in new_train_image_ids] + print(f'Number of annotations from validation to add to training: {len(new_train_annotations)}') + + # Add the original training images and annotations + print('Adding original training images and annotations...') + new_train_images.extend(train_annotations['images']) + new_train_annotations.extend(train_annotations['annotations']) + print(f'Total training images: {len(new_train_images)}') + print(f'Total training annotations: {len(new_train_annotations)}') + + # Create a new training annotation dictionary + print('Creating new training annotations dictionary...') + new_train_annotations_dict = { + 'images': new_train_images, + 'annotations': new_train_annotations, + 'categories': train_annotations['categories'] + } + print('New training annotations dictionary created.') + + # Save the new training annotations + print('Saving new training annotations...') + with open(new_train_ann_file, 'w') as f: + json.dump(new_train_annotations_dict, f) + print(f'New training annotations saved to {new_train_ann_file}') + + print('Processing completed successfully.') + +if __name__ == '__main__': + main() diff --git a/tools/dataset/resize_obj365.py b/tools/dataset/resize_obj365.py new file mode 100644 index 0000000000000000000000000000000000000000..d14fd865ef5e260d0d6e8b26f32daddc8c088b69 --- /dev/null +++ b/tools/dataset/resize_obj365.py @@ -0,0 +1,147 @@ +""" +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +""" + +import os +import json +from PIL import Image +from concurrent.futures import ThreadPoolExecutor +import argparse + + +def resize_image_and_update_annotations(image_path, annotations, max_size=640): + print(f"Processing image: {image_path}") + try: + with Image.open(image_path) as img: + w, h = img.size + if max(w, h) <= max_size: + return annotations, w, h, False # No need to resize + + scale = max_size / max(w, h) + new_w = int(w * scale) + new_h = int(h * scale) + print(f"Resizing image to width={new_w}, height={new_h}") + + img = img.resize((new_w, new_h), Image.Resampling.LANCZOS) + new_image_path = image_path.replace('.jpg', '_resized{}.jpg'.format(max_size)) + img.save(new_image_path) + print(f"Resized image saved: {new_image_path}") + print(f"Original size: ({w}, {h}), New size: ({new_w}, {new_h})") + + # Update annotations + for ann in annotations: + ann['area'] = ann['area'] * (scale ** 2) + ann['bbox'] = [coord * scale for coord in ann['bbox']] + if 'orig_size' in ann: + ann['orig_size'] = (new_w, new_h) + if 'size' in ann: + ann['size'] = (new_w, new_h) + + except Exception as e: + print(f"Error processing {image_path}: {e}") + return None + + return annotations, new_w, new_h, True + +def resize_images_and_update_annotations(base_dir, subset, max_size=640, num_workers=4): + print(f"Starting to resize images and update annotations for subset: {subset}") + json_file = os.path.join(base_dir, subset, 'new_zhiyuan_objv2_{}.json'.format(subset)) + if not os.path.isfile(json_file): + print(f'Error: JSON file not found at {json_file}') + return + + print(f"Loading JSON file: {json_file}") + with open(json_file, 'r') as f: + data = json.load(f) + print("JSON file loaded.") + + print("Preparing image annotations mapping...") + image_annotations = {img['id']: [] for img in data['images']} + for ann in data['annotations']: + image_annotations[ann['image_id']].append(ann) + print("Image annotations mapping prepared.") + + def process_image(image_info): + image_path = os.path.join(base_dir, subset, image_info['file_name']) + results = resize_image_and_update_annotations(image_path, image_annotations[image_info['id']], max_size) + if results is None: + updated_annotations, new_w, new_h, resized = None, None, None, None + else: + updated_annotations, new_w, new_h, resized = results + return image_info, updated_annotations, new_w, new_h, resized + + print(f"Processing images with {num_workers} worker threads...") + with ThreadPoolExecutor(max_workers=num_workers) as executor: + results = list(executor.map(process_image, data['images'])) + print("Image processing completed.") + + new_images = [] + new_annotations = [] + + print("Updating image and annotation data...") + for image_info, updated_annotations, new_w, new_h, resized in results: + if updated_annotations is not None: + image_info['width'] = new_w + image_info['height'] = new_h + image_annotations[image_info['id']] = updated_annotations + if resized: + image_info['file_name'] = image_info['file_name'].replace('.jpg', '_resized{}.jpg'.format(max_size)) + new_images.append(image_info) + new_annotations.extend(updated_annotations) + print(f"Total images processed: {len(new_images)}") + print(f"Total annotations updated: {len(new_annotations)}") + + new_data = { + 'images': new_images, + 'annotations': new_annotations, + 'categories': data['categories'] + } + + new_json_file = json_file.replace('.json', '_resized{}.json'.format(max_size)) + print('Saving new training annotations...') + with open(new_json_file, 'w') as f: + json.dump(new_data, f) + print(f'New JSON file saved to {new_json_file}') + +def parse_arguments(): + parser = argparse.ArgumentParser(description='Resize images and update dataset annotations for both train and val sets.') + parser.add_argument( + '--base_dir', + type=str, + default='/datassd/objects365', + help='Base directory of the dataset, e.g., /data/Objects365/data' + ) + parser.add_argument( + '--max_size', + type=int, + default=640, + help='Maximum size for the longer side of the image (default: 640)' + ) + parser.add_argument( + '--num_workers', + type=int, + default=4, + help='Number of worker threads for parallel processing (default: 4)' + ) + args = parser.parse_args() + return args + +def main(): + args = parse_arguments() + base_dir = args.base_dir + max_size = args.max_size + num_workers = args.num_workers + + subsets = ['train', 'val'] + for subset in subsets: + print(f'Processing subset: {subset}') + resize_images_and_update_annotations( + base_dir=base_dir, + subset=subset, + max_size=max_size, + num_workers=num_workers + ) + print("All subsets processed.") + +if __name__ == "__main__": + main() diff --git a/tools/deployment/export_onnx.py b/tools/deployment/export_onnx.py new file mode 100644 index 0000000000000000000000000000000000000000..186fda35f2a319dc841bcd2d38752df5c851b1b6 --- /dev/null +++ b/tools/deployment/export_onnx.py @@ -0,0 +1,109 @@ +""" +DEIMv2: Real-Time Object Detection Meets DINOv3 +Copyright (c) 2025 The DEIMv2 Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +D-FINE: Redefine Regression Task of DETRs as Fine-grained Distribution Refinement +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright (c) 2023 lyuwenyu. All Rights Reserved. +""" + +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '../..')) + +import torch +import torch.nn as nn + +from engine.core import YAMLConfig + + +def main(args, ): + """main + """ + cfg = YAMLConfig(args.config, resume=args.resume) + + if 'HGNetv2' in cfg.yaml_cfg: + cfg.yaml_cfg['HGNetv2']['pretrained'] = False + + if args.resume: + checkpoint = torch.load(args.resume, map_location='cpu') + if 'ema' in checkpoint: + state = checkpoint['ema']['module'] + else: + state = checkpoint['model'] + + # NOTE load train mode state -> convert to deploy mode + cfg.model.load_state_dict(state) + + else: + # raise AttributeError('Only support resume to load model.state_dict by now.') + print('not load model.state_dict, use default init state dict...') + + class Model(nn.Module): + def __init__(self, ) -> None: + 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 + + model = Model() + + img_size = cfg.yaml_cfg["eval_spatial_size"] + data = torch.rand(32, 3, *img_size) + size = torch.tensor([img_size]) + _ = model(data, size) + + dynamic_axes = { + 'images': {0: 'N', }, + 'orig_target_sizes': {0: 'N'} + } + + output_file = args.resume.replace('.pth', '.onnx') if args.resume else 'model.onnx' + + torch.onnx.export( + model, + (data, size), + output_file, + input_names=['images', 'orig_target_sizes'], + output_names=['labels', 'boxes', 'scores'], + dynamic_axes=dynamic_axes, + opset_version=args.opset, + verbose=False, + do_constant_folding=True, + ) + + if args.check: + import onnx + onnx_model = onnx.load(output_file) + onnx.checker.check_model(onnx_model) + print('Check export onnx model done...') + + if args.simplify: + import onnx + import onnxsim + dynamic = True + # input_shapes = {'images': [1, 3, 640, 640], 'orig_target_sizes': [1, 2]} if dynamic else None + input_shapes = {'images': data.shape, 'orig_target_sizes': size.shape} if dynamic else None + onnx_model_simplify, check = onnxsim.simplify(output_file, test_input_shapes=input_shapes) + onnx.save(onnx_model_simplify, output_file) + print(f'Simplify onnx model {check}...') + + +if __name__ == '__main__': + + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('--config', '-c', default='configs/dfine/dfine_hgnetv2_l_coco.yml', type=str, ) + parser.add_argument('--resume', '-r', type=str, ) + parser.add_argument('--opset', type=int, default=17,) + parser.add_argument('--check', action='store_true') + parser.add_argument('--simplify', action='store_true') + args = parser.parse_args() + main(args) diff --git a/tools/deployment/export_yolo_w_nms.py b/tools/deployment/export_yolo_w_nms.py new file mode 100644 index 0000000000000000000000000000000000000000..95c89b213d9436cba49c093d4427757b907a5e03 --- /dev/null +++ b/tools/deployment/export_yolo_w_nms.py @@ -0,0 +1,74 @@ +import torch +import torchvision + +import numpy as np +import onnxruntime as ort + +from utils import yolo_insert_nms + +class YOLO11(torch.nn.Module): + def __init__(self, name) -> None: + super().__init__() + from ultralytics import YOLO + # Load a model + # build a new model from scratch + # model = YOLO(f'{name}.yaml') + + # load a pretrained model (recommended for training) + model = YOLO("yolo11n.pt") + self.model = model.model + + def forward(self, x): + '''https://github.com/ultralytics/ultralytics/blob/main/ultralytics/nn/tasks.py#L216 + ''' + pred: torch.Tensor = self.model(x)[0] # n 84 8400, + pred = pred.permute(0, 2, 1) + boxes, scores = pred.split([4, 80], dim=-1) + boxes = torchvision.ops.box_convert(boxes, in_fmt='cxcywh', out_fmt='xyxy') + + return boxes, scores + + + +def export_onnx(name='yolov8n'): + '''export onnx + ''' + m = YOLO11(name) + + x = torch.rand(1, 3, 640, 640) + dynamic_axes = { + 'image': {0: '-1'} + } + torch.onnx.export(m, x, f'{name}.onnx', + input_names=['image'], + output_names=['boxes', 'scores'], + opset_version=13, + dynamic_axes=dynamic_axes) + + data = np.random.rand(1, 3, 640, 640).astype(np.float32) + sess = ort.InferenceSession(f'{name}.onnx') + _ = sess.run(output_names=None, input_feed={'image': data}) + + import onnx + import onnxslim + model_onnx = onnx.load(f'{name}.onnx') + model_onnx = onnxslim.slim(model_onnx) + onnx.save(model_onnx, f'{name}.onnx') + + +if __name__ == '__main__': + + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('--name', type=str, default='yolo11n_tuned') + parser.add_argument('--score_threshold', type=float, default=0.01) + parser.add_argument('--iou_threshold', type=float, default=0.6) + parser.add_argument('--max_output_boxes', type=int, default=300) + args = parser.parse_args() + + export_onnx(name=args.name) + + yolo_insert_nms(path=f'{args.name}.onnx', + score_threshold=args.score_threshold, + iou_threshold=args.iou_threshold, + max_output_boxes=args.max_output_boxes, ) diff --git a/tools/inference/onnx_inf.py b/tools/inference/onnx_inf.py new file mode 100644 index 0000000000000000000000000000000000000000..dd0016019cd43a9b8d64c454493c7c4283a5b0d4 --- /dev/null +++ b/tools/inference/onnx_inf.py @@ -0,0 +1,175 @@ +""" +DEIMv2: Real-Time Object Detection Meets DINOv3 +Copyright (c) 2025 The DEIMv2 Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE) +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +""" + +import cv2 +import numpy as np +import onnxruntime as ort +import torch +import torchvision.transforms as T +from PIL import Image, ImageDraw + + +def resize_with_aspect_ratio(image, size, interpolation=Image.BILINEAR): + """Resizes an image while maintaining aspect ratio and pads it.""" + original_width, original_height = image.size + ratio = min(size / original_width, size / original_height) + new_width = int(original_width * ratio) + new_height = int(original_height * ratio) + image = image.resize((new_width, new_height), interpolation) + + # Create a new image with the desired size and paste the resized image onto it + new_image = Image.new("RGB", (size, size)) + new_image.paste(image, ((size - new_width) // 2, (size - new_height) // 2)) + return new_image, ratio, (size - new_width) // 2, (size - new_height) // 2 + + +def draw(images, labels, boxes, scores, ratios, paddings, thrh=0.4): + result_images = [] + for i, im in enumerate(images): + draw = ImageDraw.Draw(im) + scr = scores[i] + lab = labels[i][scr > thrh] + box = boxes[i][scr > thrh] + scr = scr[scr > thrh] + + ratio = ratios[i] + pad_w, pad_h = paddings[i] + + for lbl, bb in zip(lab, box): + # Adjust bounding boxes according to the resizing and padding + bb = [ + (bb[0] - pad_w) / ratio, + (bb[1] - pad_h) / ratio, + (bb[2] - pad_w) / ratio, + (bb[3] - pad_h) / ratio, + ] + draw.rectangle(bb, outline='red') + draw.text((bb[0], bb[1]), text=str(lbl), fill='blue') + + result_images.append(im) + return result_images + + +def process_image(sess, im_pil, size=640, model_size='s'): + # Resize image while preserving aspect ratio + resized_im_pil, ratio, pad_w, pad_h = resize_with_aspect_ratio(im_pil, size) + orig_size = torch.tensor([[resized_im_pil.size[1], resized_im_pil.size[0]]]) + + transforms = T.Compose([ + T.ToTensor(), + T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + if model_size not in ['atto', 'femto', 'pico', 'n'] + else T.Lambda(lambda x: x) + ]) + im_data = transforms(resized_im_pil).unsqueeze(0) + + output = sess.run( + output_names=None, + input_feed={'images': im_data.numpy(), "orig_target_sizes": orig_size.numpy()} + ) + + labels, boxes, scores = output + + result_images = draw( + [im_pil], labels, boxes, scores, + [ratio], [(pad_w, pad_h)] + ) + result_images[0].save('onnx_result.jpg') + print("Image processing complete. Result saved as 'result.jpg'.") + + +def process_video(sess, video_path, size=640, model_size='s'): + cap = cv2.VideoCapture(video_path) + + # Get video properties + fps = cap.get(cv2.CAP_PROP_FPS) + orig_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + orig_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + # Define the codec and create VideoWriter object + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter('onnx_result.mp4', fourcc, fps, (orig_w, orig_h)) + + frame_count = 0 + print("Processing video frames...") + while cap.isOpened(): + ret, frame = cap.read() + if not ret: + break + + # Convert frame to PIL image + frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + + # Resize frame while preserving aspect ratio + resized_frame_pil, ratio, pad_w, pad_h = resize_with_aspect_ratio(frame_pil, size) + orig_size = torch.tensor([[resized_frame_pil.size[1], resized_frame_pil.size[0]]]) + + transforms = T.Compose([ + T.ToTensor(), + T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + if model_size not in ['atto', 'femto', 'pico', 'n'] + else T.Lambda(lambda x: x) + ]) + im_data = transforms(resized_frame_pil).unsqueeze(0) + + output = sess.run( + output_names=None, + input_feed={'images': im_data.numpy(), "orig_target_sizes": orig_size.numpy()} + ) + + labels, boxes, scores = output + + # Draw detections on the original frame + result_images = draw( + [frame_pil], labels, boxes, scores, + [ratio], [(pad_w, pad_h)] + ) + frame_with_detections = result_images[0] + + # Convert back to OpenCV image + frame = cv2.cvtColor(np.array(frame_with_detections), cv2.COLOR_RGB2BGR) + + # Write the frame + out.write(frame) + frame_count += 1 + + if frame_count % 10 == 0: + print(f"Processed {frame_count} frames...") + + cap.release() + out.release() + print("Video processing complete. Result saved as 'result.mp4'.") + + +def main(args): + """Main function.""" + # Load the ONNX model + sess = ort.InferenceSession(args.onnx) + size = sess.get_inputs()[0].shape[2] + print(f"Using device: {ort.get_device()}") + + input_path = args.input + + try: + # Try to open the input as an image + im_pil = Image.open(input_path).convert('RGB') + process_image(sess, im_pil, size, args.model_size) + except IOError: + # Not an image, process as video + process_video(sess, input_path, size, args.model_size) + + +if __name__ == '__main__': + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('--onnx', type=str, required=True, help='Path to the ONNX model file.') + parser.add_argument('--input', type=str, required=True, help='Path to the input image or video file.') + parser.add_argument('-ms', '--model-size', type=str, required=True, choices=['atto', 'femto', 'pico', 'n', 's', 'm', 'l', 'x'], + help='Model size') + args = parser.parse_args() + main(args) diff --git a/tools/inference/openvino_inf.py b/tools/inference/openvino_inf.py new file mode 100644 index 0000000000000000000000000000000000000000..4a66755a256f56a5594508d003d1820df23fd2e3 --- /dev/null +++ b/tools/inference/openvino_inf.py @@ -0,0 +1,7 @@ +""" +Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR) +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + + +# please reference: https://github.com/guojin-yan/RT-DETR-OpenVINO diff --git a/tools/inference/requirements.txt b/tools/inference/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..46a470c8805cfc80a3bf4c4cb6a87dacc237021f --- /dev/null +++ b/tools/inference/requirements.txt @@ -0,0 +1,2 @@ +onnxruntime +tensorrt diff --git a/tools/inference/torch_inf.py b/tools/inference/torch_inf.py new file mode 100644 index 0000000000000000000000000000000000000000..86e090016bc54870482d9bb67a22af5b1ef27227 --- /dev/null +++ b/tools/inference/torch_inf.py @@ -0,0 +1,167 @@ +""" +DEIMv2: Real-Time Object Detection Meets DINOv3 +Copyright (c) 2025 The DEIMv2 Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE) +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +""" + +import os +import sys + +import cv2 # Added for video processing +import numpy as np +import torch +import torch.nn as nn +import torchvision.transforms as T +from PIL import Image, ImageDraw + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../'))) +from engine.core import YAMLConfig + + +def draw(images, labels, boxes, scores, thrh=0.45): + for i, im in enumerate(images): + draw = ImageDraw.Draw(im) + + scr = scores[i] + lab = labels[i][scr > thrh] + box = boxes[i][scr > thrh] + scrs = scr[scr > thrh] + + for j, b in enumerate(box): + draw.rectangle(list(b), outline='red') + draw.text((b[0], b[1]), text=f"{lab[j].item()} {round(scrs[j].item(), 2)}", fill='blue', ) + + im.save('torch_results.jpg') + + +def process_image(model, device, file_path, size=(640, 640), vit_backbone=False): + im_pil = Image.open(file_path).convert('RGB') + w, h = im_pil.size + orig_size = torch.tensor([[w, h]]).to(device) + + transforms = T.Compose([ + T.Resize(size), + T.ToTensor(), + T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + if vit_backbone else T.Lambda(lambda x: x) + ]) + im_data = transforms(im_pil).unsqueeze(0).to(device) + + output = model(im_data, orig_size) + labels, boxes, scores = output + + draw([im_pil], labels, boxes, scores) + + +def process_video(model, device, file_path, size=(640, 640), vit_backbone=False): + cap = cv2.VideoCapture(file_path) + + # Get video properties + fps = cap.get(cv2.CAP_PROP_FPS) + orig_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + orig_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + # Define the codec and create VideoWriter object + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter('torch_results.mp4', fourcc, fps, (orig_w, orig_h)) + + transforms = T.Compose([ + T.Resize(size), + T.ToTensor(), + T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + if vit_backbone else T.Lambda(lambda x: x) + ]) + + frame_count = 0 + print("Processing video frames...") + while cap.isOpened(): + ret, frame = cap.read() + if not ret: + break + + # Convert frame to PIL image + frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + + w, h = frame_pil.size + orig_size = torch.tensor([[w, h]]).to(device) + + im_data = transforms(frame_pil).unsqueeze(0).to(device) + + output = model(im_data, orig_size) + labels, boxes, scores = output + + # Draw detections on the frame + draw([frame_pil], labels, boxes, scores) + + # Convert back to OpenCV image + frame = cv2.cvtColor(np.array(frame_pil), cv2.COLOR_RGB2BGR) + + # Write the frame + out.write(frame) + frame_count += 1 + + if frame_count % 10 == 0: + print(f"Processed {frame_count} frames...") + + cap.release() + out.release() + print("Video processing complete. Result saved as 'results_video.mp4'.") + + +def main(args): + """Main function""" + cfg = YAMLConfig(args.config, resume=args.resume) + + if 'HGNetv2' in cfg.yaml_cfg: + cfg.yaml_cfg['HGNetv2']['pretrained'] = False + + if args.resume: + checkpoint = torch.load(args.resume, map_location='cpu') + if 'ema' in checkpoint: + state = checkpoint['ema']['module'] + else: + state = checkpoint['model'] + else: + raise AttributeError('Only support resume to load model.state_dict by now.') + + # Load train mode state and convert to deploy mode + 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 + + device = args.device + model = Model().to(device) + img_size = cfg.yaml_cfg["eval_spatial_size"] + vit_backbone = cfg.yaml_cfg.get('DINOv3STAs', False) + + # Check if the input file is an image or a video + file_path = args.input + if os.path.splitext(file_path)[-1].lower() in ['.jpg', '.jpeg', '.png', '.bmp']: + # Process as image + process_image(model, device, file_path, img_size, vit_backbone) + print("Image processing complete.") + else: + # Process as video + process_video(model, device, file_path, img_size, vit_backbone) + + +if __name__ == '__main__': + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('-c', '--config', type=str, required=True) + parser.add_argument('-r', '--resume', type=str, required=True) + parser.add_argument('-i', '--input', type=str, required=True) + parser.add_argument('-d', '--device', type=str, default='cpu') + args = parser.parse_args() + main(args) diff --git a/tools/inference/torch_inf_vis.py b/tools/inference/torch_inf_vis.py new file mode 100644 index 0000000000000000000000000000000000000000..dc5ef632c84a1e1fe69441d05aa7aaceb5423f38 --- /dev/null +++ b/tools/inference/torch_inf_vis.py @@ -0,0 +1,155 @@ +""" +DEIMv2: Real-Time Object Detection Meets DINOv3 +Copyright (c) 2025 The DEIMv2 Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE) +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +""" + +import os +import random +import sys + +import cv2 # Added for video processing +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.nn as nn +import torchvision.transforms as T +from PIL import Image, ImageDraw, ImageFont + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../'))) +from engine.core import YAMLConfig + +label_map = { + 1: 'person', 2: 'bicycle', 3: 'car', 4: 'motorbike', 5: 'aeroplane', + 6: 'bus', 7: 'train', 8: 'truck', 9: 'boat', 10: 'trafficlight', + 11: 'firehydrant', 12: 'streetsign', 13: 'stopsign', 14: 'parkingmeter', + 15: 'bench', 16: 'bird', 17: 'cat', 18: 'dog', 19: 'horse', + 20: 'sheep', 21: 'cow', 22: 'elephant', 23: 'bear', 24: 'zebra', + 25: 'giraffe', 26: 'hat', 27: 'backpack', 28: 'umbrella', 29: 'shoe', + 30: 'eyeglasses', 31: 'handbag', 32: 'tie', 33: 'suitcase', 34: 'frisbee', + 35: 'skis', 36: 'snowboard', 37: 'sportsball', 38: 'kite', 39: 'baseballbat', + 40: 'baseballglove', 41: 'skateboard', 42: 'surfboard', 43: 'tennisracket', + 44: 'bottle', 45: 'plate', 46: 'wineglass', 47: 'cup', 48: 'fork', + 49: 'knife', 50: 'spoon', 51: 'bowl', 52: 'banana', 53: 'apple', + 54: 'sandwich', 55: 'orange', 56: 'broccoli', 57: 'carrot', 58: 'hotdog', + 59: 'pizza', 60: 'donut', 61: 'cake', 62: 'chair', 63: 'sofa', + 64: 'pottedplant', 65: 'bed', 66: 'mirror', 67: 'diningtable', 68: 'window', + 69: 'desk', 70: 'toilet', 71: 'door', 72: 'tv', 73: 'laptop', + 74: 'mouse', 75: 'remote', 76: 'keyboard', 77: 'cellphone', 78: 'microwave', + 79: 'oven', 80: 'toaster', 81: 'sink', 82: 'refrigerator', 83: 'blender', + 84: 'book', 85: 'clock', 86: 'vase', 87: 'scissors', 88: 'teddybear', + 89: 'hairdrier', 90: 'toothbrush', 91: 'hairbrush' +} + + +COLORS = plt.cm.tab20.colors +COLOR_MAP = {label: tuple([int(c * 255) for c in COLORS[i % len(COLORS)]]) for i, label in enumerate(label_map)} + + + +def draw(image, labels, boxes, scores, thrh=0.45): + draw = ImageDraw.Draw(image) + font = ImageFont.load_default() + labels, boxes, scores = labels[scores > thrh], boxes[scores > thrh], scores[scores > thrh] + + for j, box in enumerate(boxes): + category = labels[j].item() + color = COLOR_MAP.get(category, (255, 255, 255)) + box = list(map(int, box)) + + + draw.rectangle(box, outline=color, width=3) + + text = f"{label_map[category]} {scores[j].item():.2f}" + text_bbox = draw.textbbox((0, 0), text, font=font) + text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1] + + text_background = [box[0], box[1] - text_height - 2, box[0] + text_width + 4, box[1]] + draw.rectangle(text_background, fill=color) + + draw.text((box[0] + 2, box[1] - text_height - 2), text, fill="black", font=font) + + return image + + +def process_dataset(model, dataset_path, output_path, thrh=0.5, size=(640, 640), vit_backbone=False): + os.makedirs(output_path, exist_ok=True) + image_paths = [os.path.join(dataset_path, f) for f in os.listdir(dataset_path) if f.endswith(('.jpg', '.png'))] + + transforms = T.Compose([ + T.Resize(size), + T.ToTensor(), + T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + if vit_backbone else T.Lambda(lambda x: x) + ]) + + print(f"Found {len(image_paths)} images in validation set...") + for idx, file_path in enumerate(image_paths): + im_pil = Image.open(file_path).convert('RGB') + w, h = im_pil.size + orig_size = torch.tensor([[w, h]]).cuda() + + # 图像预处理 + im_data = transforms(im_pil).unsqueeze(0).cuda() + output = model(im_data, orig_size) + labels, boxes, scores = output[0]['labels'], output[0]['boxes'], output[0]['scores'] + + # 绘制结果 + vis_image = draw(im_pil.copy(), labels, boxes, scores, thrh) + save_path = os.path.join(output_path, f"vis_{os.path.basename(file_path)}") + vis_image.save(save_path) + + if idx % 500 == 0: + print(f"Processed {idx}/{len(image_paths)} images...") + + print("Visualization complete. Results saved in:", output_path) + + +def main(args): + """Main function""" + cfg = YAMLConfig(args.config, resume=args.resume) + + if 'HGNetv2' in cfg.yaml_cfg: + cfg.yaml_cfg['HGNetv2']['pretrained'] = False + + if args.resume: + checkpoint = torch.load(args.resume, map_location='cpu') + if 'ema' in checkpoint: + state = checkpoint['ema']['module'] + else: + state = checkpoint['model'] + else: + raise AttributeError('Only support resume to load model.state_dict by now.') + + # Load train mode state and convert to deploy mode + cfg.model.load_state_dict(state) + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.model = cfg.model.eval().cuda() + self.postprocessor = cfg.postprocessor.eval().cuda() + + def forward(self, images, orig_target_sizes): + outputs = self.model(images) + outputs = self.postprocessor(outputs, orig_target_sizes) + return outputs + + model = Model() + img_size = cfg.yaml_cfg["eval_spatial_size"] + vit_backbone = cfg.yaml_cfg.get('DINOv3STAs', False) + + process_dataset(model, args.dataset, args.output, thrh=0.45, size=img_size, vit_backbone=vit_backbone) + + +if __name__ == '__main__': + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('-c', '--config', type=str, required=True) + parser.add_argument('-r', '--resume', type=str, required=True) + parser.add_argument('-d', '--dataset', type=str, default='./data/fiftyone/validation/data') + parser.add_argument('-o', '--output', type=str, required=True, help="Path to save visualized results") + args = parser.parse_args() + main(args) diff --git a/tools/inference/trt_inf.py b/tools/inference/trt_inf.py new file mode 100644 index 0000000000000000000000000000000000000000..e98e0560b14f885df3ffd9ce86720227332fdb8c --- /dev/null +++ b/tools/inference/trt_inf.py @@ -0,0 +1,242 @@ +""" +DEIMv2: Real-Time Object Detection Meets DINOv3 +Copyright (c) 2025 The DEIMv2 Authors. All Rights Reserved. +--------------------------------------------------------------------------------- +Modified from D-FINE (https://github.com/Peterande/D-FINE) +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +""" + +import collections +import contextlib +import os +import time +from collections import OrderedDict + +import cv2 # Added for video processing +import numpy as np +import tensorrt as trt +import torch +import torchvision.transforms as T +from PIL import Image, ImageDraw + + +class TimeProfiler(contextlib.ContextDecorator): + def __init__(self): + self.total = 0 + + def __enter__(self): + self.start = self.time() + return self + + def __exit__(self, type, value, traceback): + self.total += self.time() - self.start + + def reset(self): + self.total = 0 + + def time(self): + if torch.cuda.is_available(): + torch.cuda.synchronize() + return time.time() + +class TRTInference(object): + def __init__(self, engine_path, device='cuda:0', backend='torch', max_batch_size=32, verbose=False): + self.engine_path = engine_path + self.device = device + self.backend = backend + self.max_batch_size = max_batch_size + + self.logger = trt.Logger(trt.Logger.VERBOSE) if verbose else trt.Logger(trt.Logger.INFO) + + self.engine = self.load_engine(engine_path) + self.context = self.engine.create_execution_context() + self.bindings = self.get_bindings(self.engine, self.context, self.max_batch_size, self.device) + self.bindings_addr = OrderedDict((n, v.ptr) for n, v in self.bindings.items()) + self.input_names = self.get_input_names() + self.output_names = self.get_output_names() + self.time_profile = TimeProfiler() + + def load_engine(self, path): + trt.init_libnvinfer_plugins(self.logger, '') + with open(path, 'rb') as f, trt.Runtime(self.logger) as runtime: + return runtime.deserialize_cuda_engine(f.read()) + + def get_input_names(self): + names = [] + for _, name in enumerate(self.engine): + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + names.append(name) + return names + + def get_output_names(self): + names = [] + for _, name in enumerate(self.engine): + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT: + names.append(name) + return names + + def get_bindings(self, engine, context, max_batch_size=32, device=None) -> OrderedDict: + Binding = collections.namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr')) + bindings = OrderedDict() + + for i, name in enumerate(engine): + shape = engine.get_tensor_shape(name) + dtype = trt.nptype(engine.get_tensor_dtype(name)) + + if shape[0] == -1: + shape[0] = max_batch_size + if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + context.set_input_shape(name, shape) + + data = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device) + bindings[name] = Binding(name, dtype, shape, data, data.data_ptr()) + + return bindings + + def run_torch(self, blob): + for n in self.input_names: + if self.bindings[n].shape != blob[n].shape: + self.context.set_input_shape(n, blob[n].shape) + self.bindings[n] = self.bindings[n]._replace(shape=blob[n].shape) + + assert self.bindings[n].data.dtype == blob[n].dtype, '{} dtype mismatch'.format(n) + + self.bindings_addr.update({n: blob[n].data_ptr() for n in self.input_names}) + self.context.execute_v2(list(self.bindings_addr.values())) + outputs = {n: self.bindings[n].data for n in self.output_names} + + return outputs + + def __call__(self, blob): + if self.backend == 'torch': + return self.run_torch(blob) + else: + raise NotImplementedError("Only 'torch' backend is implemented.") + + def synchronize(self): + if self.backend == 'torch' and torch.cuda.is_available(): + torch.cuda.synchronize() + +def draw(images, labels, boxes, scores, thrh=0.4): + for i, im in enumerate(images): + draw = ImageDraw.Draw(im) + scr = scores[i] + lab = labels[i][scr > thrh] + box = boxes[i][scr > thrh] + scrs = scr[scr > thrh] + + for j, b in enumerate(box): + draw.rectangle(list(b), outline='red') + draw.text( + (b[0], b[1]), + text=f"{lab[j].item()} {round(scrs[j].item(), 2)}", + fill='blue', + ) + + return images + +def process_image(m, file_path, device, size=(640, 640), model_size='s'): + im_pil = Image.open(file_path).convert('RGB') + w, h = im_pil.size + orig_size = torch.tensor([w, h])[None].to(device) + + transforms = T.Compose([ + T.Resize(size), + T.ToTensor(), + T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + if model_size not in ['atto', 'femto', 'pico', 'n'] + else T.Lambda(lambda x: x) + ]) + im_data = transforms(im_pil)[None] + + blob = { + 'images': im_data.to(device), + 'orig_target_sizes': orig_size.to(device), + } + + output = m(blob) + result_images = draw([im_pil], output['labels'], output['boxes'], output['scores']) + result_images[0].save('trt_result.jpg') + print("Image processing complete. Result saved as 'result.jpg'.") + +def process_video(m, file_path, device, size=(640, 640), model_size='s'): + cap = cv2.VideoCapture(file_path) + + # Get video properties + fps = cap.get(cv2.CAP_PROP_FPS) + orig_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + orig_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + # Define the codec and create VideoWriter object + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter('trt_result.mp4', fourcc, fps, (orig_w, orig_h)) + + transforms = T.Compose([ + T.Resize(size), + T.ToTensor(), + T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + if model_size not in ['atto', 'femto', 'pico', 'n'] + else T.Lambda(lambda x: x) + ]) + + frame_count = 0 + print("Processing video frames...") + while cap.isOpened(): + ret, frame = cap.read() + if not ret: + break + + # Convert frame to PIL image + frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + + w, h = frame_pil.size + orig_size = torch.tensor([w, h])[None].to(device) + + im_data = transforms(frame_pil)[None] + + blob = { + 'images': im_data.to(device), + 'orig_target_sizes': orig_size.to(device), + } + + output = m(blob) + + # Draw detections on the frame + result_images = draw([frame_pil], output['labels'], output['boxes'], output['scores']) + + # Convert back to OpenCV image + frame = cv2.cvtColor(np.array(result_images[0]), cv2.COLOR_RGB2BGR) + + # Write the frame + out.write(frame) + frame_count += 1 + + if frame_count % 10 == 0: + print(f"Processed {frame_count} frames...") + + cap.release() + out.release() + print("Video processing complete. Result saved as 'result_video.mp4'.") + +if __name__ == '__main__': + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('-trt', '--trt', type=str, required=True) + parser.add_argument('-i', '--input', type=str, required=True) + parser.add_argument('-d', '--device', type=str, default='cuda:0') + parser.add_argument('-s', '--size', type=int, required=True, help='input size, e.g., 640') + parser.add_argument('-ms', '--model-size', type=str, required=True, choices=['atto', 'femto', 'pico', 'n', 's', 'm', 'l', 'x']) + + + args = parser.parse_args() + + m = TRTInference(args.trt, device=args.device) + size = (args.size,) * 2 + + file_path = args.input + if os.path.splitext(file_path)[-1].lower() in ['.jpg', '.jpeg', '.png', '.bmp']: + # Process as image + process_image(m, file_path, args.device, size, args.model_size) + else: + # Process as video + process_video(m, file_path, args.device, size, args.model_size) diff --git a/tools/reference/convert_weight.py b/tools/reference/convert_weight.py new file mode 100644 index 0000000000000000000000000000000000000000..9651d19a98b181658400137a74bcaf39be088567 --- /dev/null +++ b/tools/reference/convert_weight.py @@ -0,0 +1,29 @@ +import torch +import os +import argparse + +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_dir', type=str, help="Path to the input checkpoint file.") + + args = parser.parse_args() + for file in os.listdir(args.checkpoint_dir): + if '.pth' in file and '_converted' not in file: + save_only_ema_weights(os.path.join(args.checkpoint_dir, file)) diff --git a/tools/reference/safe_training.sh b/tools/reference/safe_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..d3c752a48f27511a353d65dbb9e8f97146ad4817 --- /dev/null +++ b/tools/reference/safe_training.sh @@ -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 diff --git a/tools/visualization/fiftyone_vis.py b/tools/visualization/fiftyone_vis.py new file mode 100644 index 0000000000000000000000000000000000000000..5831293b16c8c77209e97411bee00695332b24e2 --- /dev/null +++ b/tools/visualization/fiftyone_vis.py @@ -0,0 +1,307 @@ +""" +Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. +""" + +import argparse +import os +import subprocess +import sys +import time + +import fiftyone as fo +import fiftyone.core.fields as fof +import fiftyone.core.labels as fol +import fiftyone.core.models as fom +import fiftyone.zoo as foz +import torch +import torchvision.transforms as transforms +import tqdm +from fiftyone import ViewField as F +from PIL import Image + +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '../..')) +from engine.core import YAMLConfig + + +def kill_existing_mongod(): + try: + result = subprocess.run(['ps', 'aux'], stdout=subprocess.PIPE) + processes = result.stdout.decode('utf-8').splitlines() + + for process in processes: + if 'mongod' in process and '--dbpath' in process: + # find mongod PID + pid = int(process.split()[1]) + print(f"Killing existing mongod process with PID: {pid}") + # kill mongod session + os.kill(pid, 9) + except Exception as e: + print(f"Error occurred while killing mongod: {e}") + +kill_existing_mongod() + + +label_map = { + 1: 'person', 2: 'bicycle', 3: 'car', 4: 'motorbike', 5: 'aeroplane', + 6: 'bus', 7: 'train', 8: 'truck', 9: 'boat', 10: 'trafficlight', + 11: 'firehydrant', 12: 'streetsign', 13: 'stopsign', 14: 'parkingmeter', + 15: 'bench', 16: 'bird', 17: 'cat', 18: 'dog', 19: 'horse', + 20: 'sheep', 21: 'cow', 22: 'elephant', 23: 'bear', 24: 'zebra', + 25: 'giraffe', 26: 'hat', 27: 'backpack', 28: 'umbrella', 29: 'shoe', + 30: 'eyeglasses', 31: 'handbag', 32: 'tie', 33: 'suitcase', 34: 'frisbee', + 35: 'skis', 36: 'snowboard', 37: 'sportsball', 38: 'kite', 39: 'baseballbat', + 40: 'baseballglove', 41: 'skateboard', 42: 'surfboard', 43: 'tennisracket', + 44: 'bottle', 45: 'plate', 46: 'wineglass', 47: 'cup', 48: 'fork', + 49: 'knife', 50: 'spoon', 51: 'bowl', 52: 'banana', 53: 'apple', + 54: 'sandwich', 55: 'orange', 56: 'broccoli', 57: 'carrot', 58: 'hotdog', + 59: 'pizza', 60: 'donut', 61: 'cake', 62: 'chair', 63: 'sofa', + 64: 'pottedplant', 65: 'bed', 66: 'mirror', 67: 'diningtable', 68: 'window', + 69: 'desk', 70: 'toilet', 71: 'door', 72: 'tv', 73: 'laptop', + 74: 'mouse', 75: 'remote', 76: 'keyboard', 77: 'cellphone', 78: 'microwave', + 79: 'oven', 80: 'toaster', 81: 'sink', 82: 'refrigerator', 83: 'blender', + 84: 'book', 85: 'clock', 86: 'vase', 87: 'scissors', 88: 'teddybear', + 89: 'hairdrier', 90: 'toothbrush', 91: 'hairbrush' +} + +class CustomModel(fom.Model): + def __init__(self, cfg): + super().__init__() + self.model = cfg.model.eval().cuda() + self.postprocessor = cfg.postprocessor.eval().cuda() + self.transform = transforms.Compose([ + transforms.ToTensor(), + transforms.Resize((640, 640)), # Resize to the size expected by your model + # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + ]) + + @property + def media_type(self): + return "image" + + @property + def has_logits(self): + return False + + @property + def has_embeddings(self): + return False + + @property + def ragged_batches(self): + return False + + @property + def transforms(self): + return None + + @property + def preprocess(self): + return True + + @preprocess.setter + def preprocess(self, value): + pass + + def _convert_predictions(self, predictions): + class_labels, bboxes, scores = predictions[0]['labels'], predictions[0]['boxes'], predictions[0]['scores'] + + detections = [] + for label, bbox, score in zip(class_labels, bboxes, scores): + detection = fol.Detection( + label=label_map[label.item()], + bounding_box=[ + bbox[0] / 640, # Normalized coordinates + bbox[1] / 640, + (bbox[2] - bbox[0]) / 640, + (bbox[3] - bbox[1]) / 640 + ], + confidence=score + ) + detections.append(detection) + + return fol.Detections(detections=detections) + + def predict(self, image): + image = Image.fromarray(image).convert('RGB') + image_tensor = self.transform(image).unsqueeze(0).cuda() + outputs = self.model(image_tensor) + orig_target_sizes = torch.tensor([[640, 640]]).cuda() + predictions = self.postprocessor(outputs, orig_target_sizes) + return self._convert_predictions(predictions) + + def predict_all(self, images): + image_tensors = [] + for image in images: + image = Image.fromarray(image) + image_tensor = self.transform(image) + image_tensors.append(image_tensor) + image_tensors = torch.stack(image_tensors).cuda() + outputs = self.model(image_tensors) + orig_target_sizes = torch.tensor([[640, 640] for image in images]).cuda() + predictions = self.postprocessor(outputs, orig_target_sizes) + converted_predictions = [self._convert_predictions(pred) for pred in predictions] + + # Ensure the output is a list of lists of Detections + return converted_predictions + +def filter_by_predictions5_confidence(predictions_view, confidence_threshold=0.3): + for j, sample in tqdm.tqdm(enumerate(predictions_view), total=len(predictions_view)): + has_modified = False + for i, detection in enumerate(sample["predictions0"].detections): + + if "original_confidence" not in detection: + detection["original_confidence"] = detection["confidence"] + + if (detection["confidence"] <= confidence_threshold and sample["predictions5"].detections[i]["confidence"] >= confidence_threshold) or \ + (detection["confidence"] >= confidence_threshold and sample["predictions5"].detections[i]["confidence"] <= confidence_threshold): + + sample["predictions0"].detections[i]["confidence"] = sample["predictions5"].detections[i]["confidence"] + has_modified = True + if has_modified: + sample.save() + + +def restore_confidence(predictions_view): + for j, sample in tqdm.tqdm(enumerate(predictions_view), total=len(predictions_view)): + for i, detection in enumerate(sample["predictions0"].detections): + if "original_confidence" in detection: + detection["confidence"] = detection["original_confidence"] + sample.save() + +def fast_iou(bbox1, bbox2): + x1, y1, w1, h1 = bbox1 + x2, y2, w2, h2 = bbox2 + xA = max(x1, x2) + yA = max(y1, y2) + xB = min(x1 + w1, x2 + w2) + yB = min(y1 + h1, y2 + h2) + interArea = max(0, xB - xA) * max(0, yB - yA) + boxAArea = w1 * h1 + boxBArea = w2 * h2 + iou = interArea / float(boxAArea + boxBArea - interArea) + return iou + +def assign_iou_diff(predictions_view): + for sample in predictions_view: + ious_0 = [detection.eval0_iou if 'eval0_iou' in detection else None for detection in sample["predictions0"].detections] + ious_5 = [detection.eval5_iou if 'eval5_iou' in detection else None for detection in sample["predictions5"].detections] + bbox_0 = [detection.bounding_box for detection in sample["predictions0"].detections] + bbox_5 = [detection.bounding_box for detection in sample["predictions5"].detections] + # iou_diffs = [abs(iou_5 - iou_0) if iou_0 is not None and iou_5 is not None else -1 for iou_0, iou_5 in zip(ious_0, ious_5)] + iou_inter = [fast_iou(b0, b5) for b0, b5 in zip(bbox_0, bbox_5)] + iou_diffs = [abs(iou_5 - iou_0) if iou_0 is not None and iou_5 is not None and iou_inter > 0.5 else -1 for iou_0, iou_5, iou_inter in zip(ious_0, ious_5, iou_inter)] + + for detection, iou_diff in zip(sample["predictions0"].detections, iou_diffs): + detection["iou_diff"] = iou_diff + for detection, iou_diff in zip(sample["predictions5"].detections, iou_diffs): + detection["iou_diff"] = iou_diff + # for detection, iou_diff in zip(sample["predictions100"].detections, iou_diffs): + # detection["iou_diff"] = iou_diff + sample.save() + +def main(args): + try: + if os.path.exists("saved_predictions_view") and os.path.exists("saved_filtered_view"): + print("Loading saved predictions and filtered views...") + dataset = foz.load_zoo_dataset( + "coco-2017", + split="validation", + dataset_name="evaluate-detections-tutorial", + dataset_dir="data/fiftyone" + ) + + dataset.persistent = True + session = fo.launch_app(dataset, port=args.port) + + predictions_view = fo.Dataset.from_dir( + dataset_dir="saved_predictions_view", + dataset_type=fo.types.FiftyOneDataset + ).view() + filtered_view = fo.Dataset.from_dir( + dataset_dir="saved_filtered_view", + dataset_type=fo.types.FiftyOneDataset + ).view() + else: + dataset = foz.load_zoo_dataset( + "coco-2017", + split="validation", + dataset_name="evaluate-detections-tutorial", + dataset_dir="data/fiftyone" + ) + + dataset.persistent = True + + session = fo.launch_app(dataset, port=args.port) + cfg = YAMLConfig(args.config, resume=args.resume) + if 'HGNetv2' in cfg.yaml_cfg: + cfg.yaml_cfg['HGNetv2']['pretrained'] = False + if args.resume: + checkpoint = torch.load(args.resume, map_location='cpu') + if 'ema' in checkpoint: + state = checkpoint['ema']['module'] + else: + state = checkpoint['model'] + else: + raise AttributeError('only support resume to load model.state_dict by now.') + + # NOTE load train mode state -> convert to deploy mode + cfg.model.load_state_dict(state) + predictions_view = dataset.take(500, seed=51) + + model = CustomModel(cfg) + L = model.model.decoder.decoder.eval_idx + # Apply models and save predictions in different label fields + for i in [L]: + model.model.decoder.decoder.eval_idx = i + label_field = "predictions{:d}".format(i) + predictions_view.apply_model(model, label_field=label_field) + + # filter_by_predictions5_confidence(predictions_view, confidence_threshold=0.3) + for i in [L]: + label_field = "predictions{:d}".format(i) + predictions_view = predictions_view.filter_labels(label_field, F("confidence") > 0.5, only_matches=False) + eval_key = "eval{:d}".format(i) + _ = predictions_view.evaluate_detections( + label_field, + gt_field="ground_truth", + eval_key=eval_key, + compute_mAP=True, + ) + + # assign_iou_diff(predictions_view) + + # filtered_view = predictions_view.filter_labels("predictions0", F("iou_diff") > 0.05, only_matches=True) + # filtered_view = filtered_view.filter_labels("predictions5", F("iou_diff") > 0.05, only_matches=True) + # restore_confidence(filtered_view) + + predictions_view.export( + export_dir="saved_predictions_view", + dataset_type=fo.types.FiftyOneDataset + ) + # filtered_view.export( + # export_dir="saved_filtered_view", + # dataset_type=fo.types.FiftyOneDataset + # ) + + # Display the filtered view + session.view = predictions_view + + # Keep the session open + while True: + time.sleep(1) + except Exception as e: + print(f"An error occurred: {e}") + finally: + print("Shutting down session") + if 'session' in locals(): + session.close() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--config', '-c', type=str) + parser.add_argument('--resume', '-r', type=str) + parser.add_argument('--port', '-p', type=int) + args = parser.parse_args() + + main(args)