src/codebase/notebooks/Mammo-Factor/Mammo-Factor.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f30468f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import argparse\n",
    "import os\n",
    "import pickle\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d09e0e23",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "\n",
    "from breastclip.data import DataModule\n",
    "from breastclip.model import build_model\n",
    "\n",
    "ATTRIBUTES = [\n",
    "    \"mass\",\n",
    "    \"suspicious_calcification\",\n",
    "]\n",
    "\n",
    "\n",
    "def generate_attribute_embs(out_dir, breast_clip_path, model):\n",
    "    \"\"\"\n",
    "    Generate embeddings for each attribute.\n",
    "\n",
    "    Parameters:\n",
    "        out_dir: Directory for storing attribute embeddings\n",
    "    \"\"\"\n",
    "\n",
    "    def get_prompts(attr):\n",
    "        if attr in [\"mass\"]:\n",
    "            prompts = [\n",
    "                # Replace with either the corresponding prompts or the sentences from the report.\n",
    "            ]\n",
    "\n",
    "        elif attr in [\"suspicious_calcification\"]:\n",
    "            prompts = [\n",
    "                # Replace with either the corresponding prompts or the sentences from the report.\n",
    "            ]\n",
    "    \n",
    "        return prompts\n",
    "\n",
    "    # model, preprocess = clip.load(\"RN50\", \"cuda\")\n",
    "\n",
    "    # Obtain prompts for each attribute\n",
    "    prompt_list = []\n",
    "    for attr in ATTRIBUTES:\n",
    "        print(attr)\n",
    "        prompt_list.append(get_prompts(attr))\n",
    "\n",
    "    print(len(prompt_list))\n",
    "    device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "    breast_clip_path = Path(breast_clip_path)\n",
    "    print(breast_clip_path)\n",
    "    configs = torch.load(breast_clip_path, map_location=device)\n",
    "    cfg = configs[\"config\"]\n",
    "\n",
    "    datamodule = DataModule(\n",
    "        data_config=cfg[\"data_train\"],\n",
    "        dataloader_config=cfg[\"dataloader\"],\n",
    "        tokenizer_config=cfg[\"tokenizer\"] if \"tokenizer\" in cfg else None,\n",
    "        loss_config=cfg[\"loss\"],\n",
    "        transform_config=cfg[\"transform\"],\n",
    "        mean=cfg[\"base\"][\"mean\"],\n",
    "        std=cfg[\"base\"][\"std\"],\n",
    "        image_encoder_type=cfg[\"model\"][\"image_encoder\"][\"model_type\"],\n",
    "        cur_fold=cfg[\"base\"][\"fold\"]\n",
    "    )\n",
    "\n",
    "    clip = build_model(cfg[\"model\"], cfg[\"loss\"], datamodule.tokenizer)\n",
    "    clip = clip.to(device)\n",
    "    clip.load_state_dict(configs[\"model\"], strict=True)\n",
    "    clip.eval()\n",
    "    print(clip)\n",
    "    # Compute each attribute embedding as the average of its associated prompt embeddings\n",
    "    attr_embs = []\n",
    "    with torch.no_grad():\n",
    "        for prompt in prompt_list:\n",
    "            text_token = datamodule.tokenizer(\n",
    "                prompt, padding=\"longest\", truncation=True, return_tensors=\"pt\", max_length=256\n",
    "            )\n",
    "            text_emb = clip.encode_text(text_token.to(device))\n",
    "            text_emb = clip.text_projection(text_emb) if clip.projection else text_emb\n",
    "            text_emb = text_emb.mean(dim=0, keepdim=True)\n",
    "            text_emb /= text_emb.norm(dim=-1, keepdim=True)\n",
    "            attr_embs.append(text_emb)\n",
    "\n",
    "        attr_embs = torch.stack(attr_embs).squeeze().detach().cpu().numpy()\n",
    "\n",
    "    attr_to_emb = dict(zip(ATTRIBUTES, attr_embs))\n",
    "    print(attr_embs.shape)\n",
    "    out_dir = Path(out_dir)\n",
    "    os.makedirs(out_dir, exist_ok=True)\n",
    "    torch.save(attr_to_emb, f\"{out_dir}/attr_embs_{model}.pth\")\n",
    "    print(f\"Saved {len(attr_to_emb)} attribute embeddings to {out_dir}/attr_embs_{model}.pth\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36383743",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from breastclip.model.modules import load_image_encoder\n",
    "\n",
    "\n",
    "class Mapper_model(torch.nn.Module):\n",
    "    def __init__(self, ckpt, lang_emb: int, emb_dim: int, one_proj: bool, adapter: bool, attr_embs):\n",
    "        super(Mapper_model, self).__init__()\n",
    "        self.image_encoder = load_image_encoder(ckpt[\"config\"][\"model\"][\"image_encoder\"])\n",
    "        image_encoder_weights = {}\n",
    "        for k in ckpt[\"model\"].keys():\n",
    "            if k.startswith(\"image_encoder.\"):\n",
    "                image_encoder_weights[\".\".join(k.split(\".\")[1:])] = ckpt[\"model\"][k]\n",
    "        self.image_encoder.load_state_dict(image_encoder_weights, strict=True)\n",
    "        self.image_encoder_type = ckpt[\"config\"][\"model\"][\"image_encoder\"][\"model_type\"]\n",
    "        for param in self.image_encoder.parameters():\n",
    "            param.requires_grad = False\n",
    "\n",
    "        self.emb_dim = emb_dim\n",
    "        self.lang_emb = lang_emb\n",
    "        self.one_proj = one_proj\n",
    "        self.adapter = adapter\n",
    "\n",
    "        # Initialize projection heads\n",
    "        if self.one_proj:\n",
    "            self.num_proj = 1\n",
    "        else:\n",
    "            self.num_proj = len(attr_embs)\n",
    "        self.pool = torch.nn.ModuleList(\n",
    "            [\n",
    "                torch.nn.Sequential(\n",
    "                    torch.nn.Linear(self.emb_dim, self.emb_dim),\n",
    "                    torch.nn.ReLU(),\n",
    "                    torch.nn.Linear(self.emb_dim, self.lang_emb),\n",
    "                )\n",
    "                for i in range(self.num_proj)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def encode_image(self, input):\n",
    "        image_features, raw_features = self.image_encoder(input)\n",
    "        return image_features, raw_features\n",
    "\n",
    "    def forward(self, sample: dict):\n",
    "        out_dict = {}\n",
    "\n",
    "        img_vector = sample[\"img\"].to(torch.float32).to(\"cuda\")\n",
    "        if len(img_vector.size()) == 5:\n",
    "            img_vector = img_vector.squeeze(1).permute(0, 3, 1, 2)\n",
    "        input = {\"image\": img_vector}\n",
    "\n",
    "        image_features, raw_features = self.encode_image(input)\n",
    "        bs = raw_features.size(0)\n",
    "        channel_dim = raw_features.size(1)\n",
    "        raw_features_flatten = raw_features.view(bs, channel_dim, -1)\n",
    "\n",
    "        out_img_a = []\n",
    "        for i in range(self.num_proj):\n",
    "            pool = self.pool[i](raw_features_flatten)\n",
    "            if self.adapter:\n",
    "                pool = 0.2 * pool + 0.8 * raw_features_flatten\n",
    "            out_img_a.append(pool)\n",
    "\n",
    "        region_proj_embs = torch.cat(out_img_a, dim=1).view(\n",
    "            -1, self.num_proj, self.lang_emb\n",
    "        )\n",
    "        out_dict[\"region_proj_embs\"] = region_proj_embs\n",
    "        out_dict[\"num_regions\"] = torch.tensor(channel_dim)\n",
    "        out_dict[\"image_features\"] = image_features\n",
    "        out_dict[\"raw_features\"] = raw_features\n",
    "        return out_dict\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7081a94e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import torch\n",
    "import abc\n",
    "\n",
    "\n",
    "class BaseLoss(torch.nn.Module):\n",
    "    __metaclass__ = abc.ABC\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.iteration = 0\n",
    "        self.running_loss = 0\n",
    "        self.mean_running_loss = 0\n",
    "\n",
    "    def forward(self, input):\n",
    "        return input\n",
    "\n",
    "    def update_running_loss(self, loss):\n",
    "        self.iteration += 1\n",
    "        self.running_loss += loss.item()\n",
    "        self.mean_running_loss = self.running_loss / self.iteration\n",
    "\n",
    "\n",
    "class Mapper_loss(BaseLoss):\n",
    "    def __init__(self, temp: float, one_proj: bool, attr_to_embs):\n",
    "        super().__init__()\n",
    "\n",
    "        # Initialize class variables\n",
    "        self.temperature = temp\n",
    "        self.one_proj = one_proj\n",
    "\n",
    "        # Load attribute embeddings\n",
    "        self.attr_embs = []\n",
    "        for a in attr_to_embs:\n",
    "            self.attr_embs.append(attr_to_embs[a])\n",
    "        self.attr_embs = torch.tensor(np.stack(self.attr_embs)).cuda().to(torch.float32)\n",
    "\n",
    "    def forward(self, pred: dict, sample: dict):\n",
    "        anchor_img = torch.nn.functional.normalize(\n",
    "            pred[\"region_proj_embs\"].float(), dim=2\n",
    "        )\n",
    "        labels = sample[\"labels\"].to(int)\n",
    "\n",
    "        # Get the indexes of the concepts which is at least present, for ex. mass and calcificaation\n",
    "        attr_ids = labels.sum(0).nonzero().flatten().tolist()\n",
    "        batch_size = labels.shape[0]\n",
    "\n",
    "        loss = torch.tensor(0.0).cuda()\n",
    "        num_loss_terms = 0\n",
    "\n",
    "        # Compute region-attribute similarity matrix by calculating the similarity between\n",
    "        # each attribute and the region embedding resulting from the corresponding projection head\n",
    "        if self.one_proj:\n",
    "            txt_emb = self.attr_embs[attr_ids, :].T.unsqueeze(0)\n",
    "            sim = (anchor_img @ txt_emb).squeeze() / self.temperature\n",
    "        else:\n",
    "            reg_emb = anchor_img[:, attr_ids, :]\n",
    "            txt_emb = self.attr_embs[attr_ids, :].unsqueeze(0)\n",
    "            sim = (reg_emb * txt_emb).sum(2) / self.temperature\n",
    "\n",
    "        # Convert region-attribute similarity matrix into image-attribute similarity matrix by\n",
    "        # computing the maximum pairwise similarity between all regions in the image and each attribute\n",
    "        split = torch.split(sim, pred[\"num_regions\"].tolist(), dim=0)\n",
    "        vals, _ = zip(*map(torch.max, split, [0] * batch_size))\n",
    "\n",
    "        sim = torch.stack(vals)  # Size: batch_size x len(attr_ids)\n",
    "        true_label = labels[:, attr_ids].cuda()\n",
    "        inv_true_label = (~true_label.bool()).to(int)\n",
    "\n",
    "        # Compute final contrastive loss\n",
    "        denom = torch.exp(sim) + torch.exp(sim * inv_true_label).sum(0, keepdims=True)\n",
    "        loss = ((-torch.log(torch.exp(sim) / denom)) * true_label).sum(1, keepdims=True)\n",
    "        num_loss_terms = true_label.sum()\n",
    "        loss = loss.sum() / num_loss_terms\n",
    "        self.update_running_loss(loss)\n",
    "        return loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "903bf5e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataloader_Vindr(args):\n",
    "    train_dataset = MammoDataset_Mapper(args=args, df=args.train_folds, transform=get_transforms(args))\n",
    "    valid_dataset = MammoDataset_Mapper(args=args, df=args.valid_folds)\n",
    "\n",
    "    if args.balanced_dataloader == \"y\":\n",
    "        weight_path = args.output_path / f\"random_sampler_weights_fold{str(args.cur_fold)}.pkl\"\n",
    "        if weight_path.exists():\n",
    "            weights = pickle.load(open(weight_path, \"rb\"))\n",
    "        else:\n",
    "            weight_for_positive_class = args.sampler_weights[f\"fold{str(args.cur_fold)}\"][\"pos_wt\"]\n",
    "            weight_for_negative_class = args.sampler_weights[f\"fold{str(args.cur_fold)}\"][\"neg_wt\"]\n",
    "            args.train_folds[\"weights_random_sampler\"] = args.train_folds.apply(\n",
    "                lambda row: weight_for_positive_class if row[\"cancer\"] == 1 else weight_for_negative_class, axis=1\n",
    "            )\n",
    "            weights = args.train_folds[\"weights_random_sampler\"].values\n",
    "            pickle.dump(weights, open(args.output_path / f\"random_sampler_weights_fold{args.cur_fold}.pkl\", \"wb\"))\n",
    "\n",
    "        weights = weights.tolist()\n",
    "        sampler = WeightedRandomSampler(weights, len(weights), replacement=True)\n",
    "        train_loader = DataLoader(\n",
    "            train_dataset, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=True,\n",
    "            drop_last=True, sampler=sampler\n",
    "        )\n",
    "    else:\n",
    "        train_loader = DataLoader(\n",
    "            train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True,\n",
    "            drop_last=True,\n",
    "        )\n",
    "\n",
    "    valid_loader = DataLoader(\n",
    "        valid_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True,\n",
    "        drop_last=False\n",
    "    )\n",
    "\n",
    "    return train_loader, valid_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65113c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_region_mapper(args, device):\n",
    "    print(f\"=> Training is getting started\")\n",
    "    # Initialize dataloaders\n",
    "    args.data_dir = Path(args.data_dir)\n",
    "    args.df = pd.read_csv(args.data_dir / args.csv_file)\n",
    "    args.df = args.df.fillna(0)\n",
    "    args.df = args.df[(args.df[\"Mass\"] == 1) | (args.df[\"Suspicious_Calcification\"] == 1)]\n",
    "    print(f\"df shape: {args.df.shape}\")\n",
    "    print(args.df.columns)\n",
    "    args.train_folds = args.df[args.df['split'] == \"training\"].reset_index(drop=True)\n",
    "    args.valid_folds = args.df[args.df['split'] == \"test\"].reset_index(drop=True)\n",
    "    train_loader, val_loader = get_dataloader_Vindr(args)\n",
    "    print(f'train_loader: {len(train_loader)}, valid_loader: {len(val_loader)}')\n",
    "\n",
    "    print(args.clip_chk_pt_path)\n",
    "\n",
    "    ckpt = torch.load(args.clip_chk_pt_path, map_location=\"cpu\")\n",
    "\n",
    "    # Initialize model, loss, and optimizer\n",
    "    attr_embs = torch.load(args.attr_embs_path)\n",
    "    model = Mapper_model(\n",
    "        ckpt, lang_emb=args.lang_emb, emb_dim=args.img_emb, one_proj=False, adapter=False, attr_embs=attr_embs\n",
    "    ).to(\"cuda\")\n",
    "    loss = Mapper_loss(temp=0.07, one_proj=False, attr_to_embs=attr_embs)\n",
    "\n",
    "    opt = torch.optim.AdamW(model.parameters(), lr=args.lr)\n",
    "    print(\n",
    "        f\"=> Using model {type(model).__name__} with loss {type(loss).__name__} on {torch.cuda.device_count()} GPUs\"\n",
    "    )\n",
    "\n",
    "    # Train Mapper model\n",
    "    print(f\"=> Training Mapper model\")\n",
    "    train(\n",
    "        model,\n",
    "        loss,\n",
    "        opt,\n",
    "        None,\n",
    "        train_loader,\n",
    "        val_loader,\n",
    "        args.epochs,\n",
    "        args.batch_size,\n",
    "        device,\n",
    "        args.chk_pt_path,\n",
    "        False,\n",
    "    )\n",
    "\n",
    "\n",
    "def train(\n",
    "        model,\n",
    "        loss_fn,\n",
    "        opt,\n",
    "        scheduler,\n",
    "        train_loader,\n",
    "        val_loader,\n",
    "        epochs,\n",
    "        batch_size,\n",
    "        device,\n",
    "        chk_pt_path,\n",
    "        early_stop=False,\n",
    "):\n",
    "    \"\"\"\n",
    "    Training loop.\n",
    "\n",
    "    Parameters:\n",
    "        model (torch.nn.Module): Model\n",
    "        loss_fn (torch.nn.Module): Loss function\n",
    "        opt (torch.optim): Optimizer\n",
    "        scheduler (torch.optim.lr_scheduler): LR scheduler (set to None if no scheduler)\n",
    "        train_loader (torch.utils.data.DataLoader): Dataloader for training data\n",
    "        val_loader (torch.utils.data.DataLoader): Dataloader for validation data\n",
    "        epochs (int): Number of training epochs\n",
    "        batch_size (int): Batch size\n",
    "        checkpoint_dir (str): Directory for storing model weights\n",
    "        early_stop (bool): True if early stopping based on validation loss\n",
    "    \"\"\"\n",
    "    scaler = torch.cuda.amp.GradScaler()\n",
    "    epochs_no_improvements = 0\n",
    "    best_val_loss = np.inf\n",
    "\n",
    "    for epoch in range(0, epochs):\n",
    "        model.train()\n",
    "        time_start = time.time()\n",
    "        progress_iter = tqdm(enumerate(train_loader), desc=f\"[{epoch + 1:03d}/{epochs:03d} epoch train]\",\n",
    "                             total=len(train_loader))\n",
    "        for step, sample in progress_iter:\n",
    "            opt.zero_grad()\n",
    "\n",
    "            with torch.cuda.amp.autocast(enabled=True):\n",
    "                pred = model(sample)\n",
    "\n",
    "            loss = loss_fn(pred, sample)\n",
    "            scaler.scale(loss).backward()\n",
    "            scaler.step(opt)\n",
    "            scaler.update()\n",
    "\n",
    "            # summary = (\n",
    "            #     \"\\r[Epoch {}][Step {}/{}] Loss: {}, Lr: {} - {:.2f} m remaining\".format(\n",
    "            #         epoch + 1,\n",
    "            #         step,\n",
    "            #         int(len(train_loader.dataset) / batch_size),\n",
    "            #         \"{}: {:.2f}\".format(\n",
    "            #             type(loss_fn).__name__, loss_fn.mean_running_loss\n",
    "            #         ),\n",
    "            #         *[group[\"lr\"] for group in opt.param_groups],\n",
    "            #         ((time.time() - time_start) / (step + 1))\n",
    "            #         * ((len(train_loader.dataset) / batch_size) - step)\n",
    "            #         / 60,\n",
    "            #     )\n",
    "            # )\n",
    "            # print(summary)\n",
    "\n",
    "            progress_iter.set_postfix(\n",
    "                {\n",
    "                    \"lr\": [opt.param_groups[0]['lr']],\n",
    "                    \"loss\": f\"{loss_fn.mean_running_loss:.4f}\",\n",
    "                    \"CUDA-Mem\": f\"{torch.cuda.memory_usage(device)}%\",\n",
    "                    \"CUDA-Util\": f\"{torch.cuda.utilization(device)}%\",\n",
    "                }\n",
    "            )\n",
    "        time_end = time.time()\n",
    "        elapse_time = time_end - time_start\n",
    "        print(\"Finished in {}s\".format(int(elapse_time)))\n",
    "\n",
    "        torch.save(model.state_dict(), os.path.join(chk_pt_path, f\"epoch_{epoch + 1}.pkl\"))\n",
    "\n",
    "        val_loss = evaluate(model, epoch, epochs, loss_fn, val_loader, device)\n",
    "        epochs_no_improvements += 1\n",
    "        if val_loss < best_val_loss:\n",
    "            print(\"Saving best model\")\n",
    "            torch.save(\n",
    "                model.state_dict(), os.path.join(chk_pt_path, f\"best.pkl\")\n",
    "            )\n",
    "            epochs_no_improvements = 0\n",
    "            best_val_loss = val_loss\n",
    "\n",
    "        if scheduler:\n",
    "            scheduler.step(val_loss)\n",
    "\n",
    "        if epochs_no_improvements == 5:\n",
    "            print(\"Early stop reached\")\n",
    "            return\n",
    "\n",
    "\n",
    "def evaluate(model, epoch, epochs, loss_fn, val_loader, device, split=\"val\"):\n",
    "    \"\"\"\n",
    "    Validation loop.\n",
    "\n",
    "    Parameters:\n",
    "        model (torch.nn.Module): Model\n",
    "        loss_fn (torch.nn.Module): Loss function\n",
    "        val_loader (torch.utils.data.DataLoader): Dataloader for validation data\n",
    "        split (str): Evaluation split\n",
    "    Returns:\n",
    "        loss (torch.Tensor): Validation loss\n",
    "    \"\"\"\n",
    "    print(f\"Evaluating on {split}\")\n",
    "    model.eval()\n",
    "\n",
    "    running_loss = 0\n",
    "    num_batches = 0\n",
    "    with torch.no_grad():\n",
    "        progress_iter = tqdm(enumerate(val_loader), desc=f\"[{epoch + 1:03d}/{epochs:03d} epoch train]\",\n",
    "                             total=len(val_loader))\n",
    "        for step, sample in progress_iter:\n",
    "            pred = model(sample)\n",
    "\n",
    "            running_loss += loss_fn(pred, sample)\n",
    "            num_batches += 1\n",
    "\n",
    "            progress_iter.set_postfix(\n",
    "                {\n",
    "                    \"loss\": f\"{running_loss:.4f}\",\n",
    "                    \"CUDA-Mem\": f\"{torch.cuda.memory_usage(device)}%\",\n",
    "                    \"CUDA-Util\": f\"{torch.cuda.utilization(device)}%\",\n",
    "                }\n",
    "            )\n",
    "\n",
    "    loss = running_loss / num_batches\n",
    "    print(f\"Eval Loss = {loss}\")\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34013265",
   "metadata": {},
   "source": [
    "### Step1: generate the report-sentence or prompt embedding which will localize the finding (mass)\n",
    "\n",
    "#### Replace `data_dir`, `breast_clip_path` and `model` based on your paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "386683f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = \"<Path where the sentence/prompt embeddings will be saved>\"\n",
    "breast_clip_path = \"<Mammo-CLIP checkpoint path>\"\n",
    "model = \"b5_all_clip_mapper\"\n",
    "generate_attribute_embs(data_dir, breast_clip_path, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee144be1",
   "metadata": {},
   "source": [
    "### Step2: Train Mammo-Factor\n",
    "#### 1. Get the path of `args` uploaded [here](https://github.com/batmanlab/Mammo-CLIP/blob/main/src/codebase/notebooks/Mammo-Factor/seed_10_train_configs.pkl)\n",
    "#### Replace `args.clip_chk_pt_path` and `args.attr_embs_path` as per covenience.\n",
    "####  `args.attr_embs_path` is the path of the sentence/prompt embedding generated in Stage1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f60416b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "args = pickle.load(\n",
    "        open(\n",
    "            \"seed_10_train_configs.pkl\", \n",
    "            \"rb\"))\n",
    "args.root = f\"lr_{args.lr}_epochs_{args.epochs}\"\n",
    "args.apex = True if args.apex == \"y\" else False\n",
    "args.running_interactive = True if args.running_interactive == \"y\" else False\n",
    "\n",
    "## Change this two\n",
    "args.arch = model\n",
    "args.clip_chk_pt_path = \"<Mammo-CLIP checkpoint path>\"\n",
    "args.attr_embs_path = \"<First extract the embeddings of the report sentences and give path to that embedding file>\"\n",
    "\n",
    "chk_pt_path, output_path, tb_logs_path = get_Paths(args)\n",
    "args.chk_pt_path = chk_pt_path\n",
    "args.output_path = output_path\n",
    "args.tb_logs_path = tb_logs_path\n",
    "\n",
    "os.makedirs(chk_pt_path, exist_ok=True)\n",
    "os.makedirs(output_path, exist_ok=True)\n",
    "os.makedirs(tb_logs_path, exist_ok=True)\n",
    "print(\"====================> Paths <====================\")\n",
    "print(f\"checkpoint_path: {chk_pt_path}\")\n",
    "print(f\"output_path: {output_path}\")\n",
    "print(f\"tb_logs_path: {tb_logs_path}\")\n",
    "print('device:', device)\n",
    "print('torch version:', torch.__version__)\n",
    "print(\"====================> Paths <====================\")\n",
    "\n",
    "pickle.dump(args, open(os.path.join(output_path, f\"seed_{args.seed}_train_configs.pkl\"), \"wb\"))\n",
    "train_region_mapper(args, device)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}