release iChatApp

.gitattributes

@@ -32,3 +32,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,27 @@
+# compilation and distribution
+__pycache__
+_ext
+*.pyc
+*.pyd
+*.so
+*.dll
+*.egg-info/
+build/
+dist/
+wheels/
+
+# Editor temporaries
+*.swn
+*.swo
+*.swp
+*~
+
+# editor settings
+.idea
+.vscode
+_darcs
+
+# custom files
+./image/
+./tmp_files/
+

README.md

@@ -1,13 +1,166 @@
----
 title: InternChat
-emoji: 😻
-colorFrom: green
-colorTo: indigo
+emoji: 🤖💬
+colorFrom: indigo
+colorTo: pink
 sdk: gradio
-sdk_version: 3.29.0
-app_file: app.py
+sdk_version: 3.28.1
+app_file: iChatApp.py
 pinned: false
 license: apache-2.0
----
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+[[中文文档]](README_CN.md)
+
+**The project is still under construction, we will continue to update it and welcome contributions/pull requests from the community.**
+
+
+
+<p align="center"><img src="./assets/gvlab_logo.png" width="600"></p>
+
+<a src="https://img.shields.io/discord/1099920215724277770?label=Discord&logo=discord" href="https://discord.gg/khWBFnCgAN">
+    <img src="https://img.shields.io/discord/1099920215724277770?label=Discord&logo=discord"> </a> | <a src="https://img.shields.io/badge/GPU%20Demo-Open-green?logo=alibabacloud" href="https://ichat.opengvlab.com">
+    <img src="https://img.shields.io/badge/Demo-Open-green?logo=alibabacloud"> </a> | <a src="https://img.shields.io/twitter/follow/opengvlab?style=social" href="https://twitter.com/opengvlab">
+    <img src="https://img.shields.io/twitter/follow/opengvlab?style=social">  </a> 
+    
+
+
+# InternChat [[paper](https://pjlab-gvm-data.oss-cn-shanghai.aliyuncs.com/papers/ichat.pdf)]
+
+
+<!-- ## Description -->
+**InternChat**(short for **iChat**) is pointing-language-driven visual interactive system, allowing you to interact with ChatGPT by clicking, dragging and drawing using a pointing device. The name InternChat stands for **inter**action, **n**onverbal, and **chat**bots. Different from existing interactive systems that rely on pure language, by incorporating pointing instructions, iChat significantly improves the efficiency of communication between users and chatbots, as well as the accuracy of chatbots in vision-centric tasks, especially in complicated visual scenarios. Additionally, in iChat, an auxiliary control mechanism is used to improve the control capability of LLM, and a large vision-language model termed **Husky** is fine-tuned for high-quality multi-modal dialogue (impressing ChatGPT-3.5-turbo with **93.89% GPT-4 Quality**).
+
+## Online Demo
+[**InternChat**](https://ichat.opengvlab.com/) is online. Let's try it!
+
+[**NOTE**] It is possible that you are waiting in a lengthy queue. You can clone our repo and run it with your private GPU.
+
+https://github.com/OpenGVLab/InternChat/assets/13723743/3270b05f-0823-4f13-9966-4010fd855643
+
+
+
+## Schedule
+- [ ] Support Chinese
+- [ ] Support MOSS
+- [ ] More powerful foundation models based on [InternImage](https://github.com/OpenGVLab/InternImage) and [InternVideo](https://github.com/OpenGVLab/InternVideo)
+- [ ] More accurate interactive experience
+- [ ] Web Page & Code Generation 
+- [x] Support voice assistant
+- [x] Support click interaction
+- [x] Interactive image editing
+- [x] Interactive image generation
+- [x] Interactive visual question answering
+- [x] Segment Anything
+- [x] Image inpainting
+- [x] Image caption
+- [x] image matting
+- [x] Optical character recognition
+- [x] Action recognition
+- [x] Video caption
+- [x] Video dense caption
+- [x] video highlight interpretation
+
+
+
+## System Overview
+<p align="center"><img src="./assets/arch1.png" alt="Logo"></p>
+
+## 🎁 Major Features
+<!--<!-- <p align="center"><img src="./assets/online_demo.gif" alt="Logo"></p> -->  
+<p align="center">(a) Remove the masked object</p>
+<p align="center"><img src="./assets/demo2.gif" width="500"></p>
+
+<p align="center">(b) Interactive image editing</center>
+<p align="center"><img src="./assets/demo3.gif" width="500"></p>
+
+<p align="center">(c) Image generation</p>
+<p align="center"><img src="./assets/demo4.gif" align='justify'  width="500"></p>
+
+<p align="center">(d) Interactive visual question answer</p>
+<p align="center"><img src="./assets/demo5.gif" align='justify' width="700"></p> 
+
+
+<p align="center">(e) Interactive image generation</p>
+<p align="center"><img width="800" alt="image" src="https://github.com/OpenGVLab/InternChat/assets/8529570/2b0da08e-af86-453d-99e5-1327f93aa917"></p> 
+
+
+<p align="center">(f) Video highlight interpretation</p>
+<p align="center"><img src="./assets/demo6.jpg" align='justify' width="500"></p> 
+
+<!-- ![alt]("./assets/demo5.gif" "title") -->
+
+
+## 🛠️ Installation
+
+### Basic requirements
+
+- Linux
+- Python 3.8+
+- PyTorch 1.12+
+- CUDA 11.6+ 
+- GCC & G++ 5.4+
+- GPU Memory >= 17G for loading basic tools (HuskyVQA, SegmentAnything, ImageOCRRecognition)
+
+### Install Python dependencies
+```shell
+pip install -r requirements.txt
+```
+
+### Model zoo
+Coming soon...
+
+## 👨‍🏫 Get Started 
+Running the following shell can start a gradio service:
+```shell
+python -u iChatApp.py --load "HuskyVQA_cuda:0,SegmentAnything_cuda:0,ImageOCRRecognition_cuda:0" --port 3456
+```
+
+if you want to enable the voice assistant, please use `openssl` to generate the certificate:
+```shell
+openssl req -x509 -newkey rsa:4096 -keyout ./key.pem -out ./cert.pem -sha256 -days 365 -nodes
+```
+
+and then run:
+```shell
+python -u iChatApp.py --load "HuskyVQA_cuda:0,SegmentAnything_cuda:0,ImageOCRRecognition_cuda:0" --port 3456 --https
+```
+ 
+
+
+
+## 🎫 License
+
+This project is released under the [Apache 2.0 license](LICENSE). 
+
+## 🖊️ Citation
+
+If you find this project useful in your research, please consider cite:
+```BibTeX
+@misc{2023internchat,
+    title={InternChat: Solving Vision-Centric Tasks by Interacting with Chatbots Beyond Language},
+    author={Zhaoyang Liu and Yinan He and Wenhai Wang and Weiyun Wang and Yi Wang and Shoufa Chen and Qinglong Zhang and Yang Yang and Qingyun Li and Jiashuo Yu and Kunchang Li and Zhe Chen and Xue Yang and Xizhou Zhu and Yali Wang and Limin Wang and Ping Luo and Jifeng Dai and Yu Qiao},
+    howpublished = {\url{https://arxiv.org/abs/2305.05662}},
+    year={2023}
+}
+```
+
+## 🤝 Acknowledgement
+Thanks to the open source of the following projects:
+
+[Hugging Face](https://github.com/huggingface) &#8194;
+[LangChain](https://github.com/hwchase17/langchain) &#8194;
+[TaskMatrix](https://github.com/microsoft/TaskMatrix) &#8194;
+[SAM](https://github.com/facebookresearch/segment-anything) &#8194;
+[Stable Diffusion](https://github.com/CompVis/stable-diffusion) &#8194; 
+[ControlNet](https://github.com/lllyasviel/ControlNet) &#8194; 
+[InstructPix2Pix](https://github.com/timothybrooks/instruct-pix2pix) &#8194; 
+[BLIP](https://github.com/salesforce/BLIP) &#8194;
+[Latent Diffusion Models](https://github.com/CompVis/latent-diffusion) &#8194;
+[EasyOCR](https://github.com/JaidedAI/EasyOCR) &#8194;
+
+Welcome to discuss with us and continuously improve the user experience of InternChat.
+
+WeChat QR Code
+
+<p align="center"><img width="500" alt="image" src="https://github.com/OpenGVLab/InternChat/assets/8529570/881c231d-9049-4920-a22c-680f41f0f7ee"></p> 
+

README_CN.md

@@ -0,0 +1,147 @@
+[[English Document](README.md)]
+
+**[NOTE] 该项目仍在建设中，我们将继续更新，并欢迎社区的贡献/拉取请求。**
+
+<p align="center"><img src="./assets/gvlab_logo.png" width="600"></p>
+
+<a src="https://img.shields.io/discord/1099920215724277770?label=Discord&logo=discord" href="https://discord.gg/khWBFnCgAN">
+<img src="https://img.shields.io/discord/1099920215724277770?label=Discord&logo=discord"> </a> | <a src="https://img.shields.io/badge/GPU Demo-Open-green?logo=alibabacloud" href="https://ichat.opengvlab.com">
+<img src="https://img.shields.io/badge/Demo-Open-green?logo=alibabacloud"> </a> | <a src="https://img.shields.io/twitter/follow/opengvlab?style=social" href="https://twitter.com/opengvlab">
+<img src="https://img.shields.io/twitter/follow/opengvlab?style=social">
+
+# InternChat [[论文](https://arxiv.org/pdf/2305.05662.pdf)]
+<!-- ## 描述 -->
+**InternChat**（简称 **iChat**）是一种基于指向语言驱动的视觉交互系统，允许您使用指向设备通过点击、拖动和绘制与 ChatGPT 进行互动。InternChat 的名称代表了 **inter**action（交互）、**n**onverbal（非语言）和 **chat**bots（聊天机器人）。与依赖纯语言的现有交互系统不同，通过整合指向指令，iChat 显著提高了用户与聊天机器人之间的沟通效率，以及聊天机器人在视觉为中心任务中的准确性，特别是在复杂的视觉场景中。此外，在 iChat 中，采用辅助控制机制来提高 LLM 的控制能力，并对一个大型视觉-语言模型 **Husky** 进行微调，以实现高质量的多模态对话（在ChatGPT-3.5-turbo评测中达到 **93.89% GPT-4 质量**）。
+  
+## 在线Demo
+  
+[注意] 可能会出现排队等待较长时间。您可以clone我们的仓库并使用您自己的GPU运行。
+  
+[**InternChat**已上线，尝试一下！](https://ichat.opengvlab.com)
+
+
+https://github.com/OpenGVLab/InternChat/assets/13723743/3270b05f-0823-4f13-9966-4010fd855643
+  
+## Schedule
+- [ ] 支持中文
+- [ ] 支持 MOSS
+- [ ] 基于 InternImage 和 InternVideo 的更强大的基础模型
+- [ ] 更准确的交互体验
+- [ ] 网页 & 代码生成
+- [x] 支持语音助手
+- [x] 支持点击交互
+- [x] 交互式图像编辑
+- [x] 交互式图像生成
+- [x] 交互式视觉问答
+- [x] Segment Anything模型
+- [x] 图像修复
+- [x] 图像描述
+- [x] 图像抠图
+- [x] 光学字符识别（OCR）
+- [x] 动作识别
+- [x] 视频描述
+- [x] 视频密集描述
+- [x] 视频高光时刻截取
+  
+## 系统概览
+<p align="center"><img src="./assets/arch1.png" alt="Logo"></p>
+  
+## 🎁 主要特点
+<!--<!-- <p align="center"><img src="./assets/online_demo.gif" alt="Logo"></p> -->
+
+<p align="center">(a) 移除遮盖的对象</p>
+<p align="center"><img src="./assets/demo2.gif" width="500"></p>
+
+<p align="center">(b) 交互式图像编辑</center>
+<p align="center"><img src="./assets/demo3.gif" width="500"></p>
+
+<p align="center">(c) 图像生成</p>
+<p align="center"><img src="./assets/demo4.gif" align='justify'  width="500"></p>
+
+<p align="center">(d) 交互式视觉问答</p>
+<p align="center"><img src="./assets/demo5.gif" align='justify' width="700"></p>
+
+<p align="center">(e) 交互式图像生成</p>
+<p align="center"><img width="800" alt="image" src="https://github.com/OpenGVLab/InternChat/assets/8529570/2b0da08e-af86-453d-99e5-1327f93aa917"></p>
+
+<p align="center">(f) 视频高光解释</p>
+<p align="center"><img src="./assets/demo6.jpg" align='justify' width="500"></p> 
+
+<!-- ![alt]("./assets/demo5.gif" "title") -->
+
+## 🛠️ 安装
+
+### 基本要求
+
+- Linux
+- Python 3.8+
+- PyTorch 1.12+
+- CUDA 11.6+
+- GCC & G++ 5.4+
+- GPU 内存 >= 17G 用于加载基本工具 (HuskyVQA, SegmentAnything, ImageOCRRecognition)
+
+### 安装Python的依赖项
+```shell
+pip install -r requirements.txt
+```
+
+### 模型库
+
+即将推出...
+
+## 👨‍🏫 运行指南
+
+运行以下 shell 可启动一个 gradio 服务：
+
+```shell
+python -u iChatApp.py --load "HuskyVQA_cuda:0,SegmentAnything_cuda:0,ImageOCRRecognition_cuda:0" --port 3456
+```
+
+如果您想启用语音助手，请使用 openssl 生成证书：
+
+```shell
+openssl req -x509 -newkey rsa:4096 -keyout ./key.pem -out ./cert.pem -sha256 -days 365 -nodes
+```
+然后运行：
+
+```shell
+python -u iChatApp.py --load "HuskyVQA_cuda:0,SegmentAnything_cuda:0,ImageOCRRecognition_cuda:0" --port 3456 --https
+```
+
+
+## 🎫 许可
+
+该项目根据[Apache 2.0 license](LICENSE)发布。
+
+## 🖊️ 引用
+
+如果您在研究中发现这个项目有用，请考虑引用：
+```BibTeX
+@misc{2023internchat,
+    title={InternChat: Solving Vision-Centric Tasks by Interacting with Chatbots Beyond Language},
+    author={Zhaoyang Liu and Yinan He and Wenhai Wang and Weiyun Wang and Yi Wang and Shoufa Chen and Qinglong Zhang and Yang Yang and Qingyun Li and Jiashuo Yu and Kunchang Li and Zhe Chen and Xue Yang and Xizhou Zhu and Yali Wang and Limin Wang and Ping Luo and Jifeng Dai and Yu Qiao},
+    howpublished = {\url{https://arxiv.org/abs/2305.05662}},
+    year={2023}
+}
+```
+
+## 🤝 致谢
+
+感谢以下开源项目:
+
+[Hugging Face](https://github.com/huggingface) &#8194;
+[LangChain](https://github.com/hwchase17/langchain) &#8194;
+[TaskMatrix](https://github.com/microsoft/TaskMatrix) &#8194;
+[SAM](https://github.com/facebookresearch/segment-anything) &#8194;
+[Stable Diffusion](https://github.com/CompVis/stable-diffusion) &#8194; 
+[ControlNet](https://github.com/lllyasviel/ControlNet) &#8194; 
+[InstructPix2Pix](https://github.com/timothybrooks/instruct-pix2pix) &#8194; 
+[BLIP](https://github.com/salesforce/BLIP) &#8194;
+[Latent Diffusion Models](https://github.com/CompVis/latent-diffusion) &#8194;
+[EasyOCR](https://github.com/JaidedAI/EasyOCR) &#8194;
+
+
+
+如果您在试用、运行、部署中有任何问题，欢迎加入我们的微信群讨论！如果您对项目有任何的想法和建议，欢迎加入我们的微信群讨论！
+
+<p align="center"><img width="500" alt="image" src="https://github.com/OpenGVLab/InternChat/assets/8529570/881c231d-9049-4920-a22c-680f41f0f7ee"></p> 

assets/videos/iKun.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52152a27a2c42c4dec12575c50dd6dced8db89ba041f5f90596aeb228b213e0a
+size 1055241

configs/big_lama_config.yaml

@@ -0,0 +1,157 @@
+run_title: b18_ffc075_batch8x15
+training_model:
+  kind: default
+  visualize_each_iters: 1000
+  concat_mask: true
+  store_discr_outputs_for_vis: true
+losses:
+  l1:
+    weight_missing: 0
+    weight_known: 10
+  perceptual:
+    weight: 0
+  adversarial:
+    kind: r1
+    weight: 10
+    gp_coef: 0.001
+    mask_as_fake_target: true
+    allow_scale_mask: true
+  feature_matching:
+    weight: 100
+  resnet_pl:
+    weight: 30
+    weights_path: ${env:TORCH_HOME}
+
+optimizers:
+  generator:
+    kind: adam
+    lr: 0.001
+  discriminator:
+    kind: adam
+    lr: 0.0001
+visualizer:
+  key_order:
+  - image
+  - predicted_image
+  - discr_output_fake
+  - discr_output_real
+  - inpainted
+  rescale_keys:
+  - discr_output_fake
+  - discr_output_real
+  kind: directory
+  outdir: /group-volume/User-Driven-Content-Generation/r.suvorov/inpainting/experiments/r.suvorov_2021-04-30_14-41-12_train_simple_pix2pix2_gap_sdpl_novgg_large_b18_ffc075_batch8x15/samples
+location:
+  data_root_dir: /group-volume/User-Driven-Content-Generation/datasets/inpainting_data_root_large
+  out_root_dir: /group-volume/User-Driven-Content-Generation/${env:USER}/inpainting/experiments
+  tb_dir: /group-volume/User-Driven-Content-Generation/${env:USER}/inpainting/tb_logs
+data:
+  batch_size: 15
+  val_batch_size: 2
+  num_workers: 3
+  train:
+    indir: ${location.data_root_dir}/train
+    out_size: 256
+    mask_gen_kwargs:
+      irregular_proba: 1
+      irregular_kwargs:
+        max_angle: 4
+        max_len: 200
+        max_width: 100
+        max_times: 5
+        min_times: 1
+      box_proba: 1
+      box_kwargs:
+        margin: 10
+        bbox_min_size: 30
+        bbox_max_size: 150
+        max_times: 3
+        min_times: 1
+      segm_proba: 0
+      segm_kwargs:
+        confidence_threshold: 0.5
+        max_object_area: 0.5
+        min_mask_area: 0.07
+        downsample_levels: 6
+        num_variants_per_mask: 1
+        rigidness_mode: 1
+        max_foreground_coverage: 0.3
+        max_foreground_intersection: 0.7
+        max_mask_intersection: 0.1
+        max_hidden_area: 0.1
+        max_scale_change: 0.25
+        horizontal_flip: true
+        max_vertical_shift: 0.2
+        position_shuffle: true
+    transform_variant: distortions
+    dataloader_kwargs:
+      batch_size: ${data.batch_size}
+      shuffle: true
+      num_workers: ${data.num_workers}
+  val:
+    indir: ${location.data_root_dir}/val
+    img_suffix: .png
+    dataloader_kwargs:
+      batch_size: ${data.val_batch_size}
+      shuffle: false
+      num_workers: ${data.num_workers}
+  visual_test:
+    indir: ${location.data_root_dir}/korean_test
+    img_suffix: _input.png
+    pad_out_to_modulo: 32
+    dataloader_kwargs:
+      batch_size: 1
+      shuffle: false
+      num_workers: ${data.num_workers}
+generator:
+  kind: ffc_resnet
+  input_nc: 4
+  output_nc: 3
+  ngf: 64
+  n_downsampling: 3
+  n_blocks: 18
+  add_out_act: sigmoid
+  init_conv_kwargs:
+    ratio_gin: 0
+    ratio_gout: 0
+    enable_lfu: false
+  downsample_conv_kwargs:
+    ratio_gin: ${generator.init_conv_kwargs.ratio_gout}
+    ratio_gout: ${generator.downsample_conv_kwargs.ratio_gin}
+    enable_lfu: false
+  resnet_conv_kwargs:
+    ratio_gin: 0.75
+    ratio_gout: ${generator.resnet_conv_kwargs.ratio_gin}
+    enable_lfu: false
+discriminator:
+  kind: pix2pixhd_nlayer
+  input_nc: 3
+  ndf: 64
+  n_layers: 4
+evaluator:
+  kind: default
+  inpainted_key: inpainted
+  integral_kind: ssim_fid100_f1
+trainer:
+  kwargs:
+    gpus: -1
+    accelerator: ddp
+    max_epochs: 200
+    gradient_clip_val: 1
+    log_gpu_memory: None
+    limit_train_batches: 25000
+    val_check_interval: ${trainer.kwargs.limit_train_batches}
+    log_every_n_steps: 1000
+    precision: 32
+    terminate_on_nan: false
+    check_val_every_n_epoch: 1
+    num_sanity_val_steps: 8
+    limit_val_batches: 1000
+    replace_sampler_ddp: false
+  checkpoint_kwargs:
+    verbose: true
+    save_top_k: 5
+    save_last: true
+    period: 1
+    monitor: val_ssim_fid100_f1_total_mean
+    mode: max

configs/med_config.json

@@ -0,0 +1,21 @@
+{
+  "architectures": [
+    "BertModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "type_vocab_size": 2,
+  "vocab_size": 30524,
+  "encoder_width": 768,
+  "add_cross_attention": true   
+}

configs/q2l_config.json

@@ -0,0 +1,23 @@
+{
+    "architectures": [
+      "BertModel"
+    ],
+    "attention_probs_dropout_prob": 0.1,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 768,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "layer_norm_eps": 1e-12,
+    "max_position_embeddings": 512,
+    "model_type": "bert",
+    "num_attention_heads": 4,
+    "num_hidden_layers": 2,
+    "pad_token_id": 0,
+    "type_vocab_size": 2,
+    "vocab_size": 30522,
+    "encoder_width": 768,
+    "add_cross_attention": true,
+    "add_tag_cross_attention": false
+  }
+  

configs/swin/config_swinB_224.json

@@ -0,0 +1,10 @@
+{
+    "ckpt": "pretrain_model/swin_base_patch4_window7_224_22k.pth",
+    "vision_width": 1024,
+    "image_res": 224,
+    "window_size": 7,
+    "embed_dim": 128,
+    "depths": [ 2, 2, 18, 2 ],
+    "num_heads": [ 4, 8, 16, 32 ]
+  }
+  

configs/swin/config_swinB_384.json

@@ -0,0 +1,10 @@
+{
+    "ckpt": "pretrain_model/swin_base_patch4_window7_224_22k.pth",
+    "vision_width": 1024,
+    "image_res": 384,
+    "window_size": 12,
+    "embed_dim": 128,
+    "depths": [ 2, 2, 18, 2 ],
+    "num_heads": [ 4, 8, 16, 32 ]
+  }
+  

configs/swin/config_swinB_480.json

@@ -0,0 +1,9 @@
+{
+  "ckpt": "pretrain_model/swin_base_patch4_window7_224_22k.pth",
+  "vision_width": 1024,
+  "image_res": 480,
+  "window_size": 15,
+  "embed_dim": 128,
+  "depths": [ 2, 2, 18, 2 ],
+  "num_heads": [ 4, 8, 16, 32 ]
+}

configs/swin/config_swinB_576.json

@@ -0,0 +1,9 @@
+{
+  "ckpt": "pretrain_model/swin_base_patch4_window7_224_22k.pth",
+  "vision_width": 1024,
+  "image_res": 576,
+  "window_size": 18,
+  "embed_dim": 128,
+  "depths": [ 2, 2, 18, 2 ],
+  "num_heads": [ 4, 8, 16, 32 ]
+}

configs/swin/config_swinB_608.json

@@ -0,0 +1,9 @@
+{
+  "ckpt": "pretrain_model/swin_base_patch4_window7_224_22k.pth",
+  "vision_width": 1024,
+  "image_res": 608,
+  "window_size": 19,
+  "embed_dim": 128,
+  "depths": [ 2, 2, 18, 2 ],
+  "num_heads": [ 4, 8, 16, 32 ]
+}

configs/tag2text_caption.yaml

@@ -0,0 +1,33 @@
+image_root: '/home/notebook/data/group/projects/tagging/caption/datasets/public/coco/'
+
+ann_root: 'dataset/caption_dataset'
+coco_gt_root: 'dataset/caption_dataset'
+
+pretrained: '/home/notebook/code/personal/S9049611/BLIP/output/pretrain_caption_tagtotext_v2_bert_asl'
+
+# size of vit model; base or large
+vit: 'swin_b'
+vit_grad_ckpt: False
+vit_ckpt_layer: 0
+
+batch_size: 35
+init_lr: 5e-6 
+
+image_size: 384
+
+# generation configs
+max_length: 20  
+min_length: 5
+num_beams: 3
+prompt: 'a picture of '
+
+# optimizer
+weight_decay: 0.05
+min_lr: 0
+max_epoch: 10
+
+text_pretrain: 'bert'
+
+class_num: 3429
+threshold: 0.7
+

iChat/__init__.py

@@ -0,0 +1 @@
+from .models import *

iChat/chatbot/__init__.py

@@ -0,0 +1 @@
+from .chatbot import ConversationBot

iChat/chatbot/chatbot.py

@@ -0,0 +1,440 @@
+import inspect
+import re
+import os
+import numpy as np
+import uuid
+import shutil
+import whisper
+import gradio as gr
+
+from PIL import Image
+
+from langchain.agents.initialize import initialize_agent
+from langchain.agents.tools import Tool
+from langchain.chains.conversation.memory import ConversationBufferMemory
+from langchain.llms.openai import OpenAI
+
+from ..models import *
+from iGPT.models.utils import gen_new_name
+
+GLOBAL_SEED=1912
+
+
+'''
+INTERN_CHAT_PREFIX = """InternChat is designed to be able to assist with a wide range of text and visual related tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. InternChat is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.
+
+InternChat is able to process and understand large amounts of text and images. As a language model, InternChat can not directly read images, but it has a list of tools to finish different visual tasks. Each image will have a file name formed as "image/xxx.png", and InternChat can invoke different tools to indirectly understand pictures. When talking about images, InternChat is very strict to the file name and will never fabricate nonexistent files. When using tools to generate new image files, InternChat is also known that the image may not be the same as the user's demand, and will use other visual question answering tools or description tools to observe the real image. InternChat is able to use tools in a sequence, and is loyal to the tool observation outputs rather than faking the image content and image file name. It will remember to provide the file name from the last tool observation, if a new image is generated.
+
+Human may provide new figures to InternChat with a description. The description helps InternChat to understand this image, but InternChat should use tools to finish following tasks, rather than directly imagine from the description.
+
+Overall, InternChat is a powerful visual dialogue assistant tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. 
+
+
+TOOLS:
+------
+
+InternChat has access to the following tools:"""
+
+INTERN_CHAT_FORMAT_INSTRUCTIONS = """To use a tool, please use the following format:
+
+```
+Thought: Do I need to use a tool? Yes
+Action: the action to take, should be one of [{tool_names}]
+Action Input: the input to the action
+Observation: the result of the action
+```
+
+When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:
+
+```
+Thought: Do I need to use a tool? No
+{ai_prefix}: [your response here]
+```
+"""
+
+INTERN_CHAT_SUFFIX = """You are very strict to the filename correctness and will never fake a file name if it does not exist.
+You will remember to provide the image file name loyally if it's provided in the last tool observation.
+
+Begin!
+
+Previous conversation history:
+{chat_history}
+
+New input: {input}
+Since InternChat is a text language model, InternChat must use tools to observe images rather than imagination.
+The thoughts and observations are only visible for InternChat, InternChat should remember to repeat important information in the final response for Human. 
+Thought: Do I need to use a tool? {agent_scratchpad} Let's think step by step.
+"""
+
+INTERN_CHAT_PREFIX_CN = """InternChat 旨在能够协助完成范围广泛的文本和视觉相关任务，从回答简单的问题到提供对广泛主题的深入解释和讨论。 InternChat 能够根据收到的输入生成类似人类的文本，使其能够进行听起来自然的对话，并提供连贯且与手头主题相关的响应。
+
+InternChat 能够处理和理解大量文本和图像。作为一种语言模型，InternChat 不能直接读取图像，但它有一系列工具来完成不同的视觉任务。每张图片都会有一个文件名，格式为“image/xxx.png”，InternChat可以调用不同的工具来间接理解图片。在谈论图片时，InternChat 对文件名的要求非常严格，绝不会伪造不存在的文件。在使用工具生成新的图像文件时，InternChat也知道图像可能与用户需求不一样，会使用其他视觉问答工具或描述工具来观察真实图像。 InternChat 能够按顺序使用工具，并且忠于工具观察输出，而不是伪造图像内容和图像文件名。如果生成新图像，它将记得提供上次工具观察的文件名。
+
+Human 可能会向 InternChat 提供带有描述的新图形。描述帮助 InternChat 理解这个图像，但 InternChat 应该使用工具来完成以下任务，而不是直接从描述中想象。有些工具将会返回英文描述，但你对用户的聊天应当采用中文。
+
+总的来说，InternChat 是一个强大的可视化对话辅助工具，可以帮助处理范围广泛的任务，并提供关于范围广泛的主题的有价值的见解和信息。
+
+工具列表:
+------
+
+InternChat 可以使用这些工具:"""
+
+INTERN_CHAT_FORMAT_INSTRUCTIONS_CN = """用户使用中文和你进行聊天，但是工具的参数应当使用英文。如果要调用工具，你必须遵循如下格式:
+
+```
+Thought: Do I need to use a tool? Yes
+Action: the action to take, should be one of [{tool_names}]
+Action Input: the input to the action
+Observation: the result of the action
+```
+
+当你不再需要继续调用工具，而是对观察结果进行总结回复时，你必须使用如下格式：
+
+
+```
+Thought: Do I need to use a tool? No
+{ai_prefix}: [your response here]
+```
+"""
+
+INTERN_CHAT_SUFFIX_CN = """你对文件名的正确性非常严格，而且永远不会伪造不存在的文件。
+
+开始!
+
+因为InternChat是一个文本语言模型，必须使用工具去观察图片而不是依靠想象。
+推理想法和观察结果只对InternChat可见，需要记得在最终回复时把重要的信息重复给用户，你只能给用户返回中文句子。我们一步一步思考。在你使用工具时，工具的参数只能是英文。
+
+聊天历史:
+{chat_history}
+
+新输入: {input}
+Thought: Do I need to use a tool? {agent_scratchpad}
+"""
+'''
+
+
+VISUAL_CHATGPT_PREFIX = """Visual ChatGPT is designed to be able to assist with a wide range of text and visual related tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. Visual ChatGPT is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.
+
+Visual ChatGPT is able to process and understand large amounts of text and images. As a language model, Visual ChatGPT can not directly read images, but it has a list of tools to finish different visual tasks. Each image will have a file name formed as "image/xxx.png", and Visual ChatGPT can invoke different tools to indirectly understand pictures. When talking about images, Visual ChatGPT is very strict to the file name and will never fabricate nonexistent files. When using tools to generate new image files, Visual ChatGPT is also known that the image may not be the same as the user's demand, and will use other visual question answering tools or description tools to observe the real image. Visual ChatGPT is able to use tools in a sequence, and is loyal to the tool observation outputs rather than faking the image content and image file name. It will remember to provide the file name from the last tool observation, if a new image is generated.
+
+Human may provide new figures to Visual ChatGPT with a description. The description helps Visual ChatGPT to understand this image, but Visual ChatGPT should use tools to finish following tasks, rather than directly imagine from the description.
+
+Overall, Visual ChatGPT is a powerful visual dialogue assistant tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. 
+
+
+TOOLS:
+------
+
+Visual ChatGPT  has access to the following tools:"""
+
+VISUAL_CHATGPT_FORMAT_INSTRUCTIONS = """To use a tool, please use the following format:
+
+```
+Thought: Do I need to use a tool? Yes
+Action: the action to take, should be one of [{tool_names}]
+Action Input: the input to the action
+Observation: the result of the action
+```
+
+When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:
+
+```
+Thought: Do I need to use a tool? No
+{ai_prefix}: [your response here]
+```
+"""
+
+VISUAL_CHATGPT_SUFFIX = """You are very strict to the filename correctness and will never fake a file name if it does not exist.
+You will remember to provide the image file name loyally if it's provided in the last tool observation.
+
+Begin!
+
+Previous conversation history:
+{chat_history}
+
+New input: {input}
+Since Visual ChatGPT is a text language model, Visual ChatGPT must use tools to observe images rather than imagination.
+The thoughts and observations are only visible for Visual ChatGPT, Visual ChatGPT should remember to repeat important information in the final response for Human. 
+Thought: Do I need to use a tool? {agent_scratchpad} Let's think step by step.
+"""
+
+VISUAL_CHATGPT_PREFIX_CN = """Visual ChatGPT 旨在能够协助完成范围广泛的文本和视觉相关任务，从回答简单的问题到提供对广泛主题的深入解释和讨论。 Visual ChatGPT 能够根据收到的输入生成类似人类的文本，使其能够进行听起来自然的对话，并提供连贯且与手头主题相关的响应。
+
+Visual ChatGPT 能够处理和理解大量文本和图像。作为一种语言模型，Visual ChatGPT 不能直接读取图像，但它有一系列工具来完成不同的视觉任务。每张图片都会有一个文件名，格式为“image/xxx.png”，Visual ChatGPT可以调用不同的工具来间接理解图片。在谈论图片时，Visual ChatGPT 对文件名的要求非常严格，绝不会伪造不���在的文件。在使用工具生成新的图像文件时，Visual ChatGPT也知道图像可能与用户需求不一样，会使用其他视觉问答工具或描述工具来观察真实图像。 Visual ChatGPT 能够按顺序使用工具，并且忠于工具观察输出，而不是伪造图像内容和图像文件名。如果生成新图像，它将记得提供上次工具观察的文件名。
+
+Human 可能会向 Visual ChatGPT 提供带有描述的新图形。描述帮助 Visual ChatGPT 理解这个图像，但 Visual ChatGPT 应该使用工具来完成以下任务，而不是直接从描述中想象。有些工具将会返回英文描述，但你对用户的聊天应当采用中文。
+
+总的来说，Visual ChatGPT 是一个强大的可视化对话辅助工具，可以帮助处理范围广泛的任务，并提供关于范围广泛的主题的有价值的见解和信息。
+
+工具列表:
+------
+
+Visual ChatGPT 可以使用这些工具:"""
+
+VISUAL_CHATGPT_FORMAT_INSTRUCTIONS_CN = """用户使用中文和你进行聊天，但是工具的参数应当使用英文。如果要调用工具，你必须遵循如下格式:
+
+```
+Thought: Do I need to use a tool? Yes
+Action: the action to take, should be one of [{tool_names}]
+Action Input: the input to the action
+Observation: the result of the action
+```
+
+当你不再需要继续调用工具，而是对观察结果进行总结回复时，你必须使用如下格式：
+
+
+```
+Thought: Do I need to use a tool? No
+{ai_prefix}: [your response here]
+```
+"""
+
+VISUAL_CHATGPT_SUFFIX_CN = """你对文件名的正确性非常严格，而且永远不会伪造不存在的文件。
+
+开始!
+
+因为Visual ChatGPT是一个文本语言模型，必须使用工具去观察图片而不是依靠想象。
+推理想法和观察结果只对Visual ChatGPT可见，需要记得在最终回复时把重要的信息重复给用户，你只能给用户返回中文句子。我们一步一步思考。在你使用工具时，工具的参数只能是英文。
+
+聊天历史:
+{chat_history}
+
+新输入: {input}
+Thought: Do I need to use a tool? {agent_scratchpad}
+"""
+
+
+
+def cut_dialogue_history(history_memory, keep_last_n_words=500):
+    if history_memory is None or len(history_memory) == 0:
+        return history_memory
+    tokens = history_memory.split()
+    n_tokens = len(tokens)
+    print(f"history_memory:{history_memory}, n_tokens: {n_tokens}")
+    if n_tokens < keep_last_n_words:
+        return history_memory
+    paragraphs = history_memory.split('\n')
+    last_n_tokens = n_tokens
+    while last_n_tokens >= keep_last_n_words:
+        last_n_tokens -= len(paragraphs[0].split(' '))
+        paragraphs = paragraphs[1:]
+    return '\n' + '\n'.join(paragraphs)
+
+
+class ConversationBot:
+    def __init__(self, load_dict):
+        # load_dict = {'VisualQuestionAnswering':'cuda:0', 'ImageCaptioning':'cuda:1',...}
+        print(f"Initializing VisualChatGPT, load_dict={load_dict}")
+        if 'ImageCaptioning' not in load_dict:
+            raise ValueError("You have to load ImageCaptioning as a basic function for i-GPT")
+        # if 'SegmentAnything' not in load_dict:
+        #     raise ValueError("You have to load SegmentAnything as a basic function for i-GPT")
+
+        self.models = {}
+        self.uploaded_image_filename = None
+        # self.segmented_image_filename = None
+        self.history_mask = None
+        self.load_dict = load_dict
+        # self.llm = None
+        # Load Basic Foundation Models
+        for class_name, device in load_dict.items():
+            self.models[class_name] = globals()[class_name](device=device)
+        # self.models['models'] = self.models
+
+        # Load Template Foundation Models
+        for class_name, module in globals().items():
+            if getattr(module, 'template_model', False):
+                template_required_names = {k for k in inspect.signature(module.__init__).parameters.keys() if k!='self'}
+                loaded_names = set([type(e).__name__ for e in self.models.values()])
+                if template_required_names.issubset(loaded_names):
+                    self.models[class_name] = globals()[class_name](
+                        **{name: self.models[name] for name in template_required_names})
+                # elif 'models' in template_required_names:
+                #     self.models[class_name] = globals()[class_name](
+                #         **{name: self.models[name] for name in template_required_names})
+
+        self.tools = []
+        for instance in self.models.values():
+            for e in dir(instance):
+                if e.startswith('inference'):
+                    func = getattr(instance, e)
+                    self.tools.append(Tool(name=func.name, description=func.description, func=func))
+        self.llm = None
+        self.memory = ConversationBufferMemory(memory_key="chat_history", output_key='output')
+        # self.first_init=True
+        self.audio_model = None
+
+    def init_agent(self):
+        self.memory.clear() #clear previous history
+        self.reset()
+        self.llm = OpenAI(temperature=0)
+        self.agent = initialize_agent(
+            self.tools,
+            self.llm,
+            agent="conversational-react-description",
+            verbose=True,
+            memory=self.memory,
+            return_intermediate_steps=True,
+            agent_kwargs={'prefix': VISUAL_CHATGPT_PREFIX, 'format_instructions': VISUAL_CHATGPT_FORMAT_INSTRUCTIONS,
+                          'suffix': VISUAL_CHATGPT_SUFFIX}, )
+
+    def run_text(self, text, state):
+        # print(f'text = {text}')
+        self.agent.memory.buffer = cut_dialogue_history(self.agent.memory.buffer, keep_last_n_words=500)
+        try:
+            print(f'text = {text}')
+            res = self.agent({"input": text.strip()})
+            print('ab'* 30)
+            print(res['output'])
+            print('cd'* 30)
+        except Exception as err:
+            # Human_prompt = text
+            # self.agent.memory.buffer = self.agent.memory.buffer + Human_prompt + ' AI: ' + AI_prompt
+            state += [(text, 'I can not understand your instruction. Could you provide more information?')]
+            print(err)
+            return state, state
+        
+        res['output'] = res['output'].replace("\\", "/")
+        # response = re.sub('(tmp_files/[-\w]*.[png|mp4])', lambda m: f'![](file={m.group(0)})*{m.group(0)}*', res['output'])
+
+        # print("res['output'] = ", res['output'])
+        # response = re.sub('(tmp_files/[-\w]*.(png|mp4))', replace_path, res['output'])
+        pattern = re.compile('(image/[-\\w]*.(png|mp4))')
+        out_filenames = pattern.findall(res['output'])
+        response = res['output']
+        state = state + [(text, response)]
+        for f in out_filenames:
+            state = state + [(None, f'{f[0]} is as following: ')]
+            state = state + [(None, (f[0], ))]
+        # if len(out_filenames) > 1:
+        #     state = state + [(None, (out_filenames[-1][0], ))]
+            # print('out_filename[-1][0] = ', out_filenames[-1][0])
+        print(f"\nProcessed run_text, Input text: {text}\nCurrent state: {state}\n"
+              f"Current Memory: {self.agent.memory.buffer}")
+        return state, state
+    
+    def run_audio(self, audio_path, state):
+        print(f'audio_path = {audio_path}')
+        if self.audio_model is None:
+            self.audio_model = whisper.load_model("small").to('cuda:0')
+        text = self.audio_model.transcribe(audio_path)["text"]
+        res = self.run_text(text, state)
+        print(f"\nProcessed run_audio, Input transcribed audio: {text}\nCurrent state: {state}\n"
+              f"Current Memory: {self.agent.memory.buffer}")
+        return res[0], res[1]
+
+    def upload_image(self, image, state, txt):
+        self.reset()
+        img = image['image']
+        image_filename = os.path.join('image/', f"{str(uuid.uuid4())[:6]}.png")
+        self.uploaded_image_filename=image_filename
+        img = img.convert('RGB')
+        img.save(image_filename, "PNG")
+        # print(f"Resize image form {width}x{height} to {width_new}x{height_new}")
+        # let some foundation models preprocess image
+        # NEED_PREPROCESSING_LIST = ["SegmentAnything", "ImageOCRRecognition"]
+        # for model_name in NEED_PREPROCESSING_LIST:
+        #     if model_name in self.models.keys():
+        #          self.models[model_name].preprocess(np.array(img), image_filename)
+
+        description = self.models['ImageCaptioning'].inference(image_filename)
+        # description = 'Debug'
+
+        Human_prompt = f'\nHuman: provide a figure named {image_filename}. The description is: {description}. This information helps you to understand this image, but you should use tools to finish following tasks, rather than directly imagine from my description. If you understand, say \"Received\". \n'
+        AI_prompt = "Received. "
+        self.agent.memory.buffer = self.agent.memory.buffer + Human_prompt + ' AI: ' + AI_prompt
+        state = state + [(f"![](file={image_filename})*{image_filename}*", AI_prompt)]
+        print(f"\nProcessed upload_image, Input image: {image_filename}\nCurrent state: {state}\n"
+              f"Current Memory: {self.agent.memory.buffer}")
+        return state, state, f'{txt} {image_filename} ', gr.update(visible=True), gr.update(visible=True)
+    
+    def upload_video(self, video_path, state, txt):
+        # self.cur_file = video_path
+        vid_name = os.path.basename(video_path)
+        # vid_name = gen_new_name(vid_name, '', vid_name.split('.')[-1])
+        new_video_path = os.path.join('./image/', f"{str(uuid.uuid4())[:6]}.mp4")
+        new_video_path = gen_new_name(new_video_path, '', vid_name.split('.')[-1])
+        shutil.copy(video_path, new_video_path)
+        
+        if "VideoCaption" in self.models.keys():
+            description = self.models['VideoCaption'].inference(new_video_path)
+        else:
+            description = 'A video.'
+        Human_prompt = f'\nHuman: provide a video named {new_video_path}. The description is: {description}. This information helps you to understand this video, but you should use tools to finish following tasks, rather than directly imagine from my description. If you understand, say \"Received\". \n'
+        AI_prompt = f"Received video: {new_video_path} "
+        self.agent.memory.buffer = self.agent.memory.buffer + Human_prompt + 'AI: ' + AI_prompt
+        # state = state + [(f"![](file={new_video_path})*{new_video_path}*", AI_prompt)]
+        # state = state + [(f"![](file={video_path})*{new_video_path}*", AI_prompt)]
+        state = state + [((new_video_path, ), AI_prompt)]
+        # print('exists = ', os.path.exists("./tmp_files/1e7f_f4236666_tmp.mp4"))
+        print(f"\nProcessed upload_video, Input video: {new_video_path}\nCurrent state: {state}\n"
+              f"Current Memory: {self.agent.memory.buffer}")
+        return state, state, f'{txt} {new_video_path} '
+
+    def blend_mask(self, img, mask):
+        mask = mask.astype(np.uint8)
+        transparency_ratio = mask.astype(np.float32) / 3
+        transparency_ratio = transparency_ratio[:, :, np.newaxis]
+        mask = mask[:, :, np.newaxis] * 255
+        mask= mask.repeat(3, axis=2)
+        mask[:,:,0] = 0
+        mask[:,:,2] = 0
+        new_img_arr = np.array(img) * (1 - transparency_ratio) + mask * transparency_ratio
+        new_img_arr = np.clip(new_img_arr, 0, 255).astype(np.uint8)
+        # print(new_img_arr.shape)
+        return Image.fromarray(new_img_arr)
+
+    def process_image(self, image, state):
+        img = Image.open(self.uploaded_image_filename).convert('RGB')
+        # img = image['image'].convert('RGB')
+        mask = image['mask'].convert('L')
+        mask = np.array(mask, dtype=np.uint8)
+
+        Human_prompt="Please process this image based on given mask."
+        if self.uploaded_image_filename is None:
+            AI_prompt = "Please upload an image for processing."
+            state += [(Human_prompt, AI_prompt)]
+            return state, state, None
+        if mask.sum() == 0:
+            AI_prompt = "You can click the image in the right and ask me some questions."
+            state += [(Human_prompt, AI_prompt)]
+            return state, state, image['image']
+        
+        if self.history_mask is None:
+            self.history_mask = mask
+        else:
+            self.history_mask = np.logical_or(self.history_mask, mask)
+        
+        if 'SegmentAnything' in self.models.keys():
+            self.models['SegmentAnything'].clicked_region = self.history_mask
+        if 'ImageOCRRecognition' in self.models.keys():
+            self.models['ImageOCRRecognition'].clicked_region = mask
+
+        # self.models['SegmentAnything'].mask = self.history_mask
+        # history_mask = self.history_mask.astype(np.uint8) * 255
+        res_mask = self.models['SegmentAnything'].segment_by_mask(self.history_mask)
+        
+        img = self.blend_mask(img, res_mask)
+
+        AI_prompt = f"I have finished processing. Now, you can ask me some questions."
+        state = state + [(Human_prompt, AI_prompt)]
+        # AI_prompt = f"Received. I found {ocr_text} in this position. The sgemented figure is named {seg_filename}."
+        self.agent.memory.buffer = self.agent.memory.buffer + Human_prompt + ' AI: ' + AI_prompt
+        # state = state + [(Human_prompt, f"![](file={seg_filename})*{AI_prompt}*")]
+        # print()
+        print(f"\nProcessed run_image, Input image: {self.uploaded_image_filename}\nCurrent state: {state}\n"
+              f"Current Memory: {self.agent.memory.buffer}")
+        return state, state, img
+    
+    def reset(self, clear_history_memory=False):
+        print('reset the model cache.')
+        NEED_RESET_LIST = ['SegmentAnything', 'ImageOCRRecognition']
+        for model_name in NEED_RESET_LIST:
+            if model_name in self.models.keys():
+                self.models[model_name].reset()    
+
+        self.history_mask = None 
+        self.uploaded_image_filename = None
+        if clear_history_memory:
+            self.agent.memory.clear()
+        return None

iChat/models/__init__.py

@@ -0,0 +1,37 @@
+# from .image import (MaskFormer, ImageEditing, InstructPix2Pix, \
+#     Text2Image, ImageCaptioning, Image2Canny, CannyText2Image, \
+#     Image2Line, LineText2Image, Image2Hed, HedText2Image, Image2Scribble, \
+#     ScribbleText2Image, Image2Pose, PoseText2Image, SegText2Image, \
+#     Image2Depth, DepthText2Image, Image2Normal, NormalText2Image, \
+#     VisualQuestionAnswering, InfinityOutPainting, \
+#     SegmentAnything, InpaintMaskedAnything, ExtractMaskedAnything, \
+#     ReplaceMaskedAnything, ImageOCRRecognition)
+
+from .husky import HuskyVQA
+
+from .video import (ActionRecognition, DenseCaption, VideoCaption, 
+                    Summarization, GenerateTikTokVideo)
+
+from .lang import SimpleLanguageModel
+
+from .inpainting import LDMInpainting
+
+# __all__ = [
+#     'MaskFormer', 'ImageEditing', 'InstructPix2Pix', \
+#     'Text2Image', 'ImageCaptioning', 'Image2Canny', 'CannyText2Image', \
+#     'Image2Line', 'LineText2Image', 'Image2Hed', 'HedText2Image', \
+#     'Image2Scribble', 'ScribbleText2Image', 'Image2Pose', 'PoseText2Image', \
+#     'SegText2Image', 'Image2Depth', 'DepthText2Image', 'Image2Normal', \
+#     'NormalText2Image', 'VisualQuestionAnswering', 'InfinityOutPainting', \
+#     'SegmentAnything', 'InpaintMaskedAnything', 'ExtractMaskedAnything', \
+#     'ReplaceMaskedAnything', 'ImageOCRRecognition', "SimpleLanguageModel", \
+#     'ActionRecognition', 'DenseCaption', 'VideoCaption', 'Summarization', \
+#     'GenerateTikTokVideo'
+# ]
+
+__all__ = [
+    'HuskyVQA', "SimpleLanguageModel", 'GenerateTikTokVideo', \
+    'LDMInpainting',
+    'ActionRecognition', 'DenseCaption', 'VideoCaption', 'Summarization' 
+]
+

iChat/models/grit_model.py

@@ -0,0 +1,46 @@
+import os
+import sys
+
+from .grit_src.image_dense_captions import image_caption_api, init_demo, dense_pred_to_caption, dense_pred_to_caption_only_name
+from detectron2.data.detection_utils import read_image
+
+class DenseCaptioning():
+    def __init__(self, device):
+        self.device = device
+        self.demo =  None
+
+
+    def initialize_model(self):
+        self.demo = init_demo(self.device)
+
+    def image_dense_caption_debug(self, image_src):
+        dense_caption = """
+        1. the broccoli is green, [0, 0, 333, 325]; 
+        2. a piece of broccoli, [0, 147, 143, 324]; 
+        3. silver fork on plate, [4, 547, 252, 612];
+        """
+        return dense_caption
+    
+    def image_dense_caption(self, image_src):
+        dense_caption = image_caption_api(image_src, self.device)
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
+        print("Step2, Dense Caption:\n")
+        print(dense_caption)
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
+        return dense_caption
+    
+    def run_caption_api(self,image_src):
+        img = read_image(image_src, format="BGR")
+        print(img.shape)
+        predictions, visualized_output = self.demo.run_on_image(img)
+        new_caption = dense_pred_to_caption_only_name(predictions)
+        return new_caption
+
+    def run_caption_tensor(self,img):
+        # img = read_image(image_src, format="BGR")
+        # print(img.shape)
+        predictions, visualized_output = self.demo.run_on_image(img)
+        new_caption = dense_pred_to_caption_only_name(predictions)
+        return new_caption
+    
+

iChat/models/grit_src/configs/Base.yaml

@@ -0,0 +1,77 @@
+MODEL:
+  META_ARCHITECTURE: "GRiT"
+  MASK_ON: True
+  PROPOSAL_GENERATOR:
+    NAME: "CenterNet"
+  FPN:
+    IN_FEATURES: ["layer3", "layer4", "layer5"]
+  PIXEL_MEAN: [123.675, 116.280, 103.530]
+  PIXEL_STD: [58.395, 57.12, 57.375]
+  ROI_HEADS:
+    NAME: GRiTROIHeadsAndTextDecoder
+    IN_FEATURES: ["p3", "p4", "p5"]
+    IOU_THRESHOLDS: [0.6]
+    NUM_CLASSES: 1
+    SCORE_THRESH_TEST: 0.02
+    NMS_THRESH_TEST: 0.5
+    OBJECT_FEAT_POOLER_RES: 14
+  ROI_BOX_CASCADE_HEAD:
+    IOUS: [0.6, 0.7, 0.8]
+  ROI_BOX_HEAD:
+    NAME: "FastRCNNConvFCHead"
+    NUM_FC: 2
+    POOLER_RESOLUTION: 7
+    CLS_AGNOSTIC_BBOX_REG: True
+    MULT_PROPOSAL_SCORE: True
+  ROI_MASK_HEAD:
+    NAME: "MaskRCNNConvUpsampleHead"
+    NUM_CONV: 4
+    POOLER_RESOLUTION: 14
+    CLS_AGNOSTIC_MASK: True
+  CENTERNET:
+    NUM_CLASSES: 1
+    REG_WEIGHT: 1.
+    NOT_NORM_REG: True
+    ONLY_PROPOSAL: True
+    WITH_AGN_HM: True
+    INFERENCE_TH: 0.0001
+    PRE_NMS_TOPK_TRAIN: 4000
+    POST_NMS_TOPK_TRAIN: 2000
+    PRE_NMS_TOPK_TEST: 1000
+    POST_NMS_TOPK_TEST: 256
+    NMS_TH_TRAIN: 0.9
+    NMS_TH_TEST: 0.9
+    POS_WEIGHT: 0.5
+    NEG_WEIGHT: 0.5
+    IGNORE_HIGH_FP: 0.85
+DATASETS:
+  TRAIN: ("coco_2017_train",)
+  TEST: ("coco_2017_val",)
+DATALOADER:
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  DATASET_RATIO: [1]
+  DATASET_INPUT_SIZE: [1024]
+  DATASET_INPUT_SCALE: [[0.1, 2.0]]
+  FILTER_EMPTY_ANNOTATIONS: False
+  NUM_WORKERS: 8
+TEST:
+  DETECTIONS_PER_IMAGE: 256
+SOLVER:
+  LR_SCHEDULER_NAME: "WarmupCosineLR"
+  CHECKPOINT_PERIOD: 10000
+  WARMUP_ITERS: 1000
+  WARMUP_FACTOR: 0.001
+  USE_CUSTOM_SOLVER: True
+  OPTIMIZER: "ADAMW"
+  MAX_ITER: 180000
+  IMS_PER_BATCH: 64
+  BASE_LR: 0.00008
+  VIT_LAYER_DECAY: True
+  CLIP_GRADIENTS:
+    ENABLED: True
+INPUT:
+  FORMAT: RGB
+  CUSTOM_AUG: EfficientDetResizeCrop
+  TRAIN_SIZE: 640
+USE_ACT_CHECKPOINT: True
+VERSION: 2

iChat/models/grit_src/configs/GRiT_B_DenseCap.yaml

@@ -0,0 +1,20 @@
+_BASE_: "Base.yaml"
+MODEL:
+  TRAIN_TASK: ["DenseCap"]
+  TEST_TASK: "DenseCap"
+  MASK_ON: False
+  ROI_HEADS:
+    SOFT_NMS_ENABLED: False
+  BEAM_SIZE: 1
+  WEIGHTS: "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_base.pth"
+  BACKBONE:
+    NAME: build_vit_fpn_backbone
+  VIT_LAYERS: 12
+SOLVER:
+  VIT_LAYER_DECAY_RATE: 0.7
+DATASETS:
+  TRAIN: ("vg_train",)
+  TEST: ("vg_test",)
+DATALOADER:
+  DATASET_BS: 2
+OUTPUT_DIR: "./output/GRiT_B_DenseCap"

iChat/models/grit_src/configs/GRiT_B_DenseCap_ObjectDet.yaml

@@ -0,0 +1,23 @@
+_BASE_: "Base.yaml"
+MODEL:
+  TRAIN_TASK: ["ObjectDet", "DenseCap"]
+  TEST_TASK: "DenseCap" # DenseCap or ObjectDet: Choose one for testing
+  MASK_ON: True
+  ROI_HEADS:
+    SOFT_NMS_ENABLED: False
+  BEAM_SIZE: 1
+  WEIGHTS: "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_base.pth"
+  BACKBONE:
+    NAME: build_vit_fpn_backbone
+  VIT_LAYERS: 12
+SOLVER:
+  VIT_LAYER_DECAY_RATE: 0.7
+DATASETS:
+  TRAIN: ("GRiT_coco2017_train", "vg_train")
+  TEST: ("coco_2017_test-dev",)
+DATALOADER:
+  DATASET_RATIO: [1, 1]
+  DATASET_BS: 2
+  DATASET_INPUT_SIZE: [1024, 1024]
+  DATASET_INPUT_SCALE: [[0.1, 2.0], [0.1, 2.0]]
+OUTPUT_DIR: "./output/GRiT_B_DenseCap_ObjectDet"

iChat/models/grit_src/configs/GRiT_B_ObjectDet.yaml

@@ -0,0 +1,20 @@
+_BASE_: "Base.yaml"
+MODEL:
+  TRAIN_TASK: ["ObjectDet"]
+  TEST_TASK: "ObjectDet"
+  MASK_ON: True
+  ROI_HEADS:
+    SOFT_NMS_ENABLED: True
+  BEAM_SIZE: 3
+  WEIGHTS: "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_base.pth"
+  BACKBONE:
+    NAME: build_vit_fpn_backbone
+  VIT_LAYERS: 12
+SOLVER:
+  VIT_LAYER_DECAY_RATE: 0.7
+DATASETS:
+  TRAIN: ("GRiT_coco2017_train",)
+  TEST: ("coco_2017_val",)
+DATALOADER:
+  DATASET_BS: 2
+OUTPUT_DIR: "./output/GRiT_B_ObjectDet"

iChat/models/grit_src/configs/GRiT_H_ObjectDet.yaml

@@ -0,0 +1,21 @@
+_BASE_: "Base.yaml"
+MODEL:
+  TRAIN_TASK: ["ObjectDet"]
+  TEST_TASK: "ObjectDet"
+  MASK_ON: True
+  ROI_HEADS:
+    SOFT_NMS_ENABLED: True
+  BEAM_SIZE: 3
+  WEIGHTS: "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_huge_p14to16.pth"
+  BACKBONE:
+    NAME: build_vit_fpn_backbone_huge
+  VIT_LAYERS: 32
+SOLVER:
+  MAX_ITER: 135000
+  VIT_LAYER_DECAY_RATE: 0.9
+DATASETS:
+  TRAIN: ("GRiT_coco2017_train",)
+  TEST: ("coco_2017_val",)
+DATALOADER:
+  DATASET_BS: 1
+OUTPUT_DIR: "./output/GRiT_H_ObjectDet"

iChat/models/grit_src/configs/GRiT_L_ObjectDet.yaml

@@ -0,0 +1,20 @@
+_BASE_: "Base.yaml"
+MODEL:
+  TRAIN_TASK: ["ObjectDet"]
+  TEST_TASK: "ObjectDet"
+  MASK_ON: True
+  ROI_HEADS:
+    SOFT_NMS_ENABLED: True
+  BEAM_SIZE: 3
+  WEIGHTS: "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_large.pth"
+  BACKBONE:
+    NAME: build_vit_fpn_backbone_large
+  VIT_LAYERS: 24
+SOLVER:
+  VIT_LAYER_DECAY_RATE: 0.8
+DATASETS:
+  TRAIN: ("GRiT_coco2017_train",)
+  TEST: ("coco_2017_val",)
+DATALOADER:
+  DATASET_BS: 1
+OUTPUT_DIR: "./output/GRiT_L_ObjectDet"

iChat/models/grit_src/grit/__init__.py

@@ -0,0 +1,7 @@
+from .modeling.meta_arch import grit
+from .modeling.roi_heads import grit_roi_heads
+from .modeling.backbone import vit
+
+from .data.datasets import object365
+from .data.datasets import vg
+from .data.datasets import grit_coco

iChat/models/grit_src/grit/config.py

@@ -0,0 +1,50 @@
+from detectron2.config import CfgNode as CN
+
+
+def add_grit_config(cfg):
+    _C = cfg
+
+    _C.MODEL.BEAM_SIZE = 1
+    _C.MODEL.TRAIN_TASK = ["ObjectDet", "DenseCap"]
+    _C.MODEL.TEST_TASK = "DenseCap"  # This can be varied if the model is jointly trained on multiple tasks
+
+    _C.MODEL.ROI_BOX_HEAD.USE_BIAS = 0.0 # >= 0: not use
+    _C.MODEL.ROI_BOX_HEAD.MULT_PROPOSAL_SCORE = False
+
+    _C.MODEL.ROI_HEADS.MASK_WEIGHT = 1.0
+    _C.MODEL.ROI_HEADS.OBJECT_FEAT_POOLER_RES = 14
+    _C.MODEL.ROI_HEADS.SOFT_NMS_ENABLED = False
+
+    # Backbones
+    _C.MODEL.VIT_LAYERS = 12
+
+    # Text Decoder
+    _C.TEXT_DECODER = CN()
+    _C.TEXT_DECODER.VOCAB_SIZE = 30522
+    _C.TEXT_DECODER.HIDDEN_SIZE = 768
+    _C.TEXT_DECODER.NUM_LAYERS = 6
+    _C.TEXT_DECODER.ATTENTION_HEADS = 12
+    _C.TEXT_DECODER.FEEDFORWARD_SIZE = 768 * 4
+    
+    # Multi-dataset dataloader
+    _C.DATALOADER.DATASET_RATIO = [1, 1]  # sample ratio
+    _C.DATALOADER.DATASET_BS = 1
+    _C.DATALOADER.DATASET_INPUT_SIZE = [1024, 1024]
+    _C.DATALOADER.DATASET_INPUT_SCALE = [(0.1, 2.0), (0.1, 2.0)]
+    _C.DATALOADER.DATASET_MIN_SIZES = [(640, 800), (640, 800)]
+    _C.DATALOADER.DATASET_MAX_SIZES = [1333, 1333]
+    
+    _C.SOLVER.USE_CUSTOM_SOLVER = True
+    _C.SOLVER.OPTIMIZER = 'ADAMW'
+    _C.SOLVER.VIT_LAYER_DECAY = True
+    _C.SOLVER.VIT_LAYER_DECAY_RATE = 0.7
+
+    _C.INPUT.CUSTOM_AUG = 'EfficientDetResizeCrop'
+    _C.INPUT.TRAIN_SIZE = 1024
+    _C.INPUT.TEST_SIZE = 1024
+    _C.INPUT.SCALE_RANGE = (0.1, 2.)
+    # 'default' for fixed short / long edge
+    _C.INPUT.TEST_INPUT_TYPE = 'default' 
+
+    _C.FIND_UNUSED_PARAM = True
+    _C.USE_ACT_CHECKPOINT = True

iChat/models/grit_src/grit/custom_solver.py

@@ -0,0 +1,88 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# Modified by Jialian Wu from https://github.com/facebookresearch/Detic/blob/main/detic/custom_solver.py
+import itertools
+from typing import Any, Callable, Dict, Iterable, List, Set, Type, Union
+import torch
+
+from detectron2.config import CfgNode
+
+from detectron2.solver.build import maybe_add_gradient_clipping
+
+
+def build_custom_optimizer(cfg: CfgNode, model: torch.nn.Module) -> torch.optim.Optimizer:
+    params: List[Dict[str, Any]] = []
+    memo: Set[torch.nn.parameter.Parameter] = set()
+    optimizer_type = cfg.SOLVER.OPTIMIZER
+
+    for key, value in model.named_parameters(recurse=True):
+        if not value.requires_grad:
+            continue
+        # Avoid duplicating parameters
+        if value in memo:
+            continue
+        memo.add(value)
+        lr = cfg.SOLVER.BASE_LR
+        weight_decay = cfg.SOLVER.WEIGHT_DECAY
+
+        if cfg.SOLVER.VIT_LAYER_DECAY:
+            lr = lr * get_vit_lr_decay_rate(key, cfg.SOLVER.VIT_LAYER_DECAY_RATE, cfg.MODEL.VIT_LAYERS)
+
+        param = {"params": [value], "lr": lr}
+        if optimizer_type != 'ADAMW':
+            param['weight_decay'] = weight_decay
+        params += [param]
+
+    def maybe_add_full_model_gradient_clipping(optim):  # optim: the optimizer class
+        # detectron2 doesn't have full model gradient clipping now
+        clip_norm_val = cfg.SOLVER.CLIP_GRADIENTS.CLIP_VALUE
+        enable = (
+            cfg.SOLVER.CLIP_GRADIENTS.ENABLED
+            and cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model"
+            and clip_norm_val > 0.0
+        )
+
+        class FullModelGradientClippingOptimizer(optim):
+            def step(self, closure=None):
+                all_params = itertools.chain(*[x["params"] for x in self.param_groups])
+                torch.nn.utils.clip_grad_norm_(all_params, clip_norm_val)
+                super().step(closure=closure)
+
+        return FullModelGradientClippingOptimizer if enable else optim
+
+    
+    if optimizer_type == 'SGD':
+        optimizer = maybe_add_full_model_gradient_clipping(torch.optim.SGD)(
+            params, cfg.SOLVER.BASE_LR, momentum=cfg.SOLVER.MOMENTUM, 
+            nesterov=cfg.SOLVER.NESTEROV
+        )
+    elif optimizer_type == 'ADAMW':
+        optimizer = maybe_add_full_model_gradient_clipping(torch.optim.AdamW)(
+            params, cfg.SOLVER.BASE_LR, 
+            weight_decay=cfg.SOLVER.WEIGHT_DECAY
+        )
+    else:
+        raise NotImplementedError(f"no optimizer type {optimizer_type}")
+    if not cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model":
+        optimizer = maybe_add_gradient_clipping(cfg, optimizer)
+    return optimizer
+
+
+def get_vit_lr_decay_rate(name, lr_decay_rate=1.0, num_layers=12):
+    """
+    Calculate lr decay rate for different ViT blocks.
+    Args:
+        name (string): parameter name.
+        lr_decay_rate (float): base lr decay rate.
+        num_layers (int): number of ViT blocks.
+
+    Returns:
+        lr decay rate for the given parameter.
+    """
+    layer_id = num_layers + 1
+    if name.startswith("backbone"):
+        if ".pos_embed" in name or ".patch_embed" in name:
+            layer_id = 0
+        elif ".blocks." in name and ".residual." not in name:
+            layer_id = int(name[name.find(".blocks.") :].split(".")[2]) + 1
+
+    return lr_decay_rate ** (num_layers + 1 - layer_id)

iChat/models/grit_src/grit/data/custom_build_augmentation.py

@@ -0,0 +1,44 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from detectron2.data import transforms as T
+from .transforms.custom_augmentation_impl import EfficientDetResizeCrop
+
+
+def build_custom_augmentation(cfg, is_train, scale=None, size=None, \
+    min_size=None, max_size=None):
+    """
+    Create a list of default :class:`Augmentation` from config.
+    Now it includes resizing and flipping.
+
+    Returns:
+        list[Augmentation]
+    """
+    if cfg.INPUT.CUSTOM_AUG == 'ResizeShortestEdge':
+        if is_train:
+            min_size = cfg.INPUT.MIN_SIZE_TRAIN if min_size is None else min_size
+            max_size = cfg.INPUT.MAX_SIZE_TRAIN if max_size is None else max_size
+            sample_style = cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING
+        else:
+            min_size = cfg.INPUT.MIN_SIZE_TEST
+            max_size = cfg.INPUT.MAX_SIZE_TEST
+            sample_style = "choice"
+        augmentation = [T.ResizeShortestEdge(min_size, max_size, sample_style)]
+    elif cfg.INPUT.CUSTOM_AUG == 'EfficientDetResizeCrop':
+        if is_train:
+            scale = cfg.INPUT.SCALE_RANGE if scale is None else scale
+            size = cfg.INPUT.TRAIN_SIZE if size is None else size
+        else:
+            scale = (1, 1)
+            size = cfg.INPUT.TEST_SIZE
+        augmentation = [EfficientDetResizeCrop(size, scale)]
+    else:
+        assert 0, cfg.INPUT.CUSTOM_AUG
+
+    if is_train:
+        augmentation.append(T.RandomFlip())
+    return augmentation
+
+
+build_custom_transform_gen = build_custom_augmentation
+"""
+Alias for backward-compatibility.
+"""

iChat/models/grit_src/grit/data/custom_dataset_dataloader.py

@@ -0,0 +1,250 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Jialian Wu from https://github.com/facebookresearch/Detic/blob/main/detic/data/custom_dataset_dataloader.py
+import operator
+import torch
+import torch.utils.data
+from detectron2.utils.comm import get_world_size
+
+from detectron2.config import configurable
+from torch.utils.data.sampler import BatchSampler, Sampler
+from detectron2.data.common import DatasetFromList, MapDataset
+from detectron2.data.dataset_mapper import DatasetMapper
+from detectron2.data.build import get_detection_dataset_dicts, build_batch_data_loader
+from detectron2.data.samplers import TrainingSampler
+from detectron2.data.build import worker_init_reset_seed, print_instances_class_histogram
+from detectron2.data.build import filter_images_with_only_crowd_annotations
+from detectron2.data.build import filter_images_with_few_keypoints
+from detectron2.data.build import check_metadata_consistency
+from detectron2.data.catalog import MetadataCatalog, DatasetCatalog
+from detectron2.utils import comm
+import itertools
+from typing import Optional
+
+
+def _custom_train_loader_from_config(cfg, mapper=None, *, dataset=None, sampler=None):
+    sampler_name = cfg.DATALOADER.SAMPLER_TRAIN
+    if 'MultiDataset' in sampler_name:
+        dataset_dicts = get_detection_dataset_dicts_with_source(
+            cfg.DATASETS.TRAIN,
+            filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
+            min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
+            if cfg.MODEL.KEYPOINT_ON else 0,
+            proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
+        )
+    else:
+        dataset_dicts = get_detection_dataset_dicts(
+            cfg.DATASETS.TRAIN,
+            filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
+            min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
+            if cfg.MODEL.KEYPOINT_ON else 0,
+            proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
+        )
+
+    if mapper is None:
+        mapper = DatasetMapper(cfg, True)
+
+    if sampler is not None:
+        pass
+    elif sampler_name == "TrainingSampler":
+        sampler = TrainingSampler(len(dataset))
+    elif sampler_name == "MultiDatasetSampler":
+        sampler = MultiDatasetSampler(
+            dataset_dicts,
+            dataset_ratio=cfg.DATALOADER.DATASET_RATIO,
+        )
+    else:
+        raise ValueError("Unknown training sampler: {}".format(sampler_name))
+
+    return {
+        "dataset": dataset_dicts,
+        "sampler": sampler,
+        "mapper": mapper,
+        "total_batch_size": cfg.SOLVER.IMS_PER_BATCH,
+        "num_workers": cfg.DATALOADER.NUM_WORKERS,
+        'dataset_bs': cfg.DATALOADER.DATASET_BS,
+        'num_datasets': len(cfg.DATASETS.TRAIN)
+    }
+
+
+@configurable(from_config=_custom_train_loader_from_config)
+def build_custom_train_loader(
+        dataset, *, mapper, sampler, 
+        total_batch_size=16,
+        num_workers=0,
+        num_datasets=1,
+        dataset_bs=1
+):
+
+    if isinstance(dataset, list):
+        dataset = DatasetFromList(dataset, copy=False)
+    if mapper is not None:
+        dataset = MapDataset(dataset, mapper)
+    if sampler is None:
+        sampler = TrainingSampler(len(dataset))
+    assert isinstance(sampler, torch.utils.data.sampler.Sampler)
+
+    return build_dataset_batch_data_loader(
+        dataset_bs,
+        dataset,
+        sampler,
+        total_batch_size,
+        num_datasets=num_datasets,
+        num_workers=num_workers,
+    )
+
+
+def build_dataset_batch_data_loader(
+    dataset_bs, dataset, sampler, total_batch_size, num_datasets, num_workers=0
+):
+
+    world_size = get_world_size()
+    assert (
+        total_batch_size > 0 and total_batch_size % world_size == 0
+    ), "Total batch size ({}) must be divisible by the number of gpus ({}).".format(
+        total_batch_size, world_size
+    )
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        sampler=sampler,
+        num_workers=num_workers,
+        batch_sampler=None,
+        collate_fn=operator.itemgetter(0),  # don't batch, but yield individual elements
+        worker_init_fn=worker_init_reset_seed,
+    )
+
+    if num_datasets > 1:
+        return MultiDatasets(data_loader, dataset_bs, num_datasets)
+    else:
+        return SingleDataset(data_loader, dataset_bs)
+
+
+def get_detection_dataset_dicts_with_source(
+    dataset_names, filter_empty=True, min_keypoints=0, proposal_files=None
+):
+    assert len(dataset_names)
+    dataset_dicts = [DatasetCatalog.get(dataset_name) for dataset_name in dataset_names]
+    for dataset_name, dicts in zip(dataset_names, dataset_dicts):
+        assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)
+    
+    for source_id, (dataset_name, dicts) in \
+        enumerate(zip(dataset_names, dataset_dicts)):
+        assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)
+        for d in dicts:
+            d['dataset_source'] = source_id
+
+        if "annotations" in dicts[0]:
+            try:
+                class_names = MetadataCatalog.get(dataset_name).thing_classes
+                check_metadata_consistency("thing_classes", dataset_name)
+                print_instances_class_histogram(dicts, class_names)
+            except AttributeError:  # class names are not available for this dataset
+                pass
+
+    assert proposal_files is None
+
+    dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts))
+
+    has_instances = "annotations" in dataset_dicts[0]
+    if filter_empty and has_instances:
+        dataset_dicts = filter_images_with_only_crowd_annotations(dataset_dicts)
+    if min_keypoints > 0 and has_instances:
+        dataset_dicts = filter_images_with_few_keypoints(dataset_dicts, min_keypoints)
+
+    return dataset_dicts
+
+
+class MultiDatasetSampler(Sampler):
+    def __init__(
+        self, 
+        dataset_dicts, 
+        dataset_ratio,
+        seed: Optional[int] = None,
+    ):
+        sizes = [0 for _ in range(len(dataset_ratio))]
+        for d in dataset_dicts:
+            sizes[d['dataset_source']] += 1
+        print('dataset sizes', sizes)
+        self.sizes = sizes
+        assert len(dataset_ratio) == len(sizes), \
+            'length of dataset ratio {} should be equal to number if dataset {}'.format(
+                len(dataset_ratio), len(sizes)
+            )
+        if seed is None:
+            seed = comm.shared_random_seed()
+        self._seed = int(seed)
+        self._rank = comm.get_rank()
+        self._world_size = comm.get_world_size()
+        
+        self.dataset_ids = torch.tensor(
+            [d['dataset_source'] for d in dataset_dicts], dtype=torch.long)
+        self.dataset_ratio = dataset_ratio
+
+        dataset_weight = [torch.ones(s) * max(sizes) / s * r / sum(dataset_ratio) \
+            for i, (r, s) in enumerate(zip(dataset_ratio, sizes))]
+        dataset_weight = torch.cat(dataset_weight)
+
+        self.weights = dataset_weight
+        self.sample_epoch_size = len(self.weights)
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(
+            self._infinite_indices(), start, None, self._world_size)
+
+    def _infinite_indices(self):
+        g = torch.Generator()
+        g.manual_seed(self._seed)
+        while True:
+            if len(self.dataset_ratio) > 1:
+                # multiple datasets
+                ids = torch.multinomial(
+                    self.weights, self.sample_epoch_size, generator=g,
+                    replacement=True)
+                nums = [(self.dataset_ids[ids] == i).sum().int().item() \
+                    for i in range(len(self.sizes))]
+                yield from ids
+            else:
+                # single dataset
+                yield from torch.randperm(self.sizes[0], generator=g).tolist()
+
+
+class SingleDataset(torch.utils.data.IterableDataset):
+    def __init__(self, dataset, batch_sizes):
+        self.dataset = dataset
+        self.batch_sizes = batch_sizes
+        self._buckets = [[] for _ in range(2)]
+
+    def __iter__(self):
+        for d in self.dataset:
+            w, h = d["width"], d["height"]
+            aspect_ratio_bucket_id = 0 if w > h else 1
+            bucket_id = aspect_ratio_bucket_id
+            bucket = self._buckets[bucket_id]
+            bucket.append(d)
+            if len(bucket) == self.batch_sizes:
+                yield bucket[:]
+                del bucket[:]
+
+
+class MultiDatasets(torch.utils.data.IterableDataset):
+    def __init__(self, dataset, batch_sizes, num_datasets):
+        self.dataset = dataset
+        self.batch_sizes = batch_sizes
+        self._buckets = [[] for _ in range(2 * num_datasets)]
+        self.iter_idx = 0
+        self.num_datasets = num_datasets
+
+    def __iter__(self):
+        for d in self.dataset:
+            w, h = d["width"], d["height"]
+            aspect_ratio_bucket_id = 0 if w > h else 1
+            bucket_id = d['dataset_source'] * 2 + aspect_ratio_bucket_id
+            bucket = self._buckets[bucket_id]
+            if len(bucket) < self.batch_sizes:
+                bucket.append(d)
+            selected_dataset = self.iter_idx % self.num_datasets
+            if len(bucket) == self.batch_sizes and selected_dataset == d['dataset_source']:
+                self.iter_idx += 1
+                yield bucket[:]
+                del bucket[:]

iChat/models/grit_src/grit/data/custom_dataset_mapper.py

@@ -0,0 +1,149 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# Modified by Jialian Wu from https://github.com/facebookresearch/Detic/blob/main/detic/data/custom_dataset_mapper.py
+import copy
+import numpy as np
+import torch
+
+from detectron2.config import configurable
+
+from detectron2.data import detection_utils as utils
+from detectron2.data import transforms as T
+from detectron2.data.dataset_mapper import DatasetMapper
+from .custom_build_augmentation import build_custom_augmentation
+from itertools import compress
+import logging
+
+__all__ = ["CustomDatasetMapper", "ObjDescription"]
+logger = logging.getLogger(__name__)
+
+
+class CustomDatasetMapper(DatasetMapper):
+    @configurable
+    def __init__(self, is_train: bool,
+        dataset_augs=[],
+        **kwargs):
+        if is_train:
+            self.dataset_augs = [T.AugmentationList(x) for x in dataset_augs]
+        super().__init__(is_train, **kwargs)
+
+    @classmethod
+    def from_config(cls, cfg, is_train: bool = True):
+        ret = super().from_config(cfg, is_train)
+        if is_train:
+            if cfg.INPUT.CUSTOM_AUG == 'EfficientDetResizeCrop':
+                dataset_scales = cfg.DATALOADER.DATASET_INPUT_SCALE
+                dataset_sizes = cfg.DATALOADER.DATASET_INPUT_SIZE
+                ret['dataset_augs'] = [
+                    build_custom_augmentation(cfg, True, scale, size) \
+                        for scale, size in zip(dataset_scales, dataset_sizes)]
+            else:
+                assert cfg.INPUT.CUSTOM_AUG == 'ResizeShortestEdge'
+                min_sizes = cfg.DATALOADER.DATASET_MIN_SIZES
+                max_sizes = cfg.DATALOADER.DATASET_MAX_SIZES
+                ret['dataset_augs'] = [
+                    build_custom_augmentation(
+                        cfg, True, min_size=mi, max_size=ma) \
+                        for mi, ma in zip(min_sizes, max_sizes)]
+        else:
+            ret['dataset_augs'] = []
+
+        return ret
+
+    def __call__(self, dataset_dict):
+        dataset_dict_out = self.prepare_data(dataset_dict)
+
+        # When augmented image is too small, do re-augmentation
+        retry = 0
+        while (dataset_dict_out["image"].shape[1] < 32 or dataset_dict_out["image"].shape[2] < 32):
+            retry += 1
+            if retry == 100:
+                logger.info('Retry 100 times for augmentation. Make sure the image size is not too small.')
+                logger.info('Find image information below')
+                logger.info(dataset_dict)
+            dataset_dict_out = self.prepare_data(dataset_dict)
+
+        return dataset_dict_out
+
+    def prepare_data(self, dataset_dict_in):
+        dataset_dict = copy.deepcopy(dataset_dict_in)
+        if 'file_name' in dataset_dict:
+            ori_image = utils.read_image(
+                dataset_dict["file_name"], format=self.image_format)
+        else:
+            ori_image, _, _ = self.tar_dataset[dataset_dict["tar_index"]]
+            ori_image = utils._apply_exif_orientation(ori_image)
+            ori_image = utils.convert_PIL_to_numpy(ori_image, self.image_format)
+        utils.check_image_size(dataset_dict, ori_image)
+
+        aug_input = T.AugInput(copy.deepcopy(ori_image), sem_seg=None)
+        if self.is_train:
+            transforms = \
+                self.dataset_augs[dataset_dict['dataset_source']](aug_input)
+        else:
+            transforms = self.augmentations(aug_input)
+        image, sem_seg_gt = aug_input.image, aug_input.sem_seg
+
+        image_shape = image.shape[:2]
+        dataset_dict["image"] = torch.as_tensor(
+            np.ascontiguousarray(image.transpose(2, 0, 1)))
+
+        if not self.is_train:
+            # USER: Modify this if you want to keep them for some reason.
+            dataset_dict.pop("annotations", None)
+            return dataset_dict
+
+        if "annotations" in dataset_dict:
+            if len(dataset_dict["annotations"]) > 0:
+                object_descriptions = [an['object_description'] for an in dataset_dict["annotations"]]
+            else:
+                object_descriptions = []
+            # USER: Modify this if you want to keep them for some reason.
+            for anno in dataset_dict["annotations"]:
+                if not self.use_instance_mask:
+                    anno.pop("segmentation", None)
+                if not self.use_keypoint:
+                    anno.pop("keypoints", None)
+
+            all_annos = [
+                (utils.transform_instance_annotations(
+                    obj, transforms, image_shape, 
+                    keypoint_hflip_indices=self.keypoint_hflip_indices,
+                ),  obj.get("iscrowd", 0))
+                for obj in dataset_dict.pop("annotations")
+            ]
+            annos = [ann[0] for ann in all_annos if ann[1] == 0]
+            instances = utils.annotations_to_instances(
+                annos, image_shape, mask_format=self.instance_mask_format
+            )
+
+            instances.gt_object_descriptions = ObjDescription(object_descriptions)
+            
+            del all_annos
+            if self.recompute_boxes:
+                instances.gt_boxes = instances.gt_masks.get_bounding_boxes()
+            dataset_dict["instances"] = utils.filter_empty_instances(instances)
+
+        return dataset_dict
+
+
+class ObjDescription:
+    def __init__(self, object_descriptions):
+        self.data = object_descriptions
+
+    def __getitem__(self, item):
+        assert type(item) == torch.Tensor
+        assert item.dim() == 1
+        if len(item) > 0:
+            assert item.dtype == torch.int64 or item.dtype == torch.bool
+            if item.dtype == torch.int64:
+                return ObjDescription([self.data[x.item()] for x in item])
+            elif item.dtype == torch.bool:
+                return ObjDescription(list(compress(self.data, item)))
+
+        return ObjDescription(list(compress(self.data, item)))
+
+    def __len__(self):
+        return len(self.data)
+
+    def __repr__(self):
+        return "ObjDescription({})".format(self.data)

iChat/models/grit_src/grit/data/datasets/grit_coco.py

@@ -0,0 +1,112 @@
+import logging
+import os
+from fvcore.common.timer import Timer
+from detectron2.structures import BoxMode
+from fvcore.common.file_io import PathManager
+from detectron2.data import DatasetCatalog, MetadataCatalog
+from lvis import LVIS
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["load_GRiTcoco_json", "register_GRiTcoco_instances"]
+
+
+def register_GRiTcoco_instances(name, metadata, json_file, image_root):
+    """
+    """
+    DatasetCatalog.register(name, lambda: load_GRiTcoco_json(
+        json_file, image_root, name))
+    MetadataCatalog.get(name).set(
+        json_file=json_file, image_root=image_root,
+        evaluator_type="coco", **metadata
+    )
+
+
+def get_GRiTcoco_meta():
+    categories = [{'supercategory': 'object', 'id': 1, 'name': 'object'}]
+    categories = sorted(categories, key=lambda x: x["id"])
+    thing_classes = [k["name"] for k in categories]
+    meta = {"thing_classes": thing_classes}
+    return meta
+
+
+def load_GRiTcoco_json(json_file, image_root, dataset_name=None):
+    '''
+    Load COCO class name text for object description for GRiT
+    '''
+
+    json_file = PathManager.get_local_path(json_file)
+
+    timer = Timer()
+    lvis_api = LVIS(json_file)
+    if timer.seconds() > 1:
+        logger.info("Loading {} takes {:.2f} seconds.".format(
+            json_file, timer.seconds()))
+
+    class_names = {}
+    sort_cat = sorted(lvis_api.dataset['categories'], key=lambda x: x['id'])
+    for x in sort_cat:
+        class_names[x['id']] = x['name']
+
+    img_ids = sorted(lvis_api.imgs.keys())
+    imgs = lvis_api.load_imgs(img_ids)
+    anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]
+
+    ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
+    assert len(set(ann_ids)) == len(ann_ids), \
+        "Annotation ids in '{}' are not unique".format(json_file)
+
+    imgs_anns = list(zip(imgs, anns))
+    logger.info("Loaded {} images in the LVIS v1 format from {}".format(
+        len(imgs_anns), json_file))
+
+    dataset_dicts = []
+
+    for (img_dict, anno_dict_list) in imgs_anns:
+        record = {}
+        if "file_name" in img_dict:
+            file_name = img_dict["file_name"]
+            record["file_name"] = os.path.join(image_root, file_name)
+
+        record["height"] = int(img_dict["height"])
+        record["width"] = int(img_dict["width"])
+        image_id = record["image_id"] = img_dict["id"]
+
+        objs = []
+        for anno in anno_dict_list:
+            assert anno["image_id"] == image_id
+            if anno.get('iscrowd', 0) > 0:
+                continue
+            obj = {"bbox": anno["bbox"], "bbox_mode": BoxMode.XYWH_ABS}
+            obj["category_id"] = 0
+            obj["object_description"] = class_names[anno['category_id']]
+            if 'segmentation' in anno:
+                segm = anno["segmentation"]
+                valid_segm = [poly for poly in segm \
+                    if len(poly) % 2 == 0 and len(poly) >= 6]
+                if not len(segm) == len(valid_segm):
+                    print('Annotation contains an invalid polygon with < 3 points')
+                assert len(segm) > 0
+                obj["segmentation"] = segm
+            objs.append(obj)
+        record["annotations"] = objs
+        if len(record["annotations"]) == 0:
+            continue
+        record["task"] = "ObjectDet"
+        dataset_dicts.append(record)
+
+    return dataset_dicts
+
+
+_CUSTOM_SPLITS_LVIS = {
+    "GRiT_coco2017_train": ("coco/train2017/", "coco/annotations/instances_train2017.json"),
+}
+
+
+for key, (image_root, json_file) in _CUSTOM_SPLITS_LVIS.items():
+    register_GRiTcoco_instances(
+        key,
+        get_GRiTcoco_meta(),
+        os.path.join("datasets", json_file) if "://" not in json_file else json_file,
+        os.path.join("datasets", image_root),
+    )

iChat/models/grit_src/grit/data/datasets/object365.py

@@ -0,0 +1,111 @@
+import logging
+import os
+from fvcore.common.timer import Timer
+from detectron2.structures import BoxMode
+from fvcore.common.file_io import PathManager
+from detectron2.data import DatasetCatalog, MetadataCatalog
+from lvis import LVIS
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["load_o365_json", "register_o365_instances"]
+
+
+def register_o365_instances(name, metadata, json_file, image_root):
+    DatasetCatalog.register(name, lambda: load_o365_json(
+        json_file, image_root, name))
+    MetadataCatalog.get(name).set(
+        json_file=json_file, image_root=image_root,
+        evaluator_type="lvis", **metadata
+    )
+
+
+def get_o365_meta():
+    categories = [{'supercategory': 'object', 'id': 1, 'name': 'object'}]
+    o365_categories = sorted(categories, key=lambda x: x["id"])
+    thing_classes = [k["name"] for k in o365_categories]
+    meta = {"thing_classes": thing_classes}
+    return meta
+
+
+def load_o365_json(json_file, image_root, dataset_name=None):
+    '''
+    Load Object365 class name text for object description for GRiT
+    '''
+
+    json_file = PathManager.get_local_path(json_file)
+
+    timer = Timer()
+    lvis_api = LVIS(json_file)
+    if timer.seconds() > 1:
+        logger.info("Loading {} takes {:.2f} seconds.".format(
+            json_file, timer.seconds()))
+
+    class_names = {}
+    sort_cat = sorted(lvis_api.dataset['categories'], key=lambda x: x['id'])
+    for x in sort_cat:
+        if '/' in x['name']:
+            text = ''
+            for xx in x['name'].split('/'):
+                text += xx
+                text += ' '
+            text = text[:-1]
+        else:
+            text = x['name']
+        class_names[x['id']] = text
+
+    img_ids = sorted(lvis_api.imgs.keys())
+    imgs = lvis_api.load_imgs(img_ids)
+    anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]
+
+    ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
+    assert len(set(ann_ids)) == len(ann_ids), \
+        "Annotation ids in '{}' are not unique".format(json_file)
+
+    imgs_anns = list(zip(imgs, anns))
+    logger.info("Loaded {} images in the LVIS v1 format from {}".format(
+        len(imgs_anns), json_file))
+
+    dataset_dicts = []
+
+    for (img_dict, anno_dict_list) in imgs_anns:
+        record = {}
+        if "file_name" in img_dict:
+            file_name = img_dict["file_name"]
+            record["file_name"] = os.path.join(image_root, file_name)
+
+        record["height"] = int(img_dict["height"])
+        record["width"] = int(img_dict["width"])
+        image_id = record["image_id"] = img_dict["id"]
+
+        objs = []
+        for anno in anno_dict_list:
+            assert anno["image_id"] == image_id
+            if anno.get('iscrowd', 0) > 0:
+                continue
+            obj = {"bbox": anno["bbox"], "bbox_mode": BoxMode.XYWH_ABS}
+            obj["category_id"] = 0
+            obj["object_description"] = class_names[anno['category_id']]
+
+            objs.append(obj)
+        record["annotations"] = objs
+        if len(record["annotations"]) == 0:
+            continue
+        record["task"] = "ObjectDet"
+        dataset_dicts.append(record)
+
+    return dataset_dicts
+
+
+_CUSTOM_SPLITS_LVIS = {
+    "object365_train": ("object365/images/train/", "object365/annotations/train_v1.json"),
+}
+
+
+for key, (image_root, json_file) in _CUSTOM_SPLITS_LVIS.items():
+    register_o365_instances(
+        key,
+        get_o365_meta(),
+        os.path.join("datasets", json_file) if "://" not in json_file else json_file,
+        os.path.join("datasets", image_root),
+    )

iChat/models/grit_src/grit/data/datasets/vg.py

@@ -0,0 +1,98 @@
+import logging
+import os
+from fvcore.common.timer import Timer
+from detectron2.structures import BoxMode
+from fvcore.common.file_io import PathManager
+from detectron2.data import DatasetCatalog, MetadataCatalog
+from lvis import LVIS
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["load_vg_json", "register_vg_instances"]
+
+
+def register_vg_instances(name, metadata, json_file, image_root):
+    """
+    """
+    DatasetCatalog.register(name, lambda: load_vg_json(
+        json_file, image_root, name))
+    MetadataCatalog.get(name).set(
+        json_file=json_file, image_root=image_root,
+        evaluator_type="vg", **metadata
+    )
+
+
+def get_vg_meta():
+    categories = [{'supercategory': 'object', 'id': 1, 'name': 'object'}]
+    vg_categories = sorted(categories, key=lambda x: x["id"])
+    thing_classes = [k["name"] for k in vg_categories]
+    meta = {"thing_classes": thing_classes}
+    return meta
+
+
+def load_vg_json(json_file, image_root, dataset_name=None):
+
+    json_file = PathManager.get_local_path(json_file)
+
+    timer = Timer()
+    lvis_api = LVIS(json_file)
+    if timer.seconds() > 1:
+        logger.info("Loading {} takes {:.2f} seconds.".format(
+            json_file, timer.seconds()))
+
+    img_ids = sorted(lvis_api.imgs.keys())
+    imgs = lvis_api.load_imgs(img_ids)
+    anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]
+
+    ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
+    assert len(set(ann_ids)) == len(ann_ids), \
+        "Annotation ids in '{}' are not unique".format(json_file)
+
+    imgs_anns = list(zip(imgs, anns))
+    logger.info("Loaded {} images in the LVIS v1 format from {}".format(
+        len(imgs_anns), json_file))
+
+    dataset_dicts = []
+
+    for (img_dict, anno_dict_list) in imgs_anns:
+        record = {}
+        if "file_name" in img_dict:
+            file_name = img_dict["file_name"]
+            record["file_name"] = os.path.join(image_root, file_name)
+
+        record["height"] = int(img_dict["height"])
+        record["width"] = int(img_dict["width"])
+        image_id = record["image_id"] = img_dict["id"]
+
+        objs = []
+        for anno in anno_dict_list:
+            assert anno["image_id"] == image_id
+            if anno.get('iscrowd', 0) > 0:
+                continue
+            obj = {"bbox": anno["bbox"], "bbox_mode": BoxMode.XYWH_ABS}
+            obj["category_id"] = 0
+            obj["object_description"] = anno["caption"]
+
+            objs.append(obj)
+        record["annotations"] = objs
+        if len(record["annotations"]) == 0:
+            continue
+        record["task"] = "DenseCap"
+        dataset_dicts.append(record)
+
+    return dataset_dicts
+
+
+_CUSTOM_SPLITS_LVIS = {
+    "vg_train": ("vg/images", "vg/annotations/train.json"),
+    "vg_test": ("vg/images", "vg/annotations/test.json"),
+}
+
+
+for key, (image_root, json_file) in _CUSTOM_SPLITS_LVIS.items():
+    register_vg_instances(
+        key,
+        get_vg_meta(),
+        os.path.join("datasets", json_file) if "://" not in json_file else json_file,
+        os.path.join("datasets", image_root),
+    )

iChat/models/grit_src/grit/data/transforms/custom_augmentation_impl.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# Part of the code is from https://github.com/rwightman/efficientdet-pytorch/blob/master/effdet/data/transforms.py 
+# Modified by Xingyi Zhou
+# The original code is under Apache-2.0 License
+import numpy as np
+from PIL import Image
+
+from detectron2.data.transforms.augmentation import Augmentation
+from .custom_transform import EfficientDetResizeCropTransform
+
+__all__ = [
+    "EfficientDetResizeCrop",
+]
+
+
+class EfficientDetResizeCrop(Augmentation):
+    """
+    Scale the shorter edge to the given size, with a limit of `max_size` on the longer edge.
+    If `max_size` is reached, then downscale so that the longer edge does not exceed max_size.
+    """
+
+    def __init__(
+        self, size, scale, interp=Image.BILINEAR
+    ):
+        """
+        """
+        super().__init__()
+        self.target_size = (size, size)
+        self.scale = scale
+        self.interp = interp
+
+    def get_transform(self, img):
+        # Select a random scale factor.
+        scale_factor = np.random.uniform(*self.scale)
+        scaled_target_height = scale_factor * self.target_size[0]
+        scaled_target_width = scale_factor * self.target_size[1]
+        # Recompute the accurate scale_factor using rounded scaled image size.
+        width, height = img.shape[1], img.shape[0]
+        img_scale_y = scaled_target_height / height
+        img_scale_x = scaled_target_width / width
+        img_scale = min(img_scale_y, img_scale_x)
+
+        # Select non-zero random offset (x, y) if scaled image is larger than target size
+        scaled_h = int(height * img_scale)
+        scaled_w = int(width * img_scale)
+        offset_y = scaled_h - self.target_size[0]
+        offset_x = scaled_w - self.target_size[1]
+        offset_y = int(max(0.0, float(offset_y)) * np.random.uniform(0, 1))
+        offset_x = int(max(0.0, float(offset_x)) * np.random.uniform(0, 1))
+        return EfficientDetResizeCropTransform(
+            scaled_h, scaled_w, offset_y, offset_x, img_scale, self.target_size, self.interp)

iChat/models/grit_src/grit/data/transforms/custom_transform.py

@@ -0,0 +1,115 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# Part of the code is from https://github.com/rwightman/efficientdet-pytorch/blob/master/effdet/data/transforms.py 
+# Modified by Xingyi Zhou
+# The original code is under Apache-2.0 License
+import numpy as np
+import torch
+import torch.nn.functional as F
+from fvcore.transforms.transform import (
+    CropTransform,
+    HFlipTransform,
+    NoOpTransform,
+    Transform,
+    TransformList,
+)
+from PIL import Image
+
+try:
+    import cv2  # noqa
+except ImportError:
+    # OpenCV is an optional dependency at the moment
+    pass
+
+__all__ = [
+    "EfficientDetResizeCropTransform",
+]
+
+
+class EfficientDetResizeCropTransform(Transform):
+    """
+    """
+
+    def __init__(self, scaled_h, scaled_w, offset_y, offset_x, img_scale, \
+        target_size, interp=None):
+        """
+        Args:
+            h, w (int): original image size
+            new_h, new_w (int): new image size
+            interp: PIL interpolation methods, defaults to bilinear.
+        """
+        # TODO decide on PIL vs opencv
+        super().__init__()
+        if interp is None:
+            interp = Image.BILINEAR
+        self._set_attributes(locals())
+
+    def apply_image(self, img, interp=None):
+        assert len(img.shape) <= 4
+
+        if img.dtype == np.uint8:
+            pil_image = Image.fromarray(img)
+            interp_method = interp if interp is not None else self.interp
+            pil_image = pil_image.resize((self.scaled_w, self.scaled_h), interp_method)
+            ret = np.asarray(pil_image)
+            right = min(self.scaled_w, self.offset_x + self.target_size[1])
+            lower = min(self.scaled_h, self.offset_y + self.target_size[0])
+            if len(ret.shape) <= 3:
+                ret = ret[self.offset_y: lower, self.offset_x: right]
+            else:
+                ret = ret[..., self.offset_y: lower, self.offset_x: right, :]
+        else:
+            # PIL only supports uint8
+            img = torch.from_numpy(img)
+            shape = list(img.shape)
+            shape_4d = shape[:2] + [1] * (4 - len(shape)) + shape[2:]
+            img = img.view(shape_4d).permute(2, 3, 0, 1)  # hw(c) -> nchw
+            _PIL_RESIZE_TO_INTERPOLATE_MODE = {Image.BILINEAR: "bilinear", Image.BICUBIC: "bicubic"}
+            mode = _PIL_RESIZE_TO_INTERPOLATE_MODE[self.interp]
+            img = F.interpolate(img, (self.scaled_h, self.scaled_w), mode=mode, align_corners=False)
+            shape[:2] = (self.scaled_h, self.scaled_w)
+            ret = img.permute(2, 3, 0, 1).view(shape).numpy()  # nchw -> hw(c)
+            right = min(self.scaled_w, self.offset_x + self.target_size[1])
+            lower = min(self.scaled_h, self.offset_y + self.target_size[0])
+            if len(ret.shape) <= 3:
+                ret = ret[self.offset_y: lower, self.offset_x: right]
+            else:
+                ret = ret[..., self.offset_y: lower, self.offset_x: right, :]
+        return ret
+
+
+    def apply_coords(self, coords):
+        coords[:, 0] = coords[:, 0] * self.img_scale
+        coords[:, 1] = coords[:, 1] * self.img_scale
+        coords[:, 0] -= self.offset_x
+        coords[:, 1] -= self.offset_y
+        return coords
+
+
+    def apply_segmentation(self, segmentation):
+        segmentation = self.apply_image(segmentation, interp=Image.NEAREST)
+        return segmentation
+
+
+    def inverse(self):
+        raise NotImplementedError
+
+
+    def inverse_apply_coords(self, coords):
+        coords[:, 0] += self.offset_x
+        coords[:, 1] += self.offset_y
+        coords[:, 0] = coords[:, 0] / self.img_scale
+        coords[:, 1] = coords[:, 1] / self.img_scale
+        return coords
+
+
+    def inverse_apply_box(self, box: np.ndarray) -> np.ndarray:
+        """
+        """
+        idxs = np.array([(0, 1), (2, 1), (0, 3), (2, 3)]).flatten()
+        coords = np.asarray(box).reshape(-1, 4)[:, idxs].reshape(-1, 2)
+        coords = self.inverse_apply_coords(coords).reshape((-1, 4, 2))
+        minxy = coords.min(axis=1)
+        maxxy = coords.max(axis=1)
+        trans_boxes = np.concatenate((minxy, maxxy), axis=1)
+        return trans_boxes

iChat/models/grit_src/grit/evaluation/eval.py

@@ -0,0 +1,156 @@
+import itertools
+import json
+import os
+from detectron2.structures import Boxes, BoxMode, pairwise_iou
+from detectron2.utils.file_io import PathManager
+import numpy as np
+import pycocotools.mask as mask_util
+from detectron2.evaluation.coco_evaluation import COCOEvaluator
+from detectron2.evaluation.coco_evaluation import _evaluate_predictions_on_coco
+
+
+class GRiTCOCOEvaluator(COCOEvaluator):
+    def process(self, inputs, outputs):
+        for input, output in zip(inputs, outputs):
+            prediction = {"image_id": input["image_id"]}
+
+            if "instances" in output:
+                instances = output["instances"].to(self._cpu_device)
+                prediction["instances"] = instances_to_coco_json(instances, input["image_id"])
+
+            if len(prediction) > 1:
+                self._predictions.append(prediction)
+
+    def _eval_predictions(self, predictions, img_ids=None):
+        self._logger.info("Preparing results for COCO format ...")
+        coco_results = list(itertools.chain(*[x["instances"] for x in predictions]))
+        tasks = self._tasks or self._tasks_from_predictions(coco_results)
+
+        if self._output_dir:
+            file_path = os.path.join(self._output_dir, "coco_instances_results.json")
+            self._logger.info("Saving results to {}".format(file_path))
+            with PathManager.open(file_path, "w") as f:
+                f.write(json.dumps(coco_results))
+                f.flush()
+
+        if not self._do_evaluation:
+            self._logger.info("Annotations are not available for evaluation.")
+            return
+
+        self._logger.info(
+            "Evaluating predictions with {} COCO API...".format(
+                "unofficial" if self._use_fast_impl else "official"
+            )
+        )
+
+        coco_results = self.convert_classname_to_id(coco_results)
+
+        for task in sorted(tasks):
+            assert task in {"bbox", "segm", "keypoints"}, f"Got unknown task: {task}!"
+            coco_eval = (
+                _evaluate_predictions_on_coco(
+                    self._coco_api,
+                    coco_results,
+                    task,
+                    kpt_oks_sigmas=self._kpt_oks_sigmas,
+                    use_fast_impl=self._use_fast_impl,
+                    img_ids=img_ids,
+                    max_dets_per_image=self._max_dets_per_image,
+                )
+                if len(coco_results) > 0
+                else None  # cocoapi does not handle empty results very well
+            )
+
+            res = self._derive_coco_results(
+                coco_eval, task, class_names=self._metadata.get("thing_classes")
+            )
+            self._results[task] = res
+
+    def convert_classname_to_id(self, results):
+        outputs = []
+        class_name_to_id = {}
+        categories = sorted(self._coco_api.dataset['categories'], key=lambda x: x['id'])
+
+        for cat in categories:
+            class_name_to_id[cat['name']] = cat['id']
+
+        for pred in results:
+            if pred['object_descriptions'] in class_name_to_id:
+                pred['category_id'] = class_name_to_id[pred['object_descriptions']]
+                del pred['object_descriptions']
+                outputs.append(pred)
+
+        return outputs
+
+
+class GRiTVGEvaluator(COCOEvaluator):
+    def process(self, inputs, outputs):
+        for input, output in zip(inputs, outputs):
+            assert input["image_id"] == int(input['file_name'].split('/')[-1].split('.')[0])
+            prediction = {"image_id": input["image_id"]}
+
+            if "instances" in output:
+                instances = output["instances"].to(self._cpu_device)
+                prediction["instances"] = instances_to_coco_json(instances, input["image_id"], output_logits=True)
+                h = input['height']
+                w = input['width']
+                scale = 720.0 / max(h, w)
+                scaled_inst = []
+                for inst in prediction["instances"]:
+                    inst['bbox'][0] = inst['bbox'][0] * scale
+                    inst['bbox'][1] = inst['bbox'][1] * scale
+                    inst['bbox'][2] = inst['bbox'][2] * scale
+                    inst['bbox'][3] = inst['bbox'][3] * scale
+                    scaled_inst.append(inst)
+                if len(scaled_inst) > 0:
+                    prediction["instances"] = scaled_inst
+            if len(prediction) > 1:
+                self._predictions.append(prediction)
+
+    def _eval_predictions(self, predictions, img_ids=None):
+        '''
+        This is only for saving the results to json file
+        '''
+        self._logger.info("Preparing results for COCO format ...")
+        coco_results = list(itertools.chain(*[x["instances"] for x in predictions]))
+
+        if self._output_dir:
+            file_path = os.path.join(self._output_dir, "vg_instances_results.json")
+            self._logger.info("Saving results to {}".format(file_path))
+            with PathManager.open(file_path, "w") as f:
+                f.write(json.dumps(coco_results))
+                f.flush()
+
+
+def instances_to_coco_json(instances, img_id, output_logits=False):
+    """
+        Add object_descriptions and logit (if applicable) to
+        detectron2's instances_to_coco_json
+    """
+    num_instance = len(instances)
+    if num_instance == 0:
+        return []
+
+    boxes = instances.pred_boxes.tensor.numpy()
+    boxes = BoxMode.convert(boxes, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    boxes = boxes.tolist()
+    scores = instances.scores.tolist()
+    classes = instances.pred_classes.tolist()
+    object_descriptions = instances.pred_object_descriptions.data
+    if output_logits:
+        logits = instances.logits.tolist()
+
+    results = []
+    for k in range(num_instance):
+        result = {
+            "image_id": img_id,
+            "category_id": classes[k],
+            "bbox": boxes[k],
+            "score": scores[k],
+            'object_descriptions': object_descriptions[k],
+        }
+        if output_logits:
+            result["logit"] = logits[k]
+
+        results.append(result)
+    return results