[ { "path": ".gitignore", "content": "**/__pycache__\nruns/\nmodels/\ndatasets/\ndatasets\n.vscode/\ncore*\nvis_output/\ntest_vis/\nopenai\n.DS_Store\nXMem/weights/" }, { "path": ".gitmodules", "content": "[submodule \"LLaMA-VID\"]\n\tpath = LLaMA-VID\n\turl = git@github.com:dvlab-research/LLaMA-VID.git" }, { "path": "README.md", "content": "# VISA: Reasoning Video Object Segmentation via Large Language Model\n\n
\n[![ GitHub stars ](https://badgen.net/github/stars/cilinyan/VISA)](https://github.com/cilinyan/VISA)\n[![arXiv](https://img.shields.io/badge/arXiv-2407.11325-972B24)](http://arxiv.org/abs/2407.11325)\n[![Static Badge](https://img.shields.io/badge/dataset-ReVOS-3F7EBC)](https://github.com/cilinyan/ReVOS-api)\n
\n\n
\n\n
\n\n## ๐Ÿš€ Performance\n\n
\n VISA demonstrates remarkable proficiency in handling complex segmentation tasks that require: (a) reasoning based on world knowledge; (b) inference of future events; and (c) a comprehensive understanding of video content.\n
\n\n
\n\n
\n\n## ๐Ÿ› ๏ธ Installation\n```shell\npip install -r requirements.txt\npip install flash-attn --no-build-isolation\n```\n\n## ๐Ÿฆ„ Training and Validation\n\n### 1. Training Data Preparation\n\nBefore training, please download the datasets, and then configure the path in [dataset_config.py](utils/dataset_config.py).\n\n
\n LISA's Dataset \n \nFollow [LISA](https://github.com/dvlab-research/LISA/tree/main) to prepare LISA's datasets. The dataset folder should be stored in the `$LISA_ROOT` folder.\n\n```\nLISA_ROOT\nโ”œโ”€โ”€ ade20k\nโ”œโ”€โ”€ coco\nโ”œโ”€โ”€ cocostuff\nโ”œโ”€โ”€ llava_dataset\nโ”œโ”€โ”€ mapillary\nโ”œโ”€โ”€ reason_seg\nโ”œโ”€โ”€ refer_seg\nโ””โ”€โ”€ vlpart\n```\n\n
\n\n\n
\n Chat-UniVi's Dataset \n\nFollow [Chat-UniVi/Chat-UniVi-Instruct](https://huggingface.co/datasets/Chat-UniVi/Chat-UniVi-Instruct/tree/main) to prepare `Chat-UniVi-Instruct` datasets. The dataset folder should be stored in the `$ChatUniVi_ROOT` folder.\n```\nChatUniVi_ROOT\nโ”œโ”€โ”€ Fine-tuning\nโ”‚ โ”œโ”€โ”€ MIMIC_imageonly\nโ”‚ โ””โ”€โ”€ VIDEO\nโ””โ”€โ”€ ScienceQA_tuning\n```\n\n
\n\n\n
\n RVOS's Dataset \n\n1. Reasoning Video Segmentation Datasets: [ReVOS](https://github.com/cilinyan/ReVOS-api).\n2. Referring Video Segmentation Datasets: [Ref-Youtube-VOS](https://github.com/wjn922/ReferFormer/blob/main/docs/data.md), [Ref-DAVIS17](https://github.com/wjn922/ReferFormer/blob/main/docs/data.md), [MeViS](https://github.com/henghuiding/MeViS).\n - Ref-Youtube-VOS: Download `mask_dict.pkl` from [OneDrive](https://mailsjlueducn-my.sharepoint.com/:f:/g/personal/yancl9918_mails_jlu_edu_cn/EqR9g3yWG5pPtVoil0EfsbgBJhCZ7YwaRG9w9GsYy1_N5g?e=JLaJfc) or [BaiduPan](https://pan.baidu.com/s/1mbJaDDy0UTlA7sysp0zypg?pwd=visa).\n - Ref-DAVIS17: Download `mask_dict.pkl` from [OneDrive](https://mailsjlueducn-my.sharepoint.com/:f:/g/personal/yancl9918_mails_jlu_edu_cn/Eq8bmGqNcYxGhQ1bioN65q4B_gPxIabpJUjGaV5uqcaq3w?e=2J6Ldp) or [BaiduPan](https://pan.baidu.com/s/1Gg5qPvxRZMKDp0JrVRJ75w?pwd=visa).\n3. Open-Vocabulary Video Instance Segmentation Dataset: [LV-VIS](https://github.com/haochenheheda/LVVIS/tree/main).\n Download `mask_dict.json` and `meta_expressions.json` from [OneDrive](https://mailsjlueducn-my.sharepoint.com/:f:/g/personal/yancl9918_mails_jlu_edu_cn/EttXAjMV8yFJhHMQwX3mIw0BP7dymKV-cuw4uAotDaAwYw?e=j6Y44X) or [BaiduPan](https://pan.baidu.com/s/1LOWPnuxXF_LXGSL7osRptA?pwd=visa). Then, put the annotations files in the `$RVOS_ROOT/lvvis/train` directory as follows.\n```\nRVOS_ROOT\nโ”œโ”€โ”€ ReVOS\nโ”‚ โ”œโ”€โ”€ JPEGImages \nโ”‚ โ”œโ”€โ”€ mask_dict.json \nโ”‚ โ”œโ”€โ”€ mask_dict_foreground.json \nโ”‚ โ”œโ”€โ”€ meta_expressions_train_.json \nโ”‚ โ””โ”€โ”€ meta_expressions_valid_.json \nโ”œโ”€โ”€ lvvis\nโ”‚ โ””โ”€โ”€ train\n| โ”œโ”€โ”€ JPEGImages\n| โ”œโ”€โ”€ mask_dict.json\n| โ””โ”€โ”€ meta_expressions.json\nโ”œโ”€โ”€ Ref-Youtube-VOS\nโ”‚ โ”œโ”€โ”€ meta_expressions\n| | โ”œโ”€โ”€ train/meta_expressions.json\n| | โ””โ”€โ”€ valid/meta_expressions.json\nโ”‚ โ”œโ”€โ”€ train\n| | โ”œโ”€โ”€ JPEGImages\n| | โ””โ”€โ”€ mask_dict.pkl\nโ”‚ โ””โ”€โ”€ valid\n| โ””โ”€โ”€ JPEGImages\nโ”œโ”€โ”€ davis17\nโ”‚ โ”œโ”€โ”€ meta_expressions\n| | โ”œโ”€โ”€ train/meta_expressions.json\n| | โ””โ”€โ”€ valid/meta_expressions.json\nโ”‚ โ”œโ”€โ”€ train\n| | โ”œโ”€โ”€ JPEGImages\n| | โ””โ”€โ”€ mask_dict.pkl\nโ”‚ โ””โ”€โ”€ valid\n| โ”œโ”€โ”€ JPEGImages\n| โ””โ”€โ”€ mask_dict.pkl\nโ””โ”€โ”€ mevis\n```\n\n
\n\n### 2. Pre-trained weights\n\n
\n Chat-UniVi \n\nTo train VISA-7B or 13B, you need to download Chat-UniVi weights from [Chat-UniVi-7B](https://huggingface.co/Chat-UniVi/Chat-UniVi) and [Chat-UniVi-13B](https://huggingface.co/Chat-UniVi/Chat-UniVi-13B).\n\n
\n\n
\n SAM \n\nDownload SAM ViT-H pre-trained weights from the [link](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth).\n\n
\n\n### 3. Training VISA\n\n```shell\n# Training VISA-7B\nbash scripts/train_7b.sh \n\n# Extracting fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints.\ncd /PATH/TO/VISA-7B/ckpt_model && python zero_to_fp32.py . ../pytorch_model.bin\n\n# Merge LoRA Weight\nCUDA_VISIBLE_DEVICES=\"\" python merge_lora_weights_and_save_hf_model.py \\\n --version Chat-UniVi/Chat-UniVi \\\n --weight /PATH/TO/VISA-7B/pytorch_model.bin \\\n --save_path /PATH/TO/VISA-7B/hf_model\n```\n\n### 4. Validation\n\n
\n 1. Using `VISA` to generate predicted mask of each video [demo] \n\n```shell\ndeepspeed --master_port=24999 train_ds.py \\\n --version=\"/PATH/TO/VISA-7B/hf_model\" \\\n --vision_pretrained=\"/PATH/TO/sam_vit_h_4b8939.pth\" \\\n --log_base_dir=\"/PATH/TO/LOG_BASE_DIR\" \\\n --exp_name=\"val_7b\" \\\n --balance_sample \\\n --dataset=\"reason_seg\" \\\n --sample_rates=\"13\" \\\n --val_dataset \"revos_valid\" \\\n --eval_only \n```\n\n
\n\n
\n 2. Using LLaMA-VID to generate target frame for each video \n\n> You can directly download the results of our run from [OneDrive](https://mailsjlueducn-my.sharepoint.com/:u:/g/personal/yancl9918_mails_jlu_edu_cn/ETmoJF2i8ZZBsgIwdELiL8gBfptZZoPWjx6Y0eH6Myr3sw?e=mTt6rO) or [BaiduPan](https://pan.baidu.com/s/1YWs6NLPvANfhgUBHKQwnBg?pwd=visa)\n\n- Run [http_server_mp.py](https://github.com/cilinyan/VISA/blob/main/utils_llamavid/llamavid_server.py) to build the API server for LLaMA-VID [`[demo]`](https://github.com/cilinyan/VISA/blob/c53d2cd31407eab583c5eb04f84fd95b4694f2ce/utils_llamavid/llamavid_server.py#L215-L220)\n ```shell\n python utils_llamavid/llamavid_server.py \\\n --vision_tower /PATH/TO/eva_vit_g.pth \\\n --image_processor /PATH/TO/openai/clip-vit-large-patch14 \\\n --model-path /PATH/TO/YanweiLi/llama-vid-13b-full-224-video-fps-1\n ```\n- Using the API for inference [`[demo]`](https://github.com/cilinyan/VISA/blob/c53d2cd31407eab583c5eb04f84fd95b4694f2ce/utils_llamavid/llamavid_client.py#L58-L63)\n \n ```shell\n python utils_llamavid/llamavid_client.py \\\n --video_root /PATH/TO/ReVOS/JPEGImages \\\n --data_json_file /PATH/TO/ReVOS/meta_expressions_valid_.json\n ```\n\n
\n\n
\n 3. Using XMem for mask propagation [demo] \n
\n\n
\n\n 4. Evaluate ReVOS's performance [demo] \n\n```shell\ncd tools\npython eval_revos.py /PATH/TO/FINAL_ANNOTATION [ARGS]\n```\n \n
\n\n## ๐Ÿ“‘ Todo list\n \n- [x] Release code with `Text-guided Frame Sampler`'s Local Sampling\n\n- [ ] Release VISA model weights [issue #6](https://github.com/cilinyan/VISA/issues/6)\n\n- [ ] Release code with `Text-guided Frame Sampler`'s Global-Local Sampling\n\n## โญ Cite\n\nIf you find this project useful in your research, please consider citing:\n\n```\n@article{yan2024visa,\n title={VISA: Reasoning Video Object Segmentation via Large Language Models},\n author={Yan, Cilin and Wang, Haochen and Yan, Shilin and Jiang, Xiaolong and Hu, Yao and Kang, Guoliang and Xie, Weidi and Gavves, Efstratios},\n journal={arXiv preprint arXiv:2407.11325},\n year={2024}\n}\n```\n\n## ๐ŸŽ–๏ธ Acknowledgement\nThis work is built upon the [LLaVA](https://github.com/haotian-liu/LLaVA), [SAM](https://github.com/facebookresearch/segment-anything), [LISA](https://github.com/dvlab-research/LISA), [Chat-UniVi](https://github.com/PKU-YuanGroup/Chat-UniVi), [MeViS](https://github.com/henghuiding/MeViS), [LLaMA-VID](https://github.com/dvlab-research/LLaMA-VID) and [XMem](https://github.com/hkchengrex/XMem).\n\n" }, { "path": "XMem/dataset/__init__.py", "content": "" }, { "path": "XMem/dataset/range_transform.py", "content": "import torchvision.transforms as transforms\n\nim_mean = (124, 116, 104)\n\nim_normalization = transforms.Normalize(\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]\n )\n\ninv_im_trans = transforms.Normalize(\n mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225],\n std=[1/0.229, 1/0.224, 1/0.225])\n" }, { "path": "XMem/dataset/reseed.py", "content": "import torch\nimport random\n\ndef reseed(seed):\n random.seed(seed)\n torch.manual_seed(seed)" }, { "path": "XMem/dataset/static_dataset.py", "content": "import os\nfrom os import path\n\nimport torch\nfrom torch.utils.data.dataset import Dataset\nfrom torchvision import transforms\nfrom torchvision.transforms import InterpolationMode\nfrom PIL import Image\nimport numpy as np\n\nfrom dataset.range_transform import im_normalization, im_mean\nfrom dataset.tps import random_tps_warp\nfrom dataset.reseed import reseed\n\n\nclass StaticTransformDataset(Dataset):\n \"\"\"\n Generate pseudo VOS data by applying random transforms on static images.\n Single-object only.\n\n Method 0 - FSS style (class/1.jpg class/1.png)\n Method 1 - Others style (XXX.jpg XXX.png)\n \"\"\"\n def __init__(self, parameters, num_frames=3, max_num_obj=1):\n self.num_frames = num_frames\n self.max_num_obj = max_num_obj\n\n self.im_list = []\n for parameter in parameters:\n root, method, multiplier = parameter\n if method == 0:\n # Get images\n classes = os.listdir(root)\n for c in classes:\n imgs = os.listdir(path.join(root, c))\n jpg_list = [im for im in imgs if 'jpg' in im[-3:].lower()]\n\n joint_list = [path.join(root, c, im) for im in jpg_list]\n self.im_list.extend(joint_list * multiplier)\n\n elif method == 1:\n self.im_list.extend([path.join(root, im) for im in os.listdir(root) if '.jpg' in im] * multiplier)\n\n print(f'{len(self.im_list)} images found.')\n\n # These set of transform is the same for im/gt pairs, but different among the 3 sampled frames\n self.pair_im_lone_transform = transforms.Compose([\n transforms.ColorJitter(0.1, 0.05, 0.05, 0), # No hue change here as that's not realistic\n ])\n\n self.pair_im_dual_transform = transforms.Compose([\n transforms.RandomAffine(degrees=20, scale=(0.9,1.1), shear=10, interpolation=InterpolationMode.BICUBIC, fill=im_mean),\n transforms.Resize(384, InterpolationMode.BICUBIC),\n transforms.RandomCrop((384, 384), pad_if_needed=True, fill=im_mean),\n ])\n\n self.pair_gt_dual_transform = transforms.Compose([\n transforms.RandomAffine(degrees=20, scale=(0.9,1.1), shear=10, interpolation=InterpolationMode.BICUBIC, fill=0),\n transforms.Resize(384, InterpolationMode.NEAREST),\n transforms.RandomCrop((384, 384), pad_if_needed=True, fill=0),\n ])\n\n\n # These transform are the same for all pairs in the sampled sequence\n self.all_im_lone_transform = transforms.Compose([\n transforms.ColorJitter(0.1, 0.05, 0.05, 0.05),\n transforms.RandomGrayscale(0.05),\n ])\n\n self.all_im_dual_transform = transforms.Compose([\n transforms.RandomAffine(degrees=0, scale=(0.8, 1.5), fill=im_mean),\n transforms.RandomHorizontalFlip(),\n ])\n\n self.all_gt_dual_transform = transforms.Compose([\n transforms.RandomAffine(degrees=0, scale=(0.8, 1.5), fill=0),\n transforms.RandomHorizontalFlip(),\n ])\n\n # Final transform without randomness\n self.final_im_transform = transforms.Compose([\n transforms.ToTensor(),\n im_normalization,\n ])\n\n self.final_gt_transform = transforms.Compose([\n transforms.ToTensor(),\n ])\n\n def _get_sample(self, idx):\n im = Image.open(self.im_list[idx]).convert('RGB')\n gt = Image.open(self.im_list[idx][:-3]+'png').convert('L')\n\n sequence_seed = np.random.randint(2147483647)\n\n images = []\n masks = []\n for _ in range(self.num_frames):\n reseed(sequence_seed)\n this_im = self.all_im_dual_transform(im)\n this_im = self.all_im_lone_transform(this_im)\n reseed(sequence_seed)\n this_gt = self.all_gt_dual_transform(gt)\n\n pairwise_seed = np.random.randint(2147483647)\n reseed(pairwise_seed)\n this_im = self.pair_im_dual_transform(this_im)\n this_im = self.pair_im_lone_transform(this_im)\n reseed(pairwise_seed)\n this_gt = self.pair_gt_dual_transform(this_gt)\n\n # Use TPS only some of the times\n # Not because TPS is bad -- just that it is too slow and I need to speed up data loading\n if np.random.rand() < 0.33:\n this_im, this_gt = random_tps_warp(this_im, this_gt, scale=0.02)\n\n this_im = self.final_im_transform(this_im)\n this_gt = self.final_gt_transform(this_gt)\n\n images.append(this_im)\n masks.append(this_gt)\n\n images = torch.stack(images, 0)\n masks = torch.stack(masks, 0)\n\n return images, masks.numpy()\n\n def __getitem__(self, idx):\n additional_objects = np.random.randint(self.max_num_obj)\n indices = [idx, *np.random.randint(self.__len__(), size=additional_objects)]\n\n merged_images = None\n merged_masks = np.zeros((self.num_frames, 384, 384), dtype=np.int64)\n\n for i, list_id in enumerate(indices):\n images, masks = self._get_sample(list_id)\n if merged_images is None:\n merged_images = images\n else:\n merged_images = merged_images*(1-masks) + images*masks\n merged_masks[masks[:,0]>0.5] = (i+1)\n\n masks = merged_masks\n\n labels = np.unique(masks[0])\n # Remove background\n labels = labels[labels!=0]\n target_objects = labels.tolist()\n\n # Generate one-hot ground-truth\n cls_gt = np.zeros((self.num_frames, 384, 384), dtype=np.int64)\n first_frame_gt = np.zeros((1, self.max_num_obj, 384, 384), dtype=np.int64)\n for i, l in enumerate(target_objects):\n this_mask = (masks==l)\n cls_gt[this_mask] = i+1\n first_frame_gt[0,i] = (this_mask[0])\n cls_gt = np.expand_dims(cls_gt, 1)\n\n info = {}\n info['name'] = self.im_list[idx]\n info['num_objects'] = max(1, len(target_objects))\n\n # 1 if object exist, 0 otherwise\n selector = [1 if i < info['num_objects'] else 0 for i in range(self.max_num_obj)]\n selector = torch.FloatTensor(selector)\n\n data = {\n 'rgb': merged_images,\n 'first_frame_gt': first_frame_gt,\n 'cls_gt': cls_gt,\n 'selector': selector,\n 'info': info\n }\n\n return data\n\n\n def __len__(self):\n return len(self.im_list)\n" }, { "path": "XMem/dataset/tps.py", "content": "import numpy as np\nfrom PIL import Image\nimport cv2\nimport thinplate as tps\n\ncv2.setNumThreads(0)\n\ndef pick_random_points(h, w, n_samples):\n y_idx = np.random.choice(np.arange(h), size=n_samples, replace=False)\n x_idx = np.random.choice(np.arange(w), size=n_samples, replace=False)\n return y_idx/h, x_idx/w\n\n\ndef warp_dual_cv(img, mask, c_src, c_dst):\n dshape = img.shape\n theta = tps.tps_theta_from_points(c_src, c_dst, reduced=True)\n grid = tps.tps_grid(theta, c_dst, dshape)\n mapx, mapy = tps.tps_grid_to_remap(grid, img.shape)\n return cv2.remap(img, mapx, mapy, cv2.INTER_LINEAR), cv2.remap(mask, mapx, mapy, cv2.INTER_NEAREST)\n\n\ndef random_tps_warp(img, mask, scale, n_ctrl_pts=12):\n \"\"\"\n Apply a random TPS warp of the input image and mask\n Uses randomness from numpy\n \"\"\"\n img = np.asarray(img)\n mask = np.asarray(mask)\n\n h, w = mask.shape\n points = pick_random_points(h, w, n_ctrl_pts)\n c_src = np.stack(points, 1)\n c_dst = c_src + np.random.normal(scale=scale, size=c_src.shape)\n warp_im, warp_gt = warp_dual_cv(img, mask, c_src, c_dst)\n\n return Image.fromarray(warp_im), Image.fromarray(warp_gt)\n\n" }, { "path": "XMem/dataset/util.py", "content": "import numpy as np\n\n\ndef all_to_onehot(masks, labels):\n if len(masks.shape) == 3:\n Ms = np.zeros((len(labels), masks.shape[0], masks.shape[1], masks.shape[2]), dtype=np.uint8)\n else:\n Ms = np.zeros((len(labels), masks.shape[0], masks.shape[1]), dtype=np.uint8)\n\n for ni, l in enumerate(labels):\n Ms[ni] = (masks == l).astype(np.uint8)\n \n return Ms\n" }, { "path": "XMem/dataset/vos_dataset.py", "content": "import os\nfrom os import path, replace\n\nimport torch\nfrom torch.utils.data.dataset import Dataset\nfrom torchvision import transforms\nfrom torchvision.transforms import InterpolationMode\nfrom PIL import Image\nimport numpy as np\n\nfrom dataset.range_transform import im_normalization, im_mean\nfrom dataset.reseed import reseed\n\n\nclass VOSDataset(Dataset):\n \"\"\"\n Works for DAVIS/YouTubeVOS/BL30K training\n For each sequence:\n - Pick three frames\n - Pick two objects\n - Apply some random transforms that are the same for all frames\n - Apply random transform to each of the frame\n - The distance between frames is controlled\n \"\"\"\n def __init__(self, im_root, gt_root, max_jump, is_bl, subset=None, num_frames=3, max_num_obj=3, finetune=False):\n self.im_root = im_root\n self.gt_root = gt_root\n self.max_jump = max_jump\n self.is_bl = is_bl\n self.num_frames = num_frames\n self.max_num_obj = max_num_obj\n\n self.videos = []\n self.frames = {}\n\n vid_list = sorted(os.listdir(self.im_root))\n # Pre-filtering\n for vid in vid_list:\n if subset is not None:\n if vid not in subset:\n continue\n frames = sorted(os.listdir(os.path.join(self.im_root, vid)))\n if len(frames) < num_frames:\n continue\n self.frames[vid] = frames\n self.videos.append(vid)\n\n print('%d out of %d videos accepted in %s.' % (len(self.videos), len(vid_list), im_root))\n\n # These set of transform is the same for im/gt pairs, but different among the 3 sampled frames\n self.pair_im_lone_transform = transforms.Compose([\n transforms.ColorJitter(0.01, 0.01, 0.01, 0),\n ])\n\n self.pair_im_dual_transform = transforms.Compose([\n transforms.RandomAffine(degrees=0 if finetune or self.is_bl else 15, shear=0 if finetune or self.is_bl else 10, interpolation=InterpolationMode.BILINEAR, fill=im_mean),\n ])\n\n self.pair_gt_dual_transform = transforms.Compose([\n transforms.RandomAffine(degrees=0 if finetune or self.is_bl else 15, shear=0 if finetune or self.is_bl else 10, interpolation=InterpolationMode.NEAREST, fill=0),\n ])\n\n # These transform are the same for all pairs in the sampled sequence\n self.all_im_lone_transform = transforms.Compose([\n transforms.ColorJitter(0.1, 0.03, 0.03, 0),\n transforms.RandomGrayscale(0.05),\n ])\n\n if self.is_bl:\n # Use a different cropping scheme for the blender dataset because the image size is different\n self.all_im_dual_transform = transforms.Compose([\n transforms.RandomHorizontalFlip(),\n transforms.RandomResizedCrop((384, 384), scale=(0.25, 1.00), interpolation=InterpolationMode.BILINEAR)\n ])\n\n self.all_gt_dual_transform = transforms.Compose([\n transforms.RandomHorizontalFlip(),\n transforms.RandomResizedCrop((384, 384), scale=(0.25, 1.00), interpolation=InterpolationMode.NEAREST)\n ])\n else:\n self.all_im_dual_transform = transforms.Compose([\n transforms.RandomHorizontalFlip(),\n transforms.RandomResizedCrop((384, 384), scale=(0.36,1.00), interpolation=InterpolationMode.BILINEAR)\n ])\n\n self.all_gt_dual_transform = transforms.Compose([\n transforms.RandomHorizontalFlip(),\n transforms.RandomResizedCrop((384, 384), scale=(0.36,1.00), interpolation=InterpolationMode.NEAREST)\n ])\n\n # Final transform without randomness\n self.final_im_transform = transforms.Compose([\n transforms.ToTensor(),\n im_normalization,\n ])\n\n def __getitem__(self, idx):\n video = self.videos[idx]\n info = {}\n info['name'] = video\n\n vid_im_path = path.join(self.im_root, video)\n vid_gt_path = path.join(self.gt_root, video)\n frames = self.frames[video]\n\n trials = 0\n while trials < 5:\n info['frames'] = [] # Appended with actual frames\n\n num_frames = self.num_frames\n length = len(frames)\n this_max_jump = min(len(frames), self.max_jump)\n\n # iterative sampling\n frames_idx = [np.random.randint(length)]\n acceptable_set = set(range(max(0, frames_idx[-1]-this_max_jump), min(length, frames_idx[-1]+this_max_jump+1))).difference(set(frames_idx))\n while(len(frames_idx) < num_frames):\n idx = np.random.choice(list(acceptable_set))\n frames_idx.append(idx)\n new_set = set(range(max(0, frames_idx[-1]-this_max_jump), min(length, frames_idx[-1]+this_max_jump+1)))\n acceptable_set = acceptable_set.union(new_set).difference(set(frames_idx))\n\n frames_idx = sorted(frames_idx)\n if np.random.rand() < 0.5:\n # Reverse time\n frames_idx = frames_idx[::-1]\n\n sequence_seed = np.random.randint(2147483647)\n images = []\n masks = []\n target_objects = []\n for f_idx in frames_idx:\n jpg_name = frames[f_idx][:-4] + '.jpg'\n png_name = frames[f_idx][:-4] + '.png'\n info['frames'].append(jpg_name)\n\n reseed(sequence_seed)\n this_im = Image.open(path.join(vid_im_path, jpg_name)).convert('RGB')\n this_im = self.all_im_dual_transform(this_im)\n this_im = self.all_im_lone_transform(this_im)\n reseed(sequence_seed)\n this_gt = Image.open(path.join(vid_gt_path, png_name)).convert('P')\n this_gt = self.all_gt_dual_transform(this_gt)\n\n pairwise_seed = np.random.randint(2147483647)\n reseed(pairwise_seed)\n this_im = self.pair_im_dual_transform(this_im)\n this_im = self.pair_im_lone_transform(this_im)\n reseed(pairwise_seed)\n this_gt = self.pair_gt_dual_transform(this_gt)\n\n this_im = self.final_im_transform(this_im)\n this_gt = np.array(this_gt)\n\n images.append(this_im)\n masks.append(this_gt)\n\n images = torch.stack(images, 0)\n\n labels = np.unique(masks[0])\n # Remove background\n labels = labels[labels!=0]\n\n if self.is_bl:\n # Find large enough labels\n good_lables = []\n for l in labels:\n pixel_sum = (masks[0]==l).sum()\n if pixel_sum > 10*10:\n # OK if the object is always this small\n # Not OK if it is actually much bigger\n if pixel_sum > 30*30:\n good_lables.append(l)\n elif max((masks[1]==l).sum(), (masks[2]==l).sum()) < 20*20:\n good_lables.append(l)\n labels = np.array(good_lables, dtype=np.uint8)\n \n if len(labels) == 0:\n target_objects = []\n trials += 1\n else:\n target_objects = labels.tolist()\n break\n\n if len(target_objects) > self.max_num_obj:\n target_objects = np.random.choice(target_objects, size=self.max_num_obj, replace=False)\n\n info['num_objects'] = max(1, len(target_objects))\n\n masks = np.stack(masks, 0)\n\n # Generate one-hot ground-truth\n cls_gt = np.zeros((self.num_frames, 384, 384), dtype=np.int64)\n first_frame_gt = np.zeros((1, self.max_num_obj, 384, 384), dtype=np.int64)\n for i, l in enumerate(target_objects):\n this_mask = (masks==l)\n cls_gt[this_mask] = i+1\n first_frame_gt[0,i] = (this_mask[0])\n cls_gt = np.expand_dims(cls_gt, 1)\n\n # 1 if object exist, 0 otherwise\n selector = [1 if i < info['num_objects'] else 0 for i in range(self.max_num_obj)]\n selector = torch.FloatTensor(selector)\n\n data = {\n 'rgb': images,\n 'first_frame_gt': first_frame_gt,\n 'cls_gt': cls_gt,\n 'selector': selector,\n 'info': info,\n }\n\n return data\n\n def __len__(self):\n return len(self.videos)" }, { "path": "XMem/eval.py", "content": "import os\nfrom os import path\nfrom argparse import ArgumentParser\nimport shutil\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nimport numpy as np\nfrom PIL import Image\n\nfrom inference.data.test_datasets import LongTestDataset, DAVISTestDataset, YouTubeVOSTestDataset\nfrom inference.data.mask_mapper import MaskMapper\nfrom model.network import XMem\nfrom inference.inference_core import InferenceCore\n\nfrom tqdm import tqdm\ntry:\n import hickle as hkl\nexcept ImportError:\n print('Failed to import hickle. Fine if not using multi-scale testing.')\n\nparser = ArgumentParser()\nparser.add_argument('--model', default='./saves/XMem.pth')\nparser.add_argument('--meta_exp', type=str)\n\n# Data options\nparser.add_argument('--d16_path', default='../DAVIS/2016')\nparser.add_argument('--d17_path', default='../DAVIS/2017')\nparser.add_argument('--y18_path', default='../YouTube2018')\nparser.add_argument('--y19_path', default='../YouTube')\nparser.add_argument('--lv_path', default='../long_video_set')\n# For generic (G) evaluation, point to a folder that contains \"JPEGImages\" and \"Annotations\"\nparser.add_argument('--generic_path')\nparser.add_argument('--img_dir')\nparser.add_argument('--reversed', action='store_true')\nparser.add_argument('--split_part', type=int, default=0)\n\n\nparser.add_argument('--dataset', help='D16/D17/Y18/Y19/LV1/LV3/G', default='D17')\nparser.add_argument('--split', help='val/test', default='val')\nparser.add_argument('--output', default=None)\nparser.add_argument('--save_all', action='store_true', \n help='Save all frames. Useful only in YouTubeVOS/long-time video', )\n\nparser.add_argument('--benchmark', action='store_true', help='enable to disable amp for FPS benchmarking')\n \n# Long-term memory options\nparser.add_argument('--disable_long_term', action='store_true')\nparser.add_argument('--max_mid_term_frames', help='T_max in paper, decrease to save memory', type=int, default=10)\nparser.add_argument('--min_mid_term_frames', help='T_min in paper, decrease to save memory', type=int, default=5)\nparser.add_argument('--max_long_term_elements', help='LT_max in paper, increase if objects disappear for a long time', type=int, default=10000)\nparser.add_argument('--num_prototypes', help='P in paper', type=int, default=128)\n\nparser.add_argument('--top_k', type=int, default=30)\nparser.add_argument('--mem_every', help='r in paper. Increase to improve running speed.', type=int, default=5)\nparser.add_argument('--deep_update_every', help='Leave -1 normally to synchronize with mem_every', type=int, default=-1)\n\n# Multi-scale options\nparser.add_argument('--save_scores', action='store_true')\nparser.add_argument('--flip', action='store_true')\nparser.add_argument('--size', default=480, type=int, help='Resize the shorter side to this size. -1 to use original resolution. ')\n\nargs = parser.parse_args()\nconfig = vars(args)\nconfig['enable_long_term'] = not config['disable_long_term']\n\nif args.output is None:\n args.output = f'../output/{args.dataset}_{args.split}'\n print(f'Output path not provided. Defaulting to {args.output}')\n\n\"\"\"\nData preparation\n\"\"\"\nis_youtube = args.dataset.startswith('Y')\nis_davis = args.dataset.startswith('D')\nis_lv = args.dataset.startswith('LV')\n\nif is_youtube or args.save_scores:\n out_path = path.join(args.output, 'Annotations')\nelse:\n out_path = args.output\n\nif is_youtube:\n if args.dataset == 'Y18':\n yv_path = args.y18_path\n elif args.dataset == 'Y19':\n yv_path = args.y19_path\n\n if args.split == 'val':\n args.split = 'valid'\n meta_dataset = YouTubeVOSTestDataset(data_root=yv_path, split='valid', size=args.size)\n elif args.split == 'test':\n meta_dataset = YouTubeVOSTestDataset(data_root=yv_path, split='test', size=args.size)\n else:\n raise NotImplementedError\n\nelif is_davis:\n if args.dataset == 'D16':\n if args.split == 'val':\n # Set up Dataset, a small hack to use the image set in the 2017 folder because the 2016 one is of a different format\n meta_dataset = DAVISTestDataset(args.d16_path, imset='../../2017/trainval/ImageSets/2016/val.txt', size=args.size)\n else:\n raise NotImplementedError\n palette = None\n elif args.dataset == 'D17':\n if args.split == 'val':\n meta_dataset = DAVISTestDataset(path.join(args.d17_path, 'trainval'), imset='2017/val.txt', size=args.size)\n elif args.split == 'test':\n meta_dataset = DAVISTestDataset(path.join(args.d17_path, 'test-dev'), imset='2017/test-dev.txt', size=args.size)\n else:\n raise NotImplementedError\n\nelif is_lv:\n if args.dataset == 'LV1':\n meta_dataset = LongTestDataset(args.meta_exp, path.join(args.lv_path, 'long_video'))\n elif args.dataset == 'LV3':\n meta_dataset = LongTestDataset(args.meta_exp, path.join(args.lv_path, 'long_video_x3'))\n else:\n raise NotImplementedError\nelif args.dataset == 'G':\n meta_dataset = LongTestDataset(args.meta_exp, path.join(args.generic_path), size=args.size, img_dir=args.img_dir, reversed_=args.reversed, split_part=args.split_part)\n if not args.save_all:\n args.save_all = True\n print('save_all is forced to be true in generic evaluation mode.')\nelse:\n raise NotImplementedError\n\ntorch.autograd.set_grad_enabled(False)\n\n# Set up loader\nmeta_loader = meta_dataset.get_datasets()\n\n# Load our checkpoint\nnetwork = XMem(config, args.model).cuda().eval()\nif args.model is not None:\n model_weights = torch.load(args.model)\n network.load_weights(model_weights, init_as_zero_if_needed=True)\nelse:\n print('No model loaded.')\n\ntotal_process_time = 0\ntotal_frames = 0\n# Start eval\nfor vid_reader in tqdm(meta_loader, total=len(meta_dataset)):\n\n loader = DataLoader(vid_reader, batch_size=1, shuffle=False, num_workers=2)\n vid_name = vid_reader.vid_name\n vid_length = len(loader)\n # no need to count usage for LT if the video is not that long anyway\n config['enable_long_term_count_usage'] = (\n config['enable_long_term'] and\n (vid_length\n / (config['max_mid_term_frames']-config['min_mid_term_frames'])\n * config['num_prototypes'])\n >= config['max_long_term_elements']\n )\n\n mapper = MaskMapper()\n processor = InferenceCore(network, config=config)\n first_mask_loaded = False\n\n for ti, data in enumerate(loader):\n with torch.cuda.amp.autocast(enabled=not args.benchmark):\n rgb = data['rgb'].cuda()[0]\n msk = data.get('mask')\n info = data['info']\n frame = info['frame'][0]\n shape = info['shape']\n need_resize = info['need_resize'][0]\n\n \"\"\"\n For timing see https://discuss.pytorch.org/t/how-to-measure-time-in-pytorch/26964\n Seems to be very similar in testing as my previous timing method \n with two cuda sync + time.time() in STCN though \n \"\"\"\n start = torch.cuda.Event(enable_timing=True)\n end = torch.cuda.Event(enable_timing=True)\n start.record()\n\n if not first_mask_loaded:\n if msk is not None:\n first_mask_loaded = True\n else:\n # no point to do anything without a mask\n continue\n\n if args.flip:\n rgb = torch.flip(rgb, dims=[-1])\n msk = torch.flip(msk, dims=[-1]) if msk is not None else None\n\n # Map possibly non-continuous labels to continuous ones\n if msk is not None:\n msk, labels = mapper.convert_mask(msk[0].numpy())\n msk = torch.Tensor(msk).cuda()\n if need_resize:\n if msk.shape[0] == 0:\n print(vid_name)\n msk = vid_reader.resize_mask(msk.unsqueeze(0))[0]\n processor.set_all_labels(list(mapper.remappings.values()))\n else:\n labels = None\n\n # Run the model on this frame\n prob = processor.step(rgb, msk, labels, end=(ti==vid_length-1))\n\n # Upsample to original size if needed\n if need_resize:\n prob = F.interpolate(prob.unsqueeze(1), shape, mode='bilinear', align_corners=False)[:,0]\n\n end.record()\n torch.cuda.synchronize()\n total_process_time += (start.elapsed_time(end)/1000)\n total_frames += 1\n\n if args.flip:\n prob = torch.flip(prob, dims=[-1])\n\n # Probability mask -> index mask\n out_mask = torch.max(prob, dim=0).indices\n out_mask = (out_mask.detach().cpu().numpy()).astype(np.uint8)\n\n if args.save_scores:\n prob = (prob.detach().cpu().numpy()*255).astype(np.uint8)\n\n # Save the mask\n if args.save_all or info['save'][0]:\n this_out_path = path.join(out_path, vid_name)\n os.makedirs(this_out_path, exist_ok=True)\n out_mask = mapper.remap_index_mask(out_mask)\n out_img = Image.fromarray(out_mask)\n if vid_reader.get_palette() is not None:\n out_img.putpalette(vid_reader.get_palette())\n out_img.save(os.path.join(this_out_path, frame[:-4]+'.png'))\n\n if args.save_scores:\n np_path = path.join(args.output, 'Scores', vid_name)\n os.makedirs(np_path, exist_ok=True)\n if ti==len(loader)-1:\n hkl.dump(mapper.remappings, path.join(np_path, f'backward.hkl'), mode='w')\n if args.save_all or info['save'][0]:\n hkl.dump(prob, path.join(np_path, f'{frame[:-4]}.hkl'), mode='w', compression='lzf')\n\nprint(f'Total processing time: {total_process_time}')\nprint(f'Total processed frames: {total_frames}')\nprint(f'FPS: {total_frames / total_process_time}')\nprint(f'Max allocated memory (MB): {torch.cuda.max_memory_allocated() / (2**20)}')\n\nif not args.save_scores:\n if is_youtube:\n print('Making zip for YouTubeVOS...')\n shutil.make_archive(path.join(args.output, path.basename(args.output)), 'zip', args.output, 'Annotations')\n elif is_davis and args.split == 'test':\n print('Making zip for DAVIS test-dev...')\n shutil.make_archive(args.output, 'zip', args.output)\n" }, { "path": "XMem/eval_batch.py", "content": "import os\nimport time\nimport torch\nimport argparse\nimport multiprocessing as mp\nfrom termcolor import colored\nfrom datetime import datetime\nfrom importlib.util import find_spec\nif find_spec(\"GPUtil\") is None: os.system(\"pip install gputil\")\nimport GPUtil\n\n_GPU_LIST = [_.id for _ in GPUtil.getGPUs()]\n_GPU_QUEUE = mp.Queue()\nfor _ in _GPU_LIST: _GPU_QUEUE.put(_)\n\ndef run_eval(meta_expression, temp_xmem_anno, final_xmem_anno, img_dir, split_part, cfgs=\" --reversed \"):\n gpu_id = _GPU_QUEUE.get()\n cmd = f\"CUDA_VISIBLE_DEVICES={gpu_id} python eval.py --meta_exp {meta_expression} --output {final_xmem_anno} --generic_path {temp_xmem_anno} --img_dir {img_dir} --split_part {split_part} --dataset G {cfgs}\"\n print(f\"Running: {cmd}\")\n os.system(cmd)\n _GPU_QUEUE.put(gpu_id)\n\ndef main():\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--meta_expression\", type=str, help='/PATH/TO/ReVOS/meta_expressions_valid__llamavid.json')\n parser.add_argument(\"--temp_xmem_anno\", type=str, help='/PATH/TO/VISA_exp/revos_valid_XMem_temp/Annotations')\n parser.add_argument(\"--final_xmem_anno\", type=str, help='/PATH/TO/VISA_exp/revos_valid_XMem_final/Annotations')\n parser.add_argument(\"--img_dir\", type=str, help='/PATH/TO/ReVOS/JPEGImages')\n args = parser.parse_args()\n\n p = mp.Pool(8)\n for split_part in [0, 1, 2, 3]:\n for cfgs in [\" \", \" --reversed \"]:\n p.apply_async(\n run_eval,\n args=(args.meta_expression, args.temp_xmem_anno, args.final_xmem_anno, args.img_dir, split_part, cfgs),\n error_callback=lambda e: print(colored(e, \"red\"))\n )\n p.close()\n p.join()\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "XMem/generate_xmem_data_single.py", "content": "import sys\nimport os\nimport os.path as osp\nimport glob\nimport cv2\nimport multiprocessing\nimport json\nimport argparse\nfrom tqdm import tqdm\nfrom termcolor import colored\n\n\"\"\"\npython generate_xmem_data_single.py \\\n --video_root /PATH/TO/VISA_exp/revos_valid/Annotations \\\n --output_dir /PATH/TO/VISA_exp/revos_valid_XMem_temp/Annotations \\\n --final_xmem_anno /PATH/TO/VISA_exp/revos_valid_XMem_final/Annotations \\\n --llama_vid_meta /PATH/TO/ReVOS/meta_expressions_valid__llamavid.json\n\"\"\"\n\ndef generate(obj, temp_xmem_anno, final_xmem_anno):\n\n obj_dir, video_name, obj_id, tp = obj\n img_list = glob.glob(obj_dir + '/*.png') # Mask\n img_list.sort()\n frame_id = int(len(img_list) * tp)\n if frame_id == len(img_list):\n frame_id -= 1\n \n used_img = img_list[frame_id]\n\n img_output_path = osp.join(temp_xmem_anno, video_name, obj_id, osp.basename(used_img))\n final_img_output_dir = osp.join(final_xmem_anno, video_name, obj_id)\n img_output_dir = osp.dirname(img_output_path)\n os.makedirs(img_output_dir, exist_ok=True)\n os.makedirs(final_img_output_dir, exist_ok=True)\n os.system('cp {} {}'.format(used_img, img_output_path))\n\n img = cv2.imread(img_output_path)\n if img.sum() == 0:\n target_img_list = [i.split('/')[-1] for i in img_list]\n for img_ in target_img_list:\n print(os.path.join(final_img_output_dir, img_))\n os.system('cp {} {}'.format(img_output_path, os.path.join(img_output_dir, img_)))\n os.system('cp {} {}'.format(img_output_path, os.path.join(final_img_output_dir, img_)))\n return 0\n\ndef main():\n parser = argparse.ArgumentParser(description='rgvos')\n parser.add_argument('--video_root', type=str, help='/PATH/TO/VISA_exp/revos_valid/Annotations', )\n parser.add_argument('--temp_xmem_anno', type=str, help='/PATH/TO/VISA_exp/revos_valid_XMem_temp/Annotations', ) # ไฟๅญ˜ๅ•ๅธง Mask ็š„่ทฏๅพ„\n parser.add_argument('--final_xmem_anno', type=str, help='/PATH/TO/VISA_exp/revos_valid_XMem_final/Annotations', ) # ไฟๅญ˜ XMem ๆœ€ๅŽ่พ“ๅ‡บ็ป“ๆžœ็š„่ทฏๅพ„\n parser.add_argument(\"--llama_vid_meta\", type=str, help='/PATH/TO/ReVOS/meta_expressions_valid__llamavid.json', )\n args = parser.parse_args()\n\n video_root = args.video_root\n temp_xmem_anno = args.temp_xmem_anno\n final_xmem_anno = args.final_xmem_anno\n os.makedirs(temp_xmem_anno, exist_ok=True)\n\n data = json.load(open(args.llama_vid_meta, 'r'))['videos']\n\n all_obj_list = []\n for video_name in data.keys():\n exps = data[video_name]['expressions']\n for obj_id in exps.keys():\n tp = exps[obj_id]['tp']\n obj_dir = os.path.join(video_root, video_name, obj_id)\n all_obj_list.append([obj_dir, video_name, obj_id, tp])\n\n print('start')\n cpu_num = multiprocessing.cpu_count()-1\n print(\"cpu_num:\", cpu_num)\n pool = multiprocessing.Pool(cpu_num)\n pbar = tqdm(total=len(all_obj_list))\n\n for obj in all_obj_list:\n pool.apply_async(\n generate, \n args = (obj, temp_xmem_anno, final_xmem_anno ),\n callback = lambda *a: pbar.update(1),\n error_callback = lambda e: print(colored(e, \"red\"))\n )\n \n pool.close()\n pool.join()\n pbar.close()\n\n\nif __name__ == '__main__':\n main()" }, { "path": "XMem/inference/__init__.py", "content": "" }, { "path": "XMem/inference/data/__init__.py", "content": "" }, { "path": "XMem/inference/data/mask_mapper.py", "content": "import numpy as np\nimport torch\n\nfrom dataset.util import all_to_onehot\n\n\nclass MaskMapper:\n \"\"\"\n This class is used to convert a indexed-mask to a one-hot representation.\n It also takes care of remapping non-continuous indices\n It has two modes:\n 1. Default. Only masks with new indices are supposed to go into the remapper.\n This is also the case for YouTubeVOS.\n i.e., regions with index 0 are not \"background\", but \"don't care\".\n\n 2. Exhaustive. Regions with index 0 are considered \"background\".\n Every single pixel is considered to be \"labeled\".\n \"\"\"\n def __init__(self):\n self.labels = []\n self.remappings = {}\n\n # if coherent, no mapping is required\n self.coherent = True\n\n def convert_mask(self, mask, exhaustive=False):\n # mask is in index representation, H*W numpy array\n labels = np.unique(mask).astype(np.uint8)\n labels = labels[labels!=0].tolist()\n\n\n new_labels = list(set(labels) - set(self.labels))\n if not exhaustive:\n assert len(new_labels) == len(labels), 'Old labels found in non-exhaustive mode'\n\n # add new remappings\n for i, l in enumerate(new_labels):\n self.remappings[l] = i+len(self.labels)+1\n if self.coherent and i+len(self.labels)+1 != l:\n self.coherent = False\n\n if exhaustive:\n new_mapped_labels = range(1, len(self.labels)+len(new_labels)+1)\n else:\n if self.coherent:\n new_mapped_labels = new_labels\n else:\n new_mapped_labels = range(len(self.labels)+1, len(self.labels)+len(new_labels)+1)\n\n self.labels.extend(new_labels)\n mask = torch.from_numpy(all_to_onehot(mask, self.labels)).float()\n\n # mask num_objects*H*W\n return mask, new_mapped_labels\n\n\n def remap_index_mask(self, mask):\n # mask is in index representation, H*W numpy array\n if self.coherent:\n return mask\n\n new_mask = np.zeros_like(mask)\n for l, i in self.remappings.items():\n new_mask[mask==i] = l\n return new_mask" }, { "path": "XMem/inference/data/test_datasets.py", "content": "import os\nfrom os import path\nimport json\nimport glob\n\nfrom inference.data.video_reader import VideoReader\n\n\nclass LongTestDataset:\n\n def __init__(self, meta_expression, data_root, size=-1, img_dir = '', reversed_ = False, split_part = 0):\n self.image_dir = img_dir\n self.mask_dir = data_root\n self.size = size\n self.reversed = reversed_\n self.split_part = split_part\n\n self.vid_list = []\n\n videos_names = json.load(open(meta_expression, 'r'))['videos']\n for video_name in videos_names:\n video_mask_dir = path.join(self.mask_dir, video_name)\n obj_ids = [d for d in os.listdir(video_mask_dir) if os.path.isdir(path.join(video_mask_dir, d))]\n for obj_id in obj_ids:\n obj_dir = path.join(video_mask_dir, obj_id)\n img_list = glob.glob(obj_dir + '/*')\n if len(img_list) == 1:\n self.vid_list.append(path.join(video_name, obj_id))\n \n\n self.vid_list.sort()\n self.vid_list = [i for idx, i in enumerate(self.vid_list) if idx % 4 == self.split_part]\n\n def get_datasets(self):\n for video in self.vid_list:\n yield VideoReader(video, \n path.join(self.image_dir, '/'.join(video.split('/')[:-1])), \n path.join(self.mask_dir, video),\n to_save = [\n name[:-4] for name in os.listdir(path.join(self.mask_dir, video)) # remove .png\n ],\n size=self.size,\n reversed=self.reversed,\n )\n\n def __len__(self):\n return len(self.vid_list)\n\n\nclass DAVISTestDataset:\n def __init__(self, data_root, imset='2017/val.txt', size=-1):\n if size != 480:\n self.image_dir = path.join(data_root, 'JPEGImages', 'Full-Resolution')\n self.mask_dir = path.join(data_root, 'Annotations', 'Full-Resolution')\n if not path.exists(self.image_dir):\n print(f'{self.image_dir} not found. Look at other options.')\n self.image_dir = path.join(data_root, 'JPEGImages', '1080p')\n self.mask_dir = path.join(data_root, 'Annotations', '1080p')\n assert path.exists(self.image_dir), 'path not found'\n else:\n self.image_dir = path.join(data_root, 'JPEGImages', '480p')\n self.mask_dir = path.join(data_root, 'Annotations', '480p')\n self.size_dir = path.join(data_root, 'JPEGImages', '480p')\n self.size = size\n\n with open(path.join(data_root, 'ImageSets', imset)) as f:\n self.vid_list = sorted([line.strip() for line in f])\n\n def get_datasets(self):\n for video in self.vid_list:\n yield VideoReader(video, \n path.join(self.image_dir, video), \n path.join(self.mask_dir, video),\n size=self.size,\n size_dir=path.join(self.size_dir, video),\n )\n\n def __len__(self):\n return len(self.vid_list)\n\n\nclass YouTubeVOSTestDataset:\n def __init__(self, data_root, split, size=480):\n self.image_dir = path.join(data_root, 'all_frames', split+'_all_frames', 'JPEGImages')\n self.mask_dir = path.join(data_root, split, 'Annotations')\n self.size = size\n\n self.vid_list = sorted(os.listdir(self.image_dir))\n self.req_frame_list = {}\n\n with open(path.join(data_root, split, 'meta.json')) as f:\n # read meta.json to know which frame is required for evaluation\n meta = json.load(f)['videos']\n\n for vid in self.vid_list:\n req_frames = []\n objects = meta[vid]['objects']\n for value in objects.values():\n req_frames.extend(value['frames'])\n\n req_frames = list(set(req_frames))\n self.req_frame_list[vid] = req_frames\n\n def get_datasets(self):\n for video in self.vid_list:\n yield VideoReader(video, \n path.join(self.image_dir, video), \n path.join(self.mask_dir, video),\n size=self.size,\n to_save=self.req_frame_list[video], \n use_all_mask=True\n )\n\n def __len__(self):\n return len(self.vid_list)\n" }, { "path": "XMem/inference/data/video_reader.py", "content": "import os\nfrom os import path\n\nfrom torch.utils.data.dataset import Dataset\nfrom torchvision import transforms\nfrom torchvision.transforms import InterpolationMode\nimport torch.nn.functional as F\nfrom PIL import Image\nimport numpy as np\n\nfrom dataset.range_transform import im_normalization\n\n\nclass VideoReader(Dataset):\n \"\"\"\n This class is used to read a video, one frame at a time\n \"\"\"\n def __init__(self, vid_name, image_dir, mask_dir, size=-1, to_save=None, use_all_mask=False, size_dir=None, reversed = False):\n \"\"\"\n image_dir - points to a directory of jpg images\n mask_dir - points to a directory of png masks\n size - resize min. side to size. Does nothing if <0.\n to_save - optionally contains a list of file names without extensions \n where the segmentation mask is required\n use_all_mask - when true, read all available mask in mask_dir.\n Default false. Set to true for YouTubeVOS validation.\n \"\"\"\n self.vid_name = vid_name\n self.image_dir = image_dir\n self.mask_dir = mask_dir\n self.to_save = to_save\n self.use_all_mask = use_all_mask\n self.reversed = reversed\n if size_dir is None:\n self.size_dir = self.image_dir\n else:\n self.size_dir = size_dir\n\n self.frames = sorted(os.listdir(self.image_dir))\n if self.reversed:\n self.frames = self.frames[::-1]\n self.palette = Image.open(path.join(mask_dir, sorted(os.listdir(mask_dir))[0])).getpalette()\n self.first_gt_path = path.join(self.mask_dir, sorted(os.listdir(self.mask_dir))[0])\n\n if size < 0:\n self.im_transform = transforms.Compose([\n transforms.ToTensor(),\n im_normalization,\n ])\n else:\n self.im_transform = transforms.Compose([\n transforms.ToTensor(),\n im_normalization,\n transforms.Resize(size, interpolation=InterpolationMode.BILINEAR),\n ])\n self.size = size\n\n\n def __getitem__(self, idx):\n frame = self.frames[idx]\n info = {}\n data = {}\n info['frame'] = frame\n info['save'] = (self.to_save is None) or (frame[:-4] in self.to_save)\n\n im_path = path.join(self.image_dir, frame)\n img = Image.open(im_path).convert('RGB')\n\n if self.image_dir == self.size_dir:\n shape = np.array(img).shape[:2]\n else:\n size_path = path.join(self.size_dir, frame)\n size_im = Image.open(size_path).convert('RGB')\n shape = np.array(size_im).shape[:2]\n\n gt_path = path.join(self.mask_dir, frame[:-4]+'.png')\n img = self.im_transform(img)\n\n load_mask = self.use_all_mask or (gt_path == self.first_gt_path)\n if load_mask and path.exists(gt_path):\n mask = Image.open(gt_path).convert('P')\n mask = np.array(mask, dtype=np.uint8)\n data['mask'] = mask\n\n info['shape'] = shape\n info['need_resize'] = not (self.size < 0)\n data['rgb'] = img\n data['info'] = info\n\n return data\n\n def resize_mask(self, mask):\n # mask transform is applied AFTER mapper, so we need to post-process it in eval.py\n h, w = mask.shape[-2:]\n min_hw = min(h, w)\n return F.interpolate(mask, (int(h/min_hw*self.size), int(w/min_hw*self.size)), \n mode='nearest')\n\n def get_palette(self):\n return self.palette\n\n def __len__(self):\n return len(self.frames)\n" }, { "path": "XMem/inference/inference_core.py", "content": "from inference.memory_manager import MemoryManager\nfrom model.network import XMem\nfrom model.aggregate import aggregate\n\nfrom util.tensor_util import pad_divide_by, unpad\n\n\nclass InferenceCore:\n def __init__(self, network:XMem, config):\n self.config = config\n self.network = network\n self.mem_every = config['mem_every']\n self.deep_update_every = config['deep_update_every']\n self.enable_long_term = config['enable_long_term']\n\n # if deep_update_every < 0, synchronize deep update with memory frame\n self.deep_update_sync = (self.deep_update_every < 0)\n\n self.clear_memory()\n self.all_labels = None\n\n def clear_memory(self):\n self.curr_ti = -1\n self.last_mem_ti = 0\n if not self.deep_update_sync:\n self.last_deep_update_ti = -self.deep_update_every\n self.memory = MemoryManager(config=self.config)\n\n def update_config(self, config):\n self.mem_every = config['mem_every']\n self.deep_update_every = config['deep_update_every']\n self.enable_long_term = config['enable_long_term']\n\n # if deep_update_every < 0, synchronize deep update with memory frame\n self.deep_update_sync = (self.deep_update_every < 0)\n self.memory.update_config(config)\n\n def set_all_labels(self, all_labels):\n # self.all_labels = [l.item() for l in all_labels]\n self.all_labels = all_labels\n\n def step(self, image, mask=None, valid_labels=None, end=False):\n # image: 3*H*W\n # mask: num_objects*H*W or None\n self.curr_ti += 1\n image, self.pad = pad_divide_by(image, 16)\n image = image.unsqueeze(0) # add the batch dimension\n\n is_mem_frame = ((self.curr_ti-self.last_mem_ti >= self.mem_every) or (mask is not None)) and (not end)\n need_segment = (self.curr_ti > 0) and ((valid_labels is None) or (len(self.all_labels) != len(valid_labels)))\n is_deep_update = (\n (self.deep_update_sync and is_mem_frame) or # synchronized\n (not self.deep_update_sync and self.curr_ti-self.last_deep_update_ti >= self.deep_update_every) # no-sync\n ) and (not end)\n is_normal_update = (not self.deep_update_sync or not is_deep_update) and (not end)\n\n key, shrinkage, selection, f16, f8, f4 = self.network.encode_key(image, \n need_ek=(self.enable_long_term or need_segment), \n need_sk=is_mem_frame)\n multi_scale_features = (f16, f8, f4)\n\n # segment the current frame is needed\n if need_segment:\n memory_readout = self.memory.match_memory(key, selection).unsqueeze(0)\n hidden, _, pred_prob_with_bg = self.network.segment(multi_scale_features, memory_readout, \n self.memory.get_hidden(), h_out=is_normal_update, strip_bg=False)\n # remove batch dim\n pred_prob_with_bg = pred_prob_with_bg[0]\n pred_prob_no_bg = pred_prob_with_bg[1:]\n if is_normal_update:\n self.memory.set_hidden(hidden)\n else:\n pred_prob_no_bg = pred_prob_with_bg = None\n\n # use the input mask if any\n if mask is not None:\n mask, _ = pad_divide_by(mask, 16)\n\n if pred_prob_no_bg is not None:\n # if we have a predicted mask, we work on it\n # make pred_prob_no_bg consistent with the input mask\n mask_regions = (mask.sum(0) > 0.5)\n pred_prob_no_bg[:, mask_regions] = 0\n # shift by 1 because mask/pred_prob_no_bg do not contain background\n mask = mask.type_as(pred_prob_no_bg)\n if valid_labels is not None:\n shift_by_one_non_labels = [i for i in range(pred_prob_no_bg.shape[0]) if (i+1) not in valid_labels]\n # non-labelled objects are copied from the predicted mask\n mask[shift_by_one_non_labels] = pred_prob_no_bg[shift_by_one_non_labels]\n pred_prob_with_bg = aggregate(mask, dim=0)\n\n # also create new hidden states\n self.memory.create_hidden_state(len(self.all_labels), key)\n\n # save as memory if needed\n if is_mem_frame:\n value, hidden = self.network.encode_value(image, f16, self.memory.get_hidden(), \n pred_prob_with_bg[1:].unsqueeze(0), is_deep_update=is_deep_update)\n self.memory.add_memory(key, shrinkage, value, self.all_labels, \n selection=selection if self.enable_long_term else None)\n self.last_mem_ti = self.curr_ti\n\n if is_deep_update:\n self.memory.set_hidden(hidden)\n self.last_deep_update_ti = self.curr_ti\n \n return unpad(pred_prob_with_bg, self.pad)\n" }, { "path": "XMem/inference/interact/__init__.py", "content": "" }, { "path": "XMem/inference/interact/fbrs/LICENSE", "content": "Mozilla Public License Version 2.0\n==================================\n\n1. Definitions\n--------------\n\n1.1. \"Contributor\"\n means each individual or legal entity that creates, contributes to\n the creation of, or owns Covered Software.\n\n1.2. \"Contributor Version\"\n means the combination of the Contributions of others (if any) used\n by a Contributor and that particular Contributor's Contribution.\n\n1.3. \"Contribution\"\n means Covered Software of a particular Contributor.\n\n1.4. \"Covered Software\"\n means Source Code Form to which the initial Contributor has attached\n the notice in Exhibit A, the Executable Form of such Source Code\n Form, and Modifications of such Source Code Form, in each case\n including portions thereof.\n\n1.5. \"Incompatible With Secondary Licenses\"\n means\n\n (a) that the initial Contributor has attached the notice described\n in Exhibit B to the Covered Software; or\n\n (b) that the Covered Software was made available under the terms of\n version 1.1 or earlier of the License, but not also under the\n terms of a Secondary License.\n\n1.6. \"Executable Form\"\n means any form of the work other than Source Code Form.\n\n1.7. \"Larger Work\"\n means a work that combines Covered Software with other material, in\n a separate file or files, that is not Covered Software.\n\n1.8. \"License\"\n means this document.\n\n1.9. \"Licensable\"\n means having the right to grant, to the maximum extent possible,\n whether at the time of the initial grant or subsequently, any and\n all of the rights conveyed by this License.\n\n1.10. \"Modifications\"\n means any of the following:\n\n (a) any file in Source Code Form that results from an addition to,\n deletion from, or modification of the contents of Covered\n Software; or\n\n (b) any new file in Source Code Form that contains any Covered\n Software.\n\n1.11. \"Patent Claims\" of a Contributor\n means any patent claim(s), including without limitation, method,\n process, and apparatus claims, in any patent Licensable by such\n Contributor that would be infringed, but for the grant of the\n License, by the making, using, selling, offering for sale, having\n made, import, or transfer of either its Contributions or its\n Contributor Version.\n\n1.12. \"Secondary License\"\n means either the GNU General Public License, Version 2.0, the GNU\n Lesser General Public License, Version 2.1, the GNU Affero General\n Public License, Version 3.0, or any later versions of those\n licenses.\n\n1.13. \"Source Code Form\"\n means the form of the work preferred for making modifications.\n\n1.14. \"You\" (or \"Your\")\n means an individual or a legal entity exercising rights under this\n License. For legal entities, \"You\" includes any entity that\n controls, is controlled by, or is under common control with You. For\n purposes of this definition, \"control\" means (a) the power, direct\n or indirect, to cause the direction or management of such entity,\n whether by contract or otherwise, or (b) ownership of more than\n fifty percent (50%) of the outstanding shares or beneficial\n ownership of such entity.\n\n2. License Grants and Conditions\n--------------------------------\n\n2.1. Grants\n\nEach Contributor hereby grants You a world-wide, royalty-free,\nnon-exclusive license:\n\n(a) under intellectual property rights (other than patent or trademark)\n Licensable by such Contributor to use, reproduce, make available,\n modify, display, perform, distribute, and otherwise exploit its\n Contributions, either on an unmodified basis, with Modifications, or\n as part of a Larger Work; and\n\n(b) under Patent Claims of such Contributor to make, use, sell, offer\n for sale, have made, import, and otherwise transfer either its\n Contributions or its Contributor Version.\n\n2.2. Effective Date\n\nThe licenses granted in Section 2.1 with respect to any Contribution\nbecome effective for each Contribution on the date the Contributor first\ndistributes such Contribution.\n\n2.3. Limitations on Grant Scope\n\nThe licenses granted in this Section 2 are the only rights granted under\nthis License. No additional rights or licenses will be implied from the\ndistribution or licensing of Covered Software under this License.\nNotwithstanding Section 2.1(b) above, no patent license is granted by a\nContributor:\n\n(a) for any code that a Contributor has removed from Covered Software;\n or\n\n(b) for infringements caused by: (i) Your and any other third party's\n modifications of Covered Software, or (ii) the combination of its\n Contributions with other software (except as part of its Contributor\n Version); or\n\n(c) under Patent Claims infringed by Covered Software in the absence of\n its Contributions.\n\nThis License does not grant any rights in the trademarks, service marks,\nor logos of any Contributor (except as may be necessary to comply with\nthe notice requirements in Section 3.4).\n\n2.4. Subsequent Licenses\n\nNo Contributor makes additional grants as a result of Your choice to\ndistribute the Covered Software under a subsequent version of this\nLicense (see Section 10.2) or under the terms of a Secondary License (if\npermitted under the terms of Section 3.3).\n\n2.5. Representation\n\nEach Contributor represents that the Contributor believes its\nContributions are its original creation(s) or it has sufficient rights\nto grant the rights to its Contributions conveyed by this License.\n\n2.6. Fair Use\n\nThis License is not intended to limit any rights You have under\napplicable copyright doctrines of fair use, fair dealing, or other\nequivalents.\n\n2.7. Conditions\n\nSections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted\nin Section 2.1.\n\n3. Responsibilities\n-------------------\n\n3.1. Distribution of Source Form\n\nAll distribution of Covered Software in Source Code Form, including any\nModifications that You create or to which You contribute, must be under\nthe terms of this License. You must inform recipients that the Source\nCode Form of the Covered Software is governed by the terms of this\nLicense, and how they can obtain a copy of this License. You may not\nattempt to alter or restrict the recipients' rights in the Source Code\nForm.\n\n3.2. Distribution of Executable Form\n\nIf You distribute Covered Software in Executable Form then:\n\n(a) such Covered Software must also be made available in Source Code\n Form, as described in Section 3.1, and You must inform recipients of\n the Executable Form how they can obtain a copy of such Source Code\n Form by reasonable means in a timely manner, at a charge no more\n than the cost of distribution to the recipient; and\n\n(b) You may distribute such Executable Form under the terms of this\n License, or sublicense it under different terms, provided that the\n license for the Executable Form does not attempt to limit or alter\n the recipients' rights in the Source Code Form under this License.\n\n3.3. Distribution of a Larger Work\n\nYou may create and distribute a Larger Work under terms of Your choice,\nprovided that You also comply with the requirements of this License for\nthe Covered Software. If the Larger Work is a combination of Covered\nSoftware with a work governed by one or more Secondary Licenses, and the\nCovered Software is not Incompatible With Secondary Licenses, this\nLicense permits You to additionally distribute such Covered Software\nunder the terms of such Secondary License(s), so that the recipient of\nthe Larger Work may, at their option, further distribute the Covered\nSoftware under the terms of either this License or such Secondary\nLicense(s).\n\n3.4. Notices\n\nYou may not remove or alter the substance of any license notices\n(including copyright notices, patent notices, disclaimers of warranty,\nor limitations of liability) contained within the Source Code Form of\nthe Covered Software, except that You may alter any license notices to\nthe extent required to remedy known factual inaccuracies.\n\n3.5. Application of Additional Terms\n\nYou may choose to offer, and to charge a fee for, warranty, support,\nindemnity or liability obligations to one or more recipients of Covered\nSoftware. However, You may do so only on Your own behalf, and not on\nbehalf of any Contributor. You must make it absolutely clear that any\nsuch warranty, support, indemnity, or liability obligation is offered by\nYou alone, and You hereby agree to indemnify every Contributor for any\nliability incurred by such Contributor as a result of warranty, support,\nindemnity or liability terms You offer. You may include additional\ndisclaimers of warranty and limitations of liability specific to any\njurisdiction.\n\n4. Inability to Comply Due to Statute or Regulation\n---------------------------------------------------\n\nIf it is impossible for You to comply with any of the terms of this\nLicense with respect to some or all of the Covered Software due to\nstatute, judicial order, or regulation then You must: (a) comply with\nthe terms of this License to the maximum extent possible; and (b)\ndescribe the limitations and the code they affect. Such description must\nbe placed in a text file included with all distributions of the Covered\nSoftware under this License. Except to the extent prohibited by statute\nor regulation, such description must be sufficiently detailed for a\nrecipient of ordinary skill to be able to understand it.\n\n5. Termination\n--------------\n\n5.1. The rights granted under this License will terminate automatically\nif You fail to comply with any of its terms. However, if You become\ncompliant, then the rights granted under this License from a particular\nContributor are reinstated (a) provisionally, unless and until such\nContributor explicitly and finally terminates Your grants, and (b) on an\nongoing basis, if such Contributor fails to notify You of the\nnon-compliance by some reasonable means prior to 60 days after You have\ncome back into compliance. Moreover, Your grants from a particular\nContributor are reinstated on an ongoing basis if such Contributor\nnotifies You of the non-compliance by some reasonable means, this is the\nfirst time You have received notice of non-compliance with this License\nfrom such Contributor, and You become compliant prior to 30 days after\nYour receipt of the notice.\n\n5.2. If You initiate litigation against any entity by asserting a patent\ninfringement claim (excluding declaratory judgment actions,\ncounter-claims, and cross-claims) alleging that a Contributor Version\ndirectly or indirectly infringes any patent, then the rights granted to\nYou by any and all Contributors for the Covered Software under Section\n2.1 of this License shall terminate.\n\n5.3. In the event of termination under Sections 5.1 or 5.2 above, all\nend user license agreements (excluding distributors and resellers) which\nhave been validly granted by You or Your distributors under this License\nprior to termination shall survive termination.\n\n************************************************************************\n* *\n* 6. Disclaimer of Warranty *\n* ------------------------- *\n* *\n* Covered Software is provided under this License on an \"as is\" *\n* basis, without warranty of any kind, either expressed, implied, or *\n* statutory, including, without limitation, warranties that the *\n* Covered Software is free of defects, merchantable, fit for a *\n* particular purpose or non-infringing. The entire risk as to the *\n* quality and performance of the Covered Software is with You. *\n* Should any Covered Software prove defective in any respect, You *\n* (not any Contributor) assume the cost of any necessary servicing, *\n* repair, or correction. This disclaimer of warranty constitutes an *\n* essential part of this License. No use of any Covered Software is *\n* authorized under this License except under this disclaimer. *\n* *\n************************************************************************\n\n************************************************************************\n* *\n* 7. Limitation of Liability *\n* -------------------------- *\n* *\n* Under no circumstances and under no legal theory, whether tort *\n* (including negligence), contract, or otherwise, shall any *\n* Contributor, or anyone who distributes Covered Software as *\n* permitted above, be liable to You for any direct, indirect, *\n* special, incidental, or consequential damages of any character *\n* including, without limitation, damages for lost profits, loss of *\n* goodwill, work stoppage, computer failure or malfunction, or any *\n* and all other commercial damages or losses, even if such party *\n* shall have been informed of the possibility of such damages. This *\n* limitation of liability shall not apply to liability for death or *\n* personal injury resulting from such party's negligence to the *\n* extent applicable law prohibits such limitation. Some *\n* jurisdictions do not allow the exclusion or limitation of *\n* incidental or consequential damages, so this exclusion and *\n* limitation may not apply to You. *\n* *\n************************************************************************\n\n8. Litigation\n-------------\n\nAny litigation relating to this License may be brought only in the\ncourts of a jurisdiction where the defendant maintains its principal\nplace of business and such litigation shall be governed by laws of that\njurisdiction, without reference to its conflict-of-law provisions.\nNothing in this Section shall prevent a party's ability to bring\ncross-claims or counter-claims.\n\n9. Miscellaneous\n----------------\n\nThis License represents the complete agreement concerning the subject\nmatter hereof. If any provision of this License is held to be\nunenforceable, such provision shall be reformed only to the extent\nnecessary to make it enforceable. Any law or regulation which provides\nthat the language of a contract shall be construed against the drafter\nshall not be used to construe this License against a Contributor.\n\n10. Versions of the License\n---------------------------\n\n10.1. New Versions\n\nMozilla Foundation is the license steward. Except as provided in Section\n10.3, no one other than the license steward has the right to modify or\npublish new versions of this License. Each version will be given a\ndistinguishing version number.\n\n10.2. Effect of New Versions\n\nYou may distribute the Covered Software under the terms of the version\nof the License under which You originally received the Covered Software,\nor under the terms of any subsequent version published by the license\nsteward.\n\n10.3. Modified Versions\n\nIf you create software not governed by this License, and you want to\ncreate a new license for such software, you may create and use a\nmodified version of this License if you rename the license and remove\nany references to the name of the license steward (except to note that\nsuch modified license differs from this License).\n\n10.4. Distributing Source Code Form that is Incompatible With Secondary\nLicenses\n\nIf You choose to distribute Source Code Form that is Incompatible With\nSecondary Licenses under the terms of this version of the License, the\nnotice described in Exhibit B of this License must be attached.\n\nExhibit A - Source Code Form License Notice\n-------------------------------------------\n\n This Source Code Form is subject to the terms of the Mozilla Public\n License, v. 2.0. If a copy of the MPL was not distributed with this\n file, You can obtain one at http://mozilla.org/MPL/2.0/.\n\nIf it is not possible or desirable to put the notice in a particular\nfile, then You may include the notice in a location (such as a LICENSE\nfile in a relevant directory) where a recipient would be likely to look\nfor such a notice.\n\nYou may add additional accurate notices of copyright ownership.\n\nExhibit B - \"Incompatible With Secondary Licenses\" Notice\n---------------------------------------------------------\n\n This Source Code Form is \"Incompatible With Secondary Licenses\", as\n defined by the Mozilla Public License, v. 2.0." }, { "path": "XMem/inference/interact/fbrs/__init__.py", "content": "" }, { "path": "XMem/inference/interact/fbrs/controller.py", "content": "import torch\ntry:\n from torch import mps\nexcept:\n pass\n\nfrom ..fbrs.inference import clicker\nfrom ..fbrs.inference.predictors import get_predictor\n\n\nclass InteractiveController:\n def __init__(self, net, device, predictor_params, prob_thresh=0.5):\n self.net = net.to(device)\n self.prob_thresh = prob_thresh\n self.clicker = clicker.Clicker()\n self.states = []\n self.probs_history = []\n self.object_count = 0\n self._result_mask = None\n\n self.image = None\n self.predictor = None\n self.device = device\n self.predictor_params = predictor_params\n self.reset_predictor()\n\n def set_image(self, image):\n self.image = image\n self._result_mask = torch.zeros(image.shape[-2:], dtype=torch.uint8)\n self.object_count = 0\n self.reset_last_object()\n\n def add_click(self, x, y, is_positive):\n self.states.append({\n 'clicker': self.clicker.get_state(),\n 'predictor': self.predictor.get_states()\n })\n\n click = clicker.Click(is_positive=is_positive, coords=(y, x))\n self.clicker.add_click(click)\n pred = self.predictor.get_prediction(self.clicker)\n if self.device.type == 'cuda':\n torch.cuda.empty_cache()\n elif self.device.type == 'mps':\n mps.empty_cache()\n\n if self.probs_history:\n self.probs_history.append((self.probs_history[-1][0], pred))\n else:\n self.probs_history.append((torch.zeros_like(pred), pred))\n\n def undo_click(self):\n if not self.states:\n return\n\n prev_state = self.states.pop()\n self.clicker.set_state(prev_state['clicker'])\n self.predictor.set_states(prev_state['predictor'])\n self.probs_history.pop()\n\n def partially_finish_object(self):\n object_prob = self.current_object_prob\n if object_prob is None:\n return\n\n self.probs_history.append((object_prob, torch.zeros_like(object_prob)))\n self.states.append(self.states[-1])\n\n self.clicker.reset_clicks()\n self.reset_predictor()\n\n def finish_object(self):\n object_prob = self.current_object_prob\n if object_prob is None:\n return\n\n self.object_count += 1\n object_mask = object_prob > self.prob_thresh\n self._result_mask[object_mask] = self.object_count\n self.reset_last_object()\n\n def reset_last_object(self):\n self.states = []\n self.probs_history = []\n self.clicker.reset_clicks()\n self.reset_predictor()\n\n def reset_predictor(self, predictor_params=None):\n if predictor_params is not None:\n self.predictor_params = predictor_params\n self.predictor = get_predictor(self.net, device=self.device,\n **self.predictor_params)\n if self.image is not None:\n self.predictor.set_input_image(self.image)\n\n @property\n def current_object_prob(self):\n if self.probs_history:\n current_prob_total, current_prob_additive = self.probs_history[-1]\n return torch.maximum(current_prob_total, current_prob_additive)\n else:\n return None\n\n @property\n def is_incomplete_mask(self):\n return len(self.probs_history) > 0\n\n @property\n def result_mask(self):\n return self._result_mask.clone()\n" }, { "path": "XMem/inference/interact/fbrs/inference/__init__.py", "content": "" }, { "path": "XMem/inference/interact/fbrs/inference/clicker.py", "content": "from collections import namedtuple\n\nimport numpy as np\nfrom copy import deepcopy\nfrom scipy.ndimage import distance_transform_edt\n\nClick = namedtuple('Click', ['is_positive', 'coords'])\n\n\nclass Clicker(object):\n def __init__(self, gt_mask=None, init_clicks=None, ignore_label=-1):\n if gt_mask is not None:\n self.gt_mask = gt_mask == 1\n self.not_ignore_mask = gt_mask != ignore_label\n else:\n self.gt_mask = None\n\n self.reset_clicks()\n\n if init_clicks is not None:\n for click in init_clicks:\n self.add_click(click)\n\n def make_next_click(self, pred_mask):\n assert self.gt_mask is not None\n click = self._get_click(pred_mask)\n self.add_click(click)\n\n def get_clicks(self, clicks_limit=None):\n return self.clicks_list[:clicks_limit]\n\n def _get_click(self, pred_mask, padding=True):\n fn_mask = np.logical_and(np.logical_and(self.gt_mask, np.logical_not(pred_mask)), self.not_ignore_mask)\n fp_mask = np.logical_and(np.logical_and(np.logical_not(self.gt_mask), pred_mask), self.not_ignore_mask)\n\n if padding:\n fn_mask = np.pad(fn_mask, ((1, 1), (1, 1)), 'constant')\n fp_mask = np.pad(fp_mask, ((1, 1), (1, 1)), 'constant')\n\n fn_mask_dt = distance_transform_edt(fn_mask)\n fp_mask_dt = distance_transform_edt(fp_mask)\n\n if padding:\n fn_mask_dt = fn_mask_dt[1:-1, 1:-1]\n fp_mask_dt = fp_mask_dt[1:-1, 1:-1]\n\n fn_mask_dt = fn_mask_dt * self.not_clicked_map\n fp_mask_dt = fp_mask_dt * self.not_clicked_map\n\n fn_max_dist = np.max(fn_mask_dt)\n fp_max_dist = np.max(fp_mask_dt)\n\n is_positive = fn_max_dist > fp_max_dist\n if is_positive:\n coords_y, coords_x = np.where(fn_mask_dt == fn_max_dist) # coords is [y, x]\n else:\n coords_y, coords_x = np.where(fp_mask_dt == fp_max_dist) # coords is [y, x]\n\n return Click(is_positive=is_positive, coords=(coords_y[0], coords_x[0]))\n\n def add_click(self, click):\n coords = click.coords\n\n if click.is_positive:\n self.num_pos_clicks += 1\n else:\n self.num_neg_clicks += 1\n\n self.clicks_list.append(click)\n if self.gt_mask is not None:\n self.not_clicked_map[coords[0], coords[1]] = False\n\n def _remove_last_click(self):\n click = self.clicks_list.pop()\n coords = click.coords\n\n if click.is_positive:\n self.num_pos_clicks -= 1\n else:\n self.num_neg_clicks -= 1\n\n if self.gt_mask is not None:\n self.not_clicked_map[coords[0], coords[1]] = True\n\n def reset_clicks(self):\n if self.gt_mask is not None:\n self.not_clicked_map = np.ones_like(self.gt_mask, dtype=np.bool)\n\n self.num_pos_clicks = 0\n self.num_neg_clicks = 0\n\n self.clicks_list = []\n\n def get_state(self):\n return deepcopy(self.clicks_list)\n\n def set_state(self, state):\n self.reset_clicks()\n for click in state:\n self.add_click(click)\n\n def __len__(self):\n return len(self.clicks_list)\n" }, { "path": "XMem/inference/interact/fbrs/inference/evaluation.py", "content": "from time import time\n\nimport numpy as np\nimport torch\n\nfrom ..inference import utils\nfrom ..inference.clicker import Clicker\n\ntry:\n get_ipython()\n from tqdm import tqdm_notebook as tqdm\nexcept NameError:\n from tqdm import tqdm\n\n\ndef evaluate_dataset(dataset, predictor, oracle_eval=False, **kwargs):\n all_ious = []\n\n start_time = time()\n for index in tqdm(range(len(dataset)), leave=False):\n sample = dataset.get_sample(index)\n item = dataset[index]\n\n if oracle_eval:\n gt_mask = torch.tensor(sample['instances_mask'], dtype=torch.float32)\n gt_mask = gt_mask.unsqueeze(0).unsqueeze(0)\n predictor.opt_functor.mask_loss.set_gt_mask(gt_mask)\n _, sample_ious, _ = evaluate_sample(item['images'], sample['instances_mask'], predictor, **kwargs)\n all_ious.append(sample_ious)\n end_time = time()\n elapsed_time = end_time - start_time\n\n return all_ious, elapsed_time\n\n\ndef evaluate_sample(image_nd, instances_mask, predictor, max_iou_thr,\n pred_thr=0.49, max_clicks=20):\n clicker = Clicker(gt_mask=instances_mask)\n pred_mask = np.zeros_like(instances_mask)\n ious_list = []\n\n with torch.no_grad():\n predictor.set_input_image(image_nd)\n\n for click_number in range(max_clicks):\n clicker.make_next_click(pred_mask)\n pred_probs = predictor.get_prediction(clicker)\n pred_mask = pred_probs > pred_thr\n\n iou = utils.get_iou(instances_mask, pred_mask)\n ious_list.append(iou)\n\n if iou >= max_iou_thr:\n break\n\n return clicker.clicks_list, np.array(ious_list, dtype=np.float32), pred_probs\n" }, { "path": "XMem/inference/interact/fbrs/inference/predictors/__init__.py", "content": "from .base import BasePredictor\nfrom .brs import InputBRSPredictor, FeatureBRSPredictor, HRNetFeatureBRSPredictor\nfrom .brs_functors import InputOptimizer, ScaleBiasOptimizer\nfrom ..transforms import ZoomIn\nfrom ...model.is_hrnet_model import DistMapsHRNetModel\n\n\ndef get_predictor(net, brs_mode, device,\n prob_thresh=0.49,\n with_flip=True,\n zoom_in_params=dict(),\n predictor_params=None,\n brs_opt_func_params=None,\n lbfgs_params=None):\n lbfgs_params_ = {\n 'm': 20,\n 'factr': 0,\n 'pgtol': 1e-8,\n 'maxfun': 20,\n }\n\n predictor_params_ = {\n 'optimize_after_n_clicks': 1\n }\n\n if zoom_in_params is not None:\n zoom_in = ZoomIn(**zoom_in_params)\n else:\n zoom_in = None\n\n if lbfgs_params is not None:\n lbfgs_params_.update(lbfgs_params)\n lbfgs_params_['maxiter'] = 2 * lbfgs_params_['maxfun']\n\n if brs_opt_func_params is None:\n brs_opt_func_params = dict()\n\n if brs_mode == 'NoBRS':\n if predictor_params is not None:\n predictor_params_.update(predictor_params)\n predictor = BasePredictor(net, device, zoom_in=zoom_in, with_flip=with_flip, **predictor_params_)\n elif brs_mode.startswith('f-BRS'):\n predictor_params_.update({\n 'net_clicks_limit': 8,\n })\n if predictor_params is not None:\n predictor_params_.update(predictor_params)\n\n insertion_mode = {\n 'f-BRS-A': 'after_c4',\n 'f-BRS-B': 'after_aspp',\n 'f-BRS-C': 'after_deeplab'\n }[brs_mode]\n\n opt_functor = ScaleBiasOptimizer(prob_thresh=prob_thresh,\n with_flip=with_flip,\n optimizer_params=lbfgs_params_,\n **brs_opt_func_params)\n\n if isinstance(net, DistMapsHRNetModel):\n FeaturePredictor = HRNetFeatureBRSPredictor\n insertion_mode = {'after_c4': 'A', 'after_aspp': 'A', 'after_deeplab': 'C'}[insertion_mode]\n else:\n FeaturePredictor = FeatureBRSPredictor\n\n predictor = FeaturePredictor(net, device,\n opt_functor=opt_functor,\n with_flip=with_flip,\n insertion_mode=insertion_mode,\n zoom_in=zoom_in,\n **predictor_params_)\n elif brs_mode == 'RGB-BRS' or brs_mode == 'DistMap-BRS':\n use_dmaps = brs_mode == 'DistMap-BRS'\n\n predictor_params_.update({\n 'net_clicks_limit': 5,\n })\n if predictor_params is not None:\n predictor_params_.update(predictor_params)\n\n opt_functor = InputOptimizer(prob_thresh=prob_thresh,\n with_flip=with_flip,\n optimizer_params=lbfgs_params_,\n **brs_opt_func_params)\n\n predictor = InputBRSPredictor(net, device,\n optimize_target='dmaps' if use_dmaps else 'rgb',\n opt_functor=opt_functor,\n with_flip=with_flip,\n zoom_in=zoom_in,\n **predictor_params_)\n else:\n raise NotImplementedError\n\n return predictor\n" }, { "path": "XMem/inference/interact/fbrs/inference/predictors/base.py", "content": "import torch\nimport torch.nn.functional as F\n\nfrom ..transforms import AddHorizontalFlip, SigmoidForPred, LimitLongestSide\n\n\nclass BasePredictor(object):\n def __init__(self, net, device,\n net_clicks_limit=None,\n with_flip=False,\n zoom_in=None,\n max_size=None,\n **kwargs):\n self.net = net\n self.with_flip = with_flip\n self.net_clicks_limit = net_clicks_limit\n self.original_image = None\n self.device = device\n self.zoom_in = zoom_in\n\n self.transforms = [zoom_in] if zoom_in is not None else []\n if max_size is not None:\n self.transforms.append(LimitLongestSide(max_size=max_size))\n self.transforms.append(SigmoidForPred())\n if with_flip:\n self.transforms.append(AddHorizontalFlip())\n\n def set_input_image(self, image_nd):\n for transform in self.transforms:\n transform.reset()\n self.original_image = image_nd.to(self.device)\n if len(self.original_image.shape) == 3:\n self.original_image = self.original_image.unsqueeze(0)\n\n def get_prediction(self, clicker):\n clicks_list = clicker.get_clicks()\n\n image_nd, clicks_lists, is_image_changed = self.apply_transforms(\n self.original_image, [clicks_list]\n )\n\n pred_logits = self._get_prediction(image_nd, clicks_lists, is_image_changed)\n prediction = F.interpolate(pred_logits, mode='bilinear', align_corners=True,\n size=image_nd.size()[2:])\n\n for t in reversed(self.transforms):\n prediction = t.inv_transform(prediction)\n\n if self.zoom_in is not None and self.zoom_in.check_possible_recalculation():\n print('zooming')\n return self.get_prediction(clicker)\n\n # return prediction.cpu().numpy()[0, 0]\n return prediction\n\n def _get_prediction(self, image_nd, clicks_lists, is_image_changed):\n points_nd = self.get_points_nd(clicks_lists)\n return self.net(image_nd, points_nd)['instances']\n\n def _get_transform_states(self):\n return [x.get_state() for x in self.transforms]\n\n def _set_transform_states(self, states):\n assert len(states) == len(self.transforms)\n for state, transform in zip(states, self.transforms):\n transform.set_state(state)\n\n def apply_transforms(self, image_nd, clicks_lists):\n is_image_changed = False\n for t in self.transforms:\n image_nd, clicks_lists = t.transform(image_nd, clicks_lists)\n is_image_changed |= t.image_changed\n\n return image_nd, clicks_lists, is_image_changed\n\n def get_points_nd(self, clicks_lists):\n total_clicks = []\n num_pos_clicks = [sum(x.is_positive for x in clicks_list) for clicks_list in clicks_lists]\n num_neg_clicks = [len(clicks_list) - num_pos for clicks_list, num_pos in zip(clicks_lists, num_pos_clicks)]\n num_max_points = max(num_pos_clicks + num_neg_clicks)\n if self.net_clicks_limit is not None:\n num_max_points = min(self.net_clicks_limit, num_max_points)\n num_max_points = max(1, num_max_points)\n\n for clicks_list in clicks_lists:\n clicks_list = clicks_list[:self.net_clicks_limit]\n pos_clicks = [click.coords for click in clicks_list if click.is_positive]\n pos_clicks = pos_clicks + (num_max_points - len(pos_clicks)) * [(-1, -1)]\n\n neg_clicks = [click.coords for click in clicks_list if not click.is_positive]\n neg_clicks = neg_clicks + (num_max_points - len(neg_clicks)) * [(-1, -1)]\n total_clicks.append(pos_clicks + neg_clicks)\n\n return torch.tensor(total_clicks, device=self.device)\n\n def get_states(self):\n return {'transform_states': self._get_transform_states()}\n\n def set_states(self, states):\n self._set_transform_states(states['transform_states'])\n" }, { "path": "XMem/inference/interact/fbrs/inference/predictors/brs.py", "content": "import torch\nimport torch.nn.functional as F\nimport numpy as np\nfrom scipy.optimize import fmin_l_bfgs_b\n\nfrom .base import BasePredictor\nfrom ...model.is_hrnet_model import DistMapsHRNetModel\n\n\nclass BRSBasePredictor(BasePredictor):\n def __init__(self, model, device, opt_functor, optimize_after_n_clicks=1, **kwargs):\n super().__init__(model, device, **kwargs)\n self.optimize_after_n_clicks = optimize_after_n_clicks\n self.opt_functor = opt_functor\n\n self.opt_data = None\n self.input_data = None\n\n def set_input_image(self, image_nd):\n super().set_input_image(image_nd)\n self.opt_data = None\n self.input_data = None\n\n def _get_clicks_maps_nd(self, clicks_lists, image_shape, radius=1):\n pos_clicks_map = np.zeros((len(clicks_lists), 1) + image_shape, dtype=np.float32)\n neg_clicks_map = np.zeros((len(clicks_lists), 1) + image_shape, dtype=np.float32)\n\n for list_indx, clicks_list in enumerate(clicks_lists):\n for click in clicks_list:\n y, x = click.coords\n y, x = int(round(y)), int(round(x))\n y1, x1 = y - radius, x - radius\n y2, x2 = y + radius + 1, x + radius + 1\n\n if click.is_positive:\n pos_clicks_map[list_indx, 0, y1:y2, x1:x2] = True\n else:\n neg_clicks_map[list_indx, 0, y1:y2, x1:x2] = True\n\n with torch.no_grad():\n pos_clicks_map = torch.from_numpy(pos_clicks_map).to(self.device)\n neg_clicks_map = torch.from_numpy(neg_clicks_map).to(self.device)\n\n return pos_clicks_map, neg_clicks_map\n\n def get_states(self):\n return {'transform_states': self._get_transform_states(), 'opt_data': self.opt_data}\n\n def set_states(self, states):\n self._set_transform_states(states['transform_states'])\n self.opt_data = states['opt_data']\n\n\nclass FeatureBRSPredictor(BRSBasePredictor):\n def __init__(self, model, device, opt_functor, insertion_mode='after_deeplab', **kwargs):\n super().__init__(model, device, opt_functor=opt_functor, **kwargs)\n self.insertion_mode = insertion_mode\n self._c1_features = None\n\n if self.insertion_mode == 'after_deeplab':\n self.num_channels = model.feature_extractor.ch\n elif self.insertion_mode == 'after_c4':\n self.num_channels = model.feature_extractor.aspp_in_channels\n elif self.insertion_mode == 'after_aspp':\n self.num_channels = model.feature_extractor.ch + 32\n else:\n raise NotImplementedError\n\n def _get_prediction(self, image_nd, clicks_lists, is_image_changed):\n points_nd = self.get_points_nd(clicks_lists)\n pos_mask, neg_mask = self._get_clicks_maps_nd(clicks_lists, image_nd.shape[2:])\n\n num_clicks = len(clicks_lists[0])\n bs = image_nd.shape[0] // 2 if self.with_flip else image_nd.shape[0]\n\n if self.opt_data is None or self.opt_data.shape[0] // (2 * self.num_channels) != bs:\n self.opt_data = np.zeros((bs * 2 * self.num_channels), dtype=np.float32)\n\n if num_clicks <= self.net_clicks_limit or is_image_changed or self.input_data is None:\n self.input_data = self._get_head_input(image_nd, points_nd)\n\n def get_prediction_logits(scale, bias):\n scale = scale.view(bs, -1, 1, 1)\n bias = bias.view(bs, -1, 1, 1)\n if self.with_flip:\n scale = scale.repeat(2, 1, 1, 1)\n bias = bias.repeat(2, 1, 1, 1)\n\n scaled_backbone_features = self.input_data * scale\n scaled_backbone_features = scaled_backbone_features + bias\n if self.insertion_mode == 'after_c4':\n x = self.net.feature_extractor.aspp(scaled_backbone_features)\n x = F.interpolate(x, mode='bilinear', size=self._c1_features.size()[2:],\n align_corners=True)\n x = torch.cat((x, self._c1_features), dim=1)\n scaled_backbone_features = self.net.feature_extractor.head(x)\n elif self.insertion_mode == 'after_aspp':\n scaled_backbone_features = self.net.feature_extractor.head(scaled_backbone_features)\n\n pred_logits = self.net.head(scaled_backbone_features)\n pred_logits = F.interpolate(pred_logits, size=image_nd.size()[2:], mode='bilinear',\n align_corners=True)\n return pred_logits\n\n self.opt_functor.init_click(get_prediction_logits, pos_mask, neg_mask, self.device)\n if num_clicks > self.optimize_after_n_clicks:\n opt_result = fmin_l_bfgs_b(func=self.opt_functor, x0=self.opt_data,\n **self.opt_functor.optimizer_params)\n self.opt_data = opt_result[0]\n\n with torch.no_grad():\n if self.opt_functor.best_prediction is not None:\n opt_pred_logits = self.opt_functor.best_prediction\n else:\n opt_data_nd = torch.from_numpy(self.opt_data).to(self.device)\n opt_vars, _ = self.opt_functor.unpack_opt_params(opt_data_nd)\n opt_pred_logits = get_prediction_logits(*opt_vars)\n\n return opt_pred_logits\n\n def _get_head_input(self, image_nd, points):\n with torch.no_grad():\n coord_features = self.net.dist_maps(image_nd, points)\n x = self.net.rgb_conv(torch.cat((image_nd, coord_features), dim=1))\n if self.insertion_mode == 'after_c4' or self.insertion_mode == 'after_aspp':\n c1, _, c3, c4 = self.net.feature_extractor.backbone(x)\n c1 = self.net.feature_extractor.skip_project(c1)\n\n if self.insertion_mode == 'after_aspp':\n x = self.net.feature_extractor.aspp(c4)\n x = F.interpolate(x, size=c1.size()[2:], mode='bilinear', align_corners=True)\n x = torch.cat((x, c1), dim=1)\n backbone_features = x\n else:\n backbone_features = c4\n self._c1_features = c1\n else:\n backbone_features = self.net.feature_extractor(x)[0]\n\n return backbone_features\n\n\nclass HRNetFeatureBRSPredictor(BRSBasePredictor):\n def __init__(self, model, device, opt_functor, insertion_mode='A', **kwargs):\n super().__init__(model, device, opt_functor=opt_functor, **kwargs)\n self.insertion_mode = insertion_mode\n self._c1_features = None\n\n if self.insertion_mode == 'A':\n self.num_channels = sum(k * model.feature_extractor.width for k in [1, 2, 4, 8])\n elif self.insertion_mode == 'C':\n self.num_channels = 2 * model.feature_extractor.ocr_width\n else:\n raise NotImplementedError\n\n def _get_prediction(self, image_nd, clicks_lists, is_image_changed):\n points_nd = self.get_points_nd(clicks_lists)\n pos_mask, neg_mask = self._get_clicks_maps_nd(clicks_lists, image_nd.shape[2:])\n num_clicks = len(clicks_lists[0])\n bs = image_nd.shape[0] // 2 if self.with_flip else image_nd.shape[0]\n\n if self.opt_data is None or self.opt_data.shape[0] // (2 * self.num_channels) != bs:\n self.opt_data = np.zeros((bs * 2 * self.num_channels), dtype=np.float32)\n\n if num_clicks <= self.net_clicks_limit or is_image_changed or self.input_data is None:\n self.input_data = self._get_head_input(image_nd, points_nd)\n\n def get_prediction_logits(scale, bias):\n scale = scale.view(bs, -1, 1, 1)\n bias = bias.view(bs, -1, 1, 1)\n if self.with_flip:\n scale = scale.repeat(2, 1, 1, 1)\n bias = bias.repeat(2, 1, 1, 1)\n\n scaled_backbone_features = self.input_data * scale\n scaled_backbone_features = scaled_backbone_features + bias\n if self.insertion_mode == 'A':\n out_aux = self.net.feature_extractor.aux_head(scaled_backbone_features)\n feats = self.net.feature_extractor.conv3x3_ocr(scaled_backbone_features)\n\n context = self.net.feature_extractor.ocr_gather_head(feats, out_aux)\n feats = self.net.feature_extractor.ocr_distri_head(feats, context)\n pred_logits = self.net.feature_extractor.cls_head(feats)\n elif self.insertion_mode == 'C':\n pred_logits = self.net.feature_extractor.cls_head(scaled_backbone_features)\n else:\n raise NotImplementedError\n\n pred_logits = F.interpolate(pred_logits, size=image_nd.size()[2:], mode='bilinear',\n align_corners=True)\n return pred_logits\n\n self.opt_functor.init_click(get_prediction_logits, pos_mask, neg_mask, self.device)\n if num_clicks > self.optimize_after_n_clicks:\n opt_result = fmin_l_bfgs_b(func=self.opt_functor, x0=self.opt_data,\n **self.opt_functor.optimizer_params)\n self.opt_data = opt_result[0]\n\n with torch.no_grad():\n if self.opt_functor.best_prediction is not None:\n opt_pred_logits = self.opt_functor.best_prediction\n else:\n opt_data_nd = torch.from_numpy(self.opt_data).to(self.device)\n opt_vars, _ = self.opt_functor.unpack_opt_params(opt_data_nd)\n opt_pred_logits = get_prediction_logits(*opt_vars)\n\n return opt_pred_logits\n\n def _get_head_input(self, image_nd, points):\n with torch.no_grad():\n coord_features = self.net.dist_maps(image_nd, points)\n x = self.net.rgb_conv(torch.cat((image_nd, coord_features), dim=1))\n feats = self.net.feature_extractor.compute_hrnet_feats(x)\n if self.insertion_mode == 'A':\n backbone_features = feats\n elif self.insertion_mode == 'C':\n out_aux = self.net.feature_extractor.aux_head(feats)\n feats = self.net.feature_extractor.conv3x3_ocr(feats)\n\n context = self.net.feature_extractor.ocr_gather_head(feats, out_aux)\n backbone_features = self.net.feature_extractor.ocr_distri_head(feats, context)\n else:\n raise NotImplementedError\n\n return backbone_features\n\n\nclass InputBRSPredictor(BRSBasePredictor):\n def __init__(self, model, device, opt_functor, optimize_target='rgb', **kwargs):\n super().__init__(model, device, opt_functor=opt_functor, **kwargs)\n self.optimize_target = optimize_target\n\n def _get_prediction(self, image_nd, clicks_lists, is_image_changed):\n points_nd = self.get_points_nd(clicks_lists)\n pos_mask, neg_mask = self._get_clicks_maps_nd(clicks_lists, image_nd.shape[2:])\n num_clicks = len(clicks_lists[0])\n\n if self.opt_data is None or is_image_changed:\n opt_channels = 2 if self.optimize_target == 'dmaps' else 3\n bs = image_nd.shape[0] // 2 if self.with_flip else image_nd.shape[0]\n self.opt_data = torch.zeros((bs, opt_channels, image_nd.shape[2], image_nd.shape[3]),\n device=self.device, dtype=torch.float32)\n\n def get_prediction_logits(opt_bias):\n input_image = image_nd\n if self.optimize_target == 'rgb':\n input_image = input_image + opt_bias\n dmaps = self.net.dist_maps(input_image, points_nd)\n if self.optimize_target == 'dmaps':\n dmaps = dmaps + opt_bias\n\n x = self.net.rgb_conv(torch.cat((input_image, dmaps), dim=1))\n if self.optimize_target == 'all':\n x = x + opt_bias\n\n if isinstance(self.net, DistMapsHRNetModel):\n pred_logits = self.net.feature_extractor(x)[0]\n else:\n backbone_features = self.net.feature_extractor(x)\n pred_logits = self.net.head(backbone_features[0])\n pred_logits = F.interpolate(pred_logits, size=image_nd.size()[2:], mode='bilinear', align_corners=True)\n\n return pred_logits\n\n self.opt_functor.init_click(get_prediction_logits, pos_mask, neg_mask, self.device,\n shape=self.opt_data.shape)\n if num_clicks > self.optimize_after_n_clicks:\n opt_result = fmin_l_bfgs_b(func=self.opt_functor, x0=self.opt_data.cpu().numpy().ravel(),\n **self.opt_functor.optimizer_params)\n\n self.opt_data = torch.from_numpy(opt_result[0]).view(self.opt_data.shape).to(self.device)\n\n with torch.no_grad():\n if self.opt_functor.best_prediction is not None:\n opt_pred_logits = self.opt_functor.best_prediction\n else:\n opt_vars, _ = self.opt_functor.unpack_opt_params(self.opt_data)\n opt_pred_logits = get_prediction_logits(*opt_vars)\n\n return opt_pred_logits\n" }, { "path": "XMem/inference/interact/fbrs/inference/predictors/brs_functors.py", "content": "import torch\nimport numpy as np\n\nfrom ...model.metrics import _compute_iou\nfrom .brs_losses import BRSMaskLoss\n\n\nclass BaseOptimizer:\n def __init__(self, optimizer_params,\n prob_thresh=0.49,\n reg_weight=1e-3,\n min_iou_diff=0.01,\n brs_loss=BRSMaskLoss(),\n with_flip=False,\n flip_average=False,\n **kwargs):\n self.brs_loss = brs_loss\n self.optimizer_params = optimizer_params\n self.prob_thresh = prob_thresh\n self.reg_weight = reg_weight\n self.min_iou_diff = min_iou_diff\n self.with_flip = with_flip\n self.flip_average = flip_average\n\n self.best_prediction = None\n self._get_prediction_logits = None\n self._opt_shape = None\n self._best_loss = None\n self._click_masks = None\n self._last_mask = None\n self.device = None\n\n def init_click(self, get_prediction_logits, pos_mask, neg_mask, device, shape=None):\n self.best_prediction = None\n self._get_prediction_logits = get_prediction_logits\n self._click_masks = (pos_mask, neg_mask)\n self._opt_shape = shape\n self._last_mask = None\n self.device = device\n\n def __call__(self, x):\n opt_params = torch.from_numpy(x).float().to(self.device)\n opt_params.requires_grad_(True)\n\n with torch.enable_grad():\n opt_vars, reg_loss = self.unpack_opt_params(opt_params)\n result_before_sigmoid = self._get_prediction_logits(*opt_vars)\n result = torch.sigmoid(result_before_sigmoid)\n\n pos_mask, neg_mask = self._click_masks\n if self.with_flip and self.flip_average:\n result, result_flipped = torch.chunk(result, 2, dim=0)\n result = 0.5 * (result + torch.flip(result_flipped, dims=[3]))\n pos_mask, neg_mask = pos_mask[:result.shape[0]], neg_mask[:result.shape[0]]\n\n loss, f_max_pos, f_max_neg = self.brs_loss(result, pos_mask, neg_mask)\n loss = loss + reg_loss\n\n f_val = loss.detach().cpu().numpy()\n if self.best_prediction is None or f_val < self._best_loss:\n self.best_prediction = result_before_sigmoid.detach()\n self._best_loss = f_val\n\n if f_max_pos < (1 - self.prob_thresh) and f_max_neg < self.prob_thresh:\n return [f_val, np.zeros_like(x)]\n\n current_mask = result > self.prob_thresh\n if self._last_mask is not None and self.min_iou_diff > 0:\n diff_iou = _compute_iou(current_mask, self._last_mask)\n if len(diff_iou) > 0 and diff_iou.mean() > 1 - self.min_iou_diff:\n return [f_val, np.zeros_like(x)]\n self._last_mask = current_mask\n\n loss.backward()\n f_grad = opt_params.grad.cpu().numpy().ravel().astype(np.float32)\n\n return [f_val, f_grad]\n\n def unpack_opt_params(self, opt_params):\n raise NotImplementedError\n\n\nclass InputOptimizer(BaseOptimizer):\n def unpack_opt_params(self, opt_params):\n opt_params = opt_params.view(self._opt_shape)\n if self.with_flip:\n opt_params_flipped = torch.flip(opt_params, dims=[3])\n opt_params = torch.cat([opt_params, opt_params_flipped], dim=0)\n reg_loss = self.reg_weight * torch.sum(opt_params**2)\n\n return (opt_params,), reg_loss\n\n\nclass ScaleBiasOptimizer(BaseOptimizer):\n def __init__(self, *args, scale_act=None, reg_bias_weight=10.0, **kwargs):\n super().__init__(*args, **kwargs)\n self.scale_act = scale_act\n self.reg_bias_weight = reg_bias_weight\n\n def unpack_opt_params(self, opt_params):\n scale, bias = torch.chunk(opt_params, 2, dim=0)\n reg_loss = self.reg_weight * (torch.sum(scale**2) + self.reg_bias_weight * torch.sum(bias**2))\n\n if self.scale_act == 'tanh':\n scale = torch.tanh(scale)\n elif self.scale_act == 'sin':\n scale = torch.sin(scale)\n\n return (1 + scale, bias), reg_loss\n" }, { "path": "XMem/inference/interact/fbrs/inference/predictors/brs_losses.py", "content": "import torch\n\nfrom ...model.losses import SigmoidBinaryCrossEntropyLoss\n\n\nclass BRSMaskLoss(torch.nn.Module):\n def __init__(self, eps=1e-5):\n super().__init__()\n self._eps = eps\n\n def forward(self, result, pos_mask, neg_mask):\n pos_diff = (1 - result) * pos_mask\n pos_target = torch.sum(pos_diff ** 2)\n pos_target = pos_target / (torch.sum(pos_mask) + self._eps)\n\n neg_diff = result * neg_mask\n neg_target = torch.sum(neg_diff ** 2)\n neg_target = neg_target / (torch.sum(neg_mask) + self._eps)\n \n loss = pos_target + neg_target\n\n with torch.no_grad():\n f_max_pos = torch.max(torch.abs(pos_diff)).item()\n f_max_neg = torch.max(torch.abs(neg_diff)).item()\n\n return loss, f_max_pos, f_max_neg\n\n\nclass OracleMaskLoss(torch.nn.Module):\n def __init__(self):\n super().__init__()\n self.gt_mask = None\n self.loss = SigmoidBinaryCrossEntropyLoss(from_sigmoid=True)\n self.predictor = None\n self.history = []\n\n def set_gt_mask(self, gt_mask):\n self.gt_mask = gt_mask\n self.history = []\n\n def forward(self, result, pos_mask, neg_mask):\n gt_mask = self.gt_mask.to(result.device)\n if self.predictor.object_roi is not None:\n r1, r2, c1, c2 = self.predictor.object_roi[:4]\n gt_mask = gt_mask[:, :, r1:r2 + 1, c1:c2 + 1]\n gt_mask = torch.nn.functional.interpolate(gt_mask, result.size()[2:], mode='bilinear', align_corners=True)\n\n if result.shape[0] == 2:\n gt_mask_flipped = torch.flip(gt_mask, dims=[3])\n gt_mask = torch.cat([gt_mask, gt_mask_flipped], dim=0)\n\n loss = self.loss(result, gt_mask)\n self.history.append(loss.detach().cpu().numpy()[0])\n\n if len(self.history) > 5 and abs(self.history[-5] - self.history[-1]) < 1e-5:\n return 0, 0, 0\n\n return loss, 1.0, 1.0\n" }, { "path": "XMem/inference/interact/fbrs/inference/transforms/__init__.py", "content": "from .base import SigmoidForPred\nfrom .flip import AddHorizontalFlip\nfrom .zoom_in import ZoomIn\nfrom .limit_longest_side import LimitLongestSide\nfrom .crops import Crops\n" }, { "path": "XMem/inference/interact/fbrs/inference/transforms/base.py", "content": "import torch\n\n\nclass BaseTransform(object):\n def __init__(self):\n self.image_changed = False\n\n def transform(self, image_nd, clicks_lists):\n raise NotImplementedError\n\n def inv_transform(self, prob_map):\n raise NotImplementedError\n\n def reset(self):\n raise NotImplementedError\n\n def get_state(self):\n raise NotImplementedError\n\n def set_state(self, state):\n raise NotImplementedError\n\n\nclass SigmoidForPred(BaseTransform):\n def transform(self, image_nd, clicks_lists):\n return image_nd, clicks_lists\n\n def inv_transform(self, prob_map):\n return torch.sigmoid(prob_map)\n\n def reset(self):\n pass\n\n def get_state(self):\n return None\n\n def set_state(self, state):\n pass\n" }, { "path": "XMem/inference/interact/fbrs/inference/transforms/crops.py", "content": "import math\n\nimport torch\nimport numpy as np\n\nfrom ...inference.clicker import Click\nfrom .base import BaseTransform\n\n\nclass Crops(BaseTransform):\n def __init__(self, crop_size=(320, 480), min_overlap=0.2):\n super().__init__()\n self.crop_height, self.crop_width = crop_size\n self.min_overlap = min_overlap\n\n self.x_offsets = None\n self.y_offsets = None\n self._counts = None\n\n def transform(self, image_nd, clicks_lists):\n assert image_nd.shape[0] == 1 and len(clicks_lists) == 1\n image_height, image_width = image_nd.shape[2:4]\n self._counts = None\n\n if image_height < self.crop_height or image_width < self.crop_width:\n return image_nd, clicks_lists\n\n self.x_offsets = get_offsets(image_width, self.crop_width, self.min_overlap)\n self.y_offsets = get_offsets(image_height, self.crop_height, self.min_overlap)\n self._counts = np.zeros((image_height, image_width))\n\n image_crops = []\n for dy in self.y_offsets:\n for dx in self.x_offsets:\n self._counts[dy:dy + self.crop_height, dx:dx + self.crop_width] += 1\n image_crop = image_nd[:, :, dy:dy + self.crop_height, dx:dx + self.crop_width]\n image_crops.append(image_crop)\n image_crops = torch.cat(image_crops, dim=0)\n self._counts = torch.tensor(self._counts, device=image_nd.device, dtype=torch.float32)\n\n clicks_list = clicks_lists[0]\n clicks_lists = []\n for dy in self.y_offsets:\n for dx in self.x_offsets:\n crop_clicks = [Click(is_positive=x.is_positive, coords=(x.coords[0] - dy, x.coords[1] - dx))\n for x in clicks_list]\n clicks_lists.append(crop_clicks)\n\n return image_crops, clicks_lists\n\n def inv_transform(self, prob_map):\n if self._counts is None:\n return prob_map\n\n new_prob_map = torch.zeros((1, 1, *self._counts.shape),\n dtype=prob_map.dtype, device=prob_map.device)\n\n crop_indx = 0\n for dy in self.y_offsets:\n for dx in self.x_offsets:\n new_prob_map[0, 0, dy:dy + self.crop_height, dx:dx + self.crop_width] += prob_map[crop_indx, 0]\n crop_indx += 1\n new_prob_map = torch.div(new_prob_map, self._counts)\n\n return new_prob_map\n\n def get_state(self):\n return self.x_offsets, self.y_offsets, self._counts\n\n def set_state(self, state):\n self.x_offsets, self.y_offsets, self._counts = state\n\n def reset(self):\n self.x_offsets = None\n self.y_offsets = None\n self._counts = None\n\n\ndef get_offsets(length, crop_size, min_overlap_ratio=0.2):\n if length == crop_size:\n return [0]\n\n N = (length / crop_size - min_overlap_ratio) / (1 - min_overlap_ratio)\n N = math.ceil(N)\n\n overlap_ratio = (N - length / crop_size) / (N - 1)\n overlap_width = int(crop_size * overlap_ratio)\n\n offsets = [0]\n for i in range(1, N):\n new_offset = offsets[-1] + crop_size - overlap_width\n if new_offset + crop_size > length:\n new_offset = length - crop_size\n\n offsets.append(new_offset)\n\n return offsets\n" }, { "path": "XMem/inference/interact/fbrs/inference/transforms/flip.py", "content": "import torch\n\nfrom ..clicker import Click\nfrom .base import BaseTransform\n\n\nclass AddHorizontalFlip(BaseTransform):\n def transform(self, image_nd, clicks_lists):\n assert len(image_nd.shape) == 4\n image_nd = torch.cat([image_nd, torch.flip(image_nd, dims=[3])], dim=0)\n\n image_width = image_nd.shape[3]\n clicks_lists_flipped = []\n for clicks_list in clicks_lists:\n clicks_list_flipped = [Click(is_positive=click.is_positive,\n coords=(click.coords[0], image_width - click.coords[1] - 1))\n for click in clicks_list]\n clicks_lists_flipped.append(clicks_list_flipped)\n clicks_lists = clicks_lists + clicks_lists_flipped\n\n return image_nd, clicks_lists\n\n def inv_transform(self, prob_map):\n assert len(prob_map.shape) == 4 and prob_map.shape[0] % 2 == 0\n num_maps = prob_map.shape[0] // 2\n prob_map, prob_map_flipped = prob_map[:num_maps], prob_map[num_maps:]\n\n return 0.5 * (prob_map + torch.flip(prob_map_flipped, dims=[3]))\n\n def get_state(self):\n return None\n\n def set_state(self, state):\n pass\n\n def reset(self):\n pass\n" }, { "path": "XMem/inference/interact/fbrs/inference/transforms/limit_longest_side.py", "content": "from .zoom_in import ZoomIn, get_roi_image_nd\n\n\nclass LimitLongestSide(ZoomIn):\n def __init__(self, max_size=800):\n super().__init__(target_size=max_size, skip_clicks=0)\n\n def transform(self, image_nd, clicks_lists):\n assert image_nd.shape[0] == 1 and len(clicks_lists) == 1\n image_max_size = max(image_nd.shape[2:4])\n self.image_changed = False\n\n if image_max_size <= self.target_size:\n return image_nd, clicks_lists\n self._input_image = image_nd\n\n self._object_roi = (0, image_nd.shape[2] - 1, 0, image_nd.shape[3] - 1)\n self._roi_image = get_roi_image_nd(image_nd, self._object_roi, self.target_size)\n self.image_changed = True\n\n tclicks_lists = [self._transform_clicks(clicks_lists[0])]\n return self._roi_image, tclicks_lists\n" }, { "path": "XMem/inference/interact/fbrs/inference/transforms/zoom_in.py", "content": "import torch\n\nfrom ..clicker import Click\nfrom ...utils.misc import get_bbox_iou, get_bbox_from_mask, expand_bbox, clamp_bbox\nfrom .base import BaseTransform\n\n\nclass ZoomIn(BaseTransform):\n def __init__(self,\n target_size=400,\n skip_clicks=1,\n expansion_ratio=1.4,\n min_crop_size=200,\n recompute_thresh_iou=0.5,\n prob_thresh=0.50):\n super().__init__()\n self.target_size = target_size\n self.min_crop_size = min_crop_size\n self.skip_clicks = skip_clicks\n self.expansion_ratio = expansion_ratio\n self.recompute_thresh_iou = recompute_thresh_iou\n self.prob_thresh = prob_thresh\n\n self._input_image_shape = None\n self._prev_probs = None\n self._object_roi = None\n self._roi_image = None\n\n def transform(self, image_nd, clicks_lists):\n assert image_nd.shape[0] == 1 and len(clicks_lists) == 1\n self.image_changed = False\n\n clicks_list = clicks_lists[0]\n if len(clicks_list) <= self.skip_clicks:\n return image_nd, clicks_lists\n\n self._input_image_shape = image_nd.shape\n\n current_object_roi = None\n if self._prev_probs is not None:\n current_pred_mask = (self._prev_probs > self.prob_thresh)[0, 0]\n if current_pred_mask.sum() > 0:\n current_object_roi = get_object_roi(current_pred_mask, clicks_list,\n self.expansion_ratio, self.min_crop_size)\n\n if current_object_roi is None:\n return image_nd, clicks_lists\n\n update_object_roi = False\n if self._object_roi is None:\n update_object_roi = True\n elif not check_object_roi(self._object_roi, clicks_list):\n update_object_roi = True\n elif get_bbox_iou(current_object_roi, self._object_roi) < self.recompute_thresh_iou:\n update_object_roi = True\n\n if update_object_roi:\n self._object_roi = current_object_roi\n self._roi_image = get_roi_image_nd(image_nd, self._object_roi, self.target_size)\n self.image_changed = True\n\n tclicks_lists = [self._transform_clicks(clicks_list)]\n return self._roi_image.to(image_nd.device), tclicks_lists\n\n def inv_transform(self, prob_map):\n if self._object_roi is None:\n self._prev_probs = prob_map.cpu().numpy()\n return prob_map\n\n assert prob_map.shape[0] == 1\n rmin, rmax, cmin, cmax = self._object_roi\n prob_map = torch.nn.functional.interpolate(prob_map, size=(rmax - rmin + 1, cmax - cmin + 1),\n mode='bilinear', align_corners=True)\n\n if self._prev_probs is not None:\n new_prob_map = torch.zeros(*self._prev_probs.shape, device=prob_map.device, dtype=prob_map.dtype)\n new_prob_map[:, :, rmin:rmax + 1, cmin:cmax + 1] = prob_map\n else:\n new_prob_map = prob_map\n\n self._prev_probs = new_prob_map.cpu().numpy()\n\n return new_prob_map\n\n def check_possible_recalculation(self):\n if self._prev_probs is None or self._object_roi is not None or self.skip_clicks > 0:\n return False\n\n pred_mask = (self._prev_probs > self.prob_thresh)[0, 0]\n if pred_mask.sum() > 0:\n possible_object_roi = get_object_roi(pred_mask, [],\n self.expansion_ratio, self.min_crop_size)\n image_roi = (0, self._input_image_shape[2] - 1, 0, self._input_image_shape[3] - 1)\n if get_bbox_iou(possible_object_roi, image_roi) < 0.50:\n return True\n return False\n\n def get_state(self):\n roi_image = self._roi_image.cpu() if self._roi_image is not None else None\n return self._input_image_shape, self._object_roi, self._prev_probs, roi_image, self.image_changed\n\n def set_state(self, state):\n self._input_image_shape, self._object_roi, self._prev_probs, self._roi_image, self.image_changed = state\n\n def reset(self):\n self._input_image_shape = None\n self._object_roi = None\n self._prev_probs = None\n self._roi_image = None\n self.image_changed = False\n\n def _transform_clicks(self, clicks_list):\n if self._object_roi is None:\n return clicks_list\n\n rmin, rmax, cmin, cmax = self._object_roi\n crop_height, crop_width = self._roi_image.shape[2:]\n\n transformed_clicks = []\n for click in clicks_list:\n new_r = crop_height * (click.coords[0] - rmin) / (rmax - rmin + 1)\n new_c = crop_width * (click.coords[1] - cmin) / (cmax - cmin + 1)\n transformed_clicks.append(Click(is_positive=click.is_positive, coords=(new_r, new_c)))\n return transformed_clicks\n\n\ndef get_object_roi(pred_mask, clicks_list, expansion_ratio, min_crop_size):\n pred_mask = pred_mask.copy()\n\n for click in clicks_list:\n if click.is_positive:\n pred_mask[int(click.coords[0]), int(click.coords[1])] = 1\n\n bbox = get_bbox_from_mask(pred_mask)\n bbox = expand_bbox(bbox, expansion_ratio, min_crop_size)\n h, w = pred_mask.shape[0], pred_mask.shape[1]\n bbox = clamp_bbox(bbox, 0, h - 1, 0, w - 1)\n\n return bbox\n\n\ndef get_roi_image_nd(image_nd, object_roi, target_size):\n rmin, rmax, cmin, cmax = object_roi\n\n height = rmax - rmin + 1\n width = cmax - cmin + 1\n\n if isinstance(target_size, tuple):\n new_height, new_width = target_size\n else:\n scale = target_size / max(height, width)\n new_height = int(round(height * scale))\n new_width = int(round(width * scale))\n\n with torch.no_grad():\n roi_image_nd = image_nd[:, :, rmin:rmax + 1, cmin:cmax + 1]\n roi_image_nd = torch.nn.functional.interpolate(roi_image_nd, size=(new_height, new_width),\n mode='bilinear', align_corners=True)\n\n return roi_image_nd\n\n\ndef check_object_roi(object_roi, clicks_list):\n for click in clicks_list:\n if click.is_positive:\n if click.coords[0] < object_roi[0] or click.coords[0] >= object_roi[1]:\n return False\n if click.coords[1] < object_roi[2] or click.coords[1] >= object_roi[3]:\n return False\n\n return True\n" }, { "path": "XMem/inference/interact/fbrs/inference/utils.py", "content": "from datetime import timedelta\nfrom pathlib import Path\n\nimport torch\nimport numpy as np\n\nfrom ..model.is_deeplab_model import get_deeplab_model\nfrom ..model.is_hrnet_model import get_hrnet_model\n\n\ndef get_time_metrics(all_ious, elapsed_time):\n n_images = len(all_ious)\n n_clicks = sum(map(len, all_ious))\n\n mean_spc = elapsed_time / n_clicks\n mean_spi = elapsed_time / n_images\n\n return mean_spc, mean_spi\n\n\ndef load_is_model(checkpoint, device, backbone='auto', **kwargs):\n if isinstance(checkpoint, (str, Path)):\n state_dict = torch.load(checkpoint, map_location='cpu')\n else:\n state_dict = checkpoint\n\n if backbone == 'auto':\n for k in state_dict.keys():\n if 'feature_extractor.stage2.0.branches' in k:\n return load_hrnet_is_model(state_dict, device, backbone, **kwargs)\n return load_deeplab_is_model(state_dict, device, backbone, **kwargs)\n elif 'resnet' in backbone:\n return load_deeplab_is_model(state_dict, device, backbone, **kwargs)\n elif 'hrnet' in backbone:\n return load_hrnet_is_model(state_dict, device, backbone, **kwargs)\n else:\n raise NotImplementedError('Unknown backbone')\n\n\ndef load_hrnet_is_model(state_dict, device, backbone='auto', width=48, ocr_width=256,\n small=False, cpu_dist_maps=False, norm_radius=260):\n if backbone == 'auto':\n num_fe_weights = len([x for x in state_dict.keys() if 'feature_extractor.' in x])\n small = num_fe_weights < 1800\n\n ocr_f_down = [v for k, v in state_dict.items() if 'object_context_block.f_down.1.0.bias' in k]\n assert len(ocr_f_down) == 1\n ocr_width = ocr_f_down[0].shape[0]\n\n s2_conv1_w = [v for k, v in state_dict.items() if 'stage2.0.branches.0.0.conv1.weight' in k]\n assert len(s2_conv1_w) == 1\n width = s2_conv1_w[0].shape[0]\n\n model = get_hrnet_model(width=width, ocr_width=ocr_width, small=small,\n with_aux_output=False, cpu_dist_maps=cpu_dist_maps,\n norm_radius=norm_radius)\n\n model.load_state_dict(state_dict, strict=False)\n for param in model.parameters():\n param.requires_grad = False\n model.to(device)\n model.eval()\n\n return model\n\n\ndef load_deeplab_is_model(state_dict, device, backbone='auto', deeplab_ch=128, aspp_dropout=0.2,\n cpu_dist_maps=False, norm_radius=260):\n if backbone == 'auto':\n num_backbone_params = len([x for x in state_dict.keys()\n if 'feature_extractor.backbone' in x and not('num_batches_tracked' in x)])\n\n if num_backbone_params <= 181:\n backbone = 'resnet34'\n elif num_backbone_params <= 276:\n backbone = 'resnet50'\n elif num_backbone_params <= 531:\n backbone = 'resnet101'\n else:\n raise NotImplementedError('Unknown backbone')\n\n if 'aspp_dropout' in state_dict:\n aspp_dropout = float(state_dict['aspp_dropout'].cpu().numpy())\n else:\n aspp_project_weight = [v for k, v in state_dict.items() if 'aspp.project.0.weight' in k][0]\n deeplab_ch = aspp_project_weight.size(0)\n if deeplab_ch == 256:\n aspp_dropout = 0.5\n\n model = get_deeplab_model(backbone=backbone, deeplab_ch=deeplab_ch,\n aspp_dropout=aspp_dropout, cpu_dist_maps=cpu_dist_maps,\n norm_radius=norm_radius)\n\n model.load_state_dict(state_dict, strict=False)\n for param in model.parameters():\n param.requires_grad = False\n model.to(device)\n model.eval()\n\n return model\n\n\ndef get_iou(gt_mask, pred_mask, ignore_label=-1):\n ignore_gt_mask_inv = gt_mask != ignore_label\n obj_gt_mask = gt_mask == 1\n\n intersection = np.logical_and(np.logical_and(pred_mask, obj_gt_mask), ignore_gt_mask_inv).sum()\n union = np.logical_and(np.logical_or(pred_mask, obj_gt_mask), ignore_gt_mask_inv).sum()\n\n return intersection / union\n\n\ndef compute_noc_metric(all_ious, iou_thrs, max_clicks=20):\n def _get_noc(iou_arr, iou_thr):\n vals = iou_arr >= iou_thr\n return np.argmax(vals) + 1 if np.any(vals) else max_clicks\n\n noc_list = []\n over_max_list = []\n for iou_thr in iou_thrs:\n scores_arr = np.array([_get_noc(iou_arr, iou_thr)\n for iou_arr in all_ious], dtype=np.int32)\n\n score = scores_arr.mean()\n over_max = (scores_arr == max_clicks).sum()\n\n noc_list.append(score)\n over_max_list.append(over_max)\n\n return noc_list, over_max_list\n\n\ndef find_checkpoint(weights_folder, checkpoint_name):\n weights_folder = Path(weights_folder)\n if ':' in checkpoint_name:\n model_name, checkpoint_name = checkpoint_name.split(':')\n models_candidates = [x for x in weights_folder.glob(f'{model_name}*') if x.is_dir()]\n assert len(models_candidates) == 1\n model_folder = models_candidates[0]\n else:\n model_folder = weights_folder\n\n if checkpoint_name.endswith('.pth'):\n if Path(checkpoint_name).exists():\n checkpoint_path = checkpoint_name\n else:\n checkpoint_path = weights_folder / checkpoint_name\n else:\n model_checkpoints = list(model_folder.rglob(f'{checkpoint_name}*.pth'))\n assert len(model_checkpoints) == 1\n checkpoint_path = model_checkpoints[0]\n\n return str(checkpoint_path)\n\n\ndef get_results_table(noc_list, over_max_list, brs_type, dataset_name, mean_spc, elapsed_time,\n n_clicks=20, model_name=None):\n table_header = (f'|{\"BRS Type\":^13}|{\"Dataset\":^11}|'\n f'{\"NoC@80%\":^9}|{\"NoC@85%\":^9}|{\"NoC@90%\":^9}|'\n f'{\">=\"+str(n_clicks)+\"@85%\":^9}|{\">=\"+str(n_clicks)+\"@90%\":^9}|'\n f'{\"SPC,s\":^7}|{\"Time\":^9}|')\n row_width = len(table_header)\n\n header = f'Eval results for model: {model_name}\\n' if model_name is not None else ''\n header += '-' * row_width + '\\n'\n header += table_header + '\\n' + '-' * row_width\n\n eval_time = str(timedelta(seconds=int(elapsed_time)))\n table_row = f'|{brs_type:^13}|{dataset_name:^11}|'\n table_row += f'{noc_list[0]:^9.2f}|'\n table_row += f'{noc_list[1]:^9.2f}|' if len(noc_list) > 1 else f'{\"?\":^9}|'\n table_row += f'{noc_list[2]:^9.2f}|' if len(noc_list) > 2 else f'{\"?\":^9}|'\n table_row += f'{over_max_list[1]:^9}|' if len(noc_list) > 1 else f'{\"?\":^9}|'\n table_row += f'{over_max_list[2]:^9}|' if len(noc_list) > 2 else f'{\"?\":^9}|'\n table_row += f'{mean_spc:^7.3f}|{eval_time:^9}|'\n\n return header, table_row" }, { "path": "XMem/inference/interact/fbrs/model/__init__.py", "content": "" }, { "path": "XMem/inference/interact/fbrs/model/initializer.py", "content": "import torch\nimport torch.nn as nn\nimport numpy as np\n\n\nclass Initializer(object):\n def __init__(self, local_init=True, gamma=None):\n self.local_init = local_init\n self.gamma = gamma\n\n def __call__(self, m):\n if getattr(m, '__initialized', False):\n return\n\n if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d,\n nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d,\n nn.GroupNorm, nn.SyncBatchNorm)) or 'BatchNorm' in m.__class__.__name__:\n if m.weight is not None:\n self._init_gamma(m.weight.data)\n if m.bias is not None:\n self._init_beta(m.bias.data)\n else:\n if getattr(m, 'weight', None) is not None:\n self._init_weight(m.weight.data)\n if getattr(m, 'bias', None) is not None:\n self._init_bias(m.bias.data)\n\n if self.local_init:\n object.__setattr__(m, '__initialized', True)\n\n def _init_weight(self, data):\n nn.init.uniform_(data, -0.07, 0.07)\n\n def _init_bias(self, data):\n nn.init.constant_(data, 0)\n\n def _init_gamma(self, data):\n if self.gamma is None:\n nn.init.constant_(data, 1.0)\n else:\n nn.init.normal_(data, 1.0, self.gamma)\n\n def _init_beta(self, data):\n nn.init.constant_(data, 0)\n\n\nclass Bilinear(Initializer):\n def __init__(self, scale, groups, in_channels, **kwargs):\n super().__init__(**kwargs)\n self.scale = scale\n self.groups = groups\n self.in_channels = in_channels\n\n def _init_weight(self, data):\n \"\"\"Reset the weight and bias.\"\"\"\n bilinear_kernel = self.get_bilinear_kernel(self.scale)\n weight = torch.zeros_like(data)\n for i in range(self.in_channels):\n if self.groups == 1:\n j = i\n else:\n j = 0\n weight[i, j] = bilinear_kernel\n data[:] = weight\n\n @staticmethod\n def get_bilinear_kernel(scale):\n \"\"\"Generate a bilinear upsampling kernel.\"\"\"\n kernel_size = 2 * scale - scale % 2\n scale = (kernel_size + 1) // 2\n center = scale - 0.5 * (1 + kernel_size % 2)\n\n og = np.ogrid[:kernel_size, :kernel_size]\n kernel = (1 - np.abs(og[0] - center) / scale) * (1 - np.abs(og[1] - center) / scale)\n\n return torch.tensor(kernel, dtype=torch.float32)\n\n\nclass XavierGluon(Initializer):\n def __init__(self, rnd_type='uniform', factor_type='avg', magnitude=3, **kwargs):\n super().__init__(**kwargs)\n\n self.rnd_type = rnd_type\n self.factor_type = factor_type\n self.magnitude = float(magnitude)\n\n def _init_weight(self, arr):\n fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(arr)\n\n if self.factor_type == 'avg':\n factor = (fan_in + fan_out) / 2.0\n elif self.factor_type == 'in':\n factor = fan_in\n elif self.factor_type == 'out':\n factor = fan_out\n else:\n raise ValueError('Incorrect factor type')\n scale = np.sqrt(self.magnitude / factor)\n\n if self.rnd_type == 'uniform':\n nn.init.uniform_(arr, -scale, scale)\n elif self.rnd_type == 'gaussian':\n nn.init.normal_(arr, 0, scale)\n else:\n raise ValueError('Unknown random type')\n" }, { "path": "XMem/inference/interact/fbrs/model/is_deeplab_model.py", "content": "import torch\nimport torch.nn as nn\n\nfrom .ops import DistMaps\nfrom .modeling.deeplab_v3 import DeepLabV3Plus\nfrom .modeling.basic_blocks import SepConvHead\n\n\ndef get_deeplab_model(backbone='resnet50', deeplab_ch=256, aspp_dropout=0.5,\n norm_layer=nn.BatchNorm2d, backbone_norm_layer=None,\n use_rgb_conv=True, cpu_dist_maps=False,\n norm_radius=260):\n model = DistMapsModel(\n feature_extractor=DeepLabV3Plus(backbone=backbone,\n ch=deeplab_ch,\n project_dropout=aspp_dropout,\n norm_layer=norm_layer,\n backbone_norm_layer=backbone_norm_layer),\n head=SepConvHead(1, in_channels=deeplab_ch, mid_channels=deeplab_ch // 2,\n num_layers=2, norm_layer=norm_layer),\n use_rgb_conv=use_rgb_conv,\n norm_layer=norm_layer,\n norm_radius=norm_radius,\n cpu_dist_maps=cpu_dist_maps\n )\n\n return model\n\n\nclass DistMapsModel(nn.Module):\n def __init__(self, feature_extractor, head, norm_layer=nn.BatchNorm2d, use_rgb_conv=True,\n cpu_dist_maps=False, norm_radius=260):\n super(DistMapsModel, self).__init__()\n\n if use_rgb_conv:\n self.rgb_conv = nn.Sequential(\n nn.Conv2d(in_channels=5, out_channels=8, kernel_size=1),\n nn.LeakyReLU(negative_slope=0.2),\n norm_layer(8),\n nn.Conv2d(in_channels=8, out_channels=3, kernel_size=1),\n )\n else:\n self.rgb_conv = None\n\n self.dist_maps = DistMaps(norm_radius=norm_radius, spatial_scale=1.0,\n cpu_mode=cpu_dist_maps)\n self.feature_extractor = feature_extractor\n self.head = head\n\n def forward(self, image, points):\n coord_features = self.dist_maps(image, points)\n\n if self.rgb_conv is not None:\n x = self.rgb_conv(torch.cat((image, coord_features), dim=1))\n else:\n c1, c2 = torch.chunk(coord_features, 2, dim=1)\n c3 = torch.ones_like(c1)\n coord_features = torch.cat((c1, c2, c3), dim=1)\n x = 0.8 * image * coord_features + 0.2 * image\n\n backbone_features = self.feature_extractor(x)\n instance_out = self.head(backbone_features[0])\n instance_out = nn.functional.interpolate(instance_out, size=image.size()[2:],\n mode='bilinear', align_corners=True)\n\n return {'instances': instance_out}\n\n def load_weights(self, path_to_weights):\n current_state_dict = self.state_dict()\n new_state_dict = torch.load(path_to_weights, map_location='cpu')\n current_state_dict.update(new_state_dict)\n self.load_state_dict(current_state_dict)\n\n def get_trainable_params(self):\n backbone_params = nn.ParameterList()\n other_params = nn.ParameterList()\n\n for name, param in self.named_parameters():\n if param.requires_grad:\n if 'backbone' in name:\n backbone_params.append(param)\n else:\n other_params.append(param)\n return backbone_params, other_params\n\n\n" }, { "path": "XMem/inference/interact/fbrs/model/is_hrnet_model.py", "content": "import torch\nimport torch.nn as nn\n\nfrom .ops import DistMaps\nfrom .modeling.hrnet_ocr import HighResolutionNet\n\n\ndef get_hrnet_model(width=48, ocr_width=256, small=False, norm_radius=260,\n use_rgb_conv=True, with_aux_output=False, cpu_dist_maps=False,\n norm_layer=nn.BatchNorm2d):\n model = DistMapsHRNetModel(\n feature_extractor=HighResolutionNet(width=width, ocr_width=ocr_width, small=small,\n num_classes=1, norm_layer=norm_layer),\n use_rgb_conv=use_rgb_conv,\n with_aux_output=with_aux_output,\n norm_layer=norm_layer,\n norm_radius=norm_radius,\n cpu_dist_maps=cpu_dist_maps\n )\n\n return model\n\n\nclass DistMapsHRNetModel(nn.Module):\n def __init__(self, feature_extractor, use_rgb_conv=True, with_aux_output=False,\n norm_layer=nn.BatchNorm2d, norm_radius=260, cpu_dist_maps=False):\n super(DistMapsHRNetModel, self).__init__()\n self.with_aux_output = with_aux_output\n\n if use_rgb_conv:\n self.rgb_conv = nn.Sequential(\n nn.Conv2d(in_channels=5, out_channels=8, kernel_size=1),\n nn.LeakyReLU(negative_slope=0.2),\n norm_layer(8),\n nn.Conv2d(in_channels=8, out_channels=3, kernel_size=1),\n )\n else:\n self.rgb_conv = None\n\n self.dist_maps = DistMaps(norm_radius=norm_radius, spatial_scale=1.0, cpu_mode=cpu_dist_maps)\n self.feature_extractor = feature_extractor\n\n def forward(self, image, points):\n coord_features = self.dist_maps(image, points)\n\n if self.rgb_conv is not None:\n x = self.rgb_conv(torch.cat((image, coord_features), dim=1))\n else:\n c1, c2 = torch.chunk(coord_features, 2, dim=1)\n c3 = torch.ones_like(c1)\n coord_features = torch.cat((c1, c2, c3), dim=1)\n x = 0.8 * image * coord_features + 0.2 * image\n\n feature_extractor_out = self.feature_extractor(x)\n instance_out = feature_extractor_out[0]\n instance_out = nn.functional.interpolate(instance_out, size=image.size()[2:],\n mode='bilinear', align_corners=True)\n outputs = {'instances': instance_out}\n if self.with_aux_output:\n instance_aux_out = feature_extractor_out[1]\n instance_aux_out = nn.functional.interpolate(instance_aux_out, size=image.size()[2:],\n mode='bilinear', align_corners=True)\n outputs['instances_aux'] = instance_aux_out\n\n return outputs\n\n def load_weights(self, path_to_weights):\n current_state_dict = self.state_dict()\n new_state_dict = torch.load(path_to_weights)\n current_state_dict.update(new_state_dict)\n self.load_state_dict(current_state_dict)\n\n def get_trainable_params(self):\n backbone_params = nn.ParameterList()\n other_params = nn.ParameterList()\n other_params_keys = []\n nonbackbone_keywords = ['rgb_conv', 'aux_head', 'cls_head', 'conv3x3_ocr', 'ocr_distri_head']\n\n for name, param in self.named_parameters():\n if param.requires_grad:\n if any(x in name for x in nonbackbone_keywords):\n other_params.append(param)\n other_params_keys.append(name)\n else:\n backbone_params.append(param)\n print('Nonbackbone params:', sorted(other_params_keys))\n return backbone_params, other_params\n" }, { "path": "XMem/inference/interact/fbrs/model/losses.py", "content": "import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom ..utils import misc\n\n\nclass NormalizedFocalLossSigmoid(nn.Module):\n def __init__(self, axis=-1, alpha=0.25, gamma=2,\n from_logits=False, batch_axis=0,\n weight=None, size_average=True, detach_delimeter=True,\n eps=1e-12, scale=1.0,\n ignore_label=-1):\n super(NormalizedFocalLossSigmoid, self).__init__()\n self._axis = axis\n self._alpha = alpha\n self._gamma = gamma\n self._ignore_label = ignore_label\n self._weight = weight if weight is not None else 1.0\n self._batch_axis = batch_axis\n\n self._scale = scale\n self._from_logits = from_logits\n self._eps = eps\n self._size_average = size_average\n self._detach_delimeter = detach_delimeter\n self._k_sum = 0\n\n def forward(self, pred, label, sample_weight=None):\n one_hot = label > 0\n sample_weight = label != self._ignore_label\n\n if not self._from_logits:\n pred = torch.sigmoid(pred)\n\n alpha = torch.where(one_hot, self._alpha * sample_weight, (1 - self._alpha) * sample_weight)\n pt = torch.where(one_hot, pred, 1 - pred)\n pt = torch.where(sample_weight, pt, torch.ones_like(pt))\n\n beta = (1 - pt) ** self._gamma\n\n sw_sum = torch.sum(sample_weight, dim=(-2, -1), keepdim=True)\n beta_sum = torch.sum(beta, dim=(-2, -1), keepdim=True)\n mult = sw_sum / (beta_sum + self._eps)\n if self._detach_delimeter:\n mult = mult.detach()\n beta = beta * mult\n\n ignore_area = torch.sum(label == self._ignore_label, dim=tuple(range(1, label.dim()))).cpu().numpy()\n sample_mult = torch.mean(mult, dim=tuple(range(1, mult.dim()))).cpu().numpy()\n if np.any(ignore_area == 0):\n self._k_sum = 0.9 * self._k_sum + 0.1 * sample_mult[ignore_area == 0].mean()\n\n loss = -alpha * beta * torch.log(torch.min(pt + self._eps, torch.ones(1, dtype=torch.float).to(pt.device)))\n loss = self._weight * (loss * sample_weight)\n\n if self._size_average:\n bsum = torch.sum(sample_weight, dim=misc.get_dims_with_exclusion(sample_weight.dim(), self._batch_axis))\n loss = torch.sum(loss, dim=misc.get_dims_with_exclusion(loss.dim(), self._batch_axis)) / (bsum + self._eps)\n else:\n loss = torch.sum(loss, dim=misc.get_dims_with_exclusion(loss.dim(), self._batch_axis))\n\n return self._scale * loss\n\n def log_states(self, sw, name, global_step):\n sw.add_scalar(tag=name + '_k', value=self._k_sum, global_step=global_step)\n\n\nclass FocalLoss(nn.Module):\n def __init__(self, axis=-1, alpha=0.25, gamma=2,\n from_logits=False, batch_axis=0,\n weight=None, num_class=None,\n eps=1e-9, size_average=True, scale=1.0):\n super(FocalLoss, self).__init__()\n self._axis = axis\n self._alpha = alpha\n self._gamma = gamma\n self._weight = weight if weight is not None else 1.0\n self._batch_axis = batch_axis\n\n self._scale = scale\n self._num_class = num_class\n self._from_logits = from_logits\n self._eps = eps\n self._size_average = size_average\n\n def forward(self, pred, label, sample_weight=None):\n if not self._from_logits:\n pred = F.sigmoid(pred)\n\n one_hot = label > 0\n pt = torch.where(one_hot, pred, 1 - pred)\n\n t = label != -1\n alpha = torch.where(one_hot, self._alpha * t, (1 - self._alpha) * t)\n beta = (1 - pt) ** self._gamma\n\n loss = -alpha * beta * torch.log(torch.min(pt + self._eps, torch.ones(1, dtype=torch.float).to(pt.device)))\n sample_weight = label != -1\n\n loss = self._weight * (loss * sample_weight)\n\n if self._size_average:\n tsum = torch.sum(label == 1, dim=misc.get_dims_with_exclusion(label.dim(), self._batch_axis))\n loss = torch.sum(loss, dim=misc.get_dims_with_exclusion(loss.dim(), self._batch_axis)) / (tsum + self._eps)\n else:\n loss = torch.sum(loss, dim=misc.get_dims_with_exclusion(loss.dim(), self._batch_axis))\n\n return self._scale * loss\n\n\nclass SigmoidBinaryCrossEntropyLoss(nn.Module):\n def __init__(self, from_sigmoid=False, weight=None, batch_axis=0, ignore_label=-1):\n super(SigmoidBinaryCrossEntropyLoss, self).__init__()\n self._from_sigmoid = from_sigmoid\n self._ignore_label = ignore_label\n self._weight = weight if weight is not None else 1.0\n self._batch_axis = batch_axis\n\n def forward(self, pred, label):\n label = label.view(pred.size())\n sample_weight = label != self._ignore_label\n label = torch.where(sample_weight, label, torch.zeros_like(label))\n\n if not self._from_sigmoid:\n loss = torch.relu(pred) - pred * label + F.softplus(-torch.abs(pred))\n else:\n eps = 1e-12\n loss = -(torch.log(pred + eps) * label\n + torch.log(1. - pred + eps) * (1. - label))\n\n loss = self._weight * (loss * sample_weight)\n return torch.mean(loss, dim=misc.get_dims_with_exclusion(loss.dim(), self._batch_axis))\n" }, { "path": "XMem/inference/interact/fbrs/model/metrics.py", "content": "import torch\nimport numpy as np\n\nfrom ..utils import misc\n\n\nclass TrainMetric(object):\n def __init__(self, pred_outputs, gt_outputs):\n self.pred_outputs = pred_outputs\n self.gt_outputs = gt_outputs\n\n def update(self, *args, **kwargs):\n raise NotImplementedError\n\n def get_epoch_value(self):\n raise NotImplementedError\n\n def reset_epoch_stats(self):\n raise NotImplementedError\n\n def log_states(self, sw, tag_prefix, global_step):\n pass\n\n @property\n def name(self):\n return type(self).__name__\n\n\nclass AdaptiveIoU(TrainMetric):\n def __init__(self, init_thresh=0.4, thresh_step=0.025, thresh_beta=0.99, iou_beta=0.9,\n ignore_label=-1, from_logits=True,\n pred_output='instances', gt_output='instances'):\n super().__init__(pred_outputs=(pred_output,), gt_outputs=(gt_output,))\n self._ignore_label = ignore_label\n self._from_logits = from_logits\n self._iou_thresh = init_thresh\n self._thresh_step = thresh_step\n self._thresh_beta = thresh_beta\n self._iou_beta = iou_beta\n self._ema_iou = 0.0\n self._epoch_iou_sum = 0.0\n self._epoch_batch_count = 0\n\n def update(self, pred, gt):\n gt_mask = gt > 0\n if self._from_logits:\n pred = torch.sigmoid(pred)\n\n gt_mask_area = torch.sum(gt_mask, dim=(1, 2)).detach().cpu().numpy()\n if np.all(gt_mask_area == 0):\n return\n\n ignore_mask = gt == self._ignore_label\n max_iou = _compute_iou(pred > self._iou_thresh, gt_mask, ignore_mask).mean()\n best_thresh = self._iou_thresh\n for t in [best_thresh - self._thresh_step, best_thresh + self._thresh_step]:\n temp_iou = _compute_iou(pred > t, gt_mask, ignore_mask).mean()\n if temp_iou > max_iou:\n max_iou = temp_iou\n best_thresh = t\n\n self._iou_thresh = self._thresh_beta * self._iou_thresh + (1 - self._thresh_beta) * best_thresh\n self._ema_iou = self._iou_beta * self._ema_iou + (1 - self._iou_beta) * max_iou\n self._epoch_iou_sum += max_iou\n self._epoch_batch_count += 1\n\n def get_epoch_value(self):\n if self._epoch_batch_count > 0:\n return self._epoch_iou_sum / self._epoch_batch_count\n else:\n return 0.0\n\n def reset_epoch_stats(self):\n self._epoch_iou_sum = 0.0\n self._epoch_batch_count = 0\n\n def log_states(self, sw, tag_prefix, global_step):\n sw.add_scalar(tag=tag_prefix + '_ema_iou', value=self._ema_iou, global_step=global_step)\n sw.add_scalar(tag=tag_prefix + '_iou_thresh', value=self._iou_thresh, global_step=global_step)\n\n @property\n def iou_thresh(self):\n return self._iou_thresh\n\n\ndef _compute_iou(pred_mask, gt_mask, ignore_mask=None, keep_ignore=False):\n if ignore_mask is not None:\n pred_mask = torch.where(ignore_mask, torch.zeros_like(pred_mask), pred_mask)\n\n reduction_dims = misc.get_dims_with_exclusion(gt_mask.dim(), 0)\n union = torch.mean((pred_mask | gt_mask).float(), dim=reduction_dims).detach().cpu().numpy()\n intersection = torch.mean((pred_mask & gt_mask).float(), dim=reduction_dims).detach().cpu().numpy()\n nonzero = union > 0\n\n iou = intersection[nonzero] / union[nonzero]\n if not keep_ignore:\n return iou\n else:\n result = np.full_like(intersection, -1)\n result[nonzero] = iou\n return result\n" }, { "path": "XMem/inference/interact/fbrs/model/modeling/__init__.py", "content": "" }, { "path": "XMem/inference/interact/fbrs/model/modeling/basic_blocks.py", "content": "import torch.nn as nn\n\nfrom ...model import ops\n\n\nclass ConvHead(nn.Module):\n def __init__(self, out_channels, in_channels=32, num_layers=1,\n kernel_size=3, padding=1,\n norm_layer=nn.BatchNorm2d):\n super(ConvHead, self).__init__()\n convhead = []\n\n for i in range(num_layers):\n convhead.extend([\n nn.Conv2d(in_channels, in_channels, kernel_size, padding=padding),\n nn.ReLU(),\n norm_layer(in_channels) if norm_layer is not None else nn.Identity()\n ])\n convhead.append(nn.Conv2d(in_channels, out_channels, 1, padding=0))\n\n self.convhead = nn.Sequential(*convhead)\n\n def forward(self, *inputs):\n return self.convhead(inputs[0])\n\n\nclass SepConvHead(nn.Module):\n def __init__(self, num_outputs, in_channels, mid_channels, num_layers=1,\n kernel_size=3, padding=1, dropout_ratio=0.0, dropout_indx=0,\n norm_layer=nn.BatchNorm2d):\n super(SepConvHead, self).__init__()\n\n sepconvhead = []\n\n for i in range(num_layers):\n sepconvhead.append(\n SeparableConv2d(in_channels=in_channels if i == 0 else mid_channels,\n out_channels=mid_channels,\n dw_kernel=kernel_size, dw_padding=padding,\n norm_layer=norm_layer, activation='relu')\n )\n if dropout_ratio > 0 and dropout_indx == i:\n sepconvhead.append(nn.Dropout(dropout_ratio))\n\n sepconvhead.append(\n nn.Conv2d(in_channels=mid_channels, out_channels=num_outputs, kernel_size=1, padding=0)\n )\n\n self.layers = nn.Sequential(*sepconvhead)\n\n def forward(self, *inputs):\n x = inputs[0]\n\n return self.layers(x)\n\n\nclass SeparableConv2d(nn.Module):\n def __init__(self, in_channels, out_channels, dw_kernel, dw_padding, dw_stride=1,\n activation=None, use_bias=False, norm_layer=None):\n super(SeparableConv2d, self).__init__()\n _activation = ops.select_activation_function(activation)\n self.body = nn.Sequential(\n nn.Conv2d(in_channels, in_channels, kernel_size=dw_kernel, stride=dw_stride,\n padding=dw_padding, bias=use_bias, groups=in_channels),\n nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=use_bias),\n norm_layer(out_channels) if norm_layer is not None else nn.Identity(),\n _activation()\n )\n\n def forward(self, x):\n return self.body(x)\n" }, { "path": "XMem/inference/interact/fbrs/model/modeling/deeplab_v3.py", "content": "from contextlib import ExitStack\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\n\nfrom .basic_blocks import SeparableConv2d\nfrom .resnet import ResNetBackbone\nfrom ...model import ops\n\n\nclass DeepLabV3Plus(nn.Module):\n def __init__(self, backbone='resnet50', norm_layer=nn.BatchNorm2d,\n backbone_norm_layer=None,\n ch=256,\n project_dropout=0.5,\n inference_mode=False,\n **kwargs):\n super(DeepLabV3Plus, self).__init__()\n if backbone_norm_layer is None:\n backbone_norm_layer = norm_layer\n\n self.backbone_name = backbone\n self.norm_layer = norm_layer\n self.backbone_norm_layer = backbone_norm_layer\n self.inference_mode = False\n self.ch = ch\n self.aspp_in_channels = 2048\n self.skip_project_in_channels = 256 # layer 1 out_channels\n\n self._kwargs = kwargs\n if backbone == 'resnet34':\n self.aspp_in_channels = 512\n self.skip_project_in_channels = 64\n\n self.backbone = ResNetBackbone(backbone=self.backbone_name, pretrained_base=False,\n norm_layer=self.backbone_norm_layer, **kwargs)\n\n self.head = _DeepLabHead(in_channels=ch + 32, mid_channels=ch, out_channels=ch,\n norm_layer=self.norm_layer)\n self.skip_project = _SkipProject(self.skip_project_in_channels, 32, norm_layer=self.norm_layer)\n self.aspp = _ASPP(in_channels=self.aspp_in_channels,\n atrous_rates=[12, 24, 36],\n out_channels=ch,\n project_dropout=project_dropout,\n norm_layer=self.norm_layer)\n\n if inference_mode:\n self.set_prediction_mode()\n\n def load_pretrained_weights(self):\n pretrained = ResNetBackbone(backbone=self.backbone_name, pretrained_base=True,\n norm_layer=self.backbone_norm_layer, **self._kwargs)\n backbone_state_dict = self.backbone.state_dict()\n pretrained_state_dict = pretrained.state_dict()\n\n backbone_state_dict.update(pretrained_state_dict)\n self.backbone.load_state_dict(backbone_state_dict)\n\n if self.inference_mode:\n for param in self.backbone.parameters():\n param.requires_grad = False\n\n def set_prediction_mode(self):\n self.inference_mode = True\n self.eval()\n\n def forward(self, x):\n with ExitStack() as stack:\n if self.inference_mode:\n stack.enter_context(torch.no_grad())\n\n c1, _, c3, c4 = self.backbone(x)\n c1 = self.skip_project(c1)\n\n x = self.aspp(c4)\n x = F.interpolate(x, c1.size()[2:], mode='bilinear', align_corners=True)\n x = torch.cat((x, c1), dim=1)\n x = self.head(x)\n\n return x,\n\n\nclass _SkipProject(nn.Module):\n def __init__(self, in_channels, out_channels, norm_layer=nn.BatchNorm2d):\n super(_SkipProject, self).__init__()\n _activation = ops.select_activation_function(\"relu\")\n\n self.skip_project = nn.Sequential(\n nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),\n norm_layer(out_channels),\n _activation()\n )\n\n def forward(self, x):\n return self.skip_project(x)\n\n\nclass _DeepLabHead(nn.Module):\n def __init__(self, out_channels, in_channels, mid_channels=256, norm_layer=nn.BatchNorm2d):\n super(_DeepLabHead, self).__init__()\n\n self.block = nn.Sequential(\n SeparableConv2d(in_channels=in_channels, out_channels=mid_channels, dw_kernel=3,\n dw_padding=1, activation='relu', norm_layer=norm_layer),\n SeparableConv2d(in_channels=mid_channels, out_channels=mid_channels, dw_kernel=3,\n dw_padding=1, activation='relu', norm_layer=norm_layer),\n nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=1)\n )\n\n def forward(self, x):\n return self.block(x)\n\n\nclass _ASPP(nn.Module):\n def __init__(self, in_channels, atrous_rates, out_channels=256,\n project_dropout=0.5, norm_layer=nn.BatchNorm2d):\n super(_ASPP, self).__init__()\n\n b0 = nn.Sequential(\n nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=False),\n norm_layer(out_channels),\n nn.ReLU()\n )\n\n rate1, rate2, rate3 = tuple(atrous_rates)\n b1 = _ASPPConv(in_channels, out_channels, rate1, norm_layer)\n b2 = _ASPPConv(in_channels, out_channels, rate2, norm_layer)\n b3 = _ASPPConv(in_channels, out_channels, rate3, norm_layer)\n b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer)\n\n self.concurent = nn.ModuleList([b0, b1, b2, b3, b4])\n\n project = [\n nn.Conv2d(in_channels=5*out_channels, out_channels=out_channels,\n kernel_size=1, bias=False),\n norm_layer(out_channels),\n nn.ReLU()\n ]\n if project_dropout > 0:\n project.append(nn.Dropout(project_dropout))\n self.project = nn.Sequential(*project)\n\n def forward(self, x):\n x = torch.cat([block(x) for block in self.concurent], dim=1)\n\n return self.project(x)\n\n\nclass _AsppPooling(nn.Module):\n def __init__(self, in_channels, out_channels, norm_layer):\n super(_AsppPooling, self).__init__()\n\n self.gap = nn.Sequential(\n nn.AdaptiveAvgPool2d((1, 1)),\n nn.Conv2d(in_channels=in_channels, out_channels=out_channels,\n kernel_size=1, bias=False),\n norm_layer(out_channels),\n nn.ReLU()\n )\n\n def forward(self, x):\n pool = self.gap(x)\n return F.interpolate(pool, x.size()[2:], mode='bilinear', align_corners=True)\n\n\ndef _ASPPConv(in_channels, out_channels, atrous_rate, norm_layer):\n block = nn.Sequential(\n nn.Conv2d(in_channels=in_channels, out_channels=out_channels,\n kernel_size=3, padding=atrous_rate,\n dilation=atrous_rate, bias=False),\n norm_layer(out_channels),\n nn.ReLU()\n )\n\n return block\n" }, { "path": "XMem/inference/interact/fbrs/model/modeling/hrnet_ocr.py", "content": "import os\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch._utils\nimport torch.nn.functional as F\nfrom .ocr import SpatialOCR_Module, SpatialGather_Module\nfrom .resnetv1b import BasicBlockV1b, BottleneckV1b\n\nrelu_inplace = True\n\n\nclass HighResolutionModule(nn.Module):\n def __init__(self, num_branches, blocks, num_blocks, num_inchannels,\n num_channels, fuse_method,multi_scale_output=True,\n norm_layer=nn.BatchNorm2d, align_corners=True):\n super(HighResolutionModule, self).__init__()\n self._check_branches(num_branches, num_blocks, num_inchannels, num_channels)\n\n self.num_inchannels = num_inchannels\n self.fuse_method = fuse_method\n self.num_branches = num_branches\n self.norm_layer = norm_layer\n self.align_corners = align_corners\n\n self.multi_scale_output = multi_scale_output\n\n self.branches = self._make_branches(\n num_branches, blocks, num_blocks, num_channels)\n self.fuse_layers = self._make_fuse_layers()\n self.relu = nn.ReLU(inplace=relu_inplace)\n\n def _check_branches(self, num_branches, num_blocks, num_inchannels, num_channels):\n if num_branches != len(num_blocks):\n error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(\n num_branches, len(num_blocks))\n raise ValueError(error_msg)\n\n if num_branches != len(num_channels):\n error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(\n num_branches, len(num_channels))\n raise ValueError(error_msg)\n\n if num_branches != len(num_inchannels):\n error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(\n num_branches, len(num_inchannels))\n raise ValueError(error_msg)\n\n def _make_one_branch(self, branch_index, block, num_blocks, num_channels,\n stride=1):\n downsample = None\n if stride != 1 or \\\n self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(self.num_inchannels[branch_index],\n num_channels[branch_index] * block.expansion,\n kernel_size=1, stride=stride, bias=False),\n self.norm_layer(num_channels[branch_index] * block.expansion),\n )\n\n layers = []\n layers.append(block(self.num_inchannels[branch_index],\n num_channels[branch_index], stride,\n downsample=downsample, norm_layer=self.norm_layer))\n self.num_inchannels[branch_index] = \\\n num_channels[branch_index] * block.expansion\n for i in range(1, num_blocks[branch_index]):\n layers.append(block(self.num_inchannels[branch_index],\n num_channels[branch_index],\n norm_layer=self.norm_layer))\n\n return nn.Sequential(*layers)\n\n def _make_branches(self, num_branches, block, num_blocks, num_channels):\n branches = []\n\n for i in range(num_branches):\n branches.append(\n self._make_one_branch(i, block, num_blocks, num_channels))\n\n return nn.ModuleList(branches)\n\n def _make_fuse_layers(self):\n if self.num_branches == 1:\n return None\n\n num_branches = self.num_branches\n num_inchannels = self.num_inchannels\n fuse_layers = []\n for i in range(num_branches if self.multi_scale_output else 1):\n fuse_layer = []\n for j in range(num_branches):\n if j > i:\n fuse_layer.append(nn.Sequential(\n nn.Conv2d(in_channels=num_inchannels[j],\n out_channels=num_inchannels[i],\n kernel_size=1,\n bias=False),\n self.norm_layer(num_inchannels[i])))\n elif j == i:\n fuse_layer.append(None)\n else:\n conv3x3s = []\n for k in range(i - j):\n if k == i - j - 1:\n num_outchannels_conv3x3 = num_inchannels[i]\n conv3x3s.append(nn.Sequential(\n nn.Conv2d(num_inchannels[j],\n num_outchannels_conv3x3,\n kernel_size=3, stride=2, padding=1, bias=False),\n self.norm_layer(num_outchannels_conv3x3)))\n else:\n num_outchannels_conv3x3 = num_inchannels[j]\n conv3x3s.append(nn.Sequential(\n nn.Conv2d(num_inchannels[j],\n num_outchannels_conv3x3,\n kernel_size=3, stride=2, padding=1, bias=False),\n self.norm_layer(num_outchannels_conv3x3),\n nn.ReLU(inplace=relu_inplace)))\n fuse_layer.append(nn.Sequential(*conv3x3s))\n fuse_layers.append(nn.ModuleList(fuse_layer))\n\n return nn.ModuleList(fuse_layers)\n\n def get_num_inchannels(self):\n return self.num_inchannels\n\n def forward(self, x):\n if self.num_branches == 1:\n return [self.branches[0](x[0])]\n\n for i in range(self.num_branches):\n x[i] = self.branches[i](x[i])\n\n x_fuse = []\n for i in range(len(self.fuse_layers)):\n y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])\n for j in range(1, self.num_branches):\n if i == j:\n y = y + x[j]\n elif j > i:\n width_output = x[i].shape[-1]\n height_output = x[i].shape[-2]\n y = y + F.interpolate(\n self.fuse_layers[i][j](x[j]),\n size=[height_output, width_output],\n mode='bilinear', align_corners=self.align_corners)\n else:\n y = y + self.fuse_layers[i][j](x[j])\n x_fuse.append(self.relu(y))\n\n return x_fuse\n\n\nclass HighResolutionNet(nn.Module):\n def __init__(self, width, num_classes, ocr_width=256, small=False,\n norm_layer=nn.BatchNorm2d, align_corners=True):\n super(HighResolutionNet, self).__init__()\n self.norm_layer = norm_layer\n self.width = width\n self.ocr_width = ocr_width\n self.align_corners = align_corners\n\n self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)\n self.bn1 = norm_layer(64)\n self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False)\n self.bn2 = norm_layer(64)\n self.relu = nn.ReLU(inplace=relu_inplace)\n\n num_blocks = 2 if small else 4\n\n stage1_num_channels = 64\n self.layer1 = self._make_layer(BottleneckV1b, 64, stage1_num_channels, blocks=num_blocks)\n stage1_out_channel = BottleneckV1b.expansion * stage1_num_channels\n\n self.stage2_num_branches = 2\n num_channels = [width, 2 * width]\n num_inchannels = [\n num_channels[i] * BasicBlockV1b.expansion for i in range(len(num_channels))]\n self.transition1 = self._make_transition_layer(\n [stage1_out_channel], num_inchannels)\n self.stage2, pre_stage_channels = self._make_stage(\n BasicBlockV1b, num_inchannels=num_inchannels, num_modules=1, num_branches=self.stage2_num_branches,\n num_blocks=2 * [num_blocks], num_channels=num_channels)\n\n self.stage3_num_branches = 3\n num_channels = [width, 2 * width, 4 * width]\n num_inchannels = [\n num_channels[i] * BasicBlockV1b.expansion for i in range(len(num_channels))]\n self.transition2 = self._make_transition_layer(\n pre_stage_channels, num_inchannels)\n self.stage3, pre_stage_channels = self._make_stage(\n BasicBlockV1b, num_inchannels=num_inchannels,\n num_modules=3 if small else 4, num_branches=self.stage3_num_branches,\n num_blocks=3 * [num_blocks], num_channels=num_channels)\n\n self.stage4_num_branches = 4\n num_channels = [width, 2 * width, 4 * width, 8 * width]\n num_inchannels = [\n num_channels[i] * BasicBlockV1b.expansion for i in range(len(num_channels))]\n self.transition3 = self._make_transition_layer(\n pre_stage_channels, num_inchannels)\n self.stage4, pre_stage_channels = self._make_stage(\n BasicBlockV1b, num_inchannels=num_inchannels, num_modules=2 if small else 3,\n num_branches=self.stage4_num_branches,\n num_blocks=4 * [num_blocks], num_channels=num_channels)\n\n last_inp_channels = np.int32(np.sum(pre_stage_channels))\n ocr_mid_channels = 2 * ocr_width\n ocr_key_channels = ocr_width\n\n self.conv3x3_ocr = nn.Sequential(\n nn.Conv2d(last_inp_channels, ocr_mid_channels,\n kernel_size=3, stride=1, padding=1),\n norm_layer(ocr_mid_channels),\n nn.ReLU(inplace=relu_inplace),\n )\n self.ocr_gather_head = SpatialGather_Module(num_classes)\n\n self.ocr_distri_head = SpatialOCR_Module(in_channels=ocr_mid_channels,\n key_channels=ocr_key_channels,\n out_channels=ocr_mid_channels,\n scale=1,\n dropout=0.05,\n norm_layer=norm_layer,\n align_corners=align_corners)\n self.cls_head = nn.Conv2d(\n ocr_mid_channels, num_classes, kernel_size=1, stride=1, padding=0, bias=True)\n\n self.aux_head = nn.Sequential(\n nn.Conv2d(last_inp_channels, last_inp_channels,\n kernel_size=1, stride=1, padding=0),\n norm_layer(last_inp_channels),\n nn.ReLU(inplace=relu_inplace),\n nn.Conv2d(last_inp_channels, num_classes,\n kernel_size=1, stride=1, padding=0, bias=True)\n )\n\n def _make_transition_layer(\n self, num_channels_pre_layer, num_channels_cur_layer):\n num_branches_cur = len(num_channels_cur_layer)\n num_branches_pre = len(num_channels_pre_layer)\n\n transition_layers = []\n for i in range(num_branches_cur):\n if i < num_branches_pre:\n if num_channels_cur_layer[i] != num_channels_pre_layer[i]:\n transition_layers.append(nn.Sequential(\n nn.Conv2d(num_channels_pre_layer[i],\n num_channels_cur_layer[i],\n kernel_size=3,\n stride=1,\n padding=1,\n bias=False),\n self.norm_layer(num_channels_cur_layer[i]),\n nn.ReLU(inplace=relu_inplace)))\n else:\n transition_layers.append(None)\n else:\n conv3x3s = []\n for j in range(i + 1 - num_branches_pre):\n inchannels = num_channels_pre_layer[-1]\n outchannels = num_channels_cur_layer[i] \\\n if j == i - num_branches_pre else inchannels\n conv3x3s.append(nn.Sequential(\n nn.Conv2d(inchannels, outchannels,\n kernel_size=3, stride=2, padding=1, bias=False),\n self.norm_layer(outchannels),\n nn.ReLU(inplace=relu_inplace)))\n transition_layers.append(nn.Sequential(*conv3x3s))\n\n return nn.ModuleList(transition_layers)\n\n def _make_layer(self, block, inplanes, planes, blocks, stride=1):\n downsample = None\n if stride != 1 or inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(inplanes, planes * block.expansion,\n kernel_size=1, stride=stride, bias=False),\n self.norm_layer(planes * block.expansion),\n )\n\n layers = []\n layers.append(block(inplanes, planes, stride,\n downsample=downsample, norm_layer=self.norm_layer))\n inplanes = planes * block.expansion\n for i in range(1, blocks):\n layers.append(block(inplanes, planes, norm_layer=self.norm_layer))\n\n return nn.Sequential(*layers)\n\n def _make_stage(self, block, num_inchannels,\n num_modules, num_branches, num_blocks, num_channels,\n fuse_method='SUM',\n multi_scale_output=True):\n modules = []\n for i in range(num_modules):\n # multi_scale_output is only used last module\n if not multi_scale_output and i == num_modules - 1:\n reset_multi_scale_output = False\n else:\n reset_multi_scale_output = True\n modules.append(\n HighResolutionModule(num_branches,\n block,\n num_blocks,\n num_inchannels,\n num_channels,\n fuse_method,\n reset_multi_scale_output,\n norm_layer=self.norm_layer,\n align_corners=self.align_corners)\n )\n num_inchannels = modules[-1].get_num_inchannels()\n\n return nn.Sequential(*modules), num_inchannels\n\n def forward(self, x):\n feats = self.compute_hrnet_feats(x)\n out_aux = self.aux_head(feats)\n feats = self.conv3x3_ocr(feats)\n\n context = self.ocr_gather_head(feats, out_aux)\n feats = self.ocr_distri_head(feats, context)\n out = self.cls_head(feats)\n\n return [out, out_aux]\n\n def compute_hrnet_feats(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.conv2(x)\n x = self.bn2(x)\n x = self.relu(x)\n x = self.layer1(x)\n\n x_list = []\n for i in range(self.stage2_num_branches):\n if self.transition1[i] is not None:\n x_list.append(self.transition1[i](x))\n else:\n x_list.append(x)\n y_list = self.stage2(x_list)\n\n x_list = []\n for i in range(self.stage3_num_branches):\n if self.transition2[i] is not None:\n if i < self.stage2_num_branches:\n x_list.append(self.transition2[i](y_list[i]))\n else:\n x_list.append(self.transition2[i](y_list[-1]))\n else:\n x_list.append(y_list[i])\n y_list = self.stage3(x_list)\n\n x_list = []\n for i in range(self.stage4_num_branches):\n if self.transition3[i] is not None:\n if i < self.stage3_num_branches:\n x_list.append(self.transition3[i](y_list[i]))\n else:\n x_list.append(self.transition3[i](y_list[-1]))\n else:\n x_list.append(y_list[i])\n x = self.stage4(x_list)\n\n # Upsampling\n x0_h, x0_w = x[0].size(2), x[0].size(3)\n x1 = F.interpolate(x[1], size=(x0_h, x0_w),\n mode='bilinear', align_corners=self.align_corners)\n x2 = F.interpolate(x[2], size=(x0_h, x0_w),\n mode='bilinear', align_corners=self.align_corners)\n x3 = F.interpolate(x[3], size=(x0_h, x0_w),\n mode='bilinear', align_corners=self.align_corners)\n\n return torch.cat([x[0], x1, x2, x3], 1)\n\n def load_pretrained_weights(self, pretrained_path=''):\n model_dict = self.state_dict()\n\n if not os.path.exists(pretrained_path):\n print(f'\\nFile \"{pretrained_path}\" does not exist.')\n print('You need to specify the correct path to the pre-trained weights.\\n'\n 'You can download the weights for HRNet from the repository:\\n'\n 'https://github.com/HRNet/HRNet-Image-Classification')\n exit(1)\n pretrained_dict = torch.load(pretrained_path, map_location={'cuda:0': 'cpu'})\n pretrained_dict = {k.replace('last_layer', 'aux_head').replace('model.', ''): v for k, v in\n pretrained_dict.items()}\n\n print('model_dict-pretrained_dict:', sorted(list(set(model_dict) - set(pretrained_dict))))\n print('pretrained_dict-model_dict:', sorted(list(set(pretrained_dict) - set(model_dict))))\n\n pretrained_dict = {k: v for k, v in pretrained_dict.items()\n if k in model_dict.keys()}\n\n model_dict.update(pretrained_dict)\n self.load_state_dict(model_dict)\n" }, { "path": "XMem/inference/interact/fbrs/model/modeling/ocr.py", "content": "import torch\nimport torch.nn as nn\nimport torch._utils\nimport torch.nn.functional as F\n\n\nclass SpatialGather_Module(nn.Module):\n \"\"\"\n Aggregate the context features according to the initial\n predicted probability distribution.\n Employ the soft-weighted method to aggregate the context.\n \"\"\"\n\n def __init__(self, cls_num=0, scale=1):\n super(SpatialGather_Module, self).__init__()\n self.cls_num = cls_num\n self.scale = scale\n\n def forward(self, feats, probs):\n batch_size, c, h, w = probs.size(0), probs.size(1), probs.size(2), probs.size(3)\n probs = probs.view(batch_size, c, -1)\n feats = feats.view(batch_size, feats.size(1), -1)\n feats = feats.permute(0, 2, 1) # batch x hw x c\n probs = F.softmax(self.scale * probs, dim=2) # batch x k x hw\n ocr_context = torch.matmul(probs, feats) \\\n .permute(0, 2, 1).unsqueeze(3) # batch x k x c\n return ocr_context\n\n\nclass SpatialOCR_Module(nn.Module):\n \"\"\"\n Implementation of the OCR module:\n We aggregate the global object representation to update the representation for each pixel.\n \"\"\"\n\n def __init__(self,\n in_channels,\n key_channels,\n out_channels,\n scale=1,\n dropout=0.1,\n norm_layer=nn.BatchNorm2d,\n align_corners=True):\n super(SpatialOCR_Module, self).__init__()\n self.object_context_block = ObjectAttentionBlock2D(in_channels, key_channels, scale,\n norm_layer, align_corners)\n _in_channels = 2 * in_channels\n\n self.conv_bn_dropout = nn.Sequential(\n nn.Conv2d(_in_channels, out_channels, kernel_size=1, padding=0, bias=False),\n nn.Sequential(norm_layer(out_channels), nn.ReLU(inplace=True)),\n nn.Dropout2d(dropout)\n )\n\n def forward(self, feats, proxy_feats):\n context = self.object_context_block(feats, proxy_feats)\n\n output = self.conv_bn_dropout(torch.cat([context, feats], 1))\n\n return output\n\n\nclass ObjectAttentionBlock2D(nn.Module):\n '''\n The basic implementation for object context block\n Input:\n N X C X H X W\n Parameters:\n in_channels : the dimension of the input feature map\n key_channels : the dimension after the key/query transform\n scale : choose the scale to downsample the input feature maps (save memory cost)\n bn_type : specify the bn type\n Return:\n N X C X H X W\n '''\n\n def __init__(self,\n in_channels,\n key_channels,\n scale=1,\n norm_layer=nn.BatchNorm2d,\n align_corners=True):\n super(ObjectAttentionBlock2D, self).__init__()\n self.scale = scale\n self.in_channels = in_channels\n self.key_channels = key_channels\n self.align_corners = align_corners\n\n self.pool = nn.MaxPool2d(kernel_size=(scale, scale))\n self.f_pixel = nn.Sequential(\n nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,\n kernel_size=1, stride=1, padding=0, bias=False),\n nn.Sequential(norm_layer(self.key_channels), nn.ReLU(inplace=True)),\n nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,\n kernel_size=1, stride=1, padding=0, bias=False),\n nn.Sequential(norm_layer(self.key_channels), nn.ReLU(inplace=True))\n )\n self.f_object = nn.Sequential(\n nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,\n kernel_size=1, stride=1, padding=0, bias=False),\n nn.Sequential(norm_layer(self.key_channels), nn.ReLU(inplace=True)),\n nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,\n kernel_size=1, stride=1, padding=0, bias=False),\n nn.Sequential(norm_layer(self.key_channels), nn.ReLU(inplace=True))\n )\n self.f_down = nn.Sequential(\n nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,\n kernel_size=1, stride=1, padding=0, bias=False),\n nn.Sequential(norm_layer(self.key_channels), nn.ReLU(inplace=True))\n )\n self.f_up = nn.Sequential(\n nn.Conv2d(in_channels=self.key_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0, bias=False),\n nn.Sequential(norm_layer(self.in_channels), nn.ReLU(inplace=True))\n )\n\n def forward(self, x, proxy):\n batch_size, h, w = x.size(0), x.size(2), x.size(3)\n if self.scale > 1:\n x = self.pool(x)\n\n query = self.f_pixel(x).view(batch_size, self.key_channels, -1)\n query = query.permute(0, 2, 1)\n key = self.f_object(proxy).view(batch_size, self.key_channels, -1)\n value = self.f_down(proxy).view(batch_size, self.key_channels, -1)\n value = value.permute(0, 2, 1)\n\n sim_map = torch.matmul(query, key)\n sim_map = (self.key_channels ** -.5) * sim_map\n sim_map = F.softmax(sim_map, dim=-1)\n\n # add bg context ...\n context = torch.matmul(sim_map, value)\n context = context.permute(0, 2, 1).contiguous()\n context = context.view(batch_size, self.key_channels, *x.size()[2:])\n context = self.f_up(context)\n if self.scale > 1:\n context = F.interpolate(input=context, size=(h, w),\n mode='bilinear', align_corners=self.align_corners)\n\n return context\n" }, { "path": "XMem/inference/interact/fbrs/model/modeling/resnet.py", "content": "import torch\nfrom .resnetv1b import resnet34_v1b, resnet50_v1s, resnet101_v1s, resnet152_v1s\n\n\nclass ResNetBackbone(torch.nn.Module):\n def __init__(self, backbone='resnet50', pretrained_base=True, dilated=True, **kwargs):\n super(ResNetBackbone, self).__init__()\n\n if backbone == 'resnet34':\n pretrained = resnet34_v1b(pretrained=pretrained_base, dilated=dilated, **kwargs)\n elif backbone == 'resnet50':\n pretrained = resnet50_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)\n elif backbone == 'resnet101':\n pretrained = resnet101_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)\n elif backbone == 'resnet152':\n pretrained = resnet152_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)\n else:\n raise RuntimeError(f'unknown backbone: {backbone}')\n\n self.conv1 = pretrained.conv1\n self.bn1 = pretrained.bn1\n self.relu = pretrained.relu\n self.maxpool = pretrained.maxpool\n self.layer1 = pretrained.layer1\n self.layer2 = pretrained.layer2\n self.layer3 = pretrained.layer3\n self.layer4 = pretrained.layer4\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n c1 = self.layer1(x)\n c2 = self.layer2(c1)\n c3 = self.layer3(c2)\n c4 = self.layer4(c3)\n\n return c1, c2, c3, c4\n" }, { "path": "XMem/inference/interact/fbrs/model/modeling/resnetv1b.py", "content": "import torch\nimport torch.nn as nn\nGLUON_RESNET_TORCH_HUB = 'rwightman/pytorch-pretrained-gluonresnet'\n\n\nclass BasicBlockV1b(nn.Module):\n expansion = 1\n\n def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None,\n previous_dilation=1, norm_layer=nn.BatchNorm2d):\n super(BasicBlockV1b, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride,\n padding=dilation, dilation=dilation, bias=False)\n self.bn1 = norm_layer(planes)\n self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1,\n padding=previous_dilation, dilation=previous_dilation, bias=False)\n self.bn2 = norm_layer(planes)\n\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out = out + residual\n out = self.relu(out)\n\n return out\n\n\nclass BottleneckV1b(nn.Module):\n expansion = 4\n\n def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None,\n previous_dilation=1, norm_layer=nn.BatchNorm2d):\n super(BottleneckV1b, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n self.bn1 = norm_layer(planes)\n\n self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,\n padding=dilation, dilation=dilation, bias=False)\n self.bn2 = norm_layer(planes)\n\n self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)\n self.bn3 = norm_layer(planes * self.expansion)\n\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out = out + residual\n out = self.relu(out)\n\n return out\n\n\nclass ResNetV1b(nn.Module):\n \"\"\" Pre-trained ResNetV1b Model, which produces the strides of 8 featuremaps at conv5.\n\n Parameters\n ----------\n block : Block\n Class for the residual block. Options are BasicBlockV1, BottleneckV1.\n layers : list of int\n Numbers of layers in each block\n classes : int, default 1000\n Number of classification classes.\n dilated : bool, default False\n Applying dilation strategy to pretrained ResNet yielding a stride-8 model,\n typically used in Semantic Segmentation.\n norm_layer : object\n Normalization layer used (default: :class:`nn.BatchNorm2d`)\n deep_stem : bool, default False\n Whether to replace the 7x7 conv1 with 3 3x3 convolution layers.\n avg_down : bool, default False\n Whether to use average pooling for projection skip connection between stages/downsample.\n final_drop : float, default 0.0\n Dropout ratio before the final classification layer.\n\n Reference:\n - He, Kaiming, et al. \"Deep residual learning for image recognition.\"\n Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.\n\n - Yu, Fisher, and Vladlen Koltun. \"Multi-scale context aggregation by dilated convolutions.\"\n \"\"\"\n def __init__(self, block, layers, classes=1000, dilated=True, deep_stem=False, stem_width=32,\n avg_down=False, final_drop=0.0, norm_layer=nn.BatchNorm2d):\n self.inplanes = stem_width*2 if deep_stem else 64\n super(ResNetV1b, self).__init__()\n if not deep_stem:\n self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)\n else:\n self.conv1 = nn.Sequential(\n nn.Conv2d(3, stem_width, kernel_size=3, stride=2, padding=1, bias=False),\n norm_layer(stem_width),\n nn.ReLU(True),\n nn.Conv2d(stem_width, stem_width, kernel_size=3, stride=1, padding=1, bias=False),\n norm_layer(stem_width),\n nn.ReLU(True),\n nn.Conv2d(stem_width, 2*stem_width, kernel_size=3, stride=1, padding=1, bias=False)\n )\n self.bn1 = norm_layer(self.inplanes)\n self.relu = nn.ReLU(True)\n self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)\n self.layer1 = self._make_layer(block, 64, layers[0], avg_down=avg_down,\n norm_layer=norm_layer)\n self.layer2 = self._make_layer(block, 128, layers[1], stride=2, avg_down=avg_down,\n norm_layer=norm_layer)\n if dilated:\n self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2,\n avg_down=avg_down, norm_layer=norm_layer)\n self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4,\n avg_down=avg_down, norm_layer=norm_layer)\n else:\n self.layer3 = self._make_layer(block, 256, layers[2], stride=2,\n avg_down=avg_down, norm_layer=norm_layer)\n self.layer4 = self._make_layer(block, 512, layers[3], stride=2,\n avg_down=avg_down, norm_layer=norm_layer)\n self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n self.drop = None\n if final_drop > 0.0:\n self.drop = nn.Dropout(final_drop)\n self.fc = nn.Linear(512 * block.expansion, classes)\n\n def _make_layer(self, block, planes, blocks, stride=1, dilation=1,\n avg_down=False, norm_layer=nn.BatchNorm2d):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = []\n if avg_down:\n if dilation == 1:\n downsample.append(\n nn.AvgPool2d(kernel_size=stride, stride=stride, ceil_mode=True, count_include_pad=False)\n )\n else:\n downsample.append(\n nn.AvgPool2d(kernel_size=1, stride=1, ceil_mode=True, count_include_pad=False)\n )\n downsample.extend([\n nn.Conv2d(self.inplanes, out_channels=planes * block.expansion,\n kernel_size=1, stride=1, bias=False),\n norm_layer(planes * block.expansion)\n ])\n downsample = nn.Sequential(*downsample)\n else:\n downsample = nn.Sequential(\n nn.Conv2d(self.inplanes, out_channels=planes * block.expansion,\n kernel_size=1, stride=stride, bias=False),\n norm_layer(planes * block.expansion)\n )\n\n layers = []\n if dilation in (1, 2):\n layers.append(block(self.inplanes, planes, stride, dilation=1, downsample=downsample,\n previous_dilation=dilation, norm_layer=norm_layer))\n elif dilation == 4:\n layers.append(block(self.inplanes, planes, stride, dilation=2, downsample=downsample,\n previous_dilation=dilation, norm_layer=norm_layer))\n else:\n raise RuntimeError(\"=> unknown dilation size: {}\".format(dilation))\n\n self.inplanes = planes * block.expansion\n for _ in range(1, blocks):\n layers.append(block(self.inplanes, planes, dilation=dilation,\n previous_dilation=dilation, norm_layer=norm_layer))\n\n return nn.Sequential(*layers)\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.layer4(x)\n\n x = self.avgpool(x)\n x = x.view(x.size(0), -1)\n if self.drop is not None:\n x = self.drop(x)\n x = self.fc(x)\n\n return x\n\n\ndef _safe_state_dict_filtering(orig_dict, model_dict_keys):\n filtered_orig_dict = {}\n for k, v in orig_dict.items():\n if k in model_dict_keys:\n filtered_orig_dict[k] = v\n else:\n print(f\"[ERROR] Failed to load <{k}> in backbone\")\n return filtered_orig_dict\n\n\ndef resnet34_v1b(pretrained=False, **kwargs):\n model = ResNetV1b(BasicBlockV1b, [3, 4, 6, 3], **kwargs)\n if pretrained:\n model_dict = model.state_dict()\n filtered_orig_dict = _safe_state_dict_filtering(\n torch.hub.load(GLUON_RESNET_TORCH_HUB, 'gluon_resnet34_v1b', pretrained=True).state_dict(),\n model_dict.keys()\n )\n model_dict.update(filtered_orig_dict)\n model.load_state_dict(model_dict)\n return model\n\n\ndef resnet50_v1s(pretrained=False, **kwargs):\n model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], deep_stem=True, stem_width=64, **kwargs)\n if pretrained:\n model_dict = model.state_dict()\n filtered_orig_dict = _safe_state_dict_filtering(\n torch.hub.load(GLUON_RESNET_TORCH_HUB, 'gluon_resnet50_v1s', pretrained=True).state_dict(),\n model_dict.keys()\n )\n model_dict.update(filtered_orig_dict)\n model.load_state_dict(model_dict)\n return model\n\n\ndef resnet101_v1s(pretrained=False, **kwargs):\n model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], deep_stem=True, stem_width=64, **kwargs)\n if pretrained:\n model_dict = model.state_dict()\n filtered_orig_dict = _safe_state_dict_filtering(\n torch.hub.load(GLUON_RESNET_TORCH_HUB, 'gluon_resnet101_v1s', pretrained=True).state_dict(),\n model_dict.keys()\n )\n model_dict.update(filtered_orig_dict)\n model.load_state_dict(model_dict)\n return model\n\n\ndef resnet152_v1s(pretrained=False, **kwargs):\n model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], deep_stem=True, stem_width=64, **kwargs)\n if pretrained:\n model_dict = model.state_dict()\n filtered_orig_dict = _safe_state_dict_filtering(\n torch.hub.load(GLUON_RESNET_TORCH_HUB, 'gluon_resnet152_v1s', pretrained=True).state_dict(),\n model_dict.keys()\n )\n model_dict.update(filtered_orig_dict)\n model.load_state_dict(model_dict)\n return model\n" }, { "path": "XMem/inference/interact/fbrs/model/ops.py", "content": "import torch\nfrom torch import nn as nn\nimport numpy as np\n\nfrom . import initializer as initializer\nfrom ..utils.cython import get_dist_maps\n\n\ndef select_activation_function(activation):\n if isinstance(activation, str):\n if activation.lower() == 'relu':\n return nn.ReLU\n elif activation.lower() == 'softplus':\n return nn.Softplus\n else:\n raise ValueError(f\"Unknown activation type {activation}\")\n elif isinstance(activation, nn.Module):\n return activation\n else:\n raise ValueError(f\"Unknown activation type {activation}\")\n\n\nclass BilinearConvTranspose2d(nn.ConvTranspose2d):\n def __init__(self, in_channels, out_channels, scale, groups=1):\n kernel_size = 2 * scale - scale % 2\n self.scale = scale\n\n super().__init__(\n in_channels, out_channels,\n kernel_size=kernel_size,\n stride=scale,\n padding=1,\n groups=groups,\n bias=False)\n\n self.apply(initializer.Bilinear(scale=scale, in_channels=in_channels, groups=groups))\n\n\nclass DistMaps(nn.Module):\n def __init__(self, norm_radius, spatial_scale=1.0, cpu_mode=False):\n super(DistMaps, self).__init__()\n self.spatial_scale = spatial_scale\n self.norm_radius = norm_radius\n self.cpu_mode = cpu_mode\n\n def get_coord_features(self, points, batchsize, rows, cols):\n if self.cpu_mode:\n coords = []\n for i in range(batchsize):\n norm_delimeter = self.spatial_scale * self.norm_radius\n coords.append(get_dist_maps(points[i].cpu().float().numpy(), rows, cols,\n norm_delimeter))\n coords = torch.from_numpy(np.stack(coords, axis=0)).to(points.device).float()\n else:\n num_points = points.shape[1] // 2\n points = points.view(-1, 2)\n invalid_points = torch.max(points, dim=1, keepdim=False)[0] < 0\n row_array = torch.arange(start=0, end=rows, step=1, dtype=torch.float32, device=points.device)\n col_array = torch.arange(start=0, end=cols, step=1, dtype=torch.float32, device=points.device)\n\n coord_rows, coord_cols = torch.meshgrid(row_array, col_array)\n coords = torch.stack((coord_rows, coord_cols), dim=0).unsqueeze(0).repeat(points.size(0), 1, 1, 1)\n\n add_xy = (points * self.spatial_scale).view(points.size(0), points.size(1), 1, 1)\n coords.add_(-add_xy)\n coords.div_(self.norm_radius * self.spatial_scale)\n coords.mul_(coords)\n\n coords[:, 0] += coords[:, 1]\n coords = coords[:, :1]\n\n coords[invalid_points, :, :, :] = 1e6\n\n coords = coords.view(-1, num_points, 1, rows, cols)\n coords = coords.min(dim=1)[0] # -> (bs * num_masks * 2) x 1 x h x w\n coords = coords.view(-1, 2, rows, cols)\n\n coords.sqrt_().mul_(2).tanh_()\n\n return coords\n\n def forward(self, x, coords):\n return self.get_coord_features(coords, x.shape[0], x.shape[2], x.shape[3])\n" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/LICENSE", "content": "MIT License\n\nCopyright (c) 2018 Tamaki Kojima\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/README.md", "content": "# pytorch-syncbn\n\nTamaki Kojima(tamakoji@gmail.com)\n\n## Announcement\n\n**Pytorch 1.0 support**\n\n## Overview\nThis is alternative implementation of \"Synchronized Multi-GPU Batch Normalization\" which computes global stats across gpus instead of locally computed. SyncBN are getting important for those input image is large, and must use multi-gpu to increase the minibatch-size for the training.\n\nThe code was inspired by [Pytorch-Encoding](https://github.com/zhanghang1989/PyTorch-Encoding) and [Inplace-ABN](https://github.com/mapillary/inplace_abn)\n\n## Remarks\n- Unlike [Pytorch-Encoding](https://github.com/zhanghang1989/PyTorch-Encoding), you don't need custom `nn.DataParallel`\n- Unlike [Inplace-ABN](https://github.com/mapillary/inplace_abn), you can just replace your `nn.BatchNorm2d` to this module implementation, since it will not mark for inplace operation\n- You can plug into arbitrary module written in PyTorch to enable Synchronized BatchNorm\n- Backward computation is rewritten and tested against behavior of `nn.BatchNorm2d`\n\n## Requirements\nFor PyTorch, please refer to https://pytorch.org/\n\nNOTE : The code is tested only with PyTorch v1.0.0, CUDA10/CuDNN7.4.2 on ubuntu18.04\n\nIt utilize Pytorch JIT mechanism to compile seamlessly, using ninja. Please install ninja-build before use.\n\n```\nsudo apt-get install ninja-build\n```\n\nAlso install all dependencies for python. For pip, run:\n\n\n```\npip install -U -r requirements.txt\n```\n\n## Build\n\nThere is no need to build. just run and JIT will take care.\nJIT and cpp extensions are supported after PyTorch0.4, however it is highly recommended to use PyTorch > 1.0 due to huge design changes.\n\n## Usage\n\nPlease refer to [`test.py`](./test.py) for testing the difference between `nn.BatchNorm2d` and `modules.nn.BatchNorm2d`\n\n```\nimport torch\nfrom modules import nn as NN\nnum_gpu = torch.cuda.device_count()\nmodel = nn.Sequential(\n nn.Conv2d(3, 3, 1, 1, bias=False),\n NN.BatchNorm2d(3),\n nn.ReLU(inplace=True),\n nn.Conv2d(3, 3, 1, 1, bias=False),\n NN.BatchNorm2d(3),\n).cuda()\nmodel = nn.DataParallel(model, device_ids=range(num_gpu))\nx = torch.rand(num_gpu, 3, 2, 2).cuda()\nz = model(x)\n```\n\n## Math\n\n### Forward\n1. compute in each gpu\n2. gather all from workers to master and compute where\n\n \n\n and\n\n \n\n and then above global stats to be shared to all gpus, update running_mean and running_var by moving average using global stats.\n\n3. forward batchnorm using global stats by\n\n \n\n and then\n\n \n\n where is weight parameter and is bias parameter.\n\n4. save for backward\n\n### Backward\n\n1. Restore saved \n\n2. Compute below sums on each gpu\n\n \n\n and\n\n \n\n where \n\n then gather them at master node to sum up global, and normalize with N where N is total number of elements for each channels. Global sums are then shared among all gpus.\n\n3. compute gradients using global stats\n\n \n\n where\n\n \n\n and\n\n \n\n and finally,\n\n \n\n \n\n \n\n Note that in the implementation, normalization with N is performed at step (2) and above equation and implementation is not exactly the same, but mathematically is same.\n\n You can go deeper on above explanation at [Kevin Zakka's Blog](https://kevinzakka.github.io/2016/09/14/batch_normalization/)" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/__init__.py", "content": "" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/modules/__init__.py", "content": "" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/modules/functional/__init__.py", "content": "from .syncbn import batchnorm2d_sync\n" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/modules/functional/_csrc.py", "content": "\"\"\"\n/*****************************************************************************/\n\nExtension module loader\n\ncode referenced from : https://github.com/facebookresearch/maskrcnn-benchmark\n\n/*****************************************************************************/\n\"\"\"\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport glob\nimport os.path\n\nimport torch\n\ntry:\n from torch.utils.cpp_extension import load\n from torch.utils.cpp_extension import CUDA_HOME\nexcept ImportError:\n raise ImportError(\n \"The cpp layer extensions requires PyTorch 0.4 or higher\")\n\n\ndef _load_C_extensions():\n this_dir = os.path.dirname(os.path.abspath(__file__))\n this_dir = os.path.join(this_dir, \"csrc\")\n\n main_file = glob.glob(os.path.join(this_dir, \"*.cpp\"))\n sources_cpu = glob.glob(os.path.join(this_dir, \"cpu\", \"*.cpp\"))\n sources_cuda = glob.glob(os.path.join(this_dir, \"cuda\", \"*.cu\"))\n\n sources = main_file + sources_cpu\n\n extra_cflags = []\n extra_cuda_cflags = []\n if torch.cuda.is_available() and CUDA_HOME is not None:\n sources.extend(sources_cuda)\n extra_cflags = [\"-O3\", \"-DWITH_CUDA\"]\n extra_cuda_cflags = [\"--expt-extended-lambda\"]\n sources = [os.path.join(this_dir, s) for s in sources]\n extra_include_paths = [this_dir]\n return load(\n name=\"ext_lib\",\n sources=sources,\n extra_cflags=extra_cflags,\n extra_include_paths=extra_include_paths,\n extra_cuda_cflags=extra_cuda_cflags,\n )\n\n\n_backend = _load_C_extensions()\n" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/modules/functional/csrc/bn.h", "content": "/*****************************************************************************\n\nSyncBN\n\n*****************************************************************************/\n#pragma once\n\n#ifdef WITH_CUDA\n#include \"cuda/ext_lib.h\"\n#endif\n\n/// SyncBN\n\nstd::vector syncbn_sum_sqsum(const at::Tensor& x) {\n if (x.is_cuda()) {\n#ifdef WITH_CUDA\n return syncbn_sum_sqsum_cuda(x);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n } else {\n AT_ERROR(\"CPU implementation not supported\");\n }\n}\n\nat::Tensor syncbn_forward(const at::Tensor& x, const at::Tensor& weight,\n const at::Tensor& bias, const at::Tensor& mean,\n const at::Tensor& var, bool affine, float eps) {\n if (x.is_cuda()) {\n#ifdef WITH_CUDA\n return syncbn_forward_cuda(x, weight, bias, mean, var, affine, eps);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n } else {\n AT_ERROR(\"CPU implementation not supported\");\n }\n}\n\nstd::vector syncbn_backward_xhat(const at::Tensor& dz,\n const at::Tensor& x,\n const at::Tensor& mean,\n const at::Tensor& var, float eps) {\n if (dz.is_cuda()) {\n#ifdef WITH_CUDA\n return syncbn_backward_xhat_cuda(dz, x, mean, var, eps);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n } else {\n AT_ERROR(\"CPU implementation not supported\");\n }\n}\n\nstd::vector syncbn_backward(\n const at::Tensor& dz, const at::Tensor& x, const at::Tensor& weight,\n const at::Tensor& bias, const at::Tensor& mean, const at::Tensor& var,\n const at::Tensor& sum_dz, const at::Tensor& sum_dz_xhat, bool affine,\n float eps) {\n if (dz.is_cuda()) {\n#ifdef WITH_CUDA\n return syncbn_backward_cuda(dz, x, weight, bias, mean, var, sum_dz,\n sum_dz_xhat, affine, eps);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n } else {\n AT_ERROR(\"CPU implementation not supported\");\n }\n}\n" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/modules/functional/csrc/cuda/bn_cuda.cu", "content": "/*****************************************************************************\n\nCUDA SyncBN code\n\ncode referenced from : https://github.com/mapillary/inplace_abn\n\n*****************************************************************************/\n#include \n#include \n#include \n#include \n#include \"cuda/common.h\"\n\n// Utilities\nvoid get_dims(at::Tensor x, int64_t &num, int64_t &chn, int64_t &sp) {\n num = x.size(0);\n chn = x.size(1);\n sp = 1;\n for (int64_t i = 2; i < x.ndimension(); ++i) sp *= x.size(i);\n}\n\n/// SyncBN\n\ntemplate \nstruct SqSumOp {\n __device__ SqSumOp(const T *t, int c, int s) : tensor(t), chn(c), sp(s) {}\n __device__ __forceinline__ Pair operator()(int batch, int plane, int n) {\n T x = tensor[(batch * chn + plane) * sp + n];\n return Pair(x, x * x); // x, x^2\n }\n const T *tensor;\n const int chn;\n const int sp;\n};\n\ntemplate \n__global__ void syncbn_sum_sqsum_kernel(const T *x, T *sum, T *sqsum,\n int num, int chn, int sp) {\n int plane = blockIdx.x;\n Pair res =\n reduce, SqSumOp>(SqSumOp(x, chn, sp), plane, num, chn, sp);\n __syncthreads();\n if (threadIdx.x == 0) {\n sum[plane] = res.v1;\n sqsum[plane] = res.v2;\n }\n}\n\nstd::vector syncbn_sum_sqsum_cuda(const at::Tensor &x) {\n CHECK_INPUT(x);\n\n // Extract dimensions\n int64_t num, chn, sp;\n get_dims(x, num, chn, sp);\n\n // Prepare output tensors\n auto sum = at::empty({chn}, x.options());\n auto sqsum = at::empty({chn}, x.options());\n\n // Run kernel\n dim3 blocks(chn);\n dim3 threads(getNumThreads(sp));\n AT_DISPATCH_FLOATING_TYPES(\n x.type(), \"syncbn_sum_sqsum_cuda\", ([&] {\n syncbn_sum_sqsum_kernel<<>>(\n x.data(), sum.data(),\n sqsum.data(), num, chn, sp);\n }));\n return {sum, sqsum};\n}\n\ntemplate \n__global__ void syncbn_forward_kernel(T *z, const T *x, const T *weight,\n const T *bias, const T *mean,\n const T *var, bool affine, float eps,\n int num, int chn, int sp) {\n int plane = blockIdx.x;\n T _mean = mean[plane];\n T _var = var[plane];\n T _weight = affine ? weight[plane] : T(1);\n T _bias = affine ? bias[plane] : T(0);\n float _invstd = T(0);\n if (_var || eps) {\n _invstd = rsqrt(_var + eps);\n }\n for (int batch = 0; batch < num; ++batch) {\n for (int n = threadIdx.x; n < sp; n += blockDim.x) {\n T _x = x[(batch * chn + plane) * sp + n];\n T _xhat = (_x - _mean) * _invstd;\n T _z = _xhat * _weight + _bias;\n z[(batch * chn + plane) * sp + n] = _z;\n }\n }\n}\n\nat::Tensor syncbn_forward_cuda(const at::Tensor &x, const at::Tensor &weight,\n const at::Tensor &bias, const at::Tensor &mean,\n const at::Tensor &var, bool affine, float eps) {\n CHECK_INPUT(x);\n CHECK_INPUT(weight);\n CHECK_INPUT(bias);\n CHECK_INPUT(mean);\n CHECK_INPUT(var);\n\n // Extract dimensions\n int64_t num, chn, sp;\n get_dims(x, num, chn, sp);\n\n auto z = at::zeros_like(x);\n\n // Run kernel\n dim3 blocks(chn);\n dim3 threads(getNumThreads(sp));\n AT_DISPATCH_FLOATING_TYPES(\n x.type(), \"syncbn_forward_cuda\", ([&] {\n syncbn_forward_kernel<<>>(\n z.data(), x.data(),\n weight.data(), bias.data(),\n mean.data(), var.data(),\n affine, eps, num, chn, sp);\n }));\n return z;\n}\n\ntemplate \nstruct XHatOp {\n __device__ XHatOp(T _weight, T _bias, const T *_dz, const T *_x, int c, int s)\n : weight(_weight), bias(_bias), x(_x), dz(_dz), chn(c), sp(s) {}\n __device__ __forceinline__ Pair operator()(int batch, int plane, int n) {\n // xhat = (x - bias) * weight\n T _xhat = (x[(batch * chn + plane) * sp + n] - bias) * weight;\n // dxhat * x_hat\n T _dz = dz[(batch * chn + plane) * sp + n];\n return Pair(_dz, _dz * _xhat);\n }\n const T weight;\n const T bias;\n const T *dz;\n const T *x;\n const int chn;\n const int sp;\n};\n\ntemplate \n__global__ void syncbn_backward_xhat_kernel(const T *dz, const T *x,\n const T *mean, const T *var,\n T *sum_dz, T *sum_dz_xhat,\n float eps, int num, int chn,\n int sp) {\n int plane = blockIdx.x;\n T _mean = mean[plane];\n T _var = var[plane];\n T _invstd = T(0);\n if (_var || eps) {\n _invstd = rsqrt(_var + eps);\n }\n Pair res = reduce, XHatOp>(\n XHatOp(_invstd, _mean, dz, x, chn, sp), plane, num, chn, sp);\n __syncthreads();\n if (threadIdx.x == 0) {\n // \\sum(\\frac{dJ}{dy_i})\n sum_dz[plane] = res.v1;\n // \\sum(\\frac{dJ}{dy_i}*\\hat{x_i})\n sum_dz_xhat[plane] = res.v2;\n }\n}\n\nstd::vector syncbn_backward_xhat_cuda(const at::Tensor &dz,\n const at::Tensor &x,\n const at::Tensor &mean,\n const at::Tensor &var,\n float eps) {\n CHECK_INPUT(dz);\n CHECK_INPUT(x);\n CHECK_INPUT(mean);\n CHECK_INPUT(var);\n // Extract dimensions\n int64_t num, chn, sp;\n get_dims(x, num, chn, sp);\n // Prepare output tensors\n auto sum_dz = at::empty({chn}, x.options());\n auto sum_dz_xhat = at::empty({chn}, x.options());\n // Run kernel\n dim3 blocks(chn);\n dim3 threads(getNumThreads(sp));\n AT_DISPATCH_FLOATING_TYPES(\n x.type(), \"syncbn_backward_xhat_cuda\", ([&] {\n syncbn_backward_xhat_kernel<<>>(\n dz.data(), x.data(), mean.data(),\n var.data(), sum_dz.data(),\n sum_dz_xhat.data(), eps, num, chn, sp);\n }));\n return {sum_dz, sum_dz_xhat};\n}\n\ntemplate \n__global__ void syncbn_backward_kernel(const T *dz, const T *x, const T *weight,\n const T *bias, const T *mean,\n const T *var, const T *sum_dz,\n const T *sum_dz_xhat, T *dx, T *dweight,\n T *dbias, bool affine, float eps,\n int num, int chn, int sp) {\n int plane = blockIdx.x;\n T _mean = mean[plane];\n T _var = var[plane];\n T _weight = affine ? weight[plane] : T(1);\n T _sum_dz = sum_dz[plane];\n T _sum_dz_xhat = sum_dz_xhat[plane];\n T _invstd = T(0);\n if (_var || eps) {\n _invstd = rsqrt(_var + eps);\n }\n /*\n \\frac{dJ}{dx_i} = \\frac{1}{N\\sqrt{(\\sigma^2+\\epsilon)}} (\n N\\frac{dJ}{d\\hat{x_i}} -\n \\sum_{j=1}^{N}(\\frac{dJ}{d\\hat{x_j}}) -\n \\hat{x_i}\\sum_{j=1}^{N}(\\frac{dJ}{d\\hat{x_j}}\\hat{x_j})\n )\n Note : N is omitted here since it will be accumulated and\n _sum_dz and _sum_dz_xhat expected to be already normalized\n before the call.\n */\n if (dx) {\n T _mul = _weight * _invstd;\n for (int batch = 0; batch < num; ++batch) {\n for (int n = threadIdx.x; n < sp; n += blockDim.x) {\n T _dz = dz[(batch * chn + plane) * sp + n];\n T _xhat = (x[(batch * chn + plane) * sp + n] - _mean) * _invstd;\n T _dx = (_dz - _sum_dz - _xhat * _sum_dz_xhat) * _mul;\n dx[(batch * chn + plane) * sp + n] = _dx;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n if (affine) {\n T _norm = num * sp;\n dweight[plane] += _sum_dz_xhat * _norm;\n dbias[plane] += _sum_dz * _norm;\n }\n }\n}\n\nstd::vector syncbn_backward_cuda(\n const at::Tensor &dz, const at::Tensor &x, const at::Tensor &weight,\n const at::Tensor &bias, const at::Tensor &mean, const at::Tensor &var,\n const at::Tensor &sum_dz, const at::Tensor &sum_dz_xhat, bool affine,\n float eps) {\n CHECK_INPUT(dz);\n CHECK_INPUT(x);\n CHECK_INPUT(weight);\n CHECK_INPUT(bias);\n CHECK_INPUT(mean);\n CHECK_INPUT(var);\n CHECK_INPUT(sum_dz);\n CHECK_INPUT(sum_dz_xhat);\n\n // Extract dimensions\n int64_t num, chn, sp;\n get_dims(x, num, chn, sp);\n\n // Prepare output tensors\n auto dx = at::zeros_like(dz);\n auto dweight = at::zeros_like(weight);\n auto dbias = at::zeros_like(bias);\n\n // Run kernel\n dim3 blocks(chn);\n dim3 threads(getNumThreads(sp));\n AT_DISPATCH_FLOATING_TYPES(\n x.type(), \"syncbn_backward_cuda\", ([&] {\n syncbn_backward_kernel<<>>(\n dz.data(), x.data(), weight.data(),\n bias.data(), mean.data(), var.data(),\n sum_dz.data(), sum_dz_xhat.data(),\n dx.data(), dweight.data(),\n dbias.data(), affine, eps, num, chn, sp);\n }));\n return {dx, dweight, dbias};\n}" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/modules/functional/csrc/cuda/common.h", "content": "/*****************************************************************************\n\nCUDA utility funcs\n\ncode referenced from : https://github.com/mapillary/inplace_abn\n\n*****************************************************************************/\n#pragma once\n\n#include \n\n// Checks\n#ifndef AT_CHECK\n #define AT_CHECK AT_ASSERT\n#endif\n#define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)\n\n/*\n * General settings\n */\nconst int WARP_SIZE = 32;\nconst int MAX_BLOCK_SIZE = 512;\n\ntemplate \nstruct Pair {\n T v1, v2;\n __device__ Pair() {}\n __device__ Pair(T _v1, T _v2) : v1(_v1), v2(_v2) {}\n __device__ Pair(T v) : v1(v), v2(v) {}\n __device__ Pair(int v) : v1(v), v2(v) {}\n __device__ Pair &operator+=(const Pair &a) {\n v1 += a.v1;\n v2 += a.v2;\n return *this;\n }\n};\n\n/*\n * Utility functions\n */\ntemplate \n__device__ __forceinline__ T WARP_SHFL_XOR(T value, int laneMask,\n int width = warpSize,\n unsigned int mask = 0xffffffff) {\n#if CUDART_VERSION >= 9000\n return __shfl_xor_sync(mask, value, laneMask, width);\n#else\n return __shfl_xor(value, laneMask, width);\n#endif\n}\n\n__device__ __forceinline__ int getMSB(int val) { return 31 - __clz(val); }\n\nstatic int getNumThreads(int nElem) {\n int threadSizes[5] = {32, 64, 128, 256, MAX_BLOCK_SIZE};\n for (int i = 0; i != 5; ++i) {\n if (nElem <= threadSizes[i]) {\n return threadSizes[i];\n }\n }\n return MAX_BLOCK_SIZE;\n}\n\ntemplate \nstatic __device__ __forceinline__ T warpSum(T val) {\n#if __CUDA_ARCH__ >= 300\n for (int i = 0; i < getMSB(WARP_SIZE); ++i) {\n val += WARP_SHFL_XOR(val, 1 << i, WARP_SIZE);\n }\n#else\n __shared__ T values[MAX_BLOCK_SIZE];\n values[threadIdx.x] = val;\n __threadfence_block();\n const int base = (threadIdx.x / WARP_SIZE) * WARP_SIZE;\n for (int i = 1; i < WARP_SIZE; i++) {\n val += values[base + ((i + threadIdx.x) % WARP_SIZE)];\n }\n#endif\n return val;\n}\n\ntemplate \nstatic __device__ __forceinline__ Pair warpSum(Pair value) {\n value.v1 = warpSum(value.v1);\n value.v2 = warpSum(value.v2);\n return value;\n}\n\ntemplate \n__device__ T reduce(Op op, int plane, int N, int C, int S) {\n T sum = (T)0;\n for (int batch = 0; batch < N; ++batch) {\n for (int x = threadIdx.x; x < S; x += blockDim.x) {\n sum += op(batch, plane, x);\n }\n }\n\n // sum over NumThreads within a warp\n sum = warpSum(sum);\n\n // 'transpose', and reduce within warp again\n __shared__ T shared[32];\n __syncthreads();\n if (threadIdx.x % WARP_SIZE == 0) {\n shared[threadIdx.x / WARP_SIZE] = sum;\n }\n if (threadIdx.x >= blockDim.x / WARP_SIZE && threadIdx.x < WARP_SIZE) {\n // zero out the other entries in shared\n shared[threadIdx.x] = (T)0;\n }\n __syncthreads();\n if (threadIdx.x / WARP_SIZE == 0) {\n sum = warpSum(shared[threadIdx.x]);\n if (threadIdx.x == 0) {\n shared[0] = sum;\n }\n }\n __syncthreads();\n\n // Everyone picks it up, should be broadcast into the whole gradInput\n return shared[0];\n}" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/modules/functional/csrc/cuda/ext_lib.h", "content": "/*****************************************************************************\n\nCUDA SyncBN code\n\n*****************************************************************************/\n#pragma once\n#include \n#include \n\n/// Sync-BN\nstd::vector syncbn_sum_sqsum_cuda(const at::Tensor& x);\nat::Tensor syncbn_forward_cuda(const at::Tensor& x, const at::Tensor& weight,\n const at::Tensor& bias, const at::Tensor& mean,\n const at::Tensor& var, bool affine, float eps);\nstd::vector syncbn_backward_xhat_cuda(const at::Tensor& dz,\n const at::Tensor& x,\n const at::Tensor& mean,\n const at::Tensor& var,\n float eps);\nstd::vector syncbn_backward_cuda(\n const at::Tensor& dz, const at::Tensor& x, const at::Tensor& weight,\n const at::Tensor& bias, const at::Tensor& mean, const at::Tensor& var,\n const at::Tensor& sum_dz, const at::Tensor& sum_dz_xhat, bool affine,\n float eps);\n" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/modules/functional/csrc/ext_lib.cpp", "content": "#include \"bn.h\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"syncbn_sum_sqsum\", &syncbn_sum_sqsum, \"Sum and Sum^2 computation\");\n m.def(\"syncbn_forward\", &syncbn_forward, \"SyncBN forward computation\");\n m.def(\"syncbn_backward_xhat\", &syncbn_backward_xhat,\n \"First part of SyncBN backward computation\");\n m.def(\"syncbn_backward\", &syncbn_backward,\n \"Second part of SyncBN backward computation\");\n}" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/modules/functional/syncbn.py", "content": "\"\"\"\n/*****************************************************************************/\n\nBatchNorm2dSync with multi-gpu\n\ncode referenced from : https://github.com/mapillary/inplace_abn\n\n/*****************************************************************************/\n\"\"\"\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport torch.cuda.comm as comm\nfrom torch.autograd import Function\nfrom torch.autograd.function import once_differentiable\nfrom ._csrc import _backend\n\n\ndef _count_samples(x):\n count = 1\n for i, s in enumerate(x.size()):\n if i != 1:\n count *= s\n return count\n\n\nclass BatchNorm2dSyncFunc(Function):\n\n @staticmethod\n def forward(ctx, x, weight, bias, running_mean, running_var,\n extra, compute_stats=True, momentum=0.1, eps=1e-05):\n def _parse_extra(ctx, extra):\n ctx.is_master = extra[\"is_master\"]\n if ctx.is_master:\n ctx.master_queue = extra[\"master_queue\"]\n ctx.worker_queues = extra[\"worker_queues\"]\n ctx.worker_ids = extra[\"worker_ids\"]\n else:\n ctx.master_queue = extra[\"master_queue\"]\n ctx.worker_queue = extra[\"worker_queue\"]\n # Save context\n if extra is not None:\n _parse_extra(ctx, extra)\n ctx.compute_stats = compute_stats\n ctx.momentum = momentum\n ctx.eps = eps\n ctx.affine = weight is not None and bias is not None\n if ctx.compute_stats:\n N = _count_samples(x) * (ctx.master_queue.maxsize + 1)\n assert N > 1\n # 1. compute sum(x) and sum(x^2)\n xsum, xsqsum = _backend.syncbn_sum_sqsum(x.detach())\n if ctx.is_master:\n xsums, xsqsums = [xsum], [xsqsum]\n # master : gatther all sum(x) and sum(x^2) from slaves\n for _ in range(ctx.master_queue.maxsize):\n xsum_w, xsqsum_w = ctx.master_queue.get()\n ctx.master_queue.task_done()\n xsums.append(xsum_w)\n xsqsums.append(xsqsum_w)\n xsum = comm.reduce_add(xsums)\n xsqsum = comm.reduce_add(xsqsums)\n mean = xsum / N\n sumvar = xsqsum - xsum * mean\n var = sumvar / N\n uvar = sumvar / (N - 1)\n # master : broadcast global mean, variance to all slaves\n tensors = comm.broadcast_coalesced(\n (mean, uvar, var), [mean.get_device()] + ctx.worker_ids)\n for ts, queue in zip(tensors[1:], ctx.worker_queues):\n queue.put(ts)\n else:\n # slave : send sum(x) and sum(x^2) to master\n ctx.master_queue.put((xsum, xsqsum))\n # slave : get global mean and variance\n mean, uvar, var = ctx.worker_queue.get()\n ctx.worker_queue.task_done()\n\n # Update running stats\n running_mean.mul_((1 - ctx.momentum)).add_(ctx.momentum * mean)\n running_var.mul_((1 - ctx.momentum)).add_(ctx.momentum * uvar)\n ctx.N = N\n ctx.save_for_backward(x, weight, bias, mean, var)\n else:\n mean, var = running_mean, running_var\n\n # do batch norm forward\n z = _backend.syncbn_forward(x, weight, bias, mean, var,\n ctx.affine, ctx.eps)\n return z\n\n @staticmethod\n @once_differentiable\n def backward(ctx, dz):\n x, weight, bias, mean, var = ctx.saved_tensors\n dz = dz.contiguous()\n\n # 1. compute \\sum(\\frac{dJ}{dy_i}) and \\sum(\\frac{dJ}{dy_i}*\\hat{x_i})\n sum_dz, sum_dz_xhat = _backend.syncbn_backward_xhat(\n dz, x, mean, var, ctx.eps)\n if ctx.is_master:\n sum_dzs, sum_dz_xhats = [sum_dz], [sum_dz_xhat]\n # master : gatther from slaves\n for _ in range(ctx.master_queue.maxsize):\n sum_dz_w, sum_dz_xhat_w = ctx.master_queue.get()\n ctx.master_queue.task_done()\n sum_dzs.append(sum_dz_w)\n sum_dz_xhats.append(sum_dz_xhat_w)\n # master : compute global stats\n sum_dz = comm.reduce_add(sum_dzs)\n sum_dz_xhat = comm.reduce_add(sum_dz_xhats)\n sum_dz /= ctx.N\n sum_dz_xhat /= ctx.N\n # master : broadcast global stats\n tensors = comm.broadcast_coalesced(\n (sum_dz, sum_dz_xhat), [mean.get_device()] + ctx.worker_ids)\n for ts, queue in zip(tensors[1:], ctx.worker_queues):\n queue.put(ts)\n else:\n # slave : send to master\n ctx.master_queue.put((sum_dz, sum_dz_xhat))\n # slave : get global stats\n sum_dz, sum_dz_xhat = ctx.worker_queue.get()\n ctx.worker_queue.task_done()\n\n # do batch norm backward\n dx, dweight, dbias = _backend.syncbn_backward(\n dz, x, weight, bias, mean, var, sum_dz, sum_dz_xhat,\n ctx.affine, ctx.eps)\n\n return dx, dweight, dbias, \\\n None, None, None, None, None, None\n\nbatchnorm2d_sync = BatchNorm2dSyncFunc.apply\n\n__all__ = [\"batchnorm2d_sync\"]\n" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/modules/nn/__init__.py", "content": "from .syncbn import *\n" }, { "path": "XMem/inference/interact/fbrs/model/syncbn/modules/nn/syncbn.py", "content": "\"\"\"\n/*****************************************************************************/\n\nBatchNorm2dSync with multi-gpu\n\n/*****************************************************************************/\n\"\"\"\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\ntry:\n # python 3\n from queue import Queue\nexcept ImportError:\n # python 2\n from Queue import Queue\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\nfrom torch.nn.parameter import Parameter\nfrom isegm.model.syncbn.modules.functional import batchnorm2d_sync\n\n\nclass _BatchNorm(nn.Module):\n \"\"\"\n Customized BatchNorm from nn.BatchNorm\n >> added freeze attribute to enable bn freeze.\n \"\"\"\n\n def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,\n track_running_stats=True):\n super(_BatchNorm, self).__init__()\n self.num_features = num_features\n self.eps = eps\n self.momentum = momentum\n self.affine = affine\n self.track_running_stats = track_running_stats\n self.freezed = False\n if self.affine:\n self.weight = Parameter(torch.Tensor(num_features))\n self.bias = Parameter(torch.Tensor(num_features))\n else:\n self.register_parameter('weight', None)\n self.register_parameter('bias', None)\n if self.track_running_stats:\n self.register_buffer('running_mean', torch.zeros(num_features))\n self.register_buffer('running_var', torch.ones(num_features))\n else:\n self.register_parameter('running_mean', None)\n self.register_parameter('running_var', None)\n self.reset_parameters()\n\n def reset_parameters(self):\n if self.track_running_stats:\n self.running_mean.zero_()\n self.running_var.fill_(1)\n if self.affine:\n self.weight.data.uniform_()\n self.bias.data.zero_()\n\n def _check_input_dim(self, input):\n return NotImplemented\n\n def forward(self, input):\n self._check_input_dim(input)\n\n compute_stats = not self.freezed and \\\n self.training and self.track_running_stats\n\n ret = F.batch_norm(input, self.running_mean, self.running_var,\n self.weight, self.bias, compute_stats,\n self.momentum, self.eps)\n return ret\n\n def extra_repr(self):\n return '{num_features}, eps={eps}, momentum={momentum}, '\\\n 'affine={affine}, ' \\\n 'track_running_stats={track_running_stats}'.format(\n **self.__dict__)\n\n\nclass BatchNorm2dNoSync(_BatchNorm):\n \"\"\"\n Equivalent to nn.BatchNorm2d\n \"\"\"\n\n def _check_input_dim(self, input):\n if input.dim() != 4:\n raise ValueError('expected 4D input (got {}D input)'\n .format(input.dim()))\n\n\nclass BatchNorm2dSync(BatchNorm2dNoSync):\n \"\"\"\n BatchNorm2d with automatic multi-GPU Sync\n \"\"\"\n\n def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,\n track_running_stats=True):\n super(BatchNorm2dSync, self).__init__(\n num_features, eps=eps, momentum=momentum, affine=affine,\n track_running_stats=track_running_stats)\n self.sync_enabled = True\n self.devices = list(range(torch.cuda.device_count()))\n if len(self.devices) > 1:\n # Initialize queues\n self.worker_ids = self.devices[1:]\n self.master_queue = Queue(len(self.worker_ids))\n self.worker_queues = [Queue(1) for _ in self.worker_ids]\n\n def forward(self, x):\n compute_stats = not self.freezed and \\\n self.training and self.track_running_stats\n if self.sync_enabled and compute_stats and len(self.devices) > 1:\n if x.get_device() == self.devices[0]:\n # Master mode\n extra = {\n \"is_master\": True,\n \"master_queue\": self.master_queue,\n \"worker_queues\": self.worker_queues,\n \"worker_ids\": self.worker_ids\n }\n else:\n # Worker mode\n extra = {\n \"is_master\": False,\n \"master_queue\": self.master_queue,\n \"worker_queue\": self.worker_queues[\n self.worker_ids.index(x.get_device())]\n }\n return batchnorm2d_sync(x, self.weight, self.bias,\n self.running_mean, self.running_var,\n extra, compute_stats, self.momentum,\n self.eps)\n return super(BatchNorm2dSync, self).forward(x)\n\n def __repr__(self):\n \"\"\"repr\"\"\"\n rep = '{name}({num_features}, eps={eps}, momentum={momentum},' \\\n 'affine={affine}, ' \\\n 'track_running_stats={track_running_stats},' \\\n 'devices={devices})'\n return rep.format(name=self.__class__.__name__, **self.__dict__)\n\n#BatchNorm2d = BatchNorm2dNoSync\nBatchNorm2d = BatchNorm2dSync\n" }, { "path": "XMem/inference/interact/fbrs/utils/__init__.py", "content": "" }, { "path": "XMem/inference/interact/fbrs/utils/cython/__init__.py", "content": "# noinspection PyUnresolvedReferences\nfrom .dist_maps import get_dist_maps" }, { "path": "XMem/inference/interact/fbrs/utils/cython/_get_dist_maps.pyx", "content": "import numpy as np\ncimport cython\ncimport numpy as np\nfrom libc.stdlib cimport malloc, free\n\nctypedef struct qnode:\n int row\n int col\n int layer\n int orig_row\n int orig_col\n\n@cython.infer_types(True)\n@cython.boundscheck(False)\n@cython.wraparound(False)\n@cython.nonecheck(False)\ndef get_dist_maps(np.ndarray[np.float32_t, ndim=2, mode=\"c\"] points,\n int height, int width, float norm_delimeter):\n cdef np.ndarray[np.float32_t, ndim=3, mode=\"c\"] dist_maps = \\\n np.full((2, height, width), 1e6, dtype=np.float32, order=\"C\")\n\n cdef int *dxy = [-1, 0, 0, -1, 0, 1, 1, 0]\n cdef int i, j, x, y, dx, dy\n cdef qnode v\n cdef qnode *q = malloc((4 * height * width + 1) * sizeof(qnode))\n cdef int qhead = 0, qtail = -1\n cdef float ndist\n\n for i in range(points.shape[0]):\n x, y = round(points[i, 0]), round(points[i, 1])\n if x >= 0:\n qtail += 1\n q[qtail].row = x\n q[qtail].col = y\n q[qtail].orig_row = x\n q[qtail].orig_col = y\n if i >= points.shape[0] / 2:\n q[qtail].layer = 1\n else:\n q[qtail].layer = 0\n dist_maps[q[qtail].layer, x, y] = 0\n\n while qtail - qhead + 1 > 0:\n v = q[qhead]\n qhead += 1\n\n for k in range(4):\n x = v.row + dxy[2 * k]\n y = v.col + dxy[2 * k + 1]\n\n ndist = ((x - v.orig_row)/norm_delimeter) ** 2 + ((y - v.orig_col)/norm_delimeter) ** 2\n if (x >= 0 and y >= 0 and x < height and y < width and\n dist_maps[v.layer, x, y] > ndist):\n qtail += 1\n q[qtail].orig_col = v.orig_col\n q[qtail].orig_row = v.orig_row\n q[qtail].layer = v.layer\n q[qtail].row = x\n q[qtail].col = y\n dist_maps[v.layer, x, y] = ndist\n\n free(q)\n return dist_maps\n" }, { "path": "XMem/inference/interact/fbrs/utils/cython/_get_dist_maps.pyxbld", "content": "import numpy\n\ndef make_ext(modname, pyxfilename):\n from distutils.extension import Extension\n return Extension(modname, [pyxfilename],\n include_dirs=[numpy.get_include()],\n extra_compile_args=['-O3'], language='c++')\n" }, { "path": "XMem/inference/interact/fbrs/utils/cython/dist_maps.py", "content": "import pyximport; pyximport.install(pyximport=True, language_level=3)\n# noinspection PyUnresolvedReferences\nfrom ._get_dist_maps import get_dist_maps" }, { "path": "XMem/inference/interact/fbrs/utils/misc.py", "content": "from functools import partial\n\nimport torch\nimport numpy as np\n\n\ndef get_dims_with_exclusion(dim, exclude=None):\n dims = list(range(dim))\n if exclude is not None:\n dims.remove(exclude)\n\n return dims\n\n\ndef get_unique_labels(mask):\n return np.nonzero(np.bincount(mask.flatten() + 1))[0] - 1\n\n\ndef get_bbox_from_mask(mask):\n rows = np.any(mask, axis=1)\n cols = np.any(mask, axis=0)\n rmin, rmax = np.where(rows)[0][[0, -1]]\n cmin, cmax = np.where(cols)[0][[0, -1]]\n\n return rmin, rmax, cmin, cmax\n\n\ndef expand_bbox(bbox, expand_ratio, min_crop_size=None):\n rmin, rmax, cmin, cmax = bbox\n rcenter = 0.5 * (rmin + rmax)\n ccenter = 0.5 * (cmin + cmax)\n height = expand_ratio * (rmax - rmin + 1)\n width = expand_ratio * (cmax - cmin + 1)\n if min_crop_size is not None:\n height = max(height, min_crop_size)\n width = max(width, min_crop_size)\n\n rmin = int(round(rcenter - 0.5 * height))\n rmax = int(round(rcenter + 0.5 * height))\n cmin = int(round(ccenter - 0.5 * width))\n cmax = int(round(ccenter + 0.5 * width))\n\n return rmin, rmax, cmin, cmax\n\n\ndef clamp_bbox(bbox, rmin, rmax, cmin, cmax):\n return (max(rmin, bbox[0]), min(rmax, bbox[1]),\n max(cmin, bbox[2]), min(cmax, bbox[3]))\n\n\ndef get_bbox_iou(b1, b2):\n h_iou = get_segments_iou(b1[:2], b2[:2])\n w_iou = get_segments_iou(b1[2:4], b2[2:4])\n return h_iou * w_iou\n\n\ndef get_segments_iou(s1, s2):\n a, b = s1\n c, d = s2\n intersection = max(0, min(b, d) - max(a, c) + 1)\n union = max(1e-6, max(b, d) - min(a, c) + 1)\n return intersection / union\n" }, { "path": "XMem/inference/interact/fbrs/utils/vis.py", "content": "from functools import lru_cache\n\nimport cv2\nimport numpy as np\n\n\ndef visualize_instances(imask, bg_color=255,\n boundaries_color=None, boundaries_width=1, boundaries_alpha=0.8):\n num_objects = imask.max() + 1\n palette = get_palette(num_objects)\n if bg_color is not None:\n palette[0] = bg_color\n\n result = palette[imask].astype(np.uint8)\n if boundaries_color is not None:\n boundaries_mask = get_boundaries(imask, boundaries_width=boundaries_width)\n tresult = result.astype(np.float32)\n tresult[boundaries_mask] = boundaries_color\n tresult = tresult * boundaries_alpha + (1 - boundaries_alpha) * result\n result = tresult.astype(np.uint8)\n\n return result\n\n\n@lru_cache(maxsize=16)\ndef get_palette(num_cls):\n palette = np.zeros(3 * num_cls, dtype=np.int32)\n\n for j in range(0, num_cls):\n lab = j\n i = 0\n\n while lab > 0:\n palette[j*3 + 0] |= (((lab >> 0) & 1) << (7-i))\n palette[j*3 + 1] |= (((lab >> 1) & 1) << (7-i))\n palette[j*3 + 2] |= (((lab >> 2) & 1) << (7-i))\n i = i + 1\n lab >>= 3\n\n return palette.reshape((-1, 3))\n\n\ndef visualize_mask(mask, num_cls):\n palette = get_palette(num_cls)\n mask[mask == -1] = 0\n\n return palette[mask].astype(np.uint8)\n\n\ndef visualize_proposals(proposals_info, point_color=(255, 0, 0), point_radius=1):\n proposal_map, colors, candidates = proposals_info\n\n proposal_map = draw_probmap(proposal_map)\n for x, y in candidates:\n proposal_map = cv2.circle(proposal_map, (y, x), point_radius, point_color, -1)\n\n return proposal_map\n\n\ndef draw_probmap(x):\n return cv2.applyColorMap((x * 255).astype(np.uint8), cv2.COLORMAP_HOT)\n\n\ndef draw_points(image, points, color, radius=3):\n image = image.copy()\n for p in points:\n image = cv2.circle(image, (int(p[1]), int(p[0])), radius, color, -1)\n\n return image\n\n\ndef draw_instance_map(x, palette=None):\n num_colors = x.max() + 1\n if palette is None:\n palette = get_palette(num_colors)\n\n return palette[x].astype(np.uint8)\n\n\ndef blend_mask(image, mask, alpha=0.6):\n if mask.min() == -1:\n mask = mask.copy() + 1\n\n imap = draw_instance_map(mask)\n result = (image * (1 - alpha) + alpha * imap).astype(np.uint8)\n return result\n\n\ndef get_boundaries(instances_masks, boundaries_width=1):\n boundaries = np.zeros((instances_masks.shape[0], instances_masks.shape[1]), dtype=np.bool)\n\n for obj_id in np.unique(instances_masks.flatten()):\n if obj_id == 0:\n continue\n\n obj_mask = instances_masks == obj_id\n kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))\n inner_mask = cv2.erode(obj_mask.astype(np.uint8), kernel, iterations=boundaries_width).astype(np.bool)\n\n obj_boundary = np.logical_xor(obj_mask, np.logical_and(inner_mask, obj_mask))\n boundaries = np.logical_or(boundaries, obj_boundary)\n return boundaries\n \n \ndef draw_with_blend_and_clicks(img, mask=None, alpha=0.6, clicks_list=None, pos_color=(0, 255, 0),\n neg_color=(255, 0, 0), radius=4):\n result = img.copy()\n\n if mask is not None:\n palette = get_palette(np.max(mask) + 1)\n rgb_mask = palette[mask.astype(np.uint8)]\n\n mask_region = (mask > 0).astype(np.uint8)\n result = result * (1 - mask_region[:, :, np.newaxis]) + \\\n (1 - alpha) * mask_region[:, :, np.newaxis] * result + \\\n alpha * rgb_mask\n result = result.astype(np.uint8)\n\n # result = (result * (1 - alpha) + alpha * rgb_mask).astype(np.uint8)\n\n if clicks_list is not None and len(clicks_list) > 0:\n pos_points = [click.coords for click in clicks_list if click.is_positive]\n neg_points = [click.coords for click in clicks_list if not click.is_positive]\n\n result = draw_points(result, pos_points, pos_color, radius=radius)\n result = draw_points(result, neg_points, neg_color, radius=radius)\n\n return result\n\n" }, { "path": "XMem/inference/interact/fbrs_controller.py", "content": "import torch\nfrom .fbrs.controller import InteractiveController\nfrom .fbrs.inference import utils\n\n\nclass FBRSController:\n def __init__(self, checkpoint_path, device='cuda:0', max_size=800):\n model = utils.load_is_model(checkpoint_path, device, cpu_dist_maps=True, norm_radius=260)\n\n # Predictor params\n zoomin_params = {\n 'skip_clicks': 1,\n 'target_size': 480,\n 'expansion_ratio': 1.4,\n }\n\n predictor_params = {\n 'brs_mode': 'f-BRS-B',\n 'prob_thresh': 0.5,\n 'zoom_in_params': zoomin_params,\n 'predictor_params': {\n 'net_clicks_limit': 8,\n 'max_size': 800,\n },\n 'brs_opt_func_params': {'min_iou_diff': 1e-3},\n 'lbfgs_params': {'maxfun': 20}\n }\n\n self.controller = InteractiveController(model, device, predictor_params)\n self.anchored = False\n self.device = device\n\n def unanchor(self):\n self.anchored = False\n\n def interact(self, image, x, y, is_positive):\n image = image.to(self.device, non_blocking=True)\n if not self.anchored:\n self.controller.set_image(image)\n self.controller.reset_predictor()\n self.anchored = True\n\n self.controller.add_click(x, y, is_positive)\n # return self.controller.result_mask\n # return self.controller.probs_history[-1][1]\n return (self.controller.probs_history[-1][1]>0.5).float()\n\n def undo(self):\n self.controller.undo_click()\n if len(self.controller.probs_history) == 0:\n return None\n else:\n return (self.controller.probs_history[-1][1]>0.5).float()" }, { "path": "XMem/inference/interact/gui.py", "content": "\"\"\"\nBased on https://github.com/hkchengrex/MiVOS/tree/MiVOS-STCN \n(which is based on https://github.com/seoungwugoh/ivs-demo)\n\nThis version is much simplified. \nIn this repo, we don't have\n- local control\n- fusion module\n- undo\n- timers\n\nbut with XMem as the backbone and is more memory (for both CPU and GPU) friendly\n\"\"\"\n\nimport functools\n\nimport os\nimport cv2\n# fix conflicts between qt5 and cv2\nos.environ.pop(\"QT_QPA_PLATFORM_PLUGIN_PATH\")\n\nimport numpy as np\nimport torch\ntry:\n from torch import mps\nexcept:\n print('torch.MPS not available.')\n\nfrom PySide6.QtWidgets import (QWidget, QApplication, QComboBox, QCheckBox,\n QHBoxLayout, QLabel, QPushButton, QTextEdit, QSpinBox, QFileDialog,\n QPlainTextEdit, QVBoxLayout, QSizePolicy, QButtonGroup, QSlider, QRadioButton)\n\nfrom PySide6.QtGui import QPixmap, QKeySequence, QImage, QTextCursor, QIcon, QShortcut\nfrom PySide6.QtCore import Qt, QTimer\n\nfrom model.network import XMem\n\nfrom inference.inference_core import InferenceCore\nfrom .s2m_controller import S2MController\nfrom .fbrs_controller import FBRSController\n\nfrom .interactive_utils import *\nfrom .interaction import *\nfrom .resource_manager import ResourceManager\nfrom .gui_utils import *\n\n\nclass App(QWidget):\n def __init__(self, net: XMem, \n resource_manager: ResourceManager, \n s2m_ctrl:S2MController, \n fbrs_ctrl:FBRSController, config, device):\n super().__init__()\n\n self.initialized = False\n self.num_objects = config['num_objects']\n self.s2m_controller = s2m_ctrl\n self.fbrs_controller = fbrs_ctrl\n self.config = config\n self.processor = InferenceCore(net, config)\n self.processor.set_all_labels(list(range(1, self.num_objects+1)))\n self.res_man = resource_manager\n self.device = device\n\n self.num_frames = len(self.res_man)\n self.height, self.width = self.res_man.h, self.res_man.w\n\n # set window\n self.setWindowTitle('XMem Demo')\n self.setGeometry(100, 100, self.width, self.height+100)\n self.setWindowIcon(QIcon('docs/icon.png'))\n\n # some buttons\n self.play_button = QPushButton('Play Video')\n self.play_button.clicked.connect(self.on_play_video)\n self.commit_button = QPushButton('Commit')\n self.commit_button.clicked.connect(self.on_commit)\n self.export_button = QPushButton('Export Overlays as Video')\n self.export_button.clicked.connect(self.on_export_visualization)\n\n self.forward_run_button = QPushButton('Forward Propagate')\n self.forward_run_button.clicked.connect(self.on_forward_propagation)\n self.forward_run_button.setMinimumWidth(150)\n\n self.backward_run_button = QPushButton('Backward Propagate')\n self.backward_run_button.clicked.connect(self.on_backward_propagation)\n self.backward_run_button.setMinimumWidth(150)\n\n self.reset_button = QPushButton('Reset Frame')\n self.reset_button.clicked.connect(self.on_reset_mask)\n\n # LCD\n self.lcd = QTextEdit()\n self.lcd.setReadOnly(True)\n self.lcd.setMaximumHeight(28)\n self.lcd.setMaximumWidth(120)\n self.lcd.setText('{: 4d} / {: 4d}'.format(0, self.num_frames-1))\n\n # Current Mask LCD\n self.object_dial = QSpinBox()\n self.object_dial.setReadOnly(False)\n self.object_dial.setMaximumHeight(28)\n self.object_dial.setMaximumWidth(56)\n self.object_dial.setMinimum(1)\n self.object_dial.setMaximum(self.num_objects)\n self.object_dial.editingFinished.connect(self.on_object_dial_change)\n\n # timeline slider\n self.tl_slider = QSlider(Qt.Orientation.Horizontal)\n self.tl_slider.valueChanged.connect(self.tl_slide)\n self.tl_slider.setMinimum(0)\n self.tl_slider.setMaximum(self.num_frames-1)\n self.tl_slider.setValue(0)\n self.tl_slider.setTickPosition(QSlider.TickPosition.TicksBelow)\n self.tl_slider.setTickInterval(1)\n \n # brush size slider\n self.brush_label = QLabel()\n self.brush_label.setAlignment(Qt.AlignmentFlag.AlignCenter)\n self.brush_label.setMinimumWidth(150)\n \n self.brush_slider = QSlider(Qt.Orientation.Horizontal)\n self.brush_slider.valueChanged.connect(self.brush_slide)\n self.brush_slider.setMinimum(1)\n self.brush_slider.setMaximum(100)\n self.brush_slider.setValue(3)\n self.brush_slider.setTickPosition(QSlider.TickPosition.TicksBelow)\n self.brush_slider.setTickInterval(2)\n self.brush_slider.setMinimumWidth(300)\n\n # combobox\n self.combo = QComboBox(self)\n self.combo.addItem(\"davis\")\n self.combo.addItem(\"fade\")\n self.combo.addItem(\"light\")\n self.combo.addItem(\"popup\")\n self.combo.addItem(\"layered\")\n self.combo.currentTextChanged.connect(self.set_viz_mode)\n\n self.save_visualization_checkbox = QCheckBox(self)\n self.save_visualization_checkbox.toggled.connect(self.on_save_visualization_toggle)\n self.save_visualization_checkbox.setChecked(False)\n self.save_visualization = False\n\n # Radio buttons for type of interactions\n self.curr_interaction = 'Click'\n self.interaction_group = QButtonGroup()\n self.radio_fbrs = QRadioButton('Click')\n self.radio_s2m = QRadioButton('Scribble')\n self.radio_free = QRadioButton('Free')\n self.interaction_group.addButton(self.radio_fbrs)\n self.interaction_group.addButton(self.radio_s2m)\n self.interaction_group.addButton(self.radio_free)\n self.radio_fbrs.toggled.connect(self.interaction_radio_clicked)\n self.radio_s2m.toggled.connect(self.interaction_radio_clicked)\n self.radio_free.toggled.connect(self.interaction_radio_clicked)\n self.radio_fbrs.toggle()\n\n # Main canvas -> QLabel\n self.main_canvas = QLabel()\n self.main_canvas.setSizePolicy(QSizePolicy.Policy.Expanding,\n QSizePolicy.Policy.Expanding)\n self.main_canvas.setAlignment(Qt.AlignmentFlag.AlignCenter)\n self.main_canvas.setMinimumSize(100, 100)\n\n self.main_canvas.mousePressEvent = self.on_mouse_press\n self.main_canvas.mouseMoveEvent = self.on_mouse_motion\n self.main_canvas.setMouseTracking(True) # Required for all-time tracking\n self.main_canvas.mouseReleaseEvent = self.on_mouse_release\n\n # Minimap -> Also a QLabel\n self.minimap = QLabel()\n self.minimap.setSizePolicy(QSizePolicy.Policy.Expanding,\n QSizePolicy.Policy.Expanding)\n self.minimap.setAlignment(Qt.AlignmentFlag.AlignTop)\n self.minimap.setMinimumSize(100, 100)\n\n # Zoom-in buttons\n self.zoom_p_button = QPushButton('Zoom +')\n self.zoom_p_button.clicked.connect(self.on_zoom_plus)\n self.zoom_m_button = QPushButton('Zoom -')\n self.zoom_m_button.clicked.connect(self.on_zoom_minus)\n\n # Parameters setting\n self.clear_mem_button = QPushButton('Clear memory')\n self.clear_mem_button.clicked.connect(self.on_clear_memory)\n\n self.work_mem_gauge, self.work_mem_gauge_layout = create_gauge('Working memory size')\n self.long_mem_gauge, self.long_mem_gauge_layout = create_gauge('Long-term memory size')\n self.gpu_mem_gauge, self.gpu_mem_gauge_layout = create_gauge('GPU mem. (all processes, w/ caching)')\n self.torch_mem_gauge, self.torch_mem_gauge_layout = create_gauge('GPU mem. (used by torch, w/o caching)')\n\n self.update_memory_size()\n self.update_gpu_usage()\n\n self.work_mem_min, self.work_mem_min_layout = create_parameter_box(1, 100, 'Min. working memory frames', \n callback=self.on_work_min_change)\n self.work_mem_max, self.work_mem_max_layout = create_parameter_box(2, 100, 'Max. working memory frames', \n callback=self.on_work_max_change)\n self.long_mem_max, self.long_mem_max_layout = create_parameter_box(1000, 100000, \n 'Max. long-term memory size', step=1000, callback=self.update_config)\n self.num_prototypes_box, self.num_prototypes_box_layout = create_parameter_box(32, 1280, \n 'Number of prototypes', step=32, callback=self.update_config)\n self.mem_every_box, self.mem_every_box_layout = create_parameter_box(1, 100, 'Memory frame every (r)', \n callback=self.update_config)\n\n self.work_mem_min.setValue(self.processor.memory.min_mt_frames)\n self.work_mem_max.setValue(self.processor.memory.max_mt_frames)\n self.long_mem_max.setValue(self.processor.memory.max_long_elements)\n self.num_prototypes_box.setValue(self.processor.memory.num_prototypes)\n self.mem_every_box.setValue(self.processor.mem_every)\n\n # import mask/layer\n self.import_mask_button = QPushButton('Import mask')\n self.import_mask_button.clicked.connect(self.on_import_mask)\n self.import_layer_button = QPushButton('Import layer')\n self.import_layer_button.clicked.connect(self.on_import_layer)\n\n # Console on the GUI\n self.console = QPlainTextEdit()\n self.console.setReadOnly(True)\n self.console.setMinimumHeight(100)\n self.console.setMaximumHeight(100)\n\n # navigator\n navi = QHBoxLayout()\n\n interact_subbox = QVBoxLayout()\n interact_topbox = QHBoxLayout()\n interact_botbox = QHBoxLayout()\n interact_topbox.setAlignment(Qt.AlignmentFlag.AlignCenter)\n interact_topbox.addWidget(self.lcd)\n interact_topbox.addWidget(self.play_button)\n interact_topbox.addWidget(self.radio_s2m)\n interact_topbox.addWidget(self.radio_fbrs)\n interact_topbox.addWidget(self.radio_free)\n interact_topbox.addWidget(self.reset_button)\n interact_botbox.addWidget(QLabel('Current Object ID:'))\n interact_botbox.addWidget(self.object_dial)\n interact_botbox.addWidget(self.brush_label)\n interact_botbox.addWidget(self.brush_slider)\n interact_subbox.addLayout(interact_topbox)\n interact_subbox.addLayout(interact_botbox)\n navi.addLayout(interact_subbox)\n\n apply_fixed_size_policy = lambda x: x.setSizePolicy(QSizePolicy.Policy.Fixed, \n QSizePolicy.Policy.Fixed)\n apply_to_all_children_widget(interact_topbox, apply_fixed_size_policy)\n apply_to_all_children_widget(interact_botbox, apply_fixed_size_policy)\n\n navi.addStretch(1)\n navi.addStretch(1)\n overlay_subbox = QVBoxLayout()\n overlay_topbox = QHBoxLayout()\n overlay_botbox = QHBoxLayout()\n overlay_botbox.setAlignment(Qt.AlignmentFlag.AlignRight)\n overlay_topbox.addWidget(QLabel('Overlay Mode'))\n overlay_topbox.addWidget(self.combo)\n overlay_topbox.addWidget(QLabel('Save overlay during propagation'))\n overlay_topbox.addWidget(self.save_visualization_checkbox)\n overlay_botbox.addWidget(self.export_button)\n overlay_subbox.addLayout(overlay_topbox)\n overlay_subbox.addLayout(overlay_botbox)\n navi.addLayout(overlay_subbox)\n apply_to_all_children_widget(overlay_topbox, apply_fixed_size_policy)\n apply_to_all_children_widget(overlay_botbox, apply_fixed_size_policy)\n\n navi.addStretch(1)\n navi.addWidget(self.commit_button)\n navi.addWidget(self.forward_run_button)\n navi.addWidget(self.backward_run_button)\n \n # Drawing area, main canvas and minimap\n draw_area = QHBoxLayout()\n draw_area.addWidget(self.main_canvas, 4)\n\n # Minimap area\n minimap_area = QVBoxLayout()\n minimap_area.setAlignment(Qt.AlignmentFlag.AlignTop)\n mini_label = QLabel('Minimap')\n mini_label.setAlignment(Qt.AlignmentFlag.AlignTop)\n minimap_area.addWidget(mini_label)\n\n # Minimap zooming\n minimap_ctrl = QHBoxLayout()\n minimap_ctrl.setAlignment(Qt.AlignmentFlag.AlignTop)\n minimap_ctrl.addWidget(self.zoom_p_button)\n minimap_ctrl.addWidget(self.zoom_m_button)\n minimap_area.addLayout(minimap_ctrl)\n minimap_area.addWidget(self.minimap)\n\n # Parameters \n minimap_area.addLayout(self.work_mem_gauge_layout)\n minimap_area.addLayout(self.long_mem_gauge_layout)\n minimap_area.addLayout(self.gpu_mem_gauge_layout)\n minimap_area.addLayout(self.torch_mem_gauge_layout)\n minimap_area.addWidget(self.clear_mem_button)\n minimap_area.addLayout(self.work_mem_min_layout)\n minimap_area.addLayout(self.work_mem_max_layout)\n minimap_area.addLayout(self.long_mem_max_layout)\n minimap_area.addLayout(self.num_prototypes_box_layout)\n minimap_area.addLayout(self.mem_every_box_layout)\n\n # import mask/layer\n import_area = QHBoxLayout()\n import_area.setAlignment(Qt.AlignmentFlag.AlignTop)\n import_area.addWidget(self.import_mask_button)\n import_area.addWidget(self.import_layer_button)\n minimap_area.addLayout(import_area)\n\n # console\n minimap_area.addWidget(self.console)\n\n draw_area.addLayout(minimap_area, 1)\n\n layout = QVBoxLayout()\n layout.addLayout(draw_area)\n layout.addWidget(self.tl_slider)\n layout.addLayout(navi)\n self.setLayout(layout)\n\n # timer to play video\n self.timer = QTimer()\n self.timer.setSingleShot(False)\n self.timer.timeout.connect(self.on_play_video_timer)\n\n # timer to update GPU usage\n self.gpu_timer = QTimer()\n self.gpu_timer.setSingleShot(False)\n self.gpu_timer.timeout.connect(self.on_gpu_timer)\n self.gpu_timer.setInterval(2000)\n self.gpu_timer.start()\n\n # current frame info\n self.curr_frame_dirty = False\n self.current_image = np.zeros((self.height, self.width, 3), dtype=np.uint8) \n self.current_image_torch = None\n self.current_mask = np.zeros((self.height, self.width), dtype=np.uint8)\n self.current_prob = torch.zeros((self.num_objects, self.height, self.width), dtype=torch.float).to(self.device)\n\n # initialize visualization\n self.viz_mode = 'davis'\n self.vis_map = np.zeros((self.height, self.width, 3), dtype=np.uint8)\n self.vis_alpha = np.zeros((self.height, self.width, 1), dtype=np.float32)\n self.brush_vis_map = np.zeros((self.height, self.width, 3), dtype=np.uint8)\n self.brush_vis_alpha = np.zeros((self.height, self.width, 1), dtype=np.float32)\n self.cursur = 0\n self.on_showing = None\n\n # Zoom parameters\n self.zoom_pixels = 150\n \n # initialize action\n self.interaction = None\n self.pressed = False\n self.right_click = False\n self.current_object = 1\n self.last_ex = self.last_ey = 0\n\n self.propagating = False\n\n # Objects shortcuts\n for i in range(1, self.num_objects+1):\n QShortcut(QKeySequence(str(i)), self).activated.connect(functools.partial(self.hit_number_key, i))\n QShortcut(QKeySequence(f\"Ctrl+{i}\"), self).activated.connect(functools.partial(self.hit_number_key, i))\n\n # <- and -> shortcuts\n QShortcut(QKeySequence(Qt.Key.Key_Left), self).activated.connect(self.on_prev_frame)\n QShortcut(QKeySequence(Qt.Key.Key_Right), self).activated.connect(self.on_next_frame)\n\n self.interacted_prob = None\n self.overlay_layer = None\n self.overlay_layer_torch = None\n\n # the object id used for popup/layered overlay\n self.vis_target_objects = [1]\n # try to load the default overlay\n self._try_load_layer('./docs/ECCV-logo.png')\n\n self.load_current_image_mask()\n self.show_current_frame()\n self.show()\n\n self.console_push_text('Initialized.')\n self.initialized = True\n\n def resizeEvent(self, event):\n self.show_current_frame()\n\n def console_push_text(self, text):\n self.console.moveCursor(QTextCursor.MoveOperation.End)\n self.console.insertPlainText(text+'\\n')\n\n def interaction_radio_clicked(self, event):\n self.last_interaction = self.curr_interaction\n if self.radio_s2m.isChecked():\n self.curr_interaction = 'Scribble'\n self.brush_size = 3\n self.brush_slider.setDisabled(True)\n elif self.radio_fbrs.isChecked():\n self.curr_interaction = 'Click'\n self.brush_size = 3\n self.brush_slider.setDisabled(True)\n elif self.radio_free.isChecked():\n self.brush_slider.setDisabled(False)\n self.brush_slide()\n self.curr_interaction = 'Free'\n if self.curr_interaction == 'Scribble':\n self.commit_button.setEnabled(True)\n else:\n self.commit_button.setEnabled(False)\n\n def load_current_image_mask(self, no_mask=False):\n self.current_image = self.res_man.get_image(self.cursur)\n self.current_image_torch = None\n\n if not no_mask:\n loaded_mask = self.res_man.get_mask(self.cursur)\n if loaded_mask is None:\n self.current_mask.fill(0)\n else:\n self.current_mask = loaded_mask.copy()\n self.current_prob = None\n\n def load_current_torch_image_mask(self, no_mask=False):\n if self.current_image_torch is None:\n self.current_image_torch, self.current_image_torch_no_norm = image_to_torch(self.current_image, self.device)\n\n if self.current_prob is None and not no_mask:\n self.current_prob = index_numpy_to_one_hot_torch(self.current_mask, self.num_objects+1).to(self.device)\n\n def compose_current_im(self):\n self.viz = get_visualization(self.viz_mode, self.current_image, self.current_mask, \n self.overlay_layer, self.vis_target_objects)\n\n def update_interact_vis(self):\n # Update the interactions without re-computing the overlay\n height, width, channel = self.viz.shape\n bytesPerLine = 3 * width\n\n vis_map = self.vis_map\n vis_alpha = self.vis_alpha\n brush_vis_map = self.brush_vis_map\n brush_vis_alpha = self.brush_vis_alpha\n\n self.viz_with_stroke = self.viz*(1-vis_alpha) + vis_map*vis_alpha\n self.viz_with_stroke = self.viz_with_stroke*(1-brush_vis_alpha) + brush_vis_map*brush_vis_alpha\n self.viz_with_stroke = self.viz_with_stroke.astype(np.uint8)\n\n qImg = QImage(self.viz_with_stroke.data, width, height, bytesPerLine, QImage.Format.Format_RGB888)\n self.main_canvas.setPixmap(QPixmap(qImg.scaled(self.main_canvas.size(),\n Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.FastTransformation)))\n\n self.main_canvas_size = self.main_canvas.size()\n self.image_size = qImg.size()\n\n def update_minimap(self):\n ex, ey = self.last_ex, self.last_ey\n r = self.zoom_pixels//2\n ex = int(round(max(r, min(self.width-r, ex))))\n ey = int(round(max(r, min(self.height-r, ey))))\n\n patch = self.viz_with_stroke[ey-r:ey+r, ex-r:ex+r, :].astype(np.uint8)\n\n height, width, channel = patch.shape\n bytesPerLine = 3 * width\n qImg = QImage(patch.data, width, height, bytesPerLine, QImage.Format.Format_RGB888)\n self.minimap.setPixmap(QPixmap(qImg.scaled(self.minimap.size(),\n Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.FastTransformation)))\n\n def update_current_image_fast(self):\n # fast path, uses gpu. Changes the image in-place to avoid copying\n self.viz = get_visualization_torch(self.viz_mode, self.current_image_torch_no_norm, \n self.current_prob, self.overlay_layer_torch, self.vis_target_objects)\n if self.save_visualization:\n self.res_man.save_visualization(self.cursur, self.viz)\n\n height, width, channel = self.viz.shape\n bytesPerLine = 3 * width\n\n qImg = QImage(self.viz.data, width, height, bytesPerLine, QImage.Format.Format_RGB888)\n self.main_canvas.setPixmap(QPixmap(qImg.scaled(self.main_canvas.size(),\n Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.FastTransformation)))\n\n def show_current_frame(self, fast=False):\n # Re-compute overlay and show the image\n if fast:\n self.update_current_image_fast()\n else:\n self.compose_current_im()\n self.update_interact_vis()\n self.update_minimap()\n\n self.lcd.setText('{: 3d} / {: 3d}'.format(self.cursur, self.num_frames-1))\n self.tl_slider.setValue(self.cursur)\n\n def pixel_pos_to_image_pos(self, x, y):\n # Un-scale and un-pad the label coordinates into image coordinates\n oh, ow = self.image_size.height(), self.image_size.width()\n nh, nw = self.main_canvas_size.height(), self.main_canvas_size.width()\n\n h_ratio = nh/oh\n w_ratio = nw/ow\n dominate_ratio = min(h_ratio, w_ratio)\n\n # Solve scale\n x /= dominate_ratio\n y /= dominate_ratio\n\n # Solve padding\n fh, fw = nh/dominate_ratio, nw/dominate_ratio\n x -= (fw-ow)/2\n y -= (fh-oh)/2\n\n return x, y\n\n def is_pos_out_of_bound(self, x, y):\n x, y = self.pixel_pos_to_image_pos(x, y)\n\n out_of_bound = (\n (x < 0) or\n (y < 0) or\n (x > self.width-1) or \n (y > self.height-1)\n )\n\n return out_of_bound\n\n def get_scaled_pos(self, x, y):\n x, y = self.pixel_pos_to_image_pos(x, y)\n\n x = max(0, min(self.width-1, x))\n y = max(0, min(self.height-1, y))\n\n return x, y\n\n def clear_visualization(self):\n self.vis_map.fill(0)\n self.vis_alpha.fill(0)\n\n def reset_this_interaction(self):\n self.complete_interaction()\n self.clear_visualization()\n self.interaction = None\n if self.fbrs_controller is not None:\n self.fbrs_controller.unanchor()\n\n def set_viz_mode(self):\n self.viz_mode = self.combo.currentText()\n self.show_current_frame()\n\n def save_current_mask(self):\n # save mask to hard disk\n self.res_man.save_mask(self.cursur, self.current_mask)\n\n def tl_slide(self):\n # if we are propagating, the on_run function will take care of everything\n # don't do duplicate work here\n if not self.propagating:\n if self.curr_frame_dirty:\n self.save_current_mask()\n self.curr_frame_dirty = False\n\n self.reset_this_interaction()\n self.cursur = self.tl_slider.value()\n self.load_current_image_mask()\n self.show_current_frame()\n\n def brush_slide(self):\n self.brush_size = self.brush_slider.value()\n self.brush_label.setText('Brush size (in free mode): %d' % self.brush_size)\n try:\n if type(self.interaction) == FreeInteraction:\n self.interaction.set_size(self.brush_size)\n except AttributeError:\n # Initialization, forget about it\n pass\n\n def on_forward_propagation(self):\n if self.propagating:\n # acts as a pause button\n self.propagating = False\n else:\n self.propagate_fn = self.on_next_frame\n self.backward_run_button.setEnabled(False)\n self.forward_run_button.setText('Pause Propagation')\n self.on_propagation()\n\n def on_backward_propagation(self):\n if self.propagating:\n # acts as a pause button\n self.propagating = False\n else:\n self.propagate_fn = self.on_prev_frame\n self.forward_run_button.setEnabled(False)\n self.backward_run_button.setText('Pause Propagation')\n self.on_propagation()\n\n def on_pause(self):\n self.propagating = False\n self.forward_run_button.setEnabled(True)\n self.backward_run_button.setEnabled(True)\n self.clear_mem_button.setEnabled(True)\n self.forward_run_button.setText('Forward Propagate')\n self.backward_run_button.setText('Backward Propagate')\n self.console_push_text('Propagation stopped.')\n\n def on_propagation(self):\n # start to propagate\n self.load_current_torch_image_mask()\n self.show_current_frame(fast=True)\n\n self.console_push_text('Propagation started.')\n self.current_prob = self.processor.step(self.current_image_torch, self.current_prob[1:])\n self.current_mask = torch_prob_to_numpy_mask(self.current_prob)\n # clear\n self.interacted_prob = None\n self.reset_this_interaction()\n\n self.propagating = True\n self.clear_mem_button.setEnabled(False)\n # propagate till the end\n while self.propagating:\n self.propagate_fn()\n\n self.load_current_image_mask(no_mask=True)\n self.load_current_torch_image_mask(no_mask=True)\n\n self.current_prob = self.processor.step(self.current_image_torch)\n self.current_mask = torch_prob_to_numpy_mask(self.current_prob)\n\n self.save_current_mask()\n self.show_current_frame(fast=True)\n\n self.update_memory_size()\n QApplication.processEvents()\n\n if self.cursur == 0 or self.cursur == self.num_frames-1:\n break\n\n self.propagating = False\n self.curr_frame_dirty = False\n self.on_pause()\n self.tl_slide()\n QApplication.processEvents()\n\n def pause_propagation(self):\n self.propagating = False\n\n def on_commit(self):\n self.complete_interaction()\n self.update_interacted_mask()\n\n def on_prev_frame(self):\n # self.tl_slide will trigger on setValue\n self.cursur = max(0, self.cursur-1)\n self.tl_slider.setValue(self.cursur)\n\n def on_next_frame(self):\n # self.tl_slide will trigger on setValue\n self.cursur = min(self.cursur+1, self.num_frames-1)\n self.tl_slider.setValue(self.cursur)\n\n def on_play_video_timer(self):\n self.cursur += 1\n if self.cursur > self.num_frames-1:\n self.cursur = 0\n self.tl_slider.setValue(self.cursur)\n\n def on_play_video(self):\n if self.timer.isActive():\n self.timer.stop()\n self.play_button.setText('Play Video')\n else:\n self.timer.start(1000 // 30)\n self.play_button.setText('Stop Video')\n\n def on_export_visualization(self):\n # NOTE: Save visualization at the end of propagation\n image_folder = f\"{self.config['workspace']}/visualization/\"\n save_folder = self.config['workspace']\n if os.path.exists(image_folder):\n # Sorted so frames will be in order\n self.console_push_text(f'Exporting visualization to {self.config[\"workspace\"]}/visualization.mp4')\n images = [img for img in sorted(os.listdir(image_folder)) if img.endswith(\".jpg\")]\n frame = cv2.imread(os.path.join(image_folder, images[0]))\n height, width, layers = frame.shape\n # 10 is the FPS -- change if needed\n video = cv2.VideoWriter(f\"{save_folder}/visualization.mp4\", cv2.VideoWriter_fourcc(*'mp4v'), 10, (width,height))\n for image in images:\n video.write(cv2.imread(os.path.join(image_folder, image)))\n video.release()\n self.console_push_text(f'Visualization exported to {self.config[\"workspace\"]}/visualization.mp4')\n else:\n self.console_push_text(f'No visualization images found in {image_folder}')\n\n def on_object_dial_change(self):\n object_id = self.object_dial.value()\n self.hit_number_key(object_id)\n\n def on_reset_mask(self):\n self.current_mask.fill(0)\n if self.current_prob is not None:\n self.current_prob.fill_(0)\n self.curr_frame_dirty = True\n self.save_current_mask()\n self.reset_this_interaction()\n self.show_current_frame()\n\n def on_zoom_plus(self):\n self.zoom_pixels -= 25\n self.zoom_pixels = max(50, self.zoom_pixels)\n self.update_minimap()\n\n def on_zoom_minus(self):\n self.zoom_pixels += 25\n self.zoom_pixels = min(self.zoom_pixels, 300)\n self.update_minimap()\n\n def set_navi_enable(self, boolean):\n self.zoom_p_button.setEnabled(boolean)\n self.zoom_m_button.setEnabled(boolean)\n self.run_button.setEnabled(boolean)\n self.tl_slider.setEnabled(boolean)\n self.play_button.setEnabled(boolean)\n self.export_button.setEnabled(boolean)\n self.lcd.setEnabled(boolean)\n\n def hit_number_key(self, number):\n if number == self.current_object:\n return\n self.current_object = number\n self.object_dial.setValue(number)\n if self.fbrs_controller is not None:\n self.fbrs_controller.unanchor()\n self.console_push_text(f'Current object changed to {number}.')\n self.clear_brush()\n self.vis_brush(self.last_ex, self.last_ey)\n self.update_interact_vis()\n self.show_current_frame()\n\n def clear_brush(self):\n self.brush_vis_map.fill(0)\n self.brush_vis_alpha.fill(0)\n\n def vis_brush(self, ex, ey):\n self.brush_vis_map = cv2.circle(self.brush_vis_map, \n (int(round(ex)), int(round(ey))), self.brush_size//2+1, color_map[self.current_object], thickness=-1)\n self.brush_vis_alpha = cv2.circle(self.brush_vis_alpha, \n (int(round(ex)), int(round(ey))), self.brush_size//2+1, 0.5, thickness=-1)\n\n def on_mouse_press(self, event):\n if self.is_pos_out_of_bound(event.position().x(), event.position().y()):\n return\n\n # mid-click\n if (event.button() == Qt.MouseButton.MiddleButton):\n ex, ey = self.get_scaled_pos(event.position().x(), event.position().y())\n target_object = self.current_mask[int(ey),int(ex)]\n if target_object in self.vis_target_objects:\n self.vis_target_objects.remove(target_object)\n else:\n self.vis_target_objects.append(target_object)\n self.console_push_text(f'Target objects for visualization changed to {self.vis_target_objects}')\n self.show_current_frame()\n return\n\n self.right_click = (event.button() == Qt.MouseButton.RightButton)\n self.pressed = True\n\n h, w = self.height, self.width\n\n self.load_current_torch_image_mask()\n image = self.current_image_torch\n\n last_interaction = self.interaction\n new_interaction = None\n if self.curr_interaction == 'Scribble':\n if last_interaction is None or type(last_interaction) != ScribbleInteraction:\n self.complete_interaction()\n new_interaction = ScribbleInteraction(image, torch.from_numpy(self.current_mask).float().to(self.device), \n (h, w), self.s2m_controller, self.num_objects)\n elif self.curr_interaction == 'Free':\n if last_interaction is None or type(last_interaction) != FreeInteraction:\n self.complete_interaction()\n new_interaction = FreeInteraction(image, self.current_mask, (h, w), \n self.num_objects)\n new_interaction.set_size(self.brush_size)\n elif self.curr_interaction == 'Click':\n if (last_interaction is None or type(last_interaction) != ClickInteraction \n or last_interaction.tar_obj != self.current_object):\n self.complete_interaction()\n self.fbrs_controller.unanchor()\n new_interaction = ClickInteraction(image, self.current_prob, (h, w), \n self.fbrs_controller, self.current_object)\n\n if new_interaction is not None:\n self.interaction = new_interaction\n\n # Just motion it as the first step\n self.on_mouse_motion(event)\n\n def on_mouse_motion(self, event):\n ex, ey = self.get_scaled_pos(event.position().x(), event.position().y())\n self.last_ex, self.last_ey = ex, ey\n self.clear_brush()\n # Visualize\n self.vis_brush(ex, ey)\n if self.pressed:\n if self.curr_interaction == 'Scribble' or self.curr_interaction == 'Free':\n obj = 0 if self.right_click else self.current_object\n self.vis_map, self.vis_alpha = self.interaction.push_point(\n ex, ey, obj, (self.vis_map, self.vis_alpha)\n )\n self.update_interact_vis()\n self.update_minimap()\n\n def update_interacted_mask(self):\n self.current_prob = self.interacted_prob\n self.current_mask = torch_prob_to_numpy_mask(self.interacted_prob)\n self.show_current_frame()\n self.save_current_mask()\n self.curr_frame_dirty = False\n\n def complete_interaction(self):\n if self.interaction is not None:\n self.clear_visualization()\n self.interaction = None\n\n def on_mouse_release(self, event):\n if not self.pressed:\n # this can happen when the initial press is out-of-bound\n return\n\n ex, ey = self.get_scaled_pos(event.position().x(), event.position().y())\n\n self.console_push_text('%s interaction at frame %d.' % (self.curr_interaction, self.cursur))\n interaction = self.interaction\n\n if self.curr_interaction == 'Scribble' or self.curr_interaction == 'Free':\n self.on_mouse_motion(event)\n interaction.end_path()\n if self.curr_interaction == 'Free':\n self.clear_visualization()\n elif self.curr_interaction == 'Click':\n ex, ey = self.get_scaled_pos(event.position().x(), event.position().y())\n self.vis_map, self.vis_alpha = interaction.push_point(ex, ey,\n self.right_click, (self.vis_map, self.vis_alpha))\n\n self.interacted_prob = interaction.predict().to(self.device)\n self.update_interacted_mask()\n self.update_gpu_usage()\n\n self.pressed = self.right_click = False\n\n def wheelEvent(self, event):\n ex, ey = self.get_scaled_pos(event.position().x(), event.position().y())\n if self.curr_interaction == 'Free':\n self.brush_slider.setValue(self.brush_slider.value() + event.angleDelta().y()//30)\n self.clear_brush()\n self.vis_brush(ex, ey)\n self.update_interact_vis()\n self.update_minimap()\n\n def update_gpu_usage(self):\n if self.device.type == 'cuda':\n info = torch.cuda.mem_get_info()\n elif self.device.type == 'mps':\n info = (0, mps.current_allocated_memory()) # NOTE: torch.mps does not support accessing free and total memory\n else:\n info = (0, 0)\n global_free, global_total = info\n global_free /= (2**30)\n global_total /= (2**30)\n global_used = global_total - global_free\n\n self.gpu_mem_gauge.setFormat(f'{global_used:.01f} GB / {global_total:.01f} GB')\n self.gpu_mem_gauge.setValue(round(global_used/global_total*100))\n\n used_by_torch = torch.cuda.max_memory_allocated() / (2**20)\n self.torch_mem_gauge.setFormat(f'{used_by_torch:.0f} MB / {global_total:.01f} GB')\n self.torch_mem_gauge.setValue(round(used_by_torch/global_total*100/1024))\n\n def on_gpu_timer(self):\n self.update_gpu_usage()\n\n def update_memory_size(self):\n try:\n max_work_elements = self.processor.memory.max_work_elements\n max_long_elements = self.processor.memory.max_long_elements\n\n curr_work_elements = self.processor.memory.work_mem.size\n curr_long_elements = self.processor.memory.long_mem.size\n\n self.work_mem_gauge.setFormat(f'{curr_work_elements} / {max_work_elements}')\n self.work_mem_gauge.setValue(round(curr_work_elements/max_work_elements*100))\n\n self.long_mem_gauge.setFormat(f'{curr_long_elements} / {max_long_elements}')\n self.long_mem_gauge.setValue(round(curr_long_elements/max_long_elements*100))\n\n except AttributeError:\n self.work_mem_gauge.setFormat('Unknown')\n self.long_mem_gauge.setFormat('Unknown')\n self.work_mem_gauge.setValue(0)\n self.long_mem_gauge.setValue(0)\n\n def on_work_min_change(self):\n if self.initialized:\n self.work_mem_min.setValue(min(self.work_mem_min.value(), self.work_mem_max.value()-1))\n self.update_config()\n\n def on_work_max_change(self):\n if self.initialized:\n self.work_mem_max.setValue(max(self.work_mem_max.value(), self.work_mem_min.value()+1))\n self.update_config()\n\n def update_config(self):\n if self.initialized:\n self.config['min_mid_term_frames'] = self.work_mem_min.value()\n self.config['max_mid_term_frames'] = self.work_mem_max.value()\n self.config['max_long_term_elements'] = self.long_mem_max.value()\n self.config['num_prototypes'] = self.num_prototypes_box.value()\n self.config['mem_every'] = self.mem_every_box.value()\n\n self.processor.update_config(self.config)\n\n def on_clear_memory(self):\n self.processor.clear_memory()\n if self.device.type == 'cuda':\n torch.cuda.empty_cache()\n elif self.device.type == 'mps':\n mps.empty_cache()\n self.update_gpu_usage()\n self.update_memory_size()\n\n def _open_file(self, prompt):\n options = QFileDialog.Options()\n file_name, _ = QFileDialog.getOpenFileName(self, prompt, \"\", \"Image files (*)\", options=options)\n return file_name\n\n def on_import_mask(self):\n file_name = self._open_file('Mask')\n if len(file_name) == 0:\n return\n\n mask = self.res_man.read_external_image(file_name, size=(self.height, self.width))\n\n shape_condition = (\n (len(mask.shape) == 2) and\n (mask.shape[-1] == self.width) and \n (mask.shape[-2] == self.height)\n )\n\n object_condition = (\n mask.max() <= self.num_objects\n )\n\n if not shape_condition:\n self.console_push_text(f'Expected ({self.height}, {self.width}). Got {mask.shape} instead.')\n elif not object_condition:\n self.console_push_text(f'Expected {self.num_objects} objects. Got {mask.max()} objects instead.')\n else:\n self.console_push_text(f'Mask file {file_name} loaded.')\n self.current_image_torch = self.current_prob = None\n self.current_mask = mask\n self.show_current_frame()\n self.save_current_mask()\n\n def on_import_layer(self):\n file_name = self._open_file('Layer')\n if len(file_name) == 0:\n return\n\n self._try_load_layer(file_name)\n\n def _try_load_layer(self, file_name):\n try:\n layer = self.res_man.read_external_image(file_name, size=(self.height, self.width))\n\n if layer.shape[-1] == 3:\n layer = np.concatenate([layer, np.ones_like(layer[:,:,0:1])*255], axis=-1)\n\n condition = (\n (len(layer.shape) == 3) and\n (layer.shape[-1] == 4) and \n (layer.shape[-2] == self.width) and \n (layer.shape[-3] == self.height)\n )\n\n if not condition:\n self.console_push_text(f'Expected ({self.height}, {self.width}, 4). Got {layer.shape}.')\n else:\n self.console_push_text(f'Layer file {file_name} loaded.')\n self.overlay_layer = layer\n self.overlay_layer_torch = torch.from_numpy(layer).float().to(self.device)/255\n self.show_current_frame()\n except FileNotFoundError:\n self.console_push_text(f'{file_name} not found.')\n\n def on_save_visualization_toggle(self):\n self.save_visualization = self.save_visualization_checkbox.isChecked()\n" }, { "path": "XMem/inference/interact/gui_utils.py", "content": "from PySide6.QtCore import Qt\nfrom PySide6.QtWidgets import (QBoxLayout, QHBoxLayout, QLabel, QSpinBox, QVBoxLayout, QProgressBar)\n\n\ndef create_parameter_box(min_val, max_val, text, step=1, callback=None):\n layout = QHBoxLayout()\n\n dial = QSpinBox()\n dial.setMaximumHeight(28)\n dial.setMaximumWidth(150)\n dial.setMinimum(min_val)\n dial.setMaximum(max_val)\n dial.setAlignment(Qt.AlignmentFlag.AlignRight)\n dial.setSingleStep(step)\n dial.valueChanged.connect(callback)\n\n label = QLabel(text)\n label.setAlignment(Qt.AlignmentFlag.AlignRight)\n\n layout.addWidget(label)\n layout.addWidget(dial)\n\n return dial, layout\n\n\ndef create_gauge(text):\n layout = QHBoxLayout()\n\n gauge = QProgressBar()\n gauge.setMaximumHeight(28)\n gauge.setMaximumWidth(200)\n gauge.setAlignment(Qt.AlignmentFlag.AlignCenter)\n\n label = QLabel(text)\n label.setAlignment(Qt.AlignmentFlag.AlignRight)\n\n layout.addWidget(label)\n layout.addWidget(gauge)\n\n return gauge, layout\n\n\ndef apply_to_all_children_widget(layout, func):\n # deliberately non-recursive\n for i in range(layout.count()):\n func(layout.itemAt(i).widget())" }, { "path": "XMem/inference/interact/interaction.py", "content": "\"\"\"\nContains all the types of interaction related to the GUI\nNot related to automatic evaluation in the DAVIS dataset\n\nYou can inherit the Interaction class to create new interaction types\nundo is (sometimes partially) supported\n\"\"\"\n\n\nimport torch\nimport torch.nn.functional as F\nimport numpy as np\nimport cv2\nimport time\nfrom .interactive_utils import color_map, index_numpy_to_one_hot_torch\n\n\ndef aggregate_sbg(prob, keep_bg=False, hard=False):\n device = prob.device\n k, h, w = prob.shape\n ex_prob = torch.zeros((k+1, h, w), device=device)\n ex_prob[0] = 0.5\n ex_prob[1:] = prob\n ex_prob = torch.clamp(ex_prob, 1e-7, 1-1e-7)\n logits = torch.log((ex_prob /(1-ex_prob)))\n\n if hard:\n # Very low temperature o((โŠ™๏นโŠ™))o ๐Ÿฅถ\n logits *= 1000\n\n if keep_bg:\n return F.softmax(logits, dim=0)\n else:\n return F.softmax(logits, dim=0)[1:]\n\ndef aggregate_wbg(prob, keep_bg=False, hard=False):\n k, h, w = prob.shape\n new_prob = torch.cat([\n torch.prod(1-prob, dim=0, keepdim=True),\n prob\n ], 0).clamp(1e-7, 1-1e-7)\n logits = torch.log((new_prob /(1-new_prob)))\n\n if hard:\n # Very low temperature o((โŠ™๏นโŠ™))o ๐Ÿฅถ\n logits *= 1000\n\n if keep_bg:\n return F.softmax(logits, dim=0)\n else:\n return F.softmax(logits, dim=0)[1:]\n\nclass Interaction:\n def __init__(self, image, prev_mask, true_size, controller):\n self.image = image \n self.prev_mask = prev_mask\n self.controller = controller\n self.start_time = time.time()\n\n self.h, self.w = true_size\n\n self.out_prob = None\n self.out_mask = None\n\n def predict(self):\n pass\n\n\nclass FreeInteraction(Interaction):\n def __init__(self, image, prev_mask, true_size, num_objects):\n \"\"\"\n prev_mask should be index format numpy array\n \"\"\"\n super().__init__(image, prev_mask, true_size, None)\n\n self.K = num_objects\n\n self.drawn_map = self.prev_mask.copy()\n self.curr_path = [[] for _ in range(self.K + 1)]\n\n self.size = None\n\n def set_size(self, size):\n self.size = size\n\n \"\"\"\n k - object id\n vis - a tuple (visualization map, pass through alpha). None if not needed.\n \"\"\"\n def push_point(self, x, y, k, vis=None):\n if vis is not None:\n vis_map, vis_alpha = vis\n selected = self.curr_path[k]\n selected.append((x, y))\n if len(selected) >= 2:\n cv2.line(self.drawn_map, \n (int(round(selected[-2][0])), int(round(selected[-2][1]))),\n (int(round(selected[-1][0])), int(round(selected[-1][1]))),\n k, thickness=self.size)\n\n # Plot visualization\n if vis is not None:\n # Visualization for drawing\n if k == 0:\n vis_map = cv2.line(vis_map, \n (int(round(selected[-2][0])), int(round(selected[-2][1]))),\n (int(round(selected[-1][0])), int(round(selected[-1][1]))),\n color_map[k], thickness=self.size)\n else:\n vis_map = cv2.line(vis_map, \n (int(round(selected[-2][0])), int(round(selected[-2][1]))),\n (int(round(selected[-1][0])), int(round(selected[-1][1]))),\n color_map[k], thickness=self.size)\n # Visualization on/off boolean filter\n vis_alpha = cv2.line(vis_alpha, \n (int(round(selected[-2][0])), int(round(selected[-2][1]))),\n (int(round(selected[-1][0])), int(round(selected[-1][1]))),\n 0.75, thickness=self.size)\n\n if vis is not None:\n return vis_map, vis_alpha\n\n def end_path(self):\n # Complete the drawing\n self.curr_path = [[] for _ in range(self.K + 1)]\n\n def predict(self):\n self.out_prob = index_numpy_to_one_hot_torch(self.drawn_map, self.K+1)\n # self.out_prob = torch.from_numpy(self.drawn_map).float().cuda()\n # self.out_prob, _ = pad_divide_by(self.out_prob, 16, self.out_prob.shape[-2:])\n # self.out_prob = aggregate_sbg(self.out_prob, keep_bg=True)\n return self.out_prob\n\nclass ScribbleInteraction(Interaction):\n def __init__(self, image, prev_mask, true_size, controller, num_objects):\n \"\"\"\n prev_mask should be in an indexed form\n \"\"\"\n super().__init__(image, prev_mask, true_size, controller)\n\n self.K = num_objects\n\n self.drawn_map = np.empty((self.h, self.w), dtype=np.uint8)\n self.drawn_map.fill(255)\n # background + k\n self.curr_path = [[] for _ in range(self.K + 1)]\n self.size = 3\n\n \"\"\"\n k - object id\n vis - a tuple (visualization map, pass through alpha). None if not needed.\n \"\"\"\n def push_point(self, x, y, k, vis=None):\n if vis is not None:\n vis_map, vis_alpha = vis\n selected = self.curr_path[k]\n selected.append((x, y))\n if len(selected) >= 2:\n self.drawn_map = cv2.line(self.drawn_map, \n (int(round(selected[-2][0])), int(round(selected[-2][1]))),\n (int(round(selected[-1][0])), int(round(selected[-1][1]))),\n k, thickness=self.size)\n\n # Plot visualization\n if vis is not None:\n # Visualization for drawing\n if k == 0:\n vis_map = cv2.line(vis_map, \n (int(round(selected[-2][0])), int(round(selected[-2][1]))),\n (int(round(selected[-1][0])), int(round(selected[-1][1]))),\n color_map[k], thickness=self.size)\n else:\n vis_map = cv2.line(vis_map, \n (int(round(selected[-2][0])), int(round(selected[-2][1]))),\n (int(round(selected[-1][0])), int(round(selected[-1][1]))),\n color_map[k], thickness=self.size)\n # Visualization on/off boolean filter\n vis_alpha = cv2.line(vis_alpha, \n (int(round(selected[-2][0])), int(round(selected[-2][1]))),\n (int(round(selected[-1][0])), int(round(selected[-1][1]))),\n 0.75, thickness=self.size)\n\n # Optional vis return\n if vis is not None:\n return vis_map, vis_alpha\n\n def end_path(self):\n # Complete the drawing\n self.curr_path = [[] for _ in range(self.K + 1)]\n\n def predict(self):\n self.out_prob = self.controller.interact(self.image.unsqueeze(0), self.prev_mask, self.drawn_map)\n self.out_prob = aggregate_wbg(self.out_prob, keep_bg=True, hard=True)\n return self.out_prob\n\n\nclass ClickInteraction(Interaction):\n def __init__(self, image, prev_mask, true_size, controller, tar_obj):\n \"\"\"\n prev_mask in a prob. form\n \"\"\"\n super().__init__(image, prev_mask, true_size, controller)\n self.tar_obj = tar_obj\n\n # negative/positive for each object\n self.pos_clicks = []\n self.neg_clicks = []\n\n self.out_prob = self.prev_mask.clone()\n\n \"\"\"\n neg - Negative interaction or not\n vis - a tuple (visualization map, pass through alpha). None if not needed.\n \"\"\"\n def push_point(self, x, y, neg, vis=None):\n # Clicks\n if neg:\n self.neg_clicks.append((x, y))\n else:\n self.pos_clicks.append((x, y))\n\n # Do the prediction\n self.obj_mask = self.controller.interact(self.image.unsqueeze(0), x, y, not neg)\n\n # Plot visualization\n if vis is not None:\n vis_map, vis_alpha = vis\n # Visualization for clicks\n if neg:\n vis_map = cv2.circle(vis_map, \n (int(round(x)), int(round(y))),\n 2, color_map[0], thickness=-1)\n else:\n vis_map = cv2.circle(vis_map, \n (int(round(x)), int(round(y))),\n 2, color_map[self.tar_obj], thickness=-1)\n\n vis_alpha = cv2.circle(vis_alpha, \n (int(round(x)), int(round(y))),\n 2, 1, thickness=-1)\n\n # Optional vis return\n return vis_map, vis_alpha\n\n def predict(self):\n self.out_prob = self.prev_mask.clone()\n # a small hack to allow the interacting object to overwrite existing masks\n # without remembering all the object probabilities\n self.out_prob = torch.clamp(self.out_prob, max=0.9)\n self.out_prob[self.tar_obj] = self.obj_mask\n self.out_prob = aggregate_wbg(self.out_prob[1:], keep_bg=True, hard=True)\n return self.out_prob\n" }, { "path": "XMem/inference/interact/interactive_utils.py", "content": "# Modifed from https://github.com/seoungwugoh/ivs-demo\n\nimport numpy as np\n\nimport torch\nimport torch.nn.functional as F\nfrom util.palette import davis_palette\nfrom dataset.range_transform import im_normalization\n\ndef image_to_torch(frame: np.ndarray, device='cuda'):\n # frame: H*W*3 numpy array\n frame = frame.transpose(2, 0, 1)\n frame = torch.from_numpy(frame).float().to(device)/255\n frame_norm = im_normalization(frame)\n return frame_norm, frame\n\ndef torch_prob_to_numpy_mask(prob):\n mask = torch.max(prob, dim=0).indices\n mask = mask.cpu().numpy().astype(np.uint8)\n return mask\n\ndef index_numpy_to_one_hot_torch(mask, num_classes):\n mask = torch.from_numpy(mask).long()\n return F.one_hot(mask, num_classes=num_classes).permute(2, 0, 1).float()\n\n\"\"\"\nSome constants fro visualization\n\"\"\"\ntry:\n if torch.cuda.is_available():\n device = torch.device(\"cuda\")\n elif torch.backends.mps.is_available():\n device = torch.device(\"mps\")\n else:\n device = torch.device(\"cpu\")\nexcept:\n device = torch.device(\"cpu\")\n\ncolor_map_np = np.frombuffer(davis_palette, dtype=np.uint8).reshape(-1, 3).copy()\n# scales for better visualization\ncolor_map_np = (color_map_np.astype(np.float32)*1.5).clip(0, 255).astype(np.uint8)\ncolor_map = color_map_np.tolist()\ncolor_map_torch = torch.from_numpy(color_map_np).to(device) / 255\n\ngrayscale_weights = np.array([[0.3,0.59,0.11]]).astype(np.float32)\ngrayscale_weights_torch = torch.from_numpy(grayscale_weights).to(device).unsqueeze(0)\n\ndef get_visualization(mode, image, mask, layer, target_object):\n if mode == 'fade':\n return overlay_davis(image, mask, fade=True)\n elif mode == 'davis':\n return overlay_davis(image, mask)\n elif mode == 'light':\n return overlay_davis(image, mask, 0.9)\n elif mode == 'popup':\n return overlay_popup(image, mask, target_object)\n elif mode == 'layered':\n if layer is None:\n print('Layer file not given. Defaulting to DAVIS.')\n return overlay_davis(image, mask)\n else:\n return overlay_layer(image, mask, layer, target_object)\n else:\n raise NotImplementedError\n\ndef get_visualization_torch(mode, image, prob, layer, target_object):\n if mode == 'fade':\n return overlay_davis_torch(image, prob, fade=True)\n elif mode == 'davis':\n return overlay_davis_torch(image, prob)\n elif mode == 'light':\n return overlay_davis_torch(image, prob, 0.9)\n elif mode == 'popup':\n return overlay_popup_torch(image, prob, target_object)\n elif mode == 'layered':\n if layer is None:\n print('Layer file not given. Defaulting to DAVIS.')\n return overlay_davis_torch(image, prob)\n else:\n return overlay_layer_torch(image, prob, layer, target_object)\n else:\n raise NotImplementedError\n\ndef overlay_davis(image, mask, alpha=0.5, fade=False):\n \"\"\" Overlay segmentation on top of RGB image. from davis official\"\"\"\n im_overlay = image.copy()\n\n colored_mask = color_map_np[mask]\n foreground = image*alpha + (1-alpha)*colored_mask\n binary_mask = (mask > 0)\n # Compose image\n im_overlay[binary_mask] = foreground[binary_mask]\n if fade:\n im_overlay[~binary_mask] = im_overlay[~binary_mask] * 0.6\n return im_overlay.astype(image.dtype)\n\ndef overlay_popup(image, mask, target_object):\n # Keep foreground colored. Convert background to grayscale.\n im_overlay = image.copy()\n\n binary_mask = ~(np.isin(mask, target_object))\n colored_region = (im_overlay[binary_mask]*grayscale_weights).sum(-1, keepdims=-1)\n im_overlay[binary_mask] = colored_region\n return im_overlay.astype(image.dtype)\n\ndef overlay_layer(image, mask, layer, target_object):\n # insert a layer between foreground and background\n # The CPU version is less accurate because we are using the hard mask\n # The GPU version has softer edges as it uses soft probabilities\n obj_mask = (np.isin(mask, target_object)).astype(np.float32)[:, :, np.newaxis]\n layer_alpha = layer[:, :, 3].astype(np.float32)[:, :, np.newaxis] / 255\n layer_rgb = layer[:, :, :3]\n background_alpha = np.maximum(obj_mask, layer_alpha)\n im_overlay = (image * (1 - background_alpha) + layer_rgb * (1 - obj_mask) * layer_alpha +\n image * obj_mask).clip(0, 255)\n return im_overlay.astype(image.dtype)\n\ndef overlay_davis_torch(image, mask, alpha=0.5, fade=False):\n \"\"\" Overlay segmentation on top of RGB image. from davis official\"\"\"\n # Changes the image in-place to avoid copying\n image = image.permute(1, 2, 0)\n im_overlay = image\n mask = torch.max(mask, dim=0).indices\n\n colored_mask = color_map_torch[mask]\n foreground = image*alpha + (1-alpha)*colored_mask\n binary_mask = (mask > 0)\n # Compose image\n im_overlay[binary_mask] = foreground[binary_mask]\n if fade:\n im_overlay[~binary_mask] = im_overlay[~binary_mask] * 0.6\n\n im_overlay = (im_overlay*255).cpu().numpy()\n im_overlay = im_overlay.astype(np.uint8)\n\n return im_overlay\n\ndef overlay_popup_torch(image, mask, target_object):\n # Keep foreground colored. Convert background to grayscale.\n image = image.permute(1, 2, 0)\n \n if len(target_object) == 0:\n obj_mask = torch.zeros_like(mask[0]).unsqueeze(2)\n else:\n # I should not need to convert this to numpy.\n # uUsing list works most of the time but consistently fails\n # if I include first object -> exclude it -> include it again.\n # I check everywhere and it makes absolutely no sense.\n # I am blaming this on PyTorch and calling it a day\n obj_mask = mask[np.array(target_object,dtype=np.int32)].sum(0).unsqueeze(2)\n gray_image = (image*grayscale_weights_torch).sum(-1, keepdim=True)\n im_overlay = obj_mask*image + (1-obj_mask)*gray_image\n\n im_overlay = (im_overlay*255).cpu().numpy()\n im_overlay = im_overlay.astype(np.uint8)\n\n return im_overlay\n\ndef overlay_layer_torch(image, prob, layer, target_object):\n # insert a layer between foreground and background\n # The CPU version is less accurate because we are using the hard mask\n # The GPU version has softer edges as it uses soft probabilities\n image = image.permute(1, 2, 0)\n\n if len(target_object) == 0:\n obj_mask = torch.zeros_like(prob[0]).unsqueeze(2)\n else:\n # TODO: figure out why we need to convert this to numpy array\n obj_mask = prob[np.array(target_object, dtype=np.int32)].sum(0).unsqueeze(2)\n layer_alpha = layer[:, :, 3].unsqueeze(2)\n layer_rgb = layer[:, :, :3]\n background_alpha = torch.maximum(obj_mask, layer_alpha)\n im_overlay = (image * (1 - background_alpha) + layer_rgb * (1 - obj_mask) * layer_alpha +\n image * obj_mask).clip(0, 1)\n\n im_overlay = (im_overlay * 255).cpu().numpy()\n im_overlay = im_overlay.astype(np.uint8)\n\n return im_overlay\n\n" }, { "path": "XMem/inference/interact/resource_manager.py", "content": "import os\nfrom os import path\nimport shutil\nimport collections\n\nimport cv2\nfrom PIL import Image\nif not hasattr(Image, 'Resampling'): # Pillow<9.0\n Image.Resampling = Image\nimport numpy as np\n\nfrom util.palette import davis_palette\nimport progressbar\n \n\n# https://bugs.python.org/issue28178\n# ah python ah why\nclass LRU:\n def __init__(self, func, maxsize=128):\n self.cache = collections.OrderedDict()\n self.func = func\n self.maxsize = maxsize\n \n def __call__(self, *args):\n cache = self.cache\n if args in cache:\n cache.move_to_end(args)\n return cache[args]\n result = self.func(*args)\n cache[args] = result\n if len(cache) > self.maxsize:\n cache.popitem(last=False)\n return result\n\n def invalidate(self, key):\n self.cache.pop(key, None)\n\n\nclass ResourceManager:\n def __init__(self, config):\n # determine inputs\n images = config['images']\n video = config['video']\n self.workspace = config['workspace']\n self.size = config['size']\n self.palette = davis_palette\n\n # create temporary workspace if not specified\n if self.workspace is None:\n if images is not None:\n basename = path.basename(images)\n elif video is not None:\n basename = path.basename(video)[:-4]\n else:\n raise NotImplementedError(\n 'Either images, video, or workspace has to be specified')\n\n self.workspace = path.join('./workspace', basename)\n\n print(f'Workspace is in: {self.workspace}')\n\n # determine the location of input images\n need_decoding = False\n need_resizing = False\n if path.exists(path.join(self.workspace, 'images')):\n pass\n elif images is not None:\n need_resizing = True\n elif video is not None:\n # will decode video into frames later\n need_decoding = True\n\n # create workspace subdirectories\n self.image_dir = path.join(self.workspace, 'images')\n self.mask_dir = path.join(self.workspace, 'masks')\n os.makedirs(self.image_dir, exist_ok=True)\n os.makedirs(self.mask_dir, exist_ok=True)\n\n # convert read functions to be buffered\n self.get_image = LRU(self._get_image_unbuffered, maxsize=config['buffer_size'])\n self.get_mask = LRU(self._get_mask_unbuffered, maxsize=config['buffer_size'])\n\n # extract frames from video\n if need_decoding:\n self._extract_frames(video)\n\n # copy/resize existing images to the workspace\n if need_resizing:\n self._copy_resize_frames(images)\n\n # read all frame names\n self.names = sorted(os.listdir(self.image_dir))\n self.names = [f[:-4] for f in self.names] # remove extensions\n self.length = len(self.names)\n\n assert self.length > 0, f'No images found! Check {self.workspace}/images. Remove folder if necessary.'\n\n print(f'{self.length} images found.')\n\n self.height, self.width = self.get_image(0).shape[:2]\n self.visualization_init = False\n\n def _extract_frames(self, video):\n cap = cv2.VideoCapture(video)\n frame_index = 0\n print(f'Extracting frames from {video} into {self.image_dir}...')\n bar = progressbar.ProgressBar(max_value=progressbar.UnknownLength)\n while(cap.isOpened()):\n _, frame = cap.read()\n if frame is None:\n break\n if self.size > 0:\n h, w = frame.shape[:2]\n new_w = (w*self.size//min(w, h))\n new_h = (h*self.size//min(w, h))\n if new_w != w or new_h != h:\n frame = cv2.resize(frame,dsize=(new_w,new_h),interpolation=cv2.INTER_AREA)\n cv2.imwrite(path.join(self.image_dir, f'{frame_index:07d}.jpg'), frame)\n frame_index += 1\n bar.update(frame_index)\n bar.finish()\n print('Done!')\n\n def _copy_resize_frames(self, images):\n image_list = os.listdir(images)\n print(f'Copying/resizing frames into {self.image_dir}...')\n for image_name in progressbar.progressbar(image_list):\n if self.size < 0:\n # just copy\n shutil.copy2(path.join(images, image_name), self.image_dir)\n else:\n frame = cv2.imread(path.join(images, image_name))\n h, w = frame.shape[:2]\n new_w = (w*self.size//min(w, h))\n new_h = (h*self.size//min(w, h))\n if new_w != w or new_h != h:\n frame = cv2.resize(frame,dsize=(new_w,new_h),interpolation=cv2.INTER_AREA)\n cv2.imwrite(path.join(self.image_dir, image_name), frame)\n print('Done!')\n\n def save_mask(self, ti, mask):\n # mask should be uint8 H*W without channels\n assert 0 <= ti < self.length\n assert isinstance(mask, np.ndarray)\n\n mask = Image.fromarray(mask)\n mask.putpalette(self.palette)\n mask.save(path.join(self.mask_dir, self.names[ti]+'.png'))\n self.invalidate(ti)\n\n def save_visualization(self, ti, image):\n # image should be uint8 3*H*W\n assert 0 <= ti < self.length\n assert isinstance(image, np.ndarray)\n if not self.visualization_init:\n self.visualization_dir = path.join(self.workspace, 'visualization')\n os.makedirs(self.visualization_dir, exist_ok=True)\n self.visualization_init = True\n\n image = Image.fromarray(image)\n image.save(path.join(self.visualization_dir, self.names[ti]+'.jpg'))\n\n def _get_image_unbuffered(self, ti):\n # returns H*W*3 uint8 array\n assert 0 <= ti < self.length\n\n image = Image.open(path.join(self.image_dir, self.names[ti]+'.jpg'))\n image = np.array(image)\n return image\n\n def _get_mask_unbuffered(self, ti):\n # returns H*W uint8 array\n assert 0 <= ti < self.length\n\n mask_path = path.join(self.mask_dir, self.names[ti]+'.png')\n if path.exists(mask_path):\n mask = Image.open(mask_path)\n mask = np.array(mask)\n return mask\n else:\n return None\n\n def read_external_image(self, file_name, size=None):\n image = Image.open(file_name)\n is_mask = image.mode in ['L', 'P']\n if size is not None:\n # PIL uses (width, height)\n image = image.resize((size[1], size[0]), \n resample=Image.Resampling.NEAREST if is_mask else Image.Resampling.BICUBIC)\n image = np.array(image)\n return image\n\n def invalidate(self, ti):\n # the image buffer is never invalidated\n self.get_mask.invalidate((ti,))\n\n def __len__(self):\n return self.length\n\n @property\n def h(self):\n return self.height\n\n @property\n def w(self):\n return self.width\n" }, { "path": "XMem/inference/interact/s2m/__init__.py", "content": "" }, { "path": "XMem/inference/interact/s2m/_deeplab.py", "content": "# Credit: https://github.com/VainF/DeepLabV3Plus-Pytorch\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\nfrom .utils import _SimpleSegmentationModel\n\n\n__all__ = [\"DeepLabV3\"]\n\n\nclass DeepLabV3(_SimpleSegmentationModel):\n \"\"\"\n Implements DeepLabV3 model from\n `\"Rethinking Atrous Convolution for Semantic Image Segmentation\"\n `_.\n\n Arguments:\n backbone (nn.Module): the network used to compute the features for the model.\n The backbone should return an OrderedDict[Tensor], with the key being\n \"out\" for the last feature map used, and \"aux\" if an auxiliary classifier\n is used.\n classifier (nn.Module): module that takes the \"out\" element returned from\n the backbone and returns a dense prediction.\n aux_classifier (nn.Module, optional): auxiliary classifier used during training\n \"\"\"\n pass\n\nclass DeepLabHeadV3Plus(nn.Module):\n def __init__(self, in_channels, low_level_channels, num_classes, aspp_dilate=[12, 24, 36]):\n super(DeepLabHeadV3Plus, self).__init__()\n self.project = nn.Sequential( \n nn.Conv2d(low_level_channels, 48, 1, bias=False),\n nn.BatchNorm2d(48),\n nn.ReLU(inplace=True),\n )\n\n self.aspp = ASPP(in_channels, aspp_dilate)\n\n self.classifier = nn.Sequential(\n nn.Conv2d(304, 256, 3, padding=1, bias=False),\n nn.BatchNorm2d(256),\n nn.ReLU(inplace=True),\n nn.Conv2d(256, num_classes, 1)\n )\n self._init_weight()\n\n def forward(self, feature):\n low_level_feature = self.project( feature['low_level'] )\n output_feature = self.aspp(feature['out'])\n output_feature = F.interpolate(output_feature, size=low_level_feature.shape[2:], mode='bilinear', align_corners=False)\n return self.classifier( torch.cat( [ low_level_feature, output_feature ], dim=1 ) )\n \n def _init_weight(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(m.weight)\n elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\nclass DeepLabHead(nn.Module):\n def __init__(self, in_channels, num_classes, aspp_dilate=[12, 24, 36]):\n super(DeepLabHead, self).__init__()\n\n self.classifier = nn.Sequential(\n ASPP(in_channels, aspp_dilate),\n nn.Conv2d(256, 256, 3, padding=1, bias=False),\n nn.BatchNorm2d(256),\n nn.ReLU(inplace=True),\n nn.Conv2d(256, num_classes, 1)\n )\n self._init_weight()\n\n def forward(self, feature):\n return self.classifier( feature['out'] )\n\n def _init_weight(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(m.weight)\n elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\nclass AtrousSeparableConvolution(nn.Module):\n \"\"\" Atrous Separable Convolution\n \"\"\"\n def __init__(self, in_channels, out_channels, kernel_size,\n stride=1, padding=0, dilation=1, bias=True):\n super(AtrousSeparableConvolution, self).__init__()\n self.body = nn.Sequential(\n # Separable Conv\n nn.Conv2d( in_channels, in_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias, groups=in_channels ),\n # PointWise Conv\n nn.Conv2d( in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=bias),\n )\n \n self._init_weight()\n\n def forward(self, x):\n return self.body(x)\n\n def _init_weight(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(m.weight)\n elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\nclass ASPPConv(nn.Sequential):\n def __init__(self, in_channels, out_channels, dilation):\n modules = [\n nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False),\n nn.BatchNorm2d(out_channels),\n nn.ReLU(inplace=True)\n ]\n super(ASPPConv, self).__init__(*modules)\n\nclass ASPPPooling(nn.Sequential):\n def __init__(self, in_channels, out_channels):\n super(ASPPPooling, self).__init__(\n nn.AdaptiveAvgPool2d(1),\n nn.Conv2d(in_channels, out_channels, 1, bias=False),\n nn.BatchNorm2d(out_channels),\n nn.ReLU(inplace=True))\n\n def forward(self, x):\n size = x.shape[-2:]\n x = super(ASPPPooling, self).forward(x)\n return F.interpolate(x, size=size, mode='bilinear', align_corners=False)\n\nclass ASPP(nn.Module):\n def __init__(self, in_channels, atrous_rates):\n super(ASPP, self).__init__()\n out_channels = 256\n modules = []\n modules.append(nn.Sequential(\n nn.Conv2d(in_channels, out_channels, 1, bias=False),\n nn.BatchNorm2d(out_channels),\n nn.ReLU(inplace=True)))\n\n rate1, rate2, rate3 = tuple(atrous_rates)\n modules.append(ASPPConv(in_channels, out_channels, rate1))\n modules.append(ASPPConv(in_channels, out_channels, rate2))\n modules.append(ASPPConv(in_channels, out_channels, rate3))\n modules.append(ASPPPooling(in_channels, out_channels))\n\n self.convs = nn.ModuleList(modules)\n\n self.project = nn.Sequential(\n nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),\n nn.BatchNorm2d(out_channels),\n nn.ReLU(inplace=True),\n nn.Dropout(0.1),)\n\n def forward(self, x):\n res = []\n for conv in self.convs:\n res.append(conv(x))\n res = torch.cat(res, dim=1)\n return self.project(res)\n\n\n\ndef convert_to_separable_conv(module):\n new_module = module\n if isinstance(module, nn.Conv2d) and module.kernel_size[0]>1:\n new_module = AtrousSeparableConvolution(module.in_channels,\n module.out_channels, \n module.kernel_size,\n module.stride,\n module.padding,\n module.dilation,\n module.bias)\n for name, child in module.named_children():\n new_module.add_module(name, convert_to_separable_conv(child))\n return new_module" }, { "path": "XMem/inference/interact/s2m/s2m_network.py", "content": "# Credit: https://github.com/VainF/DeepLabV3Plus-Pytorch\n\nfrom .utils import IntermediateLayerGetter\nfrom ._deeplab import DeepLabHead, DeepLabHeadV3Plus, DeepLabV3\nfrom . import s2m_resnet\n\ndef _segm_resnet(name, backbone_name, num_classes, output_stride, pretrained_backbone):\n\n if output_stride==8:\n replace_stride_with_dilation=[False, True, True]\n aspp_dilate = [12, 24, 36]\n else:\n replace_stride_with_dilation=[False, False, True]\n aspp_dilate = [6, 12, 18]\n\n backbone = s2m_resnet.__dict__[backbone_name](\n pretrained=pretrained_backbone,\n replace_stride_with_dilation=replace_stride_with_dilation)\n \n inplanes = 2048\n low_level_planes = 256\n\n if name=='deeplabv3plus':\n return_layers = {'layer4': 'out', 'layer1': 'low_level'}\n classifier = DeepLabHeadV3Plus(inplanes, low_level_planes, num_classes, aspp_dilate)\n elif name=='deeplabv3':\n return_layers = {'layer4': 'out'}\n classifier = DeepLabHead(inplanes , num_classes, aspp_dilate)\n backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)\n\n model = DeepLabV3(backbone, classifier)\n return model\n\ndef _load_model(arch_type, backbone, num_classes, output_stride, pretrained_backbone):\n\n if backbone.startswith('resnet'):\n model = _segm_resnet(arch_type, backbone, num_classes, output_stride=output_stride, pretrained_backbone=pretrained_backbone)\n else:\n raise NotImplementedError\n return model\n\n\n# Deeplab v3\ndef deeplabv3_resnet50(num_classes=1, output_stride=16, pretrained_backbone=False):\n \"\"\"Constructs a DeepLabV3 model with a ResNet-50 backbone.\n\n Args:\n num_classes (int): number of classes.\n output_stride (int): output stride for deeplab.\n pretrained_backbone (bool): If True, use the pretrained backbone.\n \"\"\"\n return _load_model('deeplabv3', 'resnet50', num_classes, output_stride=output_stride, pretrained_backbone=pretrained_backbone)\n\n\n# Deeplab v3+\ndef deeplabv3plus_resnet50(num_classes=1, output_stride=16, pretrained_backbone=False):\n \"\"\"Constructs a DeepLabV3 model with a ResNet-50 backbone.\n\n Args:\n num_classes (int): number of classes.\n output_stride (int): output stride for deeplab.\n pretrained_backbone (bool): If True, use the pretrained backbone.\n \"\"\"\n return _load_model('deeplabv3plus', 'resnet50', num_classes, output_stride=output_stride, pretrained_backbone=pretrained_backbone)\n\n" }, { "path": "XMem/inference/interact/s2m/s2m_resnet.py", "content": "import torch\nimport torch.nn as nn\ntry:\n from torchvision.models.utils import load_state_dict_from_url\nexcept ModuleNotFoundError:\n from torch.utils.model_zoo import load_url as load_state_dict_from_url\n\n\n__all__ = ['ResNet', 'resnet50']\n\n\nmodel_urls = {\n 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n}\n\n\ndef conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):\n \"\"\"3x3 convolution with padding\"\"\"\n return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n padding=dilation, groups=groups, bias=False, dilation=dilation)\n\n\ndef conv1x1(in_planes, out_planes, stride=1):\n \"\"\"1x1 convolution\"\"\"\n return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n\n\nclass Bottleneck(nn.Module):\n expansion = 4\n\n def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,\n base_width=64, dilation=1, norm_layer=None):\n super(Bottleneck, self).__init__()\n if norm_layer is None:\n norm_layer = nn.BatchNorm2d\n width = int(planes * (base_width / 64.)) * groups\n # Both self.conv2 and self.downsample layers downsample the input when stride != 1\n self.conv1 = conv1x1(inplanes, width)\n self.bn1 = norm_layer(width)\n self.conv2 = conv3x3(width, width, stride, groups, dilation)\n self.bn2 = norm_layer(width)\n self.conv3 = conv1x1(width, planes * self.expansion)\n self.bn3 = norm_layer(planes * self.expansion)\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n identity = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n out = self.relu(out)\n\n return out\n\n\nclass ResNet(nn.Module):\n\n def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,\n groups=1, width_per_group=64, replace_stride_with_dilation=None,\n norm_layer=None):\n super(ResNet, self).__init__()\n if norm_layer is None:\n norm_layer = nn.BatchNorm2d\n self._norm_layer = norm_layer\n\n self.inplanes = 64\n self.dilation = 1\n if replace_stride_with_dilation is None:\n # each element in the tuple indicates if we should replace\n # the 2x2 stride with a dilated convolution instead\n replace_stride_with_dilation = [False, False, False]\n if len(replace_stride_with_dilation) != 3:\n raise ValueError(\"replace_stride_with_dilation should be None \"\n \"or a 3-element tuple, got {}\".format(replace_stride_with_dilation))\n self.groups = groups\n self.base_width = width_per_group\n self.conv1 = nn.Conv2d(6, self.inplanes, kernel_size=7, stride=2, padding=3,\n bias=False)\n self.bn1 = norm_layer(self.inplanes)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n self.layer1 = self._make_layer(block, 64, layers[0])\n self.layer2 = self._make_layer(block, 128, layers[1], stride=2,\n dilate=replace_stride_with_dilation[0])\n self.layer3 = self._make_layer(block, 256, layers[2], stride=2,\n dilate=replace_stride_with_dilation[1])\n self.layer4 = self._make_layer(block, 512, layers[3], stride=2,\n dilate=replace_stride_with_dilation[2])\n self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n self.fc = nn.Linear(512 * block.expansion, num_classes)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n # Zero-initialize the last BN in each residual branch,\n # so that the residual branch starts with zeros, and each residual block behaves like an identity.\n # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677\n if zero_init_residual:\n for m in self.modules():\n if isinstance(m, Bottleneck):\n nn.init.constant_(m.bn3.weight, 0)\n\n def _make_layer(self, block, planes, blocks, stride=1, dilate=False):\n norm_layer = self._norm_layer\n downsample = None\n previous_dilation = self.dilation\n if dilate:\n self.dilation *= stride\n stride = 1\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n conv1x1(self.inplanes, planes * block.expansion, stride),\n norm_layer(planes * block.expansion),\n )\n\n layers = []\n layers.append(block(self.inplanes, planes, stride, downsample, self.groups,\n self.base_width, previous_dilation, norm_layer))\n self.inplanes = planes * block.expansion\n for _ in range(1, blocks):\n layers.append(block(self.inplanes, planes, groups=self.groups,\n base_width=self.base_width, dilation=self.dilation,\n norm_layer=norm_layer))\n\n return nn.Sequential(*layers)\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.layer4(x)\n\n x = self.avgpool(x)\n x = torch.flatten(x, 1)\n x = self.fc(x)\n\n return x\n\n\ndef _resnet(arch, block, layers, pretrained, progress, **kwargs):\n model = ResNet(block, layers, **kwargs)\n if pretrained:\n state_dict = load_state_dict_from_url(model_urls[arch],\n progress=progress)\n model.load_state_dict(state_dict)\n return model\n\n\ndef resnet50(pretrained=False, progress=True, **kwargs):\n r\"\"\"ResNet-50 model from\n `\"Deep Residual Learning for Image Recognition\" `_\n\n Args:\n pretrained (bool): If True, returns a model pre-trained on ImageNet\n progress (bool): If True, displays a progress bar of the download to stderr\n \"\"\"\n return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,\n **kwargs)\n" }, { "path": "XMem/inference/interact/s2m/utils.py", "content": "# Credit: https://github.com/VainF/DeepLabV3Plus-Pytorch\n\nimport torch\nimport torch.nn as nn\nimport numpy as np\nimport torch.nn.functional as F\nfrom collections import OrderedDict\n\nclass _SimpleSegmentationModel(nn.Module):\n def __init__(self, backbone, classifier):\n super(_SimpleSegmentationModel, self).__init__()\n self.backbone = backbone\n self.classifier = classifier\n \n def forward(self, x):\n input_shape = x.shape[-2:]\n features = self.backbone(x)\n x = self.classifier(features)\n x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False)\n return x\n\n\nclass IntermediateLayerGetter(nn.ModuleDict):\n \"\"\"\n Module wrapper that returns intermediate layers from a model\n\n It has a strong assumption that the modules have been registered\n into the model in the same order as they are used.\n This means that one should **not** reuse the same nn.Module\n twice in the forward if you want this to work.\n\n Additionally, it is only able to query submodules that are directly\n assigned to the model. So if `model` is passed, `model.feature1` can\n be returned, but not `model.feature1.layer2`.\n\n Arguments:\n model (nn.Module): model on which we will extract the features\n return_layers (Dict[name, new_name]): a dict containing the names\n of the modules for which the activations will be returned as\n the key of the dict, and the value of the dict is the name\n of the returned activation (which the user can specify).\n\n Examples::\n\n >>> m = torchvision.models.resnet18(pretrained=True)\n >>> # extract layer1 and layer3, giving as names `feat1` and feat2`\n >>> new_m = torchvision.models._utils.IntermediateLayerGetter(m,\n >>> {'layer1': 'feat1', 'layer3': 'feat2'})\n >>> out = new_m(torch.rand(1, 3, 224, 224))\n >>> print([(k, v.shape) for k, v in out.items()])\n >>> [('feat1', torch.Size([1, 64, 56, 56])),\n >>> ('feat2', torch.Size([1, 256, 14, 14]))]\n \"\"\"\n def __init__(self, model, return_layers):\n if not set(return_layers).issubset([name for name, _ in model.named_children()]):\n raise ValueError(\"return_layers are not present in model\")\n\n orig_return_layers = return_layers\n return_layers = {k: v for k, v in return_layers.items()}\n layers = OrderedDict()\n for name, module in model.named_children():\n layers[name] = module\n if name in return_layers:\n del return_layers[name]\n if not return_layers:\n break\n\n super(IntermediateLayerGetter, self).__init__(layers)\n self.return_layers = orig_return_layers\n\n def forward(self, x):\n out = OrderedDict()\n for name, module in self.named_children():\n x = module(x)\n if name in self.return_layers:\n out_name = self.return_layers[name]\n out[out_name] = x\n return out\n" }, { "path": "XMem/inference/interact/s2m_controller.py", "content": "import torch\nimport numpy as np\nfrom ..interact.s2m.s2m_network import deeplabv3plus_resnet50 as S2M\n\nfrom util.tensor_util import pad_divide_by, unpad\n\n\nclass S2MController:\n \"\"\"\n A controller for Scribble-to-Mask (for user interaction, not for DAVIS)\n Takes the image, previous mask, and scribbles to produce a new mask\n ignore_class is usually 255 \n 0 is NOT the ignore class -- it is the label for the background\n \"\"\"\n def __init__(self, s2m_net:S2M, num_objects, ignore_class, device='cuda:0'):\n self.s2m_net = s2m_net\n self.num_objects = num_objects\n self.ignore_class = ignore_class\n self.device = device\n\n def interact(self, image, prev_mask, scr_mask):\n print(self.device)\n image = image.to(self.device, non_blocking=True)\n prev_mask = prev_mask.unsqueeze(0)\n\n h, w = image.shape[-2:]\n unaggre_mask = torch.zeros((self.num_objects, h, w), dtype=torch.float32, device=image.device)\n\n for ki in range(1, self.num_objects+1):\n p_srb = (scr_mask==ki).astype(np.uint8)\n n_srb = ((scr_mask!=ki) * (scr_mask!=self.ignore_class)).astype(np.uint8)\n\n Rs = torch.from_numpy(np.stack([p_srb, n_srb], 0)).unsqueeze(0).float().to(image.device)\n\n inputs = torch.cat([image, (prev_mask==ki).float().unsqueeze(0), Rs], 1)\n inputs, pads = pad_divide_by(inputs, 16)\n\n unaggre_mask[ki-1] = unpad(torch.sigmoid(self.s2m_net(inputs)), pads)\n\n return unaggre_mask" }, { "path": "XMem/inference/interact/timer.py", "content": "import time\n\nclass Timer:\n def __init__(self):\n self._acc_time = 0\n self._paused = True\n\n def start(self):\n if self._paused:\n self.last_time = time.time()\n self._paused = False\n return self\n\n def pause(self):\n self.count()\n self._paused = True\n return self\n\n def count(self):\n if self._paused:\n return self._acc_time\n t = time.time()\n self._acc_time += t - self.last_time\n self.last_time = t\n return self._acc_time\n\n def format(self):\n # count = int(self.count()*100)\n # return '%02d:%02d:%02d' % (count//6000, (count//100)%60, count%100)\n return '%03.2f' % self.count()\n\n def __str__(self):\n return self.format()" }, { "path": "XMem/inference/kv_memory_store.py", "content": "import torch\nfrom typing import List\n\nclass KeyValueMemoryStore:\n \"\"\"\n Works for key/value pairs type storage\n e.g., working and long-term memory\n \"\"\"\n\n \"\"\"\n An object group is created when new objects enter the video\n Objects in the same group share the same temporal extent\n i.e., objects initialized in the same frame are in the same group\n For DAVIS/interactive, there is only one object group\n For YouTubeVOS, there can be multiple object groups\n \"\"\"\n\n def __init__(self, count_usage: bool):\n self.count_usage = count_usage\n\n # keys are stored in a single tensor and are shared between groups/objects\n # values are stored as a list indexed by object groups\n self.k = None\n self.v = []\n self.obj_groups = []\n # for debugging only\n self.all_objects = []\n\n # shrinkage and selection are also single tensors\n self.s = self.e = None\n\n # usage\n if self.count_usage:\n self.use_count = self.life_count = None\n\n def add(self, key, value, shrinkage, selection, objects: List[int]):\n new_count = torch.zeros((key.shape[0], 1, key.shape[2]), device=key.device, dtype=torch.float32)\n new_life = torch.zeros((key.shape[0], 1, key.shape[2]), device=key.device, dtype=torch.float32) + 1e-7\n\n # add the key\n if self.k is None:\n self.k = key\n self.s = shrinkage\n self.e = selection\n if self.count_usage:\n self.use_count = new_count\n self.life_count = new_life\n else:\n self.k = torch.cat([self.k, key], -1)\n if shrinkage is not None:\n self.s = torch.cat([self.s, shrinkage], -1)\n if selection is not None:\n self.e = torch.cat([self.e, selection], -1)\n if self.count_usage:\n self.use_count = torch.cat([self.use_count, new_count], -1)\n self.life_count = torch.cat([self.life_count, new_life], -1)\n\n # add the value\n if objects is not None:\n # When objects is given, v is a tensor; used in working memory\n assert isinstance(value, torch.Tensor)\n # First consume objects that are already in the memory bank\n # cannot use set here because we need to preserve order\n # shift by one as background is not part of value\n remaining_objects = [obj-1 for obj in objects]\n for gi, group in enumerate(self.obj_groups):\n for obj in group:\n # should properly raise an error if there are overlaps in obj_groups\n remaining_objects.remove(obj)\n self.v[gi] = torch.cat([self.v[gi], value[group]], -1)\n\n # If there are remaining objects, add them as a new group\n if len(remaining_objects) > 0:\n new_group = list(remaining_objects)\n self.v.append(value[new_group])\n self.obj_groups.append(new_group)\n self.all_objects.extend(new_group)\n \n assert sorted(self.all_objects) == self.all_objects, 'Objects MUST be inserted in sorted order '\n else:\n # When objects is not given, v is a list that already has the object groups sorted\n # used in long-term memory\n assert isinstance(value, list)\n for gi, gv in enumerate(value):\n if gv is None:\n continue\n if gi < self.num_groups:\n self.v[gi] = torch.cat([self.v[gi], gv], -1)\n else:\n self.v.append(gv)\n\n def update_usage(self, usage):\n # increase all life count by 1\n # increase use of indexed elements\n if not self.count_usage:\n return\n \n self.use_count += usage.view_as(self.use_count)\n self.life_count += 1\n\n def sieve_by_range(self, start: int, end: int, min_size: int):\n # keep only the elements *outside* of this range (with some boundary conditions)\n # i.e., concat (a[:start], a[end:])\n # min_size is only used for values, we do not sieve values under this size\n # (because they are not consolidated)\n\n if end == 0:\n # negative 0 would not work as the end index!\n self.k = self.k[:,:,:start]\n if self.count_usage:\n self.use_count = self.use_count[:,:,:start]\n self.life_count = self.life_count[:,:,:start]\n if self.s is not None:\n self.s = self.s[:,:,:start]\n if self.e is not None:\n self.e = self.e[:,:,:start]\n \n for gi in range(self.num_groups):\n if self.v[gi].shape[-1] >= min_size:\n self.v[gi] = self.v[gi][:,:,:start]\n else:\n self.k = torch.cat([self.k[:,:,:start], self.k[:,:,end:]], -1)\n if self.count_usage:\n self.use_count = torch.cat([self.use_count[:,:,:start], self.use_count[:,:,end:]], -1)\n self.life_count = torch.cat([self.life_count[:,:,:start], self.life_count[:,:,end:]], -1)\n if self.s is not None:\n self.s = torch.cat([self.s[:,:,:start], self.s[:,:,end:]], -1)\n if self.e is not None:\n self.e = torch.cat([self.e[:,:,:start], self.e[:,:,end:]], -1)\n \n for gi in range(self.num_groups):\n if self.v[gi].shape[-1] >= min_size:\n self.v[gi] = torch.cat([self.v[gi][:,:,:start], self.v[gi][:,:,end:]], -1)\n\n def remove_obsolete_features(self, max_size: int):\n # normalize with life duration\n usage = self.get_usage().flatten()\n\n values, _ = torch.topk(usage, k=(self.size-max_size), largest=False, sorted=True)\n survived = (usage > values[-1])\n\n self.k = self.k[:, :, survived]\n self.s = self.s[:, :, survived] if self.s is not None else None\n # Long-term memory does not store ek so this should not be needed\n self.e = self.e[:, :, survived] if self.e is not None else None\n if self.num_groups > 1:\n raise NotImplementedError(\"\"\"The current data structure does not support feature removal with \n multiple object groups (e.g., some objects start to appear later in the video)\n The indices for \"survived\" is based on keys but not all values are present for every key\n Basically we need to remap the indices for keys to values\n \"\"\")\n for gi in range(self.num_groups):\n self.v[gi] = self.v[gi][:, :, survived]\n\n self.use_count = self.use_count[:, :, survived]\n self.life_count = self.life_count[:, :, survived]\n\n def get_usage(self):\n # return normalized usage\n if not self.count_usage:\n raise RuntimeError('I did not count usage!')\n else:\n usage = self.use_count / self.life_count\n return usage\n\n def get_all_sliced(self, start: int, end: int):\n # return k, sk, ek, usage in order, sliced by start and end\n\n if end == 0:\n # negative 0 would not work as the end index!\n k = self.k[:,:,start:]\n sk = self.s[:,:,start:] if self.s is not None else None\n ek = self.e[:,:,start:] if self.e is not None else None\n usage = self.get_usage()[:,:,start:]\n else:\n k = self.k[:,:,start:end]\n sk = self.s[:,:,start:end] if self.s is not None else None\n ek = self.e[:,:,start:end] if self.e is not None else None\n usage = self.get_usage()[:,:,start:end]\n\n return k, sk, ek, usage\n\n def get_v_size(self, ni: int):\n return self.v[ni].shape[2]\n\n def engaged(self):\n return self.k is not None\n\n @property\n def size(self):\n if self.k is None:\n return 0\n else:\n return self.k.shape[-1]\n\n @property\n def num_groups(self):\n return len(self.v)\n\n @property\n def key(self):\n return self.k\n\n @property\n def value(self):\n return self.v\n\n @property\n def shrinkage(self):\n return self.s\n\n @property\n def selection(self):\n return self.e\n\n" }, { "path": "XMem/inference/memory_manager.py", "content": "import torch\nimport warnings\n\nfrom inference.kv_memory_store import KeyValueMemoryStore\nfrom model.memory_util import *\n\n\nclass MemoryManager:\n \"\"\"\n Manages all three memory stores and the transition between working/long-term memory\n \"\"\"\n def __init__(self, config):\n self.hidden_dim = config['hidden_dim']\n self.top_k = config['top_k']\n\n self.enable_long_term = config['enable_long_term']\n self.enable_long_term_usage = config['enable_long_term_count_usage']\n if self.enable_long_term:\n self.max_mt_frames = config['max_mid_term_frames']\n self.min_mt_frames = config['min_mid_term_frames']\n self.num_prototypes = config['num_prototypes']\n self.max_long_elements = config['max_long_term_elements']\n\n # dimensions will be inferred from input later\n self.CK = self.CV = None\n self.H = self.W = None\n\n # The hidden state will be stored in a single tensor for all objects\n # B x num_objects x CH x H x W\n self.hidden = None\n\n self.work_mem = KeyValueMemoryStore(count_usage=self.enable_long_term)\n if self.enable_long_term:\n self.long_mem = KeyValueMemoryStore(count_usage=self.enable_long_term_usage)\n\n self.reset_config = True\n\n def update_config(self, config):\n self.reset_config = True\n self.hidden_dim = config['hidden_dim']\n self.top_k = config['top_k']\n\n assert self.enable_long_term == config['enable_long_term'], 'cannot update this'\n assert self.enable_long_term_usage == config['enable_long_term_count_usage'], 'cannot update this'\n\n self.enable_long_term_usage = config['enable_long_term_count_usage']\n if self.enable_long_term:\n self.max_mt_frames = config['max_mid_term_frames']\n self.min_mt_frames = config['min_mid_term_frames']\n self.num_prototypes = config['num_prototypes']\n self.max_long_elements = config['max_long_term_elements']\n\n def _readout(self, affinity, v):\n # this function is for a single object group\n return v @ affinity\n\n def match_memory(self, query_key, selection):\n # query_key: B x C^k x H x W\n # selection: B x C^k x H x W\n num_groups = self.work_mem.num_groups\n h, w = query_key.shape[-2:]\n\n query_key = query_key.flatten(start_dim=2)\n selection = selection.flatten(start_dim=2) if selection is not None else None\n\n \"\"\"\n Memory readout using keys\n \"\"\"\n\n if self.enable_long_term and self.long_mem.engaged():\n # Use long-term memory\n long_mem_size = self.long_mem.size\n memory_key = torch.cat([self.long_mem.key, self.work_mem.key], -1)\n shrinkage = torch.cat([self.long_mem.shrinkage, self.work_mem.shrinkage], -1) \n\n similarity = get_similarity(memory_key, shrinkage, query_key, selection)\n work_mem_similarity = similarity[:, long_mem_size:]\n long_mem_similarity = similarity[:, :long_mem_size]\n\n # get the usage with the first group\n # the first group always have all the keys valid\n affinity, usage = do_softmax(\n torch.cat([long_mem_similarity[:, -self.long_mem.get_v_size(0):], work_mem_similarity], 1), \n top_k=self.top_k, inplace=True, return_usage=True)\n affinity = [affinity]\n\n # compute affinity group by group as later groups only have a subset of keys\n for gi in range(1, num_groups):\n if gi < self.long_mem.num_groups:\n # merge working and lt similarities before softmax\n affinity_one_group = do_softmax(\n torch.cat([long_mem_similarity[:, -self.long_mem.get_v_size(gi):], \n work_mem_similarity[:, -self.work_mem.get_v_size(gi):]], 1), \n top_k=self.top_k, inplace=True)\n else:\n # no long-term memory for this group\n affinity_one_group = do_softmax(work_mem_similarity[:, -self.work_mem.get_v_size(gi):], \n top_k=self.top_k, inplace=(gi==num_groups-1))\n affinity.append(affinity_one_group)\n\n all_memory_value = []\n for gi, gv in enumerate(self.work_mem.value):\n # merge the working and lt values before readout\n if gi < self.long_mem.num_groups:\n all_memory_value.append(torch.cat([self.long_mem.value[gi], self.work_mem.value[gi]], -1))\n else:\n all_memory_value.append(gv)\n\n \"\"\"\n Record memory usage for working and long-term memory\n \"\"\"\n # ignore the index return for long-term memory\n work_usage = usage[:, long_mem_size:]\n self.work_mem.update_usage(work_usage.flatten())\n\n if self.enable_long_term_usage:\n # ignore the index return for working memory\n long_usage = usage[:, :long_mem_size]\n self.long_mem.update_usage(long_usage.flatten())\n else:\n # No long-term memory\n similarity = get_similarity(self.work_mem.key, self.work_mem.shrinkage, query_key, selection)\n\n if self.enable_long_term:\n affinity, usage = do_softmax(similarity, inplace=(num_groups==1), \n top_k=self.top_k, return_usage=True)\n\n # Record memory usage for working memory\n self.work_mem.update_usage(usage.flatten())\n else:\n affinity = do_softmax(similarity, inplace=(num_groups==1), \n top_k=self.top_k, return_usage=False)\n\n affinity = [affinity]\n\n # compute affinity group by group as later groups only have a subset of keys\n for gi in range(1, num_groups):\n affinity_one_group = do_softmax(similarity[:, -self.work_mem.get_v_size(gi):], \n top_k=self.top_k, inplace=(gi==num_groups-1))\n affinity.append(affinity_one_group)\n \n all_memory_value = self.work_mem.value\n\n # Shared affinity within each group\n all_readout_mem = torch.cat([\n self._readout(affinity[gi], gv)\n for gi, gv in enumerate(all_memory_value)\n ], 0)\n\n return all_readout_mem.view(all_readout_mem.shape[0], self.CV, h, w)\n\n def add_memory(self, key, shrinkage, value, objects, selection=None):\n # key: 1*C*H*W\n # value: 1*num_objects*C*H*W\n # objects contain a list of object indices\n if self.H is None or self.reset_config:\n self.reset_config = False\n self.H, self.W = key.shape[-2:]\n self.HW = self.H*self.W\n if self.enable_long_term:\n # convert from num. frames to num. nodes\n self.min_work_elements = self.min_mt_frames*self.HW\n self.max_work_elements = self.max_mt_frames*self.HW\n\n # key: 1*C*N\n # value: num_objects*C*N\n key = key.flatten(start_dim=2)\n shrinkage = shrinkage.flatten(start_dim=2) \n value = value[0].flatten(start_dim=2)\n\n self.CK = key.shape[1]\n self.CV = value.shape[1]\n\n if selection is not None:\n if not self.enable_long_term:\n warnings.warn('the selection factor is only needed in long-term mode', UserWarning)\n selection = selection.flatten(start_dim=2)\n\n self.work_mem.add(key, value, shrinkage, selection, objects)\n\n # long-term memory cleanup\n if self.enable_long_term:\n # Do memory compressed if needed\n if self.work_mem.size >= self.max_work_elements:\n # Remove obsolete features if needed\n if self.long_mem.size >= (self.max_long_elements-self.num_prototypes):\n self.long_mem.remove_obsolete_features(self.max_long_elements-self.num_prototypes)\n \n self.compress_features()\n\n\n def create_hidden_state(self, n, sample_key):\n # n is the TOTAL number of objects\n h, w = sample_key.shape[-2:]\n if self.hidden is None:\n self.hidden = torch.zeros((1, n, self.hidden_dim, h, w), device=sample_key.device)\n elif self.hidden.shape[1] != n:\n self.hidden = torch.cat([\n self.hidden, \n torch.zeros((1, n-self.hidden.shape[1], self.hidden_dim, h, w), device=sample_key.device)\n ], 1)\n\n assert(self.hidden.shape[1] == n)\n\n def set_hidden(self, hidden):\n self.hidden = hidden\n\n def get_hidden(self):\n return self.hidden\n\n def compress_features(self):\n HW = self.HW\n candidate_value = []\n total_work_mem_size = self.work_mem.size\n for gv in self.work_mem.value:\n # Some object groups might be added later in the video\n # So not all keys have values associated with all objects\n # We need to keep track of the key->value validity\n mem_size_in_this_group = gv.shape[-1]\n if mem_size_in_this_group == total_work_mem_size:\n # full LT\n candidate_value.append(gv[:,:,HW:-self.min_work_elements+HW])\n else:\n # mem_size is smaller than total_work_mem_size, but at least HW\n assert HW <= mem_size_in_this_group < total_work_mem_size\n if mem_size_in_this_group > self.min_work_elements+HW:\n # part of this object group still goes into LT\n candidate_value.append(gv[:,:,HW:-self.min_work_elements+HW])\n else:\n # this object group cannot go to the LT at all\n candidate_value.append(None)\n\n # perform memory consolidation\n prototype_key, prototype_value, prototype_shrinkage = self.consolidation(\n *self.work_mem.get_all_sliced(HW, -self.min_work_elements+HW), candidate_value)\n\n # remove consolidated working memory\n self.work_mem.sieve_by_range(HW, -self.min_work_elements+HW, min_size=self.min_work_elements+HW)\n\n # add to long-term memory\n self.long_mem.add(prototype_key, prototype_value, prototype_shrinkage, selection=None, objects=None)\n\n def consolidation(self, candidate_key, candidate_shrinkage, candidate_selection, usage, candidate_value):\n # keys: 1*C*N\n # values: num_objects*C*N\n N = candidate_key.shape[-1]\n\n # find the indices with max usage\n _, max_usage_indices = torch.topk(usage, k=self.num_prototypes, dim=-1, sorted=True)\n prototype_indices = max_usage_indices.flatten()\n\n # Prototypes are invalid for out-of-bound groups\n validity = [prototype_indices >= (N-gv.shape[2]) if gv is not None else None for gv in candidate_value]\n\n prototype_key = candidate_key[:, :, prototype_indices]\n prototype_selection = candidate_selection[:, :, prototype_indices] if candidate_selection is not None else None\n\n \"\"\"\n Potentiation step\n \"\"\"\n similarity = get_similarity(candidate_key, candidate_shrinkage, prototype_key, prototype_selection)\n\n # convert similarity to affinity\n # need to do it group by group since the softmax normalization would be different\n affinity = [\n do_softmax(similarity[:, -gv.shape[2]:, validity[gi]]) if gv is not None else None\n for gi, gv in enumerate(candidate_value)\n ]\n\n # some values can be have all False validity. Weed them out.\n affinity = [\n aff if aff is None or aff.shape[-1] > 0 else None for aff in affinity\n ]\n\n # readout the values\n prototype_value = [\n self._readout(affinity[gi], gv) if affinity[gi] is not None else None\n for gi, gv in enumerate(candidate_value)\n ]\n\n # readout the shrinkage term\n prototype_shrinkage = self._readout(affinity[0], candidate_shrinkage) if candidate_shrinkage is not None else None\n\n return prototype_key, prototype_value, prototype_shrinkage\n" }, { "path": "XMem/interactive_demo.py", "content": "\"\"\"\nA simple user interface for XMem\n\"\"\"\n\nimport os\nfrom os import path\n# fix for Windows\nif 'QT_QPA_PLATFORM_PLUGIN_PATH' not in os.environ:\n os.environ['QT_QPA_PLATFORM_PLUGIN_PATH'] = ''\n\nimport signal\nsignal.signal(signal.SIGINT, signal.SIG_DFL)\n\nimport sys\nfrom argparse import ArgumentParser\n\nimport torch\n\nfrom model.network import XMem\nfrom inference.interact.s2m_controller import S2MController\nfrom inference.interact.fbrs_controller import FBRSController\nfrom inference.interact.s2m.s2m_network import deeplabv3plus_resnet50 as S2M\n\nfrom PySide6.QtWidgets import QApplication\nfrom inference.interact.gui import App\nfrom inference.interact.resource_manager import ResourceManager\nfrom contextlib import nullcontext\n\ntorch.set_grad_enabled(False)\n\nif torch.cuda.is_available():\n device = torch.device(\"cuda\")\nelif torch.backends.mps.is_available():\n device = torch.device(\"mps\")\nelse:\n device = torch.device(\"cpu\")\n\nif __name__ == '__main__':\n # Arguments parsing\n parser = ArgumentParser()\n parser.add_argument('--model', default='./saves/XMem.pth')\n parser.add_argument('--s2m_model', default='saves/s2m.pth')\n parser.add_argument('--fbrs_model', default='saves/fbrs.pth')\n\n \"\"\"\n Priority 1: If a \"images\" folder exists in the workspace, we will read from that directory\n Priority 2: If --images is specified, we will copy/resize those images to the workspace\n Priority 3: If --video is specified, we will extract the frames to the workspace (in an \"images\" folder) and read from there\n\n In any case, if a \"masks\" folder exists in the workspace, we will use that to initialize the mask\n That way, you can continue annotation from an interrupted run as long as the same workspace is used.\n \"\"\"\n parser.add_argument('--images', help='Folders containing input images.', default=None)\n parser.add_argument('--video', help='Video file readable by OpenCV.', default=None)\n parser.add_argument('--workspace', help='directory for storing buffered images (if needed) and output masks', default=None)\n\n parser.add_argument('--buffer_size', help='Correlate with CPU memory consumption', type=int, default=100)\n \n parser.add_argument('--num_objects', type=int, default=1)\n\n # Long-memory options\n # Defaults. Some can be changed in the GUI.\n parser.add_argument('--max_mid_term_frames', help='T_max in paper, decrease to save memory', type=int, default=10)\n parser.add_argument('--min_mid_term_frames', help='T_min in paper, decrease to save memory', type=int, default=5)\n parser.add_argument('--max_long_term_elements', help='LT_max in paper, increase if objects disappear for a long time', \n type=int, default=10000)\n parser.add_argument('--num_prototypes', help='P in paper', type=int, default=128) \n\n parser.add_argument('--top_k', type=int, default=30)\n parser.add_argument('--mem_every', type=int, default=10)\n parser.add_argument('--deep_update_every', help='Leave -1 normally to synchronize with mem_every', type=int, default=-1)\n parser.add_argument('--no_amp', help='Turn off AMP', action='store_true')\n parser.add_argument('--size', default=480, type=int, \n help='Resize the shorter side to this size. -1 to use original resolution. ')\n args = parser.parse_args()\n\n # create temporary workspace if not specified\n config = vars(args)\n config['enable_long_term'] = True\n config['enable_long_term_count_usage'] = True\n\n if config[\"workspace\"] is None:\n if config[\"images\"] is not None:\n basename = path.basename(config[\"images\"])\n elif config[\"video\"] is not None:\n basename = path.basename(config[\"video\"])[:-4]\n else:\n raise NotImplementedError(\n 'Either images, video, or workspace has to be specified')\n\n config[\"workspace\"] = path.join('./workspace', basename)\n\n with torch.cuda.amp.autocast(enabled=not args.no_amp) if device.type == 'cuda' else nullcontext():\n # Load our checkpoint\n network = XMem(config, args.model, map_location=device).to(device).eval()\n\n # Loads the S2M model\n if args.s2m_model is not None:\n s2m_saved = torch.load(args.s2m_model, map_location=device)\n s2m_model = S2M().to(device).eval()\n s2m_model.load_state_dict(s2m_saved)\n else:\n s2m_model = None\n\n s2m_controller = S2MController(s2m_model, args.num_objects, ignore_class=255, device=device)\n if args.fbrs_model is not None:\n fbrs_controller = FBRSController(args.fbrs_model, device=device)\n else:\n fbrs_controller = None\n\n # Manages most IO\n resource_manager = ResourceManager(config)\n\n app = QApplication(sys.argv)\n ex = App(network, resource_manager, s2m_controller, fbrs_controller, config, device)\n sys.exit(app.exec())\n" }, { "path": "XMem/merge_multi_scale.py", "content": "import os\nfrom os import path\nfrom argparse import ArgumentParser\nimport glob\nfrom collections import defaultdict\n\nimport numpy as np\nimport hickle as hkl\nfrom PIL import Image, ImagePalette\n\nfrom progressbar import progressbar\nfrom multiprocessing import Pool\nfrom util import palette\n\nfrom util.palette import davis_palette, youtube_palette\nimport shutil\n\n\ndef search_options(options, name):\n for option in options:\n if path.exists(path.join(option, name)):\n return path.join(option, name)\n else:\n return None\n\ndef process_vid(vid):\n vid_path = search_options(all_options, vid)\n if vid_path is not None:\n backward_mapping = hkl.load(path.join(vid_path, 'backward.hkl'))\n else:\n backward_mapping = None\n\n frames = os.listdir(path.join(all_options[0], vid))\n frames = [f for f in frames if 'backward' not in f]\n\n print(vid)\n if 'Y' in args.dataset:\n this_out_path = path.join(out_path, 'Annotations', vid)\n else:\n this_out_path = path.join(out_path, vid)\n os.makedirs(this_out_path, exist_ok=True)\n\n for f in frames:\n result_sum = None\n\n for option in all_options:\n if not path.exists(path.join(option, vid, f)):\n continue\n\n result = hkl.load(path.join(option, vid, f))\n if result_sum is None:\n result_sum = result.astype(np.float32)\n else:\n result_sum += result\n\n # argmax and to idx\n result_sum = np.argmax(result_sum, axis=0)\n\n # Remap the indices to the original domain\n if backward_mapping is not None:\n idx_mask = np.zeros_like(result_sum, dtype=np.uint8)\n for l, i in backward_mapping.items():\n idx_mask[result_sum==i] = l\n else:\n idx_mask = result_sum.astype(np.uint8)\n\n # Save the results\n img_E = Image.fromarray(idx_mask)\n img_E.putpalette(palette)\n img_E.save(path.join(this_out_path, f[:-4]+'.png'))\n\n\nif __name__ == '__main__':\n \"\"\"\n Arguments loading\n \"\"\"\n parser = ArgumentParser()\n parser.add_argument('--dataset', default='Y', help='D/Y, D for DAVIS; Y for YouTubeVOS')\n parser.add_argument('--list', nargs=\"+\")\n parser.add_argument('--pattern', default=None, help='Glob patten. Can be used in place of list.')\n parser.add_argument('--output')\n parser.add_argument('--num_proc', default=4, type=int)\n args = parser.parse_args()\n\n out_path = args.output\n\n # Find the input candidates\n if args.pattern is None:\n all_options = args.list\n else:\n assert args.list is None, 'cannot specify both list and pattern'\n all_options = glob.glob(args.pattern)\n\n # Get the correct palette\n if 'D' in args.dataset:\n palette = ImagePalette.ImagePalette(mode='P', palette=davis_palette)\n elif 'Y' in args.dataset:\n palette = ImagePalette.ImagePalette(mode='P', palette=youtube_palette)\n else:\n raise NotImplementedError\n\n # Count of the number of videos in each candidate\n all_options = [path.join(o, 'Scores') for o in all_options]\n vid_count = defaultdict(int)\n for option in all_options:\n vid_in_here = sorted(os.listdir(option))\n for vid in vid_in_here:\n vid_count[vid] += 1\n\n all_vid = []\n count_to_vid = defaultdict(int)\n for k, v in vid_count.items():\n count_to_vid[v] += 1\n all_vid.append(k)\n\n for k, v in count_to_vid.items():\n print('Videos with count %d: %d' % (k, v))\n\n all_vid = sorted(all_vid)\n print('Total number of videos: ', len(all_vid))\n\n pool = Pool(processes=args.num_proc)\n for _ in progressbar(pool.imap_unordered(process_vid, all_vid), max_value=len(all_vid)):\n pass\n\n pool.close()\n pool.join()\n\n if 'D' in args.dataset:\n print('Making zip for DAVIS test-dev...')\n shutil.make_archive(args.output, 'zip', args.output)\n\n if 'Y' in args.dataset:\n print('Making zip for YouTubeVOS...')\n shutil.make_archive(path.join(args.output, path.basename(args.output)), 'zip', args.output, 'Annotations')\n" }, { "path": "XMem/merge_results.py", "content": "import glob\nimport os\nfrom PIL import Image\nimport numpy as np\nimport tqdm\nimport multiprocessing\n\nmulti_dir = \"mevis_val/vis_output/\"\noutdir = \"mevis_val_merge/vis_output/\"\n\n\nall_obj_list = []\nvideo_list = glob.glob(os.path.join(multi_dir, \"0/*\"))\nfor video in video_list:\n obj_list = glob.glob(video + \"/*\")\n all_obj_list = all_obj_list + ['/'.join(i.split('/')[-2:]) for i in obj_list]\n\ndef merge(obj):\n obj_output_dir = os.path.join(outdir, obj)\n os.makedirs(obj_output_dir, exist_ok=True)\n img_list = [i.split('/')[-1] for i in glob.glob(os.path.join(multi_dir, \"0\", obj, \"*.png\"))]\n for img_name in img_list:\n agg_img = None\n for i in range(7):\n img_path = os.path.join(multi_dir, str(i), obj, img_name)\n tmp_img = (np.array(Image.open(img_path)) > 0).astype(np.uint8)\n if agg_img is not None:\n agg_img = agg_img + tmp_img\n else:\n agg_img = tmp_img\n agg_img = (agg_img >= 4).astype(np.uint8)\n agg_img = Image.fromarray(agg_img)\n img_output_path = os.path.join(obj_output_dir, img_name)\n agg_img.save(img_output_path)\n\nprint('start')\ncpu_num = multiprocessing.cpu_count()-1\nprint(\"cpu_num:\", cpu_num)\npool = multiprocessing.Pool(cpu_num)\n\nfor obj in all_obj_list:\n pool.apply_async(merge, args=(obj,))\n \npool.close()\npool.join()\n" }, { "path": "XMem/model/__init__.py", "content": "" }, { "path": "XMem/model/aggregate.py", "content": "import torch\nimport torch.nn.functional as F\n\n\n# Soft aggregation from STM\ndef aggregate(prob, dim, return_logits=False):\n new_prob = torch.cat([\n torch.prod(1-prob, dim=dim, keepdim=True),\n prob\n ], dim).clamp(1e-7, 1-1e-7)\n logits = torch.log((new_prob /(1-new_prob)))\n prob = F.softmax(logits, dim=dim)\n\n if return_logits:\n return logits, prob\n else:\n return prob" }, { "path": "XMem/model/cbam.py", "content": "# Modified from https://github.com/Jongchan/attention-module/blob/master/MODELS/cbam.py\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass BasicConv(nn.Module):\n def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):\n super(BasicConv, self).__init__()\n self.out_channels = out_planes\n self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)\n\n def forward(self, x):\n x = self.conv(x)\n return x\n\nclass Flatten(nn.Module):\n def forward(self, x):\n return x.view(x.size(0), -1)\n\nclass ChannelGate(nn.Module):\n def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):\n super(ChannelGate, self).__init__()\n self.gate_channels = gate_channels\n self.mlp = nn.Sequential(\n Flatten(),\n nn.Linear(gate_channels, gate_channels // reduction_ratio),\n nn.ReLU(),\n nn.Linear(gate_channels // reduction_ratio, gate_channels)\n )\n self.pool_types = pool_types\n def forward(self, x):\n channel_att_sum = None\n for pool_type in self.pool_types:\n if pool_type=='avg':\n avg_pool = F.avg_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))\n channel_att_raw = self.mlp( avg_pool )\n elif pool_type=='max':\n max_pool = F.max_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))\n channel_att_raw = self.mlp( max_pool )\n\n if channel_att_sum is None:\n channel_att_sum = channel_att_raw\n else:\n channel_att_sum = channel_att_sum + channel_att_raw\n\n scale = torch.sigmoid( channel_att_sum ).unsqueeze(2).unsqueeze(3).expand_as(x)\n return x * scale\n\nclass ChannelPool(nn.Module):\n def forward(self, x):\n return torch.cat( (torch.max(x,1)[0].unsqueeze(1), torch.mean(x,1).unsqueeze(1)), dim=1 )\n\nclass SpatialGate(nn.Module):\n def __init__(self):\n super(SpatialGate, self).__init__()\n kernel_size = 7\n self.compress = ChannelPool()\n self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size-1) // 2)\n def forward(self, x):\n x_compress = self.compress(x)\n x_out = self.spatial(x_compress)\n scale = torch.sigmoid(x_out) # broadcasting\n return x * scale\n\nclass CBAM(nn.Module):\n def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False):\n super(CBAM, self).__init__()\n self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types)\n self.no_spatial=no_spatial\n if not no_spatial:\n self.SpatialGate = SpatialGate()\n def forward(self, x):\n x_out = self.ChannelGate(x)\n if not self.no_spatial:\n x_out = self.SpatialGate(x_out)\n return x_out\n" }, { "path": "XMem/model/group_modules.py", "content": "\"\"\"\nGroup-specific modules\nThey handle features that also depends on the mask. \nFeatures are typically of shape\n batch_size * num_objects * num_channels * H * W\n\nAll of them are permutation equivariant w.r.t. to the num_objects dimension\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef interpolate_groups(g, ratio, mode, align_corners):\n batch_size, num_objects = g.shape[:2]\n g = F.interpolate(g.flatten(start_dim=0, end_dim=1), \n scale_factor=ratio, mode=mode, align_corners=align_corners)\n g = g.view(batch_size, num_objects, *g.shape[1:])\n return g\n\ndef upsample_groups(g, ratio=2, mode='bilinear', align_corners=False):\n return interpolate_groups(g, ratio, mode, align_corners)\n\ndef downsample_groups(g, ratio=1/2, mode='area', align_corners=None):\n return interpolate_groups(g, ratio, mode, align_corners)\n\n\nclass GConv2D(nn.Conv2d):\n def forward(self, g):\n batch_size, num_objects = g.shape[:2]\n g = super().forward(g.flatten(start_dim=0, end_dim=1))\n return g.view(batch_size, num_objects, *g.shape[1:])\n\n\nclass GroupResBlock(nn.Module):\n def __init__(self, in_dim, out_dim):\n super().__init__()\n\n if in_dim == out_dim:\n self.downsample = None\n else:\n self.downsample = GConv2D(in_dim, out_dim, kernel_size=3, padding=1)\n\n self.conv1 = GConv2D(in_dim, out_dim, kernel_size=3, padding=1)\n self.conv2 = GConv2D(out_dim, out_dim, kernel_size=3, padding=1)\n \n def forward(self, g):\n out_g = self.conv1(F.relu(g))\n out_g = self.conv2(F.relu(out_g))\n \n if self.downsample is not None:\n g = self.downsample(g)\n\n return out_g + g\n\n\nclass MainToGroupDistributor(nn.Module):\n def __init__(self, x_transform=None, method='cat', reverse_order=False):\n super().__init__()\n\n self.x_transform = x_transform\n self.method = method\n self.reverse_order = reverse_order\n\n def forward(self, x, g):\n num_objects = g.shape[1]\n\n if self.x_transform is not None:\n x = self.x_transform(x)\n\n if self.method == 'cat':\n if self.reverse_order:\n g = torch.cat([g, x.unsqueeze(1).expand(-1,num_objects,-1,-1,-1)], 2)\n else:\n g = torch.cat([x.unsqueeze(1).expand(-1,num_objects,-1,-1,-1), g], 2)\n elif self.method == 'add':\n g = x.unsqueeze(1).expand(-1,num_objects,-1,-1,-1) + g\n else:\n raise NotImplementedError\n\n return g\n" }, { "path": "XMem/model/losses.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom collections import defaultdict\n\n\ndef dice_loss(input_mask, cls_gt):\n num_objects = input_mask.shape[1]\n losses = []\n for i in range(num_objects):\n mask = input_mask[:,i].flatten(start_dim=1)\n # background not in mask, so we add one to cls_gt\n gt = (cls_gt==(i+1)).float().flatten(start_dim=1)\n numerator = 2 * (mask * gt).sum(-1)\n denominator = mask.sum(-1) + gt.sum(-1)\n loss = 1 - (numerator + 1) / (denominator + 1)\n losses.append(loss)\n return torch.cat(losses).mean()\n\n\n# https://stackoverflow.com/questions/63735255/how-do-i-compute-bootstrapped-cross-entropy-loss-in-pytorch\nclass BootstrappedCE(nn.Module):\n def __init__(self, start_warm, end_warm, top_p=0.15):\n super().__init__()\n\n self.start_warm = start_warm\n self.end_warm = end_warm\n self.top_p = top_p\n\n def forward(self, input, target, it):\n if it < self.start_warm:\n return F.cross_entropy(input, target), 1.0\n\n raw_loss = F.cross_entropy(input, target, reduction='none').view(-1)\n num_pixels = raw_loss.numel()\n\n if it > self.end_warm:\n this_p = self.top_p\n else:\n this_p = self.top_p + (1-self.top_p)*((self.end_warm-it)/(self.end_warm-self.start_warm))\n loss, _ = torch.topk(raw_loss, int(num_pixels * this_p), sorted=False)\n return loss.mean(), this_p\n\n\nclass LossComputer:\n def __init__(self, config):\n super().__init__()\n self.config = config\n self.bce = BootstrappedCE(config['start_warm'], config['end_warm'])\n\n def compute(self, data, num_objects, it):\n losses = defaultdict(int)\n\n b, t = data['rgb'].shape[:2]\n\n losses['total_loss'] = 0\n for ti in range(1, t):\n for bi in range(b):\n loss, p = self.bce(data[f'logits_{ti}'][bi:bi+1, :num_objects[bi]+1], data['cls_gt'][bi:bi+1,ti,0], it)\n losses['p'] += p / b / (t-1)\n losses[f'ce_loss_{ti}'] += loss / b\n\n losses['total_loss'] += losses['ce_loss_%d'%ti]\n losses[f'dice_loss_{ti}'] = dice_loss(data[f'masks_{ti}'], data['cls_gt'][:,ti,0])\n losses['total_loss'] += losses[f'dice_loss_{ti}']\n\n return losses\n" }, { "path": "XMem/model/memory_util.py", "content": "import math\nimport numpy as np\nimport torch\nfrom typing import Optional\n\n\ndef get_similarity(mk, ms, qk, qe):\n # used for training/inference and memory reading/memory potentiation\n # mk: B x CK x [N] - Memory keys\n # ms: B x 1 x [N] - Memory shrinkage\n # qk: B x CK x [HW/P] - Query keys\n # qe: B x CK x [HW/P] - Query selection\n # Dimensions in [] are flattened\n CK = mk.shape[1]\n mk = mk.flatten(start_dim=2)\n ms = ms.flatten(start_dim=1).unsqueeze(2) if ms is not None else None\n qk = qk.flatten(start_dim=2)\n qe = qe.flatten(start_dim=2) if qe is not None else None\n\n if qe is not None:\n # See appendix for derivation\n # or you can just trust me ใƒฝ(ใƒผ_ใƒผ )ใƒŽ\n mk = mk.transpose(1, 2)\n a_sq = (mk.pow(2) @ qe)\n two_ab = 2 * (mk @ (qk * qe))\n b_sq = (qe * qk.pow(2)).sum(1, keepdim=True)\n similarity = (-a_sq+two_ab-b_sq)\n else:\n # similar to STCN if we don't have the selection term\n a_sq = mk.pow(2).sum(1).unsqueeze(2)\n two_ab = 2 * (mk.transpose(1, 2) @ qk)\n similarity = (-a_sq+two_ab)\n\n if ms is not None:\n similarity = similarity * ms / math.sqrt(CK) # B*N*HW\n else:\n similarity = similarity / math.sqrt(CK) # B*N*HW\n\n return similarity\n\ndef do_softmax(similarity, top_k: Optional[int]=None, inplace=False, return_usage=False):\n # normalize similarity with top-k softmax\n # similarity: B x N x [HW/P]\n # use inplace with care\n if top_k is not None:\n values, indices = torch.topk(similarity, k=top_k, dim=1)\n\n x_exp = values.exp_()\n x_exp /= torch.sum(x_exp, dim=1, keepdim=True)\n if inplace:\n similarity.zero_().scatter_(1, indices, x_exp) # B*N*HW\n affinity = similarity\n else:\n affinity = torch.zeros_like(similarity).scatter_(1, indices, x_exp) # B*N*HW\n else:\n maxes = torch.max(similarity, dim=1, keepdim=True)[0]\n x_exp = torch.exp(similarity - maxes)\n x_exp_sum = torch.sum(x_exp, dim=1, keepdim=True)\n affinity = x_exp / x_exp_sum \n indices = None\n\n if return_usage:\n return affinity, affinity.sum(dim=2)\n\n return affinity\n\ndef get_affinity(mk, ms, qk, qe):\n # shorthand used in training with no top-k\n similarity = get_similarity(mk, ms, qk, qe)\n affinity = do_softmax(similarity)\n return affinity\n\ndef readout(affinity, mv):\n B, CV, T, H, W = mv.shape\n\n mo = mv.view(B, CV, T*H*W) \n mem = torch.bmm(mo, affinity)\n mem = mem.view(B, CV, H, W)\n\n return mem\n" }, { "path": "XMem/model/modules.py", "content": "\"\"\"\nmodules.py - This file stores the rather boring network blocks.\n\nx - usually means features that only depends on the image\ng - usually means features that also depends on the mask. \n They might have an extra \"group\" or \"num_objects\" dimension, hence\n batch_size * num_objects * num_channels * H * W\n\nThe trailing number of a variable usually denote the stride\n\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom model.group_modules import *\nfrom model import resnet\nfrom model.cbam import CBAM\n\n\nclass FeatureFusionBlock(nn.Module):\n def __init__(self, x_in_dim, g_in_dim, g_mid_dim, g_out_dim):\n super().__init__()\n\n self.distributor = MainToGroupDistributor()\n self.block1 = GroupResBlock(x_in_dim+g_in_dim, g_mid_dim)\n self.attention = CBAM(g_mid_dim)\n self.block2 = GroupResBlock(g_mid_dim, g_out_dim)\n\n def forward(self, x, g):\n batch_size, num_objects = g.shape[:2]\n\n g = self.distributor(x, g)\n g = self.block1(g)\n r = self.attention(g.flatten(start_dim=0, end_dim=1))\n r = r.view(batch_size, num_objects, *r.shape[1:])\n\n g = self.block2(g+r)\n\n return g\n\n\nclass HiddenUpdater(nn.Module):\n # Used in the decoder, multi-scale feature + GRU\n def __init__(self, g_dims, mid_dim, hidden_dim):\n super().__init__()\n self.hidden_dim = hidden_dim\n\n self.g16_conv = GConv2D(g_dims[0], mid_dim, kernel_size=1)\n self.g8_conv = GConv2D(g_dims[1], mid_dim, kernel_size=1)\n self.g4_conv = GConv2D(g_dims[2], mid_dim, kernel_size=1)\n\n self.transform = GConv2D(mid_dim+hidden_dim, hidden_dim*3, kernel_size=3, padding=1)\n\n nn.init.xavier_normal_(self.transform.weight)\n\n def forward(self, g, h):\n g = self.g16_conv(g[0]) + self.g8_conv(downsample_groups(g[1], ratio=1/2)) + \\\n self.g4_conv(downsample_groups(g[2], ratio=1/4))\n\n g = torch.cat([g, h], 2)\n\n # defined slightly differently than standard GRU, \n # namely the new value is generated before the forget gate.\n # might provide better gradient but frankly it was initially just an \n # implementation error that I never bothered fixing\n values = self.transform(g)\n forget_gate = torch.sigmoid(values[:,:,:self.hidden_dim])\n update_gate = torch.sigmoid(values[:,:,self.hidden_dim:self.hidden_dim*2])\n new_value = torch.tanh(values[:,:,self.hidden_dim*2:])\n new_h = forget_gate*h*(1-update_gate) + update_gate*new_value\n\n return new_h\n\n\nclass HiddenReinforcer(nn.Module):\n # Used in the value encoder, a single GRU\n def __init__(self, g_dim, hidden_dim):\n super().__init__()\n self.hidden_dim = hidden_dim\n self.transform = GConv2D(g_dim+hidden_dim, hidden_dim*3, kernel_size=3, padding=1)\n\n nn.init.xavier_normal_(self.transform.weight)\n\n def forward(self, g, h):\n g = torch.cat([g, h], 2)\n\n # defined slightly differently than standard GRU, \n # namely the new value is generated before the forget gate.\n # might provide better gradient but frankly it was initially just an \n # implementation error that I never bothered fixing\n values = self.transform(g)\n forget_gate = torch.sigmoid(values[:,:,:self.hidden_dim])\n update_gate = torch.sigmoid(values[:,:,self.hidden_dim:self.hidden_dim*2])\n new_value = torch.tanh(values[:,:,self.hidden_dim*2:])\n new_h = forget_gate*h*(1-update_gate) + update_gate*new_value\n\n return new_h\n\n\nclass ValueEncoder(nn.Module):\n def __init__(self, value_dim, hidden_dim, single_object=False):\n super().__init__()\n \n self.single_object = single_object\n network = resnet.resnet18(pretrained=True, extra_dim=1 if single_object else 2)\n self.conv1 = network.conv1\n self.bn1 = network.bn1\n self.relu = network.relu # 1/2, 64\n self.maxpool = network.maxpool\n\n self.layer1 = network.layer1 # 1/4, 64\n self.layer2 = network.layer2 # 1/8, 128\n self.layer3 = network.layer3 # 1/16, 256\n\n self.distributor = MainToGroupDistributor()\n self.fuser = FeatureFusionBlock(1024, 256, value_dim, value_dim)\n if hidden_dim > 0:\n self.hidden_reinforce = HiddenReinforcer(value_dim, hidden_dim)\n else:\n self.hidden_reinforce = None\n\n def forward(self, image, image_feat_f16, h, masks, others, is_deep_update=True):\n # image_feat_f16 is the feature from the key encoder\n if not self.single_object:\n g = torch.stack([masks, others], 2)\n else:\n g = masks.unsqueeze(2)\n g = self.distributor(image, g)\n\n batch_size, num_objects = g.shape[:2]\n g = g.flatten(start_dim=0, end_dim=1)\n\n g = self.conv1(g)\n g = self.bn1(g) # 1/2, 64\n g = self.maxpool(g) # 1/4, 64\n g = self.relu(g) \n\n g = self.layer1(g) # 1/4\n g = self.layer2(g) # 1/8\n g = self.layer3(g) # 1/16\n\n g = g.view(batch_size, num_objects, *g.shape[1:])\n g = self.fuser(image_feat_f16, g)\n\n if is_deep_update and self.hidden_reinforce is not None:\n h = self.hidden_reinforce(g, h)\n\n return g, h\n \n\nclass KeyEncoder(nn.Module):\n def __init__(self):\n super().__init__()\n network = resnet.resnet50(pretrained=True)\n self.conv1 = network.conv1\n self.bn1 = network.bn1\n self.relu = network.relu # 1/2, 64\n self.maxpool = network.maxpool\n\n self.res2 = network.layer1 # 1/4, 256\n self.layer2 = network.layer2 # 1/8, 512\n self.layer3 = network.layer3 # 1/16, 1024\n\n def forward(self, f):\n x = self.conv1(f) \n x = self.bn1(x)\n x = self.relu(x) # 1/2, 64\n x = self.maxpool(x) # 1/4, 64\n f4 = self.res2(x) # 1/4, 256\n f8 = self.layer2(f4) # 1/8, 512\n f16 = self.layer3(f8) # 1/16, 1024\n\n return f16, f8, f4\n\n\nclass UpsampleBlock(nn.Module):\n def __init__(self, skip_dim, g_up_dim, g_out_dim, scale_factor=2):\n super().__init__()\n self.skip_conv = nn.Conv2d(skip_dim, g_up_dim, kernel_size=3, padding=1)\n self.distributor = MainToGroupDistributor(method='add')\n self.out_conv = GroupResBlock(g_up_dim, g_out_dim)\n self.scale_factor = scale_factor\n\n def forward(self, skip_f, up_g):\n skip_f = self.skip_conv(skip_f)\n g = upsample_groups(up_g, ratio=self.scale_factor)\n g = self.distributor(skip_f, g)\n g = self.out_conv(g)\n return g\n\n\nclass KeyProjection(nn.Module):\n def __init__(self, in_dim, keydim):\n super().__init__()\n\n self.key_proj = nn.Conv2d(in_dim, keydim, kernel_size=3, padding=1)\n # shrinkage\n self.d_proj = nn.Conv2d(in_dim, 1, kernel_size=3, padding=1)\n # selection\n self.e_proj = nn.Conv2d(in_dim, keydim, kernel_size=3, padding=1)\n\n nn.init.orthogonal_(self.key_proj.weight.data)\n nn.init.zeros_(self.key_proj.bias.data)\n \n def forward(self, x, need_s, need_e):\n shrinkage = self.d_proj(x)**2 + 1 if (need_s) else None\n selection = torch.sigmoid(self.e_proj(x)) if (need_e) else None\n\n return self.key_proj(x), shrinkage, selection\n\n\nclass Decoder(nn.Module):\n def __init__(self, val_dim, hidden_dim):\n super().__init__()\n\n self.fuser = FeatureFusionBlock(1024, val_dim+hidden_dim, 512, 512)\n if hidden_dim > 0:\n self.hidden_update = HiddenUpdater([512, 256, 256+1], 256, hidden_dim)\n else:\n self.hidden_update = None\n \n self.up_16_8 = UpsampleBlock(512, 512, 256) # 1/16 -> 1/8\n self.up_8_4 = UpsampleBlock(256, 256, 256) # 1/8 -> 1/4\n\n self.pred = nn.Conv2d(256, 1, kernel_size=3, padding=1, stride=1)\n\n def forward(self, f16, f8, f4, hidden_state, memory_readout, h_out=True):\n batch_size, num_objects = memory_readout.shape[:2]\n\n if self.hidden_update is not None:\n g16 = self.fuser(f16, torch.cat([memory_readout, hidden_state], 2))\n else:\n g16 = self.fuser(f16, memory_readout)\n\n g8 = self.up_16_8(f8, g16)\n g4 = self.up_8_4(f4, g8)\n logits = self.pred(F.relu(g4.flatten(start_dim=0, end_dim=1)))\n\n if h_out and self.hidden_update is not None:\n g4 = torch.cat([g4, logits.view(batch_size, num_objects, 1, *logits.shape[-2:])], 2)\n hidden_state = self.hidden_update([g16, g8, g4], hidden_state)\n else:\n hidden_state = None\n \n logits = F.interpolate(logits, scale_factor=4, mode='bilinear', align_corners=False)\n logits = logits.view(batch_size, num_objects, *logits.shape[-2:])\n\n return hidden_state, logits\n" }, { "path": "XMem/model/network.py", "content": "\"\"\"\nThis file defines XMem, the highest level nn.Module interface\nDuring training, it is used by trainer.py\nDuring evaluation, it is used by inference_core.py\n\nIt further depends on modules.py which gives more detailed implementations of sub-modules\n\"\"\"\n\nimport torch\nimport torch.nn as nn\n\nfrom loguru import logger\n\nfrom model.aggregate import aggregate\nfrom model.modules import *\nfrom model.memory_util import *\n\n\nclass XMem(nn.Module):\n def __init__(self, config, model_path=None, map_location=None):\n \"\"\"\n model_path/map_location are used in evaluation only\n map_location is for converting models saved in cuda to cpu\n \"\"\"\n super().__init__()\n model_weights = self.init_hyperparameters(config, model_path, map_location)\n\n self.single_object = config.get('single_object', False)\n print(f'Single object mode: {self.single_object}')\n\n self.key_encoder = KeyEncoder()\n self.value_encoder = ValueEncoder(self.value_dim, self.hidden_dim, self.single_object)\n\n # Projection from f16 feature space to key/value space\n self.key_proj = KeyProjection(1024, self.key_dim)\n\n self.decoder = Decoder(self.value_dim, self.hidden_dim)\n\n if model_weights is not None:\n self.load_weights(model_weights, init_as_zero_if_needed=True)\n\n def encode_key(self, frame, need_sk=True, need_ek=True): \n # Determine input shape\n if len(frame.shape) == 5:\n # shape is b*t*c*h*w\n need_reshape = True\n b, t = frame.shape[:2]\n # flatten so that we can feed them into a 2D CNN\n frame = frame.flatten(start_dim=0, end_dim=1)\n elif len(frame.shape) == 4:\n # shape is b*c*h*w\n need_reshape = False\n else:\n raise NotImplementedError\n \n f16, f8, f4 = self.key_encoder(frame)\n key, shrinkage, selection = self.key_proj(f16, need_sk, need_ek)\n\n if need_reshape:\n # B*C*T*H*W\n key = key.view(b, t, *key.shape[-3:]).transpose(1, 2).contiguous()\n if shrinkage is not None:\n shrinkage = shrinkage.view(b, t, *shrinkage.shape[-3:]).transpose(1, 2).contiguous()\n if selection is not None:\n selection = selection.view(b, t, *selection.shape[-3:]).transpose(1, 2).contiguous()\n\n # B*T*C*H*W\n f16 = f16.view(b, t, *f16.shape[-3:])\n f8 = f8.view(b, t, *f8.shape[-3:])\n f4 = f4.view(b, t, *f4.shape[-3:])\n\n return key, shrinkage, selection, f16, f8, f4\n\n def encode_value(self, frame, image_feat_f16, h16, masks, is_deep_update=True): \n num_objects = masks.shape[1]\n if num_objects != 1:\n others = torch.cat([\n torch.sum(\n masks[:, [j for j in range(num_objects) if i!=j]]\n , dim=1, keepdim=True)\n for i in range(num_objects)], 1)\n else:\n others = torch.zeros_like(masks)\n\n g16, h16 = self.value_encoder(frame, image_feat_f16, h16, masks, others, is_deep_update)\n\n return g16, h16\n\n # Used in training only. \n # This step is replaced by MemoryManager in test time\n def read_memory(self, query_key, query_selection, memory_key, \n memory_shrinkage, memory_value):\n \"\"\"\n query_key : B * CK * H * W\n query_selection : B * CK * H * W\n memory_key : B * CK * T * H * W\n memory_shrinkage: B * 1 * T * H * W\n memory_value : B * num_objects * CV * T * H * W\n \"\"\"\n batch_size, num_objects = memory_value.shape[:2]\n memory_value = memory_value.flatten(start_dim=1, end_dim=2)\n\n affinity = get_affinity(memory_key, memory_shrinkage, query_key, query_selection)\n memory = readout(affinity, memory_value)\n memory = memory.view(batch_size, num_objects, self.value_dim, *memory.shape[-2:])\n\n return memory\n\n def segment(self, multi_scale_features, memory_readout,\n hidden_state, selector=None, h_out=True, strip_bg=True): \n\n hidden_state, logits = self.decoder(*multi_scale_features, hidden_state, memory_readout, h_out=h_out)\n prob = torch.sigmoid(logits)\n if selector is not None:\n prob = prob * selector\n \n logits, prob = aggregate(prob, dim=1, return_logits=True)\n if strip_bg:\n # Strip away the background\n prob = prob[:, 1:]\n\n return hidden_state, logits, prob\n\n def forward(self, mode, *args, **kwargs):\n if mode == 'encode_key':\n return self.encode_key(*args, **kwargs)\n elif mode == 'encode_value':\n return self.encode_value(*args, **kwargs)\n elif mode == 'read_memory':\n return self.read_memory(*args, **kwargs)\n elif mode == 'segment':\n return self.segment(*args, **kwargs)\n else:\n raise NotImplementedError\n\n @logger.catch()\n def init_hyperparameters(self, config, model_path=None, map_location=None):\n \"\"\"\n Init three hyperparameters: key_dim, value_dim, and hidden_dim\n If model_path is provided, we load these from the model weights\n The actual parameters are then updated to the config in-place\n\n Otherwise we load it either from the config or default\n \"\"\"\n if model_path is not None:\n # load the model and key/value/hidden dimensions with some hacks\n # config is updated with the loaded parameters\n model_weights = torch.load(model_path, map_location=map_location)\n self.key_dim = model_weights['key_proj.key_proj.weight'].shape[0]\n self.value_dim = model_weights['value_encoder.fuser.block2.conv2.weight'].shape[0]\n self.disable_hidden = 'decoder.hidden_update.transform.weight' not in model_weights\n if self.disable_hidden:\n self.hidden_dim = 0\n else:\n self.hidden_dim = model_weights['decoder.hidden_update.transform.weight'].shape[0]//3\n print(f'Hyperparameters read from the model weights: '\n f'C^k={self.key_dim}, C^v={self.value_dim}, C^h={self.hidden_dim}')\n else:\n model_weights = None\n # load dimensions from config or default\n if 'key_dim' not in config:\n self.key_dim = 64\n print(f'key_dim not found in config. Set to default {self.key_dim}')\n else:\n self.key_dim = config['key_dim']\n\n if 'value_dim' not in config:\n self.value_dim = 512\n print(f'value_dim not found in config. Set to default {self.value_dim}')\n else:\n self.value_dim = config['value_dim']\n\n if 'hidden_dim' not in config:\n self.hidden_dim = 64\n print(f'hidden_dim not found in config. Set to default {self.hidden_dim}')\n else:\n self.hidden_dim = config['hidden_dim']\n\n self.disable_hidden = (self.hidden_dim <= 0)\n\n config['key_dim'] = self.key_dim\n config['value_dim'] = self.value_dim\n config['hidden_dim'] = self.hidden_dim\n\n return model_weights\n\n @logger.catch()\n def load_weights(self, src_dict, init_as_zero_if_needed=False):\n # Maps SO weight (without other_mask) to MO weight (with other_mask)\n for k in list(src_dict.keys()):\n if k == 'value_encoder.conv1.weight':\n if src_dict[k].shape[1] == 4:\n print('Converting weights from single object to multiple objects.')\n pads = torch.zeros((64,1,7,7), device=src_dict[k].device)\n if not init_as_zero_if_needed:\n print('Randomly initialized padding.')\n nn.init.orthogonal_(pads)\n else:\n print('Zero-initialized padding.')\n src_dict[k] = torch.cat([src_dict[k], pads], 1)\n\n self.load_state_dict(src_dict)\n" }, { "path": "XMem/model/resnet.py", "content": "\"\"\"\nresnet.py - A modified ResNet structure\nWe append extra channels to the first conv by some network surgery\n\"\"\"\n\nfrom collections import OrderedDict\nimport math\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils import model_zoo\n\n\ndef load_weights_add_extra_dim(target, source_state, extra_dim=1):\n new_dict = OrderedDict()\n\n for k1, v1 in target.state_dict().items():\n if not 'num_batches_tracked' in k1:\n if k1 in source_state:\n tar_v = source_state[k1]\n\n if v1.shape != tar_v.shape:\n # Init the new segmentation channel with zeros\n # print(v1.shape, tar_v.shape)\n c, _, w, h = v1.shape\n pads = torch.zeros((c,extra_dim,w,h), device=tar_v.device)\n nn.init.orthogonal_(pads)\n tar_v = torch.cat([tar_v, pads], 1)\n\n new_dict[k1] = tar_v\n\n target.load_state_dict(new_dict)\n\n\nmodel_urls = {\n 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n}\n\n\ndef conv3x3(in_planes, out_planes, stride=1, dilation=1):\n return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n padding=dilation, dilation=dilation, bias=False)\n\n\nclass BasicBlock(nn.Module):\n expansion = 1\n\n def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):\n super(BasicBlock, self).__init__()\n self.conv1 = conv3x3(inplanes, planes, stride=stride, dilation=dilation)\n self.bn1 = nn.BatchNorm2d(planes)\n self.relu = nn.ReLU(inplace=True)\n self.conv2 = conv3x3(planes, planes, stride=1, dilation=dilation)\n self.bn2 = nn.BatchNorm2d(planes)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass Bottleneck(nn.Module):\n expansion = 4\n\n def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):\n super(Bottleneck, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n self.bn1 = nn.BatchNorm2d(planes)\n self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, dilation=dilation,\n padding=dilation, bias=False)\n self.bn2 = nn.BatchNorm2d(planes)\n self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)\n self.bn3 = nn.BatchNorm2d(planes * 4)\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass ResNet(nn.Module):\n def __init__(self, block, layers=(3, 4, 23, 3), extra_dim=0):\n self.inplanes = 64\n super(ResNet, self).__init__()\n self.conv1 = nn.Conv2d(3+extra_dim, 64, kernel_size=7, stride=2, padding=3, bias=False)\n self.bn1 = nn.BatchNorm2d(64)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n self.layer1 = self._make_layer(block, 64, layers[0])\n self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n m.weight.data.normal_(0, math.sqrt(2. / n))\n elif isinstance(m, nn.BatchNorm2d):\n m.weight.data.fill_(1)\n m.bias.data.zero_()\n\n def _make_layer(self, block, planes, blocks, stride=1, dilation=1):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(self.inplanes, planes * block.expansion,\n kernel_size=1, stride=stride, bias=False),\n nn.BatchNorm2d(planes * block.expansion),\n )\n\n layers = [block(self.inplanes, planes, stride, downsample)]\n self.inplanes = planes * block.expansion\n for i in range(1, blocks):\n layers.append(block(self.inplanes, planes, dilation=dilation))\n\n return nn.Sequential(*layers)\n\ndef resnet18(pretrained=True, extra_dim=0):\n model = ResNet(BasicBlock, [2, 2, 2, 2], extra_dim)\n if pretrained:\n load_weights_add_extra_dim(model, model_zoo.load_url(model_urls['resnet18']), extra_dim)\n return model\n\ndef resnet50(pretrained=True, extra_dim=0):\n model = ResNet(Bottleneck, [3, 4, 6, 3], extra_dim)\n if pretrained:\n load_weights_add_extra_dim(model, model_zoo.load_url(model_urls['resnet50']), extra_dim)\n return model\n\n" }, { "path": "XMem/model/trainer.py", "content": "\"\"\"\ntrainer.py - warpper and utility functions for network training\nCompute loss, back-prop, update parameters, logging, etc.\n\"\"\"\n\n\nimport os\nimport time\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\n\nfrom model.network import XMem\nfrom model.losses import LossComputer\nfrom util.log_integrator import Integrator\nfrom util.image_saver import pool_pairs\n\n\nclass XMemTrainer:\n def __init__(self, config, logger=None, save_path=None, local_rank=0, world_size=1):\n self.config = config\n self.num_frames = config['num_frames']\n self.num_ref_frames = config['num_ref_frames']\n self.deep_update_prob = config['deep_update_prob']\n self.local_rank = local_rank\n\n self.XMem = nn.parallel.DistributedDataParallel(\n XMem(config).cuda(), \n device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False)\n\n # Set up logger when local_rank=0\n self.logger = logger\n self.save_path = save_path\n if logger is not None:\n self.last_time = time.time()\n self.logger.log_string('model_size', str(sum([param.nelement() for param in self.XMem.parameters()])))\n self.train_integrator = Integrator(self.logger, distributed=True, local_rank=local_rank, world_size=world_size)\n self.loss_computer = LossComputer(config)\n\n self.train()\n self.optimizer = optim.AdamW(filter(\n lambda p: p.requires_grad, self.XMem.parameters()), lr=config['lr'], weight_decay=config['weight_decay'])\n self.scheduler = optim.lr_scheduler.MultiStepLR(self.optimizer, config['steps'], config['gamma'])\n if config['amp']:\n self.scaler = torch.cuda.amp.GradScaler()\n\n # Logging info\n self.log_text_interval = config['log_text_interval']\n self.log_image_interval = config['log_image_interval']\n self.save_network_interval = config['save_network_interval']\n self.save_checkpoint_interval = config['save_checkpoint_interval']\n if config['debug']:\n self.log_text_interval = self.log_image_interval = 1\n\n def do_pass(self, data, it=0):\n # No need to store the gradient outside training\n torch.set_grad_enabled(self._is_train)\n\n for k, v in data.items():\n if type(v) != list and type(v) != dict and type(v) != int:\n data[k] = v.cuda(non_blocking=True)\n\n out = {}\n frames = data['rgb']\n first_frame_gt = data['first_frame_gt'].float()\n b = frames.shape[0]\n num_filled_objects = [o.item() for o in data['info']['num_objects']]\n num_objects = first_frame_gt.shape[2]\n selector = data['selector'].unsqueeze(2).unsqueeze(2)\n\n with torch.cuda.amp.autocast(enabled=self.config['amp']):\n # image features never change, compute once\n key, shrinkage, selection, f16, f8, f4 = self.XMem('encode_key', frames)\n\n filler_one = torch.zeros(1, dtype=torch.int64)\n hidden = torch.zeros((b, num_objects, self.config['hidden_dim'], *key.shape[-2:]))\n v16, hidden = self.XMem('encode_value', frames[:,0], f16[:,0], hidden, first_frame_gt[:,0])\n values = v16.unsqueeze(3) # add the time dimension\n\n for ti in range(1, self.num_frames):\n if ti <= self.num_ref_frames:\n ref_values = values\n ref_keys = key[:,:,:ti]\n ref_shrinkage = shrinkage[:,:,:ti] if shrinkage is not None else None\n else:\n # pick num_ref_frames random frames\n # this is not very efficient but I think we would \n # need broadcasting in gather which we don't have\n indices = [\n torch.cat([filler_one, torch.randperm(ti-1)[:self.num_ref_frames-1]+1])\n for _ in range(b)]\n ref_values = torch.stack([\n values[bi, :, :, indices[bi]] for bi in range(b)\n ], 0)\n ref_keys = torch.stack([\n key[bi, :, indices[bi]] for bi in range(b)\n ], 0)\n ref_shrinkage = torch.stack([\n shrinkage[bi, :, indices[bi]] for bi in range(b)\n ], 0) if shrinkage is not None else None\n\n # Segment frame ti\n memory_readout = self.XMem('read_memory', key[:,:,ti], selection[:,:,ti] if selection is not None else None, \n ref_keys, ref_shrinkage, ref_values)\n hidden, logits, masks = self.XMem('segment', (f16[:,ti], f8[:,ti], f4[:,ti]), memory_readout, \n hidden, selector, h_out=(ti < (self.num_frames-1)))\n\n # No need to encode the last frame\n if ti < (self.num_frames-1):\n is_deep_update = np.random.rand() < self.deep_update_prob\n v16, hidden = self.XMem('encode_value', frames[:,ti], f16[:,ti], hidden, masks, is_deep_update=is_deep_update)\n values = torch.cat([values, v16.unsqueeze(3)], 3)\n\n out[f'masks_{ti}'] = masks\n out[f'logits_{ti}'] = logits\n\n if self._do_log or self._is_train:\n losses = self.loss_computer.compute({**data, **out}, num_filled_objects, it)\n\n # Logging\n if self._do_log:\n self.integrator.add_dict(losses)\n if self._is_train:\n if it % self.log_image_interval == 0 and it != 0:\n if self.logger is not None:\n images = {**data, **out}\n size = (384, 384)\n self.logger.log_cv2('train/pairs', pool_pairs(images, size, num_filled_objects), it)\n\n if self._is_train:\n if (it) % self.log_text_interval == 0 and it != 0:\n if self.logger is not None:\n self.logger.log_scalar('train/lr', self.scheduler.get_last_lr()[0], it)\n self.logger.log_metrics('train', 'time', (time.time()-self.last_time)/self.log_text_interval, it)\n self.last_time = time.time()\n self.train_integrator.finalize('train', it)\n self.train_integrator.reset_except_hooks()\n\n if it % self.save_network_interval == 0 and it != 0:\n if self.logger is not None:\n self.save_network(it)\n\n if it % self.save_checkpoint_interval == 0 and it != 0:\n if self.logger is not None:\n self.save_checkpoint(it)\n\n # Backward pass\n self.optimizer.zero_grad(set_to_none=True)\n if self.config['amp']:\n self.scaler.scale(losses['total_loss']).backward()\n self.scaler.step(self.optimizer)\n self.scaler.update()\n else:\n losses['total_loss'].backward() \n self.optimizer.step()\n\n self.scheduler.step()\n\n def save_network(self, it):\n if self.save_path is None:\n print('Saving has been disabled.')\n return\n \n os.makedirs(os.path.dirname(self.save_path), exist_ok=True)\n model_path = f'{self.save_path}_{it}.pth'\n torch.save(self.XMem.module.state_dict(), model_path)\n print(f'Network saved to {model_path}.')\n\n def save_checkpoint(self, it):\n if self.save_path is None:\n print('Saving has been disabled.')\n return\n\n os.makedirs(os.path.dirname(self.save_path), exist_ok=True)\n checkpoint_path = f'{self.save_path}_checkpoint_{it}.pth'\n checkpoint = { \n 'it': it,\n 'network': self.XMem.module.state_dict(),\n 'optimizer': self.optimizer.state_dict(),\n 'scheduler': self.scheduler.state_dict()}\n torch.save(checkpoint, checkpoint_path)\n print(f'Checkpoint saved to {checkpoint_path}.')\n\n def load_checkpoint(self, path):\n # This method loads everything and should be used to resume training\n map_location = 'cuda:%d' % self.local_rank\n checkpoint = torch.load(path, map_location={'cuda:0': map_location})\n\n it = checkpoint['it']\n network = checkpoint['network']\n optimizer = checkpoint['optimizer']\n scheduler = checkpoint['scheduler']\n\n map_location = 'cuda:%d' % self.local_rank\n self.XMem.module.load_state_dict(network)\n self.optimizer.load_state_dict(optimizer)\n self.scheduler.load_state_dict(scheduler)\n\n print('Network weights, optimizer states, and scheduler states loaded.')\n\n return it\n\n def load_network_in_memory(self, src_dict):\n self.XMem.module.load_weights(src_dict)\n print('Network weight loaded from memory.')\n\n def load_network(self, path):\n # This method loads only the network weight and should be used to load a pretrained model\n map_location = 'cuda:%d' % self.local_rank\n src_dict = torch.load(path, map_location={'cuda:0': map_location})\n\n self.load_network_in_memory(src_dict)\n print(f'Network weight loaded from {path}')\n\n def train(self):\n self._is_train = True\n self._do_log = True\n self.integrator = self.train_integrator\n self.XMem.eval()\n return self\n\n def val(self):\n self._is_train = False\n self._do_log = True\n self.XMem.eval()\n return self\n\n def test(self):\n self._is_train = False\n self._do_log = False\n self.XMem.eval()\n return self\n\n" }, { "path": "XMem/requirements.txt", "content": "progressbar2\ngdown\nhickle\ntensorboard\nnumpy" }, { "path": "XMem/scripts/__init__.py", "content": "" }, { "path": "XMem/scripts/download_bl30k.py", "content": "import os\nimport gdown\nimport tarfile\n\n\nLICENSE = \"\"\"\nThis dataset is a derivative of ShapeNet.\nPlease read and respect their licenses and terms before use. \nTextures and skybox image are obtained from Google image search with the \"non-commercial reuse\" flag.\nDo not use this dataset for commercial purposes.\nYou should cite both ShapeNet and our paper if you use this dataset.\n\"\"\"\n\nprint(LICENSE)\nprint('Datasets will be downloaded and extracted to ../BL30K')\nprint('The script will download and extract the segment one by one')\nprint('You are going to need ~1TB of free disk space')\nreply = input('[y] to confirm, others to exit: ')\nif reply != 'y':\n exit()\n\nlinks = [\n 'https://drive.google.com/uc?id=1z9V5zxLOJLNt1Uj7RFqaP2FZWKzyXvVc',\n 'https://drive.google.com/uc?id=11-IzgNwEAPxgagb67FSrBdzZR7OKAEdJ',\n 'https://drive.google.com/uc?id=1ZfIv6GTo-OGpXpoKen1fUvDQ0A_WoQ-Q',\n 'https://drive.google.com/uc?id=1G4eXgYS2kL7_Cc0x3N1g1x7Zl8D_aU_-',\n 'https://drive.google.com/uc?id=1Y8q0V_oBwJIY27W_6-8CD1dRqV2gNTdE',\n 'https://drive.google.com/uc?id=1nawBAazf_unMv46qGBHhWcQ4JXZ5883r',\n]\n\nnames = [\n 'BL30K_a.tar',\n 'BL30K_b.tar',\n 'BL30K_c.tar',\n 'BL30K_d.tar',\n 'BL30K_e.tar',\n 'BL30K_f.tar',\n]\n\nfor i, link in enumerate(links):\n print('Downloading segment %d/%d ...' % (i, len(links)))\n gdown.download(link, output='../%s' % names[i], quiet=False)\n print('Extracting...')\n with tarfile.open('../%s' % names[i], 'r') as tar_file:\n tar_file.extractall('../%s' % names[i])\n print('Cleaning up...')\n os.remove('../%s' % names[i])\n\n\nprint('Done.')" }, { "path": "XMem/scripts/download_datasets.py", "content": "import os\nimport gdown\nimport zipfile\nfrom scripts import resize_youtube\n\n\nLICENSE = \"\"\"\nThese are either re-distribution of the original datasets or derivatives (through simple processing) of the original datasets. \nPlease read and respect their licenses and terms before use. \nYou should cite the original papers if you use any of the datasets.\n\nFor BL30K, see download_bl30k.py\n\nLinks:\nDUTS: http://saliencydetection.net/duts\nHRSOD: https://github.com/yi94code/HRSOD\nFSS: https://github.com/HKUSTCV/FSS-1000\nECSSD: https://www.cse.cuhk.edu.hk/leojia/projects/hsaliency/dataset.html\nBIG: https://github.com/hkchengrex/CascadePSP\n\nYouTubeVOS: https://youtube-vos.org\nDAVIS: https://davischallenge.org/\nBL30K: https://github.com/hkchengrex/MiVOS\nLong-Time Video: https://github.com/xmlyqing00/AFB-URR\n\"\"\"\n\nprint(LICENSE)\nprint('Datasets will be downloaded and extracted to ../YouTube, ../YouTube2018, ../static, ../DAVIS, ../long_video_set')\nreply = input('[y] to confirm, others to exit: ')\nif reply != 'y':\n exit()\n\n\n\"\"\"\nStatic image data\n\"\"\"\nos.makedirs('../static', exist_ok=True)\nprint('Downloading static datasets...')\ngdown.download('https://drive.google.com/uc?id=1wUJq3HcLdN-z1t4CsUhjeZ9BVDb9YKLd', output='../static/static_data.zip', quiet=False)\nprint('Extracting static datasets...')\nwith zipfile.ZipFile('../static/static_data.zip', 'r') as zip_file:\n zip_file.extractall('../static/')\nprint('Cleaning up static datasets...')\nos.remove('../static/static_data.zip')\n\n\n\"\"\"\nDAVIS dataset\n\"\"\"\n# Google drive mirror: https://drive.google.com/drive/folders/1hEczGHw7qcMScbCJukZsoOW4Q9byx16A?usp=sharing\nos.makedirs('../DAVIS/2017', exist_ok=True)\n\nprint('Downloading DAVIS 2016...')\ngdown.download('https://drive.google.com/uc?id=198aRlh5CpAoFz0hfRgYbiNenn_K8DxWD', output='../DAVIS/DAVIS-data.zip', quiet=False)\n\nprint('Downloading DAVIS 2017 trainval...')\ngdown.download('https://drive.google.com/uc?id=1kiaxrX_4GuW6NmiVuKGSGVoKGWjOdp6d', output='../DAVIS/2017/DAVIS-2017-trainval-480p.zip', quiet=False)\n\nprint('Downloading DAVIS 2017 testdev...')\ngdown.download('https://drive.google.com/uc?id=1fmkxU2v9cQwyb62Tj1xFDdh2p4kDsUzD', output='../DAVIS/2017/DAVIS-2017-test-dev-480p.zip', quiet=False)\n\nprint('Downloading DAVIS 2017 scribbles...')\ngdown.download('https://drive.google.com/uc?id=1JzIQSu36h7dVM8q0VoE4oZJwBXvrZlkl', output='../DAVIS/2017/DAVIS-2017-scribbles-trainval.zip', quiet=False)\n\nprint('Extracting DAVIS datasets...')\nwith zipfile.ZipFile('../DAVIS/DAVIS-data.zip', 'r') as zip_file:\n zip_file.extractall('../DAVIS/')\nos.rename('../DAVIS/DAVIS', '../DAVIS/2016')\n\nwith zipfile.ZipFile('../DAVIS/2017/DAVIS-2017-trainval-480p.zip', 'r') as zip_file:\n zip_file.extractall('../DAVIS/2017/')\nwith zipfile.ZipFile('../DAVIS/2017/DAVIS-2017-scribbles-trainval.zip', 'r') as zip_file:\n zip_file.extractall('../DAVIS/2017/')\nos.rename('../DAVIS/2017/DAVIS', '../DAVIS/2017/trainval')\n\nwith zipfile.ZipFile('../DAVIS/2017/DAVIS-2017-test-dev-480p.zip', 'r') as zip_file:\n zip_file.extractall('../DAVIS/2017/')\nos.rename('../DAVIS/2017/DAVIS', '../DAVIS/2017/test-dev')\n\nprint('Cleaning up DAVIS datasets...')\nos.remove('../DAVIS/2017/DAVIS-2017-trainval-480p.zip')\nos.remove('../DAVIS/2017/DAVIS-2017-test-dev-480p.zip')\nos.remove('../DAVIS/2017/DAVIS-2017-scribbles-trainval.zip')\nos.remove('../DAVIS/DAVIS-data.zip')\n\n\n\"\"\"\nYouTubeVOS dataset\n\"\"\"\nos.makedirs('../YouTube', exist_ok=True)\nos.makedirs('../YouTube/all_frames', exist_ok=True)\n\nprint('Downloading YouTubeVOS train...')\ngdown.download('https://drive.google.com/uc?id=13Eqw0gVK-AO5B-cqvJ203mZ2vzWck9s4', output='../YouTube/train.zip', quiet=False)\nprint('Downloading YouTubeVOS val...')\ngdown.download('https://drive.google.com/uc?id=1o586Wjya-f2ohxYf9C1RlRH-gkrzGS8t', output='../YouTube/valid.zip', quiet=False)\nprint('Downloading YouTubeVOS all frames valid...')\ngdown.download('https://drive.google.com/uc?id=1rWQzZcMskgpEQOZdJPJ7eTmLCBEIIpEN', output='../YouTube/all_frames/valid.zip', quiet=False)\n\nprint('Extracting YouTube datasets...')\nwith zipfile.ZipFile('../YouTube/train.zip', 'r') as zip_file:\n zip_file.extractall('../YouTube/')\nwith zipfile.ZipFile('../YouTube/valid.zip', 'r') as zip_file:\n zip_file.extractall('../YouTube/')\nwith zipfile.ZipFile('../YouTube/all_frames/valid.zip', 'r') as zip_file:\n zip_file.extractall('../YouTube/all_frames')\n\nprint('Cleaning up YouTubeVOS datasets...')\nos.remove('../YouTube/train.zip')\nos.remove('../YouTube/valid.zip')\nos.remove('../YouTube/all_frames/valid.zip')\n\nprint('Resizing YouTubeVOS to 480p...')\nresize_youtube.resize_all('../YouTube/train', '../YouTube/train_480p')\n\n# YouTubeVOS 2018\nos.makedirs('../YouTube2018', exist_ok=True)\nos.makedirs('../YouTube2018/all_frames', exist_ok=True)\n\nprint('Downloading YouTubeVOS2018 val...')\ngdown.download('https://drive.google.com/uc?id=1-QrceIl5sUNTKz7Iq0UsWC6NLZq7girr', output='../YouTube2018/valid.zip', quiet=False)\nprint('Downloading YouTubeVOS2018 all frames valid...')\ngdown.download('https://drive.google.com/uc?id=1yVoHM6zgdcL348cFpolFcEl4IC1gorbV', output='../YouTube2018/all_frames/valid.zip', quiet=False)\n\nprint('Extracting YouTube2018 datasets...')\nwith zipfile.ZipFile('../YouTube2018/valid.zip', 'r') as zip_file:\n zip_file.extractall('../YouTube2018/')\nwith zipfile.ZipFile('../YouTube2018/all_frames/valid.zip', 'r') as zip_file:\n zip_file.extractall('../YouTube2018/all_frames')\n\nprint('Cleaning up YouTubeVOS2018 datasets...')\nos.remove('../YouTube2018/valid.zip')\nos.remove('../YouTube2018/all_frames/valid.zip')\n\n\n\"\"\"\nLong-Time Video dataset\n\"\"\"\nos.makedirs('../long_video_set', exist_ok=True)\nprint('Downloading long video dataset...')\ngdown.download('https://drive.google.com/uc?id=100MxAuV0_UL20ca5c-5CNpqQ5QYPDSoz', output='../long_video_set/LongTimeVideo.zip', quiet=False)\nprint('Extracting long video dataset...')\nwith zipfile.ZipFile('../long_video_set/LongTimeVideo.zip', 'r') as zip_file:\n zip_file.extractall('../long_video_set/')\nprint('Cleaning up long video dataset...')\nos.remove('../long_video_set/LongTimeVideo.zip')\n\n\nprint('Done.')" }, { "path": "XMem/scripts/download_models.sh", "content": "wget -P ./saves/ https://github.com/hkchengrex/XMem/releases/download/v1.0/XMem.pth\nwget -P ./saves/ https://github.com/hkchengrex/XMem/releases/download/v1.0/XMem-s012.pth" }, { "path": "XMem/scripts/download_models_demo.sh", "content": "wget -P ./saves/ https://github.com/hkchengrex/XMem/releases/download/v1.0/XMem.pth\nwget -P ./saves/ https://github.com/hkchengrex/XMem/releases/download/v1.0/fbrs.pth\nwget -P ./saves/ https://github.com/hkchengrex/XMem/releases/download/v1.0/s2m.pth" }, { "path": "XMem/scripts/expand_long_vid.py", "content": "import sys\nimport os\nfrom os import path\nfrom shutil import copy2\n\ninput_path = sys.argv[1]\noutput_path = sys.argv[2]\nmultiplier = int(sys.argv[3])\nimage_path = path.join(input_path, 'JPEGImages')\ngt_path = path.join(input_path, 'Annotations')\n\nvideos = sorted(os.listdir(image_path))\n\nfor vid in videos:\n os.makedirs(path.join(output_path, 'JPEGImages', vid), exist_ok=True)\n os.makedirs(path.join(output_path, 'Annotations', vid), exist_ok=True)\n frames = sorted(os.listdir(path.join(image_path, vid)))\n\n num_frames = len(frames)\n counter = 0\n output_counter = 0\n direction = 1\n for _ in range(multiplier):\n for _ in range(num_frames):\n copy2(path.join(image_path, vid, frames[counter]), \n path.join(output_path, 'JPEGImages', vid, f'{output_counter:05d}.jpg'))\n\n mask_path = path.join(gt_path, vid, frames[counter].replace('.jpg', '.png'))\n if path.exists(mask_path):\n copy2(mask_path, \n path.join(output_path, 'Annotations', vid, f'{output_counter:05d}.png'))\n\n counter += direction\n output_counter += 1\n if counter == 0 or counter == len(frames) - 1:\n direction *= -1\n" }, { "path": "XMem/scripts/resize_youtube.py", "content": "import sys\nimport os\nfrom os import path\n\nfrom PIL import Image\nimport numpy as np\nfrom progressbar import progressbar\nfrom multiprocessing import Pool\n\nnew_min_size = 480\n\ndef resize_vid_jpeg(inputs):\n vid_name, folder_path, out_path = inputs\n\n vid_path = path.join(folder_path, vid_name)\n vid_out_path = path.join(out_path, 'JPEGImages', vid_name)\n os.makedirs(vid_out_path, exist_ok=True)\n\n for im_name in os.listdir(vid_path):\n hr_im = Image.open(path.join(vid_path, im_name))\n w, h = hr_im.size\n\n ratio = new_min_size / min(w, h)\n\n lr_im = hr_im.resize((int(w*ratio), int(h*ratio)), Image.BICUBIC)\n lr_im.save(path.join(vid_out_path, im_name))\n\ndef resize_vid_anno(inputs):\n vid_name, folder_path, out_path = inputs\n\n vid_path = path.join(folder_path, vid_name)\n vid_out_path = path.join(out_path, 'Annotations', vid_name)\n os.makedirs(vid_out_path, exist_ok=True)\n\n for im_name in os.listdir(vid_path):\n hr_im = Image.open(path.join(vid_path, im_name)).convert('P')\n w, h = hr_im.size\n\n ratio = new_min_size / min(w, h)\n\n lr_im = hr_im.resize((int(w*ratio), int(h*ratio)), Image.NEAREST)\n lr_im.save(path.join(vid_out_path, im_name))\n\n\ndef resize_all(in_path, out_path):\n for folder in os.listdir(in_path):\n\n if folder not in ['JPEGImages', 'Annotations']:\n continue\n folder_path = path.join(in_path, folder)\n videos = os.listdir(folder_path)\n\n videos = [(v, folder_path, out_path) for v in videos]\n\n if folder == 'JPEGImages':\n print('Processing images')\n os.makedirs(path.join(out_path, 'JPEGImages'), exist_ok=True)\n\n pool = Pool(processes=8)\n for _ in progressbar(pool.imap_unordered(resize_vid_jpeg, videos), max_value=len(videos)):\n pass\n else:\n print('Processing annotations')\n os.makedirs(path.join(out_path, 'Annotations'), exist_ok=True)\n\n pool = Pool(processes=8)\n for _ in progressbar(pool.imap_unordered(resize_vid_anno, videos), max_value=len(videos)):\n pass\n\n\nif __name__ == '__main__':\n in_path = sys.argv[1]\n out_path = sys.argv[2]\n\n resize_all(in_path, out_path)\n\n print('Done.')" }, { "path": "XMem/tracking.py", "content": "import sys\nsys.path.insert(0, './XMem')\n\nimport os\nimport os.path as osp\nimport glob\nimport cv2\nimport json\nimport argparse\nimport multiprocessing as mp\nfrom tqdm import tqdm\nfrom termcolor import colored\n\nfrom importlib.util import find_spec\nif find_spec(\"GPUtil\") is None: os.system(\"pip install gputil\")\nimport GPUtil\n\n_GPU_LIST = [_.id for _ in GPUtil.getGPUs()]\n_GPU_QUEUE = mp.Queue()\nfor _ in _GPU_LIST: _GPU_QUEUE.put(_)\n\ndef run_eval(meta_expression, temp_xmem_anno, final_xmem_anno, img_dir, split_part, xmem_weight, cfgs=\" --reversed \", ):\n gpu_id = _GPU_QUEUE.get()\n cmd = f\"cd XMem && CUDA_VISIBLE_DEVICES={gpu_id} python eval.py --meta_exp {meta_expression} --output {final_xmem_anno} --generic_path {temp_xmem_anno} --img_dir {img_dir} --split_part {split_part} --model {xmem_weight} --dataset G {cfgs}\"\n print(f\"Running: {cmd}\")\n os.system(cmd)\n _GPU_QUEUE.put(gpu_id)\n\ndef generate(obj, temp_xmem_anno, final_xmem_anno):\n\n obj_dir, video_name, obj_id, tp = obj\n img_list = glob.glob(obj_dir + '/*.png') # Mask\n img_list.sort()\n frame_id = int(len(img_list) * tp)\n if frame_id == len(img_list):\n frame_id -= 1\n \n used_img = img_list[frame_id]\n\n img_output_path = osp.join(temp_xmem_anno, video_name, obj_id, osp.basename(used_img))\n final_img_output_dir = osp.join(final_xmem_anno, video_name, obj_id)\n img_output_dir = osp.dirname(img_output_path)\n os.makedirs(img_output_dir, exist_ok=True)\n os.makedirs(final_img_output_dir, exist_ok=True)\n os.system('cp {} {}'.format(used_img, img_output_path))\n\n img = cv2.imread(img_output_path)\n if img.sum() == 0:\n target_img_list = [i.split('/')[-1] for i in img_list]\n for img_ in target_img_list:\n print(os.path.join(final_img_output_dir, img_))\n os.system('cp {} {}'.format(img_output_path, os.path.join(img_output_dir, img_)))\n os.system('cp {} {}'.format(img_output_path, os.path.join(final_img_output_dir, img_)))\n return 0\n\ndef prepare(args):\n video_root = args.video_root\n temp_xmem_anno = args.temp_xmem_anno\n final_xmem_anno = args.final_xmem_anno\n os.makedirs(temp_xmem_anno, exist_ok=True)\n\n data = json.load(open(args.llama_vid_meta, 'r'))['videos']\n\n all_obj_list = []\n for video_name in data.keys():\n exps = data[video_name]['expressions']\n for obj_id in exps.keys():\n tp = exps[obj_id]['tp']\n obj_dir = os.path.join(video_root, video_name, obj_id)\n all_obj_list.append([obj_dir, video_name, obj_id, tp])\n\n print('start')\n cpu_num = mp.cpu_count()-1\n print(\"cpu_num:\", cpu_num)\n pool = mp.Pool(cpu_num)\n pbar = tqdm(total=len(all_obj_list))\n\n for obj in all_obj_list:\n pool.apply_async(\n generate, \n args = (obj, temp_xmem_anno, final_xmem_anno ),\n callback = lambda *a: pbar.update(1),\n error_callback = lambda e: print(colored(e, \"red\"))\n )\n \n pool.close()\n pool.join()\n pbar.close()\n\ndef inference(args):\n p = mp.Pool(8)\n for split_part in [0, 1, 2, 3]:\n for cfgs in [\" \", \" --reversed \"]:\n p.apply_async(\n run_eval,\n args=(args.llama_vid_meta, args.temp_xmem_anno, args.final_xmem_anno, args.img_dir, split_part, args.xmem_weight, cfgs),\n error_callback=lambda e: print(colored(e, \"red\"))\n )\n p.close()\n p.join()\n\n\"\"\"\npython XMem/tracking.py \\\n --video_root /mnt/public03/dataset/ovis/rgvos/visa7b/val_7b/revos_valid/Annotations \\\n --temp_xmem_anno /mnt/public03/dataset/ovis/rgvos/visa7b/val_7b/revos_valid/revos_valid_XMem_temp/Annotations \\\n --final_xmem_anno /mnt/public03/dataset/ovis/rgvos/visa7b/val_7b/revos_valid/revos_valid_XMem_final/Annotations \\\n --llama_vid_meta /mnt/public02/usr/yancilin/clyan_data/other-datasets/ReVOS/meta_expressions_valid__llamavid.json \\\n --img_dir /mnt/public02/usr/yancilin/clyan_data/other-datasets/ReVOS/JPEGImages \\\n --xmem_weight /mnt/public02/usr/yancilin/VISA/XMem/weights/XMem.pth\n\"\"\"\n\ndef main():\n parser = argparse.ArgumentParser(description='rgvos')\n parser.add_argument('--video_root', type=str, help='/PATH/TO/VISA_exp/revos_valid/Annotations', )\n parser.add_argument('--temp_xmem_anno', type=str, help='/PATH/TO/VISA_exp/revos_valid_XMem_temp/Annotations', ) # ไฟๅญ˜ๅ•ๅธง Mask ็š„่ทฏๅพ„\n parser.add_argument('--final_xmem_anno', type=str, help='/PATH/TO/VISA_exp/revos_valid_XMem_final/Annotations', ) # ไฟๅญ˜ XMem ๆœ€ๅŽ่พ“ๅ‡บ็ป“ๆžœ็š„่ทฏๅพ„\n parser.add_argument(\"--llama_vid_meta\", type=str, help='/PATH/TO/ReVOS/meta_expressions_valid__llamavid.json', )\n parser.add_argument(\"--img_dir\", type=str, help='/PATH/TO/ReVOS/JPEGImages')\n parser.add_argument(\"--xmem_weight\", type=str, help='/PATH/TO/XMEM_WEIGHT')\n args = parser.parse_args()\n\n prepare(args)\n inference(args)\n\n print('Done.')\n\n\nif __name__ == '__main__':\n main()\n" }, { "path": "XMem/train.py", "content": "import datetime\nfrom os import path\nimport math\nimport git\n\nimport random\nimport numpy as np\nimport torch\nfrom torch.utils.data import DataLoader, ConcatDataset\nimport torch.distributed as distributed\n\nfrom model.trainer import XMemTrainer\nfrom dataset.static_dataset import StaticTransformDataset\nfrom dataset.vos_dataset import VOSDataset\n\nfrom util.logger import TensorboardLogger\nfrom util.configuration import Configuration\nfrom util.load_subset import load_sub_davis, load_sub_yv\n\n\n\"\"\"\nInitial setup\n\"\"\"\n# Init distributed environment\ndistributed.init_process_group(backend=\"nccl\")\nprint(f'CUDA Device count: {torch.cuda.device_count()}')\n\n# Parse command line arguments\nraw_config = Configuration()\nraw_config.parse()\n\nif raw_config['benchmark']:\n torch.backends.cudnn.benchmark = True\n\n# Get current git info\nrepo = git.Repo(\".\")\ngit_info = str(repo.active_branch)+' '+str(repo.head.commit.hexsha)\n\nlocal_rank = torch.distributed.get_rank()\nworld_size = torch.distributed.get_world_size()\ntorch.cuda.set_device(local_rank)\n\nprint(f'I am rank {local_rank} in this world of size {world_size}!')\n\nnetwork_in_memory = None\nstages = raw_config['stages']\nstages_to_perform = list(stages)\nfor si, stage in enumerate(stages_to_perform):\n\n # Set seed to ensure the same initialization\n torch.manual_seed(14159265)\n np.random.seed(14159265)\n random.seed(14159265)\n\n # Pick stage specific hyperparameters out\n stage_config = raw_config.get_stage_parameters(stage)\n config = dict(**raw_config.args, **stage_config)\n if config['exp_id'] != 'NULL':\n config['exp_id'] = config['exp_id']+'_s%s'%stages[:si+1]\n\n config['single_object'] = (stage == '0')\n\n config['num_gpus'] = world_size\n if config['batch_size']//config['num_gpus']*config['num_gpus'] != config['batch_size']:\n raise ValueError('Batch size must be divisible by the number of GPUs.')\n config['batch_size'] //= config['num_gpus']\n config['num_workers'] //= config['num_gpus']\n print(f'We are assuming {config[\"num_gpus\"]} GPUs.')\n\n print(f'We are now starting stage {stage}')\n\n \"\"\"\n Model related\n \"\"\"\n if local_rank == 0:\n # Logging\n if config['exp_id'].lower() != 'null':\n print('I will take the role of logging!')\n long_id = '%s_%s' % (datetime.datetime.now().strftime('%b%d_%H.%M.%S'), config['exp_id'])\n else:\n long_id = None\n logger = TensorboardLogger(config['exp_id'], long_id, git_info)\n logger.log_string('hyperpara', str(config))\n\n # Construct the rank 0 model\n model = XMemTrainer(config, logger=logger, \n save_path=path.join('saves', long_id, long_id) if long_id is not None else None, \n local_rank=local_rank, world_size=world_size).train()\n else:\n # Construct model for other ranks\n model = XMemTrainer(config, local_rank=local_rank, world_size=world_size).train()\n\n # Load pertrained model if needed\n if raw_config['load_checkpoint'] is not None:\n total_iter = model.load_checkpoint(raw_config['load_checkpoint'])\n raw_config['load_checkpoint'] = None\n print('Previously trained model loaded!')\n else:\n total_iter = 0\n\n if network_in_memory is not None:\n print('I am loading network from the previous stage')\n model.load_network_in_memory(network_in_memory)\n network_in_memory = None\n elif raw_config['load_network'] is not None:\n print('I am loading network from a disk, as listed in configuration')\n model.load_network(raw_config['load_network'])\n raw_config['load_network'] = None\n\n \"\"\"\n Dataloader related\n \"\"\"\n # To re-seed the randomness everytime we start a worker\n def worker_init_fn(worker_id): \n worker_seed = torch.initial_seed()%(2**31) + worker_id + local_rank*100\n np.random.seed(worker_seed)\n random.seed(worker_seed)\n\n def construct_loader(dataset):\n train_sampler = torch.utils.data.distributed.DistributedSampler(dataset, rank=local_rank, shuffle=True)\n train_loader = DataLoader(dataset, config['batch_size'], sampler=train_sampler, num_workers=config['num_workers'],\n worker_init_fn=worker_init_fn, drop_last=True)\n return train_sampler, train_loader\n\n def renew_vos_loader(max_skip, finetune=False):\n # //5 because we only have annotation for every five frames\n yv_dataset = VOSDataset(path.join(yv_root, 'JPEGImages'), \n path.join(yv_root, 'Annotations'), max_skip//5, is_bl=False, subset=load_sub_yv(), num_frames=config['num_frames'], finetune=finetune)\n davis_dataset = VOSDataset(path.join(davis_root, 'JPEGImages', '480p'), \n path.join(davis_root, 'Annotations', '480p'), max_skip, is_bl=False, subset=load_sub_davis(), num_frames=config['num_frames'], finetune=finetune)\n train_dataset = ConcatDataset([davis_dataset]*5 + [yv_dataset])\n\n print(f'YouTube dataset size: {len(yv_dataset)}')\n print(f'DAVIS dataset size: {len(davis_dataset)}')\n print(f'Concat dataset size: {len(train_dataset)}')\n print(f'Renewed with {max_skip=}')\n\n return construct_loader(train_dataset)\n\n def renew_bl_loader(max_skip, finetune=False):\n train_dataset = VOSDataset(path.join(bl_root, 'JPEGImages'), \n path.join(bl_root, 'Annotations'), max_skip, is_bl=True, num_frames=config['num_frames'], finetune=finetune)\n\n print(f'Blender dataset size: {len(train_dataset)}')\n print(f'Renewed with {max_skip=}')\n\n return construct_loader(train_dataset)\n\n \"\"\"\n Dataset related\n \"\"\"\n\n \"\"\"\n These define the training schedule of the distance between frames\n We will switch to max_skip_values[i] once we pass the percentage specified by increase_skip_fraction[i]\n Not effective for stage 0 training\n The initial value is not listed here but in renew_vos_loader(X)\n \"\"\"\n max_skip_values = [10, 15, 5, 5]\n\n if stage == '0':\n static_root = path.expanduser(config['static_root'])\n # format: path, method (style of storing images), mutliplier\n train_dataset = StaticTransformDataset(\n [\n (path.join(static_root, 'fss'), 0, 1),\n (path.join(static_root, 'DUTS-TR'), 1, 1),\n (path.join(static_root, 'DUTS-TE'), 1, 1),\n (path.join(static_root, 'ecssd'), 1, 1),\n (path.join(static_root, 'BIG_small'), 1, 5),\n (path.join(static_root, 'HRSOD_small'), 1, 5),\n ], num_frames=config['num_frames'])\n train_sampler, train_loader = construct_loader(train_dataset)\n\n print(f'Static dataset size: {len(train_dataset)}')\n elif stage == '1':\n increase_skip_fraction = [0.1, 0.3, 0.8, 100]\n bl_root = path.join(path.expanduser(config['bl_root']))\n\n train_sampler, train_loader = renew_bl_loader(5)\n renew_loader = renew_bl_loader\n else:\n # stage 2 or 3\n increase_skip_fraction = [0.1, 0.3, 0.9, 100]\n # VOS dataset, 480p is used for both datasets\n yv_root = path.join(path.expanduser(config['yv_root']), 'train_480p')\n davis_root = path.join(path.expanduser(config['davis_root']), '2017', 'trainval')\n\n train_sampler, train_loader = renew_vos_loader(5)\n renew_loader = renew_vos_loader\n\n\n \"\"\"\n Determine max epoch\n \"\"\"\n total_epoch = math.ceil(config['iterations']/len(train_loader))\n current_epoch = total_iter // len(train_loader)\n print(f'We approximately use {total_epoch} epochs.')\n if stage != '0':\n change_skip_iter = [round(config['iterations']*f) for f in increase_skip_fraction]\n # Skip will only change after an epoch, not in the middle\n print(f'The skip value will change approximately at the following iterations: {change_skip_iter[:-1]}')\n\n \"\"\"\n Starts training\n \"\"\"\n finetuning = False\n # Need this to select random bases in different workers\n np.random.seed(np.random.randint(2**30-1) + local_rank*100)\n try:\n while total_iter < config['iterations'] + config['finetune']:\n \n # Crucial for randomness! \n train_sampler.set_epoch(current_epoch)\n current_epoch += 1\n print(f'Current epoch: {current_epoch}')\n\n # Train loop\n model.train()\n for data in train_loader:\n # Update skip if needed\n if stage!='0' and total_iter >= change_skip_iter[0]:\n while total_iter >= change_skip_iter[0]:\n cur_skip = max_skip_values[0]\n max_skip_values = max_skip_values[1:]\n change_skip_iter = change_skip_iter[1:]\n print(f'Changing skip to {cur_skip=}')\n train_sampler, train_loader = renew_loader(cur_skip)\n break\n\n # fine-tune means fewer augmentations to train the sensory memory\n if config['finetune'] > 0 and not finetuning and total_iter >= config['iterations']:\n train_sampler, train_loader = renew_loader(cur_skip, finetune=True)\n finetuning = True\n model.save_network_interval = 1000\n break\n\n model.do_pass(data, total_iter)\n total_iter += 1\n\n if total_iter >= config['iterations'] + config['finetune']:\n break\n finally:\n if not config['debug'] and model.logger is not None and total_iter>5000:\n model.save_network(total_iter)\n model.save_checkpoint(total_iter)\n\n network_in_memory = model.XMem.module.state_dict()\n\ndistributed.destroy_process_group()\n" }, { "path": "XMem/util/__init__.py", "content": "" }, { "path": "XMem/util/configuration.py", "content": "from argparse import ArgumentParser\n\n\ndef none_or_default(x, default):\n return x if x is not None else default\n\nclass Configuration():\n def parse(self, unknown_arg_ok=False):\n parser = ArgumentParser()\n\n # Enable torch.backends.cudnn.benchmark -- Faster in some cases, test in your own environment\n parser.add_argument('--benchmark', action='store_true')\n parser.add_argument('--no_amp', action='store_true')\n\n # Data parameters\n parser.add_argument('--static_root', help='Static training data root', default='../static')\n parser.add_argument('--bl_root', help='Blender training data root', default='../BL30K')\n parser.add_argument('--yv_root', help='YouTubeVOS data root', default='../YouTube')\n parser.add_argument('--davis_root', help='DAVIS data root', default='../DAVIS')\n parser.add_argument('--num_workers', help='Total number of dataloader workers across all GPUs processes', type=int, default=16)\n\n parser.add_argument('--key_dim', default=64, type=int)\n parser.add_argument('--value_dim', default=512, type=int)\n parser.add_argument('--hidden_dim', default=64, help='Set to =0 to disable', type=int)\n\n parser.add_argument('--deep_update_prob', default=0.2, type=float)\n\n parser.add_argument('--stages', help='Training stage (0-static images, 1-Blender dataset, 2-DAVIS+YouTubeVOS)', default='02')\n\n \"\"\"\n Stage-specific learning parameters\n Batch sizes are effective -- you don't have to scale them when you scale the number processes\n \"\"\"\n # Stage 0, static images\n parser.add_argument('--s0_batch_size', default=16, type=int)\n parser.add_argument('--s0_iterations', default=150000, type=int)\n parser.add_argument('--s0_finetune', default=0, type=int)\n parser.add_argument('--s0_steps', nargs=\"*\", default=[], type=int)\n parser.add_argument('--s0_lr', help='Initial learning rate', default=1e-5, type=float)\n parser.add_argument('--s0_num_ref_frames', default=2, type=int)\n parser.add_argument('--s0_num_frames', default=3, type=int)\n parser.add_argument('--s0_start_warm', default=20000, type=int)\n parser.add_argument('--s0_end_warm', default=70000, type=int)\n\n # Stage 1, BL30K\n parser.add_argument('--s1_batch_size', default=8, type=int)\n parser.add_argument('--s1_iterations', default=250000, type=int)\n # fine-tune means fewer augmentations to train the sensory memory\n parser.add_argument('--s1_finetune', default=0, type=int)\n parser.add_argument('--s1_steps', nargs=\"*\", default=[200000], type=int)\n parser.add_argument('--s1_lr', help='Initial learning rate', default=1e-5, type=float)\n parser.add_argument('--s1_num_ref_frames', default=3, type=int)\n parser.add_argument('--s1_num_frames', default=8, type=int)\n parser.add_argument('--s1_start_warm', default=20000, type=int)\n parser.add_argument('--s1_end_warm', default=70000, type=int)\n\n # Stage 2, DAVIS+YoutubeVOS, longer\n parser.add_argument('--s2_batch_size', default=8, type=int)\n parser.add_argument('--s2_iterations', default=150000, type=int)\n # fine-tune means fewer augmentations to train the sensory memory\n parser.add_argument('--s2_finetune', default=10000, type=int)\n parser.add_argument('--s2_steps', nargs=\"*\", default=[120000], type=int)\n parser.add_argument('--s2_lr', help='Initial learning rate', default=1e-5, type=float)\n parser.add_argument('--s2_num_ref_frames', default=3, type=int)\n parser.add_argument('--s2_num_frames', default=8, type=int)\n parser.add_argument('--s2_start_warm', default=20000, type=int)\n parser.add_argument('--s2_end_warm', default=70000, type=int)\n\n # Stage 3, DAVIS+YoutubeVOS, shorter\n parser.add_argument('--s3_batch_size', default=8, type=int)\n parser.add_argument('--s3_iterations', default=100000, type=int)\n # fine-tune means fewer augmentations to train the sensory memory\n parser.add_argument('--s3_finetune', default=10000, type=int)\n parser.add_argument('--s3_steps', nargs=\"*\", default=[80000], type=int)\n parser.add_argument('--s3_lr', help='Initial learning rate', default=1e-5, type=float)\n parser.add_argument('--s3_num_ref_frames', default=3, type=int)\n parser.add_argument('--s3_num_frames', default=8, type=int)\n parser.add_argument('--s3_start_warm', default=20000, type=int)\n parser.add_argument('--s3_end_warm', default=70000, type=int)\n\n parser.add_argument('--gamma', help='LR := LR*gamma at every decay step', default=0.1, type=float)\n parser.add_argument('--weight_decay', default=0.05, type=float)\n\n # Loading\n parser.add_argument('--load_network', help='Path to pretrained network weight only')\n parser.add_argument('--load_checkpoint', help='Path to the checkpoint file, including network, optimizer and such')\n\n # Logging information\n parser.add_argument('--log_text_interval', default=100, type=int)\n parser.add_argument('--log_image_interval', default=1000, type=int)\n parser.add_argument('--save_network_interval', default=25000, type=int)\n parser.add_argument('--save_checkpoint_interval', default=50000, type=int)\n parser.add_argument('--exp_id', help='Experiment UNIQUE id, use NULL to disable logging to tensorboard', default='NULL')\n parser.add_argument('--debug', help='Debug mode which logs information more often', action='store_true')\n\n # # Multiprocessing parameters, not set by users\n # parser.add_argument('--local_rank', default=0, type=int, help='Local rank of this process')\n\n if unknown_arg_ok:\n args, _ = parser.parse_known_args()\n self.args = vars(args)\n else:\n self.args = vars(parser.parse_args())\n\n self.args['amp'] = not self.args['no_amp']\n\n # check if the stages are valid\n stage_to_perform = list(self.args['stages'])\n for s in stage_to_perform:\n if s not in ['0', '1', '2', '3']:\n raise NotImplementedError\n\n def get_stage_parameters(self, stage):\n parameters = {\n 'batch_size': self.args['s%s_batch_size'%stage],\n 'iterations': self.args['s%s_iterations'%stage],\n 'finetune': self.args['s%s_finetune'%stage],\n 'steps': self.args['s%s_steps'%stage],\n 'lr': self.args['s%s_lr'%stage],\n 'num_ref_frames': self.args['s%s_num_ref_frames'%stage],\n 'num_frames': self.args['s%s_num_frames'%stage],\n 'start_warm': self.args['s%s_start_warm'%stage],\n 'end_warm': self.args['s%s_end_warm'%stage],\n }\n\n return parameters\n\n def __getitem__(self, key):\n return self.args[key]\n\n def __setitem__(self, key, value):\n self.args[key] = value\n\n def __str__(self):\n return str(self.args)\n" }, { "path": "XMem/util/davis_subset.txt", "content": "bear\nbmx-bumps\nboat\nboxing-fisheye\nbreakdance-flare\nbus\ncar-turn\ncat-girl\nclassic-car\ncolor-run\ncrossing\ndance-jump\ndancing\ndisc-jockey\ndog-agility\ndog-gooses\ndogs-scale\ndrift-turn\ndrone\nelephant\nflamingo\nhike\nhockey\nhorsejump-low\nkid-football\nkite-walk\nkoala\nlady-running\nlindy-hop\nlongboard\nlucia\nmallard-fly\nmallard-water\nmiami-surf\nmotocross-bumps\nmotorbike\nnight-race\nparagliding\nplanes-water\nrallye\nrhino\nrollerblade\nschoolgirls\nscooter-board\nscooter-gray\nsheep\nskate-park\nsnowboard\nsoccerball\nstroller\nstunt\nsurf\nswing\ntennis\ntractor-sand\ntrain\ntuk-tuk\nupside-down\nvaranus-cage\nwalking" }, { "path": "XMem/util/image_saver.py", "content": "import cv2\nimport numpy as np\n\nimport torch\nfrom dataset.range_transform import inv_im_trans\nfrom collections import defaultdict\n\ndef tensor_to_numpy(image):\n image_np = (image.numpy() * 255).astype('uint8')\n return image_np\n\ndef tensor_to_np_float(image):\n image_np = image.numpy().astype('float32')\n return image_np\n\ndef detach_to_cpu(x):\n return x.detach().cpu()\n\ndef transpose_np(x):\n return np.transpose(x, [1,2,0])\n\ndef tensor_to_gray_im(x):\n x = detach_to_cpu(x)\n x = tensor_to_numpy(x)\n x = transpose_np(x)\n return x\n\ndef tensor_to_im(x):\n x = detach_to_cpu(x)\n x = inv_im_trans(x).clamp(0, 1)\n x = tensor_to_numpy(x)\n x = transpose_np(x)\n return x\n\n# Predefined key <-> caption dict\nkey_captions = {\n 'im': 'Image', \n 'gt': 'GT', \n}\n\n\"\"\"\nReturn an image array with captions\nkeys in dictionary will be used as caption if not provided\nvalues should contain lists of cv2 images\n\"\"\"\ndef get_image_array(images, grid_shape, captions={}):\n h, w = grid_shape\n cate_counts = len(images)\n rows_counts = len(next(iter(images.values())))\n\n font = cv2.FONT_HERSHEY_SIMPLEX\n\n output_image = np.zeros([w*cate_counts, h*(rows_counts+1), 3], dtype=np.uint8)\n col_cnt = 0\n for k, v in images.items():\n\n # Default as key value itself\n caption = captions.get(k, k)\n\n # Handles new line character\n dy = 40\n for i, line in enumerate(caption.split('\\n')):\n cv2.putText(output_image, line, (10, col_cnt*w+100+i*dy),\n font, 0.8, (255,255,255), 2, cv2.LINE_AA)\n\n # Put images\n for row_cnt, img in enumerate(v):\n im_shape = img.shape\n if len(im_shape) == 2:\n img = img[..., np.newaxis]\n\n img = (img * 255).astype('uint8')\n\n output_image[(col_cnt+0)*w:(col_cnt+1)*w,\n (row_cnt+1)*h:(row_cnt+2)*h, :] = img\n \n col_cnt += 1\n\n return output_image\n\ndef base_transform(im, size):\n im = tensor_to_np_float(im)\n if len(im.shape) == 3:\n im = im.transpose((1, 2, 0))\n else:\n im = im[:, :, None]\n\n # Resize\n if im.shape[1] != size:\n im = cv2.resize(im, size, interpolation=cv2.INTER_NEAREST)\n\n return im.clip(0, 1)\n\ndef im_transform(im, size):\n return base_transform(inv_im_trans(detach_to_cpu(im)), size=size)\n\ndef mask_transform(mask, size):\n return base_transform(detach_to_cpu(mask), size=size)\n\ndef out_transform(mask, size):\n return base_transform(detach_to_cpu(torch.sigmoid(mask)), size=size)\n\ndef pool_pairs(images, size, num_objects):\n req_images = defaultdict(list)\n\n b, t = images['rgb'].shape[:2]\n\n # limit the number of images saved\n b = min(2, b)\n\n # find max num objects\n max_num_objects = max(num_objects[:b])\n\n GT_suffix = ''\n for bi in range(b):\n GT_suffix += ' \\n%s' % images['info']['name'][bi][-25:-4]\n\n for bi in range(b):\n for ti in range(t):\n req_images['RGB'].append(im_transform(images['rgb'][bi,ti], size))\n for oi in range(max_num_objects):\n if ti == 0 or oi >= num_objects[bi]:\n req_images['Mask_%d'%oi].append(mask_transform(images['first_frame_gt'][bi][0,oi], size))\n # req_images['Mask_X8_%d'%oi].append(mask_transform(images['first_frame_gt'][bi][0,oi], size))\n # req_images['Mask_X16_%d'%oi].append(mask_transform(images['first_frame_gt'][bi][0,oi], size))\n else:\n req_images['Mask_%d'%oi].append(mask_transform(images['masks_%d'%ti][bi][oi], size))\n # req_images['Mask_%d'%oi].append(mask_transform(images['masks_%d'%ti][bi][oi][2], size))\n # req_images['Mask_X8_%d'%oi].append(mask_transform(images['masks_%d'%ti][bi][oi][1], size))\n # req_images['Mask_X16_%d'%oi].append(mask_transform(images['masks_%d'%ti][bi][oi][0], size))\n req_images['GT_%d_%s'%(oi, GT_suffix)].append(mask_transform(images['cls_gt'][bi,ti,0]==(oi+1), size))\n # print((images['cls_gt'][bi,ti,0]==(oi+1)).shape)\n # print(mask_transform(images['cls_gt'][bi,ti,0]==(oi+1), size).shape)\n\n\n return get_image_array(req_images, size, key_captions)" }, { "path": "XMem/util/load_subset.py", "content": "\"\"\"\nload_subset.py - Presents a subset of data\nDAVIS - only the training set\nYouTubeVOS - I manually filtered some erroneous ones out but I haven't checked all\n\"\"\"\n\n\ndef load_sub_davis(path='util/davis_subset.txt'):\n with open(path, mode='r') as f:\n subset = set(f.read().splitlines())\n return subset\n\ndef load_sub_yv(path='util/yv_subset.txt'):\n with open(path, mode='r') as f:\n subset = set(f.read().splitlines())\n return subset\n" }, { "path": "XMem/util/log_integrator.py", "content": "\"\"\"\nIntegrate numerical values for some iterations\nTypically used for loss computation / logging to tensorboard\nCall finalize and create a new Integrator when you want to display/log\n\"\"\"\n\nimport torch\n\n\nclass Integrator:\n def __init__(self, logger, distributed=True, local_rank=0, world_size=1):\n self.values = {}\n self.counts = {}\n self.hooks = [] # List is used here to maintain insertion order\n\n self.logger = logger\n\n self.distributed = distributed\n self.local_rank = local_rank\n self.world_size = world_size\n\n def add_tensor(self, key, tensor):\n if key not in self.values:\n self.counts[key] = 1\n if type(tensor) == float or type(tensor) == int:\n self.values[key] = tensor\n else:\n self.values[key] = tensor.mean().item()\n else:\n self.counts[key] += 1\n if type(tensor) == float or type(tensor) == int:\n self.values[key] += tensor\n else:\n self.values[key] += tensor.mean().item()\n\n def add_dict(self, tensor_dict):\n for k, v in tensor_dict.items():\n self.add_tensor(k, v)\n\n def add_hook(self, hook):\n \"\"\"\n Adds a custom hook, i.e. compute new metrics using values in the dict\n The hook takes the dict as argument, and returns a (k, v) tuple\n e.g. for computing IoU\n \"\"\"\n if type(hook) == list:\n self.hooks.extend(hook)\n else:\n self.hooks.append(hook)\n\n def reset_except_hooks(self):\n self.values = {}\n self.counts = {}\n\n # Average and output the metrics\n def finalize(self, prefix, it, f=None):\n\n for hook in self.hooks:\n k, v = hook(self.values)\n self.add_tensor(k, v)\n\n for k, v in self.values.items():\n\n if k[:4] == 'hide':\n continue\n\n avg = v / self.counts[k]\n\n if self.distributed:\n # Inplace operation\n avg = torch.tensor(avg).cuda()\n torch.distributed.reduce(avg, dst=0)\n\n if self.local_rank == 0:\n avg = (avg/self.world_size).cpu().item()\n self.logger.log_metrics(prefix, k, avg, it, f)\n else:\n # Simple does it\n self.logger.log_metrics(prefix, k, avg, it, f)\n\n" }, { "path": "XMem/util/logger.py", "content": "\"\"\"\nDumps things to tensorboard and console\n\"\"\"\n\nimport os\nimport warnings\n\nimport torchvision.transforms as transforms\nfrom torch.utils.tensorboard import SummaryWriter\n\n\ndef tensor_to_numpy(image):\n image_np = (image.numpy() * 255).astype('uint8')\n return image_np\n\ndef detach_to_cpu(x):\n return x.detach().cpu()\n\ndef fix_width_trunc(x):\n return ('{:.9s}'.format('{:0.9f}'.format(x)))\n\nclass TensorboardLogger:\n def __init__(self, short_id, id, git_info):\n self.short_id = short_id\n if self.short_id == 'NULL':\n self.short_id = 'DEBUG'\n\n if id is None:\n self.no_log = True\n warnings.warn('Logging has been disbaled.')\n else:\n self.no_log = False\n\n self.inv_im_trans = transforms.Normalize(\n mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225],\n std=[1/0.229, 1/0.224, 1/0.225])\n\n self.inv_seg_trans = transforms.Normalize(\n mean=[-0.5/0.5],\n std=[1/0.5])\n\n log_path = os.path.join('.', 'saves', '%s' % id)\n self.logger = SummaryWriter(log_path)\n\n self.log_string('git', git_info)\n\n def log_scalar(self, tag, x, step):\n if self.no_log:\n warnings.warn('Logging has been disabled.')\n return\n self.logger.add_scalar(tag, x, step)\n\n def log_metrics(self, l1_tag, l2_tag, val, step, f=None):\n tag = l1_tag + '/' + l2_tag\n text = '{:s} - It {:6d} [{:5s}] [{:13}]: {:s}'.format(self.short_id, step, l1_tag.upper(), l2_tag, fix_width_trunc(val))\n print(text)\n if f is not None:\n f.write(text + '\\n')\n f.flush()\n self.log_scalar(tag, val, step)\n\n def log_im(self, tag, x, step):\n if self.no_log:\n warnings.warn('Logging has been disabled.')\n return\n x = detach_to_cpu(x)\n x = self.inv_im_trans(x)\n x = tensor_to_numpy(x)\n self.logger.add_image(tag, x, step)\n\n def log_cv2(self, tag, x, step):\n if self.no_log:\n warnings.warn('Logging has been disabled.')\n return\n x = x.transpose((2, 0, 1))\n self.logger.add_image(tag, x, step)\n\n def log_seg(self, tag, x, step):\n if self.no_log:\n warnings.warn('Logging has been disabled.')\n return\n x = detach_to_cpu(x)\n x = self.inv_seg_trans(x)\n x = tensor_to_numpy(x)\n self.logger.add_image(tag, x, step)\n\n def log_gray(self, tag, x, step):\n if self.no_log:\n warnings.warn('Logging has been disabled.')\n return\n x = detach_to_cpu(x)\n x = tensor_to_numpy(x)\n self.logger.add_image(tag, x, step)\n\n def log_string(self, tag, x):\n print(tag, x)\n if self.no_log:\n warnings.warn('Logging has been disabled.')\n return\n self.logger.add_text(tag, x)\n " }, { "path": "XMem/util/palette.py", "content": "davis_palette = b'\\x00\\x00\\x00\\x80\\x00\\x00\\x00\\x80\\x00\\x80\\x80\\x00\\x00\\x00\\x80\\x80\\x00\\x80\\x00\\x80\\x80\\x80\\x80\\x80@\\x00\\x00\\xc0\\x00\\x00@\\x80\\x00\\xc0\\x80\\x00@\\x00\\x80\\xc0\\x00\\x80@\\x80\\x80\\xc0\\x80\\x80\\x00@\\x00\\x80@\\x00\\x00\\xc0\\x00\\x80\\xc0\\x00\\x00@\\x80\\x80@\\x80\\x00\\xc0\\x80\\x80\\xc0\\x80@@\\x00\\xc0@\\x00@\\xc0\\x00\\xc0\\xc0\\x00@@\\x80\\xc0@\\x80@\\xc0\\x80\\xc0\\xc0\\x80\\x00\\x00@\\x80\\x00@\\x00\\x80@\\x80\\x80@\\x00\\x00\\xc0\\x80\\x00\\xc0\\x00\\x80\\xc0\\x80\\x80\\xc0@\\x00@\\xc0\\x00@@\\x80@\\xc0\\x80@@\\x00\\xc0\\xc0\\x00\\xc0@\\x80\\xc0\\xc0\\x80\\xc0\\x00@@\\x80@@\\x00\\xc0@\\x80\\xc0@\\x00@\\xc0\\x80@\\xc0\\x00\\xc0\\xc0\\x80\\xc0\\xc0@@@\\xc0@@@\\xc0@\\xc0\\xc0@@@\\xc0\\xc0@\\xc0@\\xc0\\xc0\\xc0\\xc0\\xc0 \\x00\\x00\\xa0\\x00\\x00 \\x80\\x00\\xa0\\x80\\x00 \\x00\\x80\\xa0\\x00\\x80 \\x80\\x80\\xa0\\x80\\x80`\\x00\\x00\\xe0\\x00\\x00`\\x80\\x00\\xe0\\x80\\x00`\\x00\\x80\\xe0\\x00\\x80`\\x80\\x80\\xe0\\x80\\x80 @\\x00\\xa0@\\x00 \\xc0\\x00\\xa0\\xc0\\x00 @\\x80\\xa0@\\x80 \\xc0\\x80\\xa0\\xc0\\x80`@\\x00\\xe0@\\x00`\\xc0\\x00\\xe0\\xc0\\x00`@\\x80\\xe0@\\x80`\\xc0\\x80\\xe0\\xc0\\x80 \\x00@\\xa0\\x00@ \\x80@\\xa0\\x80@ \\x00\\xc0\\xa0\\x00\\xc0 \\x80\\xc0\\xa0\\x80\\xc0`\\x00@\\xe0\\x00@`\\x80@\\xe0\\x80@`\\x00\\xc0\\xe0\\x00\\xc0`\\x80\\xc0\\xe0\\x80\\xc0 @@\\xa0@@ \\xc0@\\xa0\\xc0@ @\\xc0\\xa0@\\xc0 \\xc0\\xc0\\xa0\\xc0\\xc0`@@\\xe0@@`\\xc0@\\xe0\\xc0@`@\\xc0\\xe0@\\xc0`\\xc0\\xc0\\xe0\\xc0\\xc0\\x00 \\x00\\x80 \\x00\\x00\\xa0\\x00\\x80\\xa0\\x00\\x00 \\x80\\x80 \\x80\\x00\\xa0\\x80\\x80\\xa0\\x80@ \\x00\\xc0 \\x00@\\xa0\\x00\\xc0\\xa0\\x00@ \\x80\\xc0 \\x80@\\xa0\\x80\\xc0\\xa0\\x80\\x00`\\x00\\x80`\\x00\\x00\\xe0\\x00\\x80\\xe0\\x00\\x00`\\x80\\x80`\\x80\\x00\\xe0\\x80\\x80\\xe0\\x80@`\\x00\\xc0`\\x00@\\xe0\\x00\\xc0\\xe0\\x00@`\\x80\\xc0`\\x80@\\xe0\\x80\\xc0\\xe0\\x80\\x00 @\\x80 @\\x00\\xa0@\\x80\\xa0@\\x00 \\xc0\\x80 \\xc0\\x00\\xa0\\xc0\\x80\\xa0\\xc0@ @\\xc0 @@\\xa0@\\xc0\\xa0@@ \\xc0\\xc0 \\xc0@\\xa0\\xc0\\xc0\\xa0\\xc0\\x00`@\\x80`@\\x00\\xe0@\\x80\\xe0@\\x00`\\xc0\\x80`\\xc0\\x00\\xe0\\xc0\\x80\\xe0\\xc0@`@\\xc0`@@\\xe0@\\xc0\\xe0@@`\\xc0\\xc0`\\xc0@\\xe0\\xc0\\xc0\\xe0\\xc0 \\x00\\xa0 \\x00 \\xa0\\x00\\xa0\\xa0\\x00 \\x80\\xa0 \\x80 \\xa0\\x80\\xa0\\xa0\\x80` \\x00\\xe0 \\x00`\\xa0\\x00\\xe0\\xa0\\x00` \\x80\\xe0 \\x80`\\xa0\\x80\\xe0\\xa0\\x80 `\\x00\\xa0`\\x00 \\xe0\\x00\\xa0\\xe0\\x00 `\\x80\\xa0`\\x80 \\xe0\\x80\\xa0\\xe0\\x80``\\x00\\xe0`\\x00`\\xe0\\x00\\xe0\\xe0\\x00``\\x80\\xe0`\\x80`\\xe0\\x80\\xe0\\xe0\\x80 @\\xa0 @ \\xa0@\\xa0\\xa0@ \\xc0\\xa0 \\xc0 \\xa0\\xc0\\xa0\\xa0\\xc0` @\\xe0 @`\\xa0@\\xe0\\xa0@` \\xc0\\xe0 \\xc0`\\xa0\\xc0\\xe0\\xa0\\xc0 `@\\xa0`@ \\xe0@\\xa0\\xe0@ `\\xc0\\xa0`\\xc0 \\xe0\\xc0\\xa0\\xe0\\xc0``@\\xe0`@`\\xe0@\\xe0\\xe0@``\\xc0\\xe0`\\xc0`\\xe0\\xc0\\xe0\\xe0\\xc0'\n\nyoutube_palette = b'\\x00\\x00\\x00\\xec_g\\xf9\\x91W\\xfa\\xc8c\\x99\\xc7\\x94b\\xb3\\xb2f\\x99\\xcc\\xc5\\x94\\xc5\\xabyg\\xff\\xff\\xffes~\\x0b\\x0b\\x0b\\x0c\\x0c\\x0c\\r\\r\\r\\x0e\\x0e\\x0e\\x0f\\x0f\\x0f'\n" }, { "path": "XMem/util/tensor_util.py", "content": "import torch.nn.functional as F\n\n\ndef compute_tensor_iu(seg, gt):\n intersection = (seg & gt).float().sum()\n union = (seg | gt).float().sum()\n\n return intersection, union\n\ndef compute_tensor_iou(seg, gt):\n intersection, union = compute_tensor_iu(seg, gt)\n iou = (intersection + 1e-6) / (union + 1e-6)\n \n return iou \n\n# STM\ndef pad_divide_by(in_img, d):\n h, w = in_img.shape[-2:]\n\n if h % d > 0:\n new_h = h + d - h % d\n else:\n new_h = h\n if w % d > 0:\n new_w = w + d - w % d\n else:\n new_w = w\n lh, uh = int((new_h-h) / 2), int(new_h-h) - int((new_h-h) / 2)\n lw, uw = int((new_w-w) / 2), int(new_w-w) - int((new_w-w) / 2)\n pad_array = (int(lw), int(uw), int(lh), int(uh))\n out = F.pad(in_img, pad_array)\n return out, pad_array\n\ndef unpad(img, pad):\n if len(img.shape) == 4:\n if pad[2]+pad[3] > 0:\n img = img[:,:,pad[2]:-pad[3],:]\n if pad[0]+pad[1] > 0:\n img = img[:,:,:,pad[0]:-pad[1]]\n elif len(img.shape) == 3:\n if pad[2]+pad[3] > 0:\n img = img[:,pad[2]:-pad[3],:]\n if pad[0]+pad[1] > 0:\n img = img[:,:,pad[0]:-pad[1]]\n else:\n raise NotImplementedError\n return img" }, { "path": "XMem/util/yv_subset.txt", "content": "003234408d\n0043f083b5\n0044fa5fba\n005a527edd\n0065b171f9\n00917dcfc4\n00a23ccf53\n00ad5016a4\n01082ae388\n011ac0a06f\n013099c098\n0155498c85\n01694ad9c8\n017ac35701\n01b80e8e1a\n01baa5a4e1\n01c3111683\n01c4cb5ffe\n01c76f0a82\n01c783268c\n01ed275c6e\n01ff60d1fa\n020cd28cd2\n02264db755\n0248626d9a\n02668dbffa\n0274193026\n02d28375aa\n02f3a5c4df\n031ccc99b1\n0321b18c10\n0348a45bca\n0355e92655\n0358b938c1\n0368107cf1\n0379ddf557\n038b2cc71d\n038c15a5dd\n03a06cc98a\n03a63e187f\n03c95b4dae\n03e2b57b0e\n04194e1248\n0444918a5f\n04460a7a52\n04474174a4\n0450095513\n045f00aed2\n04667fabaa\n04735c5030\n04990d1915\n04d62d9d98\n04f21da964\n04fbad476e\n04fe256562\n0503bf89c9\n0536c9eed0\n054acb238f\n05579ca250\n056c200404\n05774f3a2c\n058a7592c8\n05a0a513df\n05a569d8aa\n05aa652648\n05d7715782\n05e0b0f28f\n05fdbbdd7a\n05ffcfed85\n0630391881\n06840b2bbe\n068f7dce6f\n0693719753\n06ce2b51fb\n06e224798e\n06ee361788\n06fbb3fa2c\n0700264286\n070c918ca7\n07129e14a4\n07177017e9\n07238ffc58\n07353b2a89\n0738493cbf\n075926c651\n075c701292\n0762ea9a30\n07652ee4af\n076f206928\n077d32af19\n079049275c\n07913cdda7\n07a11a35e8\n07ac33b6df\n07b6e8fda8\n07c62c3d11\n07cc1c7d74\n080196ef01\n081207976e\n081ae4fa44\n081d8250cb\n082900c5d4\n0860df21e2\n0866d4c5e3\n0891ac2eb6\n08931bc458\n08aa2705d5\n08c8450db7\n08d50b926c\n08e1e4de15\n08e48c1a48\n08f561c65e\n08feb87790\n09049f6fe3\n092e4ff450\n09338adea8\n093c335ccc\n0970d28339\n0974a213dc\n097b471ed8\n0990941758\n09a348f4fa\n09a6841288\n09c5bad17b\n09c9ce80c7\n09ff54fef4\n0a23765d15\n0a275e7f12\n0a2f2bd294\n0a7a2514aa\n0a7b27fde9\n0a8c467cc3\n0ac8c560ae\n0b1627e896\n0b285c47f6\n0b34ec1d55\n0b5b5e8e5a\n0b68535614\n0b6f9105fc\n0b7dbfa3cb\n0b9cea51ca\n0b9d012be8\n0bcfc4177d\n0bd37b23c1\n0bd864064c\n0c11c6bf7b\n0c26bc77ac\n0c3a04798c\n0c44a9d545\n0c817cc390\n0ca839ee9a\n0cd7ac0ac0\n0ce06e0121\n0cfe974a89\n0d2fcc0dcd\n0d3aad05d2\n0d40b015f4\n0d97fba242\n0d9cc80d7e\n0dab85b6d3\n0db5c427a5\n0dbaf284f1\n0de4923598\n0df28a9101\n0e04f636c4\n0e05f0e232\n0e0930474b\n0e27472bea\n0e30020549\n0e621feb6c\n0e803c7d73\n0e9ebe4e3c\n0e9f2785ec\n0ea68d418b\n0eb403a222\n0ee92053d6\n0eefca067f\n0f17fa6fcb\n0f1ac8e9a3\n0f202e9852\n0f2ab8b1ff\n0f51a78756\n0f5fbe16b0\n0f6072077b\n0f6b69b2f4\n0f6c2163de\n0f74ec5599\n0f9683715b\n0fa7b59356\n0fb173695b\n0fc958cde2\n0fe7b1a621\n0ffcdb491c\n101caff7d4\n1022fe8417\n1032e80b37\n103f501680\n104e64565f\n104f1ab997\n106242403f\n10b31f5431\n10eced835e\n110d26fa3a\n1122c1d16a\n1145b49a5f\n11485838c2\n114e7676ec\n1157472b95\n115ee1072c\n1171141012\n117757b4b8\n1178932d2f\n117cc76bda\n1180cbf814\n1187bbd0e3\n1197e44b26\n119cf20728\n119dd54871\n11a0c3b724\n11a6ba8c94\n11c722a456\n11cbcb0b4d\n11ccf5e99d\n11ce6f452e\n11e53de6f2\n11feabe596\n120cb9514d\n12156b25b3\n122896672d\n1232b2f1d4\n1233ac8596\n1239c87234\n1250423f7c\n1257a1bc67\n125d1b19dd\n126d203967\n1295e19071\n12ad198c54\n12bddb2bcb\n12ec9b93ee\n12eebedc35\n132852e094\n1329409f2a\n13325cfa14\n134d06dbf9\n135625b53d\n13870016f9\n13960b3c84\n13adaad9d9\n13ae097e20\n13e3070469\n13f6a8c20d\n1416925cf2\n142d2621f5\n145d5d7c03\n145fdc3ac5\n1471274fa7\n14a6b5a139\n14c21cea0d\n14dae0dc93\n14f9bd22b5\n14fd28ae99\n15097d5d4e\n150ea711f2\n1514e3563f\n152aaa3a9e\n152b7d3bd7\n15617297cc\n15abbe0c52\n15d1fb3de5\n15f67b0fab\n161eb59aad\n16288ea47f\n164410ce62\n165c3c8cd4\n165c42b41b\n165ec9e22b\n1669502269\n16763cccbb\n16adde065e\n16af445362\n16afd538ad\n16c3fa4d5d\n16d1d65c27\n16e8599e94\n16fe9fb444\n1705796b02\n1724db7671\n17418e81ea\n175169edbb\n17622326fd\n17656bae77\n17b0d94172\n17c220e4f6\n17c7bcd146\n17cb4afe89\n17cd79a434\n17d18604c3\n17d8ca1a37\n17e33f4330\n17f7a6d805\n180abc8378\n183ba3d652\n185bf64702\n18913cc690\n1892651815\n189ac8208a\n189b44e92c\n18ac264b76\n18b245ab49\n18b5cebc34\n18bad52083\n18bb5144d5\n18c6f205c5\n1903f9ea15\n1917b209f2\n191e74c01d\n19367bb94e\n193ffaa217\n19696b67d3\n197f3ab6f3\n1981e763cc\n198afe39ae\n19a6e62b9b\n19b60d5335\n19c00c11f9\n19e061eb88\n19e8bc6178\n19ee80dac6\n1a25a9170a\n1a359a6c1a\n1a3e87c566\n1a5fe06b00\n1a6c0fbd1e\n1a6f3b5a4b\n1a8afbad92\n1a8bdc5842\n1a95752aca\n1a9c131cb7\n1aa3da3ee3\n1ab27ec7ea\n1abf16d21d\n1acd0f993b\n1ad202e499\n1af8d2395d\n1afd39a1fa\n1b2d31306f\n1b3fa67f0e\n1b43fa74b4\n1b73ea9fc2\n1b7e8bb255\n1b8680f8cd\n1b883843c0\n1b8898785b\n1b88ba1aa4\n1b96a498e5\n1bbc4c274f\n1bd87fe9ab\n1c4090c75b\n1c41934f84\n1c72b04b56\n1c87955a3a\n1c9f9eb792\n1ca240fede\n1ca5673803\n1cada35274\n1cb44b920d\n1cd10e62be\n1d3087d5e5\n1d3685150a\n1d6ff083aa\n1d746352a6\n1da256d146\n1da4e956b1\n1daf812218\n1dba687bce\n1dce57d05d\n1de4a9e537\n1dec5446c8\n1dfbe6f586\n1e1a18c45a\n1e1e42529d\n1e4be70796\n1eb60959c8\n1ec8b2566b\n1ecdc2941c\n1ee0ac70ff\n1ef8e17def\n1f1a2a9fc0\n1f1beb8daa\n1f2609ee13\n1f3876f8d0\n1f4ec0563d\n1f64955634\n1f7d31b5b2\n1f8014b7fd\n1f9c7d10f1\n1fa350df76\n1fc9538993\n1fe2f0ec59\n2000c02f9d\n20142b2f05\n201a8d75e5\n2023b3ee4f\n202b767bbc\n203594a418\n2038987336\n2039c3aecb\n204a90d81f\n207bc6cf01\n208833d1d1\n20c6d8b362\n20e3e52e0a\n2117fa0c14\n211bc5d102\n2120d9c3c3\n2125235a49\n21386f5978\n2142af8795\n215dfc0f73\n217bae91e5\n217c0d44e4\n219057c87b\n21d0edbf81\n21df87ad76\n21f1d089f5\n21f4019116\n222597030f\n222904eb5b\n223a0e0657\n223bd973ab\n22472f7395\n224e7c833e\n225aba51d9\n2261d421ea\n2263a8782b\n2268cb1ffd\n2268e93b0a\n2293c99f3f\n22a1141970\n22b13084b2\n22d9f5ab0c\n22f02efe3a\n232c09b75b\n2350d71b4b\n2376440551\n2383d8aafd\n238b84e67f\n238d4b86f6\n238d947c6b\n23993ce90d\n23b0c8a9ab\n23b3beafcc\n23d80299fe\n23f404a9fc\n240118e58a\n2431dec2fd\n24440e0ac7\n2457274dbc\n2465bf515d\n246b142c4d\n247d729e36\n2481ceafeb\n24866b4e6a\n2489d78320\n24ab0b83e8\n24b0868d92\n24b5207cd9\n24ddf05c03\n250116161c\n256ad2e3fc\n256bd83d5e\n256dcc8ab8\n2589956baa\n258b3b33c6\n25ad437e29\n25ae395636\n25c750c6db\n25d2c3fe5d\n25dc80db7c\n25f97e926f\n26011bc28b\n260846ffbe\n260dd9ad33\n267964ee57\n2680861931\n268ac7d3fc\n26b895d91e\n26bc786d4f\n26ddd2ef12\n26de3d18ca\n26f7784762\n2703e52a6a\n270ed80c12\n2719b742ab\n272f4163d0\n27303333e1\n27659fa7d6\n279214115d\n27a5f92a9c\n27cf2af1f3\n27f0d5f8a2\n28075f33c1\n281629cb41\n282b0d51f5\n282fcab00b\n28449fa0dc\n28475208ca\n285580b7c4\n285b69e223\n288c117201\n28a8eb9623\n28bf9c3cf3\n28c6b8f86a\n28c972dacd\n28d9fa6016\n28e392de91\n28f4a45190\n298c844fc9\n29a0356a2b\n29d779f9e3\n29dde5f12b\n29de7b6579\n29e630bdd0\n29f2332d30\n2a18873352\n2a3824ff31\n2a559dd27f\n2a5c09acbd\n2a63eb1524\n2a6a30a4ea\n2a6d9099d1\n2a821394e3\n2a8c5b1342\n2abc8d66d2\n2ac9ef904a\n2b08f37364\n2b351bfd7d\n2b659a49d7\n2b69ee5c26\n2b6c30bbbd\n2b88561cf2\n2b8b14954e\n2ba621c750\n2bab50f9a7\n2bb00c2434\n2bbde474ef\n2bdd82fb86\n2be06fb855\n2bf545c2f5\n2bffe4cf9a\n2c04b887b7\n2c05209105\n2c0ad8cf39\n2c11fedca8\n2c1a94ebfb\n2c1e8c8e2f\n2c29fabcf1\n2c2c076c01\n2c3ea7ee7d\n2c41fa0648\n2c44bb6d1c\n2c54cfbb78\n2c5537eddf\n2c6e63b7de\n2cb10c6a7e\n2cbcd5ccd1\n2cc5d9c5f6\n2cd01cf915\n2cdbf5f0a7\n2ce660f123\n2cf114677e\n2d01eef98e\n2d03593bdc\n2d183ac8c4\n2d33ad3935\n2d3991d83e\n2d4333577b\n2d4d015c64\n2d8f5e5025\n2d900bdb8e\n2d9a1a1d49\n2db0576a5c\n2dc0838721\n2dcc417f82\n2df005b843\n2df356de14\n2e00393d96\n2e03b8127a\n2e0f886168\n2e2bf37e6d\n2e42410932\n2ea78f46e4\n2ebb017a26\n2ee2edba2a\n2efb07554a\n2f17e4fc1e\n2f2c65c2f3\n2f2d9b33be\n2f309c206b\n2f53822e88\n2f53998171\n2f5b0c89b1\n2f680909e6\n2f710f66bd\n2f724132b9\n2f7e3517ae\n2f96f5fc6f\n2f97d9fecb\n2fbfa431ec\n2fc9520b53\n2fcd9f4c62\n2feb30f208\n2ff7f5744f\n30085a2cc6\n30176e3615\n301f72ee11\n3026bb2f61\n30318465dc\n3054ca937d\n306121e726\n3064ad91e8\n307444a47f\n307bbb7409\n30a20194ab\n30c35c64a4\n30dbdb2cd6\n30fc77d72f\n310021b58b\n3113140ee8\n3150b2ee57\n31539918c4\n318dfe2ce2\n3193da4835\n319f725ad9\n31bbd0d793\n322505c47f\n322b237865\n322da43910\n3245e049fb\n324c4c38f6\n324e35111a\n3252398f09\n327dc4cabf\n328d918c7d\n3290c0de97\n3299ae3116\n32a7cd687b\n33098cedb4\n3332334ac4\n334cb835ac\n3355e056eb\n33639a2847\n3373891cdc\n337975816b\n33e29d7e91\n34046fe4f2\n3424f58959\n34370a710f\n343bc6a65a\n3450382ef7\n3454303a08\n346aacf439\n346e92ff37\n34a5ece7dd\n34b109755a\n34d1b37101\n34dd2c70a7\n34efa703df\n34fbee00a6\n3504df2fda\n35195a56a1\n351c822748\n351cfd6bc5\n3543d8334c\n35573455c7\n35637a827f\n357a710863\n358bf16f9e\n35ab34cc34\n35c6235b8d\n35d01a438a\n3605019d3b\n3609bc3f88\n360e25da17\n36299c687c\n362c5bc56e\n3649228783\n365b0501ea\n365f459863\n369893f3ad\n369c9977e1\n369dde050a\n36c7dac02f\n36d5b1493b\n36f5cc68fd\n3735480d18\n374b479880\n375a49d38f\n375a5c0e09\n376bda9651\n377db65f60\n37c19d1087\n37d4ae24fc\n37ddce7f8b\n37e10d33af\n37e45c6247\n37fa0001e8\n3802d458c0\n382caa3cb4\n383bb93111\n388843df90\n38924f4a7f\n38b00f93d7\n38c197c10e\n38c9c3d801\n38eb2bf67f\n38fe9b3ed1\n390352cced\n390c51b987\n390ca6f1d6\n392bc0f8a1\n392ecb43bd\n3935291688\n3935e63b41\n394454fa9c\n394638fc8b\n39545e20b7\n397abeae8f\n3988074b88\n398f5d5f19\n39bc49a28c\n39befd99fb\n39c3c7bf55\n39d584b09f\n39f6f6ffb1\n3a079fb484\n3a0d3a81b7\n3a1d55d22b\n3a20a7583e\n3a2c1f66e5\n3a33f4d225\n3a3bf84b13\n3a4565e5ec\n3a4e32ed5e\n3a7ad86ce0\n3a7bdde9b8\n3a98867cbe\n3aa3f1c9e8\n3aa7fce8b6\n3aa876887d\n3ab807ded6\n3ab9b1a85a\n3adac8d7da\n3ae1a4016f\n3ae2deaec2\n3ae81609d6\n3af847e62f\n3b23792b84\n3b3b0af2ee\n3b512dad74\n3b6c7988f6\n3b6e983b5b\n3b74a0fc20\n3b7a50b80d\n3b96d3492f\n3b9ad0c5a9\n3b9ba0894a\n3bb4e10ed7\n3bd9a9b515\n3beef45388\n3c019c0a24\n3c090704aa\n3c2784fc0d\n3c47ab95f8\n3c4db32d74\n3c5ff93faf\n3c700f073e\n3c713cbf2f\n3c8320669c\n3c90d225ee\n3cadbcc404\n3cb9be84a5\n3cc37fd487\n3cc6f90cb2\n3cd5e035ef\n3cdf03531b\n3cdf828f59\n3d254b0bca\n3d5aeac5ba\n3d690473e1\n3d69fed2fb\n3d8997aeb6\n3db0d6b07e\n3db1ddb8cf\n3db907ac77\n3dcbc0635b\n3dd48ed55f\n3de4ac4ec4\n3decd63d88\n3e04a6be11\n3e108fb65a\n3e1448b01c\n3e16c19634\n3e2845307e\n3e38336da5\n3e3a819865\n3e3e4be915\n3e680622d7\n3e7d2aeb07\n3e7d8f363d\n3e91f10205\n3ea4c49bbe\n3eb39d11ab\n3ec273c8d5\n3ed3f91271\n3ee062a2fd\n3eede9782c\n3ef2fa99cb\n3efc6e9892\n3f0b0dfddd\n3f0c860359\n3f18728586\n3f3b15f083\n3f45a470ad\n3f4f3bc803\n3fd96c5267\n3fea675fab\n3fee8cbc9f\n3fff16d112\n401888b36c\n4019231330\n402316532d\n402680df52\n404d02e0c0\n40709263a8\n4083cfbe15\n40a96c5cb1\n40b8e50f82\n40f4026bf5\n4100b57a3a\n41059fdd0b\n41124e36de\n4122aba5f9\n413bab0f0d\n4164faee0b\n418035eec9\n4182d51532\n418bb97e10\n41a34c20e7\n41dab05200\n41ff6d5e2a\n420caf0859\n42264230ba\n425a0c96e0\n42da96b87c\n42eb5a5b0f\n42f17cd14d\n42f5c61c49\n42ffdcdee9\n432f9884f9\n43326d9940\n4350f3ab60\n4399ffade3\n43a6c21f37\n43b5555faa\n43d63b752a\n4416bdd6ac\n4444753edd\n444aa274e7\n444d4e0596\n446b8b5f7a\n4478f694bb\n44b1da0d87\n44b4dad8c9\n44b5ece1b9\n44d239b24e\n44eaf8f51e\n44f4f57099\n44f7422af2\n450787ac97\n4523656564\n4536c882e5\n453b65daa4\n454f227427\n45636d806a\n456fb9362e\n457e717a14\n45a89f35e1\n45bf0e947d\n45c36a9eab\n45d9fc1357\n45f8128b97\n4607f6c03c\n46146dfd39\n4620e66b1e\n4625f3f2d3\n462b22f263\n4634736113\n463c0f4fdd\n46565a75f8\n46630b55ae\n466839cb37\n466ba4ae0c\n4680236c9d\n46bf4e8709\n46e18e42f1\n46f5093c59\n47269e0499\n472da1c484\n47354fab09\n4743bb84a7\n474a796272\n4783d2ab87\n479cad5da3\n479f5d7ef6\n47a05fbd1d\n4804ee2767\n4810c3fbca\n482fb439c2\n48375af288\n484ab44de4\n485f3944cd\n4867b84887\n486a8ac57e\n486e69c5bd\n48812cf33e\n4894b3b9ea\n48bd66517d\n48d83b48a4\n49058178b8\n4918d10ff0\n4932911f80\n49405b7900\n49972c2d14\n499bf07002\n49b16e9377\n49c104258e\n49c879f82d\n49e7326789\n49ec3e406a\n49fbf0c98a\n4a0255c865\n4a088fe99a\n4a341402d0\n4a3471bdf5\n4a4b50571c\n4a50f3d2e9\n4a6e3faaa1\n4a7191f08a\n4a86fcfc30\n4a885fa3ef\n4a8af115de\n4aa2e0f865\n4aa9d6527f\n4abb74bb52\n4ae13de1cd\n4af8cb323f\n4b02c272b3\n4b19c529fb\n4b2974eff4\n4b3154c159\n4b54d2587f\n4b556740ff\n4b67aa9ef6\n4b97cc7b8d\n4baa1ed4aa\n4bc8c676bb\n4beaea4dbe\n4bf5763d24\n4bffa92b67\n4c25dfa8ec\n4c397b6fd4\n4c51e75d66\n4c7710908f\n4c9b5017be\n4ca2ffc361\n4cad2e93bc\n4cd427b535\n4cd9a4b1ef\n4cdfe3c2b2\n4cef87b649\n4cf208e9b3\n4cf5bc3e60\n4cfdd73249\n4cff5c9e42\n4d26d41091\n4d5c23c554\n4d67c59727\n4d983cad9f\n4da0d00b55\n4daa179861\n4dadd57153\n4db117e6c5\n4de4ce4dea\n4dfaee19e5\n4dfdd7fab0\n4e3f346aa5\n4e49c2a9c7\n4e4e06a749\n4e70279712\n4e72856cc7\n4e752f8075\n4e7a28907f\n4e824b9247\n4e82b1df57\n4e87a639bc\n4ea77bfd15\n4eb6fc23a2\n4ec9da329e\n4efb9a0720\n4f062fbc63\n4f35be0e0b\n4f37e86797\n4f414dd6e7\n4f424abded\n4f470cc3ae\n4f601d255a\n4f7386a1ab\n4f824d3dcd\n4f827b0751\n4f8db33a13\n4fa160f8a3\n4fa9c30a45\n4facd8f0e8\n4fca07ad01\n4fded94004\n4fdfef4dea\n4feb3ac01f\n4fffec8479\n500c835a86\n50168342bf\n50243cffdc\n5031d5a036\n504dd9c0fd\n50568fbcfb\n5069c7c5b3\n508189ac91\n50b6b3d4b7\n50c6f4fe3e\n50cce40173\n50efbe152f\n50f290b95d\n5104aa1fea\n5110dc72c0\n511e8ecd7f\n513aada14e\n5158d6e985\n5161e1fa57\n51794ddd58\n517d276725\n51a597ee04\n51b37b6d97\n51b5dc30a0\n51e85b347b\n51eea1fdac\n51eef778af\n51f384721c\n521cfadcb4\n52355da42f\n5247d4b160\n524b470fd0\n524cee1534\n5252195e8a\n5255c9ca97\n525928f46f\n526df007a7\n529b12de78\n52c7a3d653\n52c8ec0373\n52d225ed52\n52ee406d9e\n52ff1ccd4a\n53143511e8\n5316d11eb7\n53253f2362\n534a560609\n5352c4a70e\n536096501f\n536b17bcea\n5380eaabff\n5390a43a54\n53af427bb2\n53bf5964ce\n53c30110b5\n53cad8e44a\n53d9c45013\n53e274f1b5\n53e32d21ea\n540850e1c7\n540cb31cfe\n541c4da30f\n541d7935d7\n545468262b\n5458647306\n54657855cd\n547b3fb23b\n5497dc3712\n549c56f1d4\n54a4260bb1\n54b98b8d5e\n54e1054b0f\n54e8867b83\n54ebe34f6e\n5519b4ad13\n551acbffd5\n55341f42da\n5566ab97e1\n556c79bbf2\n5589637cc4\n558aa072f0\n559824b6f6\n55c1764e90\n55eda6c77e\n562d173565\n5665c024cb\n566cef4959\n5675d78833\n5678a91bd8\n567a2b4bd0\n569c282890\n56cc449917\n56e71f3e07\n56f09b9d92\n56fc0e8cf9\n571ca79c71\n57243657cf\n57246af7d1\n57427393e9\n574b682c19\n578f211b86\n5790ac295d\n579393912d\n57a344ab1a\n57bd3bcda4\n57bfb7fa4c\n57c010175e\n57c457cc75\n57c7fc2183\n57d5289a01\n58045fde85\n58163c37cd\n582d463e5c\n5851739c15\n585dd0f208\n587250f3c3\n589e4cc1de\n589f65f5d5\n58a07c17d5\n58adc6d8b6\n58b9bcf656\n58c374917e\n58fc75fd42\n5914c30f05\n59323787d5\n5937b08d69\n594065ddd7\n595a0ceea6\n59623ec40b\n597ff7ef78\n598935ef05\n598c2ad3b2\n59a6459751\n59b175e138\n59bf0a149f\n59d53d1649\n59e3e6fae7\n59fe33e560\n5a13a73fe5\n5a25c22770\n5a4a785006\n5a50640995\n5a75f7a1cf\n5a841e59ad\n5a91c5ab6d\n5ab49d9de0\n5aba1057fe\n5abe46ba6d\n5ac7c88d0c\n5aeb95cc7d\n5af15e4fc3\n5afe381ae4\n5b07b4229d\n5b1001cc4f\n5b1df237d2\n5b263013bf\n5b27d19f0b\n5b48ae16c5\n5b5babc719\n5baaebdf00\n5bab55cdbe\n5bafef6e79\n5bd1f84545\n5bddc3ba25\n5bdf7c20d2\n5bf23bc9d3\n5c01f6171a\n5c021681b7\n5c185cff1d\n5c42aba280\n5c44bf8ab6\n5c4c574894\n5c52fa4662\n5c6ea7dac3\n5c74315dc2\n5c7668855e\n5c83e96778\n5ca36173e4\n5cac477371\n5cb0cb1b2f\n5cb0cfb98f\n5cb49a19cf\n5cbf7dc388\n5d0e07d126\n5d1e24b6e3\n5d663000ff\n5da6b2dc5d\n5de9b90f24\n5e08de0ed7\n5e1011df9a\n5e1ce354fd\n5e35512dd7\n5e418b25f9\n5e4849935a\n5e4ee19663\n5e886ef78f\n5e8d00b974\n5e8d59dc31\n5ed838bd5c\n5edda6ee5a\n5ede4d2f7a\n5ede9767da\n5eec4d9fe5\n5eecf07824\n5eef7ed4f4\n5ef5860ac6\n5ef6573a99\n5f1193e72b\n5f29ced797\n5f32cf521e\n5f51876986\n5f6ebe94a9\n5f6f14977c\n5f808d0d2d\n5fb8aded6a\n5fba90767d\n5fd1c7a3df\n5fd3da9f68\n5fee2570ae\n5ff66140d6\n5ff8b85b53\n600803c0f6\n600be7f53e\n6024888af8\n603189a03c\n6057307f6e\n6061ddbb65\n606c86c455\n60c61cc2e5\n60e51ff1ae\n610e38b751\n61344be2f6\n6135e27185\n614afe7975\n614e571886\n614e7078db\n619812a1a7\n61b481a78b\n61c7172650\n61cf7e40d2\n61d08ef5a1\n61da008958\n61ed178ecb\n61f5d1282c\n61fd977e49\n621584cffe\n625817a927\n625892cf0b\n625b89d28a\n629995af95\n62a0840bb5\n62ad6e121c\n62d6ece152\n62ede7b2da\n62f025e1bc\n6316faaebc\n63281534dc\n634058dda0\n6353f09384\n6363c87314\n636e4872e0\n637681cd6b\n6376d49f31\n6377809ec2\n63936d7de5\n639bddef11\n63d37e9fd3\n63d90c2bae\n63e544a5d6\n63ebbcf874\n63fff40b31\n6406c72e4d\n64148128be\n6419386729\n643092bc41\n644081b88d\n64453cf61d\n644bad9729\n6454f548fd\n645913b63a\n64750b825f\n64a43876b7\n64dd6c83e3\n64e05bf46e\n64f55f1478\n650b0165e4\n651066ed39\n652b67d960\n653821d680\n6538d00d73\n65866dce22\n6589565c8c\n659832db64\n65ab7e1d98\n65b7dda462\n65bd5eb4f5\n65dcf115ab\n65e9825801\n65f9afe51c\n65ff12bcb5\n666b660284\n6671643f31\n668364b372\n66852243cb\n6693a52081\n669b572898\n66e98e78f5\n670f12e88f\n674c12c92d\n675c27208a\n675ed3e1ca\n67741db50a\n678a2357eb\n67b0f4d562\n67cfbff9b1\n67e717d6bd\n67ea169a3b\n67ea809e0e\n681249baa3\n683de643d9\n6846ac20df\n6848e012ef\n684bcd8812\n684dc1c40c\n685a1fa9cf\n686dafaac9\n68807d8601\n6893778c77\n6899d2dabe\n68a2fad4ab\n68cb45fda3\n68cc4a1970\n68dcb40675\n68ea4a8c3d\n68f6e7fbf0\n68fa8300b4\n69023db81f\n6908ccf557\n691a111e7c\n6927723ba5\n692ca0e1a2\n692eb57b63\n69340faa52\n693cbf0c9d\n6942f684ad\n6944fc833b\n69491c0ebf\n695b61a2b0\n6979b4d83f\n697d4fdb02\n69910460a4\n6997636670\n69a436750b\n69aebf7669\n69b8c17047\n69c67f109f\n69e0e7b868\n69ea9c09d1\n69f0af42a6\n6a078cdcc7\n6a37a91708\n6a42176f2e\n6a48e4aea8\n6a5977be3a\n6a5de0535f\n6a80d2e2e5\n6a96c8815d\n6a986084e2\n6aa8e50445\n6ab9dce449\n6abf0ba6b2\n6acc6049d9\n6adb31756c\n6ade215eb0\n6afb7d50e4\n6afd692f1a\n6b0b1044fe\n6b17c67633\n6b1b6ef28b\n6b1e04d00d\n6b2261888d\n6b25d6528a\n6b3a24395c\n6b685eb75b\n6b79be238c\n6b928b7ba6\n6b9c43c25a\n6ba99cc41f\n6bdab62bcd\n6bf2e853b1\n6bf584200f\n6bf95df2b9\n6c0949c51c\n6c11a5f11f\n6c23d89189\n6c4387daf5\n6c4ce479a4\n6c5123e4bc\n6c54265f16\n6c56848429\n6c623fac5f\n6c81b014e9\n6c99ea7c31\n6c9d29d509\n6c9e3b7d1a\n6ca006e283\n6caeb928d6\n6cb2ee722a\n6cbfd32c5e\n6cc791250b\n6cccc985e0\n6d12e30c48\n6d4bf200ad\n6d6d2b8843\n6d6eea5682\n6d7a3d0c21\n6d7efa9b9e\n6da21f5c91\n6da6adabc0\n6dd2827fbb\n6dd36705b9\n6df3637557\n6dfe55e9e5\n6e1a21ba55\n6e2f834767\n6e36e4929a\n6e4f460caf\n6e618d26b6\n6ead4670f7\n6eaff19b9f\n6eb2e1cd9e\n6eb30b3b5a\n6eca26c202\n6ecad29e52\n6ef0b44654\n6efcfe9275\n6f4789045c\n6f49f522ef\n6f67d7c4c4\n6f96e91d81\n6fc6fce380\n6fc9b44c00\n6fce7f3226\n6fdf1ca888\n702fd8b729\n70405185d2\n7053e4f41e\n707bf4ce41\n7082544248\n708535b72a\n7094ac0f60\n70a6b875fa\n70c3e97e41\n7106b020ab\n711dce6fe2\n7136a4453f\n7143fb084f\n714d902095\n7151c53b32\n715357be94\n7163b8085f\n716df1aa59\n71caded286\n71d2665f35\n71d67b9e19\n71e06dda39\n720b398b9c\n720e3fa04c\n720e7a5f1e\n721bb6f2cb\n722803f4f2\n72552a07c9\n726243a205\n72690ef572\n728cda9b65\n728e81c319\n72a810a799\n72acb8cdf6\n72b01281f9\n72cac683e4\n72cadebbce\n72cae058a5\n72d8dba870\n72e8d1c1ff\n72edc08285\n72f04f1a38\n731b825695\n7320b49b13\n732626383b\n732df1eb05\n73329902ab\n733798921e\n733824d431\n734ea0d7fb\n735a7cf7b9\n7367a42892\n7368d5c053\n73c6ae7711\n73e1852735\n73e4e5cc74\n73eac9156b\n73f8441a88\n7419e2ab3f\n74267f68b9\n7435690c8c\n747c44785c\n747f1b1f2f\n748b2d5c01\n74d4cee0a4\n74ec2b3073\n74ef677020\n750be4c4d8\n75172d4ac8\n75285a7eb1\n75504539c3\n7550949b1d\n7551cbd537\n75595b453d\n7559b4b0ec\n755bd1fbeb\n756f76f74d\n7570ca7f3c\n757a69746e\n757cac96c6\n7584129dc3\n75a058dbcd\n75b09ce005\n75cae39a8f\n75cee6caf0\n75cf58fb2c\n75d5c2f32a\n75eaf5669d\n75f7937438\n75f99bd3b3\n75fa586876\n7613df1f84\n762e1b3487\n76379a3e69\n764271f0f3\n764503c499\n7660005554\n7666351b84\n76693db153\n767856368b\n768671f652\n768802b80d\n76962c7ed2\n76a75f4eee\n76b90809f7\n770a441457\n772a0fa402\n772f2ffc3e\n774f6c2175\n77610860e0\n777e58ff3d\n77920f1708\n7799df28e7\n779e847a9a\n77ba4edc72\n77c834dc43\n77d8aa8691\n77e7f38f4d\n77eea6845e\n7806308f33\n78254660ea\n7828af8bff\n784398620a\n784d201b12\n78613981ed\n78896c6baf\n78aff3ebc0\n78c7c03716\n78d3676361\n78e29dd4c3\n78f1a1a54f\n79208585cd\n792218456c\n7923bad550\n794e6fc49f\n796e6762ce\n797cd21f71\n79921b21c2\n79a5778027\n79bc006280\n79bf95e624\n79d9e00c55\n79e20fc008\n79e9db913e\n79f014085e\n79fcbb433a\n7a13a5dfaa\n7a14bc9a36\n7a3c535f70\n7a446a51e9\n7a56e759c5\n7a5f46198d\n7a626ec98d\n7a802264c4\n7a8b5456ca\n7abdff3086\n7aecf9f7ac\n7b0fd09c28\n7b18b3db87\n7b39fe7371\n7b49e03d4c\n7b5388c9f1\n7b5cf7837f\n7b733d31d8\n7b74fd7b98\n7b918ccb8a\n7ba3ce3485\n7bb0abc031\n7bb5bb25cd\n7bb7dac673\n7bc7761b8c\n7bf3820566\n7c03a18ec1\n7c078f211b\n7c37d7991a\n7c4ec17eff\n7c649c2aaf\n7c73340ab7\n7c78a2266d\n7c88ce3c5b\n7ca6843a72\n7cc9258dee\n7cec7296ae\n7d0ffa68a4\n7d11b4450f\n7d1333fcbe\n7d18074fef\n7d18c8c716\n7d508fb027\n7d55f791f0\n7d74e3c2f6\n7d783f67a9\n7d83a5d854\n7dd409947e\n7de45f75e5\n7e0cd25696\n7e1922575c\n7e1e3bbcc1\n7e24023274\n7e2f212fd3\n7e6d1cc1f4\n7e7cdcb284\n7e9b6bef69\n7ea5b49283\n7eb2605d96\n7eb26b8485\n7ecd1f0c69\n7f02b3cfe2\n7f1723f0d5\n7f21063c3a\n7f3658460e\n7f54132e48\n7f559f9d4a\n7f5faedf8b\n7f838baf2b\n7fa5f527e3\n7ff84d66dd\n802b45c8c4\n804382b1ad\n804c558adb\n804f6338a4\n8056117b89\n806b6223ab\n8088bda461\n80b790703b\n80c4a94706\n80ce2e351b\n80db581acd\n80e12193df\n80e41b608f\n80f16b016d\n81541b3725\n8175486e6a\n8179095000\n8193671178\n81a58d2c6b\n81aa1286fb\n81dffd30fb\n8200245704\n823e7a86e8\n824973babb\n824ca5538f\n827171a845\n8273a03530\n827cf4f886\n82b865c7dd\n82c1517708\n82d15514d6\n82e117b900\n82fec06574\n832b5ef379\n83424c9fbf\n8345358fb8\n834b50b31b\n835e3b67d7\n836ea92b15\n837c618777\n838eb3bd89\n839381063f\n839bc71489\n83a8151377\n83ae88d217\n83ca8bcad0\n83ce590d7f\n83d3130ba0\n83d40bcba5\n83daba503a\n83de906ec0\n84044f37f3\n84696b5a5e\n84752191a3\n847eeeb2e0\n848e7835a0\n84a4b29286\n84a4bf147d\n84be115c09\n84d95c4350\n84e0922cf7\n84f0cfc665\n8515f6db22\n851f2f32c1\n852a4d6067\n854c48b02a\n857a387c86\n859633d56a\n85a4f4a639\n85ab85510c\n85b1eda0d9\n85dc1041c6\n85e081f3c7\n85f75187ad\n8604bb2b75\n860745b042\n863b4049d7\n8643de22d0\n8647d06439\n864ffce4fe\n8662d9441a\n8666521b13\n868d6a0685\n869fa45998\n86a40b655d\n86a8ae4223\n86b2180703\n86c85d27df\n86d3755680\n86e61829a1\n871015806c\n871e409c5c\n8744b861ce\n8749369ba0\n878a299541\n8792c193a0\n8799ab0118\n87d1f7d741\n882b9e4500\n885673ea17\n8859dedf41\n8873ab2806\n887a93b198\n8883e991a9\n8891aa6dfa\n8899d8cbcd\n88b8274d67\n88d3b80af6\n88ede83da2\n88f345941b\n890976d6da\n8909bde9ab\n8929c7d5d9\n89363acf76\n89379487e0\n8939db6354\n893f658345\n8953138465\n895c96d671\n895cbf96f9\n895e8b29a7\n898fa256c8\n89986c60be\n89b874547b\n89bdb021d5\n89c802ff9c\n89d6336c2b\n89ebb27334\n8a27e2407c\n8a31f7bca5\n8a4a2fc105\n8a5d6c619c\n8a75ad7924\n8aa817e4ed\n8aad0591eb\n8aca214360\n8ae168c71b\n8b0cfbab97\n8b3645d826\n8b3805dbd4\n8b473f0f5d\n8b4f6d1186\n8b4fb018b7\n8b518ee936\n8b523bdfd6\n8b52fb5fba\n8b91036e5c\n8b99a77ac5\n8ba04b1e7b\n8ba782192f\n8bbeaad78b\n8bd1b45776\n8bd7a2dda6\n8bdb091ccf\n8be56f165d\n8be950d00f\n8bf84e7d45\n8bffc4374b\n8bfff50747\n8c09867481\n8c0a3251c3\n8c3015cccb\n8c469815cf\n8c9ccfedc7\n8ca1af9f3c\n8ca3f6e6c1\n8ca6a4f60f\n8cac6900fe\n8cba221a1e\n8cbbe62ccd\n8d064b29e2\n8d167e7c08\n8d4ab94e1c\n8d81f6f899\n8d87897d66\n8dcccd2bd2\n8dcfb878a8\n8dd3ab71b9\n8dda6bf10f\n8ddd51ca94\n8dea22c533\n8def5bd3bf\n8e1848197c\n8e3a83cf2d\n8e478e73f3\n8e98ae3c84\n8ea6687ab0\n8eb0d315c1\n8ec10891f9\n8ec3065ec2\n8ecf51a971\n8eddbab9f7\n8ee198467a\n8ee2368f40\n8ef595ce82\n8f0a653ad7\n8f1204a732\n8f1600f7f6\n8f16366707\n8f1ce0a411\n8f2e05e814\n8f320d0e09\n8f3b4a84ad\n8f3fdad3da\n8f5d3622d8\n8f62a2c633\n8f81c9405a\n8f8c974d53\n8f918598b6\n8ff61619f6\n9002761b41\n90107941f3\n90118a42ee\n902bc16b37\n903e87e0d6\n9041a0f489\n9047bf3222\n9057bfa502\n90617b0954\n9076f4b6db\n9077e69b08\n909655b4a6\n909c2eca88\n909dbd1b76\n90bc4a319a\n90c7a87887\n90cc785ddd\n90d300f09b\n9101ea9b1b\n9108130458\n911ac9979b\n9151cad9b5\n9153762797\n91634ee0c9\n916942666f\n9198cfb4ea\n919ac864d6\n91b67d58d4\n91bb8df281\n91be106477\n91c33b4290\n91ca7dd9f3\n91d095f869\n91f107082e\n920329dd5e\n920c959958\n92128fbf4b\n9223dacb40\n923137bb7f\n9268e1f88a\n927647fe08\n9276f5ba47\n92a28cd233\n92b5c1fc6d\n92c46be756\n92dabbe3a0\n92e3159361\n92ebab216a\n934bdc2893\n9359174efc\n935d97dd2f\n935feaba1b\n93901858ee\n939378f6d6\n939bdf742e\n93a22bee7e\n93da9aeddf\n93e2feacce\n93e6f1fdf9\n93e811e393\n93e85d8fd3\n93f623d716\n93ff35e801\n94031f12f2\n94091a4873\n94125907e3\n9418653742\n941c870569\n94209c86f0\n9437c715eb\n9445c3eca2\n9467c8617c\n946d71fb5d\n948f3ae6fb\n9498baa359\n94a33abeab\n94bf1af5e3\n94cf3a8025\n94db712ac8\n94e4b66cff\n94e76cbaf6\n950be91db1\n952058e2d0\n952633c37f\n952ec313fe\n9533fc037c\n9574b81269\n9579b73761\n957f7bc48b\n958073d2b0\n9582e0eb33\n9584092d0b\n95b58b8004\n95bd88da55\n95f74a9959\n962781c601\n962f045bf5\n964ad23b44\n967b90590e\n967bffe201\n96825c4714\n968492136a\n9684ef9d64\n968c41829e\n96a856ef9a\n96dfc49961\n96e1a5b4f8\n96e6ff0917\n96fb88e9d7\n96fbe5fc23\n96fc924050\n9715cc83dc\n9720eff40f\n972c187c0d\n97476eb38d\n97659ed431\n9773492949\n97756b264f\n977bff0d10\n97ab569ff3\n97ba838008\n97d9d008c7\n97e59f09fa\n97eb642e56\n98043e2d14\n981ff580cf\n983e66cbfc\n984f0f1c36\n98595f2bb4\n985c3be474\n9869a12362\n986b5a5e18\n9877af5063\n98911292da\n9893a3cf77\n9893d9202d\n98a8b06e7f\n98ac6f93d9\n98b6974d12\n98ba3c9417\n98c7c00a19\n98d044f206\n98e909f9d1\n98fe7f0410\n990f2742c7\n992bd0779a\n994b9b47ba\n9955b76bf5\n9966f3adac\n997117a654\n999d53d841\n99c04108d3\n99c4277aee\n99c6b1acf2\n99dc8bb20b\n99fcba71e5\n99fecd4efb\n9a02c70ba2\n9a08e7a6f8\n9a2f2c0f86\n9a3254a76e\n9a3570a020\n9a39112493\n9a4e9fd399\n9a50af4bfb\n9a68631d24\n9a72318dbf\n9a767493b7\n9a7fc1548b\n9a84ccf6a7\n9a9c0e15b7\n9adf06d89b\n9b22b54ee4\n9b473fc8fe\n9b4f081782\n9b997664ba\n9bc454e109\n9bccfd04de\n9bce4583a2\n9bebf1b87f\n9bfc50d261\n9c166c86ff\n9c293ef4d7\n9c29c047b0\n9c3bc2e2a7\n9c3ce23bd1\n9c404cac0c\n9c5180d23a\n9c7feca6e4\n9caa49d3ff\n9cb2f1b646\n9ce6f765c3\n9cfee34031\n9d01f08ec6\n9d04c280b8\n9d12ceaddc\n9d15f8cb3c\n9d2101e9bf\n9d407c3aeb\n9ddefc6165\n9df0b1e298\n9e16f115d8\n9e249b4982\n9e29b1982c\n9e493e4773\n9e4c752cd0\n9e4de40671\n9e6319faeb\n9e6ddbb52d\n9eadcea74f\n9ecec5f8ea\n9efb47b595\n9f30bfe61e\n9f3734c3a4\n9f5b858101\n9f66640cda\n9f913803e9\n9f97bc74c8\n9fbad86e20\n9fc2bad316\n9fc5c3af78\n9fcb310255\n9fcc256871\n9fd2fd4d47\na0071ae316\na023141022\na046399a74\na066e739c1\na06722ba82\na07a15dd64\na07b47f694\na09c39472e\na0b208fe2e\na0b61c959e\na0bc6c611d\na0e6da5ba2\na1193d6490\na14ef483ff\na14f709908\na15ccc5658\na16062456f\na174e8d989\na177c2733c\na17c62e764\na18ad065fc\na1aaf63216\na1bb65fb91\na1bd8e5349\na1dfdd0cac\na2052e4f6c\na20fd34693\na21ffe4d81\na22349e647\na235d01ec1\na24f63e8a2\na2554c9f6d\na263ce8a87\na29bfc29ec\na2a80072d4\na2a800ab63\na2bcd10a33\na2bdaff3b0\na2c146ab0d\na2c996e429\na2dc51ebe8\na2e6608bfa\na2f2a55f01\na301869dea\na31fccd2cc\na34f440f33\na35e0206da\na36bdc4cab\na36e8c79d8\na378053b20\na37db3a2b3\na38950ebc2\na39a0eb433\na39c9bca52\na3a945dc8c\na3b40a0c1e\na3b8588550\na3c502bec3\na3f2878017\na3f4d58010\na3f51855c3\na402dc0dfe\na4065a7eda\na412bb2fef\na416b56b53\na41ec95906\na43299e362\na4757bd7af\na48c53c454\na49dcf9ad5\na4a506521f\na4ba7753d9\na4bac06849\na4f05d681c\na50c10060f\na50eb5a0ea\na5122c6ec6\na522b1aa79\na590915345\na5b5b59139\na5b77abe43\na5c2b2c3e1\na5cd17bb11\na5da03aef1\na5dd11de0d\na5ea2b93b6\na5eaeac80b\na5ec5b0265\na5f350a87e\na5f472caf4\na6027a53cf\na61715bb1b\na61cf4389d\na61d9bbd9b\na6470dbbf5\na64a40f3eb\na653d5c23b\na65bd23cb5\na66e0b7ad4\na66fc5053c\na68259572b\na6a810a92c\na6bc36937f\na6c3a374e9\na6d8a4228d\na6f4e0817f\na71e0481f5\na7203deb2d\na7392d4438\na73d3c3902\na7491f1578\na74b9ca19c\na77b7a91df\na78195a5f5\na78758d4ce\na7e6d6c29a\na800d85e88\na832fa8790\na83d06410d\na8999af004\na8f78125b9\na907b18df1\na919392446\na965504e88\na96b84b8d2\na973f239cd\na977126596\na9804f2a08\na984e56893\na99738f24c\na99bdd0079\na9c9c1517e\na9cbf9c41b\na9e42e3c0c\naa07b7c1c0\naa175e5ec7\naa1a338630\naa27d7b868\naa45f1caaf\naa49e46432\naa51934e1b\naa6287bb6c\naa6d999971\naa85278334\naab33f0e2a\naaba004362\naade4cf385\naae78feda4\naaed233bf3\naaff16c2db\nab199e8dfb\nab23b78715\nab2e1b5577\nab33a18ded\nab45078265\nab56201494\nab90f0d24b\nabab2e6c20\nabb50c8697\nabbe2d15a0\nabbe73cd21\nabe61a11bb\nabeae8ce21\nac2b431d5f\nac2cb1b9eb\nac31fcd6d0\nac3d3a126d\nac46bd8087\nac783ef388\nacb73e4297\nacbf581760\naccafc3531\nacf2c4b745\nacf44293a2\nacf736a27b\nacff336758\nad1fe56886\nad28f9b9d9\nad2de9f80e\nad397527b2\nad3d1cfbcb\nad3fada9d9\nad4108ee8e\nad54468654\nad573f7d31\nad6255bc29\nad65ebaa07\nad97cc064a\nadabbd1cc4\nadb0b5a270\nadc648f890\nadd21ee467\nadfd15ceef\nadfdd52eac\nae01cdab63\nae0b50ff4f\nae13ee3d70\nae1bcbd423\nae20d09dea\nae2cecf5f6\nae3bc4a0ef\nae499c7514\nae628f2cd4\nae8545d581\nae93214fe6\nae9cd16dbf\naeba9ac967\naebb242b5c\naed4e0b4c4\naedd71f125\naef3e2cb0e\naf0b54cee3\naf3de54c7a\naf5fd24a36\naf8826d084\naf8ad72057\nafb71e22c5\nafcb331e1f\nafe1a35c1e\nb01080b5d3\nb05ad0d345\nb0623a6232\nb064dbd4b7\nb06ed37831\nb06f5888e6\nb08dcc490e\nb0a68228dc\nb0aece727f\nb0b0731606\nb0c7f11f9f\nb0cca8b830\nb0dd580a89\nb0de66ca08\nb0df7c5c5c\nb0f5295608\nb11099eb09\nb132a53086\nb1399fac64\nb13abc0c69\nb1457e3b5e\nb15bf4453b\nb179c4a82d\nb17ee70e8c\nb190b1aa65\nb19b3e22c0\nb19c561fab\nb1d1cd2e6e\nb1d7c03927\nb1d7fe2753\nb1f540a4bd\nb1fc9c64e1\nb1fcbb3ced\nb220939e93\nb22099b419\nb241e95235\nb2432ae86d\nb2456267df\nb247940d01\nb24af1c35c\nb24f600420\nb24fe36b2a\nb258fb0b7d\nb26b219919\nb26d9904de\nb274456ce1\nb27b28d581\nb2a26bc912\nb2a9c51e1b\nb2b0baf470\nb2b2756fe7\nb2ce7699e3\nb2edc76bd2\nb2f6b52100\nb30bf47bcd\nb34105a4e9\nb372a82edf\nb3779a1962\nb379ab4ff5\nb37a1d69e3\nb37c01396e\nb382b09e25\nb3996e4ba5\nb3d9ca2aee\nb3dde1e1e9\nb3eb7f05eb\nb40b25055c\nb41e0f1f19\nb44e32a42b\nb4805ae9cd\nb4807569a5\nb48efceb3e\nb493c25c7f\nb4b565aba1\nb4b715a15b\nb4d0c90bf4\nb4d84bc371\nb4e5ad97aa\nb4eaea9e6b\nb50f4b90d5\nb53f675641\nb54278cd43\nb554843889\nb573c0677a\nb58d853734\nb5943b18ab\nb5a09a83f3\nb5aae1fe25\nb5b9da5364\nb5eb64d419\nb5ebb1d000\nb5f1c0c96a\nb5f7fece90\nb6070de1bb\nb60a76fe73\nb61f998772\nb62c943664\nb63094ba0c\nb64fca8100\nb673e7dcfb\nb678b7db00\nb68fc1b217\nb69926d9fa\nb6a1df3764\nb6a4859528\nb6b4738b78\nb6b4f847b7\nb6b8d502d4\nb6bb00e366\nb6d65a9eef\nb6d79a0845\nb6e9ec577f\nb6ec609f7b\nb6f92a308d\nb70a2c0ab1\nb70a5a0d50\nb70c052f2f\nb70d231781\nb72ac6e10b\nb7302d8226\nb73867d769\nb751e767f2\nb76df6e059\nb77e5eddef\nb7a2c2c83c\nb7bcbe6466\nb7c2a469c4\nb7d69da8f0\nb7f31b7c36\nb7f675fb98\nb7fb871660\nb82e5ad1c9\nb841cfb932\nb84b8ae665\nb85b78ac2b\nb86c17caa6\nb86e50d82d\nb871db031a\nb87d56925a\nb8aaa59b75\nb8c03d1091\nb8c3210036\nb8e16df00b\nb8f34cf72e\nb8fb75864e\nb9004db86c\nb9166cbae9\nb920b256a6\nb938d79dff\nb93963f214\nb941aef1a0\nb94d34d14e\nb964c57da4\nb96a95bc7a\nb96c57d2c7\nb9b6bdde0c\nb9bcb3e0f2\nb9d3b92169\nb9dd4b306c\nb9f43ef41e\nba1f03c811\nba3a775d7b\nba3c7f2a31\nba3fcd417d\nba5e1f4faa\nba795f3089\nba8a291e6a\nba98512f97\nbac9db04f5\nbaedae3442\nbaff40d29d\nbb04e28695\nbb1b0ee89f\nbb1c770fe7\nbb1fc34f99\nbb2d220506\nbb334e5cdb\nbb337f9830\nbb721eb9aa\nbb87ff58bd\nbb89a6b18a\nbbaa9a036a\nbbb4302dda\nbbd31510cf\nbbe0256a75\nbc141b9ad5\nbc17ab8a99\nbc318160de\nbc3b9ee033\nbc4240b43c\nbc4ce49105\nbc4f71372d\nbc6b8d6371\nbcaad44ad7\nbcc241b081\nbcc5d8095e\nbcd1d39afb\nbd0d849da4\nbd0e9ed437\nbd2c94730f\nbd321d2be6\nbd3ec46511\nbd5b2e2848\nbd7e02b139\nbd96f9943a\nbda224cb25\nbda4a82837\nbdb74e333f\nbdccd69dde\nbddcc15521\nbe116aab29\nbe15e18f1e\nbe1a284edb\nbe2a367a7b\nbe376082d0\nbe3e3cffbd\nbe5d1d89a0\nbe8b72fe37\nbe9b29e08e\nbea1f6e62c\nbea83281b5\nbeb921a4c9\nbec5e9edcd\nbeeb8a3f92\nbf2232b58d\nbf28751739\nbf443804e8\nbf461df850\nbf5374f122\nbf551a6f60\nbf8d0f5ada\nbf961167a6\nbfab1ad8f9\nbfcb05d88d\nbfd8f6e6c9\nbfd91d0742\nbfe262322f\nc013f42ed7\nc01878083f\nc01faff1ed\nc046fd0edb\nc053e35f97\nc079a6482d\nc0847b521a\nc0a1e06710\nc0e8d4635c\nc0e973ad85\nc0f49c6579\nc0f5b222d7\nc10d07c90d\nc1268d998c\nc130c3fc0c\nc14826ad5e\nc15b922281\nc16f09cb63\nc18e19d922\nc1c830a735\nc1e8aeea45\nc20a5ccc99\nc20fd5e597\nc219d6f8dc\nc2406ae462\nc26f7b5824\nc279e641ee\nc27adaeac5\nc2a35c1cda\nc2a9903b8b\nc2b62567c1\nc2b974ec8c\nc2baaff7bf\nc2be6900f2\nc304dd44d5\nc307f33da2\nc30a7b62c9\nc3128733ee\nc31fa6c598\nc325c8201e\nc32d4aa5d1\nc33f28249a\nc34365e2d7\nc3457af795\nc34d120a88\nc3509e728d\nc35e4fa6c4\nc36240d96f\nc3641dfc5a\nc37b17a4a9\nc39559ddf6\nc3b0c6e180\nc3b3d82e6c\nc3be369fdb\nc3bf1e40c2\nc3c760b015\nc3dd38bf98\nc3e4274614\nc3edc48cbd\nc41e6587f5\nc4272227b0\nc42917fe82\nc438858117\nc44676563f\nc44beb7472\nc45411dacb\nc4571bedc8\nc46deb2956\nc479ee052e\nc47d551843\nc49f07d46d\nc4cc40c1fc\nc4f256f5d5\nc4f5b1ddcc\nc4ff9b4885\nc52bce43db\nc544da6854\nc55784c766\nc557b69fbf\nc593a3f7ab\nc598faa682\nc5ab1f09c8\nc5b6da8602\nc5b9128d94\nc5e845c6b7\nc5fba7b341\nc60897f093\nc61fe6ed7c\nc62188c536\nc64035b2e2\nc69689f177\nc6a12c131f\nc6bb6d2d5c\nc6c18e860f\nc6d9526e0d\nc6e55c33f0\nc7030b28bd\nc70682c7cc\nc70f9be8c5\nc71f30d7b6\nc73c8e747f\nc760eeb8b3\nc7637cab0a\nc7a1a17308\nc7bf937af5\nc7c2860db3\nc7cef4aee2\nc7ebfc5d57\nc813dcf13c\nc82235a49a\nc82a7619a1\nc82ecb90cb\nc844f03dc7\nc8557963f3\nc89147e6e8\nc8a46ff0c8\nc8ab107dd5\nc8b869a04a\nc8c7b306a6\nc8c8b28781\nc8d79e3163\nc8edab0415\nc8f494f416\nc8f6cba9fd\nc909ceea97\nc9188f4980\nc922365dd4\nc92c8c3c75\nc937eb0b83\nc94b31b5e5\nc95cd17749\nc96379c03c\nc96465ee65\nc965afa713\nc9734b451f\nc9862d82dc\nc98b6fe013\nc9999b7c48\nc99e92aaf0\nc9b3a8fbda\nc9bf64e965\nc9c3cb3797\nc9d1c60cd0\nc9de9c22c4\nca1828fa54\nca346f17eb\nca3787d3d3\nca4b99cbac\nca91c69e3b\nca91e99105\ncaa8e97f81\ncaac5807f8\ncabba242c2\ncad5a656a9\ncad673e375\ncad8a85930\ncae7b0a02b\ncae7ef3184\ncaeb6b6cbb\ncaecf0a5db\ncb15312003\ncb2e35d610\ncb35a87504\ncb3f22b0cf\ncbb410da64\ncc8728052e\ncc892997b8\ncce03c2a9b\ncd47a23e31\ncd4dc03dc0\ncd5ae611da\ncd603bb9d1\ncd8f49734c\ncdc6b1c032\ncdcfe008ad\ncdd57027c2\nce1af99b4b\nce1bc5743a\nce25872021\nce2776f78f\nce49b1f474\nce4f0a266f\nce5641b195\nce6866aa19\nce712ed3c9\nce7d1c8117\nce7dbeaa88\nce9b015a5e\ncea7697b25\ncebbd826cf\ncec3415361\ncec41ad4f4\nced49d26df\nced7705ab2\ncef824a1e1\ncf13f5c95a\ncf4376a52d\ncf85ab28b5\ncfc2e50b9d\ncfcd571fff\ncfd9d4ae47\ncfda2dcce5\ncff035928b\ncff8191891\nd01608c2a5\nd01a8f1f83\nd021d68bca\nd04258ca14\nd0483573dc\nd04a90aaff\nd05279c0bd\nd0696bd5fc\nd072fda75b\nd0a83bcd9f\nd0ab39112e\nd0acde820f\nd0b4442c71\nd0c65e9e95\nd0fb600c73\nd107a1457c\nd123d674c1\nd14d1e9289\nd154e3388e\nd177e9878a\nd1802f69f8\nd182c4483a\nd195d31128\nd200838929\nd205e3cff5\nd247420c4c\nd2484bff33\nd26f6ed9b0\nd280fcd1cb\nd2857f0faa\nd292a50c7f\nd295ea2dc7\nd2a58b4fa6\nd2b026739a\nd2ebe0890f\nd2ede5d862\nd301ca58cc\nd3069da8bb\nd343d4a77d\nd355e634ef\nd367fb5253\nd36d16358e\nd38bc77e2c\nd38d1679e2\nd3932ad4bd\nd3987b2930\nd39934abe3\nd3ae1c3f4c\nd3b088e593\nd3e6e05e16\nd3eefae7c5\nd3f55f5ab8\nd3f5c309cc\nd4034a7fdf\nd4193011f3\nd429c67630\nd42c0ff975\nd44a764409\nd44e6acd1d\nd45158c175\nd454e8444f\nd45f62717e\nd48ebdcf74\nd49ab52a25\nd4a607ad81\nd4b063c7db\nd4da13e9ba\nd4dd1a7d00\nd4f4f7c9c3\nd521aba02e\nd535bb1b97\nd53b955f78\nd55cb7a205\nd55f247a45\nd5695544d8\nd5853d9b8b\nd5b6c6d94a\nd5cae12834\nd5df027f0c\nd5ee40e5d0\nd600046f73\nd632fd3510\nd6476cad55\nd65a7bae86\nd664c89912\nd689658f06\nd6917db4be\nd69967143e\nd699d3d798\nd69f757a3f\nd6ac0e065c\nd6c02bfda5\nd6c1b5749e\nd6e12ef6cc\nd6eed152c4\nd6faaaf726\nd704766646\nd708e1350c\nd7135cf104\nd7157a9f44\nd719cf9316\nd724134cfd\nd73a60a244\nd7411662da\nd74875ea7c\nd756f5a694\nd7572b7d8a\nd763bd6d96\nd7697c8b13\nd7797196b4\nd79c834768\nd7b34e5d73\nd7bb6b37a7\nd7c7e064a6\nd7fbf545b3\nd82a0aa15b\nd847e24abd\nd8596701b7\nd86101499c\nd87069ba86\nd87160957b\nd874654b52\nd88a403092\nd8aee40f3f\nd8e77a222d\nd8eb07c381\nd9010348a1\nd90e3cf281\nd92532c7b2\nd927fae122\nd95707bca8\nd973b31c00\nd991cb471d\nd992c69d37\nd99d770820\nd9b63abc11\nd9db6f1983\nd9e52be2d2\nd9edc82650\nda01070697\nda070ea4b7\nda080507b9\nda0e944cc4\nda28d94ff4\nda5d78b9d1\nda6003fc72\nda690fee9f\nda6c68708f\nda7a816676\ndac361e828\ndac71659b8\ndad980385d\ndaebc12b77\ndb0968cdd3\ndb231a7100\ndb59282ace\ndb7f267c3f\ndba35b87fd\ndbba735a50\ndbca076acd\ndbd66dc3ac\ndbdc3c292b\ndbf4a5b32b\ndbfc417d28\ndc1745e0a2\ndc32a44804\ndc34b35e30\ndc504a4f79\ndc704dd647\ndc71bc6918\ndc7771b3be\ndcf8c93617\ndd0f4c9fb9\ndd415df125\ndd601f9a3f\ndd61d903df\ndd77583736\ndd8636bd8b\ndd9fe6c6ac\nddb2da4c14\nddcd450d47\ndde8e67fb4\nddfc3f04d3\nde2ab79dfa\nde2f35b2fd\nde30990a51\nde36b216da\nde37403340\nde46e4943b\nde4ddbccb1\nde5e480f05\nde6a9382ca\nde74a601d3\nde827c510d\nded6069f7b\ndefb71c741\ndf01f277f1\ndf05214b82\ndf0638b0a0\ndf11931ffe\ndf1b0e4620\ndf20a8650d\ndf2bc56d7c\ndf365282c6\ndf39a0d9df\ndf3c430c24\ndf5536cfb9\ndf59cfd91d\ndf5e2152b3\ndf741313c9\ndf7626172f\ndf8ad5deb9\ndf96aa609a\ndf9705605c\ndf9c91c4da\ndfc0d3d27a\ndfdbf91a99\ne00baaae9b\ne0a938c6e7\ne0b2ceee6f\ne0bdb5dfae\ne0be1f6e17\ne0c478f775\ne0de82caa7\ne0f217dd59\ne0f7208874\ne0fb58395e\ne1194c2e9d\ne11adcd05d\ne128124b9d\ne1495354e4\ne1561d6d4b\ne158805399\ne16945b951\ne19edcd34b\ne1a1544285\ne1ab7957f4\ne1d26d35be\ne1e957085b\ne1f14510fa\ne214b160f4\ne2167379b8\ne21acb20ab\ne221105579\ne22ddf8a1b\ne22de45950\ne22ffc469b\ne23cca5244\ne252f46f0b\ne25fa6cf39\ne26e486026\ne275760245\ne27bbedbfe\ne29e9868a8\ne2b37ff8af\ne2b608d309\ne2bef4da9a\ne2c87a6421\ne2ea25542c\ne2fb1d6497\ne2fcc99117\ne33c18412a\ne348377191\ne352cb59c8\ne36ac982f0\ne391bc981e\ne39e3e0a06\ne3bf38265f\ne3d5b2cd21\ne3d60e82d5\ne3e3245492\ne3e4134877\ne3f4635e03\ne4004ee048\ne402d1afa5\ne415093d27\ne41ceb5d81\ne424653b78\ne42b6d3dbb\ne42d60f0d4\ne436d0ff1e\ne43d7ae2c5\ne4428801bc\ne44e0b4917\ne470345ede\ne48e8b4263\ne4922e3726\ne4936852bb\ne495f32c60\ne499228f26\ne4af66e163\ne4b2095f58\ne4d19c8283\ne4d4872dab\ne4e2983570\ne4eaa63aab\ne4ef0a3a34\ne4f8e5f46e\ne4ffb6d0dd\ne53e21aa02\ne57f4f668b\ne588433c1e\ne597442c99\ne5abc0e96b\ne5be628030\ne5ce96a55d\ne5d6b70a9f\ne5fde1574c\ne625e1d27b\ne6261d2348\ne6267d46bc\ne6295f223f\ne63463d8c6\ne6387bd1e0\ne653883384\ne65f134e0b\ne668ef5664\ne672ccd250\ne674510b20\ne676107765\ne699da0cdf\ne6be243065\ne6deab5e0b\ne6f065f2b9\ne71629e7b5\ne72a7d7b0b\ne72f6104e1\ne75a466eea\ne76c55933f\ne7784ec8ad\ne78922e5e6\ne78d450a9c\ne7c6354e77\ne7c8de1fce\ne7ea10db28\ne803918710\ne8073a140b\ne828dd02db\ne845994987\ne8485a2615\ne85c5118a7\ne88b6736e4\ne8962324e3\ne8b3018d36\ne8cee8bf0b\ne8d97ebece\ne8da49ea6a\ne8ed1a3ccf\ne8f7904326\ne8f8341dec\ne8fa21eb13\ne90c10fc4c\ne914b8cac8\ne92b6bfea4\ne92e1b7623\ne93f83e512\ne9422ad240\ne9460b55f9\ne9502628f6\ne950befd5f\ne9582bdd1b\ne95e5afe0f\ne97cfac475\ne98d57d99c\ne98eda8978\ne99706b555\ne9bc0760ba\ne9d3c78bf3\ne9ec1b7ea8\nea065cc205\nea138b6617\nea16d3fd48\nea2545d64b\nea286a581c\nea320da917\nea345f3627\nea3b94a591\nea444a37eb\nea4a01216b\nea5672ffa8\neaa99191cb\neaab4d746c\neac7a59bc1\nead5d3835a\neaec65cfa7\neaed1a87be\neb2f821c6f\neb383cb82e\neb6992fe02\neb6ac20a01\neb6d7ab39e\neb7921facd\neb8fce51a6\nebbb90e9f9\nebbf5c9ee1\nebc4ec32e6\nebe56e5ef8\nec1299aee4\nec139ff675\nec193e1a01\nec28252938\nec387be051\nec3d4fac00\nec4186ce12\nec579c2f96\necae59b782\necb33a0448\nece6bc9e92\necfedd4035\necfff22fd6\ned3291c3d6\ned3cd5308d\ned3e6fc1a5\ned72ae8825\ned7455da68\ned844e879f\ned8f814b2b\ned911a1f63\ned9ff4f649\neda8ab984b\nedb8878849\nedbfdfe1b4\nedd22c46a2\nedd663afa3\nede3552eae\nedeab61ee0\nee07583fc0\nee316eaed6\nee3f509537\nee40a1e491\nee4bf100f1\nee6f9b01f9\nee947ed771\nee9706ac7f\nee9a7840ae\neeb90cb569\neebf45e5c5\neeed0c7d73\nef0061a309\nef07f1a655\nef0a8e8f35\nef232a2aed\nef308ad2e9\nef44945428\nef45ce3035\nef5dde449d\nef5e770988\nef6359cea3\nef65268834\nef6cb5eae0\nef78972bc2\nef8cfcfc4f\nef96501dd0\nef9a2e976b\nefb24f950f\nefce0c1868\nefe5ac6901\nefe828affa\nefea4e0523\nf0268aa627\nf0483250c8\nf04cf99ee6\nf05b189097\nf08928c6d3\nf09d74856f\nf0a7607d63\nf0ad38da27\nf0c34e1213\nf0c7f86c29\nf0dfa18ba7\nf0eb3179f7\nf119bab27d\nf14409b6a3\nf1489baff4\nf14c18cf6a\nf15c607b92\nf1af214222\nf1b77bd309\nf1ba9e1a3e\nf1d99239eb\nf1dc710cf4\nf1ec5c08fa\nf22648fe12\nf22d21f1f1\nf233257395\nf23e95dbe5\nf2445b1572\nf253b3486d\nf277c7a6a4\nf2ab2b84d6\nf2b7c9b1f3\nf2b83d5ce5\nf2c276018f\nf2cfd94d64\nf2dd6e3add\nf2e7653f16\nf2f333ad06\nf2f55d6713\nf2fdb6abec\nf305a56d9f\nf3085d6570\nf3325c3338\nf3400f1204\nf34497c932\nf34a56525e\nf36483c824\nf3704d5663\nf3734c4913\nf38e5aa5b4\nf3986fba44\nf3a0ffc7d9\nf3b24a7d28\nf3e6c35ec3\nf3fc0ea80b\nf40a683fbe\nf4207ca554\nf4377499c2\nf46184f393\nf46c2d0a6d\nf46c364dca\nf46f7a0b63\nf46fe141b0\nf470b9aeb0\nf47eb7437f\nf48b535719\nf49e4866ac\nf4aa882cfd\nf4daa3dbd5\nf4dd51ac35\nf507a1b9dc\nf51c5ac84b\nf52104164b\nf54c67b9bb\nf5966cadd2\nf5bddf5598\nf5d85cfd17\nf5e2e7d6a0\nf5f051e9b4\nf5f8a93a76\nf6283e8af5\nf635e9568b\nf6474735be\nf659251be2\nf66981af4e\nf6708fa398\nf697fe8e8f\nf6adb12c42\nf6c7906ca4\nf6cd0a8016\nf6d6f15ae7\nf6e501892c\nf6f59d986f\nf6fe8c90a5\nf714160545\nf74c3888d7\nf7782c430e\nf7783ae5f2\nf77ab47923\nf788a98327\nf7961ac1f0\nf7a71e7574\nf7a8521432\nf7afbf4947\nf7b7cd5f44\nf7cf4b4a39\nf7d49799ad\nf7e0c9bb83\nf7e5b84928\nf7e6bd58be\nf7f2a38ac6\nf7f6cb2d6d\nf83f19e796\nf85796a921\nf8603c26b2\nf8819b42ec\nf891f8eaa1\nf89288d10c\nf895ae8cc1\nf8b4ac12f1\nf8c3fb2b01\nf8c8de2764\nf8db369b40\nf8fcb6a78c\nf94aafdeef\nf95d217b70\nf9681d5103\nf9750192a4\nf9823a32c2\nf991ddb4c2\nf99d535567\nf9ae3d98b7\nf9b6217959\nf9bd1fabf5\nf9c68eaa64\nf9d3e04c4f\nf9daf64494\nf9e4cc5a0a\nf9ea6b7f31\nf9f3852526\nfa04c615cf\nfa08e00a56\nfa4370d74d\nfa67744af3\nfa88d48a92\nfa8b904cc9\nfa9526bdf1\nfa9b9d2426\nfad633fbe1\nfaf5222dc3\nfaff0e15f1\nfb08c64e8c\nfb23455a7f\nfb2e19fa6e\nfb34dfbb77\nfb47fcea1e\nfb49738155\nfb4cbc514b\nfb4e6062f7\nfb5ba7ad6e\nfb63cd1236\nfb81157a07\nfb92abdaeb\nfba22a6848\nfbaca0c9df\nfbc645f602\nfbd77444cd\nfbe53dc8e8\nfbe541dd73\nfbe8488798\nfbfd25174f\nfc28cb305e\nfc33b1ffd6\nfc6186f0bb\nfc918e3a40\nfc96cda9d8\nfc9832eea4\nfcb10d0f81\nfcd20a2509\nfcf637e3ab\nfcfd81727f\nfd31890379\nfd33551c28\nfd542da05e\nfd6789b3fe\nfd77828200\nfd7af75f4d\nfdb28d0fbb\nfdb3d1fb1e\nfdb8b04124\nfdc6e3d581\nfdfce7e6fc\nfe0f76d41b\nfe24b0677d\nfe3c02699d\nfe58b48235\nfe6a5596b8\nfe6c244f63\nfe7afec086\nfe985d510a\nfe9db35d15\nfea8ffcd36\nfeb1080388\nfed208bfca\nfeda5ad1c2\nfeec95b386\nff15a5eff6\nff204daf4b\nff25f55852\nff2ada194f\nff2ce142e8\nff49d36d20\nff5a1ec4f3\nff66152b25\nff692fdc56\nff773b1a1e\nff97129478\nffb904207d\nffc43fc345\nfffe5f8df6\n" }, { "path": "merge_lora_weights_and_save_hf_model.py", "content": "import argparse\nimport glob\nimport os\nimport sys\n\nimport cv2\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nimport transformers\nfrom peft import LoraConfig, get_peft_model\nfrom transformers import AutoTokenizer\n\nfrom model.VISA import VISAForCausalLM\nfrom utils.utils import DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN\n\n\"\"\"\npython merge_lora_weights_and_save_hf_model.py \\\n --version /mnt/nlp-ali/usr/yancilin/clyan-data-2/video-llm/Chat-UniVi/Chat-UniVi \\\n --weight /mnt/public03/dataset/ovis/rgvos/visa7b/ckpt_model/pytorch_model15000.bin \\\n --save_path /mnt/public03/dataset/ovis/rgvos/visa7b/ckpt_model/hf_model\n\"\"\"\n\ndef parse_args(args):\n parser = argparse.ArgumentParser(\n description=\"merge lora weights and save model with hf format\"\n )\n parser.add_argument(\n \"--version\", default=\"Chat-UniVi/Chat-UniVi\", type=str, required=True\n )\n parser.add_argument(\n \"--weight\", default=\"/path/to/visa/pytorch_model.bin\", type=str, required=True\n )\n parser.add_argument(\n \"--save_path\", default=\"/path/to/hf_model\", type=str, required=True\n )\n parser.add_argument(\"--precision\", default=\"bf16\", type=str, choices=[\"fp32\", \"bf16\", \"fp16\"], help=\"precision for inference\", )\n parser.add_argument(\"--vision_pretrained\", default=\"PATH_TO_SAM_ViT-H\", type=str)\n parser.add_argument(\"--out_dim\", default=256, type=int)\n parser.add_argument(\"--image_size\", default=1024, type=int, help=\"image size\")\n parser.add_argument(\"--model_max_length\", default=512, type=int)\n parser.add_argument(\"--vision-tower\", default=\"openai/clip-vit-large-patch14\", type=str, )\n parser.add_argument(\"--lora_r\", default=8, type=int)\n parser.add_argument(\"--lora_alpha\", default=16, type=int)\n parser.add_argument(\"--lora_dropout\", default=0.05, type=float)\n parser.add_argument(\"--lora_target_modules\", default=\"q_proj,v_proj\", type=str)\n parser.add_argument(\"--local-rank\", default=0, type=int, help=\"node rank\")\n parser.add_argument(\"--train_mask_decoder\", action=\"store_true\", default=True)\n parser.add_argument(\"--use_mm_start_end\", action=\"store_true\", default=False)\n parser.add_argument(\"--conv_type\", default=\"llava_v1\", type=str, choices=[\"llava_v1\", \"llava_llama_2\"], )\n return parser.parse_args(args)\n\n\ndef main(args):\n args = parse_args(args)\n\n # Create model\n tokenizer = transformers.AutoTokenizer.from_pretrained(\n args.version,\n cache_dir = None,\n model_max_length = args.model_max_length,\n padding_side = \"right\",\n use_fast = False,\n )\n tokenizer.pad_token = tokenizer.unk_token\n num_added_tokens = tokenizer.add_tokens(\"[SEG]\")\n args.seg_token_idx = tokenizer(\"[SEG]\", add_special_tokens=False).input_ids[0]\n\n if args.use_mm_start_end:\n tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)\n\n model_args = {\n \"train_mask_decoder\": args.train_mask_decoder,\n \"out_dim\" : args.out_dim,\n \"seg_token_idx\" : args.seg_token_idx,\n \"vision_tower\" : args.vision_tower,\n }\n\n torch_dtype = torch.float32\n if args.precision == \"bf16\":\n torch_dtype = torch.bfloat16\n elif args.precision == \"fp16\":\n torch_dtype = torch.half\n model = VISAForCausalLM.from_pretrained(args.version, torch_dtype=torch_dtype, low_cpu_mem_usage=True, **model_args, )\n model.config.eos_token_id = tokenizer.eos_token_id\n model.config.bos_token_id = tokenizer.bos_token_id\n model.config.pad_token_id = tokenizer.pad_token_id\n\n model.get_model().initialize_vision_modules(model.get_model().config)\n vision_tower = model.get_model().get_vision_tower()\n vision_tower.to(dtype=torch_dtype)\n model.get_model().initialize_lisa_modules(model.get_model().config)\n\n lora_r = args.lora_r\n if lora_r > 0:\n\n def find_linear_layers(model, lora_target_modules):\n cls = torch.nn.Linear\n lora_module_names = set()\n for name, module in model.named_modules():\n if (\n isinstance(module, cls)\n and all([x not in name for x in [\"visual_model\", \"vision_tower\", \"mm_projector\", \"text_hidden_fcs\", ]])\n and any([x in name for x in lora_target_modules])\n ):\n lora_module_names.add(name)\n return sorted(list(lora_module_names))\n\n lora_alpha = args.lora_alpha\n lora_dropout = args.lora_dropout\n lora_target_modules = find_linear_layers(model, args.lora_target_modules.split(\",\"), )\n lora_config = LoraConfig(\n r = lora_r,\n lora_alpha = lora_alpha,\n target_modules = lora_target_modules,\n lora_dropout = lora_dropout,\n bias = \"none\",\n task_type = \"CAUSAL_LM\",\n )\n model = get_peft_model(model, lora_config)\n model.print_trainable_parameters()\n\n model.resize_token_embeddings(len(tokenizer))\n\n state_dict = torch.load(args.weight, map_location=\"cpu\")\n model.load_state_dict(state_dict, strict=True)\n\n model = model.merge_and_unload()\n state_dict = {}\n for k, v in model.state_dict().items():\n if \"vision_tower\" not in k:\n state_dict[k] = v\n model.save_pretrained(args.save_path, state_dict=state_dict)\n tokenizer.save_pretrained(args.save_path)\n\n\nif __name__ == \"__main__\":\n main(sys.argv[1:])\n" }, { "path": "model/VISA.py", "content": "from typing import List\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom transformers import BitsAndBytesConfig, CLIPVisionModel\n\nfrom utils.utils import DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IMAGE_PATCH_TOKEN\n\nfrom .univi.model.language_model.llama import ChatUniViLlamaForCausalLM, ChatUniViLlamaModel\n\nfrom .segment_anything import build_sam_vit_h\nfrom model.univi.constants import IMAGE_TOKEN_INDEX\n\nimport time\n\n\ndef dice_loss(\n inputs : torch.Tensor,\n targets : torch.Tensor,\n num_masks: float,\n scale : float =1000,\n eps : float =1e-6,\n):\n \"\"\"\n Compute the DICE loss, similar to generalized IOU for masks\n Args:\n inputs: A float tensor of arbitrary shape.\n The predictions for each example.\n targets: A float tensor with the same shape as inputs. Stores the binary classification label for each element in inputs (0 for the negative class and 1 for the positive class).\n \"\"\"\n inputs = inputs.sigmoid()\n inputs = inputs.flatten(1, 2)\n targets = targets.flatten(1, 2)\n numerator = 2 * (inputs / scale * targets).sum(-1)\n denominator = (inputs / scale).sum(-1) + (targets / scale).sum(-1)\n loss = 1 - (numerator + eps) / (denominator + eps)\n loss = loss.sum() / (num_masks + 1e-8)\n return loss\n\n\ndef sigmoid_ce_loss(\n inputs: torch.Tensor,\n targets: torch.Tensor,\n num_masks: float,\n):\n \"\"\"\n Args:\n inputs: A float tensor of arbitrary shape.\n The predictions for each example.\n targets: A float tensor with the same shape as inputs. Stores the binary classification label for each element in inputs (0 for the negative class and 1 for the positive class).\n Returns:\n Loss tensor\n \"\"\"\n loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction=\"none\")\n loss = loss.flatten(1, 2).mean(1).sum() / (num_masks + 1e-8)\n return loss\n\n\nclass VisaMetaModel:\n def __init__(\n self,\n config,\n **kwargs,\n ):\n super(VisaMetaModel, self).__init__(config)\n\n self.config = config\n if not hasattr(self.config, \"train_mask_decoder\"):\n self.config.train_mask_decoder = kwargs[\"train_mask_decoder\"]\n self.config.out_dim = kwargs[\"out_dim\"]\n self.vision_pretrained = kwargs.get(\"vision_pretrained\", None)\n else:\n self.vision_pretrained = kwargs.get(\"vision_pretrained\", None)\n self.initialize_lisa_modules(self.config)\n\n def initialize_lisa_modules(self, config):\n # SAM\n self.visual_model = build_sam_vit_h(self.vision_pretrained)\n for param in self.visual_model.parameters():\n param.requires_grad = False\n if config.train_mask_decoder:\n self.visual_model.mask_decoder.train()\n for param in self.visual_model.mask_decoder.parameters():\n param.requires_grad = True\n\n # Projection layer\n in_dim = config.hidden_size\n out_dim = config.out_dim\n text_fc = [\n nn.Linear(in_dim, in_dim),\n nn.ReLU(inplace=True),\n nn.Linear(in_dim, out_dim),\n nn.Dropout(0.0),\n ]\n self.text_hidden_fcs = nn.ModuleList([nn.Sequential(*text_fc)])\n self.text_hidden_fcs.train()\n for param in self.text_hidden_fcs.parameters():\n param.requires_grad = True\n\n\nclass VisaModel(VisaMetaModel, ChatUniViLlamaModel):\n def __init__(\n self,\n config,\n **kwargs,\n ):\n super(VisaModel, self).__init__(config, **kwargs)\n\n self.config.use_cache = False\n self.config.vision_tower = self.config.mm_vision_tower\n self.config.mm_vision_select_feature = \"patch\"\n self.config.image_aspect_ratio = \"square\"\n self.config.image_grid_pinpoints = None\n self.config.tune_mm_mlp_adapter = False\n self.config.freeze_mm_mlp_adapter = True\n self.config.pretrain_mm_mlp_adapter = None\n self.config.mm_use_im_patch_token = False\n\n\nclass VISAForCausalLM(ChatUniViLlamaForCausalLM):\n def __init__(\n self,\n config,\n **kwargs,\n ):\n if not hasattr(config, \"train_mask_decoder\"):\n config.mm_use_im_start_end = kwargs.pop(\"use_mm_start_end\", True)\n config.mm_vision_tower = kwargs.get(\"vision_tower\", \"openai/clip-vit-large-patch14\")\n self.ce_loss_weight = kwargs.pop(\"ce_loss_weight\", None)\n self.dice_loss_weight = kwargs.pop(\"dice_loss_weight\", None)\n self.bce_loss_weight = kwargs.pop(\"bce_loss_weight\", None)\n else:\n config.mm_vision_tower = config.vision_tower\n \n self.seg_token_idx = kwargs.pop(\"seg_token_idx\")\n\n super().__init__(config)\n\n self.model = VisaModel(config, **kwargs)\n\n self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)\n\n # Initialize weights and apply final processing\n self.post_init()\n\n def get_visual_embs(self, pixel_values: torch.FloatTensor):\n with torch.no_grad():\n image_embeddings = self.model.visual_model.image_encoder(pixel_values)\n return image_embeddings\n\n def forward(self, **kwargs):\n if \"past_key_values\" in kwargs:\n return super().forward(**kwargs)\n return self.model_forward(**kwargs)\n\n def model_forward(\n self,\n images: torch.FloatTensor,\n images_clip: torch.FloatTensor,\n input_ids: torch.LongTensor,\n labels: torch.LongTensor,\n attention_masks: torch.LongTensor,\n offset: torch.LongTensor,\n masks_list: List[torch.FloatTensor],\n label_list: List[torch.Tensor],\n resize_list: List[tuple],\n conversation_list: List[str], \n num_frame_list: List[int],\n num_conv_list: List[int],\n inference: bool = False,\n **kwargs,\n ):\n\n batch_size = len(images)\n image_embeddings = self.get_visual_embs(torch.cat(images,dim=0))\n assert image_embeddings.shape[0] == batch_size\n assert batch_size == len(offset) - 1\n for batch_idx in range(batch_size):\n assert num_conv_list[batch_idx] == offset[batch_idx + 1] - offset[batch_idx]\n\n if inference:\n length = input_ids.shape[0]\n assert len(images_clip) == 1, f'Inference only supports one video, but got {len(images_clip)} videos.'\n images_clip = [\n images_clip[0].unsqueeze(0).expand(length, -1, -1, -1, -1).contiguous().flatten(0,1)\n ]\n\n output_i = super().forward(\n images=images_clip,\n attention_mask=attention_masks,\n input_ids=input_ids,\n output_hidden_states=True,\n )\n torch.cuda.empty_cache()\n\n output_hidden_states = output_i.hidden_states\n output = None\n\n num_image_ori_token = (input_ids[0] == IMAGE_TOKEN_INDEX).sum()\n assert all(\n [\n (input_ids[i] == IMAGE_TOKEN_INDEX).sum() == num_image_ori_token for i in range(length)\n ]\n )\n token_add = 111 * num_image_ori_token\n\n seg_token_mask = input_ids[:, 1:] == self.seg_token_idx\n seg_token_mask = torch.cat([seg_token_mask, torch.zeros((seg_token_mask.shape[0], 1)).bool().cuda(), ], dim=1, )\n seg_token_mask = torch.cat([torch.zeros((seg_token_mask.shape[0], token_add)).bool().cuda(), seg_token_mask], dim=1, )\n all_conv_seg_token_num = seg_token_mask.sum(dim=1).tolist()\n\n else:\n images_clip_list = []\n for batch_idx in range(batch_size):\n bs_conv_num = num_conv_list[batch_idx]\n images_clip_i = images_clip[batch_idx].unsqueeze(0).expand(bs_conv_num, -1, -1, -1, -1).contiguous()\n images_clip_list.append(images_clip_i)\n images_clip_list = [i.flatten(0,1) for i in images_clip_list]\n\n output = super().forward(\n images=images_clip_list,\n attention_mask=attention_masks,\n input_ids=input_ids,\n labels=labels,\n output_hidden_states=True,\n )\n output_hidden_states = output.hidden_states\n\n\n seg_token_mask = output.labels[..., 1:] == self.seg_token_idx\n seg_token_mask = torch.cat([seg_token_mask, torch.zeros((seg_token_mask.shape[0], 1), device=output.labels.device).bool(), ], dim=1, )\n all_conv_seg_token_num = seg_token_mask.sum(dim=1).tolist()\n\n assert len(self.model.text_hidden_fcs) == 1\n\n pred_embeddings = self.model.text_hidden_fcs[0](output_hidden_states[-1][seg_token_mask])\n\n seg_token_counts = seg_token_mask.int().sum(-1) # [bs, ]\n\n seg_token_offset = seg_token_counts.cumsum(-1)\n seg_token_offset = torch.cat(\n [torch.zeros(1).long().cuda(), seg_token_offset], dim=0\n )\n\n seg_token_offset = seg_token_offset[offset]\n\n pred_embeddings_ = []\n for i in range(len(seg_token_offset) - 1):\n start_i, end_i = seg_token_offset[i], seg_token_offset[i + 1]\n pred_embeddings_.append(pred_embeddings[start_i:end_i])\n pred_embeddings = pred_embeddings_\n\n assert len(pred_embeddings) == batch_size\n\n multimask_output = False\n pred_masks = []\n for i in range(batch_size):\n\n (\n sparse_embeddings,\n dense_embeddings,\n ) = self.model.visual_model.prompt_encoder(\n points=None,\n boxes=None,\n masks=None,\n text_embeds=pred_embeddings[i].unsqueeze(1),\n )\n sparse_embeddings = sparse_embeddings.to(pred_embeddings[i].dtype)\n low_res_masks, iou_predictions = self.model.visual_model.mask_decoder(\n image_embeddings=image_embeddings[i].unsqueeze(0),\n image_pe=self.model.visual_model.prompt_encoder.get_dense_pe(),\n sparse_prompt_embeddings=sparse_embeddings,\n dense_prompt_embeddings=dense_embeddings,\n multimask_output=multimask_output,\n )\n pred_mask = self.model.visual_model.postprocess_masks(\n low_res_masks,\n input_size=resize_list[i],\n original_size=label_list[i].shape,\n )\n pred_masks.append(pred_mask[:, 0])\n\n\n\n model_output = output\n gt_masks = [mm.flatten(0, 1) for mm in masks_list]\n\n if inference:\n return {\n \"pred_masks\": pred_masks,\n \"gt_masks\": gt_masks,\n }\n\n output = model_output.logits\n\n ce_loss = model_output.loss\n ce_loss = ce_loss * self.ce_loss_weight\n mask_bce_loss = 0\n mask_dice_loss = 0\n num_masks = 0\n for batch_idx in range(batch_size):\n gt_mask = gt_masks[batch_idx]\n pred_mask = pred_masks[batch_idx]\n\n assert (\n gt_mask.shape[0] == pred_mask.shape[0]\n ), \"gt_mask.shape: {}, pred_mask.shape: {}\".format(\n gt_mask.shape, pred_mask.shape\n )\n mask_bce_loss += (\n sigmoid_ce_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])\n * gt_mask.shape[0]\n )\n mask_dice_loss += (\n dice_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])\n * gt_mask.shape[0]\n )\n num_masks += gt_mask.shape[0]\n\n mask_bce_loss = self.bce_loss_weight * mask_bce_loss / (num_masks + 1e-8)\n mask_dice_loss = self.dice_loss_weight * mask_dice_loss / (num_masks + 1e-8)\n mask_loss = mask_bce_loss + mask_dice_loss\n\n loss = ce_loss + mask_loss\n return {\n \"loss\": loss,\n \"ce_loss\": ce_loss,\n \"mask_bce_loss\": mask_bce_loss,\n \"mask_dice_loss\": mask_dice_loss,\n \"mask_loss\": mask_loss,\n }\n\n def evaluate(self, *args, **kwargs):\n raise NotImplementedError(\"This method is not implemented.\")" }, { "path": "model/llava/__init__.py", "content": "from .model import LlavaLlamaForCausalLM\n" }, { "path": "model/llava/constants.py", "content": "CONTROLLER_HEART_BEAT_EXPIRATION = 30\nWORKER_HEART_BEAT_INTERVAL = 15\n\nLOGDIR = \".\"\n\n# Model Constants\nIGNORE_INDEX = -100\nIMAGE_TOKEN_INDEX = -200\nDEFAULT_IMAGE_TOKEN = \"\"\nDEFAULT_IMAGE_PATCH_TOKEN = \"\"\nDEFAULT_IM_START_TOKEN = \"\"\nDEFAULT_IM_END_TOKEN = \"\"\n" }, { "path": "model/llava/conversation.py", "content": "import dataclasses\nfrom enum import Enum, auto\nfrom typing import List, Tuple\n\n\nclass SeparatorStyle(Enum):\n \"\"\"Different separator style.\"\"\"\n\n SINGLE = auto()\n TWO = auto()\n MPT = auto()\n PLAIN = auto()\n LLAMA_2 = auto()\n\n\n@dataclasses.dataclass\nclass Conversation:\n \"\"\"A class that keeps all conversation history.\"\"\"\n\n system: str\n roles: List[str]\n messages: List[List[str]]\n offset: int\n sep_style: SeparatorStyle = SeparatorStyle.SINGLE\n sep: str = \"###\"\n sep2: str = None\n version: str = \"Unknown\"\n\n skip_next: bool = False\n\n def get_prompt(self):\n messages = self.messages\n if len(messages) > 0 and type(messages[0][1]) is tuple:\n messages = self.messages.copy()\n init_role, init_msg = messages[0].copy()\n init_msg = init_msg[0].replace(\"\", \"\").strip()\n if \"mmtag\" in self.version:\n messages[0] = (init_role, init_msg)\n messages.insert(0, (self.roles[0], \"\"))\n messages.insert(1, (self.roles[1], \"Received.\"))\n else:\n messages[0] = (init_role, \"\\n\" + init_msg)\n\n if self.sep_style == SeparatorStyle.SINGLE:\n ret = self.system + self.sep\n for role, message in messages:\n if message:\n if type(message) is tuple:\n message, _, _ = message\n ret += role + \": \" + message + self.sep\n else:\n ret += role + \":\"\n elif self.sep_style == SeparatorStyle.TWO:\n seps = [self.sep, self.sep2]\n ret = self.system + seps[0]\n for i, (role, message) in enumerate(messages):\n if message:\n if type(message) is tuple:\n message, _, _ = message\n ret += role + \": \" + message + seps[i % 2]\n else:\n ret += role + \":\"\n elif self.sep_style == SeparatorStyle.MPT:\n ret = self.system + self.sep\n for role, message in messages:\n if message:\n if type(message) is tuple:\n message, _, _ = message\n ret += role + message + self.sep\n else:\n ret += role\n elif self.sep_style == SeparatorStyle.LLAMA_2:\n wrap_sys = lambda msg: f\"<>\\n{msg}\\n<>\\n\\n\"\n wrap_inst = lambda msg: f\"[INST] {msg} [/INST]\"\n ret = \"\"\n\n for i, (role, message) in enumerate(messages):\n if i == 0:\n assert message, \"first message should not be none\"\n assert role == self.roles[0], \"first message should come from user\"\n if message:\n if type(message) is tuple:\n message, _, _ = message\n if i == 0:\n message = wrap_sys(self.system) + message\n if i % 2 == 0:\n message = wrap_inst(message)\n ret += self.sep + message\n else:\n ret += \" \" + message + \" \" + self.sep2\n else:\n ret += \"\"\n ret = ret.lstrip(self.sep)\n elif self.sep_style == SeparatorStyle.PLAIN:\n seps = [self.sep, self.sep2]\n ret = self.system\n for i, (role, message) in enumerate(messages):\n if message:\n if type(message) is tuple:\n message, _, _ = message\n ret += message + seps[i % 2]\n else:\n ret += \"\"\n else:\n raise ValueError(f\"Invalid style: {self.sep_style}\")\n\n return ret\n\n def append_message(self, role, message):\n self.messages.append([role, message])\n\n def get_images(self, return_pil=False):\n images = []\n for i, (role, msg) in enumerate(self.messages[self.offset :]):\n if i % 2 == 0:\n if type(msg) is tuple:\n import base64\n from io import BytesIO\n\n from PIL import Image\n\n msg, image, image_process_mode = msg\n if image_process_mode == \"Pad\":\n\n def expand2square(pil_img, background_color=(122, 116, 104)):\n width, height = pil_img.size\n if width == height:\n return pil_img\n elif width > height:\n result = Image.new(\n pil_img.mode, (width, width), background_color\n )\n result.paste(pil_img, (0, (width - height) // 2))\n return result\n else:\n result = Image.new(\n pil_img.mode, (height, height), background_color\n )\n result.paste(pil_img, ((height - width) // 2, 0))\n return result\n\n image = expand2square(image)\n elif image_process_mode == \"Crop\":\n pass\n elif image_process_mode == \"Resize\":\n image = image.resize((336, 336))\n else:\n raise ValueError(\n f\"Invalid image_process_mode: {image_process_mode}\"\n )\n max_hw, min_hw = max(image.size), min(image.size)\n aspect_ratio = max_hw / min_hw\n max_len, min_len = 800, 400\n shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))\n longest_edge = int(shortest_edge * aspect_ratio)\n W, H = image.size\n if H > W:\n H, W = longest_edge, shortest_edge\n else:\n H, W = shortest_edge, longest_edge\n image = image.resize((W, H))\n if return_pil:\n images.append(image)\n else:\n buffered = BytesIO()\n image.save(buffered, format=\"PNG\")\n img_b64_str = base64.b64encode(buffered.getvalue()).decode()\n images.append(img_b64_str)\n return images\n\n def to_gradio_chatbot(self):\n ret = []\n for i, (role, msg) in enumerate(self.messages[self.offset :]):\n if i % 2 == 0:\n if type(msg) is tuple:\n import base64\n from io import BytesIO\n\n msg, image, image_process_mode = msg\n max_hw, min_hw = max(image.size), min(image.size)\n aspect_ratio = max_hw / min_hw\n max_len, min_len = 800, 400\n shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))\n longest_edge = int(shortest_edge * aspect_ratio)\n W, H = image.size\n if H > W:\n H, W = longest_edge, shortest_edge\n else:\n H, W = shortest_edge, longest_edge\n image = image.resize((W, H))\n buffered = BytesIO()\n image.save(buffered, format=\"JPEG\")\n img_b64_str = base64.b64encode(buffered.getvalue()).decode()\n img_str = f'\"user'\n ret.append([img_str, None])\n msg = msg.replace(\"\", \"\").strip()\n if len(msg) > 0:\n ret.append([msg, None])\n else:\n ret.append([msg, None])\n else:\n ret[-1][-1] = msg\n return ret\n\n def copy(self):\n return Conversation(\n system=self.system,\n roles=self.roles,\n messages=[[x, y] for x, y in self.messages],\n offset=self.offset,\n sep_style=self.sep_style,\n sep=self.sep,\n sep2=self.sep2,\n version=self.version,\n )\n\n def dict(self):\n if len(self.get_images()) > 0:\n return {\n \"system\": self.system,\n \"roles\": self.roles,\n \"messages\": [\n [x, y[0] if type(y) is tuple else y] for x, y in self.messages\n ],\n \"offset\": self.offset,\n \"sep\": self.sep,\n \"sep2\": self.sep2,\n }\n return {\n \"system\": self.system,\n \"roles\": self.roles,\n \"messages\": self.messages,\n \"offset\": self.offset,\n \"sep\": self.sep,\n \"sep2\": self.sep2,\n }\n\n\nconv_vicuna_v0 = Conversation(\n system=\"A chat between a curious human and an artificial intelligence assistant. \"\n \"The assistant gives helpful, detailed, and polite answers to the human's questions.\",\n roles=(\"Human\", \"Assistant\"),\n messages=(\n (\n \"Human\",\n \"What are the key differences between renewable and non-renewable energy sources?\",\n ),\n (\n \"Assistant\",\n \"Renewable energy sources are those that can be replenished naturally in a relatively \"\n \"short amount of time, such as solar, wind, hydro, geothermal, and biomass. \"\n \"Non-renewable energy sources, on the other hand, are finite and will eventually be \"\n \"depleted, such as coal, oil, and natural gas. Here are some key differences between \"\n \"renewable and non-renewable energy sources:\\n\"\n \"1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable \"\n \"energy sources are finite and will eventually run out.\\n\"\n \"2. Environmental impact: Renewable energy sources have a much lower environmental impact \"\n \"than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, \"\n \"and other negative effects.\\n\"\n \"3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically \"\n \"have lower operational costs than non-renewable sources.\\n\"\n \"4. Reliability: Renewable energy sources are often more reliable and can be used in more remote \"\n \"locations than non-renewable sources.\\n\"\n \"5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different \"\n \"situations and needs, while non-renewable sources are more rigid and inflexible.\\n\"\n \"6. Sustainability: Renewable energy sources are more sustainable over the long term, while \"\n \"non-renewable sources are not, and their depletion can lead to economic and social instability.\\n\",\n ),\n ),\n offset=2,\n sep_style=SeparatorStyle.SINGLE,\n sep=\"###\",\n)\n\nconv_vicuna_v1 = Conversation(\n system=\"A chat between a curious user and an artificial intelligence assistant. \"\n \"The assistant gives helpful, detailed, and polite answers to the user's questions.\",\n roles=(\"USER\", \"ASSISTANT\"),\n version=\"v1\",\n messages=(),\n offset=0,\n sep_style=SeparatorStyle.TWO,\n sep=\" \",\n sep2=\"\",\n)\n\nconv_llama_2 = Conversation(\n system=\"\"\"You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\"\"\",\n roles=(\"USER\", \"ASSISTANT\"),\n version=\"llama_v2\",\n messages=(),\n offset=0,\n sep_style=SeparatorStyle.LLAMA_2,\n sep=\"\",\n sep2=\"\",\n)\n\nconv_llava_llama_2 = Conversation(\n system=\"You are a helpful language and vision assistant. \"\n \"You are able to understand the visual content that the user provides, \"\n \"and assist the user with a variety of tasks using natural language.\",\n roles=(\"USER\", \"ASSISTANT\"),\n version=\"llama_v2\",\n messages=(),\n offset=0,\n sep_style=SeparatorStyle.LLAMA_2,\n sep=\"\",\n sep2=\"\",\n)\n\nconv_mpt = Conversation(\n system=\"\"\"<|im_start|>system\nA conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.\"\"\",\n roles=(\"<|im_start|>user\\n\", \"<|im_start|>assistant\\n\"),\n version=\"mpt\",\n messages=(),\n offset=0,\n sep_style=SeparatorStyle.MPT,\n sep=\"<|im_end|>\",\n)\n\nconv_llava_plain = Conversation(\n system=\"\",\n roles=(\"\", \"\"),\n messages=(),\n offset=0,\n sep_style=SeparatorStyle.PLAIN,\n sep=\"\\n\",\n)\n\nconv_llava_v0 = Conversation(\n system=\"A chat between a curious human and an artificial intelligence assistant. \"\n \"The assistant gives helpful, detailed, and polite answers to the human's questions.\",\n roles=(\"Human\", \"Assistant\"),\n messages=((\"Human\", \"Hi!\"), (\"Assistant\", \"Hi there! How can I help you today?\")),\n offset=2,\n sep_style=SeparatorStyle.SINGLE,\n sep=\"###\",\n)\n\nconv_llava_v0_mmtag = Conversation(\n system=\"A chat between a curious user and an artificial intelligence assistant. \"\n \"The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.\"\n \"The visual content will be provided with the following format: visual content.\",\n roles=(\"Human\", \"Assistant\"),\n messages=(),\n offset=0,\n sep_style=SeparatorStyle.SINGLE,\n sep=\"###\",\n version=\"v0_mmtag\",\n)\n\nconv_llava_v1 = Conversation(\n system=\"A chat between a curious human and an artificial intelligence assistant. \"\n \"The assistant gives helpful, detailed, and polite answers to the human's questions.\",\n roles=(\"USER\", \"ASSISTANT\"),\n version=\"v1\",\n messages=(),\n offset=0,\n sep_style=SeparatorStyle.TWO,\n sep=\" \",\n sep2=\"\",\n)\n\nconv_llava_v1_mmtag = Conversation(\n system=\"A chat between a curious user and an artificial intelligence assistant. \"\n \"The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.\"\n \"The visual content will be provided with the following format: visual content.\",\n roles=(\"USER\", \"ASSISTANT\"),\n messages=(),\n offset=0,\n sep_style=SeparatorStyle.TWO,\n sep=\" \",\n sep2=\"\",\n version=\"v1_mmtag\",\n)\n\ndefault_conversation = conv_vicuna_v0\nconv_templates = {\n \"default\": conv_vicuna_v0,\n \"v0\": conv_vicuna_v0,\n \"v1\": conv_vicuna_v1,\n \"vicuna_v1\": conv_vicuna_v1,\n \"llama_2\": conv_llama_2,\n \"plain\": conv_llava_plain,\n \"v0_plain\": conv_llava_plain,\n \"llava_v0\": conv_llava_v0,\n \"v0_mmtag\": conv_llava_v0_mmtag,\n \"llava_v1\": conv_llava_v1,\n \"v1_mmtag\": conv_llava_v1_mmtag,\n \"llava_llama_2\": conv_llava_llama_2,\n \"mpt\": conv_mpt,\n}\n\n\nif __name__ == \"__main__\":\n print(default_conversation.get_prompt())\n" }, { "path": "model/llava/mm_utils.py", "content": "import base64\nfrom io import BytesIO\n\nimport torch\nfrom PIL import Image\nfrom transformers import StoppingCriteria\n\nfrom .constants import IMAGE_TOKEN_INDEX\n\n\ndef load_image_from_base64(image):\n return Image.open(BytesIO(base64.b64decode(image)))\n\n\ndef process_images(images, image_processor, model_cfg):\n return image_processor(images, return_tensors=\"pt\")[\"pixel_values\"]\n\n\ndef tokenizer_image_token(\n prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None\n):\n prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split(\"\")]\n\n def insert_separator(X, sep):\n return [ele for sublist in zip(X, [sep] * len(X)) for ele in sublist][:-1]\n\n input_ids = []\n offset = 0\n if (\n len(prompt_chunks) > 0\n and len(prompt_chunks[0]) > 0\n and prompt_chunks[0][0] == tokenizer.bos_token_id\n ):\n offset = 1\n input_ids.append(prompt_chunks[0][0])\n\n for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):\n input_ids.extend(x[offset:])\n\n if return_tensors is not None:\n if return_tensors == \"pt\":\n return torch.tensor(input_ids, dtype=torch.long)\n raise ValueError(f\"Unsupported tensor type: {return_tensors}\")\n return input_ids\n\n\ndef get_model_name_from_path(model_path):\n model_path = model_path.strip(\"/\")\n model_paths = model_path.split(\"/\")\n if model_paths[-1].startswith(\"checkpoint-\"):\n return model_paths[-2] + \"_\" + model_paths[-1]\n else:\n return model_paths[-1]\n\n\nclass KeywordsStoppingCriteria(StoppingCriteria):\n def __init__(self, keywords, tokenizer, input_ids):\n self.keywords = keywords\n self.keyword_ids = []\n for keyword in keywords:\n cur_keyword_ids = tokenizer(keyword).input_ids\n if (\n len(cur_keyword_ids) > 1\n and cur_keyword_ids[0] == tokenizer.bos_token_id\n ):\n cur_keyword_ids = cur_keyword_ids[1:]\n self.keyword_ids.append(torch.tensor(cur_keyword_ids))\n self.tokenizer = tokenizer\n self.start_len = input_ids.shape[1]\n\n def __call__(\n self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs\n ) -> bool:\n assert output_ids.shape[0] == 1, \"Only support batch size 1 (yet)\" # TODO\n offset = min(output_ids.shape[1] - self.start_len, 3)\n self.keyword_ids = [\n keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids\n ]\n for keyword_id in self.keyword_ids:\n if output_ids[0, -keyword_id.shape[0] :] == keyword_id:\n return True\n outputs = self.tokenizer.batch_decode(\n output_ids[:, -offset:], skip_special_tokens=True\n )[0]\n for keyword in self.keywords:\n if keyword in outputs:\n return True\n return False\n" }, { "path": "model/llava/model/__init__.py", "content": "from .language_model.llava_llama import LlavaConfig, LlavaLlamaForCausalLM\nfrom .language_model.llava_mpt import LlavaMPTConfig, LlavaMPTForCausalLM\n" }, { "path": "model/llava/model/apply_delta.py", "content": "\"\"\"\nUsage:\npython3 -m fastchat.model.apply_delta --base ~/model_weights/llama-7b --target ~/model_weights/vicuna-7b --delta lmsys/vicuna-7b-delta\n\"\"\"\nimport argparse\n\nimport torch\nfrom llava import LlavaLlamaForCausalLM\nfrom tqdm import tqdm\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\n\ndef apply_delta(base_model_path, target_model_path, delta_path):\n print(\"Loading base model\")\n base = AutoModelForCausalLM.from_pretrained(\n base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True\n )\n\n print(\"Loading delta\")\n delta = LlavaLlamaForCausalLM.from_pretrained(\n delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True\n )\n delta_tokenizer = AutoTokenizer.from_pretrained(delta_path)\n\n print(\"Applying delta\")\n for name, param in tqdm(delta.state_dict().items(), desc=\"Applying delta\"):\n if name not in base.state_dict():\n assert name in [\n \"model.mm_projector.weight\",\n \"model.mm_projector.bias\",\n ], f\"{name} not in base model\"\n continue\n if param.data.shape == base.state_dict()[name].shape:\n param.data += base.state_dict()[name]\n else:\n assert name in [\n \"model.embed_tokens.weight\",\n \"lm_head.weight\",\n ], f\"{name} dimension mismatch: {param.data.shape} vs {base.state_dict()[name].shape}\"\n bparam = base.state_dict()[name]\n param.data[: bparam.shape[0], : bparam.shape[1]] += bparam\n\n print(\"Saving target model\")\n delta.save_pretrained(target_model_path)\n delta_tokenizer.save_pretrained(target_model_path)\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--base-model-path\", type=str, required=True)\n parser.add_argument(\"--target-model-path\", type=str, required=True)\n parser.add_argument(\"--delta-path\", type=str, required=True)\n\n args = parser.parse_args()\n\n apply_delta(args.base_model_path, args.target_model_path, args.delta_path)\n" }, { "path": "model/llava/model/builder.py", "content": "# Copyright 2023 Haotian Liu\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport os\nimport shutil\n\nimport torch\nfrom llava.constants import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,\n DEFAULT_IMAGE_PATCH_TOKEN)\nfrom llava.model import *\nfrom transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer,\n BitsAndBytesConfig)\n\n\ndef load_pretrained_model(\n model_path,\n model_base,\n model_name,\n load_8bit=False,\n load_4bit=False,\n device_map=\"auto\",\n):\n kwargs = {\"device_map\": device_map}\n\n if load_8bit:\n kwargs[\"load_in_8bit\"] = True\n elif load_4bit:\n kwargs[\"load_in_4bit\"] = True\n kwargs[\"quantization_config\"] = BitsAndBytesConfig(\n load_in_4bit=True,\n bnb_4bit_compute_dtype=torch.float16,\n bnb_4bit_use_double_quant=True,\n bnb_4bit_quant_type=\"nf4\",\n )\n else:\n kwargs[\"torch_dtype\"] = torch.float16\n\n if \"llava\" in model_name.lower():\n # Load LLaVA model\n if \"lora\" in model_name.lower() and model_base is not None:\n lora_cfg_pretrained = AutoConfig.from_pretrained(model_path)\n tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)\n print(\"Loading LLaVA from base model...\")\n model = LlavaLlamaForCausalLM.from_pretrained(\n model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, **kwargs\n )\n token_num, tokem_dim = model.lm_head.out_features, model.lm_head.in_features\n if model.lm_head.weight.shape[0] != token_num:\n model.lm_head.weight = torch.nn.Parameter(\n torch.empty(\n token_num, tokem_dim, device=model.device, dtype=model.dtype\n )\n )\n model.model.embed_tokens.weight = torch.nn.Parameter(\n torch.empty(\n token_num, tokem_dim, device=model.device, dtype=model.dtype\n )\n )\n\n print(\"Loading additional LLaVA weights...\")\n if os.path.exists(os.path.join(model_path, \"non_lora_trainables.bin\")):\n non_lora_trainables = torch.load(\n os.path.join(model_path, \"non_lora_trainables.bin\"),\n map_location=\"cpu\",\n )\n else:\n # this is probably from HF Hub\n from huggingface_hub import hf_hub_download\n\n def load_from_hf(repo_id, filename, subfolder=None):\n cache_file = hf_hub_download(\n repo_id=repo_id, filename=filename, subfolder=subfolder\n )\n return torch.load(cache_file, map_location=\"cpu\")\n\n non_lora_trainables = load_from_hf(\n model_path, \"non_lora_trainables.bin\"\n )\n non_lora_trainables = {\n (k[11:] if k.startswith(\"base_model.\") else k): v\n for k, v in non_lora_trainables.items()\n }\n if any(k.startswith(\"model.model.\") for k in non_lora_trainables):\n non_lora_trainables = {\n (k[6:] if k.startswith(\"model.\") else k): v\n for k, v in non_lora_trainables.items()\n }\n model.load_state_dict(non_lora_trainables, strict=False)\n\n from peft import PeftModel\n\n print(\"Loading LoRA weights...\")\n model = PeftModel.from_pretrained(model, model_path)\n print(\"Merging LoRA weights...\")\n model = model.merge_and_unload()\n print(\"Model is loaded...\")\n elif model_base is not None:\n # this may be mm projector only\n print(\"Loading LLaVA from base model...\")\n if \"mpt\" in model_name.lower():\n if not os.path.isfile(os.path.join(model_path, \"configuration_mpt.py\")):\n shutil.copyfile(\n os.path.join(model_base, \"configuration_mpt.py\"),\n os.path.join(model_path, \"configuration_mpt.py\"),\n )\n tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=True)\n cfg_pretrained = AutoConfig.from_pretrained(\n model_path, trust_remote_code=True\n )\n model = LlavaMPTForCausalLM.from_pretrained(\n model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs\n )\n else:\n tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)\n cfg_pretrained = AutoConfig.from_pretrained(model_path)\n model = LlavaLlamaForCausalLM.from_pretrained(\n model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs\n )\n\n mm_projector_weights = torch.load(\n os.path.join(model_path, \"mm_projector.bin\"), map_location=\"cpu\"\n )\n mm_projector_weights = {\n k: v.to(torch.float16) for k, v in mm_projector_weights.items()\n }\n model.load_state_dict(mm_projector_weights, strict=False)\n else:\n if \"mpt\" in model_name.lower():\n tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)\n model = LlavaMPTForCausalLM.from_pretrained(\n model_path, low_cpu_mem_usage=True, **kwargs\n )\n else:\n tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)\n model = LlavaLlamaForCausalLM.from_pretrained(\n model_path, low_cpu_mem_usage=True, **kwargs\n )\n else:\n # Load language model\n if model_base is not None:\n # PEFT model\n from peft import PeftModel\n\n tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)\n model = AutoModelForCausalLM.from_pretrained(\n model_base,\n torch_dtype=torch.float16,\n low_cpu_mem_usage=True,\n device_map=\"auto\",\n )\n print(f\"Loading LoRA weights from {model_path}\")\n model = PeftModel.from_pretrained(model, model_path)\n print(f\"Merging weights\")\n model = model.merge_and_unload()\n print(\"Convert to FP16...\")\n model.to(torch.float16)\n else:\n use_fast = False\n if \"mpt\" in model_name.lower():\n tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)\n model = AutoModelForCausalLM.from_pretrained(\n model_path, low_cpu_mem_usage=True, trust_remote_code=True, **kwargs\n )\n else:\n tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)\n model = AutoModelForCausalLM.from_pretrained(\n model_path, low_cpu_mem_usage=True, **kwargs\n )\n\n image_processor = None\n\n if \"llava\" in model_name.lower():\n mm_use_im_start_end = getattr(model.config, \"mm_use_im_start_end\", False)\n mm_use_im_patch_token = getattr(model.config, \"mm_use_im_patch_token\", True)\n if mm_use_im_patch_token:\n tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)\n if mm_use_im_start_end:\n tokenizer.add_tokens(\n [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True\n )\n model.resize_token_embeddings(len(tokenizer))\n\n vision_tower = model.get_vision_tower()\n if not vision_tower.is_loaded:\n vision_tower.load_model()\n vision_tower.to(device=\"cuda\", dtype=torch.float16)\n image_processor = vision_tower.image_processor\n\n if hasattr(model.config, \"max_sequence_length\"):\n context_len = model.config.max_sequence_length\n else:\n context_len = 2048\n\n return tokenizer, model, image_processor, context_len\n" }, { "path": "model/llava/model/consolidate.py", "content": "\"\"\"\nUsage:\npython3 -m llava.model.consolidate --src ~/model_weights/llava-7b --dst ~/model_weights/llava-7b_consolidate\n\"\"\"\nimport argparse\n\nimport torch\nfrom llava.model import *\nfrom llava.model.utils import auto_upgrade\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\n\ndef consolidate_ckpt(src_path, dst_path):\n print(\"Loading model\")\n auto_upgrade(src_path)\n src_model = AutoModelForCausalLM.from_pretrained(\n src_path, torch_dtype=torch.float16, low_cpu_mem_usage=True\n )\n src_tokenizer = AutoTokenizer.from_pretrained(src_path, use_fast=False)\n src_model.save_pretrained(dst_path)\n src_tokenizer.save_pretrained(dst_path)\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--src\", type=str, required=True)\n parser.add_argument(\"--dst\", type=str, required=True)\n\n args = parser.parse_args()\n\n consolidate_ckpt(args.src, args.dst)\n" }, { "path": "model/llava/model/language_model/llava_llama.py", "content": "# Copyright 2023 Haotian Liu\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom typing import List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn import CrossEntropyLoss\nfrom transformers import (AutoConfig, AutoModelForCausalLM, LlamaConfig,\n LlamaForCausalLM, LlamaModel)\nfrom transformers.modeling_outputs import CausalLMOutputWithPast\n\nfrom ..llava_arch import LlavaMetaForCausalLM, LlavaMetaModel\n\n\nclass LlavaConfig(LlamaConfig):\n model_type = \"llava\"\n\n\nclass LlavaLlamaModel(LlavaMetaModel, LlamaModel):\n config_class = LlavaConfig\n\n def __init__(self, config: LlamaConfig):\n super(LlavaLlamaModel, self).__init__(config)\n\n\nclass LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):\n config_class = LlavaConfig\n\n def __init__(self, config):\n super(LlamaForCausalLM, self).__init__(config)\n\n self.model = LlavaLlamaModel(config)\n\n self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)\n\n # Initialize weights and apply final processing\n self.post_init()\n\n def get_model(self):\n return self.model\n\n def forward(\n self,\n input_ids: torch.LongTensor = None,\n attention_mask: Optional[torch.Tensor] = None,\n past_key_values: Optional[List[torch.FloatTensor]] = None,\n inputs_embeds: Optional[torch.FloatTensor] = None,\n labels: Optional[torch.LongTensor] = None,\n use_cache: Optional[bool] = None,\n output_attentions: Optional[bool] = None,\n output_hidden_states: Optional[bool] = None,\n images: Optional[torch.FloatTensor] = None,\n return_dict: Optional[bool] = None,\n ) -> Union[Tuple, CausalLMOutputWithPast]:\n output_attentions = (\n output_attentions\n if output_attentions is not None\n else self.config.output_attentions\n )\n output_hidden_states = (\n output_hidden_states\n if output_hidden_states is not None\n else self.config.output_hidden_states\n )\n return_dict = (\n return_dict if return_dict is not None else self.config.use_return_dict\n )\n\n (\n input_ids,\n attention_mask,\n past_key_values,\n inputs_embeds,\n labels,\n ) = self.prepare_inputs_labels_for_multimodal(\n input_ids, attention_mask, past_key_values, labels, images\n )\n # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)\n\n outputs = self.model(\n input_ids=input_ids,\n attention_mask=attention_mask,\n past_key_values=past_key_values,\n inputs_embeds=inputs_embeds,\n use_cache=use_cache,\n output_attentions=output_attentions,\n output_hidden_states=output_hidden_states,\n return_dict=return_dict,\n )\n\n hidden_states = outputs[0]\n logits = self.lm_head(hidden_states)\n\n loss = None\n if labels is not None:\n # Shift so that tokens < n predict n\n shift_logits = logits[..., :-1, :].contiguous()\n shift_labels = labels[..., 1:].contiguous()\n # Flatten the tokens\n loss_fct = CrossEntropyLoss()\n shift_logits = shift_logits.view(-1, self.config.vocab_size)\n shift_labels = shift_labels.view(-1)\n # Enable model/pipeline parallelism\n shift_labels = shift_labels.to(shift_logits.device)\n loss = loss_fct(shift_logits, shift_labels)\n\n if not return_dict:\n output = (logits,) + outputs[1:]\n return (loss,) + output if loss is not None else output\n\n if self.training:\n output_hidden_states = outputs.hidden_states\n else:\n output_hidden_states = hidden_states\n\n return CausalLMOutputWithPast(\n loss=loss,\n logits=logits,\n past_key_values=outputs.past_key_values,\n hidden_states=output_hidden_states, # outputs.hidden_states,\n attentions=outputs.attentions,\n )\n\n def prepare_inputs_for_generation(\n self,\n input_ids,\n past_key_values=None,\n attention_mask=None,\n inputs_embeds=None,\n images=None,\n **kwargs\n ):\n if past_key_values:\n input_ids = input_ids[:, -1:]\n\n # if `inputs_embeds` are passed, we only want to use them in the 1st generation step\n if inputs_embeds is not None and past_key_values is None:\n model_inputs = {\"inputs_embeds\": inputs_embeds}\n else:\n model_inputs = {\"input_ids\": input_ids}\n\n model_inputs.update(\n {\n \"past_key_values\": past_key_values,\n \"use_cache\": kwargs.get(\"use_cache\"),\n \"attention_mask\": attention_mask,\n \"images\": images,\n }\n )\n return model_inputs\n\n\nAutoConfig.register(\"llava\", LlavaConfig)\nAutoModelForCausalLM.register(LlavaConfig, LlavaLlamaForCausalLM)\n" }, { "path": "model/llava/model/language_model/llava_mpt.py", "content": "# Copyright 2023 Haotian Liu\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport math\nimport warnings\nfrom typing import List, Optional, Tuple\n\nimport torch\nimport torch.nn.functional as F\nfrom transformers import AutoConfig, AutoModelForCausalLM\nfrom transformers.modeling_outputs import CausalLMOutputWithPast\n\nfrom ..llava_arch import LlavaMetaForCausalLM, LlavaMetaModel\nfrom .mpt.modeling_mpt import MPTConfig, MPTForCausalLM, MPTModel\n\n\nclass LlavaMPTConfig(MPTConfig):\n model_type = \"llava_mpt\"\n\n\nclass LlavaMPTModel(LlavaMetaModel, MPTModel):\n config_class = LlavaMPTConfig\n\n def __init__(self, config: MPTConfig):\n config.hidden_size = config.d_model\n super(LlavaMPTModel, self).__init__(config)\n\n def embed_tokens(self, x):\n return self.wte(x)\n\n\nclass LlavaMPTForCausalLM(MPTForCausalLM, LlavaMetaForCausalLM):\n config_class = LlavaMPTConfig\n supports_gradient_checkpointing = True\n\n def __init__(self, config):\n super(MPTForCausalLM, self).__init__(config)\n\n if not config.tie_word_embeddings:\n raise ValueError(\"MPTForCausalLM only supports tied word embeddings\")\n self.transformer = LlavaMPTModel(config)\n self.logit_scale = None\n if config.logit_scale is not None:\n logit_scale = config.logit_scale\n if isinstance(logit_scale, str):\n if logit_scale == \"inv_sqrt_d_model\":\n logit_scale = 1 / math.sqrt(config.d_model)\n else:\n raise ValueError(\n f\"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.\"\n )\n self.logit_scale = logit_scale\n\n def get_model(self):\n return self.transformer\n\n def _set_gradient_checkpointing(self, module, value=False):\n if isinstance(module, LlavaMPTModel):\n module.gradient_checkpointing = value\n\n def forward(\n self,\n input_ids: torch.LongTensor,\n past_key_values: Optional[List[Tuple[torch.FloatTensor]]] = None,\n attention_mask: Optional[torch.ByteTensor] = None,\n prefix_mask: Optional[torch.ByteTensor] = None,\n sequence_id: Optional[torch.LongTensor] = None,\n labels: Optional[torch.LongTensor] = None,\n return_dict: Optional[bool] = None,\n output_attentions: Optional[bool] = None,\n output_hidden_states: Optional[bool] = None,\n use_cache: Optional[bool] = None,\n images=None,\n ):\n return_dict = (\n return_dict if return_dict is not None else self.config.return_dict\n )\n use_cache = use_cache if use_cache is not None else self.config.use_cache\n\n (\n input_ids,\n attention_mask,\n past_key_values,\n inputs_embeds,\n labels,\n ) = self.prepare_inputs_labels_for_multimodal(\n input_ids, attention_mask, past_key_values, labels, images\n )\n outputs = self.transformer(\n input_ids=input_ids,\n inputs_embeds=inputs_embeds,\n past_key_values=past_key_values,\n attention_mask=attention_mask,\n prefix_mask=prefix_mask,\n sequence_id=sequence_id,\n return_dict=return_dict,\n output_attentions=output_attentions,\n output_hidden_states=output_hidden_states,\n use_cache=use_cache,\n )\n # FIXME: this is a hack to fix the multiple gpu inference issue in https://github.com/haotian-liu/LLaVA/issues/338\n logits = F.linear(\n outputs.last_hidden_state.to(self.transformer.wte.weight.device),\n self.transformer.wte.weight,\n )\n if self.logit_scale is not None:\n if self.logit_scale == 0:\n warnings.warn(\n f\"Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.\"\n )\n logits *= self.logit_scale\n loss = None\n if labels is not None:\n labels = torch.roll(labels, shifts=-1)\n labels[:, -1] = -100\n loss = F.cross_entropy(\n logits.view(-1, logits.size(-1)), labels.to(logits.device).view(-1)\n )\n return CausalLMOutputWithPast(\n loss=loss,\n logits=logits,\n past_key_values=outputs.past_key_values,\n hidden_states=outputs.hidden_states,\n )\n\n def prepare_inputs_for_generation(\n self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs\n ):\n if inputs_embeds is not None:\n raise NotImplementedError(\"inputs_embeds is not implemented for MPT yet\")\n attention_mask = kwargs[\"attention_mask\"].bool()\n if attention_mask[:, -1].sum() != attention_mask.shape[0]:\n raise NotImplementedError(\n \"MPT does not support generation with right padding.\"\n )\n if self.transformer.attn_uses_sequence_id and self.training:\n sequence_id = torch.zeros_like(input_ids[:1])\n else:\n sequence_id = None\n if past_key_values is not None:\n input_ids = input_ids[:, -1].unsqueeze(-1)\n if self.transformer.prefix_lm:\n prefix_mask = torch.ones_like(attention_mask)\n if kwargs.get(\"use_cache\") == False:\n raise NotImplementedError(\n \"MPT with prefix_lm=True does not support use_cache=False.\"\n )\n else:\n prefix_mask = None\n return {\n \"input_ids\": input_ids,\n \"attention_mask\": attention_mask,\n \"prefix_mask\": prefix_mask,\n \"sequence_id\": sequence_id,\n \"past_key_values\": past_key_values,\n \"use_cache\": kwargs.get(\"use_cache\", True),\n \"images\": kwargs.get(\"images\", None),\n }\n\n\nAutoConfig.register(\"llava_mpt\", LlavaMPTConfig)\nAutoModelForCausalLM.register(LlavaMPTConfig, LlavaMPTForCausalLM)\n" }, { "path": "model/llava/model/language_model/mpt/adapt_tokenizer.py", "content": "from typing import Union\n\nfrom transformers import (AutoTokenizer, PreTrainedTokenizer,\n PreTrainedTokenizerFast)\n\nTokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]\nNUM_SENTINEL_TOKENS: int = 100\n\n\ndef adapt_tokenizer_for_denoising(tokenizer: Tokenizer):\n \"\"\"Adds sentinel tokens and padding token (if missing).\n\n Expands the tokenizer vocabulary to include sentinel tokens\n used in mixture-of-denoiser tasks as well as a padding token.\n\n All added tokens are added as special tokens. No tokens are\n added if sentinel tokens and padding token already exist.\n \"\"\"\n sentinels_to_add = [f\"\" for i in range(NUM_SENTINEL_TOKENS)]\n tokenizer.add_tokens(sentinels_to_add, special_tokens=True)\n if tokenizer.pad_token is None:\n tokenizer.add_tokens(\"\", special_tokens=True)\n tokenizer.pad_token = \"\"\n assert tokenizer.pad_token_id is not None\n sentinels = \"\".join([f\"\" for i in range(NUM_SENTINEL_TOKENS)])\n _sentinel_token_ids = tokenizer(sentinels, add_special_tokens=False).input_ids\n tokenizer.sentinel_token_ids = _sentinel_token_ids\n\n\nclass AutoTokenizerForMOD(AutoTokenizer):\n \"\"\"AutoTokenizer + Adaptation for MOD.\n\n A simple wrapper around AutoTokenizer to make instantiating\n an MOD-adapted tokenizer a bit easier.\n\n MOD-adapted tokenizers have sentinel tokens (e.g., ),\n a padding token, and a property to get the token ids of the\n sentinel tokens.\n \"\"\"\n\n @classmethod\n def from_pretrained(cls, *args, **kwargs):\n \"\"\"See `AutoTokenizer.from_pretrained` docstring.\"\"\"\n tokenizer = super().from_pretrained(*args, **kwargs)\n adapt_tokenizer_for_denoising(tokenizer)\n return tokenizer\n" }, { "path": "model/llava/model/language_model/mpt/attention.py", "content": "\"\"\"Attention layers.\"\"\"\nimport math\nimport warnings\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\nfrom einops import rearrange\nfrom packaging import version\nfrom torch import nn\n\nfrom .norm import LPLayerNorm\n\n\ndef _reset_is_causal(\n num_query_tokens: int, num_key_tokens: int, original_is_causal: bool\n):\n if original_is_causal and num_query_tokens != num_key_tokens:\n if num_query_tokens != 1:\n raise NotImplementedError(\n \"MPT does not support query and key with different number of tokens, unless number of query tokens is 1.\"\n )\n else:\n return False\n return original_is_causal\n\n\ndef scaled_multihead_dot_product_attention(\n query,\n key,\n value,\n n_heads,\n past_key_value=None,\n softmax_scale=None,\n attn_bias=None,\n key_padding_mask=None,\n is_causal=False,\n dropout_p=0.0,\n training=False,\n needs_weights=False,\n multiquery=False,\n):\n q = rearrange(query, \"b s (h d) -> b h s d\", h=n_heads)\n kv_n_heads = 1 if multiquery else n_heads\n k = rearrange(key, \"b s (h d) -> b h d s\", h=kv_n_heads)\n v = rearrange(value, \"b s (h d) -> b h s d\", h=kv_n_heads)\n if past_key_value is not None:\n if len(past_key_value) != 0:\n k = torch.cat([past_key_value[0], k], dim=3)\n v = torch.cat([past_key_value[1], v], dim=2)\n past_key_value = (k, v)\n (b, _, s_q, d) = q.shape\n s_k = k.size(-1)\n if softmax_scale is None:\n softmax_scale = 1 / math.sqrt(d)\n attn_weight = q.matmul(k) * softmax_scale\n if attn_bias is not None:\n _s_q = max(0, attn_bias.size(2) - s_q)\n _s_k = max(0, attn_bias.size(3) - s_k)\n attn_bias = attn_bias[:, :, _s_q:, _s_k:]\n if (\n attn_bias.size(-1) != 1\n and attn_bias.size(-1) != s_k\n or (attn_bias.size(-2) != 1 and attn_bias.size(-2) != s_q)\n ):\n raise RuntimeError(\n f\"attn_bias (shape: {attn_bias.shape}) is expected to broadcast to shape: {attn_weight.shape}.\"\n )\n attn_weight = attn_weight + attn_bias\n min_val = torch.finfo(q.dtype).min\n if key_padding_mask is not None:\n if attn_bias is not None:\n warnings.warn(\n \"Propogating key_padding_mask to the attention module \"\n + \"and applying it within the attention module can cause \"\n + \"unneccessary computation/memory usage. Consider integrating \"\n + \"into attn_bias once and passing that to each attention \"\n + \"module instead.\"\n )\n attn_weight = attn_weight.masked_fill(\n ~key_padding_mask.view((b, 1, 1, s_k)), min_val\n )\n if is_causal and (not q.size(2) == 1):\n s = max(s_q, s_k)\n causal_mask = attn_weight.new_ones(s, s, dtype=torch.float16)\n causal_mask = causal_mask.tril()\n causal_mask = causal_mask.to(torch.bool)\n causal_mask = ~causal_mask\n causal_mask = causal_mask[-s_q:, -s_k:]\n attn_weight = attn_weight.masked_fill(causal_mask.view(1, 1, s_q, s_k), min_val)\n attn_weight = torch.softmax(attn_weight, dim=-1)\n if dropout_p:\n attn_weight = torch.nn.functional.dropout(\n attn_weight, p=dropout_p, training=training, inplace=True\n )\n out = attn_weight.to(v.dtype).matmul(v)\n out = rearrange(out, \"b h s d -> b s (h d)\")\n if needs_weights:\n return (out, attn_weight, past_key_value)\n return (out, None, past_key_value)\n\n\ndef check_valid_inputs(*tensors, valid_dtypes=[torch.float16, torch.bfloat16]):\n for tensor in tensors:\n if tensor.dtype not in valid_dtypes:\n raise TypeError(\n f\"tensor.dtype={tensor.dtype!r} must be in valid_dtypes={valid_dtypes!r}.\"\n )\n if not tensor.is_cuda:\n raise TypeError(\n f\"Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).\"\n )\n\n\ndef flash_attn_fn(\n query,\n key,\n value,\n n_heads,\n past_key_value=None,\n softmax_scale=None,\n attn_bias=None,\n key_padding_mask=None,\n is_causal=False,\n dropout_p=0.0,\n training=False,\n needs_weights=False,\n multiquery=False,\n):\n try:\n from flash_attn import bert_padding, flash_attn_interface\n except:\n raise RuntimeError(\"Please install flash-attn==1.0.3.post0\")\n check_valid_inputs(query, key, value)\n if past_key_value is not None:\n if len(past_key_value) != 0:\n key = torch.cat([past_key_value[0], key], dim=1)\n value = torch.cat([past_key_value[1], value], dim=1)\n past_key_value = (key, value)\n if attn_bias is not None:\n _s_q = max(0, attn_bias.size(2) - query.size(1))\n _s_k = max(0, attn_bias.size(3) - key.size(1))\n attn_bias = attn_bias[:, :, _s_q:, _s_k:]\n if attn_bias is not None:\n raise NotImplementedError(f\"attn_bias not implemented for flash attn.\")\n (batch_size, seqlen) = query.shape[:2]\n if key_padding_mask is None:\n key_padding_mask = torch.ones_like(key[:, :, 0], dtype=torch.bool)\n query_padding_mask = key_padding_mask[:, -query.size(1) :]\n (query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(\n query, query_padding_mask\n )\n query_unpad = rearrange(query_unpad, \"nnz (h d) -> nnz h d\", h=n_heads)\n (key_unpad, _, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(\n key, key_padding_mask\n )\n key_unpad = rearrange(\n key_unpad, \"nnz (h d) -> nnz h d\", h=1 if multiquery else n_heads\n )\n (value_unpad, _, _, _) = bert_padding.unpad_input(value, key_padding_mask)\n value_unpad = rearrange(\n value_unpad, \"nnz (h d) -> nnz h d\", h=1 if multiquery else n_heads\n )\n if multiquery:\n key_unpad = key_unpad.expand(key_unpad.size(0), n_heads, key_unpad.size(-1))\n value_unpad = value_unpad.expand(\n value_unpad.size(0), n_heads, value_unpad.size(-1)\n )\n dropout_p = dropout_p if training else 0.0\n reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)\n output_unpad = flash_attn_interface.flash_attn_unpadded_func(\n query_unpad,\n key_unpad,\n value_unpad,\n cu_seqlens_q,\n cu_seqlens_k,\n max_seqlen_q,\n max_seqlen_k,\n dropout_p,\n softmax_scale=softmax_scale,\n causal=reset_is_causal,\n return_attn_probs=needs_weights,\n )\n output = bert_padding.pad_input(\n rearrange(output_unpad, \"nnz h d -> nnz (h d)\"), indices_q, batch_size, seqlen\n )\n return (output, None, past_key_value)\n\n\ndef triton_flash_attn_fn(\n query,\n key,\n value,\n n_heads,\n past_key_value=None,\n softmax_scale=None,\n attn_bias=None,\n key_padding_mask=None,\n is_causal=False,\n dropout_p=0.0,\n training=False,\n needs_weights=False,\n multiquery=False,\n):\n try:\n from .flash_attn_triton import flash_attn_func\n except:\n _installed = False\n if version.parse(torch.__version__) < version.parse(\"2.0.0\"):\n _installed = True\n try:\n from flash_attn.flash_attn_triton import flash_attn_func\n except:\n _installed = False\n if not _installed:\n raise RuntimeError(\n \"Requirements for `attn_impl: triton` not installed. Either (1) have a CUDA-compatible GPU and `pip install .[gpu]` if installing from llm-foundry source or `pip install triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir#subdirectory=python` if installing from pypi, or (2) use torch attn model.attn_config.attn_impl=torch (torch attn_impl will be slow). Note: (1) requires you have CMake and PyTorch already installed.\"\n )\n check_valid_inputs(query, key, value)\n if past_key_value is not None:\n if len(past_key_value) != 0:\n key = torch.cat([past_key_value[0], key], dim=1)\n value = torch.cat([past_key_value[1], value], dim=1)\n past_key_value = (key, value)\n if attn_bias is not None:\n _s_q = max(0, attn_bias.size(2) - query.size(1))\n _s_k = max(0, attn_bias.size(3) - key.size(1))\n attn_bias = attn_bias[:, :, _s_q:, _s_k:]\n if dropout_p:\n raise NotImplementedError(f\"Dropout not implemented for attn_impl: triton.\")\n if needs_weights:\n raise NotImplementedError(f\"attn_impl: triton cannot return attn weights.\")\n if key_padding_mask is not None:\n warnings.warn(\n \"Propagating key_padding_mask to the attention module \"\n + \"and applying it within the attention module can cause \"\n + \"unnecessary computation/memory usage. Consider integrating \"\n + \"into attn_bias once and passing that to each attention \"\n + \"module instead.\"\n )\n (b_size, s_k) = key_padding_mask.shape[:2]\n if attn_bias is None:\n attn_bias = query.new_zeros(b_size, 1, 1, s_k)\n attn_bias = attn_bias.masked_fill(\n ~key_padding_mask.view((b_size, 1, 1, s_k)), torch.finfo(query.dtype).min\n )\n query = rearrange(query, \"b s (h d) -> b s h d\", h=n_heads)\n key = rearrange(key, \"b s (h d) -> b s h d\", h=1 if multiquery else n_heads)\n value = rearrange(value, \"b s (h d) -> b s h d\", h=1 if multiquery else n_heads)\n if multiquery:\n key = key.expand(*key.shape[:2], n_heads, key.size(-1))\n value = value.expand(*value.shape[:2], n_heads, value.size(-1))\n reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)\n attn_output = flash_attn_func(\n query, key, value, attn_bias, reset_is_causal, softmax_scale\n )\n output = attn_output.view(*attn_output.shape[:2], -1)\n return (output, None, past_key_value)\n\n\nclass MultiheadAttention(nn.Module):\n \"\"\"Multi-head self attention.\n\n Using torch or triton attention implemetation enables user to also use\n additive bias.\n \"\"\"\n\n def __init__(\n self,\n d_model: int,\n n_heads: int,\n attn_impl: str = \"triton\",\n clip_qkv: Optional[float] = None,\n qk_ln: bool = False,\n softmax_scale: Optional[float] = None,\n attn_pdrop: float = 0.0,\n low_precision_layernorm: bool = False,\n verbose: int = 0,\n device: Optional[str] = None,\n ):\n super().__init__()\n self.attn_impl = attn_impl\n self.clip_qkv = clip_qkv\n self.qk_ln = qk_ln\n self.d_model = d_model\n self.n_heads = n_heads\n self.softmax_scale = softmax_scale\n if self.softmax_scale is None:\n self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)\n self.attn_dropout_p = attn_pdrop\n self.Wqkv = nn.Linear(self.d_model, 3 * self.d_model, device=device)\n fuse_splits = (d_model, 2 * d_model)\n self.Wqkv._fused = (0, fuse_splits)\n if self.qk_ln:\n layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm\n self.q_ln = layernorm_class(self.d_model, device=device)\n self.k_ln = layernorm_class(self.d_model, device=device)\n if self.attn_impl == \"flash\":\n self.attn_fn = flash_attn_fn\n elif self.attn_impl == \"triton\":\n self.attn_fn = triton_flash_attn_fn\n if verbose:\n warnings.warn(\n \"While `attn_impl: triton` can be faster than `attn_impl: flash` \"\n + \"it uses more memory. When training larger models this can trigger \"\n + \"alloc retries which hurts performance. If encountered, we recommend \"\n + \"using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.\"\n )\n elif self.attn_impl == \"torch\":\n self.attn_fn = scaled_multihead_dot_product_attention\n if torch.cuda.is_available() and verbose:\n warnings.warn(\n \"Using `attn_impl: torch`. If your model does not use `alibi` or \"\n + \"`prefix_lm` we recommend using `attn_impl: flash` otherwise \"\n + \"we recommend using `attn_impl: triton`.\"\n )\n else:\n raise ValueError(f\"attn_impl={attn_impl!r} is an invalid setting.\")\n self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)\n self.out_proj._is_residual = True\n\n def forward(\n self,\n x,\n past_key_value=None,\n attn_bias=None,\n attention_mask=None,\n is_causal=True,\n needs_weights=False,\n ):\n qkv = self.Wqkv(x)\n if self.clip_qkv:\n qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)\n (query, key, value) = qkv.chunk(3, dim=2)\n key_padding_mask = attention_mask\n if self.qk_ln:\n dtype = query.dtype\n query = self.q_ln(query).to(dtype)\n key = self.k_ln(key).to(dtype)\n (context, attn_weights, past_key_value) = self.attn_fn(\n query,\n key,\n value,\n self.n_heads,\n past_key_value=past_key_value,\n softmax_scale=self.softmax_scale,\n attn_bias=attn_bias,\n key_padding_mask=key_padding_mask,\n is_causal=is_causal,\n dropout_p=self.attn_dropout_p,\n training=self.training,\n needs_weights=needs_weights,\n )\n return (self.out_proj(context), attn_weights, past_key_value)\n\n\nclass MultiQueryAttention(nn.Module):\n \"\"\"Multi-Query self attention.\n\n Using torch or triton attention implemetation enables user to also use\n additive bias.\n \"\"\"\n\n def __init__(\n self,\n d_model: int,\n n_heads: int,\n attn_impl: str = \"triton\",\n clip_qkv: Optional[float] = None,\n qk_ln: bool = False,\n softmax_scale: Optional[float] = None,\n attn_pdrop: float = 0.0,\n low_precision_layernorm: bool = False,\n verbose: int = 0,\n device: Optional[str] = None,\n ):\n super().__init__()\n self.attn_impl = attn_impl\n self.clip_qkv = clip_qkv\n self.qk_ln = qk_ln\n self.d_model = d_model\n self.n_heads = n_heads\n self.head_dim = d_model // n_heads\n self.softmax_scale = softmax_scale\n if self.softmax_scale is None:\n self.softmax_scale = 1 / math.sqrt(self.head_dim)\n self.attn_dropout_p = attn_pdrop\n self.Wqkv = nn.Linear(d_model, d_model + 2 * self.head_dim, device=device)\n fuse_splits = (d_model, d_model + self.head_dim)\n self.Wqkv._fused = (0, fuse_splits)\n if self.qk_ln:\n layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm\n self.q_ln = layernorm_class(d_model, device=device)\n self.k_ln = layernorm_class(self.head_dim, device=device)\n if self.attn_impl == \"flash\":\n self.attn_fn = flash_attn_fn\n elif self.attn_impl == \"triton\":\n self.attn_fn = triton_flash_attn_fn\n if verbose:\n warnings.warn(\n \"While `attn_impl: triton` can be faster than `attn_impl: flash` \"\n + \"it uses more memory. When training larger models this can trigger \"\n + \"alloc retries which hurts performance. If encountered, we recommend \"\n + \"using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.\"\n )\n elif self.attn_impl == \"torch\":\n self.attn_fn = scaled_multihead_dot_product_attention\n if torch.cuda.is_available() and verbose:\n warnings.warn(\n \"Using `attn_impl: torch`. If your model does not use `alibi` or \"\n + \"`prefix_lm` we recommend using `attn_impl: flash` otherwise \"\n + \"we recommend using `attn_impl: triton`.\"\n )\n else:\n raise ValueError(f\"attn_impl={attn_impl!r} is an invalid setting.\")\n self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)\n self.out_proj._is_residual = True\n\n def forward(\n self,\n x,\n past_key_value=None,\n attn_bias=None,\n attention_mask=None,\n is_causal=True,\n needs_weights=False,\n ):\n qkv = self.Wqkv(x)\n if self.clip_qkv:\n qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)\n (query, key, value) = qkv.split(\n [self.d_model, self.head_dim, self.head_dim], dim=2\n )\n key_padding_mask = attention_mask\n if self.qk_ln:\n dtype = query.dtype\n query = self.q_ln(query).to(dtype)\n key = self.k_ln(key).to(dtype)\n (context, attn_weights, past_key_value) = self.attn_fn(\n query,\n key,\n value,\n self.n_heads,\n past_key_value=past_key_value,\n softmax_scale=self.softmax_scale,\n attn_bias=attn_bias,\n key_padding_mask=key_padding_mask,\n is_causal=is_causal,\n dropout_p=self.attn_dropout_p,\n training=self.training,\n needs_weights=needs_weights,\n multiquery=True,\n )\n return (self.out_proj(context), attn_weights, past_key_value)\n\n\ndef attn_bias_shape(\n attn_impl, n_heads, seq_len, alibi, prefix_lm, causal, use_sequence_id\n):\n if attn_impl == \"flash\":\n return None\n elif attn_impl in [\"torch\", \"triton\"]:\n if alibi:\n if (prefix_lm or not causal) or use_sequence_id:\n return (1, n_heads, seq_len, seq_len)\n return (1, n_heads, 1, seq_len)\n elif prefix_lm or use_sequence_id:\n return (1, 1, seq_len, seq_len)\n return None\n else:\n raise ValueError(f\"attn_impl={attn_impl!r} is an invalid setting.\")\n\n\ndef build_attn_bias(\n attn_impl, attn_bias, n_heads, seq_len, causal=False, alibi=False, alibi_bias_max=8\n):\n if attn_impl == \"flash\":\n return None\n elif attn_impl in [\"torch\", \"triton\"]:\n if alibi:\n (device, dtype) = (attn_bias.device, attn_bias.dtype)\n attn_bias = attn_bias.add(\n build_alibi_bias(\n n_heads,\n seq_len,\n full=not causal,\n alibi_bias_max=alibi_bias_max,\n device=device,\n dtype=dtype,\n )\n )\n return attn_bias\n else:\n raise ValueError(f\"attn_impl={attn_impl!r} is an invalid setting.\")\n\n\ndef gen_slopes(n_heads, alibi_bias_max=8, device=None):\n _n_heads = 2 ** math.ceil(math.log2(n_heads))\n m = torch.arange(1, _n_heads + 1, dtype=torch.float32, device=device)\n m = m.mul(alibi_bias_max / _n_heads)\n slopes = 1.0 / torch.pow(2, m)\n if _n_heads != n_heads:\n slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]\n return slopes.view(1, n_heads, 1, 1)\n\n\ndef build_alibi_bias(\n n_heads, seq_len, full=False, alibi_bias_max=8, device=None, dtype=None\n):\n alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(\n 1, 1, 1, seq_len\n )\n if full:\n alibi_bias = alibi_bias - torch.arange(\n 1 - seq_len, 1, dtype=torch.int32, device=device\n ).view(1, 1, seq_len, 1)\n alibi_bias = alibi_bias.abs().mul(-1)\n slopes = gen_slopes(n_heads, alibi_bias_max, device=device)\n alibi_bias = alibi_bias * slopes\n return alibi_bias.to(dtype=dtype)\n\n\nATTN_CLASS_REGISTRY = {\n \"multihead_attention\": MultiheadAttention,\n \"multiquery_attention\": MultiQueryAttention,\n}\n" }, { "path": "model/llava/model/language_model/mpt/blocks.py", "content": "\"\"\"GPT Blocks used for the GPT Model.\"\"\"\nfrom typing import Dict, Optional, Tuple\n\nimport torch\nimport torch.nn as nn\n\nfrom .attention import ATTN_CLASS_REGISTRY\nfrom .norm import NORM_CLASS_REGISTRY\n\n\nclass MPTMLP(nn.Module):\n def __init__(\n self, d_model: int, expansion_ratio: int, device: Optional[str] = None\n ):\n super().__init__()\n self.up_proj = nn.Linear(d_model, expansion_ratio * d_model, device=device)\n self.act = nn.GELU(approximate=\"none\")\n self.down_proj = nn.Linear(expansion_ratio * d_model, d_model, device=device)\n self.down_proj._is_residual = True\n\n def forward(self, x):\n return self.down_proj(self.act(self.up_proj(x)))\n\n\nclass MPTBlock(nn.Module):\n def __init__(\n self,\n d_model: int,\n n_heads: int,\n expansion_ratio: int,\n attn_config: Dict = {\n \"attn_type\": \"multihead_attention\",\n \"attn_pdrop\": 0.0,\n \"attn_impl\": \"triton\",\n \"qk_ln\": False,\n \"clip_qkv\": None,\n \"softmax_scale\": None,\n \"prefix_lm\": False,\n \"attn_uses_sequence_id\": False,\n \"alibi\": False,\n \"alibi_bias_max\": 8,\n },\n resid_pdrop: float = 0.0,\n norm_type: str = \"low_precision_layernorm\",\n verbose: int = 0,\n device: Optional[str] = None,\n **kwargs\n ):\n del kwargs\n super().__init__()\n norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]\n attn_class = ATTN_CLASS_REGISTRY[attn_config[\"attn_type\"]]\n self.norm_1 = norm_class(d_model, device=device)\n self.attn = attn_class(\n attn_impl=attn_config[\"attn_impl\"],\n clip_qkv=attn_config[\"clip_qkv\"],\n qk_ln=attn_config[\"qk_ln\"],\n softmax_scale=attn_config[\"softmax_scale\"],\n attn_pdrop=attn_config[\"attn_pdrop\"],\n d_model=d_model,\n n_heads=n_heads,\n verbose=verbose,\n device=device,\n )\n self.norm_2 = norm_class(d_model, device=device)\n self.ffn = MPTMLP(\n d_model=d_model, expansion_ratio=expansion_ratio, device=device\n )\n self.resid_attn_dropout = nn.Dropout(resid_pdrop)\n self.resid_ffn_dropout = nn.Dropout(resid_pdrop)\n\n def forward(\n self,\n x: torch.Tensor,\n past_key_value: Optional[Tuple[torch.Tensor]] = None,\n attn_bias: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.ByteTensor] = None,\n is_causal: bool = True,\n ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor]]]:\n a = self.norm_1(x)\n (b, attn_weights, past_key_value) = self.attn(\n a,\n past_key_value=past_key_value,\n attn_bias=attn_bias,\n attention_mask=attention_mask,\n is_causal=is_causal,\n )\n x = x + self.resid_attn_dropout(b)\n m = self.norm_2(x)\n n = self.ffn(m)\n x = x + self.resid_ffn_dropout(n)\n return (x, attn_weights, past_key_value)\n" }, { "path": "model/llava/model/language_model/mpt/configuration_mpt.py", "content": "\"\"\"A HuggingFace-style model configuration.\"\"\"\nfrom typing import Dict, Optional, Union\n\nfrom transformers import PretrainedConfig\n\nattn_config_defaults: Dict = {\n \"attn_type\": \"multihead_attention\",\n \"attn_pdrop\": 0.0,\n \"attn_impl\": \"triton\",\n \"qk_ln\": False,\n \"clip_qkv\": None,\n \"softmax_scale\": None,\n \"prefix_lm\": False,\n \"attn_uses_sequence_id\": False,\n \"alibi\": False,\n \"alibi_bias_max\": 8,\n}\ninit_config_defaults: Dict = {\n \"name\": \"kaiming_normal_\",\n \"fan_mode\": \"fan_in\",\n \"init_nonlinearity\": \"relu\",\n \"init_div_is_residual\": True,\n \"emb_init_std\": None,\n \"emb_init_uniform_lim\": None,\n \"init_std\": None,\n \"init_gain\": 0.0,\n}\n\n\nclass MPTConfig(PretrainedConfig):\n model_type = \"mpt\"\n\n def __init__(\n self,\n d_model: int = 2048,\n n_heads: int = 16,\n n_layers: int = 24,\n expansion_ratio: int = 4,\n max_seq_len: int = 2048,\n vocab_size: int = 50368,\n resid_pdrop: float = 0.0,\n emb_pdrop: float = 0.0,\n learned_pos_emb: bool = True,\n attn_config: Dict = attn_config_defaults,\n init_device: str = \"cpu\",\n logit_scale: Optional[Union[float, str]] = None,\n no_bias: bool = False,\n verbose: int = 0,\n embedding_fraction: float = 1.0,\n norm_type: str = \"low_precision_layernorm\",\n use_cache: bool = False,\n init_config: Dict = init_config_defaults,\n **kwargs,\n ):\n \"\"\"The MPT configuration class.\n\n Args:\n d_model (int): The size of the embedding dimension of the model.\n n_heads (int): The number of attention heads.\n n_layers (int): The number of layers in the model.\n expansion_ratio (int): The ratio of the up/down scale in the MLP.\n max_seq_len (int): The maximum sequence length of the model.\n vocab_size (int): The size of the vocabulary.\n resid_pdrop (float): The dropout probability applied to the attention output before combining with residual.\n emb_pdrop (float): The dropout probability for the embedding layer.\n learned_pos_emb (bool): Whether to use learned positional embeddings\n attn_config (Dict): A dictionary used to configure the model's attention module:\n attn_type (str): type of attention to use. Options: multihead_attention, multiquery_attention\n attn_pdrop (float): The dropout probability for the attention layers.\n attn_impl (str): The attention implementation to use. One of 'torch', 'flash', or 'triton'.\n qk_ln (bool): Whether to apply layer normalization to the queries and keys in the attention layer.\n clip_qkv (Optional[float]): If not None, clip the queries, keys, and values in the attention layer to\n this value.\n softmax_scale (Optional[float]): If not None, scale the softmax in the attention layer by this value. If None,\n use the default scale of ``1/sqrt(d_keys)``.\n prefix_lm (Optional[bool]): Whether the model should operate as a Prefix LM. This requires passing an\n extra `prefix_mask` argument which indicates which tokens belong to the prefix. Tokens in the prefix\n can attend to one another bi-directionally. Tokens outside the prefix use causal attention.\n attn_uses_sequence_id (Optional[bool]): Whether to restrict attention to tokens that have the same sequence_id.\n When the model is in `train` mode, this requires passing an extra `sequence_id` argument which indicates\n which sub-sequence each token belongs to.\n Defaults to ``False`` meaning any provided `sequence_id` will be ignored.\n alibi (bool): Whether to use the alibi bias instead of position embeddings.\n alibi_bias_max (int): The maximum value of the alibi bias.\n init_device (str): The device to use for parameter initialization.\n logit_scale (Optional[Union[float, str]]): If not None, scale the logits by this value.\n no_bias (bool): Whether to use bias in all layers.\n verbose (int): The verbosity level. 0 is silent.\n embedding_fraction (float): The fraction to scale the gradients of the embedding layer by.\n norm_type (str): choose type of norm to use\n multiquery_attention (bool): Whether to use multiquery attention implementation.\n use_cache (bool): Whether or not the model should return the last key/values attentions\n init_config (Dict): A dictionary used to configure the model initialization:\n init_config.name: The parameter initialization scheme to use. Options: 'default_', 'baseline_',\n 'kaiming_uniform_', 'kaiming_normal_', 'neox_init_', 'small_init_', 'xavier_uniform_', or\n 'xavier_normal_'. These mimic the parameter initialization methods in PyTorch.\n init_div_is_residual (Union[int, float, str, bool]): Value to divide initial weights by if ``module._is_residual`` is True.\n emb_init_std (Optional[float]): The standard deviation of the normal distribution used to initialize the embedding layer.\n emb_init_uniform_lim (Optional[Union[Tuple[float, float], float]]): The lower and upper limits of the uniform distribution\n used to initialize the embedding layer. Mutually exclusive with ``emb_init_std``.\n init_std (float): The standard deviation of the normal distribution used to initialize the model,\n if using the baseline_ parameter initialization scheme.\n init_gain (float): The gain to use for parameter initialization with kaiming or xavier initialization schemes.\n fan_mode (str): The fan mode to use for parameter initialization with kaiming initialization schemes.\n init_nonlinearity (str): The nonlinearity to use for parameter initialization with kaiming initialization schemes.\n ---\n See llmfoundry.models.utils.param_init_fns.py for info on other param init config options\n \"\"\"\n self.d_model = d_model\n self.n_heads = n_heads\n self.n_layers = n_layers\n self.expansion_ratio = expansion_ratio\n self.max_seq_len = max_seq_len\n self.vocab_size = vocab_size\n self.resid_pdrop = resid_pdrop\n self.emb_pdrop = emb_pdrop\n self.learned_pos_emb = learned_pos_emb\n self.attn_config = attn_config\n self.init_device = init_device\n self.logit_scale = logit_scale\n self.no_bias = no_bias\n self.verbose = verbose\n self.embedding_fraction = embedding_fraction\n self.norm_type = norm_type\n self.use_cache = use_cache\n self.init_config = init_config\n if \"name\" in kwargs:\n del kwargs[\"name\"]\n if \"loss_fn\" in kwargs:\n del kwargs[\"loss_fn\"]\n super().__init__(**kwargs)\n self._validate_config()\n\n def _set_config_defaults(self, config, config_defaults):\n for k, v in config_defaults.items():\n if k not in config:\n config[k] = v\n return config\n\n def _validate_config(self):\n self.attn_config = self._set_config_defaults(\n self.attn_config, attn_config_defaults\n )\n self.init_config = self._set_config_defaults(\n self.init_config, init_config_defaults\n )\n if self.d_model % self.n_heads != 0:\n raise ValueError(\"d_model must be divisible by n_heads\")\n if any(\n (\n prob < 0 or prob > 1\n for prob in [\n self.attn_config[\"attn_pdrop\"],\n self.resid_pdrop,\n self.emb_pdrop,\n ]\n )\n ):\n raise ValueError(\n \"self.attn_config['attn_pdrop'], resid_pdrop, emb_pdrop are probabilities and must be between 0 and 1\"\n )\n if self.attn_config[\"attn_impl\"] not in [\"torch\", \"flash\", \"triton\"]:\n raise ValueError(f\"Unknown attn_impl={self.attn_config['attn_impl']}\")\n if self.attn_config[\"prefix_lm\"] and self.attn_config[\"attn_impl\"] not in [\n \"torch\",\n \"triton\",\n ]:\n raise NotImplementedError(\n \"prefix_lm only implemented with torch and triton attention.\"\n )\n if self.attn_config[\"alibi\"] and self.attn_config[\"attn_impl\"] not in [\n \"torch\",\n \"triton\",\n ]:\n raise NotImplementedError(\n \"alibi only implemented with torch and triton attention.\"\n )\n if self.attn_config[\"attn_uses_sequence_id\"] and self.attn_config[\n \"attn_impl\"\n ] not in [\"torch\", \"triton\"]:\n raise NotImplementedError(\n \"attn_uses_sequence_id only implemented with torch and triton attention.\"\n )\n if self.embedding_fraction > 1 or self.embedding_fraction <= 0:\n raise ValueError(\n \"model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!\"\n )\n if isinstance(self.logit_scale, str) and self.logit_scale != \"inv_sqrt_d_model\":\n raise ValueError(\n f\"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.\"\n )\n if self.init_config.get(\"name\", None) is None:\n raise ValueError(\n f\"self.init_config={self.init_config!r} 'name' needs to be set.\"\n )\n if not self.learned_pos_emb and (not self.attn_config[\"alibi\"]):\n raise ValueError(\n f\"Positional information must be provided to the model using either learned_pos_emb or alibi.\"\n )\n" }, { "path": "model/llava/model/language_model/mpt/custom_embedding.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch import Tensor\n\n\nclass SharedEmbedding(nn.Embedding):\n def forward(self, input: Tensor, unembed: bool = False) -> Tensor:\n if unembed:\n return F.linear(input, self.weight)\n return super().forward(input)\n" }, { "path": "model/llava/model/language_model/mpt/flash_attn_triton.py", "content": "\"\"\"\nCopied from https://github.com/HazyResearch/flash-attention/blob/eff9fe6b8076df59d64d7a3f464696738a3c7c24/flash_attn/flash_attn_triton.py\nupdate imports to use 'triton_pre_mlir'\n\n*Experimental* implementation of FlashAttention in Triton.\nTested with triton==2.0.0.dev20221202.\nTriton 2.0 has a new backend (MLIR) but seems like it doesn't yet work for head dimensions\nother than 64:\nhttps://github.com/openai/triton/blob/d376020f90002757eea3ea9475d4f7cfc2ec5ead/python/triton/ops/flash_attention.py#L207\nWe'll update this implementation with the new Triton backend once this is fixed.\n\nWe use the FlashAttention implementation from Phil Tillet a starting point.\nhttps://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py\n\nChanges:\n- Implement both causal and non-causal attention.\n- Implement both self-attention and cross-attention.\n- Support arbitrary seqlens (not just multiples of 128), for both forward and backward.\n- Support all head dimensions up to 128 (not just 16, 32, 64, 128), for both forward and backward.\n- Support attention bias.\n- Speed up the forward pass a bit, and only store the LSE instead of m and l.\n- Make the backward for d=128 much faster by reducing register spilling.\n- Optionally parallelize the backward pass across seqlen_k, to deal with the case of\nsmall batch size * nheads.\n\nCaution:\n- This is an *experimental* implementation. The forward pass should be quite robust but\nI'm not 100% sure that the backward pass doesn't have race conditions (due to the Triton compiler).\n- This implementation has only been tested on A100.\n- If you plan to use headdim other than 64 and 128, you should test for race conditions\n(due to the Triton compiler), as done in tests/test_flash_attn.py\n\"test_flash_attn_triton_race_condition\". I've tested and fixed many race conditions\nfor different head dimensions (40, 48, 64, 128, 80, 88, 96), but I'm still not 100% confident\nthat there are none left for other head dimensions.\n\nDifferences between this Triton version and the CUDA version:\n- Triton version doesn't support dropout.\n- Triton forward is generally faster than CUDA forward, while Triton backward is\ngenerally slower than CUDA backward. Overall Triton forward + backward is slightly slower\nthan CUDA forward + backward.\n- Triton version doesn't support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor).\n- Triton version supports attention bias, while CUDA version doesn't.\n\"\"\"\nimport math\n\nimport torch\nimport triton_pre_mlir as triton\nimport triton_pre_mlir.language as tl\n\n\n@triton.heuristics(\n {\n \"EVEN_M\": lambda args: args[\"seqlen_q\"] % args[\"BLOCK_M\"] == 0,\n \"EVEN_N\": lambda args: args[\"seqlen_k\"] % args[\"BLOCK_N\"] == 0,\n \"EVEN_HEADDIM\": lambda args: args[\"headdim\"] == args[\"BLOCK_HEADDIM\"],\n }\n)\n@triton.jit\ndef _fwd_kernel(\n Q,\n K,\n V,\n Bias,\n Out,\n Lse,\n TMP,\n softmax_scale,\n stride_qb,\n stride_qh,\n stride_qm,\n stride_kb,\n stride_kh,\n stride_kn,\n stride_vb,\n stride_vh,\n stride_vn,\n stride_bb,\n stride_bh,\n stride_bm,\n stride_ob,\n stride_oh,\n stride_om,\n nheads,\n seqlen_q,\n seqlen_k,\n seqlen_q_rounded,\n headdim,\n CACHE_KEY_SEQLEN_Q,\n CACHE_KEY_SEQLEN_K,\n BIAS_TYPE: tl.constexpr,\n IS_CAUSAL: tl.constexpr,\n BLOCK_HEADDIM: tl.constexpr,\n EVEN_M: tl.constexpr,\n EVEN_N: tl.constexpr,\n EVEN_HEADDIM: tl.constexpr,\n BLOCK_M: tl.constexpr,\n BLOCK_N: tl.constexpr,\n):\n start_m = tl.program_id(0)\n off_hb = tl.program_id(1)\n off_b = off_hb // nheads\n off_h = off_hb % nheads\n offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)\n offs_n = tl.arange(0, BLOCK_N)\n offs_d = tl.arange(0, BLOCK_HEADDIM)\n q_ptrs = (\n Q\n + off_b * stride_qb\n + off_h * stride_qh\n + (offs_m[:, None] * stride_qm + offs_d[None, :])\n )\n k_ptrs = (\n K\n + off_b * stride_kb\n + off_h * stride_kh\n + (offs_n[:, None] * stride_kn + offs_d[None, :])\n )\n v_ptrs = (\n V\n + off_b * stride_vb\n + off_h * stride_vh\n + (offs_n[:, None] * stride_vn + offs_d[None, :])\n )\n if BIAS_TYPE == \"vector\":\n b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + offs_n\n elif BIAS_TYPE == \"matrix\":\n b_ptrs = (\n Bias\n + off_b * stride_bb\n + off_h * stride_bh\n + (offs_m[:, None] * stride_bm + offs_n[None, :])\n )\n t_ptrs = TMP + off_hb * seqlen_q_rounded + offs_m\n lse_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float(\"inf\")\n m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float(\"inf\")\n acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)\n if EVEN_M & EVEN_N:\n if EVEN_HEADDIM:\n q = tl.load(q_ptrs)\n else:\n q = tl.load(q_ptrs, mask=offs_d[None, :] < headdim, other=0.0)\n elif EVEN_HEADDIM:\n q = tl.load(q_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0)\n else:\n q = tl.load(\n q_ptrs,\n mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim),\n other=0.0,\n )\n end_n = seqlen_k if not IS_CAUSAL else tl.minimum((start_m + 1) * BLOCK_M, seqlen_k)\n for start_n in range(0, end_n, BLOCK_N):\n start_n = tl.multiple_of(start_n, BLOCK_N)\n if EVEN_N & EVEN_M:\n if EVEN_HEADDIM:\n k = tl.load(k_ptrs + start_n * stride_kn)\n else:\n k = tl.load(\n k_ptrs + start_n * stride_kn,\n mask=offs_d[None, :] < headdim,\n other=0.0,\n )\n elif EVEN_HEADDIM:\n k = tl.load(\n k_ptrs + start_n * stride_kn,\n mask=(start_n + offs_n)[:, None] < seqlen_k,\n other=0.0,\n )\n else:\n k = tl.load(\n k_ptrs + start_n * stride_kn,\n mask=((start_n + offs_n)[:, None] < seqlen_k)\n & (offs_d[None, :] < headdim),\n other=0.0,\n )\n qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)\n qk += tl.dot(q, k, trans_b=True)\n if not EVEN_N:\n qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0, float(\"-inf\"))\n if IS_CAUSAL:\n qk += tl.where(\n offs_m[:, None] >= (start_n + offs_n)[None, :], 0, float(\"-inf\")\n )\n if BIAS_TYPE != \"none\":\n if BIAS_TYPE == \"vector\":\n if EVEN_N:\n bias = tl.load(b_ptrs + start_n).to(tl.float32)\n else:\n bias = tl.load(\n b_ptrs + start_n, mask=start_n + offs_n < seqlen_k, other=0.0\n ).to(tl.float32)\n bias = bias[None, :]\n elif BIAS_TYPE == \"matrix\":\n if EVEN_M & EVEN_N:\n bias = tl.load(b_ptrs + start_n).to(tl.float32)\n else:\n bias = tl.load(\n b_ptrs + start_n,\n mask=(offs_m[:, None] < seqlen_q)\n & ((start_n + offs_n)[None, :] < seqlen_k),\n other=0.0,\n ).to(tl.float32)\n qk = qk * softmax_scale + bias\n m_ij = tl.maximum(tl.max(qk, 1), lse_i)\n p = tl.exp(qk - m_ij[:, None])\n else:\n m_ij = tl.maximum(tl.max(qk, 1) * softmax_scale, lse_i)\n p = tl.exp(qk * softmax_scale - m_ij[:, None])\n l_ij = tl.sum(p, 1)\n acc_o_scale = tl.exp(m_i - m_ij)\n tl.store(t_ptrs, acc_o_scale)\n acc_o_scale = tl.load(t_ptrs)\n acc_o = acc_o * acc_o_scale[:, None]\n if EVEN_N & EVEN_M:\n if EVEN_HEADDIM:\n v = tl.load(v_ptrs + start_n * stride_vn)\n else:\n v = tl.load(\n v_ptrs + start_n * stride_vn,\n mask=offs_d[None, :] < headdim,\n other=0.0,\n )\n elif EVEN_HEADDIM:\n v = tl.load(\n v_ptrs + start_n * stride_vn,\n mask=(start_n + offs_n)[:, None] < seqlen_k,\n other=0.0,\n )\n else:\n v = tl.load(\n v_ptrs + start_n * stride_vn,\n mask=((start_n + offs_n)[:, None] < seqlen_k)\n & (offs_d[None, :] < headdim),\n other=0.0,\n )\n p = p.to(v.dtype)\n acc_o += tl.dot(p, v)\n m_i = m_ij\n l_i_new = tl.exp(lse_i - m_ij) + l_ij\n lse_i = m_ij + tl.log(l_i_new)\n o_scale = tl.exp(m_i - lse_i)\n tl.store(t_ptrs, o_scale)\n o_scale = tl.load(t_ptrs)\n acc_o = acc_o * o_scale[:, None]\n start_m = tl.program_id(0)\n offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)\n lse_ptrs = Lse + off_hb * seqlen_q_rounded + offs_m\n tl.store(lse_ptrs, lse_i)\n offs_d = tl.arange(0, BLOCK_HEADDIM)\n out_ptrs = (\n Out\n + off_b * stride_ob\n + off_h * stride_oh\n + (offs_m[:, None] * stride_om + offs_d[None, :])\n )\n if EVEN_M:\n if EVEN_HEADDIM:\n tl.store(out_ptrs, acc_o)\n else:\n tl.store(out_ptrs, acc_o, mask=offs_d[None, :] < headdim)\n elif EVEN_HEADDIM:\n tl.store(out_ptrs, acc_o, mask=offs_m[:, None] < seqlen_q)\n else:\n tl.store(\n out_ptrs,\n acc_o,\n mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim),\n )\n\n\n@triton.jit\ndef _bwd_preprocess_do_o_dot(\n Out,\n DO,\n Delta,\n stride_ob,\n stride_oh,\n stride_om,\n stride_dob,\n stride_doh,\n stride_dom,\n nheads,\n seqlen_q,\n seqlen_q_rounded,\n headdim,\n BLOCK_M: tl.constexpr,\n BLOCK_HEADDIM: tl.constexpr,\n):\n start_m = tl.program_id(0)\n off_hb = tl.program_id(1)\n off_b = off_hb // nheads\n off_h = off_hb % nheads\n offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)\n offs_d = tl.arange(0, BLOCK_HEADDIM)\n o = tl.load(\n Out\n + off_b * stride_ob\n + off_h * stride_oh\n + offs_m[:, None] * stride_om\n + offs_d[None, :],\n mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim),\n other=0.0,\n ).to(tl.float32)\n do = tl.load(\n DO\n + off_b * stride_dob\n + off_h * stride_doh\n + offs_m[:, None] * stride_dom\n + offs_d[None, :],\n mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim),\n other=0.0,\n ).to(tl.float32)\n delta = tl.sum(o * do, axis=1)\n tl.store(Delta + off_hb * seqlen_q_rounded + offs_m, delta)\n\n\n@triton.jit\ndef _bwd_store_dk_dv(\n dk_ptrs,\n dv_ptrs,\n dk,\n dv,\n offs_n,\n offs_d,\n seqlen_k,\n headdim,\n EVEN_M: tl.constexpr,\n EVEN_N: tl.constexpr,\n EVEN_HEADDIM: tl.constexpr,\n):\n if EVEN_N & EVEN_M:\n if EVEN_HEADDIM:\n tl.store(dv_ptrs, dv)\n tl.store(dk_ptrs, dk)\n else:\n tl.store(dv_ptrs, dv, mask=offs_d[None, :] < headdim)\n tl.store(dk_ptrs, dk, mask=offs_d[None, :] < headdim)\n elif EVEN_HEADDIM:\n tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k)\n tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k)\n else:\n tl.store(\n dv_ptrs, dv, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim)\n )\n tl.store(\n dk_ptrs, dk, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim)\n )\n\n\n@triton.jit\ndef _bwd_kernel_one_col_block(\n start_n,\n Q,\n K,\n V,\n Bias,\n DO,\n DQ,\n DK,\n DV,\n LSE,\n D,\n softmax_scale,\n stride_qm,\n stride_kn,\n stride_vn,\n stride_bm,\n stride_dom,\n stride_dqm,\n stride_dkn,\n stride_dvn,\n seqlen_q,\n seqlen_k,\n headdim,\n ATOMIC_ADD: tl.constexpr,\n BIAS_TYPE: tl.constexpr,\n IS_CAUSAL: tl.constexpr,\n BLOCK_HEADDIM: tl.constexpr,\n EVEN_M: tl.constexpr,\n EVEN_N: tl.constexpr,\n EVEN_HEADDIM: tl.constexpr,\n BLOCK_M: tl.constexpr,\n BLOCK_N: tl.constexpr,\n):\n begin_m = 0 if not IS_CAUSAL else start_n * BLOCK_N // BLOCK_M * BLOCK_M\n offs_qm = begin_m + tl.arange(0, BLOCK_M)\n offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)\n offs_m = tl.arange(0, BLOCK_M)\n offs_d = tl.arange(0, BLOCK_HEADDIM)\n q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_d[None, :])\n k_ptrs = K + (offs_n[:, None] * stride_kn + offs_d[None, :])\n v_ptrs = V + (offs_n[:, None] * stride_vn + offs_d[None, :])\n do_ptrs = DO + (offs_qm[:, None] * stride_dom + offs_d[None, :])\n dq_ptrs = DQ + (offs_qm[:, None] * stride_dqm + offs_d[None, :])\n if BIAS_TYPE == \"vector\":\n b_ptrs = Bias + offs_n\n elif BIAS_TYPE == \"matrix\":\n b_ptrs = Bias + (offs_qm[:, None] * stride_bm + offs_n[None, :])\n dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)\n dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)\n if begin_m >= seqlen_q:\n dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])\n dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])\n _bwd_store_dk_dv(\n dk_ptrs,\n dv_ptrs,\n dk,\n dv,\n offs_n,\n offs_d,\n seqlen_k,\n headdim,\n EVEN_M=EVEN_M,\n EVEN_N=EVEN_N,\n EVEN_HEADDIM=EVEN_HEADDIM,\n )\n return\n if EVEN_N & EVEN_M:\n if EVEN_HEADDIM:\n k = tl.load(k_ptrs)\n v = tl.load(v_ptrs)\n else:\n k = tl.load(k_ptrs, mask=offs_d[None, :] < headdim, other=0.0)\n v = tl.load(v_ptrs, mask=offs_d[None, :] < headdim, other=0.0)\n elif EVEN_HEADDIM:\n k = tl.load(k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)\n v = tl.load(v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)\n else:\n k = tl.load(\n k_ptrs,\n mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim),\n other=0.0,\n )\n v = tl.load(\n v_ptrs,\n mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim),\n other=0.0,\n )\n num_block_m = tl.cdiv(seqlen_q, BLOCK_M)\n for start_m in range(begin_m, num_block_m * BLOCK_M, BLOCK_M):\n start_m = tl.multiple_of(start_m, BLOCK_M)\n offs_m_curr = start_m + offs_m\n if EVEN_M & EVEN_HEADDIM:\n q = tl.load(q_ptrs)\n elif EVEN_HEADDIM:\n q = tl.load(q_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0)\n else:\n q = tl.load(\n q_ptrs,\n mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim),\n other=0.0,\n )\n qk = tl.dot(q, k, trans_b=True)\n if not EVEN_N:\n qk = tl.where(offs_n[None, :] < seqlen_k, qk, float(\"-inf\"))\n if IS_CAUSAL:\n qk = tl.where(offs_m_curr[:, None] >= offs_n[None, :], qk, float(\"-inf\"))\n if BIAS_TYPE != \"none\":\n tl.debug_barrier()\n if BIAS_TYPE == \"vector\":\n if EVEN_N:\n bias = tl.load(b_ptrs).to(tl.float32)\n else:\n bias = tl.load(b_ptrs, mask=offs_n < seqlen_k, other=0.0).to(\n tl.float32\n )\n bias = bias[None, :]\n elif BIAS_TYPE == \"matrix\":\n if EVEN_M & EVEN_N:\n bias = tl.load(b_ptrs).to(tl.float32)\n else:\n bias = tl.load(\n b_ptrs,\n mask=(offs_m_curr[:, None] < seqlen_q)\n & (offs_n[None, :] < seqlen_k),\n other=0.0,\n ).to(tl.float32)\n qk = qk * softmax_scale + bias\n if not EVEN_M & EVEN_HEADDIM:\n tl.debug_barrier()\n lse_i = tl.load(LSE + offs_m_curr)\n if BIAS_TYPE == \"none\":\n p = tl.exp(qk * softmax_scale - lse_i[:, None])\n else:\n p = tl.exp(qk - lse_i[:, None])\n if EVEN_M & EVEN_HEADDIM:\n do = tl.load(do_ptrs)\n else:\n do = tl.load(\n do_ptrs,\n mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim),\n other=0.0,\n )\n dv += tl.dot(p.to(do.dtype), do, trans_a=True)\n if not EVEN_M & EVEN_HEADDIM:\n tl.debug_barrier()\n dp = tl.dot(do, v, trans_b=True)\n if not EVEN_HEADDIM:\n tl.debug_barrier()\n Di = tl.load(D + offs_m_curr)\n ds = (p * (dp - Di[:, None]) * softmax_scale).to(q.dtype)\n dk += tl.dot(ds, q, trans_a=True)\n if not EVEN_M & EVEN_HEADDIM:\n tl.debug_barrier()\n if not ATOMIC_ADD:\n if EVEN_M & EVEN_HEADDIM:\n dq = tl.load(dq_ptrs, eviction_policy=\"evict_last\")\n dq += tl.dot(ds, k)\n tl.store(dq_ptrs, dq, eviction_policy=\"evict_last\")\n elif EVEN_HEADDIM:\n dq = tl.load(\n dq_ptrs,\n mask=offs_m_curr[:, None] < seqlen_q,\n other=0.0,\n eviction_policy=\"evict_last\",\n )\n dq += tl.dot(ds, k)\n tl.store(\n dq_ptrs,\n dq,\n mask=offs_m_curr[:, None] < seqlen_q,\n eviction_policy=\"evict_last\",\n )\n else:\n dq = tl.load(\n dq_ptrs,\n mask=(offs_m_curr[:, None] < seqlen_q)\n & (offs_d[None, :] < headdim),\n other=0.0,\n eviction_policy=\"evict_last\",\n )\n dq += tl.dot(ds, k)\n tl.store(\n dq_ptrs,\n dq,\n mask=(offs_m_curr[:, None] < seqlen_q)\n & (offs_d[None, :] < headdim),\n eviction_policy=\"evict_last\",\n )\n else:\n dq = tl.dot(ds, k)\n if EVEN_M & EVEN_HEADDIM:\n tl.atomic_add(dq_ptrs, dq)\n elif EVEN_HEADDIM:\n tl.atomic_add(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q)\n else:\n tl.atomic_add(\n dq_ptrs,\n dq,\n mask=(offs_m_curr[:, None] < seqlen_q)\n & (offs_d[None, :] < headdim),\n )\n dq_ptrs += BLOCK_M * stride_dqm\n q_ptrs += BLOCK_M * stride_qm\n do_ptrs += BLOCK_M * stride_dom\n if BIAS_TYPE == \"matrix\":\n b_ptrs += BLOCK_M * stride_bm\n dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])\n dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])\n _bwd_store_dk_dv(\n dk_ptrs,\n dv_ptrs,\n dk,\n dv,\n offs_n,\n offs_d,\n seqlen_k,\n headdim,\n EVEN_M=EVEN_M,\n EVEN_N=EVEN_N,\n EVEN_HEADDIM=EVEN_HEADDIM,\n )\n\n\ndef init_to_zero(name):\n return lambda nargs: nargs[name].zero_()\n\n\n@triton.autotune(\n configs=[\n triton.Config(\n {\"BLOCK_M\": 128, \"BLOCK_N\": 128, \"SEQUENCE_PARALLEL\": False},\n num_warps=8,\n num_stages=1,\n pre_hook=init_to_zero(\"DQ\"),\n ),\n triton.Config(\n {\"BLOCK_M\": 128, \"BLOCK_N\": 128, \"SEQUENCE_PARALLEL\": True},\n num_warps=8,\n num_stages=1,\n pre_hook=init_to_zero(\"DQ\"),\n ),\n ],\n key=[\n \"CACHE_KEY_SEQLEN_Q\",\n \"CACHE_KEY_SEQLEN_K\",\n \"BIAS_TYPE\",\n \"IS_CAUSAL\",\n \"BLOCK_HEADDIM\",\n ],\n)\n@triton.heuristics(\n {\n \"EVEN_M\": lambda args: args[\"seqlen_q\"] % args[\"BLOCK_M\"] == 0,\n \"EVEN_N\": lambda args: args[\"seqlen_k\"] % args[\"BLOCK_N\"] == 0,\n \"EVEN_HEADDIM\": lambda args: args[\"headdim\"] == args[\"BLOCK_HEADDIM\"],\n }\n)\n@triton.jit\ndef _bwd_kernel(\n Q,\n K,\n V,\n Bias,\n DO,\n DQ,\n DK,\n DV,\n LSE,\n D,\n softmax_scale,\n stride_qb,\n stride_qh,\n stride_qm,\n stride_kb,\n stride_kh,\n stride_kn,\n stride_vb,\n stride_vh,\n stride_vn,\n stride_bb,\n stride_bh,\n stride_bm,\n stride_dob,\n stride_doh,\n stride_dom,\n stride_dqb,\n stride_dqh,\n stride_dqm,\n stride_dkb,\n stride_dkh,\n stride_dkn,\n stride_dvb,\n stride_dvh,\n stride_dvn,\n nheads,\n seqlen_q,\n seqlen_k,\n seqlen_q_rounded,\n headdim,\n CACHE_KEY_SEQLEN_Q,\n CACHE_KEY_SEQLEN_K,\n BIAS_TYPE: tl.constexpr,\n IS_CAUSAL: tl.constexpr,\n BLOCK_HEADDIM: tl.constexpr,\n SEQUENCE_PARALLEL: tl.constexpr,\n EVEN_M: tl.constexpr,\n EVEN_N: tl.constexpr,\n EVEN_HEADDIM: tl.constexpr,\n BLOCK_M: tl.constexpr,\n BLOCK_N: tl.constexpr,\n):\n off_hb = tl.program_id(1)\n off_b = off_hb // nheads\n off_h = off_hb % nheads\n Q += off_b * stride_qb + off_h * stride_qh\n K += off_b * stride_kb + off_h * stride_kh\n V += off_b * stride_vb + off_h * stride_vh\n DO += off_b * stride_dob + off_h * stride_doh\n DQ += off_b * stride_dqb + off_h * stride_dqh\n DK += off_b * stride_dkb + off_h * stride_dkh\n DV += off_b * stride_dvb + off_h * stride_dvh\n if BIAS_TYPE != \"none\":\n Bias += off_b * stride_bb + off_h * stride_bh\n D += off_hb * seqlen_q_rounded\n LSE += off_hb * seqlen_q_rounded\n if not SEQUENCE_PARALLEL:\n num_block_n = tl.cdiv(seqlen_k, BLOCK_N)\n for start_n in range(0, num_block_n):\n _bwd_kernel_one_col_block(\n start_n,\n Q,\n K,\n V,\n Bias,\n DO,\n DQ,\n DK,\n DV,\n LSE,\n D,\n softmax_scale,\n stride_qm,\n stride_kn,\n stride_vn,\n stride_bm,\n stride_dom,\n stride_dqm,\n stride_dkn,\n stride_dvn,\n seqlen_q,\n seqlen_k,\n headdim,\n ATOMIC_ADD=False,\n BIAS_TYPE=BIAS_TYPE,\n IS_CAUSAL=IS_CAUSAL,\n BLOCK_HEADDIM=BLOCK_HEADDIM,\n EVEN_M=EVEN_M,\n EVEN_N=EVEN_N,\n EVEN_HEADDIM=EVEN_HEADDIM,\n BLOCK_M=BLOCK_M,\n BLOCK_N=BLOCK_N,\n )\n else:\n start_n = tl.program_id(0)\n _bwd_kernel_one_col_block(\n start_n,\n Q,\n K,\n V,\n Bias,\n DO,\n DQ,\n DK,\n DV,\n LSE,\n D,\n softmax_scale,\n stride_qm,\n stride_kn,\n stride_vn,\n stride_bm,\n stride_dom,\n stride_dqm,\n stride_dkn,\n stride_dvn,\n seqlen_q,\n seqlen_k,\n headdim,\n ATOMIC_ADD=True,\n BIAS_TYPE=BIAS_TYPE,\n IS_CAUSAL=IS_CAUSAL,\n BLOCK_HEADDIM=BLOCK_HEADDIM,\n EVEN_M=EVEN_M,\n EVEN_N=EVEN_N,\n EVEN_HEADDIM=EVEN_HEADDIM,\n BLOCK_M=BLOCK_M,\n BLOCK_N=BLOCK_N,\n )\n\n\ndef _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None):\n (batch, seqlen_q, nheads, d) = q.shape\n (_, seqlen_k, _, _) = k.shape\n assert k.shape == (batch, seqlen_k, nheads, d)\n assert v.shape == (batch, seqlen_k, nheads, d)\n assert d <= 128, \"FlashAttention only support head dimensions up to 128\"\n assert q.dtype == k.dtype == v.dtype, \"All tensors must have the same type\"\n assert q.dtype in [torch.float16, torch.bfloat16], \"Only support fp16 and bf16\"\n assert q.is_cuda and k.is_cuda and v.is_cuda\n softmax_scale = softmax_scale or 1.0 / math.sqrt(d)\n has_bias = bias is not None\n bias_type = \"none\"\n if has_bias:\n assert bias.dtype in [q.dtype, torch.float]\n assert bias.is_cuda\n assert bias.dim() == 4\n if bias.stride(-1) != 1:\n bias = bias.contiguous()\n if bias.shape[2:] == (1, seqlen_k):\n bias_type = \"vector\"\n elif bias.shape[2:] == (seqlen_q, seqlen_k):\n bias_type = \"matrix\"\n else:\n raise RuntimeError(\n \"Last 2 dimensions of bias must be (1, seqlen_k) or (seqlen_q, seqlen_k)\"\n )\n bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)\n bias_strides = (\n (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0)\n )\n seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128\n lse = torch.empty(\n (batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32\n )\n tmp = torch.empty(\n (batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32\n )\n o = torch.empty_like(q)\n BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)\n BLOCK = 128\n num_warps = 4 if d <= 64 else 8\n grid = lambda META: (triton.cdiv(seqlen_q, META[\"BLOCK_M\"]), batch * nheads)\n _fwd_kernel[grid](\n q,\n k,\n v,\n bias,\n o,\n lse,\n tmp,\n softmax_scale,\n q.stride(0),\n q.stride(2),\n q.stride(1),\n k.stride(0),\n k.stride(2),\n k.stride(1),\n v.stride(0),\n v.stride(2),\n v.stride(1),\n *bias_strides,\n o.stride(0),\n o.stride(2),\n o.stride(1),\n nheads,\n seqlen_q,\n seqlen_k,\n seqlen_q_rounded,\n d,\n seqlen_q // 32,\n seqlen_k // 32,\n bias_type,\n causal,\n BLOCK_HEADDIM,\n BLOCK_M=BLOCK,\n BLOCK_N=BLOCK,\n num_warps=num_warps,\n num_stages=1\n )\n return (o, lse, softmax_scale)\n\n\ndef _flash_attn_backward(\n do, q, k, v, o, lse, dq, dk, dv, bias=None, causal=False, softmax_scale=None\n):\n if do.stride(-1) != 1:\n do = do.contiguous()\n (batch, seqlen_q, nheads, d) = q.shape\n (_, seqlen_k, _, _) = k.shape\n assert d <= 128\n seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128\n assert lse.shape == (batch, nheads, seqlen_q_rounded)\n assert q.stride(-1) == k.stride(-1) == v.stride(-1) == o.stride(-1) == 1\n assert dq.stride(-1) == dk.stride(-1) == dv.stride(-1) == 1\n softmax_scale = softmax_scale or 1.0 / math.sqrt(d)\n dq_accum = torch.empty_like(q, dtype=torch.float32)\n delta = torch.empty_like(lse)\n BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)\n grid = lambda META: (triton.cdiv(seqlen_q, META[\"BLOCK_M\"]), batch * nheads)\n _bwd_preprocess_do_o_dot[grid](\n o,\n do,\n delta,\n o.stride(0),\n o.stride(2),\n o.stride(1),\n do.stride(0),\n do.stride(2),\n do.stride(1),\n nheads,\n seqlen_q,\n seqlen_q_rounded,\n d,\n BLOCK_M=128,\n BLOCK_HEADDIM=BLOCK_HEADDIM,\n )\n has_bias = bias is not None\n bias_type = \"none\"\n if has_bias:\n assert bias.dtype in [q.dtype, torch.float]\n assert bias.is_cuda\n assert bias.dim() == 4\n assert bias.stride(-1) == 1\n if bias.shape[2:] == (1, seqlen_k):\n bias_type = \"vector\"\n elif bias.shape[2:] == (seqlen_q, seqlen_k):\n bias_type = \"matrix\"\n else:\n raise RuntimeError(\n \"Last 2 dimensions of bias must be (1, seqlen_k) or (seqlen_q, seqlen_k)\"\n )\n bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)\n bias_strides = (\n (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0)\n )\n grid = lambda META: (\n triton.cdiv(seqlen_k, META[\"BLOCK_N\"]) if META[\"SEQUENCE_PARALLEL\"] else 1,\n batch * nheads,\n )\n _bwd_kernel[grid](\n q,\n k,\n v,\n bias,\n do,\n dq_accum,\n dk,\n dv,\n lse,\n delta,\n softmax_scale,\n q.stride(0),\n q.stride(2),\n q.stride(1),\n k.stride(0),\n k.stride(2),\n k.stride(1),\n v.stride(0),\n v.stride(2),\n v.stride(1),\n *bias_strides,\n do.stride(0),\n do.stride(2),\n do.stride(1),\n dq_accum.stride(0),\n dq_accum.stride(2),\n dq_accum.stride(1),\n dk.stride(0),\n dk.stride(2),\n dk.stride(1),\n dv.stride(0),\n dv.stride(2),\n dv.stride(1),\n nheads,\n seqlen_q,\n seqlen_k,\n seqlen_q_rounded,\n d,\n seqlen_q // 32,\n seqlen_k // 32,\n bias_type,\n causal,\n BLOCK_HEADDIM\n )\n dq.copy_(dq_accum)\n\n\nclass FlashAttnQKVPackedFunc(torch.autograd.Function):\n @staticmethod\n def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None):\n \"\"\"\n qkv: (batch, seqlen, 3, nheads, headdim)\n bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).\n For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).\n ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)\n \"\"\"\n if qkv.stride(-1) != 1:\n qkv = qkv.contiguous()\n (o, lse, ctx.softmax_scale) = _flash_attn_forward(\n qkv[:, :, 0],\n qkv[:, :, 1],\n qkv[:, :, 2],\n bias=bias,\n causal=causal,\n softmax_scale=softmax_scale,\n )\n ctx.save_for_backward(qkv, o, lse, bias)\n ctx.causal = causal\n return o\n\n @staticmethod\n def backward(ctx, do):\n (qkv, o, lse, bias) = ctx.saved_tensors\n assert not ctx.needs_input_grad[\n 1\n ], \"FlashAttention does not support bias gradient yet\"\n with torch.inference_mode():\n dqkv = torch.empty_like(qkv)\n _flash_attn_backward(\n do,\n qkv[:, :, 0],\n qkv[:, :, 1],\n qkv[:, :, 2],\n o,\n lse,\n dqkv[:, :, 0],\n dqkv[:, :, 1],\n dqkv[:, :, 2],\n bias=bias,\n causal=ctx.causal,\n softmax_scale=ctx.softmax_scale,\n )\n return (dqkv, None, None, None)\n\n\nflash_attn_qkvpacked_func = FlashAttnQKVPackedFunc.apply\n\n\nclass FlashAttnKVPackedFunc(torch.autograd.Function):\n @staticmethod\n def forward(ctx, q, kv, bias=None, causal=False, softmax_scale=None):\n \"\"\"\n q: (batch, seqlen_q, nheads, headdim)\n kv: (batch, seqlen_k, 2, nheads, headdim)\n bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).\n For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).\n ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)\n \"\"\"\n (q, kv) = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, kv]]\n (o, lse, ctx.softmax_scale) = _flash_attn_forward(\n q,\n kv[:, :, 0],\n kv[:, :, 1],\n bias=bias,\n causal=causal,\n softmax_scale=softmax_scale,\n )\n ctx.save_for_backward(q, kv, o, lse, bias)\n ctx.causal = causal\n return o\n\n @staticmethod\n def backward(ctx, do):\n (q, kv, o, lse, bias) = ctx.saved_tensors\n if len(ctx.needs_input_grad) >= 3:\n assert not ctx.needs_input_grad[\n 2\n ], \"FlashAttention does not support bias gradient yet\"\n with torch.inference_mode():\n dq = torch.empty_like(q)\n dkv = torch.empty_like(kv)\n _flash_attn_backward(\n do,\n q,\n kv[:, :, 0],\n kv[:, :, 1],\n o,\n lse,\n dq,\n dkv[:, :, 0],\n dkv[:, :, 1],\n bias=bias,\n causal=ctx.causal,\n softmax_scale=ctx.softmax_scale,\n )\n return (dq, dkv, None, None, None)\n\n\nflash_attn_kvpacked_func = FlashAttnKVPackedFunc.apply\n\n\nclass FlashAttnFunc(torch.autograd.Function):\n @staticmethod\n def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None):\n \"\"\"\n q: (batch_size, seqlen_q, nheads, headdim)\n k, v: (batch_size, seqlen_k, nheads, headdim)\n bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).\n For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).\n ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)\n \"\"\"\n (q, k, v) = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, k, v]]\n (o, lse, ctx.softmax_scale) = _flash_attn_forward(\n q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale\n )\n ctx.save_for_backward(q, k, v, o, lse, bias)\n ctx.causal = causal\n return o\n\n @staticmethod\n def backward(ctx, do):\n (q, k, v, o, lse, bias) = ctx.saved_tensors\n assert not ctx.needs_input_grad[\n 3\n ], \"FlashAttention does not support bias gradient yet\"\n with torch.inference_mode():\n dq = torch.empty_like(q)\n dk = torch.empty_like(k)\n dv = torch.empty_like(v)\n _flash_attn_backward(\n do,\n q,\n k,\n v,\n o,\n lse,\n dq,\n dk,\n dv,\n bias=bias,\n causal=ctx.causal,\n softmax_scale=ctx.softmax_scale,\n )\n return (dq, dk, dv, None, None, None)\n\n\nflash_attn_func = FlashAttnFunc.apply\n" }, { "path": "model/llava/model/language_model/mpt/hf_prefixlm_converter.py", "content": "\"\"\"Converts Huggingface Causal LM to Prefix LM.\n\nConversion does lightweight surgery on a HuggingFace\nCausal LM to convert it to a Prefix LM.\n\nPrefix LMs accepts a `bidirectional_mask` input in `forward`\nand treat the input prompt as the prefix in `generate`.\n\"\"\"\nimport math\nimport warnings\nfrom types import MethodType\nfrom typing import Any, Dict, List, Optional, Tuple, Union\n\nimport torch\nfrom transformers.models.bloom.modeling_bloom import (\n BaseModelOutputWithPastAndCrossAttentions, BloomForCausalLM, BloomModel,\n CausalLMOutputWithCrossAttentions, CrossEntropyLoss)\nfrom transformers.models.bloom.modeling_bloom import \\\n _expand_mask as _expand_mask_bloom\nfrom transformers.models.bloom.modeling_bloom import \\\n _make_causal_mask as _make_causal_mask_bloom\nfrom transformers.models.bloom.modeling_bloom import logging\nfrom transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel\nfrom transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoForCausalLM\nfrom transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM\nfrom transformers.models.gptj.modeling_gptj import GPTJForCausalLM\nfrom transformers.models.opt.modeling_opt import OPTForCausalLM\nfrom transformers.models.opt.modeling_opt import \\\n _expand_mask as _expand_mask_opt\nfrom transformers.models.opt.modeling_opt import \\\n _make_causal_mask as _make_causal_mask_opt\n\nlogger = logging.get_logger(__name__)\n_SUPPORTED_GPT_MODELS = (\n GPT2LMHeadModel,\n GPTJForCausalLM,\n GPTNeoForCausalLM,\n GPTNeoXForCausalLM,\n)\nCAUSAL_GPT_TYPES = Union[\n GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM\n]\n\n\ndef _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_TYPES:\n \"\"\"Converts a GPT-style Causal LM to a Prefix LM.\n\n Supported HuggingFace model classes:\n - `GPT2LMHeadModel`\n - `GPTNeoForCausalLM`\n - `GPTNeoXForCausalLM`\n - `GPTJForCausalLM`\n\n See `convert_hf_causal_lm_to_prefix_lm` for more details.\n \"\"\"\n if hasattr(model, \"_prefix_lm_converted\"):\n return model\n assert isinstance(model, _SUPPORTED_GPT_MODELS)\n assert (\n model.config.add_cross_attention == False\n ), \"Only supports GPT-style decoder-only models\"\n\n def _get_attn_modules(model: CAUSAL_GPT_TYPES) -> List[torch.nn.Module]:\n \"\"\"Helper that gets a list of the model's attention modules.\n\n Each module has a `bias` buffer used for causal masking. The Prefix LM\n conversion adds logic to dynamically manipulate these biases to support\n Prefix LM attention masking.\n \"\"\"\n attn_modules = []\n if isinstance(model, GPTNeoXForCausalLM):\n blocks = model.gpt_neox.layers\n else:\n blocks = model.transformer.h\n for block in blocks:\n if isinstance(model, GPTNeoForCausalLM):\n if block.attn.attention_type != \"global\":\n continue\n attn_module = block.attn.attention\n elif isinstance(model, GPTNeoXForCausalLM):\n attn_module = block.attention\n else:\n attn_module = block.attn\n attn_modules.append(attn_module)\n return attn_modules\n\n setattr(model, \"_original_forward\", getattr(model, \"forward\"))\n setattr(model, \"_original_generate\", getattr(model, \"generate\"))\n\n def forward(\n self: CAUSAL_GPT_TYPES,\n input_ids: Optional[torch.LongTensor] = None,\n past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n bidirectional_mask: Optional[torch.Tensor] = None,\n token_type_ids: Optional[torch.LongTensor] = None,\n position_ids: Optional[torch.LongTensor] = None,\n head_mask: Optional[torch.FloatTensor] = None,\n inputs_embeds: Optional[torch.FloatTensor] = None,\n labels: Optional[torch.LongTensor] = None,\n use_cache: Optional[bool] = None,\n output_attentions: Optional[bool] = None,\n output_hidden_states: Optional[bool] = None,\n return_dict: Optional[bool] = None,\n ):\n \"\"\"Wraps original forward to enable PrefixLM attention.\"\"\"\n\n def call_og_forward():\n if isinstance(self, GPTNeoXForCausalLM):\n return self._original_forward(\n input_ids=input_ids,\n past_key_values=past_key_values,\n attention_mask=attention_mask,\n head_mask=head_mask,\n inputs_embeds=inputs_embeds,\n labels=labels,\n use_cache=use_cache,\n output_attentions=output_attentions,\n output_hidden_states=output_hidden_states,\n return_dict=return_dict,\n )\n else:\n return self._original_forward(\n input_ids=input_ids,\n past_key_values=past_key_values,\n attention_mask=attention_mask,\n token_type_ids=token_type_ids,\n position_ids=position_ids,\n head_mask=head_mask,\n inputs_embeds=inputs_embeds,\n labels=labels,\n use_cache=use_cache,\n output_attentions=output_attentions,\n output_hidden_states=output_hidden_states,\n return_dict=return_dict,\n )\n\n if bidirectional_mask is None:\n return call_og_forward()\n assert isinstance(bidirectional_mask, torch.Tensor)\n attn_modules = _get_attn_modules(model)\n (b, s) = bidirectional_mask.shape\n max_length = attn_modules[0].bias.shape[-1]\n if s > max_length:\n raise ValueError(\n f\"bidirectional_mask sequence length (={s}) exceeds the \"\n + f\"max length allowed by the model ({max_length}).\"\n )\n assert s <= max_length\n if s < max_length:\n pad = torch.zeros(\n (int(b), int(max_length - s)),\n dtype=bidirectional_mask.dtype,\n device=bidirectional_mask.device,\n )\n bidirectional_mask = torch.cat([bidirectional_mask, pad], dim=1)\n bidirectional = bidirectional_mask.unsqueeze(1).unsqueeze(1)\n for attn_module in attn_modules:\n attn_module.bias.data = torch.logical_or(\n attn_module.bias.data, bidirectional\n )\n output = call_og_forward()\n for attn_module in attn_modules:\n attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]\n return output\n\n def generate(self: CAUSAL_GPT_TYPES, *args: tuple, **kwargs: Dict[str, Any]):\n \"\"\"Wraps original generate to enable PrefixLM attention.\"\"\"\n attn_modules = _get_attn_modules(model)\n for attn_module in attn_modules:\n attn_module.bias.data[:] = 1\n output = self._original_generate(*args, **kwargs)\n for attn_module in attn_modules:\n attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]\n return output\n\n setattr(model, \"forward\", MethodType(forward, model))\n setattr(model, \"generate\", MethodType(generate, model))\n setattr(model, \"_prefix_lm_converted\", True)\n return model\n\n\ndef _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCausalLM:\n \"\"\"Converts a BLOOM Causal LM to a Prefix LM.\n\n Supported HuggingFace model classes:\n - `BloomForCausalLM`\n\n See `convert_hf_causal_lm_to_prefix_lm` for more details.\n \"\"\"\n if hasattr(model, \"_prefix_lm_converted\"):\n return model\n assert isinstance(model, BloomForCausalLM)\n assert (\n model.config.add_cross_attention == False\n ), \"Only supports BLOOM decoder-only models\"\n\n def _prepare_attn_mask(\n self: BloomModel,\n attention_mask: torch.Tensor,\n bidirectional_mask: Optional[torch.Tensor],\n input_shape: Tuple[int, int],\n past_key_values_length: int,\n ) -> torch.BoolTensor:\n combined_attention_mask = None\n device = attention_mask.device\n (_, src_length) = input_shape\n if src_length > 1:\n combined_attention_mask = _make_causal_mask_bloom(\n input_shape,\n device=device,\n past_key_values_length=past_key_values_length,\n )\n if bidirectional_mask is not None:\n assert attention_mask.shape == bidirectional_mask.shape\n expanded_bidirectional_mask = _expand_mask_bloom(\n bidirectional_mask, tgt_length=src_length\n )\n combined_attention_mask = torch.logical_and(\n combined_attention_mask, expanded_bidirectional_mask\n )\n expanded_attn_mask = _expand_mask_bloom(attention_mask, tgt_length=src_length)\n combined_attention_mask = (\n expanded_attn_mask\n if combined_attention_mask is None\n else expanded_attn_mask | combined_attention_mask\n )\n return combined_attention_mask\n\n def _build_alibi_tensor(\n self: BloomModel,\n batch_size: int,\n query_length: int,\n key_length: int,\n dtype: torch.dtype,\n device: torch.device,\n ) -> torch.Tensor:\n num_heads = self.config.n_head\n closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))\n base = torch.tensor(\n 2 ** (-(2 ** (-(math.log2(closest_power_of_2) - 3)))),\n device=device,\n dtype=torch.float32,\n )\n powers = torch.arange(\n 1, 1 + closest_power_of_2, device=device, dtype=torch.int32\n )\n slopes = torch.pow(base, powers)\n if closest_power_of_2 != num_heads:\n extra_base = torch.tensor(\n 2 ** (-(2 ** (-(math.log2(2 * closest_power_of_2) - 3)))),\n device=device,\n dtype=torch.float32,\n )\n num_remaining_heads = min(\n closest_power_of_2, num_heads - closest_power_of_2\n )\n extra_powers = torch.arange(\n 1, 1 + 2 * num_remaining_heads, 2, device=device, dtype=torch.int32\n )\n slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)\n qa = torch.arange(query_length, device=device, dtype=torch.int32).view(-1, 1)\n ka = torch.arange(key_length, device=device, dtype=torch.int32).view(1, -1)\n diffs = qa - ka + key_length - query_length\n diffs = -diffs.abs()\n alibi = slopes.view(1, num_heads, 1, 1) * diffs.view(\n 1, 1, query_length, key_length\n )\n alibi = alibi.expand(batch_size, -1, -1, -1).reshape(\n -1, query_length, key_length\n )\n return alibi.to(dtype)\n\n KeyValueT = Tuple[torch.Tensor, torch.Tensor]\n\n def forward(\n self: BloomModel,\n input_ids: Optional[torch.LongTensor] = None,\n past_key_values: Optional[Tuple[KeyValueT, ...]] = None,\n attention_mask: Optional[torch.Tensor] = None,\n bidirectional_mask: Optional[torch.Tensor] = None,\n head_mask: Optional[torch.LongTensor] = None,\n inputs_embeds: Optional[torch.LongTensor] = None,\n use_cache: Optional[bool] = None,\n output_attentions: Optional[bool] = None,\n output_hidden_states: Optional[bool] = None,\n return_dict: Optional[bool] = None,\n **deprecated_arguments,\n ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:\n if deprecated_arguments.pop(\"position_ids\", False) is not False:\n warnings.warn(\n \"`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. \"\n + \"You can safely ignore passing `position_ids`.\",\n FutureWarning,\n )\n if len(deprecated_arguments) > 0:\n raise ValueError(f\"Got unexpected arguments: {deprecated_arguments}\")\n output_attentions = (\n output_attentions\n if output_attentions is not None\n else self.config.output_attentions\n )\n output_hidden_states = (\n output_hidden_states\n if output_hidden_states is not None\n else self.config.output_hidden_states\n )\n use_cache = use_cache if use_cache is not None else self.config.use_cache\n return_dict = (\n return_dict if return_dict is not None else self.config.use_return_dict\n )\n if input_ids is not None and inputs_embeds is not None:\n raise ValueError(\n \"You cannot specify both input_ids and inputs_embeds at the same time\"\n )\n elif input_ids is not None:\n (batch_size, seq_length) = input_ids.shape\n elif inputs_embeds is not None:\n (batch_size, seq_length, _) = inputs_embeds.shape\n else:\n raise ValueError(\"You have to specify either input_ids or inputs_embeds\")\n if past_key_values is None:\n past_key_values = tuple([None] * len(self.h))\n head_mask = self.get_head_mask(head_mask, self.config.n_layer)\n if inputs_embeds is None:\n inputs_embeds = self.word_embeddings(input_ids)\n hidden_states = self.word_embeddings_layernorm(inputs_embeds)\n presents = () if use_cache else None\n all_self_attentions = () if output_attentions else None\n all_hidden_states = () if output_hidden_states else None\n seq_length_with_past = seq_length\n past_key_values_length = 0\n if past_key_values[0] is not None:\n tmp = past_key_values[0][0]\n past_key_values_length = tmp.shape[2]\n seq_length_with_past = seq_length_with_past + past_key_values_length\n if attention_mask is None:\n attention_mask = torch.ones(\n (batch_size, seq_length_with_past), device=hidden_states.device\n )\n else:\n attention_mask = attention_mask.to(hidden_states.device)\n alibi = self._build_alibi_tensor(\n batch_size=batch_size,\n query_length=seq_length,\n key_length=seq_length_with_past,\n dtype=hidden_states.dtype,\n device=hidden_states.device,\n )\n causal_mask = self._prepare_attn_mask(\n attention_mask,\n bidirectional_mask,\n input_shape=(batch_size, seq_length),\n past_key_values_length=past_key_values_length,\n )\n for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):\n if output_hidden_states:\n hst = (hidden_states,)\n all_hidden_states = all_hidden_states + hst\n if self.gradient_checkpointing and self.training:\n if use_cache:\n logger.warning(\n \"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\"\n )\n use_cache = False\n\n def create_custom_forward(module):\n def custom_forward(*inputs):\n return module(\n *inputs,\n use_cache=use_cache,\n output_attentions=output_attentions,\n )\n\n return custom_forward\n\n outputs = torch.utils.checkpoint.checkpoint(\n create_custom_forward(block),\n hidden_states,\n alibi,\n causal_mask,\n head_mask[i],\n )\n else:\n outputs = block(\n hidden_states,\n layer_past=layer_past,\n attention_mask=causal_mask,\n head_mask=head_mask[i],\n use_cache=use_cache,\n output_attentions=output_attentions,\n alibi=alibi,\n )\n hidden_states = outputs[0]\n if use_cache is True:\n presents = presents + (outputs[1],)\n if output_attentions:\n oa = (outputs[2 if use_cache else 1],)\n all_self_attentions = all_self_attentions + oa\n hidden_states = self.ln_f(hidden_states)\n if output_hidden_states:\n hst = (hidden_states,)\n all_hidden_states = all_hidden_states + hst\n if not return_dict:\n return tuple(\n (\n v\n for v in [\n hidden_states,\n presents,\n all_hidden_states,\n all_self_attentions,\n ]\n if v is not None\n )\n )\n return BaseModelOutputWithPastAndCrossAttentions(\n last_hidden_state=hidden_states,\n past_key_values=presents,\n hidden_states=all_hidden_states,\n attentions=all_self_attentions,\n )\n\n setattr(\n model.transformer,\n \"_prepare_attn_mask\",\n MethodType(_prepare_attn_mask, model.transformer),\n )\n setattr(\n model.transformer,\n \"_build_alibi_tensor\",\n MethodType(_build_alibi_tensor, model.transformer),\n )\n setattr(model.transformer, \"forward\", MethodType(forward, model.transformer))\n KeyValueT = Tuple[torch.Tensor, torch.Tensor]\n\n def forward(\n self: BloomForCausalLM,\n input_ids: Optional[torch.LongTensor] = None,\n past_key_values: Optional[Tuple[KeyValueT, ...]] = None,\n attention_mask: Optional[torch.Tensor] = None,\n bidirectional_mask: Optional[torch.Tensor] = None,\n head_mask: Optional[torch.Tensor] = None,\n inputs_embeds: Optional[torch.Tensor] = None,\n labels: Optional[torch.Tensor] = None,\n use_cache: Optional[bool] = None,\n output_attentions: Optional[bool] = None,\n output_hidden_states: Optional[bool] = None,\n return_dict: Optional[bool] = None,\n **deprecated_arguments,\n ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:\n \"\"\"Replacement forward method for BloomCausalLM.\"\"\"\n if deprecated_arguments.pop(\"position_ids\", False) is not False:\n warnings.warn(\n \"`position_ids` have no functionality in BLOOM and will be removed \"\n + \"in v5.0.0. You can safely ignore passing `position_ids`.\",\n FutureWarning,\n )\n if len(deprecated_arguments) > 0:\n raise ValueError(f\"Got unexpected arguments: {deprecated_arguments}\")\n return_dict = (\n return_dict if return_dict is not None else self.config.use_return_dict\n )\n transformer_outputs = self.transformer(\n input_ids,\n past_key_values=past_key_values,\n attention_mask=attention_mask,\n bidirectional_mask=bidirectional_mask,\n head_mask=head_mask,\n inputs_embeds=inputs_embeds,\n use_cache=use_cache,\n output_attentions=output_attentions,\n output_hidden_states=output_hidden_states,\n return_dict=return_dict,\n )\n hidden_states = transformer_outputs[0]\n lm_logits = self.lm_head(hidden_states)\n loss = None\n if labels is not None:\n shift_logits = lm_logits[..., :-1, :].contiguous()\n shift_labels = labels[..., 1:].contiguous()\n (batch_size, seq_length, vocab_size) = shift_logits.shape\n loss_fct = CrossEntropyLoss()\n loss = loss_fct(\n shift_logits.view(batch_size * seq_length, vocab_size),\n shift_labels.view(batch_size * seq_length),\n )\n if not return_dict:\n output = (lm_logits,) + transformer_outputs[1:]\n return (loss,) + output if loss is not None else output\n return CausalLMOutputWithCrossAttentions(\n loss=loss,\n logits=lm_logits,\n past_key_values=transformer_outputs.past_key_values,\n hidden_states=transformer_outputs.hidden_states,\n attentions=transformer_outputs.attentions,\n )\n\n def prepare_inputs_for_generation(\n self: BloomForCausalLM,\n input_ids: torch.LongTensor,\n past: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n **kwargs,\n ) -> dict:\n if past:\n input_ids = input_ids[:, -1].unsqueeze(-1)\n bidirectional_mask = None\n if past[0][0].shape[0] == input_ids.shape[0]:\n past = self._convert_to_bloom_cache(past)\n else:\n bidirectional_mask = torch.ones_like(input_ids)\n return {\n \"input_ids\": input_ids,\n \"past_key_values\": past,\n \"use_cache\": True,\n \"attention_mask\": attention_mask,\n \"bidirectional_mask\": bidirectional_mask,\n }\n\n setattr(model, \"forward\", MethodType(forward, model))\n setattr(\n model,\n \"prepare_inputs_for_generation\",\n MethodType(prepare_inputs_for_generation, model),\n )\n setattr(model, \"_prefix_lm_converted\", True)\n return model\n\n\ndef _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM:\n \"\"\"Converts an OPT Causal LM to a Prefix LM.\n\n Supported HuggingFace model classes:\n - `OPTForCausalLM`\n\n See `convert_hf_causal_lm_to_prefix_lm` for more details.\n \"\"\"\n if hasattr(model, \"_prefix_lm_converted\"):\n return model\n assert isinstance(model, OPTForCausalLM)\n assert (\n model.config.add_cross_attention == False\n ), \"Only supports OPT decoder-only models\"\n setattr(model, \"_original_forward\", getattr(model, \"forward\"))\n setattr(model, \"_original_generate\", getattr(model, \"generate\"))\n model.model.decoder.bidirectional_mask = None\n\n def _prepare_decoder_attention_mask(\n self, attention_mask, input_shape, inputs_embeds, past_key_values_length\n ):\n combined_attention_mask = None\n if input_shape[-1] > 1:\n if self.bidirectional_mask == \"g\":\n (bsz, src_length) = input_shape\n combined_attention_mask = torch.zeros(\n (bsz, 1, src_length, src_length + past_key_values_length),\n dtype=inputs_embeds.dtype,\n device=inputs_embeds.device,\n )\n else:\n combined_attention_mask = _make_causal_mask_opt(\n input_shape,\n inputs_embeds.dtype,\n past_key_values_length=past_key_values_length,\n ).to(inputs_embeds.device)\n if self.bidirectional_mask is not None:\n assert attention_mask.shape == self.bidirectional_mask.shape\n expanded_bidirectional_mask = _expand_mask_opt(\n self.bidirectional_mask,\n inputs_embeds.dtype,\n tgt_len=input_shape[-1],\n ).to(inputs_embeds.device)\n combined_attention_mask = torch.maximum(\n expanded_bidirectional_mask, combined_attention_mask\n )\n if attention_mask is not None:\n expanded_attn_mask = _expand_mask_opt(\n attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]\n ).to(inputs_embeds.device)\n combined_attention_mask = (\n expanded_attn_mask\n if combined_attention_mask is None\n else expanded_attn_mask + combined_attention_mask\n )\n return combined_attention_mask\n\n setattr(\n model.model.decoder,\n \"_prepare_decoder_attention_mask\",\n MethodType(_prepare_decoder_attention_mask, model.model.decoder),\n )\n\n def forward(\n self: OPTForCausalLM,\n input_ids: Optional[torch.LongTensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n bidirectional_mask: Optional[torch.ByteTensor] = None,\n head_mask: Optional[torch.Tensor] = None,\n past_key_values: Optional[List[torch.FloatTensor]] = None,\n inputs_embeds: Optional[torch.FloatTensor] = None,\n labels: Optional[torch.LongTensor] = None,\n use_cache: Optional[bool] = None,\n output_attentions: Optional[bool] = None,\n output_hidden_states: Optional[bool] = None,\n return_dict: Optional[bool] = None,\n ):\n def call_og_forward():\n return self._original_forward(\n input_ids=input_ids,\n attention_mask=attention_mask,\n head_mask=head_mask,\n past_key_values=past_key_values,\n inputs_embeds=inputs_embeds,\n labels=labels,\n use_cache=use_cache,\n output_attentions=output_attentions,\n output_hidden_states=output_hidden_states,\n return_dict=return_dict,\n )\n\n if bidirectional_mask is None:\n return call_og_forward()\n self.model.decoder.bidirectional_mask = bidirectional_mask\n try:\n outputs = call_og_forward()\n except:\n self.model.decoder.bidirectional_mask = None\n raise\n self.model.decoder.bidirectional_mask = None\n return outputs\n\n def generate(self: OPTForCausalLM, *args: tuple, **kwargs: Dict[str, Any]):\n \"\"\"Wraps original generate to enable PrefixLM-style attention.\"\"\"\n self.model.decoder.bidirectional_mask = \"g\"\n try:\n output = self._original_generate(*args, **kwargs)\n except:\n self.model.decoder.bidirectional_mask = None\n raise\n self.model.decoder.bidirectional_mask = None\n return output\n\n setattr(model, \"forward\", MethodType(forward, model))\n setattr(model, \"generate\", MethodType(generate, model))\n setattr(model, \"_prefix_lm_converted\", True)\n return model\n\n\n_SUPPORTED_HF_MODELS = _SUPPORTED_GPT_MODELS + (BloomForCausalLM, OPTForCausalLM)\nCAUSAL_LM_TYPES = Union[\n GPT2LMHeadModel,\n GPTJForCausalLM,\n GPTNeoForCausalLM,\n GPTNeoXForCausalLM,\n BloomForCausalLM,\n OPTForCausalLM,\n]\n\n\ndef convert_hf_causal_lm_to_prefix_lm(model: CAUSAL_LM_TYPES) -> CAUSAL_LM_TYPES:\n \"\"\"Converts a HuggingFace Causal LM to a Prefix LM.\n\n Supported HuggingFace model classes:\n - `GPT2LMHeadModel`\n - `GPTNeoForCausalLM`\n - `GPTNeoXForCausalLM`\n - `GPTJForCausalLM`\n - `BloomForCausalLM`\n - `OPTForCausalLM`\n\n Conversion to a Prefix LM is done by modifying the `forward` method, and possibly also the\n `generate` method and/or select underlying methods depending on the model class.\n\n These changes preserve the model API, but add a new input to `forward`: \"bidirectional_mask\".\n\n Notes on training:\n To actually train the converted model as a Prefix LM, training batches will need to indicate\n the prefix/target structure by including `bidirectional_mask` as part of the batch inputs.\n\n **This is not a standard input and requires custom layers either within or after your dataloader.**\n\n In addition to adding `bidirectional_mask` to the batch, this custom code should modify `labels`\n such that `batch['labels'][batch['bidirectional_mask'] == 1] == -100`.\n That is, the prefix portion of the sequence should not generate any loss. Loss should only be\n generated by the target portion of the sequence.\n\n Notes on `GPTNeoForCausalLM`:\n To simplify the implementation, \"global\" and \"local\" attention layers are handled differently.\n For \"global\" layers, we handle conversion as described above. For \"local\" layers, which use a\n causal attention mask within a restricted local window, we do not alter the masking.\n\n Notes on `forward` method conversion:\n After conversion, the `forward` method will handle a new input, `bidirectional_mask`,\n which should be a [batch_size, seq_length] byte tensor, where 1 indicates token positions\n belonging to the prefix (prefix tokens can attend to one another bidirectionally), and\n 0 indicates token positions belonging to the target.\n\n The new `forward` method will incorporate `bidirectional_mask` (if supplied) into the existing\n causal mask, call the original `forward` method, and (if the causal mask is a buffer) reset\n the causal masks before returning the result.\n\n Notes on `generate` method conversion:\n After conversion, the `generate` method will have the same signature but will internally\n convert all causal masks to be purely bidirectional, call the original `generate` method, and\n (where appropriate) reset the causal masks before returning the result.\n\n This works thanks to the logic of the HuggingFace `generate` API, which first encodes the token\n \"prompt\" passed to `generate` (which is treated as the prefix) and then sequentially generates\n each new token. Encodings are cached as generation happens, so all prefix tokens can attend to one\n another (as expected in a Prefix LM) and generated tokens can only attend to prefix tokens and\n previously-generated tokens (also as expected in a Prefix LM).\n\n To preserve the API, the original methods are renamed to `_original_forward` and\n `_original_generate`, and replaced with new `forward` and `generate` methods that wrap\n them, respectively. Although implementation details vary by model class.\n \"\"\"\n if isinstance(model, _SUPPORTED_GPT_MODELS):\n return _convert_gpt_causal_lm_to_prefix_lm(model)\n elif isinstance(model, BloomForCausalLM):\n return _convert_bloom_causal_lm_to_prefix_lm(model)\n elif isinstance(model, OPTForCausalLM):\n return _convert_opt_causal_lm_to_prefix_lm(model)\n else:\n raise TypeError(\n f\"Cannot convert model to Prefix LM. \"\n + f\"Model does not belong to set of supported HF models:\"\n + f\"\\n{_SUPPORTED_HF_MODELS}\"\n )\n\n\ndef add_bidirectional_mask_if_missing(batch: Dict[str, Any]):\n \"\"\"Attempts to add bidirectional_mask to batch if missing.\n\n Raises:\n KeyError if bidirectional_mask is missing and can't be inferred\n \"\"\"\n if \"bidirectional_mask\" not in batch:\n if batch.get(\"mode\", None) == \"icl_task\":\n batch[\"bidirectional_mask\"] = batch[\"attention_mask\"].clone()\n for i, continuation_indices in enumerate(batch[\"continuation_indices\"]):\n batch[\"bidirectional_mask\"][i, continuation_indices] = 0\n elif \"labels\" in batch and \"attention_mask\" in batch:\n batch[\"bidirectional_mask\"] = torch.logical_and(\n torch.eq(batch[\"attention_mask\"], 1), torch.eq(batch[\"labels\"], -100)\n ).type_as(batch[\"attention_mask\"])\n else:\n raise KeyError(\n \"No bidirectional_mask in batch and not sure how to construct one.\"\n )\n" }, { "path": "model/llava/model/language_model/mpt/meta_init_context.py", "content": "from contextlib import contextmanager\n\nimport torch\nimport torch.nn as nn\n\n\n@contextmanager\ndef init_empty_weights(include_buffers: bool = False):\n \"\"\"Meta initialization context manager.\n\n A context manager under which models are initialized with all parameters\n on the meta device, therefore creating an empty model. Useful when just\n initializing the model would blow the available RAM.\n\n Args:\n include_buffers (`bool`, *optional*, defaults to `False`): Whether or\n not to also put all buffers on the meta device while initializing.\n\n Example:\n ```python\n import torch.nn as nn\n\n # Initialize a model with 100 billions parameters in no time and without using any RAM.\n with init_empty_weights():\n tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])\n ```\n\n \n\n Any model created under this context manager has no weights. As such you can't do something like\n `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`].\n\n \n \"\"\"\n with init_on_device(torch.device(\"meta\"), include_buffers=include_buffers) as f:\n yield f\n\n\n@contextmanager\ndef init_on_device(device: torch.device, include_buffers: bool = False):\n \"\"\"Device initialization context manager.\n\n A context manager under which models are initialized with all parameters\n on the specified device.\n\n Args:\n device (`torch.device`): Device to initialize all parameters on.\n include_buffers (`bool`, *optional*, defaults to `False`): Whether or\n not to also put all buffers on the meta device while initializing.\n\n Example:\n ```python\n import torch.nn as nn\n\n with init_on_device(device=torch.device(\"cuda\")):\n tst = nn.Liner(100, 100) # on `cuda` device\n ```\n \"\"\"\n old_register_parameter = nn.Module.register_parameter\n if include_buffers:\n old_register_buffer = nn.Module.register_buffer\n\n def register_empty_parameter(module, name, param):\n old_register_parameter(module, name, param)\n if param is not None:\n param_cls = type(module._parameters[name])\n kwargs = module._parameters[name].__dict__\n module._parameters[name] = param_cls(\n module._parameters[name].to(device), **kwargs\n )\n\n def register_empty_buffer(module, name, buffer):\n old_register_buffer(module, name, buffer)\n if buffer is not None:\n module._buffers[name] = module._buffers[name].to(device)\n\n if include_buffers:\n tensor_constructors_to_patch = {\n torch_function_name: getattr(torch, torch_function_name)\n for torch_function_name in [\"empty\", \"zeros\", \"ones\", \"full\"]\n }\n else:\n tensor_constructors_to_patch = {}\n\n def patch_tensor_constructor(fn):\n def wrapper(*args, **kwargs):\n kwargs[\"device\"] = device\n return fn(*args, **kwargs)\n\n return wrapper\n\n try:\n nn.Module.register_parameter = register_empty_parameter\n if include_buffers:\n nn.Module.register_buffer = register_empty_buffer\n for torch_function_name in tensor_constructors_to_patch.keys():\n setattr(\n torch,\n torch_function_name,\n patch_tensor_constructor(getattr(torch, torch_function_name)),\n )\n yield\n finally:\n nn.Module.register_parameter = old_register_parameter\n if include_buffers:\n nn.Module.register_buffer = old_register_buffer\n for (\n torch_function_name,\n old_torch_function,\n ) in tensor_constructors_to_patch.items():\n setattr(torch, torch_function_name, old_torch_function)\n" }, { "path": "model/llava/model/language_model/mpt/modeling_mpt.py", "content": "\"\"\"A simple, flexible implementation of a GPT model.\n\nInspired by https://github.com/karpathy/minGPT/blob/master/mingpt/model.py\n\"\"\"\nimport math\nimport warnings\nfrom typing import List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom transformers import (PreTrainedModel, PreTrainedTokenizer,\n PreTrainedTokenizerFast)\nfrom transformers.modeling_outputs import (BaseModelOutputWithPast,\n CausalLMOutputWithPast)\n\nfrom .adapt_tokenizer import AutoTokenizerForMOD, adapt_tokenizer_for_denoising\nfrom .attention import attn_bias_shape, build_attn_bias\nfrom .blocks import MPTBlock\nfrom .configuration_mpt import MPTConfig\nfrom .custom_embedding import SharedEmbedding\nfrom .hf_prefixlm_converter import (add_bidirectional_mask_if_missing,\n convert_hf_causal_lm_to_prefix_lm)\nfrom .meta_init_context import init_empty_weights\nfrom .norm import NORM_CLASS_REGISTRY\nfrom .param_init_fns import MODEL_INIT_REGISTRY, generic_param_init_fn_\n\ntry:\n from .flash_attn_triton import flash_attn_func\nexcept:\n pass\nTokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]\n\n\nclass MPTPreTrainedModel(PreTrainedModel):\n config_class = MPTConfig\n base_model_prefix = \"model\"\n _no_split_modules = [\"MPTBlock\"]\n\n\nclass MPTModel(MPTPreTrainedModel):\n def __init__(self, config: MPTConfig):\n config._validate_config()\n super().__init__(config)\n self.attn_impl = config.attn_config[\"attn_impl\"]\n self.prefix_lm = config.attn_config[\"prefix_lm\"]\n self.attn_uses_sequence_id = config.attn_config[\"attn_uses_sequence_id\"]\n self.alibi = config.attn_config[\"alibi\"]\n self.alibi_bias_max = config.attn_config[\"alibi_bias_max\"]\n if config.init_device == \"mixed\":\n if dist.get_local_rank() == 0:\n config.init_device = \"cpu\"\n else:\n config.init_device = \"meta\"\n if config.norm_type.lower() not in NORM_CLASS_REGISTRY.keys():\n norm_options = \" | \".join(NORM_CLASS_REGISTRY.keys())\n raise NotImplementedError(\n f\"Requested norm type ({config.norm_type}) is not implemented within this repo (Options: {norm_options}).\"\n )\n norm_class = NORM_CLASS_REGISTRY[config.norm_type.lower()]\n self.embedding_fraction = config.embedding_fraction\n self.wte = SharedEmbedding(\n config.vocab_size, config.d_model, device=config.init_device\n )\n if not self.alibi:\n self.wpe = torch.nn.Embedding(\n config.max_seq_len, config.d_model, device=config.init_device\n )\n self.emb_drop = nn.Dropout(config.emb_pdrop)\n self.blocks = nn.ModuleList(\n [\n MPTBlock(device=config.init_device, **config.to_dict())\n for _ in range(config.n_layers)\n ]\n )\n self.norm_f = norm_class(config.d_model, device=config.init_device)\n if config.init_device != \"meta\":\n print(\n f'You are using config.init_device={config.init_device!r}, but you can also use config.init_device=\"meta\" with Composer + FSDP for fast initialization.'\n )\n self.apply(self.param_init_fn)\n self.is_causal = not self.prefix_lm\n self._attn_bias_initialized = False\n self.attn_bias = None\n self.attn_bias_shape = attn_bias_shape(\n self.attn_impl,\n config.n_heads,\n config.max_seq_len,\n self.alibi,\n prefix_lm=self.prefix_lm,\n causal=self.is_causal,\n use_sequence_id=self.attn_uses_sequence_id,\n )\n if config.no_bias:\n for module in self.modules():\n if hasattr(module, \"bias\") and isinstance(module.bias, nn.Parameter):\n if config.verbose:\n warnings.warn(f\"Removing bias ({module.bias}) from {module}.\")\n module.register_parameter(\"bias\", None)\n if config.verbose and config.verbose > 2:\n print(self)\n if \"verbose\" not in self.config.init_config:\n self.config.init_config[\"verbose\"] = self.config.verbose\n if self.config.init_config[\"verbose\"] > 1:\n init_fn_name = self.config.init_config[\"name\"]\n warnings.warn(f\"Using {init_fn_name} initialization.\")\n self.gradient_checkpointing = False\n\n def get_input_embeddings(self):\n return self.wte\n\n def set_input_embeddings(self, value):\n self.wte = value\n\n @torch.no_grad()\n def _attn_bias(\n self,\n device,\n dtype,\n attention_mask: Optional[torch.ByteTensor] = None,\n prefix_mask: Optional[torch.ByteTensor] = None,\n sequence_id: Optional[torch.LongTensor] = None,\n ):\n if not self._attn_bias_initialized:\n if self.attn_bias_shape:\n self.attn_bias = torch.zeros(\n self.attn_bias_shape, device=device, dtype=dtype\n )\n self.attn_bias = build_attn_bias(\n self.attn_impl,\n self.attn_bias,\n self.config.n_heads,\n self.config.max_seq_len,\n causal=self.is_causal,\n alibi=self.alibi,\n alibi_bias_max=self.alibi_bias_max,\n )\n self._attn_bias_initialized = True\n if self.attn_impl == \"flash\":\n return (self.attn_bias, attention_mask)\n if self.attn_bias is not None:\n self.attn_bias = self.attn_bias.to(dtype=dtype, device=device)\n attn_bias = self.attn_bias\n if self.prefix_lm:\n assert isinstance(attn_bias, torch.Tensor)\n assert isinstance(prefix_mask, torch.Tensor)\n attn_bias = self._apply_prefix_mask(attn_bias, prefix_mask)\n if self.attn_uses_sequence_id and sequence_id is not None:\n assert isinstance(attn_bias, torch.Tensor)\n attn_bias = self._apply_sequence_id(attn_bias, sequence_id)\n if attention_mask is not None:\n s_k = attention_mask.shape[-1]\n if attn_bias is None:\n attn_bias = torch.zeros((1, 1, 1, s_k), device=device, dtype=dtype)\n else:\n _s_k = max(0, attn_bias.size(-1) - s_k)\n attn_bias = attn_bias[:, :, :, _s_k:]\n if prefix_mask is not None and attention_mask.shape != prefix_mask.shape:\n raise ValueError(\n f\"attention_mask shape={attention_mask.shape} \"\n + f\"and prefix_mask shape={prefix_mask.shape} are not equal.\"\n )\n min_val = torch.finfo(attn_bias.dtype).min\n attn_bias = attn_bias.masked_fill(\n ~attention_mask.view(-1, 1, 1, s_k), min_val\n )\n return (attn_bias, None)\n\n def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor):\n (s_k, s_q) = attn_bias.shape[-2:]\n if s_k != self.config.max_seq_len or s_q != self.config.max_seq_len:\n raise ValueError(\n \"attn_bias does not match the expected shape. \"\n + f\"The last two dimensions should both be {self.config.max_length} \"\n + f\"but are {s_k} and {s_q}.\"\n )\n seq_len = prefix_mask.shape[-1]\n if seq_len > self.config.max_seq_len:\n raise ValueError(\n f\"prefix_mask sequence length cannot exceed max_seq_len={self.config.max_seq_len}\"\n )\n attn_bias = attn_bias[..., :seq_len, :seq_len]\n causal = torch.tril(\n torch.ones((seq_len, seq_len), dtype=torch.bool, device=prefix_mask.device)\n ).view(1, 1, seq_len, seq_len)\n prefix = prefix_mask.view(-1, 1, 1, seq_len)\n cannot_attend = ~torch.logical_or(causal, prefix.bool())\n min_val = torch.finfo(attn_bias.dtype).min\n attn_bias = attn_bias.masked_fill(cannot_attend, min_val)\n return attn_bias\n\n def _apply_sequence_id(\n self, attn_bias: torch.Tensor, sequence_id: torch.LongTensor\n ):\n seq_len = sequence_id.shape[-1]\n if seq_len > self.config.max_seq_len:\n raise ValueError(\n f\"sequence_id sequence length cannot exceed max_seq_len={self.config.max_seq_len}\"\n )\n attn_bias = attn_bias[..., :seq_len, :seq_len]\n cannot_attend = torch.logical_not(\n torch.eq(sequence_id.view(-1, seq_len, 1), sequence_id.view(-1, 1, seq_len))\n ).unsqueeze(1)\n min_val = torch.finfo(attn_bias.dtype).min\n attn_bias = attn_bias.masked_fill(cannot_attend, min_val)\n return attn_bias\n\n def forward(\n self,\n input_ids: torch.LongTensor,\n past_key_values: Optional[List[Tuple[torch.FloatTensor]]] = None,\n attention_mask: Optional[torch.ByteTensor] = None,\n prefix_mask: Optional[torch.ByteTensor] = None,\n sequence_id: Optional[torch.LongTensor] = None,\n return_dict: Optional[bool] = None,\n output_attentions: Optional[bool] = None,\n output_hidden_states: Optional[bool] = None,\n use_cache: Optional[bool] = None,\n inputs_embeds: Optional[torch.Tensor] = None,\n ):\n return_dict = (\n return_dict if return_dict is not None else self.config.return_dict\n )\n use_cache = use_cache if use_cache is not None else self.config.use_cache\n if attention_mask is not None:\n attention_mask = attention_mask.bool()\n if prefix_mask is not None:\n prefix_mask = prefix_mask.bool()\n if not return_dict:\n raise NotImplementedError(\n \"return_dict False is not implemented yet for MPT\"\n )\n if output_attentions:\n if self.attn_impl != \"torch\":\n raise NotImplementedError(\n \"output_attentions is not implemented for MPT when using attn_impl `flash` or `triton`.\"\n )\n if (\n attention_mask is not None\n and attention_mask[:, 0].sum() != attention_mask.shape[0]\n and self.training\n ):\n raise NotImplementedError(\n \"MPT does not support training with left padding.\"\n )\n if self.prefix_lm and prefix_mask is None:\n raise ValueError(\n \"prefix_mask is a required argument when MPT is configured with prefix_lm=True.\"\n )\n if self.training:\n if self.attn_uses_sequence_id and sequence_id is None:\n raise ValueError(\n \"sequence_id is a required argument when MPT is configured with attn_uses_sequence_id=True \"\n + \"and the model is in train mode.\"\n )\n elif self.attn_uses_sequence_id is False and sequence_id is not None:\n warnings.warn(\n \"MPT received non-None input for `sequence_id` but is configured with attn_uses_sequence_id=False. \"\n + \"This input will be ignored. If you want the model to use `sequence_id`, set attn_uses_sequence_id to True.\"\n )\n if input_ids is not None:\n S = input_ids.size(1)\n assert (\n S <= self.config.max_seq_len\n ), f\"Cannot forward input with seq_len={S}, this model only supports seq_len<={self.config.max_seq_len}\"\n tok_emb = self.wte(input_ids)\n else:\n assert inputs_embeds is not None\n assert (\n self.alibi\n ), \"inputs_embeds is not implemented for MPT unless for alibi.\"\n S = inputs_embeds.size(1)\n tok_emb = inputs_embeds\n if self.alibi:\n x = tok_emb\n else:\n past_position = 0\n if past_key_values is not None:\n if len(past_key_values) != self.config.n_layers:\n raise ValueError(\n f\"past_key_values must provide a past_key_value for each attention \"\n + f\"layer in the network (len(past_key_values)={len(past_key_values)!r}; self.config.n_layers={self.config.n_layers!r}).\"\n )\n past_position = past_key_values[0][0].size(1)\n if self.attn_impl == \"torch\":\n past_position = past_key_values[0][0].size(3)\n if S + past_position > self.config.max_seq_len:\n raise ValueError(\n f\"Cannot forward input with past sequence length {past_position} and current sequence length {S + 1}, this model only supports total sequence length <= {self.config.max_seq_len}.\"\n )\n pos = torch.arange(\n past_position,\n S + past_position,\n dtype=torch.long,\n device=input_ids.device,\n ).unsqueeze(0)\n if attention_mask is not None:\n pos = torch.clamp(\n pos\n - torch.cumsum((~attention_mask).to(torch.int32), dim=1)[\n :, past_position:\n ],\n min=0,\n )\n pos_emb = self.wpe(pos)\n x = tok_emb + pos_emb\n if self.embedding_fraction == 1:\n x = self.emb_drop(x)\n else:\n x_shrunk = x * self.embedding_fraction + x.detach() * (\n 1 - self.embedding_fraction\n )\n assert isinstance(self.emb_drop, nn.Module)\n x = self.emb_drop(x_shrunk)\n (attn_bias, attention_mask) = self._attn_bias(\n device=x.device,\n dtype=torch.float32,\n attention_mask=attention_mask,\n prefix_mask=prefix_mask,\n sequence_id=sequence_id,\n )\n if use_cache and past_key_values is None:\n past_key_values = [() for _ in range(self.config.n_layers)]\n all_hidden_states = () if output_hidden_states else None\n all_self_attns = () if output_attentions else None\n for b_idx, block in enumerate(self.blocks):\n if output_hidden_states:\n assert all_hidden_states is not None\n all_hidden_states = all_hidden_states + (x,)\n past_key_value = (\n past_key_values[b_idx] if past_key_values is not None else None\n )\n if self.gradient_checkpointing and self.training:\n (x, attn_weights, past_key_value) = torch.utils.checkpoint.checkpoint(\n block, x, past_key_value, attn_bias, attention_mask, self.is_causal\n )\n else:\n (x, attn_weights, past_key_value) = block(\n x,\n past_key_value=past_key_value,\n attn_bias=attn_bias,\n attention_mask=attention_mask,\n is_causal=self.is_causal,\n )\n if past_key_values is not None:\n past_key_values[b_idx] = past_key_value\n if output_attentions:\n assert all_self_attns is not None\n all_self_attns = all_self_attns + (attn_weights,)\n x = self.norm_f(x)\n if output_hidden_states:\n assert all_hidden_states is not None\n all_hidden_states = all_hidden_states + (x,)\n return BaseModelOutputWithPast(\n last_hidden_state=x,\n past_key_values=past_key_values,\n hidden_states=all_hidden_states,\n attentions=all_self_attns,\n )\n\n def param_init_fn(self, module):\n init_fn_name = self.config.init_config[\"name\"]\n MODEL_INIT_REGISTRY[init_fn_name](\n module=module,\n n_layers=self.config.n_layers,\n d_model=self.config.d_model,\n **self.config.init_config,\n )\n\n def fsdp_wrap_fn(self, module):\n return isinstance(module, MPTBlock)\n\n def activation_checkpointing_fn(self, module):\n return isinstance(module, MPTBlock)\n\n\nclass MPTForCausalLM(MPTPreTrainedModel):\n def __init__(self, config: MPTConfig):\n super().__init__(config)\n if not config.tie_word_embeddings:\n raise ValueError(\"MPTForCausalLM only supports tied word embeddings\")\n print(f\"Instantiating an MPTForCausalLM model from {__file__}\")\n self.transformer = MPTModel(config)\n for child in self.transformer.children():\n if isinstance(child, torch.nn.ModuleList):\n continue\n if isinstance(child, torch.nn.Module):\n child._fsdp_wrap = True\n self.logit_scale = None\n if config.logit_scale is not None:\n logit_scale = config.logit_scale\n if isinstance(logit_scale, str):\n if logit_scale == \"inv_sqrt_d_model\":\n logit_scale = 1 / math.sqrt(config.d_model)\n else:\n raise ValueError(\n f\"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.\"\n )\n self.logit_scale = logit_scale\n\n def get_input_embeddings(self):\n return self.transformer.wte\n\n def set_input_embeddings(self, value):\n self.transformer.wte = value\n\n def get_output_embeddings(self):\n return self.transformer.wte\n\n def set_output_embeddings(self, new_embeddings):\n self.transformer.wte = new_embeddings\n\n def set_decoder(self, decoder):\n self.transformer = decoder\n\n def get_decoder(self):\n return self.transformer\n\n def forward(\n self,\n input_ids: torch.LongTensor,\n past_key_values: Optional[List[Tuple[torch.FloatTensor]]] = None,\n attention_mask: Optional[torch.ByteTensor] = None,\n prefix_mask: Optional[torch.ByteTensor] = None,\n sequence_id: Optional[torch.LongTensor] = None,\n labels: Optional[torch.LongTensor] = None,\n return_dict: Optional[bool] = None,\n output_attentions: Optional[bool] = None,\n output_hidden_states: Optional[bool] = None,\n use_cache: Optional[bool] = None,\n inputs_embeds: Optional[torch.FloatTensor] = None,\n ):\n return_dict = (\n return_dict if return_dict is not None else self.config.return_dict\n )\n use_cache = use_cache if use_cache is not None else self.config.use_cache\n if inputs_embeds is not None:\n raise NotImplementedError(\n \"inputs_embeds has to be None (for hf/peft support).\"\n )\n outputs = self.transformer(\n input_ids=input_ids,\n past_key_values=past_key_values,\n attention_mask=attention_mask,\n prefix_mask=prefix_mask,\n sequence_id=sequence_id,\n return_dict=return_dict,\n output_attentions=output_attentions,\n output_hidden_states=output_hidden_states,\n use_cache=use_cache,\n )\n logits = self.transformer.wte(\n outputs.last_hidden_state.to(self.transformer.wte.weight.device), True\n )\n if self.logit_scale is not None:\n if self.logit_scale == 0:\n warnings.warn(\n f\"Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.\"\n )\n logits *= self.logit_scale\n loss = None\n if labels is not None:\n labels = torch.roll(labels, shifts=-1)\n labels[:, -1] = -100\n loss = F.cross_entropy(\n logits.view(-1, logits.size(-1)), labels.to(logits.device).view(-1)\n )\n return CausalLMOutputWithPast(\n loss=loss,\n logits=logits,\n past_key_values=outputs.past_key_values,\n hidden_states=outputs.hidden_states,\n attentions=outputs.attentions,\n )\n\n def param_init_fn(self, module):\n init_fn_name = self.config.init_config[\"name\"]\n MODEL_INIT_REGISTRY[init_fn_name](\n module=module,\n n_layers=self.config.n_layers,\n d_model=self.config.d_model,\n **self.config.init_config,\n )\n\n def fsdp_wrap_fn(self, module):\n return isinstance(module, MPTBlock)\n\n def activation_checkpointing_fn(self, module):\n return isinstance(module, MPTBlock)\n\n def prepare_inputs_for_generation(\n self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs\n ):\n if inputs_embeds is not None:\n raise NotImplementedError(\"inputs_embeds is not implemented for MPT yet\")\n attention_mask = kwargs[\"attention_mask\"].bool()\n if attention_mask[:, -1].sum() != attention_mask.shape[0]:\n raise NotImplementedError(\n \"MPT does not support generation with right padding.\"\n )\n if self.transformer.attn_uses_sequence_id and self.training:\n sequence_id = torch.zeros_like(input_ids[:1])\n else:\n sequence_id = None\n if past_key_values is not None:\n input_ids = input_ids[:, -1].unsqueeze(-1)\n if self.transformer.prefix_lm:\n prefix_mask = torch.ones_like(attention_mask)\n if kwargs.get(\"use_cache\") == False:\n raise NotImplementedError(\n \"MPT with prefix_lm=True does not support use_cache=False.\"\n )\n else:\n prefix_mask = None\n return {\n \"input_ids\": input_ids,\n \"attention_mask\": attention_mask,\n \"prefix_mask\": prefix_mask,\n \"sequence_id\": sequence_id,\n \"past_key_values\": past_key_values,\n \"use_cache\": kwargs.get(\"use_cache\", True),\n }\n\n @staticmethod\n def _reorder_cache(past_key_values, beam_idx):\n \"\"\"Used by HuggingFace generate when using beam search with kv-caching.\n\n See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133\n for an example in transformers.\n \"\"\"\n reordered_past = []\n for layer_past in past_key_values:\n reordered_past += [\n tuple(\n (past_state.index_select(0, beam_idx) for past_state in layer_past)\n )\n ]\n return reordered_past\n" }, { "path": "model/llava/model/language_model/mpt/norm.py", "content": "import torch\n\n\ndef _cast_if_autocast_enabled(tensor):\n if torch.is_autocast_enabled():\n if tensor.device.type == \"cuda\":\n dtype = torch.get_autocast_gpu_dtype()\n elif tensor.device.type == \"cpu\":\n dtype = torch.get_autocast_cpu_dtype()\n else:\n raise NotImplementedError()\n return tensor.to(dtype=dtype)\n return tensor\n\n\nclass LPLayerNorm(torch.nn.LayerNorm):\n def __init__(\n self,\n normalized_shape,\n eps=1e-05,\n elementwise_affine=True,\n device=None,\n dtype=None,\n ):\n super().__init__(\n normalized_shape=normalized_shape,\n eps=eps,\n elementwise_affine=elementwise_affine,\n device=device,\n dtype=dtype,\n )\n\n def forward(self, x):\n module_device = x.device\n downcast_x = _cast_if_autocast_enabled(x)\n downcast_weight = (\n _cast_if_autocast_enabled(self.weight)\n if self.weight is not None\n else self.weight\n )\n downcast_bias = (\n _cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias\n )\n with torch.autocast(enabled=False, device_type=module_device.type):\n return torch.nn.functional.layer_norm(\n downcast_x,\n self.normalized_shape,\n downcast_weight,\n downcast_bias,\n self.eps,\n )\n\n\ndef rms_norm(x, weight=None, eps=1e-05):\n output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)\n if weight is not None:\n return output * weight\n return output\n\n\nclass RMSNorm(torch.nn.Module):\n def __init__(\n self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None\n ):\n super().__init__()\n self.eps = eps\n if weight:\n self.weight = torch.nn.Parameter(\n torch.ones(normalized_shape, dtype=dtype, device=device)\n )\n else:\n self.register_parameter(\"weight\", None)\n\n def forward(self, x):\n return rms_norm(x.float(), self.weight, self.eps).to(dtype=x.dtype)\n\n\nclass LPRMSNorm(RMSNorm):\n def __init__(\n self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None\n ):\n super().__init__(\n normalized_shape=normalized_shape,\n eps=eps,\n weight=weight,\n dtype=dtype,\n device=device,\n )\n\n def forward(self, x):\n downcast_x = _cast_if_autocast_enabled(x)\n downcast_weight = (\n _cast_if_autocast_enabled(self.weight)\n if self.weight is not None\n else self.weight\n )\n with torch.autocast(enabled=False, device_type=x.device.type):\n return rms_norm(downcast_x, downcast_weight, self.eps).to(dtype=x.dtype)\n\n\nNORM_CLASS_REGISTRY = {\n \"layernorm\": torch.nn.LayerNorm,\n \"low_precision_layernorm\": LPLayerNorm,\n \"rmsnorm\": RMSNorm,\n \"low_precision_rmsnorm\": LPRMSNorm,\n}\n" }, { "path": "model/llava/model/language_model/mpt/param_init_fns.py", "content": "import math\nimport warnings\nfrom collections.abc import Sequence\nfrom functools import partial\nfrom typing import Optional, Tuple, Union\n\nimport torch\nfrom torch import nn\n\nfrom .norm import NORM_CLASS_REGISTRY\n\n\ndef torch_default_param_init_fn_(module: nn.Module, verbose: int = 0, **kwargs):\n del kwargs\n if verbose > 1:\n warnings.warn(f\"Initializing network using module's reset_parameters attribute\")\n if hasattr(module, \"reset_parameters\"):\n module.reset_parameters()\n\n\ndef fused_init_helper_(module: nn.Module, init_fn_):\n _fused = getattr(module, \"_fused\", None)\n if _fused is None:\n raise RuntimeError(f\"Internal logic error\")\n (dim, splits) = _fused\n splits = (0, *splits, module.weight.size(dim))\n for s, e in zip(splits[:-1], splits[1:]):\n slice_indices = [slice(None)] * module.weight.ndim\n slice_indices[dim] = slice(s, e)\n init_fn_(module.weight[slice_indices])\n\n\ndef generic_param_init_fn_(\n module: nn.Module,\n init_fn_,\n n_layers: int,\n d_model: Optional[int] = None,\n init_div_is_residual: Union[int, float, str, bool] = True,\n emb_init_std: Optional[float] = None,\n emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,\n verbose: int = 0,\n **kwargs,\n):\n del kwargs\n if verbose > 1:\n warnings.warn(f\"If model has bias parameters they are initialized to 0.\")\n init_div_is_residual = init_div_is_residual\n if init_div_is_residual is False:\n div_is_residual = 1.0\n elif init_div_is_residual is True:\n div_is_residual = math.sqrt(2 * n_layers)\n elif isinstance(init_div_is_residual, float) or isinstance(\n init_div_is_residual, int\n ):\n div_is_residual = init_div_is_residual\n elif isinstance(init_div_is_residual, str) and init_div_is_residual.isnumeric():\n div_is_residual = float(init_div_is_residual)\n else:\n div_is_residual = 1.0\n raise ValueError(\n f\"Expected init_div_is_residual to be boolean or numeric, got {init_div_is_residual}\"\n )\n if init_div_is_residual is not False:\n if verbose > 1:\n warnings.warn(\n f\"Initializing _is_residual layers then dividing them by {div_is_residual:.3f}. \"\n + f\"Set `init_div_is_residual: false` in init config to disable this.\"\n )\n if isinstance(module, nn.Linear):\n if hasattr(module, \"_fused\"):\n fused_init_helper_(module, init_fn_)\n else:\n init_fn_(module.weight)\n if module.bias is not None:\n torch.nn.init.zeros_(module.bias)\n if init_div_is_residual is not False and getattr(module, \"_is_residual\", False):\n with torch.no_grad():\n module.weight.div_(div_is_residual)\n elif isinstance(module, nn.Embedding):\n if emb_init_std is not None:\n std = emb_init_std\n if std == 0:\n warnings.warn(f\"Embedding layer initialized to 0.\")\n emb_init_fn_ = partial(torch.nn.init.normal_, mean=0.0, std=std)\n if verbose > 1:\n warnings.warn(\n f\"Embedding layer initialized using normal distribution with mean=0 and std={std!r}.\"\n )\n elif emb_init_uniform_lim is not None:\n lim = emb_init_uniform_lim\n if isinstance(lim, Sequence):\n if len(lim) > 2:\n raise ValueError(\n f\"Uniform init requires a min and a max limit. User input: {lim}.\"\n )\n if lim[0] == lim[1]:\n warnings.warn(f\"Embedding layer initialized to {lim[0]}.\")\n else:\n if lim == 0:\n warnings.warn(f\"Embedding layer initialized to 0.\")\n lim = [-lim, lim]\n (a, b) = lim\n emb_init_fn_ = partial(torch.nn.init.uniform_, a=a, b=b)\n if verbose > 1:\n warnings.warn(\n f\"Embedding layer initialized using uniform distribution in range {lim}.\"\n )\n else:\n emb_init_fn_ = init_fn_\n emb_init_fn_(module.weight)\n elif isinstance(module, tuple(set(NORM_CLASS_REGISTRY.values()))):\n if verbose > 1:\n warnings.warn(\n f\"Norm weights are set to 1. If norm layer has a bias it is initialized to 0.\"\n )\n if hasattr(module, \"weight\") and module.weight is not None:\n torch.nn.init.ones_(module.weight)\n if hasattr(module, \"bias\") and module.bias is not None:\n torch.nn.init.zeros_(module.bias)\n elif isinstance(module, nn.MultiheadAttention):\n if module._qkv_same_embed_dim:\n assert module.in_proj_weight is not None\n assert (\n module.q_proj_weight is None\n and module.k_proj_weight is None\n and (module.v_proj_weight is None)\n )\n assert d_model is not None\n _d = d_model\n splits = (0, _d, 2 * _d, 3 * _d)\n for s, e in zip(splits[:-1], splits[1:]):\n init_fn_(module.in_proj_weight[s:e])\n else:\n assert (\n module.q_proj_weight is not None\n and module.k_proj_weight is not None\n and (module.v_proj_weight is not None)\n )\n assert module.in_proj_weight is None\n init_fn_(module.q_proj_weight)\n init_fn_(module.k_proj_weight)\n init_fn_(module.v_proj_weight)\n if module.in_proj_bias is not None:\n torch.nn.init.zeros_(module.in_proj_bias)\n if module.bias_k is not None:\n torch.nn.init.zeros_(module.bias_k)\n if module.bias_v is not None:\n torch.nn.init.zeros_(module.bias_v)\n init_fn_(module.out_proj.weight)\n if init_div_is_residual is not False and getattr(\n module.out_proj, \"_is_residual\", False\n ):\n with torch.no_grad():\n module.out_proj.weight.div_(div_is_residual)\n if module.out_proj.bias is not None:\n torch.nn.init.zeros_(module.out_proj.bias)\n else:\n for _ in module.parameters(recurse=False):\n raise NotImplementedError(\n f\"{module.__class__.__name__} parameters are not initialized by param_init_fn.\"\n )\n\n\ndef _normal_init_(std, mean=0.0):\n return partial(torch.nn.init.normal_, mean=mean, std=std)\n\n\ndef _normal_param_init_fn_(\n module: nn.Module,\n std: float,\n n_layers: int,\n d_model: Optional[int] = None,\n init_div_is_residual: Union[int, float, str, bool] = True,\n emb_init_std: Optional[float] = None,\n emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,\n verbose: int = 0,\n **kwargs,\n):\n del kwargs\n init_fn_ = _normal_init_(std=std)\n if verbose > 1:\n warnings.warn(f\"Using torch.nn.init.normal_ init fn mean=0.0, std={std}\")\n generic_param_init_fn_(\n module=module,\n init_fn_=init_fn_,\n d_model=d_model,\n n_layers=n_layers,\n init_div_is_residual=init_div_is_residual,\n emb_init_std=emb_init_std,\n emb_init_uniform_lim=emb_init_uniform_lim,\n verbose=verbose,\n )\n\n\ndef baseline_param_init_fn_(\n module: nn.Module,\n init_std: float,\n n_layers: int,\n d_model: Optional[int] = None,\n init_div_is_residual: Union[int, float, str, bool] = True,\n emb_init_std: Optional[float] = None,\n emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,\n verbose: int = 0,\n **kwargs,\n):\n del kwargs\n if init_std is None:\n raise ValueError(\n \"You must set model.init_config['init_std'] to a float value to use the default initialization scheme.\"\n )\n _normal_param_init_fn_(\n module=module,\n std=init_std,\n d_model=d_model,\n n_layers=n_layers,\n init_div_is_residual=init_div_is_residual,\n emb_init_std=emb_init_std,\n emb_init_uniform_lim=emb_init_uniform_lim,\n verbose=verbose,\n )\n\n\ndef small_param_init_fn_(\n module: nn.Module,\n n_layers: int,\n d_model: int,\n init_div_is_residual: Union[int, float, str, bool] = True,\n emb_init_std: Optional[float] = None,\n emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,\n verbose: int = 0,\n **kwargs,\n):\n del kwargs\n std = math.sqrt(2 / (5 * d_model))\n _normal_param_init_fn_(\n module=module,\n std=std,\n d_model=d_model,\n n_layers=n_layers,\n init_div_is_residual=init_div_is_residual,\n emb_init_std=emb_init_std,\n emb_init_uniform_lim=emb_init_uniform_lim,\n verbose=verbose,\n )\n\n\ndef neox_param_init_fn_(\n module: nn.Module,\n n_layers: int,\n d_model: int,\n emb_init_std: Optional[float] = None,\n emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,\n verbose: int = 0,\n **kwargs,\n):\n \"\"\"From section 2.3.1 of GPT-NeoX-20B:\n\n An Open-Source AutoregressiveLanguage Model โ€” Black et. al. (2022)\n see https://github.com/EleutherAI/gpt-neox/blob/9610391ab319403cef079b438edd016a2443af54/megatron/model/init_functions.py#L151\n and https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/transformer.py\n \"\"\"\n del kwargs\n residual_div = n_layers / math.sqrt(10)\n if verbose > 1:\n warnings.warn(f\"setting init_div_is_residual to {residual_div}\")\n small_param_init_fn_(\n module=module,\n d_model=d_model,\n n_layers=n_layers,\n init_div_is_residual=residual_div,\n emb_init_std=emb_init_std,\n emb_init_uniform_lim=emb_init_uniform_lim,\n verbose=verbose,\n )\n\n\ndef kaiming_uniform_param_init_fn_(\n module: nn.Module,\n n_layers: int,\n d_model: Optional[int] = None,\n init_div_is_residual: Union[int, float, str, bool] = True,\n emb_init_std: Optional[float] = None,\n emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,\n init_gain: float = 0,\n fan_mode: str = \"fan_in\",\n init_nonlinearity: str = \"leaky_relu\",\n verbose: int = 0,\n **kwargs,\n):\n del kwargs\n if verbose > 1:\n warnings.warn(\n f\"Using nn.init.kaiming_uniform_ init fn with parameters: \"\n + f\"a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}\"\n )\n kaiming_uniform_ = partial(\n nn.init.kaiming_uniform_,\n a=init_gain,\n mode=fan_mode,\n nonlinearity=init_nonlinearity,\n )\n generic_param_init_fn_(\n module=module,\n init_fn_=kaiming_uniform_,\n d_model=d_model,\n n_layers=n_layers,\n init_div_is_residual=init_div_is_residual,\n emb_init_std=emb_init_std,\n emb_init_uniform_lim=emb_init_uniform_lim,\n verbose=verbose,\n )\n\n\ndef kaiming_normal_param_init_fn_(\n module: nn.Module,\n n_layers: int,\n d_model: Optional[int] = None,\n init_div_is_residual: Union[int, float, str, bool] = True,\n emb_init_std: Optional[float] = None,\n emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,\n init_gain: float = 0,\n fan_mode: str = \"fan_in\",\n init_nonlinearity: str = \"leaky_relu\",\n verbose: int = 0,\n **kwargs,\n):\n del kwargs\n if verbose > 1:\n warnings.warn(\n f\"Using nn.init.kaiming_normal_ init fn with parameters: \"\n + f\"a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}\"\n )\n kaiming_normal_ = partial(\n torch.nn.init.kaiming_normal_,\n a=init_gain,\n mode=fan_mode,\n nonlinearity=init_nonlinearity,\n )\n generic_param_init_fn_(\n module=module,\n init_fn_=kaiming_normal_,\n d_model=d_model,\n n_layers=n_layers,\n init_div_is_residual=init_div_is_residual,\n emb_init_std=emb_init_std,\n emb_init_uniform_lim=emb_init_uniform_lim,\n verbose=verbose,\n )\n\n\ndef xavier_uniform_param_init_fn_(\n module: nn.Module,\n n_layers: int,\n d_model: Optional[int] = None,\n init_div_is_residual: Union[int, float, str, bool] = True,\n emb_init_std: Optional[float] = None,\n emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,\n init_gain: float = 0,\n verbose: int = 0,\n **kwargs,\n):\n del kwargs\n xavier_uniform_ = partial(torch.nn.init.xavier_uniform_, gain=init_gain)\n if verbose > 1:\n warnings.warn(\n f\"Using torch.nn.init.xavier_uniform_ init fn with parameters: \"\n + f\"gain={init_gain}\"\n )\n generic_param_init_fn_(\n module=module,\n init_fn_=xavier_uniform_,\n d_model=d_model,\n n_layers=n_layers,\n init_div_is_residual=init_div_is_residual,\n emb_init_std=emb_init_std,\n emb_init_uniform_lim=emb_init_uniform_lim,\n verbose=verbose,\n )\n\n\ndef xavier_normal_param_init_fn_(\n module: nn.Module,\n n_layers: int,\n d_model: Optional[int] = None,\n init_div_is_residual: Union[int, float, str, bool] = True,\n emb_init_std: Optional[float] = None,\n emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,\n init_gain: float = 0,\n verbose: int = 0,\n **kwargs,\n):\n xavier_normal_ = partial(torch.nn.init.xavier_normal_, gain=init_gain)\n if verbose > 1:\n warnings.warn(\n f\"Using torch.nn.init.xavier_normal_ init fn with parameters: \"\n + f\"gain={init_gain}\"\n )\n generic_param_init_fn_(\n module=module,\n init_fn_=xavier_normal_,\n d_model=d_model,\n n_layers=n_layers,\n init_div_is_residual=init_div_is_residual,\n emb_init_std=emb_init_std,\n emb_init_uniform_lim=emb_init_uniform_lim,\n verbose=verbose,\n )\n\n\nMODEL_INIT_REGISTRY = {\n \"default_\": torch_default_param_init_fn_,\n \"baseline_\": baseline_param_init_fn_,\n \"kaiming_uniform_\": kaiming_uniform_param_init_fn_,\n \"kaiming_normal_\": kaiming_normal_param_init_fn_,\n \"neox_init_\": neox_param_init_fn_,\n \"small_init_\": small_param_init_fn_,\n \"xavier_uniform_\": xavier_uniform_param_init_fn_,\n \"xavier_normal_\": xavier_normal_param_init_fn_,\n}\n" }, { "path": "model/llava/model/llava_arch.py", "content": "# Copyright 2023 Haotian Liu\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom abc import ABC, abstractmethod\n\nimport torch\nimport torch.nn as nn\n\n# from llava.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN\nfrom utils.utils import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,\n DEFAULT_IMAGE_PATCH_TOKEN, IGNORE_INDEX,\n IMAGE_TOKEN_INDEX)\n\nfrom .multimodal_encoder.builder import build_vision_tower\n\n\nclass LlavaMetaModel:\n def __init__(self, config):\n super(LlavaMetaModel, self).__init__(config)\n\n if hasattr(config, \"mm_vision_tower\"):\n self.vision_tower = build_vision_tower(config, delay_load=True)\n self.mm_projector = nn.Linear(config.mm_hidden_size, config.hidden_size)\n\n def get_vision_tower(self):\n vision_tower = getattr(self, \"vision_tower\", None)\n if type(vision_tower) is list:\n vision_tower = vision_tower[0]\n return vision_tower\n\n def initialize_vision_modules(self, model_args, fsdp=None):\n vision_tower = model_args.vision_tower\n mm_vision_select_layer = model_args.mm_vision_select_layer\n mm_vision_select_feature = model_args.mm_vision_select_feature\n pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter\n\n self.config.mm_vision_tower = vision_tower\n\n vision_tower = build_vision_tower(model_args)\n\n if fsdp is not None and len(fsdp) > 0:\n self.vision_tower = [vision_tower]\n else:\n self.vision_tower = vision_tower\n\n self.config.use_mm_proj = True\n self.config.mm_hidden_size = vision_tower.hidden_size\n self.config.mm_vision_select_layer = mm_vision_select_layer\n self.config.mm_vision_select_feature = mm_vision_select_feature\n\n if not hasattr(self, \"mm_projector\"):\n self.mm_projector = nn.Linear(\n self.config.mm_hidden_size, self.config.hidden_size\n )\n\n if pretrain_mm_mlp_adapter is not None:\n mm_projector_weights = torch.load(\n pretrain_mm_mlp_adapter, map_location=\"cpu\"\n )\n\n def get_w(weights, keyword):\n return {\n k.split(keyword + \".\")[1]: v\n for k, v in weights.items()\n if keyword in k\n }\n\n self.mm_projector.load_state_dict(\n get_w(mm_projector_weights, \"mm_projector\")\n )\n\n\nclass LlavaMetaForCausalLM(ABC):\n @abstractmethod\n def get_model(self):\n pass\n\n def get_vision_tower(self):\n return self.get_model().get_vision_tower()\n\n def encode_images(self, images):\n image_features = self.get_model().get_vision_tower()(images)\n image_features = self.get_model().mm_projector(image_features)\n return image_features\n\n def prepare_inputs_labels_for_multimodal(\n self, input_ids, attention_mask, past_key_values, labels, images\n ):\n vision_tower = self.get_vision_tower()\n if vision_tower is None or images is None or input_ids.shape[1] == 1:\n if (\n past_key_values is not None\n and vision_tower is not None\n and images is not None\n and input_ids.shape[1] == 1\n ):\n attention_mask = torch.ones(\n (attention_mask.shape[0], past_key_values[-1][-1].shape[-2] + 1),\n dtype=attention_mask.dtype,\n device=attention_mask.device,\n )\n return input_ids, attention_mask, past_key_values, None, labels\n\n if type(images) is list or images.ndim == 5:\n concat_images = torch.cat([image for image in images], dim=0)\n image_features = self.encode_images(concat_images)\n split_sizes = [image.shape[0] for image in images]\n image_features = torch.split(image_features, split_sizes, dim=0)\n image_features = [x.flatten(0, 1) for x in image_features]\n else:\n image_features = self.encode_images(images)\n\n new_input_embeds = []\n new_labels = [] if labels is not None else None\n cur_image_idx = 0\n for batch_idx, cur_input_ids in enumerate(input_ids):\n if (cur_input_ids == IMAGE_TOKEN_INDEX).sum() == 0:\n # multimodal LLM, but the current sample is not multimodal\n cur_input_embeds = self.get_model().embed_tokens(cur_input_ids)\n cur_input_embeds = (\n cur_input_embeds\n + (\n 0.0 * self.get_model().mm_projector(vision_tower.dummy_feature)\n ).sum()\n )\n new_input_embeds.append(cur_input_embeds)\n if labels is not None:\n new_labels.append(labels[batch_idx])\n cur_image_idx += 1\n continue\n image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]\n cur_new_input_embeds = []\n if labels is not None:\n cur_labels = labels[batch_idx]\n cur_new_labels = []\n assert cur_labels.shape == cur_input_ids.shape\n while image_token_indices.numel() > 0:\n cur_image_features = image_features[cur_image_idx]\n image_token_start = image_token_indices[0]\n if getattr(self.config, \"tune_mm_mlp_adapter\", False) and getattr(\n self.config, \"mm_use_im_start_end\", False\n ):\n cur_new_input_embeds.append(\n self.get_model()\n .embed_tokens(cur_input_ids[: image_token_start - 1])\n .detach()\n )\n cur_new_input_embeds.append(\n self.get_model().embed_tokens(\n cur_input_ids[image_token_start - 1 : image_token_start]\n )\n )\n cur_new_input_embeds.append(cur_image_features)\n cur_new_input_embeds.append(\n self.get_model().embed_tokens(\n cur_input_ids[image_token_start + 1 : image_token_start + 2]\n )\n )\n if labels is not None:\n cur_new_labels.append(cur_labels[:image_token_start])\n cur_new_labels.append(\n torch.full(\n (cur_image_features.shape[0],),\n IGNORE_INDEX,\n device=labels.device,\n dtype=labels.dtype,\n )\n )\n cur_new_labels.append(\n cur_labels[image_token_start : image_token_start + 1]\n )\n cur_labels = cur_labels[image_token_start + 2 :]\n elif getattr(self.config, \"mm_use_im_start_end\", False):\n cur_new_input_embeds.append(\n self.get_model().embed_tokens(cur_input_ids[:image_token_start])\n )\n cur_new_input_embeds.append(cur_image_features)\n cur_new_input_embeds.append(\n self.get_model().embed_tokens(\n cur_input_ids[image_token_start + 1 : image_token_start + 2]\n )\n )\n if labels is not None:\n cur_new_labels.append(cur_labels[:image_token_start])\n cur_new_labels.append(\n torch.full(\n (cur_image_features.shape[0],),\n IGNORE_INDEX,\n device=labels.device,\n dtype=labels.dtype,\n )\n )\n cur_new_labels.append(\n cur_labels[image_token_start + 1 : image_token_start + 2]\n )\n cur_labels = cur_labels[image_token_start + 2 :]\n else:\n cur_new_input_embeds.append(\n self.get_model().embed_tokens(cur_input_ids[:image_token_start])\n )\n cur_new_input_embeds.append(cur_image_features)\n if labels is not None:\n cur_new_labels.append(cur_labels[:image_token_start])\n cur_new_labels.append(\n torch.full(\n (cur_image_features.shape[0],),\n IGNORE_INDEX,\n device=labels.device,\n dtype=labels.dtype,\n )\n )\n cur_labels = cur_labels[image_token_start + 1 :]\n cur_image_idx += 1\n if getattr(self.config, \"tune_mm_mlp_adapter\", False) and getattr(\n self.config, \"mm_use_im_start_end\", False\n ):\n cur_input_ids = cur_input_ids[image_token_start + 2 :]\n elif getattr(self.config, \"mm_use_im_start_end\", False):\n cur_input_ids = cur_input_ids[image_token_start + 2 :]\n else:\n cur_input_ids = cur_input_ids[image_token_start + 1 :]\n image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]\n if cur_input_ids.numel() > 0:\n if getattr(self.config, \"tune_mm_mlp_adapter\", False) and getattr(\n self.config, \"mm_use_im_start_end\", False\n ):\n cur_new_input_embeds.append(\n self.get_model().embed_tokens(cur_input_ids).detach()\n )\n elif getattr(self.config, \"mm_use_im_start_end\", False):\n cur_new_input_embeds.append(\n self.get_model().embed_tokens(cur_input_ids)\n )\n else:\n cur_new_input_embeds.append(\n self.get_model().embed_tokens(cur_input_ids)\n )\n if labels is not None:\n cur_new_labels.append(cur_labels)\n cur_new_input_embeds = [\n x.to(device=self.device) for x in cur_new_input_embeds\n ]\n cur_new_input_embeds = torch.cat(cur_new_input_embeds, dim=0)\n new_input_embeds.append(cur_new_input_embeds)\n if labels is not None:\n cur_new_labels = torch.cat(cur_new_labels, dim=0)\n new_labels.append(cur_new_labels)\n\n if any(x.shape != new_input_embeds[0].shape for x in new_input_embeds):\n max_len = max(x.shape[0] for x in new_input_embeds)\n\n new_input_embeds_align = []\n for cur_new_embed in new_input_embeds:\n cur_new_embed = torch.cat(\n (\n cur_new_embed,\n torch.zeros(\n (max_len - cur_new_embed.shape[0], cur_new_embed.shape[1]),\n dtype=cur_new_embed.dtype,\n device=cur_new_embed.device,\n ),\n ),\n dim=0,\n )\n new_input_embeds_align.append(cur_new_embed)\n new_input_embeds = torch.stack(new_input_embeds_align, dim=0)\n\n if labels is not None:\n new_labels_align = []\n _new_labels = new_labels\n for cur_new_label in new_labels:\n cur_new_label = torch.cat(\n (\n cur_new_label,\n torch.full(\n (max_len - cur_new_label.shape[0],),\n IGNORE_INDEX,\n dtype=cur_new_label.dtype,\n device=cur_new_label.device,\n ),\n ),\n dim=0,\n )\n new_labels_align.append(cur_new_label)\n new_labels = torch.stack(new_labels_align, dim=0)\n\n if attention_mask is not None:\n new_attention_mask = []\n for cur_attention_mask, cur_new_labels, cur_new_labels_align in zip(\n attention_mask, _new_labels, new_labels\n ):\n new_attn_mask_pad_left = torch.full(\n (cur_new_labels.shape[0] - labels.shape[1],),\n True,\n dtype=attention_mask.dtype,\n device=attention_mask.device,\n )\n new_attn_mask_pad_right = torch.full(\n (cur_new_labels_align.shape[0] - cur_new_labels.shape[0],),\n False,\n dtype=attention_mask.dtype,\n device=attention_mask.device,\n )\n cur_new_attention_mask = torch.cat(\n (\n new_attn_mask_pad_left,\n cur_attention_mask,\n new_attn_mask_pad_right,\n ),\n dim=0,\n )\n new_attention_mask.append(cur_new_attention_mask)\n attention_mask = torch.stack(new_attention_mask, dim=0)\n assert attention_mask.shape == new_labels.shape\n else:\n new_input_embeds = torch.stack(new_input_embeds, dim=0)\n if labels is not None:\n new_labels = torch.stack(new_labels, dim=0)\n\n if attention_mask is not None:\n new_attn_mask_pad_left = torch.full(\n (\n attention_mask.shape[0],\n new_input_embeds.shape[1] - input_ids.shape[1],\n ),\n True,\n dtype=attention_mask.dtype,\n device=attention_mask.device,\n )\n attention_mask = torch.cat(\n (new_attn_mask_pad_left, attention_mask), dim=1\n )\n assert attention_mask.shape == new_input_embeds.shape[:2]\n\n return None, attention_mask, past_key_values, new_input_embeds, new_labels\n\n # def initialize_vision_tokenizer(self, model_args, tokenizer):\n def initialize_vision_tokenizer(self, model_args, num_new_tokens):\n # if model_args.mm_use_im_patch_token:\n # tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)\n # self.resize_token_embeddings(len(tokenizer))\n\n if model_args.mm_use_im_start_end:\n # num_new_tokens = tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)\n # self.resize_token_embeddings(len(tokenizer))\n\n # if num_new_tokens > 0:\n # input_embeddings = self.get_input_embeddings().weight.data\n # output_embeddings = self.get_output_embeddings().weight.data\n\n # input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(\n # dim=0, keepdim=True)\n # output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(\n # dim=0, keepdim=True)\n\n # input_embeddings[-num_new_tokens:] = input_embeddings_avg\n # output_embeddings[-num_new_tokens:] = output_embeddings_avg\n\n if model_args.tune_mm_mlp_adapter:\n for p in self.get_input_embeddings().parameters():\n p.requires_grad = True\n for p in self.get_output_embeddings().parameters():\n p.requires_grad = False\n\n if model_args.pretrain_mm_mlp_adapter:\n mm_projector_weights = torch.load(\n model_args.pretrain_mm_mlp_adapter, map_location=\"cpu\"\n )\n embed_tokens_weight = mm_projector_weights[\"model.embed_tokens.weight\"]\n assert num_new_tokens == 2\n if input_embeddings.shape == embed_tokens_weight.shape:\n input_embeddings[-num_new_tokens:] = embed_tokens_weight[\n -num_new_tokens:\n ]\n elif embed_tokens_weight.shape[0] == num_new_tokens:\n input_embeddings[-num_new_tokens:] = embed_tokens_weight\n else:\n raise ValueError(\n f\"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}.\"\n )\n elif model_args.mm_use_im_patch_token:\n if model_args.tune_mm_mlp_adapter:\n for p in self.get_input_embeddings().parameters():\n p.requires_grad = False\n for p in self.get_output_embeddings().parameters():\n p.requires_grad = False\n" }, { "path": "model/llava/model/make_delta.py", "content": "\"\"\"\nUsage:\npython3 -m llava.model.make_delta --base ~/model_weights/llama-7b --target ~/model_weights/llava-7b --delta ~/model_weights/llava-7b-delta --hub-repo-id liuhaotian/llava-7b-delta\n\"\"\"\nimport argparse\n\nimport torch\nfrom llava.model.utils import auto_upgrade\nfrom tqdm import tqdm\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\n\ndef make_delta(base_model_path, target_model_path, delta_path, hub_repo_id):\n print(\"Loading base model\")\n base = AutoModelForCausalLM.from_pretrained(\n base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True\n )\n\n print(\"Loading target model\")\n auto_upgrade(target_model_path)\n target = AutoModelForCausalLM.from_pretrained(\n target_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True\n )\n\n print(\"Calculating delta\")\n for name, param in tqdm(target.state_dict().items(), desc=\"Calculating delta\"):\n if name not in base.state_dict():\n assert name in [\n \"model.mm_projector.weight\",\n \"model.mm_projector.bias\",\n ], f\"{name} not in base model\"\n continue\n if param.data.shape == base.state_dict()[name].shape:\n param.data -= base.state_dict()[name]\n else:\n assert name in [\n \"model.embed_tokens.weight\",\n \"lm_head.weight\",\n ], f\"{name} dimension mismatch: {param.data.shape} vs {base.state_dict()[name].shape}\"\n bparam = base.state_dict()[name]\n param.data[: bparam.shape[0], : bparam.shape[1]] -= bparam\n\n print(\"Saving delta\")\n if hub_repo_id:\n kwargs = {\"push_to_hub\": True, \"repo_id\": hub_repo_id}\n else:\n kwargs = {}\n target.save_pretrained(delta_path, **kwargs)\n target_tokenizer = AutoTokenizer.from_pretrained(target_model_path)\n target_tokenizer.save_pretrained(delta_path, **kwargs)\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--base-model-path\", type=str, required=True)\n parser.add_argument(\"--target-model-path\", type=str, required=True)\n parser.add_argument(\"--delta-path\", type=str, required=True)\n parser.add_argument(\"--hub-repo-id\", type=str, default=None)\n args = parser.parse_args()\n\n make_delta(\n args.base_model_path, args.target_model_path, args.delta_path, args.hub_repo_id\n )\n" }, { "path": "model/llava/model/multimodal_encoder/builder.py", "content": "from .clip_encoder import CLIPVisionTower\n\n\ndef build_vision_tower(vision_tower_cfg, **kwargs):\n vision_tower = getattr(\n vision_tower_cfg,\n \"mm_vision_tower\",\n getattr(vision_tower_cfg, \"vision_tower\", None),\n )\n if (\n vision_tower.startswith(\"openai\")\n or vision_tower.startswith(\"laion\")\n or \"clip\" in vision_tower\n ):\n return CLIPVisionTower(vision_tower, args=vision_tower_cfg, **kwargs)\n\n raise ValueError(f\"Unknown vision tower: {vision_tower}\")\n" }, { "path": "model/llava/model/multimodal_encoder/clip_encoder.py", "content": "import torch\nimport torch.nn as nn\nfrom transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel\n\n\nclass CLIPVisionTower(nn.Module):\n def __init__(self, vision_tower, args, delay_load=False):\n super().__init__()\n\n self.is_loaded = False\n\n self.vision_tower_name = vision_tower\n self.select_layer = args.mm_vision_select_layer\n self.select_feature = getattr(args, \"mm_vision_select_feature\", \"patch\")\n\n if not delay_load:\n self.load_model()\n else:\n self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name)\n\n def load_model(self):\n self.image_processor = CLIPImageProcessor.from_pretrained(\n self.vision_tower_name\n )\n self.vision_tower = CLIPVisionModel.from_pretrained(\n self.vision_tower_name, low_cpu_mem_usage=True\n )\n self.vision_tower.requires_grad_(False)\n self.is_loaded = True\n\n def feature_select(self, image_forward_outs):\n image_features = image_forward_outs.hidden_states[self.select_layer]\n if self.select_feature == \"patch\":\n image_features = image_features[:, 1:]\n elif self.select_feature == \"cls_patch\":\n image_features = image_features\n else:\n raise ValueError(f\"Unexpected select feature: {self.select_feature}\")\n return image_features\n\n @torch.no_grad()\n def forward(self, images):\n if type(images) is list:\n image_features = []\n for image in images:\n image_forward_out = self.vision_tower(\n image.to(device=self.device, dtype=self.dtype).unsqueeze(0),\n output_hidden_states=True,\n )\n image_feature = self.feature_select(image_forward_out).to(image.dtype)\n image_features.append(image_feature)\n else:\n image_forward_outs = self.vision_tower(\n images.to(device=self.device, dtype=self.dtype),\n output_hidden_states=True,\n )\n image_features = self.feature_select(image_forward_outs).to(images.dtype)\n\n torch.cuda.empty_cache()\n return image_features\n\n @property\n def dummy_feature(self):\n return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)\n\n @property\n def dtype(self):\n return self.vision_tower.dtype\n\n @property\n def device(self):\n return self.vision_tower.device\n\n @property\n def config(self):\n if self.is_loaded:\n return self.vision_tower.config\n else:\n return self.cfg_only\n\n @property\n def hidden_size(self):\n return self.config.hidden_size\n\n @property\n def num_patches(self):\n return (self.config.image_size // self.config.patch_size) ** 2\n" }, { "path": "model/llava/model/utils.py", "content": "from transformers import AutoConfig\n\n\ndef auto_upgrade(config):\n cfg = AutoConfig.from_pretrained(config)\n if \"llava\" in config and \"llava\" not in cfg.model_type:\n assert cfg.model_type == \"llama\"\n print(\n \"You are using newer LLaVA code base, while the checkpoint of v0 is from older code base.\"\n )\n print(\n \"You must upgrade the checkpoint to the new code base (this can be done automatically).\"\n )\n confirm = input(\"Please confirm that you want to upgrade the checkpoint. [Y/N]\")\n if confirm.lower() in [\"y\", \"yes\"]:\n print(\"Upgrading checkpoint...\")\n assert len(cfg.architectures) == 1\n setattr(cfg.__class__, \"model_type\", \"llava\")\n cfg.architectures[0] = \"LlavaLlamaForCausalLM\"\n cfg.save_pretrained(config)\n print(\"Checkpoint upgraded.\")\n else:\n print(\"Checkpoint upgrade aborted.\")\n exit(1)\n" }, { "path": "model/llava/train/llama_flash_attn_monkey_patch.py", "content": "import logging\nfrom typing import List, Optional, Tuple\n\nimport torch\nimport transformers\nfrom einops import rearrange\nfrom torch import nn\nfrom transformers.models.llama.modeling_llama import apply_rotary_pos_emb\n\ntry:\n from flash_attn.flash_attn_interface import \\\n flash_attn_unpadded_qkvpacked_func\nexcept ImportError:\n from flash_attn.flash_attn_interface import (\n flash_attn_varlen_qkvpacked_func as flash_attn_unpadded_qkvpacked_func,\n )\n\nfrom flash_attn.bert_padding import pad_input, unpad_input\n\n\ndef forward(\n self,\n hidden_states: torch.Tensor,\n attention_mask: Optional[torch.Tensor] = None,\n position_ids: Optional[torch.Tensor] = None,\n past_key_value: Optional[Tuple[torch.Tensor]] = None,\n output_attentions: bool = False,\n use_cache: bool = False,\n) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:\n \"\"\"Input shape: Batch x Time x Channel\n\n attention_mask: [bsz, q_len]\n \"\"\"\n bsz, q_len, _ = hidden_states.size()\n\n query_states = (\n self.q_proj(hidden_states)\n .view(bsz, q_len, self.num_heads, self.head_dim)\n .transpose(1, 2)\n )\n key_states = (\n self.k_proj(hidden_states)\n .view(bsz, q_len, self.num_heads, self.head_dim)\n .transpose(1, 2)\n )\n value_states = (\n self.v_proj(hidden_states)\n .view(bsz, q_len, self.num_heads, self.head_dim)\n .transpose(1, 2)\n )\n # [bsz, q_len, nh, hd]\n # [bsz, nh, q_len, hd]\n\n kv_seq_len = key_states.shape[-2]\n assert past_key_value is None, \"past_key_value is not supported\"\n\n cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)\n query_states, key_states = apply_rotary_pos_emb(\n query_states, key_states, cos, sin, position_ids\n )\n # [bsz, nh, t, hd]\n assert not output_attentions, \"output_attentions is not supported\"\n assert not use_cache, \"use_cache is not supported\"\n\n # Flash attention codes from\n # https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attention.py\n\n # transform the data into the format required by flash attention\n qkv = torch.stack(\n [query_states, key_states, value_states], dim=2\n ) # [bsz, nh, 3, q_len, hd]\n qkv = qkv.transpose(1, 3) # [bsz, q_len, 3, nh, hd]\n # We have disabled _prepare_decoder_attention_mask in LlamaModel\n # the attention_mask should be the same as the key_padding_mask\n key_padding_mask = attention_mask\n\n if key_padding_mask is None:\n qkv = rearrange(qkv, \"b s ... -> (b s) ...\")\n max_s = q_len\n cu_q_lens = torch.arange(\n 0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32, device=qkv.device\n )\n output = flash_attn_unpadded_qkvpacked_func(\n qkv, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True\n )\n output = rearrange(output, \"(b s) ... -> b s ...\", b=bsz)\n else:\n nheads = qkv.shape[-2]\n x = rearrange(qkv, \"b s three h d -> b s (three h d)\")\n x_unpad, indices, cu_q_lens, max_s = unpad_input(x, key_padding_mask)\n x_unpad = rearrange(\n x_unpad, \"nnz (three h d) -> nnz three h d\", three=3, h=nheads\n )\n output_unpad = flash_attn_unpadded_qkvpacked_func(\n x_unpad, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True\n )\n output = rearrange(\n pad_input(\n rearrange(output_unpad, \"nnz h d -> nnz (h d)\"), indices, bsz, q_len\n ),\n \"b s (h d) -> b s h d\",\n h=nheads,\n )\n return self.o_proj(rearrange(output, \"b s h d -> b s (h d)\")), None, None\n\n\n# Disable the transformation of the attention mask in LlamaModel as the flash attention\n# requires the attention mask to be the same as the key_padding_mask\ndef _prepare_decoder_attention_mask(\n self, attention_mask, input_shape, inputs_embeds, past_key_values_length\n):\n # [bsz, seq_len]\n return attention_mask\n\n\ndef replace_llama_attn_with_flash_attn():\n cuda_major, cuda_minor = torch.cuda.get_device_capability()\n if cuda_major < 8:\n logging.warning(\n \"Flash attention is only supported on A100 or H100 GPU during training due to head dim > 64 backward.\"\n \"ref: https://github.com/HazyResearch/flash-attention/issues/190#issuecomment-1523359593\"\n )\n transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = (\n _prepare_decoder_attention_mask\n )\n transformers.models.llama.modeling_llama.LlamaAttention.forward = forward\n" }, { "path": "model/llava/train/llava_trainer.py", "content": "import os\nfrom typing import Optional\n\nimport torch\nfrom transformers import Trainer\n\n\ndef maybe_zero_3(param, ignore_status=False, name=None):\n from deepspeed import zero\n from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus\n\n if hasattr(param, \"ds_id\"):\n if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:\n if not ignore_status:\n print(name, \"no ignore status\")\n with zero.GatheredParameters([param]):\n param = param.data.detach().cpu().clone()\n else:\n param = param.detach().cpu().clone()\n return param\n\n\ndef get_mm_adapter_state_maybe_zero_3(named_params, keys_to_match):\n to_return = {\n k: t\n for k, t in named_params\n if any(key_match in k for key_match in keys_to_match)\n }\n to_return = {\n k: maybe_zero_3(v, ignore_status=True, name=k).cpu()\n for k, v in to_return.items()\n }\n return to_return\n\n\nclass LLaVATrainer(Trainer):\n def _save_checkpoint(self, model, trial, metrics=None):\n if getattr(self.args, \"tune_mm_mlp_adapter\", False):\n from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR\n\n checkpoint_folder = f\"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}\"\n\n run_dir = self._get_output_dir(trial=trial)\n output_dir = os.path.join(run_dir, checkpoint_folder)\n\n # Only save Adapter\n keys_to_match = [\"mm_projector\"]\n if getattr(self.args, \"use_im_start_end\", False):\n keys_to_match.extend([\"embed_tokens\", \"embed_in\"])\n\n weight_to_save = get_mm_adapter_state_maybe_zero_3(\n self.model.named_parameters(), keys_to_match\n )\n\n if self.args.local_rank == 0 or self.args.local_rank == -1:\n self.model.config.save_pretrained(output_dir)\n torch.save(\n weight_to_save, os.path.join(output_dir, f\"mm_projector.bin\")\n )\n else:\n super(LLaVATrainer, self)._save_checkpoint(model, trial, metrics)\n\n def _save(self, output_dir: Optional[str] = None, state_dict=None):\n if getattr(self.args, \"tune_mm_mlp_adapter\", False):\n pass\n else:\n super(LLaVATrainer, self)._save(output_dir, state_dict)\n" }, { "path": "model/llava/train/train.py", "content": "# Adopted from https://github.com/lm-sys/FastChat. Below is the original copyright:\n# Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:\n# Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport copy\nimport json\nimport logging\nimport os\nimport pathlib\nfrom dataclasses import dataclass, field\nfrom typing import Dict, List, Optional, Sequence\n\nimport torch\nimport transformers\nfrom llava import conversation as conversation_lib\nfrom llava.constants import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,\n DEFAULT_IMAGE_TOKEN, IGNORE_INDEX,\n IMAGE_TOKEN_INDEX)\nfrom llava.mm_utils import tokenizer_image_token\nfrom llava.model import *\nfrom llava.train.llava_trainer import LLaVATrainer\nfrom PIL import Image\nfrom torch.utils.data import Dataset\n\nlocal_rank = None\n\n\ndef rank0_print(*args):\n if local_rank == 0:\n print(*args)\n\n\n@dataclass\nclass ModelArguments:\n model_name_or_path: Optional[str] = field(default=\"facebook/opt-125m\")\n version: Optional[str] = field(default=\"v0\")\n freeze_backbone: bool = field(default=False)\n tune_mm_mlp_adapter: bool = field(default=False)\n vision_tower: Optional[str] = field(default=None)\n mm_vision_select_layer: Optional[int] = field(\n default=-1\n ) # default to the last layer\n pretrain_mm_mlp_adapter: Optional[str] = field(default=None)\n mm_use_im_start_end: bool = field(default=False)\n mm_use_im_patch_token: bool = field(default=True)\n mm_vision_select_feature: Optional[str] = field(default=\"patch\")\n\n\n@dataclass\nclass DataArguments:\n data_path: str = field(\n default=None, metadata={\"help\": \"Path to the training data.\"}\n )\n lazy_preprocess: bool = False\n is_multimodal: bool = False\n image_folder: Optional[str] = field(default=None)\n image_aspect_ratio: str = \"square\"\n image_grid_pinpoints: Optional[str] = field(default=None)\n\n\n@dataclass\nclass TrainingArguments(transformers.TrainingArguments):\n cache_dir: Optional[str] = field(default=None)\n optim: str = field(default=\"adamw_torch\")\n remove_unused_columns: bool = field(default=False)\n freeze_mm_mlp_adapter: bool = field(default=False)\n mpt_attn_impl: Optional[str] = field(default=\"triton\")\n model_max_length: int = field(\n default=512,\n metadata={\n \"help\": \"Maximum sequence length. Sequences will be right padded (and possibly truncated).\"\n },\n )\n double_quant: bool = field(\n default=True,\n metadata={\n \"help\": \"Compress the quantization statistics through double quantization.\"\n },\n )\n quant_type: str = field(\n default=\"nf4\",\n metadata={\n \"help\": \"Quantization data type to use. Should be one of `fp4` or `nf4`.\"\n },\n )\n bits: int = field(default=16, metadata={\"help\": \"How many bits to use.\"})\n lora_enable: bool = False\n lora_r: int = 64\n lora_alpha: int = 16\n lora_dropout: float = 0.05\n lora_weight_path: str = \"\"\n lora_bias: str = \"none\"\n\n\ndef maybe_zero_3(param, ignore_status=False, name=None):\n from deepspeed import zero\n from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus\n\n if hasattr(param, \"ds_id\"):\n if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:\n if not ignore_status:\n logging.warning(\n f\"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}\"\n )\n with zero.GatheredParameters([param]):\n param = param.data.detach().cpu().clone()\n else:\n param = param.detach().cpu().clone()\n return param\n\n\n# Borrowed from peft.utils.get_peft_model_state_dict\ndef get_peft_state_maybe_zero_3(named_params, bias):\n if bias == \"none\":\n to_return = {k: t for k, t in named_params if \"lora_\" in k}\n elif bias == \"all\":\n to_return = {k: t for k, t in named_params if \"lora_\" in k or \"bias\" in k}\n elif bias == \"lora_only\":\n to_return = {}\n maybe_lora_bias = {}\n lora_bias_names = set()\n for k, t in named_params:\n if \"lora_\" in k:\n to_return[k] = t\n bias_name = k.split(\"lora_\")[0] + \"bias\"\n lora_bias_names.add(bias_name)\n elif \"bias\" in k:\n maybe_lora_bias[k] = t\n for k, t in maybe_lora_bias:\n if bias_name in lora_bias_names:\n to_return[bias_name] = t\n else:\n raise NotImplementedError\n to_return = {k: maybe_zero_3(v, name=k) for k, v in to_return.items()}\n return to_return\n\n\ndef get_peft_state_non_lora_maybe_zero_3(named_params, require_grad_only=True):\n to_return = {k: t for k, t in named_params if \"lora_\" not in k}\n if require_grad_only:\n to_return = {k: t for k, t in to_return.items() if t.requires_grad}\n to_return = {\n k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()\n }\n return to_return\n\n\ndef get_mm_adapter_state_maybe_zero_3(named_params, keys_to_match):\n to_return = {\n k: t\n for k, t in named_params\n if any(key_match in k for key_match in keys_to_match)\n }\n to_return = {\n k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()\n }\n return to_return\n\n\ndef find_all_linear_names(model):\n cls = torch.nn.Linear\n lora_module_names = set()\n for name, module in model.named_modules():\n if isinstance(module, cls):\n names = name.split(\".\")\n lora_module_names.add(names[0] if len(names) == 1 else names[-1])\n\n if \"lm_head\" in lora_module_names: # needed for 16-bit\n lora_module_names.remove(\"lm_head\")\n return list(lora_module_names)\n\n\ndef safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):\n \"\"\"Collects the state dict and dump to disk.\"\"\"\n\n if getattr(trainer.args, \"tune_mm_mlp_adapter\", False):\n # Only save Adapter\n keys_to_match = [\"mm_projector\"]\n if getattr(trainer.args, \"use_im_start_end\", False):\n keys_to_match.extend([\"embed_tokens\", \"embed_in\"])\n\n weight_to_save = get_mm_adapter_state_maybe_zero_3(\n trainer.model.named_parameters(), keys_to_match\n )\n trainer.model.config.save_pretrained(output_dir)\n\n current_folder = output_dir.split(\"/\")[-1]\n parent_folder = os.path.dirname(output_dir)\n if trainer.args.local_rank == 0 or trainer.args.local_rank == -1:\n if current_folder.startswith(\"checkpoint-\"):\n mm_projector_folder = os.path.join(parent_folder, \"mm_projector\")\n os.makedirs(mm_projector_folder, exist_ok=True)\n torch.save(\n weight_to_save,\n os.path.join(mm_projector_folder, f\"{current_folder}.bin\"),\n )\n else:\n torch.save(\n weight_to_save, os.path.join(output_dir, f\"mm_projector.bin\")\n )\n return\n\n if trainer.deepspeed:\n torch.cuda.synchronize()\n trainer.save_model(output_dir)\n return\n\n state_dict = trainer.model.state_dict()\n if trainer.args.should_save:\n cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}\n del state_dict\n trainer._save(output_dir, state_dict=cpu_state_dict) # noqa\n\n\ndef smart_tokenizer_and_embedding_resize(\n special_tokens_dict: Dict,\n tokenizer: transformers.PreTrainedTokenizer,\n model: transformers.PreTrainedModel,\n):\n \"\"\"Resize tokenizer and embedding.\n\n Note: This is the unoptimized version that may make your embedding size not be divisible by 64.\n \"\"\"\n num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)\n model.resize_token_embeddings(len(tokenizer))\n\n if num_new_tokens > 0:\n input_embeddings = model.get_input_embeddings().weight.data\n output_embeddings = model.get_output_embeddings().weight.data\n\n input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(\n dim=0, keepdim=True\n )\n output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(\n dim=0, keepdim=True\n )\n\n input_embeddings[-num_new_tokens:] = input_embeddings_avg\n output_embeddings[-num_new_tokens:] = output_embeddings_avg\n\n\ndef _tokenize_fn(\n strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer\n) -> Dict:\n \"\"\"Tokenize a list of strings.\"\"\"\n tokenized_list = [\n tokenizer(\n text,\n return_tensors=\"pt\",\n padding=\"longest\",\n max_length=tokenizer.model_max_length,\n truncation=True,\n )\n for text in strings\n ]\n input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]\n input_ids_lens = labels_lens = [\n tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item()\n for tokenized in tokenized_list\n ]\n return dict(\n input_ids=input_ids,\n labels=labels,\n input_ids_lens=input_ids_lens,\n labels_lens=labels_lens,\n )\n\n\ndef _mask_targets(target, tokenized_lens, speakers):\n # cur_idx = 0\n cur_idx = tokenized_lens[0]\n tokenized_lens = tokenized_lens[1:]\n target[:cur_idx] = IGNORE_INDEX\n for tokenized_len, speaker in zip(tokenized_lens, speakers):\n if speaker == \"human\":\n target[cur_idx + 2 : cur_idx + tokenized_len] = IGNORE_INDEX\n cur_idx += tokenized_len\n\n\ndef _add_speaker_and_signal(header, source, get_conversation=True):\n \"\"\"Add speaker and start/end signal on each round.\"\"\"\n BEGIN_SIGNAL = \"### \"\n END_SIGNAL = \"\\n\"\n conversation = header\n for sentence in source:\n from_str = sentence[\"from\"]\n if from_str.lower() == \"human\":\n from_str = conversation_lib.default_conversation.roles[0]\n elif from_str.lower() == \"gpt\":\n from_str = conversation_lib.default_conversation.roles[1]\n else:\n from_str = \"unknown\"\n sentence[\"value\"] = (\n BEGIN_SIGNAL + from_str + \": \" + sentence[\"value\"] + END_SIGNAL\n )\n if get_conversation:\n conversation += sentence[\"value\"]\n conversation += BEGIN_SIGNAL\n return conversation\n\n\ndef preprocess_multimodal(sources: Sequence[str], data_args: DataArguments) -> Dict:\n is_multimodal = data_args.is_multimodal\n if not is_multimodal:\n return sources\n\n for source in sources:\n for sentence in source:\n if DEFAULT_IMAGE_TOKEN in sentence[\"value\"]:\n sentence[\"value\"] = (\n sentence[\"value\"].replace(DEFAULT_IMAGE_TOKEN, \"\").strip()\n )\n sentence[\"value\"] = DEFAULT_IMAGE_TOKEN + \"\\n\" + sentence[\"value\"]\n sentence[\"value\"] = sentence[\"value\"].strip()\n if \"mmtag\" in conversation_lib.default_conversation.version:\n sentence[\"value\"] = sentence[\"value\"].replace(\n DEFAULT_IMAGE_TOKEN,\n \"\" + DEFAULT_IMAGE_TOKEN + \"\",\n )\n replace_token = DEFAULT_IMAGE_TOKEN\n if data_args.mm_use_im_start_end:\n replace_token = (\n DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN\n )\n sentence[\"value\"] = sentence[\"value\"].replace(\n DEFAULT_IMAGE_TOKEN, replace_token\n )\n\n return sources\n\n\ndef preprocess_llama_2(\n sources, tokenizer: transformers.PreTrainedTokenizer, has_image: bool = False\n) -> Dict:\n conv = conversation_lib.default_conversation.copy()\n roles = {\"human\": conv.roles[0], \"gpt\": conv.roles[1]}\n\n # Apply prompt templates\n conversations = []\n for i, source in enumerate(sources):\n if roles[source[0][\"from\"]] != conv.roles[0]:\n # Skip the first one if it is not from human\n source = source[1:]\n\n conv.messages = []\n for j, sentence in enumerate(source):\n role = roles[sentence[\"from\"]]\n assert role == conv.roles[j % 2], f\"{i}\"\n conv.append_message(role, sentence[\"value\"])\n conversations.append(conv.get_prompt())\n\n # Tokenize conversations\n\n if has_image:\n input_ids = torch.stack(\n [\n tokenizer_image_token(prompt, tokenizer, return_tensors=\"pt\")\n for prompt in conversations\n ],\n dim=0,\n )\n else:\n input_ids = tokenizer(\n conversations,\n return_tensors=\"pt\",\n padding=\"longest\",\n max_length=tokenizer.model_max_length,\n truncation=True,\n ).input_ids\n\n targets = input_ids.clone()\n\n assert conv.sep_style == conversation_lib.SeparatorStyle.LLAMA_2\n\n # Mask targets\n sep = \"[/INST] \"\n for conversation, target in zip(conversations, targets):\n total_len = int(target.ne(tokenizer.pad_token_id).sum())\n\n rounds = conversation.split(conv.sep2)\n cur_len = 1\n target[:cur_len] = IGNORE_INDEX\n for i, rou in enumerate(rounds):\n if rou == \"\":\n break\n\n parts = rou.split(sep)\n if len(parts) != 2:\n break\n parts[0] += sep\n\n if has_image:\n round_len = len(tokenizer_image_token(rou, tokenizer))\n instruction_len = len(tokenizer_image_token(parts[0], tokenizer)) - 2\n else:\n round_len = len(tokenizer(rou).input_ids)\n instruction_len = len(tokenizer(parts[0]).input_ids) - 2\n\n target[cur_len : cur_len + instruction_len] = IGNORE_INDEX\n\n cur_len += round_len\n target[cur_len:] = IGNORE_INDEX\n\n if cur_len < tokenizer.model_max_length:\n if cur_len != total_len:\n target[:] = IGNORE_INDEX\n print(\n f\"WARNING: tokenization mismatch: {cur_len} vs. {total_len}.\"\n f\" (ignored)\"\n )\n\n return dict(\n input_ids=input_ids,\n labels=targets,\n )\n\n\ndef preprocess_v1(\n sources, tokenizer: transformers.PreTrainedTokenizer, has_image: bool = False\n) -> Dict:\n conv = conversation_lib.default_conversation.copy()\n roles = {\"human\": conv.roles[0], \"gpt\": conv.roles[1]}\n\n # Apply prompt templates\n conversations = []\n for i, source in enumerate(sources):\n if roles[source[0][\"from\"]] != conv.roles[0]:\n # Skip the first one if it is not from human\n source = source[1:]\n\n conv.messages = []\n for j, sentence in enumerate(source):\n role = roles[sentence[\"from\"]]\n assert role == conv.roles[j % 2], f\"{i}\"\n conv.append_message(role, sentence[\"value\"])\n conversations.append(conv.get_prompt())\n\n # Tokenize conversations\n\n if has_image:\n input_ids = torch.stack(\n [\n tokenizer_image_token(prompt, tokenizer, return_tensors=\"pt\")\n for prompt in conversations\n ],\n dim=0,\n )\n else:\n input_ids = tokenizer(\n conversations,\n return_tensors=\"pt\",\n padding=\"longest\",\n max_length=tokenizer.model_max_length,\n truncation=True,\n ).input_ids\n\n targets = input_ids.clone()\n\n assert conv.sep_style == conversation_lib.SeparatorStyle.TWO\n\n # Mask targets\n sep = conv.sep + conv.roles[1] + \": \"\n for conversation, target in zip(conversations, targets):\n total_len = int(target.ne(tokenizer.pad_token_id).sum())\n\n rounds = conversation.split(conv.sep2)\n cur_len = 1\n target[:cur_len] = IGNORE_INDEX\n for i, rou in enumerate(rounds):\n if rou == \"\":\n break\n\n parts = rou.split(sep)\n if len(parts) != 2:\n break\n parts[0] += sep\n\n if has_image:\n round_len = len(tokenizer_image_token(rou, tokenizer))\n instruction_len = len(tokenizer_image_token(parts[0], tokenizer)) - 2\n else:\n round_len = len(tokenizer(rou).input_ids)\n instruction_len = len(tokenizer(parts[0]).input_ids) - 2\n\n target[cur_len : cur_len + instruction_len] = IGNORE_INDEX\n\n cur_len += round_len\n target[cur_len:] = IGNORE_INDEX\n\n if cur_len < tokenizer.model_max_length:\n if cur_len != total_len:\n target[:] = IGNORE_INDEX\n print(\n f\"WARNING: tokenization mismatch: {cur_len} vs. {total_len}.\"\n f\" (ignored)\"\n )\n\n return dict(\n input_ids=input_ids,\n labels=targets,\n )\n\n\ndef preprocess_mpt(\n sources,\n tokenizer: transformers.PreTrainedTokenizer,\n) -> Dict:\n conv = conversation_lib.default_conversation.copy()\n roles = {\"human\": conv.roles[0], \"gpt\": conv.roles[1]}\n\n # Apply prompt templates\n conversations = []\n for i, source in enumerate(sources):\n if roles[source[0][\"from\"]] != conv.roles[0]:\n # Skip the first one if it is not from human\n source = source[1:]\n\n conv.messages = []\n for j, sentence in enumerate(source):\n role = roles[sentence[\"from\"]]\n assert role == conv.roles[j % 2], f\"{i}\"\n conv.append_message(role, sentence[\"value\"])\n conversations.append(conv.get_prompt())\n\n # Tokenize conversations\n input_ids = torch.stack(\n [\n tokenizer_image_token(prompt, tokenizer, return_tensors=\"pt\")\n for prompt in conversations\n ],\n dim=0,\n )\n targets = input_ids.clone()\n assert conv.sep_style == conversation_lib.SeparatorStyle.MPT\n\n # Mask targets\n sep = conv.sep + conv.roles[1]\n for conversation, target in zip(conversations, targets):\n total_len = int(target.ne(tokenizer.pad_token_id).sum())\n\n rounds = conversation.split(conv.sep)\n re_rounds = [conv.sep.join(rounds[:3])] # system + user + gpt\n for conv_idx in range(3, len(rounds), 2):\n re_rounds.append(\n conv.sep.join(rounds[conv_idx : conv_idx + 2])\n ) # user + gpt\n cur_len = 0\n target[:cur_len] = IGNORE_INDEX\n for i, rou in enumerate(re_rounds):\n if rou == \"\":\n break\n\n parts = rou.split(sep)\n if len(parts) != 2:\n break\n parts[0] += sep\n round_len = len(tokenizer_image_token(rou, tokenizer)) + len(\n tokenizer_image_token(conv.sep, tokenizer)\n )\n instruction_len = len(tokenizer_image_token(parts[0], tokenizer))\n target[cur_len : cur_len + instruction_len] = IGNORE_INDEX\n\n cur_len += round_len\n target[cur_len:] = IGNORE_INDEX\n\n if cur_len < tokenizer.model_max_length:\n if cur_len != total_len:\n target[:] = IGNORE_INDEX\n print(\n f\"WARNING: tokenization mismatch: {cur_len} vs. {total_len}.\"\n f\" (ignored)\"\n )\n\n return dict(\n input_ids=input_ids,\n labels=targets,\n )\n\n\ndef preprocess_plain(\n sources: Sequence[str],\n tokenizer: transformers.PreTrainedTokenizer,\n) -> Dict:\n # add end signal and concatenate together\n conversations = []\n for source in sources:\n assert len(source) == 2\n assert DEFAULT_IMAGE_TOKEN in source[0][\"value\"]\n source[0][\"value\"] = DEFAULT_IMAGE_TOKEN\n conversation = (\n source[0][\"value\"]\n + source[1][\"value\"]\n + conversation_lib.default_conversation.sep\n )\n conversations.append(conversation)\n # tokenize conversations\n input_ids = [\n tokenizer_image_token(prompt, tokenizer, return_tensors=\"pt\")\n for prompt in conversations\n ]\n targets = copy.deepcopy(input_ids)\n for target, source in zip(targets, sources):\n tokenized_len = len(tokenizer_image_token(source[0][\"value\"], tokenizer))\n target[:tokenized_len] = IGNORE_INDEX\n\n return dict(input_ids=input_ids, labels=targets)\n\n\ndef preprocess(\n sources: Sequence[str],\n tokenizer: transformers.PreTrainedTokenizer,\n has_image: bool = False,\n) -> Dict:\n \"\"\"\n Given a list of sources, each is a conversation list. This transform:\n 1. Add signal '### ' at the beginning each sentence, with end signal '\\n';\n 2. Concatenate conversations together;\n 3. Tokenize the concatenated conversation;\n 4. Make a deepcopy as the target. Mask human words with IGNORE_INDEX.\n \"\"\"\n if (\n conversation_lib.default_conversation.sep_style\n == conversation_lib.SeparatorStyle.PLAIN\n ):\n return preprocess_plain(sources, tokenizer)\n if (\n conversation_lib.default_conversation.sep_style\n == conversation_lib.SeparatorStyle.LLAMA_2\n ):\n return preprocess_llama_2(sources, tokenizer, has_image=has_image)\n if conversation_lib.default_conversation.version.startswith(\"v1\"):\n return preprocess_v1(sources, tokenizer, has_image=has_image)\n if conversation_lib.default_conversation.version == \"mpt\":\n return preprocess_mpt(sources, tokenizer)\n # add end signal and concatenate together\n conversations = []\n for source in sources:\n header = f\"{conversation_lib.default_conversation.system}\\n\\n\"\n conversation = _add_speaker_and_signal(header, source)\n conversations.append(conversation)\n\n # tokenize conversations\n def get_tokenize_len(prompts):\n return [len(tokenizer_image_token(prompt, tokenizer)) for prompt in prompts]\n\n if has_image:\n input_ids = [\n tokenizer_image_token(prompt, tokenizer, return_tensors=\"pt\")\n for prompt in conversations\n ]\n else:\n conversations_tokenized = _tokenize_fn(conversations, tokenizer)\n input_ids = conversations_tokenized[\"input_ids\"]\n\n targets = copy.deepcopy(input_ids)\n for target, source in zip(targets, sources):\n if has_image:\n tokenized_lens = get_tokenize_len([header] + [s[\"value\"] for s in source])\n else:\n tokenized_lens = _tokenize_fn(\n [header] + [s[\"value\"] for s in source], tokenizer\n )[\"input_ids_lens\"]\n speakers = [sentence[\"from\"] for sentence in source]\n _mask_targets(target, tokenized_lens, speakers)\n\n return dict(input_ids=input_ids, labels=targets)\n\n\nclass LazySupervisedDataset(Dataset):\n \"\"\"Dataset for supervised fine-tuning.\"\"\"\n\n def __init__(\n self,\n data_path: str,\n tokenizer: transformers.PreTrainedTokenizer,\n data_args: DataArguments,\n ):\n super(LazySupervisedDataset, self).__init__()\n list_data_dict = json.load(open(data_path, \"r\"))\n\n rank0_print(\"Formatting inputs...Skip in lazy mode\")\n self.tokenizer = tokenizer\n self.list_data_dict = list_data_dict\n self.data_args = data_args\n\n def __len__(self):\n return len(self.list_data_dict)\n\n def __getitem__(self, i) -> Dict[str, torch.Tensor]:\n sources = self.list_data_dict[i]\n if isinstance(i, int):\n sources = [sources]\n assert len(sources) == 1, \"Don't know why it is wrapped to a list\" # FIXME\n if \"image\" in sources[0]:\n image_file = self.list_data_dict[i][\"image\"]\n image_folder = self.data_args.image_folder\n processor = self.data_args.image_processor\n image = Image.open(os.path.join(image_folder, image_file)).convert(\"RGB\")\n if self.data_args.image_aspect_ratio == \"pad\":\n\n def expand2square(pil_img, background_color):\n width, height = pil_img.size\n if width == height:\n return pil_img\n elif width > height:\n result = Image.new(\n pil_img.mode, (width, width), background_color\n )\n result.paste(pil_img, (0, (width - height) // 2))\n return result\n else:\n result = Image.new(\n pil_img.mode, (height, height), background_color\n )\n result.paste(pil_img, ((height - width) // 2, 0))\n return result\n\n image = expand2square(\n image, tuple(int(x * 255) for x in processor.image_mean)\n )\n image = processor.preprocess(image, return_tensors=\"pt\")[\n \"pixel_values\"\n ][0]\n else:\n image = processor.preprocess(image, return_tensors=\"pt\")[\n \"pixel_values\"\n ][0]\n sources = preprocess_multimodal(\n copy.deepcopy([e[\"conversations\"] for e in sources]), self.data_args\n )\n else:\n sources = copy.deepcopy([e[\"conversations\"] for e in sources])\n data_dict = preprocess(\n sources, self.tokenizer, has_image=(\"image\" in self.list_data_dict[i])\n )\n if isinstance(i, int):\n data_dict = dict(\n input_ids=data_dict[\"input_ids\"][0], labels=data_dict[\"labels\"][0]\n )\n\n # image exist in the data\n if \"image\" in self.list_data_dict[i]:\n data_dict[\"image\"] = image\n elif self.data_args.is_multimodal:\n # image does not exist in the data, but the model is multimodal\n crop_size = self.data_args.image_processor.crop_size\n data_dict[\"image\"] = torch.zeros(3, crop_size[\"height\"], crop_size[\"width\"])\n return data_dict\n\n\n@dataclass\nclass DataCollatorForSupervisedDataset(object):\n \"\"\"Collate examples for supervised fine-tuning.\"\"\"\n\n tokenizer: transformers.PreTrainedTokenizer\n\n def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:\n input_ids, labels = tuple(\n [instance[key] for instance in instances] for key in (\"input_ids\", \"labels\")\n )\n input_ids = torch.nn.utils.rnn.pad_sequence(\n input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id\n )\n labels = torch.nn.utils.rnn.pad_sequence(\n labels, batch_first=True, padding_value=IGNORE_INDEX\n )\n input_ids = input_ids[:, : self.tokenizer.model_max_length]\n labels = labels[:, : self.tokenizer.model_max_length]\n batch = dict(\n input_ids=input_ids,\n labels=labels,\n attention_mask=input_ids.ne(self.tokenizer.pad_token_id),\n )\n\n if \"image\" in instances[0]:\n images = [instance[\"image\"] for instance in instances]\n if all(x is not None and x.shape == images[0].shape for x in images):\n batch[\"images\"] = torch.stack(images)\n else:\n batch[\"images\"] = images\n\n return batch\n\n\ndef make_supervised_data_module(\n tokenizer: transformers.PreTrainedTokenizer, data_args\n) -> Dict:\n \"\"\"Make dataset and collator for supervised fine-tuning.\"\"\"\n train_dataset = LazySupervisedDataset(\n tokenizer=tokenizer, data_path=data_args.data_path, data_args=data_args\n )\n data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)\n return dict(\n train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator\n )\n\n\ndef train():\n global local_rank\n\n parser = transformers.HfArgumentParser(\n (ModelArguments, DataArguments, TrainingArguments)\n )\n model_args, data_args, training_args = parser.parse_args_into_dataclasses()\n local_rank = training_args.local_rank\n compute_dtype = (\n torch.float16\n if training_args.fp16\n else (torch.bfloat16 if training_args.bf16 else torch.float32)\n )\n\n bnb_model_from_pretrained_args = {}\n if training_args.bits in [4, 8]:\n from transformers import BitsAndBytesConfig\n\n bnb_model_from_pretrained_args.update(\n dict(\n device_map={\"\": training_args.device},\n load_in_4bit=training_args.bits == 4,\n load_in_8bit=training_args.bits == 8,\n quantization_config=BitsAndBytesConfig(\n load_in_4bit=training_args.bits == 4,\n load_in_8bit=training_args.bits == 8,\n llm_int8_threshold=6.0,\n llm_int8_has_fp16_weight=False,\n bnb_4bit_compute_dtype=compute_dtype,\n bnb_4bit_use_double_quant=training_args.double_quant,\n bnb_4bit_quant_type=training_args.quant_type, # {'fp4', 'nf4'}\n ),\n )\n )\n\n if model_args.vision_tower is not None:\n if \"mpt\" in model_args.model_name_or_path:\n config = transformers.AutoConfig.from_pretrained(\n model_args.model_name_or_path, trust_remote_code=True\n )\n config.attn_config[\"attn_impl\"] = training_args.mpt_attn_impl\n model = LlavaMPTForCausalLM.from_pretrained(\n model_args.model_name_or_path,\n config=config,\n cache_dir=training_args.cache_dir,\n **bnb_model_from_pretrained_args,\n )\n else:\n model = LlavaLlamaForCausalLM.from_pretrained(\n model_args.model_name_or_path,\n cache_dir=training_args.cache_dir,\n **bnb_model_from_pretrained_args,\n )\n else:\n model = transformers.LlamaForCausalLM.from_pretrained(\n model_args.model_name_or_path,\n cache_dir=training_args.cache_dir,\n **bnb_model_from_pretrained_args,\n )\n model.config.use_cache = False\n\n if model_args.freeze_backbone:\n model.model.requires_grad_(False)\n\n if training_args.bits in [4, 8]:\n from peft import prepare_model_for_kbit_training\n\n model.config.torch_dtype = (\n torch.float32\n if training_args.fp16\n else (torch.bfloat16 if training_args.bf16 else torch.float32)\n )\n model = prepare_model_for_kbit_training(\n model, use_gradient_checkpointing=training_args.gradient_checkpointing\n )\n\n if training_args.gradient_checkpointing:\n if hasattr(model, \"enable_input_require_grads\"):\n model.enable_input_require_grads()\n else:\n\n def make_inputs_require_grad(module, input, output):\n output.requires_grad_(True)\n\n model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n if training_args.lora_enable:\n from peft import LoraConfig, get_peft_model\n\n lora_config = LoraConfig(\n r=training_args.lora_r,\n lora_alpha=training_args.lora_alpha,\n target_modules=find_all_linear_names(model),\n lora_dropout=training_args.lora_dropout,\n bias=training_args.lora_bias,\n task_type=\"CAUSAL_LM\",\n )\n if training_args.bits == 16:\n if training_args.bf16:\n model.to(torch.bfloat16)\n if training_args.fp16:\n model.to(torch.float16)\n rank0_print(\"Adding LoRA adapters...\")\n model = get_peft_model(model, lora_config)\n\n if \"mpt\" in model_args.model_name_or_path:\n tokenizer = transformers.AutoTokenizer.from_pretrained(\n model_args.model_name_or_path,\n cache_dir=training_args.cache_dir,\n model_max_length=training_args.model_max_length,\n padding_side=\"right\",\n )\n else:\n tokenizer = transformers.AutoTokenizer.from_pretrained(\n model_args.model_name_or_path,\n cache_dir=training_args.cache_dir,\n model_max_length=training_args.model_max_length,\n padding_side=\"right\",\n use_fast=False,\n )\n\n if model_args.version == \"v0\":\n if tokenizer.pad_token is None:\n smart_tokenizer_and_embedding_resize(\n special_tokens_dict=dict(pad_token=\"[PAD]\"),\n tokenizer=tokenizer,\n model=model,\n )\n elif model_args.version == \"v0.5\":\n tokenizer.pad_token = tokenizer.unk_token\n else:\n tokenizer.pad_token = tokenizer.unk_token\n if model_args.version in conversation_lib.conv_templates:\n conversation_lib.default_conversation = conversation_lib.conv_templates[\n model_args.version\n ]\n else:\n conversation_lib.default_conversation = conversation_lib.conv_templates[\n \"vicuna_v1\"\n ]\n\n if model_args.vision_tower is not None:\n model.get_model().initialize_vision_modules(\n model_args=model_args, fsdp=training_args.fsdp\n )\n\n vision_tower = model.get_vision_tower()\n vision_tower.to(dtype=torch.float16, device=training_args.device)\n\n data_args.image_processor = vision_tower.image_processor\n data_args.is_multimodal = True\n\n model.config.image_aspect_ratio = data_args.image_aspect_ratio\n model.config.image_grid_pinpoints = data_args.image_grid_pinpoints\n\n model.config.tune_mm_mlp_adapter = (\n training_args.tune_mm_mlp_adapter\n ) = model_args.tune_mm_mlp_adapter\n if model_args.tune_mm_mlp_adapter:\n model.requires_grad_(False)\n for p in model.get_model().mm_projector.parameters():\n p.requires_grad = True\n\n model.config.freeze_mm_mlp_adapter = training_args.freeze_mm_mlp_adapter\n if training_args.freeze_mm_mlp_adapter:\n for p in model.get_model().mm_projector.parameters():\n p.requires_grad = False\n\n if training_args.bits in [4, 8]:\n model.get_model().mm_projector.to(\n dtype=compute_dtype, device=training_args.device\n )\n\n model.config.mm_use_im_start_end = (\n data_args.mm_use_im_start_end\n ) = model_args.mm_use_im_start_end\n training_args.use_im_start_end = model_args.mm_use_im_start_end\n model.config.mm_use_im_patch_token = model_args.mm_use_im_patch_token\n model.initialize_vision_tokenizer(model_args, tokenizer=tokenizer)\n\n if training_args.bits in [4, 8]:\n from peft.tuners.lora import LoraLayer\n\n for name, module in model.named_modules():\n if isinstance(module, LoraLayer):\n if training_args.bf16:\n module = module.to(torch.bfloat16)\n if \"norm\" in name:\n module = module.to(torch.float32)\n if \"lm_head\" in name or \"embed_tokens\" in name:\n if hasattr(module, \"weight\"):\n if training_args.bf16 and module.weight.dtype == torch.float32:\n module = module.to(torch.bfloat16)\n\n data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)\n trainer = LLaVATrainer(\n model=model, tokenizer=tokenizer, args=training_args, **data_module\n )\n\n if list(pathlib.Path(training_args.output_dir).glob(\"checkpoint-*\")):\n trainer.train(resume_from_checkpoint=True)\n else:\n trainer.train()\n trainer.save_state()\n\n model.config.use_cache = True\n\n if training_args.lora_enable:\n state_dict = get_peft_state_maybe_zero_3(\n model.named_parameters(), training_args.lora_bias\n )\n non_lora_state_dict = get_peft_state_non_lora_maybe_zero_3(\n model.named_parameters()\n )\n if training_args.local_rank == 0 or training_args.local_rank == -1:\n model.config.save_pretrained(training_args.output_dir)\n model.save_pretrained(training_args.output_dir, state_dict=state_dict)\n torch.save(\n non_lora_state_dict,\n os.path.join(training_args.output_dir, \"non_lora_trainables.bin\"),\n )\n else:\n safe_save_model_for_hf_trainer(\n trainer=trainer, output_dir=training_args.output_dir\n )\n\n\nif __name__ == \"__main__\":\n train()\n" }, { "path": "model/llava/train/train_mem.py", "content": "# Adopted from https://github.com/lm-sys/FastChat. Below is the original copyright:\n# Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:\n# Make it more memory efficient by monkey patching the LLaMA model with FlashAttn.\n\n# Need to call this before importing transformers.\nfrom llava.train.llama_flash_attn_monkey_patch import \\\n replace_llama_attn_with_flash_attn\n\nreplace_llama_attn_with_flash_attn()\n\nfrom llava.train.train import train\n\nif __name__ == \"__main__\":\n train()\n" }, { "path": "model/llava/utils.py", "content": "import datetime\nimport logging\nimport logging.handlers\nimport os\nimport sys\n\nimport requests\nfrom llava.constants import LOGDIR\n\nserver_error_msg = (\n \"**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**\"\n)\nmoderation_msg = (\n \"YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN.\"\n)\n\nhandler = None\n\n\ndef build_logger(logger_name, logger_filename):\n global handler\n\n formatter = logging.Formatter(\n fmt=\"%(asctime)s | %(levelname)s | %(name)s | %(message)s\",\n datefmt=\"%Y-%m-%d %H:%M:%S\",\n )\n\n # Set the format of root handlers\n if not logging.getLogger().handlers:\n logging.basicConfig(level=logging.INFO)\n logging.getLogger().handlers[0].setFormatter(formatter)\n\n # Redirect stdout and stderr to loggers\n stdout_logger = logging.getLogger(\"stdout\")\n stdout_logger.setLevel(logging.INFO)\n sl = StreamToLogger(stdout_logger, logging.INFO)\n sys.stdout = sl\n\n stderr_logger = logging.getLogger(\"stderr\")\n stderr_logger.setLevel(logging.ERROR)\n sl = StreamToLogger(stderr_logger, logging.ERROR)\n sys.stderr = sl\n\n # Get logger\n logger = logging.getLogger(logger_name)\n logger.setLevel(logging.INFO)\n\n # Add a file handler for all loggers\n if handler is None:\n os.makedirs(LOGDIR, exist_ok=True)\n filename = os.path.join(LOGDIR, logger_filename)\n handler = logging.handlers.TimedRotatingFileHandler(\n filename, when=\"D\", utc=True\n )\n handler.setFormatter(formatter)\n\n for name, item in logging.root.manager.loggerDict.items():\n if isinstance(item, logging.Logger):\n item.addHandler(handler)\n\n return logger\n\n\nclass StreamToLogger(object):\n \"\"\"\n Fake file-like stream object that redirects writes to a logger instance.\n \"\"\"\n\n def __init__(self, logger, log_level=logging.INFO):\n self.terminal = sys.stdout\n self.logger = logger\n self.log_level = log_level\n self.linebuf = \"\"\n\n def __getattr__(self, attr):\n return getattr(self.terminal, attr)\n\n def write(self, buf):\n temp_linebuf = self.linebuf + buf\n self.linebuf = \"\"\n for line in temp_linebuf.splitlines(True):\n # From the io.TextIOWrapper docs:\n # On output, if newline is None, any '\\n' characters written\n # are translated to the system default line separator.\n # By default sys.stdout.write() expects '\\n' newlines and then\n # translates them so this is still cross platform.\n if line[-1] == \"\\n\":\n self.logger.log(self.log_level, line.rstrip())\n else:\n self.linebuf += line\n\n def flush(self):\n if self.linebuf != \"\":\n self.logger.log(self.log_level, self.linebuf.rstrip())\n self.linebuf = \"\"\n\n\ndef disable_torch_init():\n \"\"\"\n Disable the redundant torch default initialization to accelerate model creation.\n \"\"\"\n import torch\n\n setattr(torch.nn.Linear, \"reset_parameters\", lambda self: None)\n setattr(torch.nn.LayerNorm, \"reset_parameters\", lambda self: None)\n\n\ndef violates_moderation(text):\n \"\"\"\n Check whether the text violates OpenAI moderation API.\n \"\"\"\n url = \"https://api.openai.com/v1/moderations\"\n headers = {\n \"Content-Type\": \"application/json\",\n \"Authorization\": \"Bearer \" + os.environ[\"OPENAI_API_KEY\"],\n }\n text = text.replace(\"\\n\", \"\")\n data = \"{\" + '\"input\": ' + f'\"{text}\"' + \"}\"\n data = data.encode(\"utf-8\")\n try:\n ret = requests.post(url, headers=headers, data=data, timeout=5)\n flagged = ret.json()[\"results\"][0][\"flagged\"]\n except requests.exceptions.RequestException as e:\n flagged = False\n except KeyError as e:\n flagged = False\n\n return flagged\n\n\ndef pretty_print_semaphore(semaphore):\n if semaphore is None:\n return \"None\"\n return f\"Semaphore(value={semaphore._value}, locked={semaphore.locked()})\"\n" }, { "path": "model/segment_anything/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom .automatic_mask_generator import SamAutomaticMaskGenerator\nfrom .build_sam import (build_sam, build_sam_vit_b, build_sam_vit_h,\n build_sam_vit_l, sam_model_registry)\nfrom .predictor import SamPredictor\n" }, { "path": "model/segment_anything/automatic_mask_generator.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Any, Dict, List, Optional, Tuple\n\nimport numpy as np\nimport torch\nfrom torchvision.ops.boxes import batched_nms, box_area # type: ignore\n\nfrom .modeling import Sam\nfrom .predictor import SamPredictor\nfrom .utils.amg import (MaskData, area_from_rle, batch_iterator,\n batched_mask_to_box, box_xyxy_to_xywh,\n build_all_layer_point_grids, calculate_stability_score,\n coco_encode_rle, generate_crop_boxes,\n is_box_near_crop_edge, mask_to_rle_pytorch,\n remove_small_regions, rle_to_mask, uncrop_boxes_xyxy,\n uncrop_masks, uncrop_points)\n\n\nclass SamAutomaticMaskGenerator:\n def __init__(\n self,\n model: Sam,\n points_per_side: Optional[int] = 32,\n points_per_batch: int = 64,\n pred_iou_thresh: float = 0.88,\n stability_score_thresh: float = 0.95,\n stability_score_offset: float = 1.0,\n box_nms_thresh: float = 0.7,\n crop_n_layers: int = 0,\n crop_nms_thresh: float = 0.7,\n crop_overlap_ratio: float = 512 / 1500,\n crop_n_points_downscale_factor: int = 1,\n point_grids: Optional[List[np.ndarray]] = None,\n min_mask_region_area: int = 0,\n output_mode: str = \"binary_mask\",\n ) -> None:\n \"\"\"\n Using a SAM model, generates masks for the entire image.\n Generates a grid of point prompts over the image, then filters\n low quality and duplicate masks. The default settings are chosen\n for SAM with a ViT-H backbone.\n\n Arguments:\n model (Sam): The SAM model to use for mask prediction.\n points_per_side (int or None): The number of points to be sampled\n along one side of the image. The total number of points is\n points_per_side**2. If None, 'point_grids' must provide explicit\n point sampling.\n points_per_batch (int): Sets the number of points run simultaneously\n by the model. Higher numbers may be faster but use more GPU memory.\n pred_iou_thresh (float): A filtering threshold in [0,1], using the\n model's predicted mask quality.\n stability_score_thresh (float): A filtering threshold in [0,1], using\n the stability of the mask under changes to the cutoff used to binarize\n the model's mask predictions.\n stability_score_offset (float): The amount to shift the cutoff when\n calculated the stability score.\n box_nms_thresh (float): The box IoU cutoff used by non-maximal\n suppression to filter duplicate masks.\n crop_n_layers (int): If >0, mask prediction will be run again on\n crops of the image. Sets the number of layers to run, where each\n layer has 2**i_layer number of image crops.\n crop_nms_thresh (float): The box IoU cutoff used by non-maximal\n suppression to filter duplicate masks between different crops.\n crop_overlap_ratio (float): Sets the degree to which crops overlap.\n In the first crop layer, crops will overlap by this fraction of\n the image length. Later layers with more crops scale down this overlap.\n crop_n_points_downscale_factor (int): The number of points-per-side\n sampled in layer n is scaled down by crop_n_points_downscale_factor**n.\n point_grids (list(np.ndarray) or None): A list over explicit grids\n of points used for sampling, normalized to [0,1]. The nth grid in the\n list is used in the nth crop layer. Exclusive with points_per_side.\n min_mask_region_area (int): If >0, postprocessing will be applied\n to remove disconnected regions and holes in masks with area smaller\n than min_mask_region_area. Requires opencv.\n output_mode (str): The form masks are returned in. Can be 'binary_mask',\n 'uncompressed_rle', or 'coco_rle'. 'coco_rle' requires pycocotools.\n For large resolutions, 'binary_mask' may consume large amounts of\n memory.\n \"\"\"\n\n assert (points_per_side is None) != (\n point_grids is None\n ), \"Exactly one of points_per_side or point_grid must be provided.\"\n if points_per_side is not None:\n self.point_grids = build_all_layer_point_grids(\n points_per_side,\n crop_n_layers,\n crop_n_points_downscale_factor,\n )\n elif point_grids is not None:\n self.point_grids = point_grids\n else:\n raise ValueError(\"Can't have both points_per_side and point_grid be None.\")\n\n assert output_mode in [\n \"binary_mask\",\n \"uncompressed_rle\",\n \"coco_rle\",\n ], f\"Unknown output_mode {output_mode}.\"\n if output_mode == \"coco_rle\":\n from pycocotools import \\\n mask as mask_utils # type: ignore # noqa: F401\n\n if min_mask_region_area > 0:\n import cv2 # type: ignore # noqa: F401\n\n self.predictor = SamPredictor(model)\n self.points_per_batch = points_per_batch\n self.pred_iou_thresh = pred_iou_thresh\n self.stability_score_thresh = stability_score_thresh\n self.stability_score_offset = stability_score_offset\n self.box_nms_thresh = box_nms_thresh\n self.crop_n_layers = crop_n_layers\n self.crop_nms_thresh = crop_nms_thresh\n self.crop_overlap_ratio = crop_overlap_ratio\n self.crop_n_points_downscale_factor = crop_n_points_downscale_factor\n self.min_mask_region_area = min_mask_region_area\n self.output_mode = output_mode\n\n @torch.no_grad()\n def generate(self, image: np.ndarray) -> List[Dict[str, Any]]:\n \"\"\"\n Generates masks for the given image.\n\n Arguments:\n image (np.ndarray): The image to generate masks for, in HWC uint8 format.\n\n Returns:\n list(dict(str, any)): A list over records for masks. Each record is\n a dict containing the following keys:\n segmentation (dict(str, any) or np.ndarray): The mask. If\n output_mode='binary_mask', is an array of shape HW. Otherwise,\n is a dictionary containing the RLE.\n bbox (list(float)): The box around the mask, in XYWH format.\n area (int): The area in pixels of the mask.\n predicted_iou (float): The model's own prediction of the mask's\n quality. This is filtered by the pred_iou_thresh parameter.\n point_coords (list(list(float))): The point coordinates input\n to the model to generate this mask.\n stability_score (float): A measure of the mask's quality. This\n is filtered on using the stability_score_thresh parameter.\n crop_box (list(float)): The crop of the image used to generate\n the mask, given in XYWH format.\n \"\"\"\n\n # Generate masks\n mask_data = self._generate_masks(image)\n\n # Filter small disconnected regions and holes in masks\n if self.min_mask_region_area > 0:\n mask_data = self.postprocess_small_regions(\n mask_data,\n self.min_mask_region_area,\n max(self.box_nms_thresh, self.crop_nms_thresh),\n )\n\n # Encode masks\n if self.output_mode == \"coco_rle\":\n mask_data[\"segmentations\"] = [\n coco_encode_rle(rle) for rle in mask_data[\"rles\"]\n ]\n elif self.output_mode == \"binary_mask\":\n mask_data[\"segmentations\"] = [rle_to_mask(rle) for rle in mask_data[\"rles\"]]\n else:\n mask_data[\"segmentations\"] = mask_data[\"rles\"]\n\n # Write mask records\n curr_anns = []\n for idx in range(len(mask_data[\"segmentations\"])):\n ann = {\n \"segmentation\": mask_data[\"segmentations\"][idx],\n \"area\": area_from_rle(mask_data[\"rles\"][idx]),\n \"bbox\": box_xyxy_to_xywh(mask_data[\"boxes\"][idx]).tolist(),\n \"predicted_iou\": mask_data[\"iou_preds\"][idx].item(),\n \"point_coords\": [mask_data[\"points\"][idx].tolist()],\n \"stability_score\": mask_data[\"stability_score\"][idx].item(),\n \"crop_box\": box_xyxy_to_xywh(mask_data[\"crop_boxes\"][idx]).tolist(),\n }\n curr_anns.append(ann)\n\n return curr_anns\n\n def _generate_masks(self, image: np.ndarray) -> MaskData:\n orig_size = image.shape[:2]\n crop_boxes, layer_idxs = generate_crop_boxes(\n orig_size, self.crop_n_layers, self.crop_overlap_ratio\n )\n\n # Iterate over image crops\n data = MaskData()\n for crop_box, layer_idx in zip(crop_boxes, layer_idxs):\n crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)\n data.cat(crop_data)\n\n # Remove duplicate masks between crops\n if len(crop_boxes) > 1:\n # Prefer masks from smaller crops\n scores = 1 / box_area(data[\"crop_boxes\"])\n scores = scores.to(data[\"boxes\"].device)\n keep_by_nms = batched_nms(\n data[\"boxes\"].float(),\n scores,\n torch.zeros_like(data[\"boxes\"][:, 0]), # categories\n iou_threshold=self.crop_nms_thresh,\n )\n data.filter(keep_by_nms)\n\n data.to_numpy()\n return data\n\n def _process_crop(\n self,\n image: np.ndarray,\n crop_box: List[int],\n crop_layer_idx: int,\n orig_size: Tuple[int, ...],\n ) -> MaskData:\n # Crop the image and calculate embeddings\n x0, y0, x1, y1 = crop_box\n cropped_im = image[y0:y1, x0:x1, :]\n cropped_im_size = cropped_im.shape[:2]\n self.predictor.set_image(cropped_im)\n\n # Get points for this crop\n points_scale = np.array(cropped_im_size)[None, ::-1]\n points_for_image = self.point_grids[crop_layer_idx] * points_scale\n\n # Generate masks for this crop in batches\n data = MaskData()\n for (points,) in batch_iterator(self.points_per_batch, points_for_image):\n batch_data = self._process_batch(\n points, cropped_im_size, crop_box, orig_size\n )\n data.cat(batch_data)\n del batch_data\n self.predictor.reset_image()\n\n # Remove duplicates within this crop.\n keep_by_nms = batched_nms(\n data[\"boxes\"].float(),\n data[\"iou_preds\"],\n torch.zeros_like(data[\"boxes\"][:, 0]), # categories\n iou_threshold=self.box_nms_thresh,\n )\n data.filter(keep_by_nms)\n\n # Return to the original image frame\n data[\"boxes\"] = uncrop_boxes_xyxy(data[\"boxes\"], crop_box)\n data[\"points\"] = uncrop_points(data[\"points\"], crop_box)\n data[\"crop_boxes\"] = torch.tensor([crop_box for _ in range(len(data[\"rles\"]))])\n\n return data\n\n def _process_batch(\n self,\n points: np.ndarray,\n im_size: Tuple[int, ...],\n crop_box: List[int],\n orig_size: Tuple[int, ...],\n ) -> MaskData:\n orig_h, orig_w = orig_size\n\n # Run model on this batch\n transformed_points = self.predictor.transform.apply_coords(points, im_size)\n in_points = torch.as_tensor(transformed_points, device=self.predictor.device)\n in_labels = torch.ones(\n in_points.shape[0], dtype=torch.int, device=in_points.device\n )\n masks, iou_preds, _ = self.predictor.predict_torch(\n in_points[:, None, :],\n in_labels[:, None],\n multimask_output=True,\n return_logits=True,\n )\n\n # Serialize predictions and store in MaskData\n data = MaskData(\n masks=masks.flatten(0, 1),\n iou_preds=iou_preds.flatten(0, 1),\n points=torch.as_tensor(points.repeat(masks.shape[1], axis=0)),\n )\n del masks\n\n # Filter by predicted IoU\n if self.pred_iou_thresh > 0.0:\n keep_mask = data[\"iou_preds\"] > self.pred_iou_thresh\n data.filter(keep_mask)\n\n # Calculate stability score\n data[\"stability_score\"] = calculate_stability_score(\n data[\"masks\"],\n self.predictor.model.mask_threshold,\n self.stability_score_offset,\n )\n if self.stability_score_thresh > 0.0:\n keep_mask = data[\"stability_score\"] >= self.stability_score_thresh\n data.filter(keep_mask)\n\n # Threshold masks and calculate boxes\n data[\"masks\"] = data[\"masks\"] > self.predictor.model.mask_threshold\n data[\"boxes\"] = batched_mask_to_box(data[\"masks\"])\n\n # Filter boxes that touch crop boundaries\n keep_mask = ~is_box_near_crop_edge(\n data[\"boxes\"], crop_box, [0, 0, orig_w, orig_h]\n )\n if not torch.all(keep_mask):\n data.filter(keep_mask)\n\n # Compress to RLE\n data[\"masks\"] = uncrop_masks(data[\"masks\"], crop_box, orig_h, orig_w)\n data[\"rles\"] = mask_to_rle_pytorch(data[\"masks\"])\n del data[\"masks\"]\n\n return data\n\n @staticmethod\n def postprocess_small_regions(\n mask_data: MaskData, min_area: int, nms_thresh: float\n ) -> MaskData:\n \"\"\"\n Removes small disconnected regions and holes in masks, then reruns\n box NMS to remove any new duplicates.\n\n Edits mask_data in place.\n\n Requires open-cv as a dependency.\n \"\"\"\n if len(mask_data[\"rles\"]) == 0:\n return mask_data\n\n # Filter small disconnected regions and holes\n new_masks = []\n scores = []\n for rle in mask_data[\"rles\"]:\n mask = rle_to_mask(rle)\n\n mask, changed = remove_small_regions(mask, min_area, mode=\"holes\")\n unchanged = not changed\n mask, changed = remove_small_regions(mask, min_area, mode=\"islands\")\n unchanged = unchanged and not changed\n\n new_masks.append(torch.as_tensor(mask).unsqueeze(0))\n # Give score=0 to changed masks and score=1 to unchanged masks\n # so NMS will prefer ones that didn't need postprocessing\n scores.append(float(unchanged))\n\n # Recalculate boxes and remove any new duplicates\n masks = torch.cat(new_masks, dim=0)\n boxes = batched_mask_to_box(masks)\n keep_by_nms = batched_nms(\n boxes.float(),\n torch.as_tensor(scores),\n torch.zeros_like(boxes[:, 0]), # categories\n iou_threshold=nms_thresh,\n )\n\n # Only recalculate RLEs for masks that have changed\n for i_mask in keep_by_nms:\n if scores[i_mask] == 0.0:\n mask_torch = masks[i_mask].unsqueeze(0)\n mask_data[\"rles\"][i_mask] = mask_to_rle_pytorch(mask_torch)[0]\n mask_data[\"boxes\"][i_mask] = boxes[i_mask] # update res directly\n mask_data.filter(keep_by_nms)\n\n return mask_data\n" }, { "path": "model/segment_anything/build_sam.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom functools import partial\n\nimport torch\n\nfrom .modeling import (ImageEncoderViT, MaskDecoder, PromptEncoder, Sam,\n TwoWayTransformer)\n\n\ndef build_sam_vit_h(checkpoint=None):\n return _build_sam(\n encoder_embed_dim=1280,\n encoder_depth=32,\n encoder_num_heads=16,\n encoder_global_attn_indexes=[7, 15, 23, 31],\n checkpoint=checkpoint,\n )\n\n\nbuild_sam = build_sam_vit_h\n\n\ndef build_sam_vit_l(checkpoint=None):\n return _build_sam(\n encoder_embed_dim=1024,\n encoder_depth=24,\n encoder_num_heads=16,\n encoder_global_attn_indexes=[5, 11, 17, 23],\n checkpoint=checkpoint,\n )\n\n\ndef build_sam_vit_b(checkpoint=None):\n return _build_sam(\n encoder_embed_dim=768,\n encoder_depth=12,\n encoder_num_heads=12,\n encoder_global_attn_indexes=[2, 5, 8, 11],\n checkpoint=checkpoint,\n )\n\n\nsam_model_registry = {\n \"default\": build_sam_vit_h,\n \"vit_h\": build_sam_vit_h,\n \"vit_l\": build_sam_vit_l,\n \"vit_b\": build_sam_vit_b,\n}\n\n\ndef _build_sam(\n encoder_embed_dim,\n encoder_depth,\n encoder_num_heads,\n encoder_global_attn_indexes,\n checkpoint=None,\n):\n prompt_embed_dim = 256\n image_size = 1024\n vit_patch_size = 16\n image_embedding_size = image_size // vit_patch_size\n sam = Sam(\n image_encoder=ImageEncoderViT(\n depth=encoder_depth,\n embed_dim=encoder_embed_dim,\n img_size=image_size,\n mlp_ratio=4,\n norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),\n num_heads=encoder_num_heads,\n patch_size=vit_patch_size,\n qkv_bias=True,\n use_rel_pos=True,\n global_attn_indexes=encoder_global_attn_indexes,\n window_size=14,\n out_chans=prompt_embed_dim,\n ),\n prompt_encoder=PromptEncoder(\n embed_dim=prompt_embed_dim,\n image_embedding_size=(image_embedding_size, image_embedding_size),\n input_image_size=(image_size, image_size),\n mask_in_chans=16,\n ),\n mask_decoder=MaskDecoder(\n num_multimask_outputs=3,\n transformer=TwoWayTransformer(\n depth=2,\n embedding_dim=prompt_embed_dim,\n mlp_dim=2048,\n num_heads=8,\n ),\n transformer_dim=prompt_embed_dim,\n iou_head_depth=3,\n iou_head_hidden_dim=256,\n ),\n pixel_mean=[123.675, 116.28, 103.53],\n pixel_std=[58.395, 57.12, 57.375],\n )\n sam.eval()\n if checkpoint is not None:\n with open(checkpoint, \"rb\") as f:\n state_dict = torch.load(f)\n sam.load_state_dict(state_dict, strict=False)\n return sam\n" }, { "path": "model/segment_anything/modeling/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom .image_encoder import ImageEncoderViT\nfrom .mask_decoder import MaskDecoder\nfrom .prompt_encoder import PromptEncoder\nfrom .sam import Sam\nfrom .transformer import TwoWayTransformer\n" }, { "path": "model/segment_anything/modeling/common.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Type\n\nimport torch\nimport torch.nn as nn\n\n\nclass MLPBlock(nn.Module):\n def __init__(\n self,\n embedding_dim: int,\n mlp_dim: int,\n act: Type[nn.Module] = nn.GELU,\n ) -> None:\n super().__init__()\n self.lin1 = nn.Linear(embedding_dim, mlp_dim)\n self.lin2 = nn.Linear(mlp_dim, embedding_dim)\n self.act = act()\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.lin2(self.act(self.lin1(x)))\n\n\n# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa\n# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa\nclass LayerNorm2d(nn.Module):\n def __init__(self, num_channels: int, eps: float = 1e-6) -> None:\n super().__init__()\n self.weight = nn.Parameter(torch.ones(num_channels))\n self.bias = nn.Parameter(torch.zeros(num_channels))\n self.eps = eps\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n u = x.mean(1, keepdim=True)\n s = (x - u).pow(2).mean(1, keepdim=True)\n x = (x - u) / torch.sqrt(s + self.eps)\n x = self.weight[:, None, None] * x + self.bias[:, None, None]\n return x\n" }, { "path": "model/segment_anything/modeling/image_encoder.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Optional, Tuple, Type\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom .common import LayerNorm2d, MLPBlock\n\n\n# This class and its supporting functions below lightly adapted from the ViTDet backbone available at: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py # noqa\nclass ImageEncoderViT(nn.Module):\n def __init__(\n self,\n img_size: int = 1024,\n patch_size: int = 16,\n in_chans: int = 3,\n embed_dim: int = 768,\n depth: int = 12,\n num_heads: int = 12,\n mlp_ratio: float = 4.0,\n out_chans: int = 256,\n qkv_bias: bool = True,\n norm_layer: Type[nn.Module] = nn.LayerNorm,\n act_layer: Type[nn.Module] = nn.GELU,\n use_abs_pos: bool = True,\n use_rel_pos: bool = False,\n rel_pos_zero_init: bool = True,\n window_size: int = 0,\n global_attn_indexes: Tuple[int, ...] = (),\n ) -> None:\n \"\"\"\n Args:\n img_size (int): Input image size.\n patch_size (int): Patch size.\n in_chans (int): Number of input image channels.\n embed_dim (int): Patch embedding dimension.\n depth (int): Depth of ViT.\n num_heads (int): Number of attention heads in each ViT block.\n mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n qkv_bias (bool): If True, add a learnable bias to query, key, value.\n norm_layer (nn.Module): Normalization layer.\n act_layer (nn.Module): Activation layer.\n use_abs_pos (bool): If True, use absolute positional embeddings.\n use_rel_pos (bool): If True, add relative positional embeddings to the attention map.\n rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.\n window_size (int): Window size for window attention blocks.\n global_attn_indexes (list): Indexes for blocks using global attention.\n \"\"\"\n super().__init__()\n self.img_size = img_size\n self.embed_dim = embed_dim\n self.out_chans = out_chans\n\n self.patch_embed = PatchEmbed(\n kernel_size=(patch_size, patch_size),\n stride=(patch_size, patch_size),\n in_chans=in_chans,\n embed_dim=embed_dim,\n )\n\n self.pos_embed: Optional[nn.Parameter] = None\n if use_abs_pos:\n # Initialize absolute positional embedding with pretrain image size.\n self.pos_embed = nn.Parameter(\n torch.zeros(\n 1, img_size // patch_size, img_size // patch_size, embed_dim\n )\n )\n\n self.blocks = nn.ModuleList()\n for i in range(depth):\n block = Block(\n dim=embed_dim,\n num_heads=num_heads,\n mlp_ratio=mlp_ratio,\n qkv_bias=qkv_bias,\n norm_layer=norm_layer,\n act_layer=act_layer,\n use_rel_pos=use_rel_pos,\n rel_pos_zero_init=rel_pos_zero_init,\n window_size=window_size if i not in global_attn_indexes else 0,\n input_size=(img_size // patch_size, img_size // patch_size),\n )\n self.blocks.append(block)\n\n self.neck = nn.Sequential(\n nn.Conv2d(\n embed_dim,\n out_chans,\n kernel_size=1,\n bias=False,\n ),\n LayerNorm2d(out_chans),\n nn.Conv2d(\n out_chans,\n out_chans,\n kernel_size=3,\n padding=1,\n bias=False,\n ),\n LayerNorm2d(out_chans),\n )\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n x = self.patch_embed(x)\n if self.pos_embed is not None:\n x = x + self.pos_embed\n\n for blk in self.blocks:\n x = blk(x)\n\n dtype = x.dtype\n if dtype == torch.float16: # prevent overflow\n with torch.autocast(device_type=\"cuda\", dtype=torch.float32):\n x = self.neck(x.permute(0, 3, 1, 2))\n x = x.to(dtype)\n else:\n x = self.neck(x.permute(0, 3, 1, 2))\n return x\n\n\nclass Block(nn.Module):\n \"\"\"Transformer blocks with support of window attention and residual propagation blocks\"\"\"\n\n def __init__(\n self,\n dim: int,\n num_heads: int,\n mlp_ratio: float = 4.0,\n qkv_bias: bool = True,\n norm_layer: Type[nn.Module] = nn.LayerNorm,\n act_layer: Type[nn.Module] = nn.GELU,\n use_rel_pos: bool = False,\n rel_pos_zero_init: bool = True,\n window_size: int = 0,\n input_size: Optional[Tuple[int, int]] = None,\n ) -> None:\n \"\"\"\n Args:\n dim (int): Number of input channels.\n num_heads (int): Number of attention heads in each ViT block.\n mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n qkv_bias (bool): If True, add a learnable bias to query, key, value.\n norm_layer (nn.Module): Normalization layer.\n act_layer (nn.Module): Activation layer.\n use_rel_pos (bool): If True, add relative positional embeddings to the attention map.\n rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.\n window_size (int): Window size for window attention blocks. If it equals 0, then\n use global attention.\n input_size (tuple(int, int) or None): Input resolution for calculating the relative\n positional parameter size.\n \"\"\"\n super().__init__()\n self.norm1 = norm_layer(dim)\n self.attn = Attention(\n dim,\n num_heads=num_heads,\n qkv_bias=qkv_bias,\n use_rel_pos=use_rel_pos,\n rel_pos_zero_init=rel_pos_zero_init,\n input_size=input_size if window_size == 0 else (window_size, window_size),\n )\n\n self.norm2 = norm_layer(dim)\n self.mlp = MLPBlock(\n embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer\n )\n\n self.window_size = window_size\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n shortcut = x\n x = self.norm1(x)\n # Window partition\n if self.window_size > 0:\n H, W = x.shape[1], x.shape[2]\n x, pad_hw = window_partition(x, self.window_size)\n\n x = self.attn(x)\n # Reverse window partition\n if self.window_size > 0:\n x = window_unpartition(x, self.window_size, pad_hw, (H, W))\n\n x = shortcut + x\n x = x + self.mlp(self.norm2(x))\n\n return x\n\n\nclass Attention(nn.Module):\n \"\"\"Multi-head Attention block with relative position embeddings.\"\"\"\n\n def __init__(\n self,\n dim: int,\n num_heads: int = 8,\n qkv_bias: bool = True,\n use_rel_pos: bool = False,\n rel_pos_zero_init: bool = True,\n input_size: Optional[Tuple[int, int]] = None,\n ) -> None:\n \"\"\"\n Args:\n dim (int): Number of input channels.\n num_heads (int): Number of attention heads.\n qkv_bias (bool): If True, add a learnable bias to query, key, value.\n rel_pos (bool): If True, add relative positional embeddings to the attention map.\n rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.\n input_size (tuple(int, int) or None): Input resolution for calculating the relative\n positional parameter size.\n \"\"\"\n super().__init__()\n self.num_heads = num_heads\n head_dim = dim // num_heads\n self.scale = head_dim**-0.5\n\n self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n self.proj = nn.Linear(dim, dim)\n\n self.use_rel_pos = use_rel_pos\n if self.use_rel_pos:\n assert (\n input_size is not None\n ), \"Input size must be provided if using relative positional encoding.\"\n # initialize relative positional embeddings\n self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))\n self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n B, H, W, _ = x.shape\n # qkv with shape (3, B, nHead, H * W, C)\n qkv = (\n self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)\n )\n # q, k, v with shape (B * nHead, H * W, C)\n q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)\n\n attn = (q * self.scale) @ k.transpose(-2, -1)\n\n if self.use_rel_pos:\n attn = add_decomposed_rel_pos(\n attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W)\n )\n\n attn = attn.softmax(dim=-1)\n x = (\n (attn @ v)\n .view(B, self.num_heads, H, W, -1)\n .permute(0, 2, 3, 1, 4)\n .reshape(B, H, W, -1)\n )\n x = self.proj(x)\n\n return x\n\n\ndef window_partition(\n x: torch.Tensor, window_size: int\n) -> Tuple[torch.Tensor, Tuple[int, int]]:\n \"\"\"\n Partition into non-overlapping windows with padding if needed.\n Args:\n x (tensor): input tokens with [B, H, W, C].\n window_size (int): window size.\n\n Returns:\n windows: windows after partition with [B * num_windows, window_size, window_size, C].\n (Hp, Wp): padded height and width before partition\n \"\"\"\n B, H, W, C = x.shape\n\n pad_h = (window_size - H % window_size) % window_size\n pad_w = (window_size - W % window_size) % window_size\n if pad_h > 0 or pad_w > 0:\n x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))\n Hp, Wp = H + pad_h, W + pad_w\n\n x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)\n windows = (\n x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)\n )\n return windows, (Hp, Wp)\n\n\ndef window_unpartition(\n windows: torch.Tensor,\n window_size: int,\n pad_hw: Tuple[int, int],\n hw: Tuple[int, int],\n) -> torch.Tensor:\n \"\"\"\n Window unpartition into original sequences and removing padding.\n Args:\n windows (tensor): input tokens with [B * num_windows, window_size, window_size, C].\n window_size (int): window size.\n pad_hw (Tuple): padded height and width (Hp, Wp).\n hw (Tuple): original height and width (H, W) before padding.\n\n Returns:\n x: unpartitioned sequences with [B, H, W, C].\n \"\"\"\n Hp, Wp = pad_hw\n H, W = hw\n B = windows.shape[0] // (Hp * Wp // window_size // window_size)\n x = windows.view(\n B, Hp // window_size, Wp // window_size, window_size, window_size, -1\n )\n x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1)\n\n if Hp > H or Wp > W:\n x = x[:, :H, :W, :].contiguous()\n return x\n\n\ndef get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Get relative positional embeddings according to the relative positions of\n query and key sizes.\n Args:\n q_size (int): size of query q.\n k_size (int): size of key k.\n rel_pos (Tensor): relative position embeddings (L, C).\n\n Returns:\n Extracted positional embeddings according to relative positions.\n \"\"\"\n max_rel_dist = int(2 * max(q_size, k_size) - 1)\n # Interpolate rel pos if needed.\n if rel_pos.shape[0] != max_rel_dist:\n # Interpolate rel pos.\n rel_pos_resized = F.interpolate(\n rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),\n size=max_rel_dist,\n mode=\"linear\",\n )\n rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)\n else:\n rel_pos_resized = rel_pos\n\n # Scale the coords with short length if shapes for q and k are different.\n q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)\n k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)\n relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)\n\n return rel_pos_resized[relative_coords.long()]\n\n\ndef add_decomposed_rel_pos(\n attn: torch.Tensor,\n q: torch.Tensor,\n rel_pos_h: torch.Tensor,\n rel_pos_w: torch.Tensor,\n q_size: Tuple[int, int],\n k_size: Tuple[int, int],\n) -> torch.Tensor:\n \"\"\"\n Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.\n https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py # noqa B950\n Args:\n attn (Tensor): attention map.\n q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).\n rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.\n rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.\n q_size (Tuple): spatial sequence size of query q with (q_h, q_w).\n k_size (Tuple): spatial sequence size of key k with (k_h, k_w).\n\n Returns:\n attn (Tensor): attention map with added relative positional embeddings.\n \"\"\"\n q_h, q_w = q_size\n k_h, k_w = k_size\n Rh = get_rel_pos(q_h, k_h, rel_pos_h)\n Rw = get_rel_pos(q_w, k_w, rel_pos_w)\n\n B, _, dim = q.shape\n r_q = q.reshape(B, q_h, q_w, dim)\n rel_h = torch.einsum(\"bhwc,hkc->bhwk\", r_q, Rh)\n rel_w = torch.einsum(\"bhwc,wkc->bhwk\", r_q, Rw)\n\n attn = (\n attn.view(B, q_h, q_w, k_h, k_w)\n + rel_h[:, :, :, :, None]\n + rel_w[:, :, :, None, :]\n ).view(B, q_h * q_w, k_h * k_w)\n\n return attn\n\n\nclass PatchEmbed(nn.Module):\n \"\"\"\n Image to Patch Embedding.\n \"\"\"\n\n def __init__(\n self,\n kernel_size: Tuple[int, int] = (16, 16),\n stride: Tuple[int, int] = (16, 16),\n padding: Tuple[int, int] = (0, 0),\n in_chans: int = 3,\n embed_dim: int = 768,\n ) -> None:\n \"\"\"\n Args:\n kernel_size (Tuple): kernel size of the projection layer.\n stride (Tuple): stride of the projection layer.\n padding (Tuple): padding size of the projection layer.\n in_chans (int): Number of input image channels.\n embed_dim (int): Patch embedding dimension.\n \"\"\"\n super().__init__()\n\n self.proj = nn.Conv2d(\n in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding\n )\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n x = self.proj(x)\n # B C H W -> B H W C\n x = x.permute(0, 2, 3, 1)\n return x\n" }, { "path": "model/segment_anything/modeling/mask_decoder.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import List, Tuple, Type\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\nfrom .common import LayerNorm2d\n\n\nclass MaskDecoder(nn.Module):\n def __init__(\n self,\n *,\n transformer_dim: int,\n transformer: nn.Module,\n num_multimask_outputs: int = 3,\n activation: Type[nn.Module] = nn.GELU,\n iou_head_depth: int = 3,\n iou_head_hidden_dim: int = 256,\n ) -> None:\n \"\"\"\n Predicts masks given an image and prompt embeddings, using a\n transformer architecture.\n\n Arguments:\n transformer_dim (int): the channel dimension of the transformer\n transformer (nn.Module): the transformer used to predict masks\n num_multimask_outputs (int): the number of masks to predict\n when disambiguating masks\n activation (nn.Module): the type of activation to use when\n upscaling masks\n iou_head_depth (int): the depth of the MLP used to predict\n mask quality\n iou_head_hidden_dim (int): the hidden dimension of the MLP\n used to predict mask quality\n \"\"\"\n super().__init__()\n self.transformer_dim = transformer_dim\n self.transformer = transformer\n\n self.num_multimask_outputs = num_multimask_outputs\n\n self.iou_token = nn.Embedding(1, transformer_dim)\n self.num_mask_tokens = num_multimask_outputs + 1\n self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)\n\n self.output_upscaling = nn.Sequential(\n nn.ConvTranspose2d(\n transformer_dim, transformer_dim // 4, kernel_size=2, stride=2\n ),\n LayerNorm2d(transformer_dim // 4),\n activation(),\n nn.ConvTranspose2d(\n transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2\n ),\n activation(),\n )\n self.output_hypernetworks_mlps = nn.ModuleList(\n [\n MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3)\n for i in range(self.num_mask_tokens)\n ]\n )\n\n self.iou_prediction_head = MLP(\n transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth\n )\n\n def forward(\n self,\n image_embeddings: torch.Tensor,\n image_pe: torch.Tensor,\n sparse_prompt_embeddings: torch.Tensor,\n dense_prompt_embeddings: torch.Tensor,\n multimask_output: bool,\n ) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"\n Predict masks given image and prompt embeddings.\n\n Arguments:\n image_embeddings (torch.Tensor): the embeddings from the image encoder\n image_pe (torch.Tensor): positional encoding with the shape of image_embeddings\n sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes\n dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs\n multimask_output (bool): Whether to return multiple masks or a single\n mask.\n\n Returns:\n torch.Tensor: batched predicted masks\n torch.Tensor: batched predictions of mask quality\n \"\"\"\n masks, iou_pred = self.predict_masks(\n image_embeddings=image_embeddings,\n image_pe=image_pe,\n sparse_prompt_embeddings=sparse_prompt_embeddings,\n dense_prompt_embeddings=dense_prompt_embeddings,\n )\n\n # Select the correct mask or masks for output\n if multimask_output:\n mask_slice = slice(1, None)\n else:\n mask_slice = slice(0, 1)\n masks = masks[:, mask_slice, :, :]\n iou_pred = iou_pred[:, mask_slice]\n\n # Prepare output\n return masks, iou_pred\n\n def predict_masks(\n self,\n image_embeddings: torch.Tensor,\n image_pe: torch.Tensor,\n sparse_prompt_embeddings: torch.Tensor,\n dense_prompt_embeddings: torch.Tensor,\n ) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Predicts masks. See 'forward' for more details.\"\"\"\n # Concatenate output tokens\n output_tokens = torch.cat(\n [self.iou_token.weight, self.mask_tokens.weight], dim=0\n )\n output_tokens = output_tokens.unsqueeze(0).expand(\n sparse_prompt_embeddings.size(0), -1, -1\n )\n\n tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)\n\n # image_embeddings: [1, C, H, W], tokens: [B, N, C]\n # dense_prompt_embeddings: [B, C, H, W]\n # Expand per-image data in batch direction to be per-mask\n src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)\n src = src + dense_prompt_embeddings\n pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)\n b, c, h, w = src.shape\n\n # Run the transformer\n hs, src = self.transformer(src, pos_src, tokens)\n iou_token_out = hs[:, 0, :]\n mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :]\n\n # Upscale mask embeddings and predict masks using the mask tokens\n src = src.transpose(1, 2).view(b, c, h, w)\n upscaled_embedding = self.output_upscaling(src)\n hyper_in_list: List[torch.Tensor] = []\n for i in range(self.num_mask_tokens):\n hyper_in_list.append(\n self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])\n )\n hyper_in = torch.stack(hyper_in_list, dim=1)\n b, c, h, w = upscaled_embedding.shape\n masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(\n b, self.num_mask_tokens, h, w\n )\n\n # Generate mask quality predictions\n iou_pred = self.iou_prediction_head(iou_token_out)\n\n return masks, iou_pred\n\n\n def forward_modified_v3(\n self, \n image_embeddings : torch.Tensor, # [b, 256, 64, 64]\n image_pe : torch.Tensor, # [1, 256, 64, 64]\n sparse_prompt_embeddings: torch.Tensor, # [b, 1, 256]\n dense_prompt_embeddings : torch.Tensor, # [b, 256, 64, 64]\n ): # -> [b, 256, 256]\n b, _, _, _ = image_embeddings.shape\n \n # Concatenate output tokens\n output_tokens = torch.cat([self.iou_token.weight, self.mask_tokens.weight], dim=0) # [num_mask + 1, 256]\n output_tokens = output_tokens.unsqueeze(0).expand(b, -1, -1) # [b, num_mask + 1, 256]\n\n tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1) # [b, num_mask + 1 + 1, 256]\n\n # image_embeddings: [1, C, H, W], tokens: [B, N, C]\n # dense_prompt_embeddings: [B, C, H, W]\n # Expand per-image data in batch direction to be per-mask\n src = image_embeddings\n src = src + dense_prompt_embeddings # [b, 256, 64, 64]\n pos_src = torch.repeat_interleave(image_pe, b, dim=0) # [b, 256, 64, 64]\n _, c, h, w = src.shape \n\n # Run the transformer\n hs, src = self.transformer(src, pos_src, tokens)\n mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :]\n\n # Upscale mask embeddings and predict masks using the mask tokens\n src = src.transpose(1, 2).view(b, c, h, w)\n upscaled_embedding = self.output_upscaling(src)\n hyper_in_list: List[torch.Tensor] = []\n num_mask = 1\n for i in range(num_mask): # we only need the first mask\n hyper_in_list.append(self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :]))\n hyper_in = torch.stack(hyper_in_list, dim=1)\n b, c, h, w = upscaled_embedding.shape\n masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, h, w) # [b, 256, 256]\n\n return masks\n\n# Lightly adapted from\n# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa\nclass MLP(nn.Module):\n def __init__(\n self,\n input_dim: int,\n hidden_dim: int,\n output_dim: int,\n num_layers: int,\n sigmoid_output: bool = False,\n ) -> None:\n super().__init__()\n self.num_layers = num_layers\n h = [hidden_dim] * (num_layers - 1)\n self.layers = nn.ModuleList(\n nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])\n )\n self.sigmoid_output = sigmoid_output\n\n def forward(self, x):\n for i, layer in enumerate(self.layers):\n x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)\n if self.sigmoid_output:\n x = F.sigmoid(x)\n return x\n" }, { "path": "model/segment_anything/modeling/prompt_encoder.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Any, Optional, Tuple, Type\n\nimport numpy as np\nimport torch\nfrom torch import nn\n\nfrom .common import LayerNorm2d\n\n\nclass PromptEncoder(nn.Module):\n def __init__(\n self,\n embed_dim: int,\n image_embedding_size: Tuple[int, int],\n input_image_size: Tuple[int, int],\n mask_in_chans: int,\n activation: Type[nn.Module] = nn.GELU,\n ) -> None:\n \"\"\"\n Encodes prompts for input to SAM's mask decoder.\n\n Arguments:\n embed_dim (int): The prompts' embedding dimension\n image_embedding_size (tuple(int, int)): The spatial size of the\n image embedding, as (H, W).\n input_image_size (int): The padded size of the image as input\n to the image encoder, as (H, W).\n mask_in_chans (int): The number of hidden channels used for\n encoding input masks.\n activation (nn.Module): The activation to use when encoding\n input masks.\n \"\"\"\n super().__init__()\n self.embed_dim = embed_dim\n self.input_image_size = input_image_size\n self.image_embedding_size = image_embedding_size\n self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)\n\n self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners\n point_embeddings = [\n nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)\n ]\n self.point_embeddings = nn.ModuleList(point_embeddings)\n self.not_a_point_embed = nn.Embedding(1, embed_dim)\n\n self.mask_input_size = (\n 4 * image_embedding_size[0],\n 4 * image_embedding_size[1],\n )\n self.mask_downscaling = nn.Sequential(\n nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),\n LayerNorm2d(mask_in_chans // 4),\n activation(),\n nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),\n LayerNorm2d(mask_in_chans),\n activation(),\n nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1),\n )\n self.no_mask_embed = nn.Embedding(1, embed_dim)\n\n def get_dense_pe(self) -> torch.Tensor:\n \"\"\"\n Returns the positional encoding used to encode point prompts,\n applied to a dense set of points the shape of the image encoding.\n\n Returns:\n torch.Tensor: Positional encoding with shape\n 1x(embed_dim)x(embedding_h)x(embedding_w)\n \"\"\"\n return self.pe_layer(self.image_embedding_size).unsqueeze(0)\n\n def _embed_points(\n self,\n points: torch.Tensor,\n labels: torch.Tensor,\n pad: bool,\n ) -> torch.Tensor:\n \"\"\"Embeds point prompts.\"\"\"\n points = points + 0.5 # Shift to center of pixel\n if pad:\n padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)\n padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)\n points = torch.cat([points, padding_point], dim=1)\n labels = torch.cat([labels, padding_label], dim=1)\n point_embedding = self.pe_layer.forward_with_coords(\n points, self.input_image_size\n )\n point_embedding[labels == -1] = 0.0\n point_embedding[labels == -1] += self.not_a_point_embed.weight\n point_embedding[labels == 0] += self.point_embeddings[0].weight\n point_embedding[labels == 1] += self.point_embeddings[1].weight\n return point_embedding\n\n def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:\n \"\"\"Embeds box prompts.\"\"\"\n boxes = boxes + 0.5 # Shift to center of pixel\n coords = boxes.reshape(-1, 2, 2)\n corner_embedding = self.pe_layer.forward_with_coords(\n coords, self.input_image_size\n )\n corner_embedding[:, 0, :] += self.point_embeddings[2].weight\n corner_embedding[:, 1, :] += self.point_embeddings[3].weight\n return corner_embedding\n\n def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:\n \"\"\"Embeds mask inputs.\"\"\"\n mask_embedding = self.mask_downscaling(masks)\n return mask_embedding\n\n def _get_batch_size(\n self,\n points: Optional[Tuple[torch.Tensor, torch.Tensor]],\n boxes: Optional[torch.Tensor],\n masks: Optional[torch.Tensor],\n text_embeds: Optional[torch.Tensor],\n ) -> int:\n \"\"\"\n Gets the batch size of the output given the batch size of the input prompts.\n \"\"\"\n if points is not None:\n return points[0].shape[0]\n elif boxes is not None:\n return boxes.shape[0]\n elif masks is not None:\n return masks.shape[0]\n elif text_embeds is not None:\n return text_embeds.shape[0]\n else:\n return 1\n\n def _get_device(self) -> torch.device:\n return self.point_embeddings[0].weight.device\n\n def forward(\n self,\n points: Optional[Tuple[torch.Tensor, torch.Tensor]],\n boxes: Optional[torch.Tensor],\n masks: Optional[torch.Tensor],\n text_embeds: Optional[torch.Tensor],\n ) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"\n Embeds different types of prompts, returning both sparse and dense\n embeddings.\n\n Arguments:\n points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates\n and labels to embed.\n boxes (torch.Tensor or none): boxes to embed\n masks (torch.Tensor or none): masks to embed\n\n Returns:\n torch.Tensor: sparse embeddings for the points and boxes, with shape\n BxNx(embed_dim), where N is determined by the number of input points\n and boxes.\n torch.Tensor: dense embeddings for the masks, in the shape\n Bx(embed_dim)x(embed_H)x(embed_W)\n \"\"\"\n bs = self._get_batch_size(points, boxes, masks, text_embeds)\n sparse_embeddings = torch.empty(\n (bs, 0, self.embed_dim), device=self._get_device()\n )\n if points is not None:\n coords, labels = points\n point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))\n sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)\n if boxes is not None:\n box_embeddings = self._embed_boxes(boxes)\n sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)\n\n if text_embeds is not None:\n sparse_embeddings = torch.cat([sparse_embeddings, text_embeds], dim=1)\n\n if masks is not None:\n dense_embeddings = self._embed_masks(masks)\n else:\n dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(\n bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]\n )\n\n return sparse_embeddings, dense_embeddings\n\n\nclass PositionEmbeddingRandom(nn.Module):\n \"\"\"\n Positional encoding using random spatial frequencies.\n \"\"\"\n\n def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:\n super().__init__()\n if scale is None or scale <= 0.0:\n scale = 1.0\n self.register_buffer(\n \"positional_encoding_gaussian_matrix\",\n scale * torch.randn((2, num_pos_feats)),\n )\n\n def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:\n \"\"\"Positionally encode points that are normalized to [0,1].\"\"\"\n # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape\n coords = 2 * coords - 1\n\n if coords.dtype != self.positional_encoding_gaussian_matrix.dtype:\n coords = coords.to(self.positional_encoding_gaussian_matrix.dtype)\n\n coords = coords @ self.positional_encoding_gaussian_matrix\n coords = 2 * np.pi * coords\n # outputs d_1 x ... x d_n x C shape\n return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)\n\n def forward(self, size: Tuple[int, int]) -> torch.Tensor:\n \"\"\"Generate positional encoding for a grid of the specified size.\"\"\"\n h, w = size\n device: Any = self.positional_encoding_gaussian_matrix.device\n grid = torch.ones(\n (h, w), device=device, dtype=self.positional_encoding_gaussian_matrix.dtype\n )\n y_embed = grid.cumsum(dim=0) - 0.5\n x_embed = grid.cumsum(dim=1) - 0.5\n y_embed = y_embed / h\n x_embed = x_embed / w\n\n pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))\n return pe.permute(2, 0, 1) # C x H x W\n\n def forward_with_coords(\n self, coords_input: torch.Tensor, image_size: Tuple[int, int]\n ) -> torch.Tensor:\n \"\"\"Positionally encode points that are not normalized to [0,1].\"\"\"\n coords = coords_input.clone()\n coords[:, :, 0] = coords[:, :, 0] / image_size[1]\n coords[:, :, 1] = coords[:, :, 1] / image_size[0]\n return self._pe_encoding(coords.to(torch.float)) # B x N x C\n" }, { "path": "model/segment_anything/modeling/sam.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Any, Dict, List, Tuple\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\nfrom .image_encoder import ImageEncoderViT\nfrom .mask_decoder import MaskDecoder\nfrom .prompt_encoder import PromptEncoder\n\n\nclass Sam(nn.Module):\n mask_threshold: float = 0.0\n image_format: str = \"RGB\"\n\n def __init__(\n self,\n image_encoder: ImageEncoderViT,\n prompt_encoder: PromptEncoder,\n mask_decoder: MaskDecoder,\n pixel_mean: List[float] = [123.675, 116.28, 103.53],\n pixel_std: List[float] = [58.395, 57.12, 57.375],\n ) -> None:\n \"\"\"\n SAM predicts object masks from an image and input prompts.\n\n Arguments:\n image_encoder (ImageEncoderViT): The backbone used to encode the\n image into image embeddings that allow for efficient mask prediction.\n prompt_encoder (PromptEncoder): Encodes various types of input prompts.\n mask_decoder (MaskDecoder): Predicts masks from the image embeddings\n and encoded prompts.\n pixel_mean (list(float)): Mean values for normalizing pixels in the input image.\n pixel_std (list(float)): Std values for normalizing pixels in the input image.\n \"\"\"\n super().__init__()\n self.image_encoder = image_encoder\n self.prompt_encoder = prompt_encoder\n self.mask_decoder = mask_decoder\n self.register_buffer(\n \"pixel_mean\", torch.Tensor(pixel_mean).view(-1, 1, 1), False\n )\n self.register_buffer(\"pixel_std\", torch.Tensor(pixel_std).view(-1, 1, 1), False)\n\n @property\n def device(self) -> Any:\n return self.pixel_mean.device\n\n @torch.no_grad()\n def forward(\n self,\n batched_input: List[Dict[str, Any]],\n multimask_output: bool,\n ) -> List[Dict[str, torch.Tensor]]:\n \"\"\"\n Predicts masks end-to-end from provided images and prompts.\n If prompts are not known in advance, using SamPredictor is\n recommended over calling the model directly.\n\n Arguments:\n batched_input (list(dict)): A list over input images, each a\n dictionary with the following keys. A prompt key can be\n excluded if it is not present.\n 'image': The image as a torch tensor in 3xHxW format,\n already transformed for input to the model.\n 'original_size': (tuple(int, int)) The original size of\n the image before transformation, as (H, W).\n 'point_coords': (torch.Tensor) Batched point prompts for\n this image, with shape BxNx2. Already transformed to the\n input frame of the model.\n 'point_labels': (torch.Tensor) Batched labels for point prompts,\n with shape BxN.\n 'boxes': (torch.Tensor) Batched box inputs, with shape Bx4.\n Already transformed to the input frame of the model.\n 'mask_inputs': (torch.Tensor) Batched mask inputs to the model,\n in the form Bx1xHxW.\n multimask_output (bool): Whether the model should predict multiple\n disambiguating masks, or return a single mask.\n\n Returns:\n (list(dict)): A list over input images, where each element is\n as dictionary with the following keys.\n 'masks': (torch.Tensor) Batched binary mask predictions,\n with shape BxCxHxW, where B is the number of input prompts,\n C is determined by multimask_output, and (H, W) is the\n original size of the image.\n 'iou_predictions': (torch.Tensor) The model's predictions\n of mask quality, in shape BxC.\n 'low_res_logits': (torch.Tensor) Low resolution logits with\n shape BxCxHxW, where H=W=256. Can be passed as mask input\n to subsequent iterations of prediction.\n \"\"\"\n input_images = torch.stack(\n [self.preprocess(x[\"image\"]) for x in batched_input], dim=0\n )\n image_embeddings = self.image_encoder(input_images)\n\n outputs = []\n for image_record, curr_embedding in zip(batched_input, image_embeddings):\n if \"point_coords\" in image_record:\n points = (image_record[\"point_coords\"], image_record[\"point_labels\"])\n else:\n points = None\n sparse_embeddings, dense_embeddings = self.prompt_encoder(\n points=points,\n boxes=image_record.get(\"boxes\", None),\n masks=image_record.get(\"mask_inputs\", None),\n )\n low_res_masks, iou_predictions = self.mask_decoder(\n image_embeddings=curr_embedding.unsqueeze(0),\n image_pe=self.prompt_encoder.get_dense_pe(),\n sparse_prompt_embeddings=sparse_embeddings,\n dense_prompt_embeddings=dense_embeddings,\n multimask_output=multimask_output,\n )\n masks = self.postprocess_masks(\n low_res_masks,\n input_size=image_record[\"image\"].shape[-2:],\n original_size=image_record[\"original_size\"],\n )\n masks = masks > self.mask_threshold\n outputs.append(\n {\n \"masks\": masks,\n \"iou_predictions\": iou_predictions,\n \"low_res_logits\": low_res_masks,\n }\n )\n return outputs\n\n def postprocess_masks(\n self,\n masks: torch.Tensor,\n input_size: Tuple[int, ...],\n original_size: Tuple[int, ...],\n ) -> torch.Tensor:\n \"\"\"\n Remove padding and upscale masks to the original image size.\n\n Arguments:\n masks (torch.Tensor): Batched masks from the mask_decoder,\n in BxCxHxW format.\n input_size (tuple(int, int)): The size of the image input to the\n model, in (H, W) format. Used to remove padding.\n original_size (tuple(int, int)): The original size of the image\n before resizing for input to the model, in (H, W) format.\n\n Returns:\n (torch.Tensor): Batched masks in BxCxHxW format, where (H, W)\n is given by original_size.\n \"\"\"\n\n dtype = masks.dtype\n\n masks = F.interpolate(\n masks.float(),\n (self.image_encoder.img_size, self.image_encoder.img_size),\n mode=\"bilinear\",\n align_corners=False,\n )\n # masks = masks.to(dtype)\n masks = masks[..., : input_size[0], : input_size[1]]\n masks = F.interpolate(\n masks, original_size, mode=\"bilinear\", align_corners=False\n )\n return masks\n\n def preprocess(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"Normalize pixel values and pad to a square input.\"\"\"\n # Normalize colors\n x = (x - self.pixel_mean) / self.pixel_std\n\n # Pad\n h, w = x.shape[-2:]\n padh = self.image_encoder.img_size - h\n padw = self.image_encoder.img_size - w\n x = F.pad(x, (0, padw, 0, padh))\n return x\n" }, { "path": "model/segment_anything/modeling/transformer.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nfrom typing import Tuple, Type\n\nimport torch\nfrom torch import Tensor, nn\n\nfrom .common import MLPBlock\n\n\nclass TwoWayTransformer(nn.Module):\n def __init__(\n self,\n depth: int,\n embedding_dim: int,\n num_heads: int,\n mlp_dim: int,\n activation: Type[nn.Module] = nn.ReLU,\n attention_downsample_rate: int = 2,\n ) -> None:\n \"\"\"\n A transformer decoder that attends to an input image using\n queries whose positional embedding is supplied.\n\n Args:\n depth (int): number of layers in the transformer\n embedding_dim (int): the channel dimension for the input embeddings\n num_heads (int): the number of heads for multihead attention. Must\n divide embedding_dim\n mlp_dim (int): the channel dimension internal to the MLP block\n activation (nn.Module): the activation to use in the MLP block\n \"\"\"\n super().__init__()\n self.depth = depth\n self.embedding_dim = embedding_dim\n self.num_heads = num_heads\n self.mlp_dim = mlp_dim\n self.layers = nn.ModuleList()\n\n for i in range(depth):\n self.layers.append(\n TwoWayAttentionBlock(\n embedding_dim=embedding_dim,\n num_heads=num_heads,\n mlp_dim=mlp_dim,\n activation=activation,\n attention_downsample_rate=attention_downsample_rate,\n skip_first_layer_pe=(i == 0),\n )\n )\n\n self.final_attn_token_to_image = Attention(\n embedding_dim, num_heads, downsample_rate=attention_downsample_rate\n )\n self.norm_final_attn = nn.LayerNorm(embedding_dim)\n\n def forward(\n self,\n image_embedding: Tensor,\n image_pe: Tensor,\n point_embedding: Tensor,\n ) -> Tuple[Tensor, Tensor]:\n \"\"\"\n Args:\n image_embedding (torch.Tensor): image to attend to. Should be shape\n B x embedding_dim x h x w for any h and w.\n image_pe (torch.Tensor): the positional encoding to add to the image. Must\n have the same shape as image_embedding.\n point_embedding (torch.Tensor): the embedding to add to the query points.\n Must have shape B x N_points x embedding_dim for any N_points.\n\n Returns:\n torch.Tensor: the processed point_embedding\n torch.Tensor: the processed image_embedding\n \"\"\"\n # BxCxHxW -> BxHWxC == B x N_image_tokens x C\n bs, c, h, w = image_embedding.shape\n image_embedding = image_embedding.flatten(2).permute(0, 2, 1)\n image_pe = image_pe.flatten(2).permute(0, 2, 1)\n\n # Prepare queries\n queries = point_embedding\n keys = image_embedding\n\n # Apply transformer blocks and final layernorm\n for layer in self.layers:\n queries, keys = layer(\n queries=queries,\n keys=keys,\n query_pe=point_embedding,\n key_pe=image_pe,\n )\n\n # Apply the final attention layer from the points to the image\n q = queries + point_embedding\n k = keys + image_pe\n attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)\n queries = queries + attn_out\n queries = self.norm_final_attn(queries)\n\n return queries, keys\n\n\nclass TwoWayAttentionBlock(nn.Module):\n def __init__(\n self,\n embedding_dim: int,\n num_heads: int,\n mlp_dim: int = 2048,\n activation: Type[nn.Module] = nn.ReLU,\n attention_downsample_rate: int = 2,\n skip_first_layer_pe: bool = False,\n ) -> None:\n \"\"\"\n A transformer block with four layers: (1) self-attention of sparse\n inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp\n block on sparse inputs, and (4) cross attention of dense inputs to sparse\n inputs.\n\n Arguments:\n embedding_dim (int): the channel dimension of the embeddings\n num_heads (int): the number of heads in the attention layers\n mlp_dim (int): the hidden dimension of the mlp block\n activation (nn.Module): the activation of the mlp block\n skip_first_layer_pe (bool): skip the PE on the first layer\n \"\"\"\n super().__init__()\n self.self_attn = Attention(embedding_dim, num_heads)\n self.norm1 = nn.LayerNorm(embedding_dim)\n\n self.cross_attn_token_to_image = Attention(\n embedding_dim, num_heads, downsample_rate=attention_downsample_rate\n )\n self.norm2 = nn.LayerNorm(embedding_dim)\n\n self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)\n self.norm3 = nn.LayerNorm(embedding_dim)\n\n self.norm4 = nn.LayerNorm(embedding_dim)\n self.cross_attn_image_to_token = Attention(\n embedding_dim, num_heads, downsample_rate=attention_downsample_rate\n )\n\n self.skip_first_layer_pe = skip_first_layer_pe\n\n def forward(\n self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor\n ) -> Tuple[Tensor, Tensor]:\n # Self attention block\n if self.skip_first_layer_pe:\n queries = self.self_attn(q=queries, k=queries, v=queries)\n else:\n q = queries + query_pe\n attn_out = self.self_attn(q=q, k=q, v=queries)\n queries = queries + attn_out\n queries = self.norm1(queries)\n\n # Cross attention block, tokens attending to image embedding\n q = queries + query_pe\n k = keys + key_pe\n attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)\n queries = queries + attn_out\n queries = self.norm2(queries)\n\n # MLP block\n mlp_out = self.mlp(queries)\n queries = queries + mlp_out\n queries = self.norm3(queries)\n\n # Cross attention block, image embedding attending to tokens\n q = queries + query_pe\n k = keys + key_pe\n attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)\n keys = keys + attn_out\n keys = self.norm4(keys)\n\n return queries, keys\n\n\nclass Attention(nn.Module):\n \"\"\"\n An attention layer that allows for downscaling the size of the embedding\n after projection to queries, keys, and values.\n \"\"\"\n\n def __init__(\n self,\n embedding_dim: int,\n num_heads: int,\n downsample_rate: int = 1,\n ) -> None:\n super().__init__()\n self.embedding_dim = embedding_dim\n self.internal_dim = embedding_dim // downsample_rate\n self.num_heads = num_heads\n assert (\n self.internal_dim % num_heads == 0\n ), \"num_heads must divide embedding_dim.\"\n\n self.q_proj = nn.Linear(embedding_dim, self.internal_dim)\n self.k_proj = nn.Linear(embedding_dim, self.internal_dim)\n self.v_proj = nn.Linear(embedding_dim, self.internal_dim)\n self.out_proj = nn.Linear(self.internal_dim, embedding_dim)\n\n def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:\n b, n, c = x.shape\n x = x.reshape(b, n, num_heads, c // num_heads)\n return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head\n\n def _recombine_heads(self, x: Tensor) -> Tensor:\n b, n_heads, n_tokens, c_per_head = x.shape\n x = x.transpose(1, 2)\n return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C\n\n def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:\n # Input projections\n q = self.q_proj(q)\n k = self.k_proj(k)\n v = self.v_proj(v)\n\n # Separate into heads\n q = self._separate_heads(q, self.num_heads)\n k = self._separate_heads(k, self.num_heads)\n v = self._separate_heads(v, self.num_heads)\n\n # Attention\n _, _, _, c_per_head = q.shape\n attn = q @ k.permute(0, 1, 3, 2) # B x N_heads x N_tokens x N_tokens\n attn = attn / math.sqrt(c_per_head)\n attn = torch.softmax(attn, dim=-1)\n\n # Get output\n out = attn @ v\n out = self._recombine_heads(out)\n out = self.out_proj(out)\n\n return out\n" }, { "path": "model/segment_anything/predictor.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Optional, Tuple\n\nimport numpy as np\nimport torch\n\nfrom .modeling import Sam\nfrom .utils.transforms import ResizeLongestSide\n\n\nclass SamPredictor:\n def __init__(\n self,\n sam_model: Sam,\n ) -> None:\n \"\"\"\n Uses SAM to calculate the image embedding for an image, and then\n allow repeated, efficient mask prediction given prompts.\n\n Arguments:\n sam_model (Sam): The model to use for mask prediction.\n \"\"\"\n super().__init__()\n self.model = sam_model\n self.transform = ResizeLongestSide(sam_model.image_encoder.img_size)\n self.reset_image()\n\n def set_image(\n self,\n image: np.ndarray,\n image_format: str = \"RGB\",\n ) -> None:\n \"\"\"\n Calculates the image embeddings for the provided image, allowing\n masks to be predicted with the 'predict' method.\n\n Arguments:\n image (np.ndarray): The image for calculating masks. Expects an\n image in HWC uint8 format, with pixel values in [0, 255].\n image_format (str): The color format of the image, in ['RGB', 'BGR'].\n \"\"\"\n assert image_format in [\n \"RGB\",\n \"BGR\",\n ], f\"image_format must be in ['RGB', 'BGR'], is {image_format}.\"\n if image_format != self.model.image_format:\n image = image[..., ::-1]\n\n # Transform the image to the form expected by the model\n input_image = self.transform.apply_image(image)\n input_image_torch = torch.as_tensor(input_image, device=self.device)\n input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[\n None, :, :, :\n ]\n\n self.set_torch_image(input_image_torch, image.shape[:2])\n\n @torch.no_grad()\n def set_torch_image(\n self,\n transformed_image: torch.Tensor,\n original_image_size: Tuple[int, ...],\n ) -> None:\n \"\"\"\n Calculates the image embeddings for the provided image, allowing\n masks to be predicted with the 'predict' method. Expects the input\n image to be already transformed to the format expected by the model.\n\n Arguments:\n transformed_image (torch.Tensor): The input image, with shape\n 1x3xHxW, which has been transformed with ResizeLongestSide.\n original_image_size (tuple(int, int)): The size of the image\n before transformation, in (H, W) format.\n \"\"\"\n assert (\n len(transformed_image.shape) == 4\n and transformed_image.shape[1] == 3\n and max(*transformed_image.shape[2:]) == self.model.image_encoder.img_size\n ), f\"set_torch_image input must be BCHW with long side {self.model.image_encoder.img_size}.\"\n self.reset_image()\n\n self.original_size = original_image_size\n self.input_size = tuple(transformed_image.shape[-2:])\n input_image = self.model.preprocess(transformed_image)\n self.features = self.model.image_encoder(input_image)\n self.is_image_set = True\n\n def predict(\n self,\n point_coords: Optional[np.ndarray] = None,\n point_labels: Optional[np.ndarray] = None,\n box: Optional[np.ndarray] = None,\n mask_input: Optional[np.ndarray] = None,\n multimask_output: bool = True,\n return_logits: bool = False,\n ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n \"\"\"\n Predict masks for the given input prompts, using the currently set image.\n\n Arguments:\n point_coords (np.ndarray or None): A Nx2 array of point prompts to the\n model. Each point is in (X,Y) in pixels.\n point_labels (np.ndarray or None): A length N array of labels for the\n point prompts. 1 indicates a foreground point and 0 indicates a\n background point.\n box (np.ndarray or None): A length 4 array given a box prompt to the\n model, in XYXY format.\n mask_input (np.ndarray): A low resolution mask input to the model, typically\n coming from a previous prediction iteration. Has form 1xHxW, where\n for SAM, H=W=256.\n multimask_output (bool): If true, the model will return three masks.\n For ambiguous input prompts (such as a single click), this will often\n produce better masks than a single prediction. If only a single\n mask is needed, the model's predicted quality score can be used\n to select the best mask. For non-ambiguous prompts, such as multiple\n input prompts, multimask_output=False can give better results.\n return_logits (bool): If true, returns un-thresholded masks logits\n instead of a binary mask.\n\n Returns:\n (np.ndarray): The output masks in CxHxW format, where C is the\n number of masks, and (H, W) is the original image size.\n (np.ndarray): An array of length C containing the model's\n predictions for the quality of each mask.\n (np.ndarray): An array of shape CxHxW, where C is the number\n of masks and H=W=256. These low resolution logits can be passed to\n a subsequent iteration as mask input.\n \"\"\"\n if not self.is_image_set:\n raise RuntimeError(\n \"An image must be set with .set_image(...) before mask prediction.\"\n )\n\n # Transform input prompts\n coords_torch, labels_torch, box_torch, mask_input_torch = None, None, None, None\n if point_coords is not None:\n assert (\n point_labels is not None\n ), \"point_labels must be supplied if point_coords is supplied.\"\n point_coords = self.transform.apply_coords(point_coords, self.original_size)\n coords_torch = torch.as_tensor(\n point_coords, dtype=torch.float, device=self.device\n )\n labels_torch = torch.as_tensor(\n point_labels, dtype=torch.int, device=self.device\n )\n coords_torch, labels_torch = coords_torch[None, :, :], labels_torch[None, :]\n if box is not None:\n box = self.transform.apply_boxes(box, self.original_size)\n box_torch = torch.as_tensor(box, dtype=torch.float, device=self.device)\n box_torch = box_torch[None, :]\n if mask_input is not None:\n mask_input_torch = torch.as_tensor(\n mask_input, dtype=torch.float, device=self.device\n )\n mask_input_torch = mask_input_torch[None, :, :, :]\n\n masks, iou_predictions, low_res_masks = self.predict_torch(\n coords_torch,\n labels_torch,\n box_torch,\n mask_input_torch,\n multimask_output,\n return_logits=return_logits,\n )\n\n masks_np = masks[0].detach().cpu().numpy()\n iou_predictions_np = iou_predictions[0].detach().cpu().numpy()\n low_res_masks_np = low_res_masks[0].detach().cpu().numpy()\n return masks_np, iou_predictions_np, low_res_masks_np\n\n @torch.no_grad()\n def predict_torch(\n self,\n point_coords: Optional[torch.Tensor],\n point_labels: Optional[torch.Tensor],\n boxes: Optional[torch.Tensor] = None,\n mask_input: Optional[torch.Tensor] = None,\n multimask_output: bool = True,\n return_logits: bool = False,\n ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"\n Predict masks for the given input prompts, using the currently set image.\n Input prompts are batched torch tensors and are expected to already be\n transformed to the input frame using ResizeLongestSide.\n\n Arguments:\n point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the\n model. Each point is in (X,Y) in pixels.\n point_labels (torch.Tensor or None): A BxN array of labels for the\n point prompts. 1 indicates a foreground point and 0 indicates a\n background point.\n boxes (np.ndarray or None): A Bx4 array given a box prompt to the\n model, in XYXY format.\n mask_input (np.ndarray): A low resolution mask input to the model, typically\n coming from a previous prediction iteration. Has form Bx1xHxW, where\n for SAM, H=W=256. Masks returned by a previous iteration of the\n predict method do not need further transformation.\n multimask_output (bool): If true, the model will return three masks.\n For ambiguous input prompts (such as a single click), this will often\n produce better masks than a single prediction. If only a single\n mask is needed, the model's predicted quality score can be used\n to select the best mask. For non-ambiguous prompts, such as multiple\n input prompts, multimask_output=False can give better results.\n return_logits (bool): If true, returns un-thresholded masks logits\n instead of a binary mask.\n\n Returns:\n (torch.Tensor): The output masks in BxCxHxW format, where C is the\n number of masks, and (H, W) is the original image size.\n (torch.Tensor): An array of shape BxC containing the model's\n predictions for the quality of each mask.\n (torch.Tensor): An array of shape BxCxHxW, where C is the number\n of masks and H=W=256. These low res logits can be passed to\n a subsequent iteration as mask input.\n \"\"\"\n if not self.is_image_set:\n raise RuntimeError(\n \"An image must be set with .set_image(...) before mask prediction.\"\n )\n\n if point_coords is not None:\n points = (point_coords, point_labels)\n else:\n points = None\n\n # Embed prompts\n sparse_embeddings, dense_embeddings = self.model.prompt_encoder(\n points=points,\n boxes=boxes,\n masks=mask_input,\n )\n\n # Predict masks\n low_res_masks, iou_predictions = self.model.mask_decoder(\n image_embeddings=self.features,\n image_pe=self.model.prompt_encoder.get_dense_pe(),\n sparse_prompt_embeddings=sparse_embeddings,\n dense_prompt_embeddings=dense_embeddings,\n multimask_output=multimask_output,\n )\n\n # Upscale the masks to the original image resolution\n masks = self.model.postprocess_masks(\n low_res_masks, self.input_size, self.original_size\n )\n\n if not return_logits:\n masks = masks > self.model.mask_threshold\n\n return masks, iou_predictions, low_res_masks\n\n def get_image_embedding(self) -> torch.Tensor:\n \"\"\"\n Returns the image embeddings for the currently set image, with\n shape 1xCxHxW, where C is the embedding dimension and (H,W) are\n the embedding spatial dimension of SAM (typically C=256, H=W=64).\n \"\"\"\n if not self.is_image_set:\n raise RuntimeError(\n \"An image must be set with .set_image(...) to generate an embedding.\"\n )\n assert (\n self.features is not None\n ), \"Features must exist if an image has been set.\"\n return self.features\n\n @property\n def device(self) -> torch.device:\n return self.model.device\n\n def reset_image(self) -> None:\n \"\"\"Resets the currently set image.\"\"\"\n self.is_image_set = False\n self.features = None\n self.orig_h = None\n self.orig_w = None\n self.input_h = None\n self.input_w = None\n" }, { "path": "model/segment_anything/utils/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "model/segment_anything/utils/amg.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nfrom copy import deepcopy\nfrom itertools import product\nfrom typing import Any, Dict, Generator, ItemsView, List, Tuple\n\nimport numpy as np\nimport torch\n\n\nclass MaskData:\n \"\"\"\n A structure for storing masks and their related data in batched format.\n Implements basic filtering and concatenation.\n \"\"\"\n\n def __init__(self, **kwargs) -> None:\n for v in kwargs.values():\n assert isinstance(\n v, (list, np.ndarray, torch.Tensor)\n ), \"MaskData only supports list, numpy arrays, and torch tensors.\"\n self._stats = dict(**kwargs)\n\n def __setitem__(self, key: str, item: Any) -> None:\n assert isinstance(\n item, (list, np.ndarray, torch.Tensor)\n ), \"MaskData only supports list, numpy arrays, and torch tensors.\"\n self._stats[key] = item\n\n def __delitem__(self, key: str) -> None:\n del self._stats[key]\n\n def __getitem__(self, key: str) -> Any:\n return self._stats[key]\n\n def items(self) -> ItemsView[str, Any]:\n return self._stats.items()\n\n def filter(self, keep: torch.Tensor) -> None:\n for k, v in self._stats.items():\n if v is None:\n self._stats[k] = None\n elif isinstance(v, torch.Tensor):\n self._stats[k] = v[torch.as_tensor(keep, device=v.device)]\n elif isinstance(v, np.ndarray):\n self._stats[k] = v[keep.detach().cpu().numpy()]\n elif isinstance(v, list) and keep.dtype == torch.bool:\n self._stats[k] = [a for i, a in enumerate(v) if keep[i]]\n elif isinstance(v, list):\n self._stats[k] = [v[i] for i in keep]\n else:\n raise TypeError(f\"MaskData key {k} has an unsupported type {type(v)}.\")\n\n def cat(self, new_stats: \"MaskData\") -> None:\n for k, v in new_stats.items():\n if k not in self._stats or self._stats[k] is None:\n self._stats[k] = deepcopy(v)\n elif isinstance(v, torch.Tensor):\n self._stats[k] = torch.cat([self._stats[k], v], dim=0)\n elif isinstance(v, np.ndarray):\n self._stats[k] = np.concatenate([self._stats[k], v], axis=0)\n elif isinstance(v, list):\n self._stats[k] = self._stats[k] + deepcopy(v)\n else:\n raise TypeError(f\"MaskData key {k} has an unsupported type {type(v)}.\")\n\n def to_numpy(self) -> None:\n for k, v in self._stats.items():\n if isinstance(v, torch.Tensor):\n self._stats[k] = v.detach().cpu().numpy()\n\n\ndef is_box_near_crop_edge(\n boxes: torch.Tensor, crop_box: List[int], orig_box: List[int], atol: float = 20.0\n) -> torch.Tensor:\n \"\"\"Filter masks at the edge of a crop, but not at the edge of the original image.\"\"\"\n crop_box_torch = torch.as_tensor(crop_box, dtype=torch.float, device=boxes.device)\n orig_box_torch = torch.as_tensor(orig_box, dtype=torch.float, device=boxes.device)\n boxes = uncrop_boxes_xyxy(boxes, crop_box).float()\n near_crop_edge = torch.isclose(boxes, crop_box_torch[None, :], atol=atol, rtol=0)\n near_image_edge = torch.isclose(boxes, orig_box_torch[None, :], atol=atol, rtol=0)\n near_crop_edge = torch.logical_and(near_crop_edge, ~near_image_edge)\n return torch.any(near_crop_edge, dim=1)\n\n\ndef box_xyxy_to_xywh(box_xyxy: torch.Tensor) -> torch.Tensor:\n box_xywh = deepcopy(box_xyxy)\n box_xywh[2] = box_xywh[2] - box_xywh[0]\n box_xywh[3] = box_xywh[3] - box_xywh[1]\n return box_xywh\n\n\ndef batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]:\n assert len(args) > 0 and all(\n len(a) == len(args[0]) for a in args\n ), \"Batched iteration must have inputs of all the same size.\"\n n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0)\n for b in range(n_batches):\n yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args]\n\n\ndef mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]:\n \"\"\"\n Encodes masks to an uncompressed RLE, in the format expected by\n pycoco tools.\n \"\"\"\n # Put in fortran order and flatten h,w\n b, h, w = tensor.shape\n tensor = tensor.permute(0, 2, 1).flatten(1)\n\n # Compute change indices\n diff = tensor[:, 1:] ^ tensor[:, :-1]\n change_indices = diff.nonzero()\n\n # Encode run length\n out = []\n for i in range(b):\n cur_idxs = change_indices[change_indices[:, 0] == i, 1]\n cur_idxs = torch.cat(\n [\n torch.tensor([0], dtype=cur_idxs.dtype, device=cur_idxs.device),\n cur_idxs + 1,\n torch.tensor([h * w], dtype=cur_idxs.dtype, device=cur_idxs.device),\n ]\n )\n btw_idxs = cur_idxs[1:] - cur_idxs[:-1]\n counts = [] if tensor[i, 0] == 0 else [0]\n counts.extend(btw_idxs.detach().cpu().tolist())\n out.append({\"size\": [h, w], \"counts\": counts})\n return out\n\n\ndef rle_to_mask(rle: Dict[str, Any]) -> np.ndarray:\n \"\"\"Compute a binary mask from an uncompressed RLE.\"\"\"\n h, w = rle[\"size\"]\n mask = np.empty(h * w, dtype=bool)\n idx = 0\n parity = False\n for count in rle[\"counts\"]:\n mask[idx : idx + count] = parity\n idx += count\n parity ^= True\n mask = mask.reshape(w, h)\n return mask.transpose() # Put in C order\n\n\ndef area_from_rle(rle: Dict[str, Any]) -> int:\n return sum(rle[\"counts\"][1::2])\n\n\ndef calculate_stability_score(\n masks: torch.Tensor, mask_threshold: float, threshold_offset: float\n) -> torch.Tensor:\n \"\"\"\n Computes the stability score for a batch of masks. The stability\n score is the IoU between the binary masks obtained by thresholding\n the predicted mask logits at high and low values.\n \"\"\"\n # One mask is always contained inside the other.\n # Save memory by preventing unnecessary cast to torch.int64\n intersections = (\n (masks > (mask_threshold + threshold_offset))\n .sum(-1, dtype=torch.int16)\n .sum(-1, dtype=torch.int32)\n )\n unions = (\n (masks > (mask_threshold - threshold_offset))\n .sum(-1, dtype=torch.int16)\n .sum(-1, dtype=torch.int32)\n )\n return intersections / unions\n\n\ndef build_point_grid(n_per_side: int) -> np.ndarray:\n \"\"\"Generates a 2D grid of points evenly spaced in [0,1]x[0,1].\"\"\"\n offset = 1 / (2 * n_per_side)\n points_one_side = np.linspace(offset, 1 - offset, n_per_side)\n points_x = np.tile(points_one_side[None, :], (n_per_side, 1))\n points_y = np.tile(points_one_side[:, None], (1, n_per_side))\n points = np.stack([points_x, points_y], axis=-1).reshape(-1, 2)\n return points\n\n\ndef build_all_layer_point_grids(\n n_per_side: int, n_layers: int, scale_per_layer: int\n) -> List[np.ndarray]:\n \"\"\"Generates point grids for all crop layers.\"\"\"\n points_by_layer = []\n for i in range(n_layers + 1):\n n_points = int(n_per_side / (scale_per_layer**i))\n points_by_layer.append(build_point_grid(n_points))\n return points_by_layer\n\n\ndef generate_crop_boxes(\n im_size: Tuple[int, ...], n_layers: int, overlap_ratio: float\n) -> Tuple[List[List[int]], List[int]]:\n \"\"\"\n Generates a list of crop boxes of different sizes. Each layer\n has (2**i)**2 boxes for the ith layer.\n \"\"\"\n crop_boxes, layer_idxs = [], []\n im_h, im_w = im_size\n short_side = min(im_h, im_w)\n\n # Original image\n crop_boxes.append([0, 0, im_w, im_h])\n layer_idxs.append(0)\n\n def crop_len(orig_len, n_crops, overlap):\n return int(math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops))\n\n for i_layer in range(n_layers):\n n_crops_per_side = 2 ** (i_layer + 1)\n overlap = int(overlap_ratio * short_side * (2 / n_crops_per_side))\n\n crop_w = crop_len(im_w, n_crops_per_side, overlap)\n crop_h = crop_len(im_h, n_crops_per_side, overlap)\n\n crop_box_x0 = [int((crop_w - overlap) * i) for i in range(n_crops_per_side)]\n crop_box_y0 = [int((crop_h - overlap) * i) for i in range(n_crops_per_side)]\n\n # Crops in XYWH format\n for x0, y0 in product(crop_box_x0, crop_box_y0):\n box = [x0, y0, min(x0 + crop_w, im_w), min(y0 + crop_h, im_h)]\n crop_boxes.append(box)\n layer_idxs.append(i_layer + 1)\n\n return crop_boxes, layer_idxs\n\n\ndef uncrop_boxes_xyxy(boxes: torch.Tensor, crop_box: List[int]) -> torch.Tensor:\n x0, y0, _, _ = crop_box\n offset = torch.tensor([[x0, y0, x0, y0]], device=boxes.device)\n # Check if boxes has a channel dimension\n if len(boxes.shape) == 3:\n offset = offset.unsqueeze(1)\n return boxes + offset\n\n\ndef uncrop_points(points: torch.Tensor, crop_box: List[int]) -> torch.Tensor:\n x0, y0, _, _ = crop_box\n offset = torch.tensor([[x0, y0]], device=points.device)\n # Check if points has a channel dimension\n if len(points.shape) == 3:\n offset = offset.unsqueeze(1)\n return points + offset\n\n\ndef uncrop_masks(\n masks: torch.Tensor, crop_box: List[int], orig_h: int, orig_w: int\n) -> torch.Tensor:\n x0, y0, x1, y1 = crop_box\n if x0 == 0 and y0 == 0 and x1 == orig_w and y1 == orig_h:\n return masks\n # Coordinate transform masks\n pad_x, pad_y = orig_w - (x1 - x0), orig_h - (y1 - y0)\n pad = (x0, pad_x - x0, y0, pad_y - y0)\n return torch.nn.functional.pad(masks, pad, value=0)\n\n\ndef remove_small_regions(\n mask: np.ndarray, area_thresh: float, mode: str\n) -> Tuple[np.ndarray, bool]:\n \"\"\"\n Removes small disconnected regions and holes in a mask. Returns the\n mask and an indicator of if the mask has been modified.\n \"\"\"\n import cv2 # type: ignore\n\n assert mode in [\"holes\", \"islands\"]\n correct_holes = mode == \"holes\"\n working_mask = (correct_holes ^ mask).astype(np.uint8)\n n_labels, regions, stats, _ = cv2.connectedComponentsWithStats(working_mask, 8)\n sizes = stats[:, -1][1:] # Row 0 is background label\n small_regions = [i + 1 for i, s in enumerate(sizes) if s < area_thresh]\n if len(small_regions) == 0:\n return mask, False\n fill_labels = [0] + small_regions\n if not correct_holes:\n fill_labels = [i for i in range(n_labels) if i not in fill_labels]\n # If every region is below threshold, keep largest\n if len(fill_labels) == 0:\n fill_labels = [int(np.argmax(sizes)) + 1]\n mask = np.isin(regions, fill_labels)\n return mask, True\n\n\ndef coco_encode_rle(uncompressed_rle: Dict[str, Any]) -> Dict[str, Any]:\n from pycocotools import mask as mask_utils # type: ignore\n\n h, w = uncompressed_rle[\"size\"]\n rle = mask_utils.frPyObjects(uncompressed_rle, h, w)\n rle[\"counts\"] = rle[\"counts\"].decode(\"utf-8\") # Necessary to serialize with json\n return rle\n\n\ndef batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Calculates boxes in XYXY format around masks. Return [0,0,0,0] for\n an empty mask. For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4.\n \"\"\"\n # torch.max below raises an error on empty inputs, just skip in this case\n if torch.numel(masks) == 0:\n return torch.zeros(*masks.shape[:-2], 4, device=masks.device)\n\n # Normalize shape to CxHxW\n shape = masks.shape\n h, w = shape[-2:]\n if len(shape) > 2:\n masks = masks.flatten(0, -3)\n else:\n masks = masks.unsqueeze(0)\n\n # Get top and bottom edges\n in_height, _ = torch.max(masks, dim=-1)\n in_height_coords = in_height * torch.arange(h, device=in_height.device)[None, :]\n bottom_edges, _ = torch.max(in_height_coords, dim=-1)\n in_height_coords = in_height_coords + h * (~in_height)\n top_edges, _ = torch.min(in_height_coords, dim=-1)\n\n # Get left and right edges\n in_width, _ = torch.max(masks, dim=-2)\n in_width_coords = in_width * torch.arange(w, device=in_width.device)[None, :]\n right_edges, _ = torch.max(in_width_coords, dim=-1)\n in_width_coords = in_width_coords + w * (~in_width)\n left_edges, _ = torch.min(in_width_coords, dim=-1)\n\n # If the mask is empty the right edge will be to the left of the left edge.\n # Replace these boxes with [0, 0, 0, 0]\n empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges)\n out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1)\n out = out * (~empty_filter).unsqueeze(-1)\n\n # Return to original shape\n if len(shape) > 2:\n out = out.reshape(*shape[:-2], 4)\n else:\n out = out[0]\n\n return out\n" }, { "path": "model/segment_anything/utils/onnx.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Tuple\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\n\nfrom ..modeling import Sam\nfrom .amg import calculate_stability_score\n\n\nclass SamOnnxModel(nn.Module):\n \"\"\"\n This model should not be called directly, but is used in ONNX export.\n It combines the prompt encoder, mask decoder, and mask postprocessing of Sam,\n with some functions modified to enable model tracing. Also supports extra\n options controlling what information. See the ONNX export script for details.\n \"\"\"\n\n def __init__(\n self,\n model: Sam,\n return_single_mask: bool,\n use_stability_score: bool = False,\n return_extra_metrics: bool = False,\n ) -> None:\n super().__init__()\n self.mask_decoder = model.mask_decoder\n self.model = model\n self.img_size = model.image_encoder.img_size\n self.return_single_mask = return_single_mask\n self.use_stability_score = use_stability_score\n self.stability_score_offset = 1.0\n self.return_extra_metrics = return_extra_metrics\n\n @staticmethod\n def resize_longest_image_size(\n input_image_size: torch.Tensor, longest_side: int\n ) -> torch.Tensor:\n input_image_size = input_image_size.to(torch.float32)\n scale = longest_side / torch.max(input_image_size)\n transformed_size = scale * input_image_size\n transformed_size = torch.floor(transformed_size + 0.5).to(torch.int64)\n return transformed_size\n\n def _embed_points(\n self, point_coords: torch.Tensor, point_labels: torch.Tensor\n ) -> torch.Tensor:\n point_coords = point_coords + 0.5\n point_coords = point_coords / self.img_size\n point_embedding = self.model.prompt_encoder.pe_layer._pe_encoding(point_coords)\n point_labels = point_labels.unsqueeze(-1).expand_as(point_embedding)\n\n point_embedding = point_embedding * (point_labels != -1)\n point_embedding = (\n point_embedding\n + self.model.prompt_encoder.not_a_point_embed.weight * (point_labels == -1)\n )\n\n for i in range(self.model.prompt_encoder.num_point_embeddings):\n point_embedding = (\n point_embedding\n + self.model.prompt_encoder.point_embeddings[i].weight\n * (point_labels == i)\n )\n\n return point_embedding\n\n def _embed_masks(\n self, input_mask: torch.Tensor, has_mask_input: torch.Tensor\n ) -> torch.Tensor:\n mask_embedding = has_mask_input * self.model.prompt_encoder.mask_downscaling(\n input_mask\n )\n mask_embedding = mask_embedding + (\n 1 - has_mask_input\n ) * self.model.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1)\n return mask_embedding\n\n def mask_postprocessing(\n self, masks: torch.Tensor, orig_im_size: torch.Tensor\n ) -> torch.Tensor:\n masks = F.interpolate(\n masks,\n size=(self.img_size, self.img_size),\n mode=\"bilinear\",\n align_corners=False,\n )\n\n prepadded_size = self.resize_longest_image_size(orig_im_size, self.img_size).to(\n torch.int64\n )\n masks = masks[..., : prepadded_size[0], : prepadded_size[1]] # type: ignore\n\n orig_im_size = orig_im_size.to(torch.int64)\n h, w = orig_im_size[0], orig_im_size[1]\n masks = F.interpolate(masks, size=(h, w), mode=\"bilinear\", align_corners=False)\n return masks\n\n def select_masks(\n self, masks: torch.Tensor, iou_preds: torch.Tensor, num_points: int\n ) -> Tuple[torch.Tensor, torch.Tensor]:\n # Determine if we should return the multiclick mask or not from the number of points.\n # The reweighting is used to avoid control flow.\n score_reweight = torch.tensor(\n [[1000] + [0] * (self.model.mask_decoder.num_mask_tokens - 1)]\n ).to(iou_preds.device)\n score = iou_preds + (num_points - 2.5) * score_reweight\n best_idx = torch.argmax(score, dim=1)\n masks = masks[torch.arange(masks.shape[0]), best_idx, :, :].unsqueeze(1)\n iou_preds = iou_preds[torch.arange(masks.shape[0]), best_idx].unsqueeze(1)\n\n return masks, iou_preds\n\n @torch.no_grad()\n def forward(\n self,\n image_embeddings: torch.Tensor,\n point_coords: torch.Tensor,\n point_labels: torch.Tensor,\n mask_input: torch.Tensor,\n has_mask_input: torch.Tensor,\n orig_im_size: torch.Tensor,\n ):\n sparse_embedding = self._embed_points(point_coords, point_labels)\n dense_embedding = self._embed_masks(mask_input, has_mask_input)\n\n masks, scores = self.model.mask_decoder.predict_masks(\n image_embeddings=image_embeddings,\n image_pe=self.model.prompt_encoder.get_dense_pe(),\n sparse_prompt_embeddings=sparse_embedding,\n dense_prompt_embeddings=dense_embedding,\n )\n\n if self.use_stability_score:\n scores = calculate_stability_score(\n masks, self.model.mask_threshold, self.stability_score_offset\n )\n\n if self.return_single_mask:\n masks, scores = self.select_masks(masks, scores, point_coords.shape[1])\n\n upscaled_masks = self.mask_postprocessing(masks, orig_im_size)\n\n if self.return_extra_metrics:\n stability_scores = calculate_stability_score(\n upscaled_masks, self.model.mask_threshold, self.stability_score_offset\n )\n areas = (upscaled_masks > self.model.mask_threshold).sum(-1).sum(-1)\n return upscaled_masks, scores, stability_scores, areas, masks\n\n return upscaled_masks, scores, masks\n" }, { "path": "model/segment_anything/utils/transforms.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom copy import deepcopy\nfrom typing import Tuple\n\nimport numpy as np\nimport torch\nfrom torch.nn import functional as F\nfrom torchvision.transforms.functional import resize # type: ignore\nfrom torchvision.transforms.functional import to_pil_image\n\n\nclass ResizeLongestSide:\n \"\"\"\n Resizes images to the longest side 'target_length', as well as provides\n methods for resizing coordinates and boxes. Provides methods for\n transforming both numpy array and batched torch tensors.\n \"\"\"\n\n def __init__(self, target_length: int) -> None:\n self.target_length = target_length\n\n def apply_image(self, image: np.ndarray) -> np.ndarray:\n \"\"\"\n Expects a numpy array with shape HxWxC in uint8 format.\n \"\"\"\n target_size = self.get_preprocess_shape(\n image.shape[0], image.shape[1], self.target_length\n )\n return np.array(resize(to_pil_image(image), target_size))\n\n def apply_coords(\n self, coords: np.ndarray, original_size: Tuple[int, ...]\n ) -> np.ndarray:\n \"\"\"\n Expects a numpy array of length 2 in the final dimension. Requires the\n original image size in (H, W) format.\n \"\"\"\n old_h, old_w = original_size\n new_h, new_w = self.get_preprocess_shape(\n original_size[0], original_size[1], self.target_length\n )\n coords = deepcopy(coords).astype(float)\n coords[..., 0] = coords[..., 0] * (new_w / old_w)\n coords[..., 1] = coords[..., 1] * (new_h / old_h)\n return coords\n\n def apply_boxes(\n self, boxes: np.ndarray, original_size: Tuple[int, ...]\n ) -> np.ndarray:\n \"\"\"\n Expects a numpy array shape Bx4. Requires the original image size\n in (H, W) format.\n \"\"\"\n boxes = self.apply_coords(boxes.reshape(-1, 2, 2), original_size)\n return boxes.reshape(-1, 4)\n\n def apply_image_torch(self, image: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Expects batched images with shape BxCxHxW and float format. This\n transformation may not exactly match apply_image. apply_image is\n the transformation expected by the model.\n \"\"\"\n # Expects an image in BCHW format. May not exactly match apply_image.\n target_size = self.get_preprocess_shape(\n image.shape[0], image.shape[1], self.target_length\n )\n return F.interpolate(\n image, target_size, mode=\"bilinear\", align_corners=False, antialias=True\n )\n\n def apply_coords_torch(\n self, coords: torch.Tensor, original_size: Tuple[int, ...]\n ) -> torch.Tensor:\n \"\"\"\n Expects a torch tensor with length 2 in the last dimension. Requires the\n original image size in (H, W) format.\n \"\"\"\n old_h, old_w = original_size\n new_h, new_w = self.get_preprocess_shape(\n original_size[0], original_size[1], self.target_length\n )\n coords = deepcopy(coords).to(torch.float)\n coords[..., 0] = coords[..., 0] * (new_w / old_w)\n coords[..., 1] = coords[..., 1] * (new_h / old_h)\n return coords\n\n def apply_boxes_torch(\n self, boxes: torch.Tensor, original_size: Tuple[int, ...]\n ) -> torch.Tensor:\n \"\"\"\n Expects a torch tensor with shape Bx4. Requires the original image\n size in (H, W) format.\n \"\"\"\n boxes = self.apply_coords_torch(boxes.reshape(-1, 2, 2), original_size)\n return boxes.reshape(-1, 4)\n\n @staticmethod\n def get_preprocess_shape(\n oldh: int, oldw: int, long_side_length: int\n ) -> Tuple[int, int]:\n \"\"\"\n Compute the output size given input size and target long side length.\n \"\"\"\n scale = long_side_length * 1.0 / max(oldh, oldw)\n newh, neww = oldh * scale, oldw * scale\n neww = int(neww + 0.5)\n newh = int(newh + 0.5)\n return (newh, neww)\n" }, { "path": "model/tf/modeling_outputs.py", "content": "import torch\nimport warnings\nfrom dataclasses import dataclass\nfrom typing import Optional, Tuple, Dict, List\nfrom transformers.utils import ModelOutput\n\n@dataclass\nclass CausalLMOutputWithPastAndLabel(ModelOutput):\n \"\"\"\n Base class for causal language model (or autoregressive) outputs.\n\n Args:\n loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):\n Language modeling loss (for next-token prediction).\n labels (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*, returned when `labels` is provided):\n logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):\n Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).\n past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):\n Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape\n `(batch_size, num_heads, sequence_length, embed_size_per_head)`)\n\n Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see\n `past_key_values` input) to speed up sequential decoding.\n hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):\n Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +\n one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.\n\n Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.\n attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):\n Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,\n sequence_length)`.\n\n Attentions weights after the attention softmax, used to compute the weighted average in the self-attention\n heads.\n \"\"\"\n\n loss: Optional[torch.FloatTensor] = None\n labels: Optional[torch.FloatTensor] = None\n logits: torch.FloatTensor = None\n past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None\n hidden_states: Optional[Tuple[torch.FloatTensor]] = None\n attentions: Optional[Tuple[torch.FloatTensor]] = None\n bs2imgs_token_list: List[List[int]] = None" }, { "path": "model/univi/__init__.py", "content": "from .model import ChatUniViLlamaForCausalLM\n" }, { "path": "model/univi/config/__init__.py", "content": "from .dataset_config import *\nfrom .model_config import *\n\n\nModelConfig = {\n \"PRETUNE\": model_config_pretune,\n \"FINETUNE\": model_config_finetune,\n}\n\n\nDataConfig = {\n \"Pretrain\": [Pretrain, COCO_CAP, COCO_REG, COCO_REC],\n \"SQA\": [SQA],\n \"FINETUNE\": [VIT, MIMIC_imageonly, VIDEO],\n}" }, { "path": "model/univi/config/dataset_config.py", "content": "Pretrain = {\n \"chat_path\": \"${PATH}/CC3M-595K/chat.json\",\n \"CC3M\": \"${PATH}/CC3M-595K\",\n}\n\nVIT = {\n \"chat_path\": \"${PATH}/llava_instruct_150k.json\",\n \"COCO2017\": \"${PATH}/COCO2017/train2017\",\n}\n\nMIMIC_imageonly = {\n \"chat_path\": \"${PATH}/MIMIC-IT-imageonly.json\",\n \"CDG\": \"${PATH}/CGD/images\",\n \"LA\": \"${PATH}/LA/images\",\n \"SD\": \"${PATH}/SD/images\",\n}\n\nCOCO_CAP = {\n \"chat_path\": \"${PATH}/COCO/coco_cap_chat.json\",\n \"COCO2014\": \"${PATH}/COCO2014/train2014\",\n}\n\nCOCO_REG = {\n \"chat_path\": \"${PATH}/COCO/coco_reg_chat.json\",\n \"COCO2014\": \"${PATH}/COCO2014/train2014\",\n}\n\nCOCO_REC = {\n \"chat_path\": \"${PATH}/COCO/coco_rec_chat.json\",\n \"COCO2014\": \"${PATH}/COCO2014/train2014\",\n}\n\nVIDEO = {\n \"chat_path\": \"${PATH}/video_chat.json\",\n \"VIDEO\": \"${PATH}/Activity_Videos\",\n}\n\nSQA = {\n \"chat_path\": \"${PATH}/llava_train_QCM-LEA.json\",\n \"ScienceQA\": \"${PATH}/scienceqa/train\",\n}" }, { "path": "model/univi/config/model_config.py", "content": "model_config_pretune = {\n \"use_cluster\": True,\n \"freeze\": False,\n \"vision_tune\": False,\n\n \"spatial_cluster_rate0\": 64, # 0.25\n \"spatial_cluster_rate1\": 32, # 0.5\n \"spatial_cluster_rate2\": 16, # 0.5\n\n \"temporal_cluster_rate\": 1/16,\n}\n\nmodel_config_finetune = {\n \"use_cluster\": True,\n \"freeze\": False,\n \"mm_tune\": True,\n \"vision_tune\": False,\n\n \"spatial_cluster_rate0\": 64, # 0.25\n \"spatial_cluster_rate1\": 32, # 0.5\n \"spatial_cluster_rate2\": 16, # 0.5\n\n \"temporal_cluster_rate\": 1/16,\n}" }, { "path": "model/univi/constants.py", "content": "CONTROLLER_HEART_BEAT_EXPIRATION = 30\nWORKER_HEART_BEAT_INTERVAL = 15\n\nLOGDIR = \".\"\n\n# Model Constants\nMAX_IMAGE_LENGTH = 64\nIGNORE_INDEX = -100\nIMAGE_TOKEN_INDEX = -200\nDEFAULT_IMAGE_TOKEN = \"\"\nDEFAULT_VIDEO_TOKEN = \"