{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":11244260,"sourceType":"datasetVersion","datasetId":6993770}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#ライブラリのimport\nimport numpy as np\nimport pandas as pd\nimport time\nimport argparse\nfrom pathlib import Path\n\nfrom tqdm.auto import tqdm\nfrom torch.utils.data import DataLoader, default_collate\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nimport torchvision.transforms.functional as TTF\nimport timm\nimport yaml\n\nimport multiprocessing as mp\nfrom queue import Empty","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-06T15:33:44.160406Z","iopub.execute_input":"2025-04-06T15:33:44.160847Z","iopub.status.idle":"2025-04-06T15:33:56.638804Z","shell.execute_reply.started":"2025-04-06T15:33:44.160804Z","shell.execute_reply":"2025-04-06T15:33:56.637824Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"INPUT_PATH = Path('/kaggle/input/byu-locating-bacterial-flagellar-motors-2025')\nMODEL_PATH = Path('/kaggle/input/bacterial-public/weights/object/baseline')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-06T15:33:56.639852Z","iopub.execute_input":"2025-04-06T15:33:56.640141Z","iopub.status.idle":"2025-04-06T15:33:56.643757Z","shell.execute_reply.started":"2025-04-06T15:33:56.640118Z","shell.execute_reply":"2025-04-06T15:33:56.642872Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#指定したフォルダ内のtomogramのpathを取得し、tomo_pathsに格納する\ndef get_tomos_path(input_path:Path, data_type:str,*,n=None) -> list[Path]:\n\n    #input_path='input/', data_type='train'ならばdata_path = input/train/\n    data_path = input_path / data_type\n\n    #data_pathにあるフォルダまたはファイルを全件取得し、ソートしてリストに格納\n    tomo_paths = sorted(data_path.glob('*'))\n\n    #nが指定されていて、data_typeがtrainの場合、先頭のn件のtomoを使用\n    if(n is not None) and (data_type == 'train'):\n        tomo_path = tomo_paths[:n]\n\n    return tomo_paths","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-06T15:33:56.645198Z","iopub.execute_input":"2025-04-06T15:33:56.645479Z","iopub.status.idle":"2025-04-06T15:33:56.700122Z","shell.execute_reply.started":"2025-04-06T15:33:56.645454Z","shell.execute_reply":"2025-04-06T15:33:56.699410Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#引数の画像を640×640にリサイズし、パーセンタイルに基づいて正規化する\ndef preprocess(img : torch.Tensor) -> torch.Tensor:\n\n    #サイズの統一\n    size = (640, 640)\n\n    #型変換(PyTorchモデルではfloat32が基本)\n    img = img.to(dtype=torch.float32)\n\n    #一括でリサイズ\n    img = TTF.resize(img, size)\n\n    #quantile基準点の定義\n    #テンソルを画像と同じでデバイスに載せる\n    q = torch.Tensor([0.05, 0.95]).to(img.device)\n\n    #shape(バッチサイズ、チャンネル数、高さ、幅)\n    batch_size, nch, h, w = img.shape\n\n    #viewで一枚の画像を一次元ベクトルにする\n    #各画像の5%・95%の値を求める\n    x_min, x_max = torch.quantile(img.view(batch_size, nch*h*w), q, dim=1)\n\n    x_min = x_min.view(batch_size, 1, 1, 1)\n    x_max = x_max.view(batch_size, 1, 1, 1)\n\n    #スケーリング\n    img = (img - x_min) / (x_max - x_min)\n\n    #クリッピング\n    img = torch.clamp(img, 0, 1)\n\n    return img","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-06T15:33:56.701259Z","iopub.execute_input":"2025-04-06T15:33:56.701575Z","iopub.status.idle":"2025-04-06T15:33:56.719079Z","shell.execute_reply.started":"2025-04-06T15:33:56.701546Z","shell.execute_reply":"2025-04-06T15:33:56.718518Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#tomo_pathに含まれる画像を一枚ずつ読み込んで、DataLoaderで使える形式にする\nclass Dataset(torch.utils.data.Dataset):\n\n    #画像ファイル一覧を取得\n    def __init__(self, tomo_path: Path):\n        self.filenames = sorted(tomo_path.glob('*'))\n        \n    #このデータセットに含まれる「データ数（スライス画像の数）」\n    def __len__(self) -> int:\n        return len(self.filenames)\n\n    #インデックス番号の画像を読込み、辞書形式で返す\n    def __getitem__(self, i: int) -> dict:\n        filename = self.filenames[i]      # Path\n        filebase = filename.stem          #ファイル名取得\n        assert filebase[:6] == 'slice_'   #ファイル名チェック\n        slice_number = int(filebase[6:])  # slice_0000 -> int(0000)\n\n        #画像を開いて numpy 配列にし、チャンネル軸（C=1）を追加（PyTorchの形式に合わせる）\n        img = Image.open(filename)\n        W, H = img.size\n        img = np.expand_dims(np.array(img), axis=0)\n\n        #辞書で返す値を設定\n        ret = {'img': img,\n               'slice_number': slice_number,\n               'shape': np.array((H, W), dtype=int),\n        }\n\n        return ret\n\n    def loader(self, batch_size: int, num_workers: int):\n        loader = DataLoader(self, batch_size=batch_size, num_workers=num_workers)\n        return loader","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-06T15:33:56.719755Z","iopub.execute_input":"2025-04-06T15:33:56.719945Z","iopub.status.idle":"2025-04-06T15:33:56.746249Z","shell.execute_reply.started":"2025-04-06T15:33:56.719928Z","shell.execute_reply":"2025-04-06T15:33:56.745387Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Model(nn.Module):\n    #構造を作る\n    def __init__(self, cfg_model: dict, *, pretrained=True, verbose=True):\n        super().__init__()\n\n        # Timm encoder\n        name = cfg_model['encoder']   #使うモデルの種類を指定\n        in_channels = 1               #チャンネル数(白黒)\n        out_channels = 1              #マスク画像のチャンネル数\n\n        self.encoder = timm.create_model(name,                    #モデル名\n                                         in_chans=in_channels,    #入力チャンネル数\n                                         features_only=True,      #中間の特徴マップを取り出す\n                                         pretrained=pretrained)   #事前学習済みの重みを使用する?\n\n\n        #encoderから出てくる各ステージの特徴マップのチャンネル数を取得する\n        encoder_channels = self.encoder.feature_info.channels()\n\n        self.segmentation_head = nn.Conv2d(encoder_channels[-1],  #入力チャンネル数\n                                           out_channels,          #出力チャンネル数\n                                           kernel_size=3,         #カーネルサイズ\n                                           padding=1)             #出力サイズを維持する\n\n        self.regression_head = nn.Conv2d(encoder_channels[-1], out_channels=2,\n                                         kernel_size=3, padding=1)\n\n        self.criterion_seg = nn.BCEWithLogitsLoss()\n        self.criterion_reg = nn.MSELoss()\n\n        if verbose:\n            print(name)\n\n    #データの流れを定義する\n    def forward(self, img: torch.Tensor):\n\n        features = self.encoder(img)\n        out = features[-1]  # (batch_size, embed_dim, h, w)\n        y_pred = self.segmentation_head(out)  # (batch_size, 1, h, w)\n        t_pred = self.regression_head(out)    # (batch_size, 2, h, w)\n\n        return y_pred, t_pred","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-06T15:33:56.747182Z","iopub.execute_input":"2025-04-06T15:33:56.747429Z","iopub.status.idle":"2025-04-06T15:33:56.767502Z","shell.execute_reply.started":"2025-04-06T15:33:56.747409Z","shell.execute_reply":"2025-04-06T15:33:56.766781Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#モーターが一番ありそうなスライスと位置を特定\ndef predict(tomo_path: Path,models: list[nn.Module],cfg: dict) -> dict:\n    \n    #各種設定\n    assert len(models) > 0\n    tomo_id = tomo_path.name                       #ファイルの名前\n    batch_size = cfg['batch_size']                 #バッチサイズ\n    num_workers = cfg['num_workers']               #データ読込みの並列数\n    use_amp = cfg['use_amp']                       #AMPを使うか？\n    preprocess_device = cfg['preprocess_device']   #デバイス\n    assert preprocess_device == 'cpu' or preprocess_device.startswith('cuda')\n\n    #DataLoaderを生成(スライス画像をバッチで取り出す準備)\n    dataset = Dataset(tomo_path)\n    loader = dataset.loader(batch_size=batch_size, num_workers=num_workers)\n\n    #モデルの載っているデバイスを確認(GPUを優先する)\n    device = next(models[0].parameters()).device\n\n    #予測値の格納場所\n    best = (0, None)\n\n    #バッチごとの処理\n    for d in loader:        \n        if preprocess_device.startswith('cuda'):\n            img = d['img'].to(device)  # 画像をまずGPUへ\n            img = preprocess(img)      # GPU上で処理\n        elif preprocess_device == 'cpu':\n            img = preprocess(d['img']) # CPU上で処理\n            img = img.to(device)       # 処理後にGPUへ\n        else:\n            raise ValueError(f\"Unknown preprocess_device: {preprocess_device}\")\n\n\n        #アンサンブル平均をとる\n        y_pred_sum, t_pred_sum = None, None\n        for model in models:\n            with torch.no_grad():                            #勾配計算なし\n                with torch.amp.autocast(device_type='cuda',  #AMP\n                                        enabled=use_amp,\n                                        dtype=torch.float16):\n                    y_pred, t_pred = model(img) \n\n            y_pred = y_pred.sigmoid()  \n\n            if y_pred_sum is None:\n                y_pred_sum = y_pred\n                t_pred_sum = t_pred\n            else:\n                y_pred_sum += y_pred\n                t_pred_sum += t_pred\n\n        #バッチの中で一番モーターっぽいやつ\n        y_pred_max = y_pred_sum.max().item() / len(models)\n        del y_pred, t_pred\n\n        # マスクの中から一番スコアが高い場所を探し、(スライス番号+Y座標+X座標)を返す\n        if y_pred_max > best[0]:             #best=(スコア, スライス番号, y, x)\n            bs, _, h, w = y_pred_sum.shape   #バッチサイズ, チャンネル数, 特徴マップの高さと幅\n        \n            argmax = torch.unravel_index(y_pred_sum.argmax(), y_pred_sum.shape)  # b, ch, iy, ix\n            i, _, iy, ix = [t.item() for t in argmax]    \n            slice_number = d['slice_number'][i].item()   #スライス番号を取得\n            offset = t_pred_sum[i, :, iy, ix].cpu().numpy() / len(models)  \n\n            #元画像のピクセル単位に変換\n            H, W = d['shape'][i].numpy()    \n            x = (ix + offset[0]) * (W / w)\n            y = (iy + offset[1]) * (H / h)\n\n            #結果の更新\n            best = (y_pred_max, slice_number, y, x)\n\n    assert best[1] is not None\n\n    # 予測結果を辞書にまとめて返す\n    n_slices = len(dataset.filenames)\n    pred = {'tomo_id': tomo_id,\n            'n_slices': n_slices,\n            'y_pred': best[0],\n            'zyx': best[1:]}\n    return pred","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-06T15:33:56.768231Z","iopub.execute_input":"2025-04-06T15:33:56.768460Z","iopub.status.idle":"2025-04-06T15:33:56.788096Z","shell.execute_reply.started":"2025-04-06T15:33:56.768440Z","shell.execute_reply":"2025-04-06T15:33:56.787194Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#スコアを判定\ndef create_submission(preds: list, th: float, ofilename: str) -> pd.DataFrame:\n\n    rows = []             #結果をためる\n    count_positive = 0    #モーターが見つかった件数\n    for pred in preds:\n        if pred['y_pred'] < th:\n            zyx = (-1, -1, -1)\n        else:\n            count_positive += 1\n            zyx = pred['zyx']\n\n        row = {'tomo_id': pred['tomo_id'],\n               'Motor axis 0': zyx[0],\n               'Motor axis 1': zyx[1],\n               'Motor axis 2': zyx[2]}\n        rows.append(row)\n\n    submit = pd.DataFrame(rows)\n    submit.to_csv(ofilename, float_format='%.8e', index=False)\n\n    print('Submit %s: %d positives / %d tomo_ids' % (ofilename, count_positive, len(rows)))\n\n    return submit","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-06T15:33:56.790415Z","iopub.execute_input":"2025-04-06T15:33:56.790728Z","iopub.status.idle":"2025-04-06T15:33:56.821027Z","shell.execute_reply.started":"2025-04-06T15:33:56.790704Z","shell.execute_reply":"2025-04-06T15:33:56.819625Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def process_fn(process_id: int,  #GPUの番号\n               tomo_queue,       #推論すべきtomoのリストが入ったキュー\n               pred_queue,       #予測結果を格納するキュー\n               cfg: dict):       #各種設定(モデル・バッチサイズ・AMPなどの設定が入った辞書)\n\n\n\n\n    cfg = cfg.copy()\n    cfg['preprocess_device'] = 'cuda' \n\n    device = torch.device('cuda:%d' % process_id)\n\n    #モデルの読込み\n    model_path = cfg['model_path']\n    folds = cfg['folds']\n\n\n    with open(model_path / 'config.yml', 'r') as f:\n        cfg_model = yaml.safe_load(f)\n            \n    models = []\n    for ifold in folds:\n        model_filename = '%s/model%d.pytorch' % (model_path, ifold)            \n        model = Model(cfg_model['model'], pretrained=False, verbose=False)     #モデル構造をインスタンス化\n        model.load_state_dict(torch.load(model_filename, weights_only=True))   #重みを読込み\n        model.to(device)                                                       #GPUに載せる\n        model.eval()                                                           #推論モードに切り替え\n        models.append(model)\n\n        if process_id == 0:\n            print('Load model', model_filename)\n        \n    #キューから一つずつとって推論\n    while not tomo_queue.empty():\n        try:\n            tomo_path = tomo_queue.get(timeout=1)\n            pred = predict(tomo_path, models, cfg)\n            pred_queue.put(pred)\n\n        except Empty:\n            break\n\ntb = time.time()\n\n\ncfg = {\n    'model_path': MODEL_PATH,    #モデルの保存先\n    'folds': [0,1,2,3,4],        #使用するモデルの番号リスト\n    'batch_size': 16,            #一度に処理するスライス数\n    'num_workers': 1,            #DataLoaderの並列データ読込みスレッド数\n    'use_amp': True,             #AMPをON\n    'preprocess_device': 'cuda', #GPUを使う\n}\n\ntomo_paths = get_tomos_path(INPUT_PATH, 'test')\n\n#テストデータの場合trainからデータを足す\nif len(tomo_paths) == 3:\n     #Some experiment when test is dummy (optional)\n     #tomo_paths = get_tomos(INPUT_PATH, 'train', n=20, random_sample=True)\n    pass\n\nprint('Data %d' % len(tomo_paths))\n\nmanager = mp.Manager()\n#予測したいtomoを詰めるキュー\ntomo_queue = manager.Queue()\n#各GPUが出した予測結果を詰めるキュー\npred_queue = manager.Queue()\n#処理対象のtomoを予測したいtomoキューに格納\nfor tomo_path in tomo_paths:\n    tomo_queue.put(tomo_path)\n\ntime.sleep(1)\nassert not tomo_queue.empty()\n\n#GPUの数だけprcess_fnを同時実行するために、プロセスを作ってリストにまとめる\nnum_processes = 2\ntb = time.time()\n\nworkers = [mp.Process(target=process_fn,\n                      args=(i, tomo_queue, pred_queue, cfg))\n           for i in range(num_processes)]\n\nfor w in workers:\n    w.start()\n\nfor w in workers:\n    w.join()\n\n\ndt = time.time() - tb\nprint('%.2f sec for %d tomos' % (dt, len(tomo_paths)))\n\n#処理を集めてキューに格納\npreds = []\ntry:\n    while not pred_queue.empty():\n        preds.append(pred_queue.get(timeout=1))\nexcept Empty:\n    pass\n\nassert len(preds) == len(tomo_paths)\n\n\nth = 0.5        #閾値\nofilename = 'submission.csv'\ncreate_submission(preds, th, 'submission.csv')\nprint(ofilename, 'written')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-06T15:33:56.822593Z","iopub.execute_input":"2025-04-06T15:33:56.822999Z","iopub.status.idle":"2025-04-06T15:34:44.741941Z","shell.execute_reply.started":"2025-04-06T15:33:56.822957Z","shell.execute_reply":"2025-04-06T15:34:44.740788Z"}},"outputs":[],"execution_count":null}]}