IDEA-Research/ChatRex-7B

1---2language:3- en4base_model:5- lmsys/vicuna-7b-v1.56- openai/clip-vit-large-patch147- laion/CLIP-convnext_large_d.laion2B-s26B-b102K-augreg8pipeline_tag: image-text-to-text9tags:10- chatrex11- upn12---13 14arxiv.org/abs/2411.1836315 16<div align=center>17  <img src="assets/teaser.jpg" width=600 >18</div>19 20----21 22# 1. Introduction 📚23**TL;DR: ChatRex is an MLLM skilled in perception that can respond to questions while simultaneously grounding its answers to the referenced objects.**24 25ChatRex is a Multimodal Large Language Model (MLLM) designed to seamlessly integrate fine-grained object perception and robust language understanding. By adopting a decoupled architecture with a retrieval-based approach for object detection and leveraging high-resolution visual inputs, ChatRex addresses key challenges in perception tasks. It is powered by the Rexverse-2M dataset with diverse image-region-text annotations. ChatRex can be applied to various scenarios requiring fine-grained perception, such as object detection, grounded conversation, grounded image captioning and region26understanding.27 28<div align=center>29  <img src="assets/capability_overview.jpg" width=800 >30</div>31 32----33 34# 2. Installation 🛠️35```bash36conda install -n chatrex python=3.937pip install torch==2.1.2 torchvision==0.16.2 --index-url https://download.pytorch.org/whl/cu12138git clone https://github.com/IDEA-Research/ChatRex.git39cd ChatRex40pip install -v -e .41# install deformable attention for universal proposal network42cd chatrex/upn/ops43pip install -v -e .44```45 46## 2.1 Download Pre-trained UPN Models47We provide model checkpoints for both the ***Universal Proposal Network (UPN)*** and the ***ChatRex model***. You can download the pre-trained models from the following links:48- [UPN Checkpoint](https://github.com/IDEA-Research/ChatRex/releases/download/upn-large/upn_large.pth)49- [ChatRex-7B Checkpoint](https://huggingface.co/IDEA-Research/ChatRex-7B)50 51Or you can also using the following command to download the pre-trained models:52```bash53mkdir checkpoints54mkdir checkpoints/upn55# download UPN checkpoint56wget -O checkpoints/upn/upn_large.pth https://github.com/IDEA-Research/ChatRex/releases/download/upn-large/upn_large.pth57```58 59## 2.2 Verify Installation60To verify the ***installation of the Universal Proposal Network (UPN)***, run the following command:61```bash62python tests/test_upn_install.py63```64 65If the installation is successful, you will get two visualization images of both fine-grained proposal and coarse-grained proposal in `tests` folder.66 67To verify the ***installation of the ChatRex model***, run the following command:68```bash69python tests/test_chatrex_install.py70```71 72If the installation is successful, you will get an output like this:73```text74prediction: <obj0> shows a brown dog lying on a bed. The dog is resting comfortably, possibly sleeping, and is positioned on the left side of the bed75```76 77# 3. Usage 🚀78## 3.1 Use UPN for Object Proposal Generation79 80Universal Proposal Network (UPN) is a robust object proposal model designed as part of ChatRex to enable comprehensive and accurate object detection across diverse granularities and domains. Built upon T-Rex2, UPN is a DETR-based model with a dual-granularity prompt tuning strategy, combining fine-grained (e.g., part-level) and coarse-grained (e.g., instance-level) detection.81 82<div align=center>83  <img src="assets/upn_res.jpg" width=600 >84</div>85 86----87 88<details close>89<summary><strong>Example Code for UPN</strong></summary>90 91```python92import torch93from PIL import Image94from tools.visualize import plot_boxes_to_image95from chatrex.upn import UPNWrapper96 97ckpt_path = "checkpoints/upn_checkpoints/upn_large.pth"98test_image_path = "tests/images/test_upn.jpeg"99 100model = UPNWrapper(ckpt_path)101# fine-grained prompt102fine_grained_proposals = model.inference(103    test_image_path, prompt_type="fine_grained_prompt"104)105# filter by score (default: 0.3) and nms (default: 0.8)106fine_grained_filtered_proposals = model.filter(107    fine_grained_proposals, min_score=0.3, nms_value=0.8108)109## output is a dict with keys: "original_xyxy_boxes", "scores"110## - "original_xyxy_boxes": list of boxes in xyxy format in shape (B, N, 4)111## - "scores": list of scores for each box in shape (B, N)112 113# coarse-grained prompt114coarse_grained_proposals = model.inference(115    test_image_path, prompt_type="coarse_grained_prompt"116)117coarse_grained_filtered_proposals = model.filter(118    coarse_grained_proposals, min_score=0.3, nms_value=0.8119)120 121## output is a dict with keys: "original_xyxy_boxes", "scores"122## - "original_xyxy_boxes": list of boxes in xyxy format in shape (B, N, 4)123## - "scores": list of scores for each box in shape (B, N)124```125 126</details>127 128We also provide a visualization tool to visualize the object proposals generated by UPN. You can use the following code to visualize the object proposals:129 130<details close>131<summary><strong>Example Code for UPN Visualization</strong></summary>132 133```python134 135from chatrex.tools.visualize import plot_boxes_to_image136image = Image.open(test_image_path)137fine_grained_vis_image, _ = plot_boxes_to_image(138    image.copy(),139    fine_grained_filtered_proposals["original_xyxy_boxes"][0],140    fine_grained_filtered_proposals["scores"][0],141)142fine_grained_vis_image.save("tests/test_image_fine_grained.jpeg")143print(f"fine-grained proposal is saved at tests/test_image_fine_grained.jpeg")144 145coarse_grained_vis_image, _ = plot_boxes_to_image(146    image.copy(),147    coarse_grained_filtered_proposals["original_xyxy_boxes"][0],148    coarse_grained_filtered_proposals["scores"][0],149)150coarse_grained_vis_image.save("tests/test_image_coarse_grained.jpeg")151print(f"coarse-grained proposal is saved at tests/test_image_coarse_grained.jpeg")152 153```154</details>155 156## 3.2 Usage of ChatRex157 158ChatRex takes three inputs: image, text prompt, and box input. For the box input, you can either use the object proposals generated by UPN or provide your own box input (user drawn boxes). We have wrapped the ChatRex model to huggingface transformers format for easy usage. ChatRex can be used for various tasks and we provide example code for each task below.159 160### 3.2.1 ChatRex for Object Detection & Grounding & Referring161 162Example Prompt for detection, grounding, referring tasks:163```text164# Single Object Detection165Please detect dog in this image. Answer the question with object indexes.166Please detect the man in yellow shirt in this image. Answer the question with object indexes.167 168# multiple object detection, use ; to separate the objects169Please detect person; pigeon in this image. Answer the question with object indexes.170Please detect person in the car; cat below the table in this image. Answer the question with object indexes.171```172 173<details close>174<summary><strong>Example Code</strong></summary>175 176```python177import torch178from PIL import Image179from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig180 181from chatrex.tools.visualize import visualize_chatrex_output182from chatrex.upn import UPNWrapper183 184if __name__ == "__main__":185    # load the processor186    processor = AutoProcessor.from_pretrained(187        "IDEA-Research/ChatRex-7B",188        trust_remote_code=True,189        device_map="cuda",190    )191 192    print(f"loading chatrex model...")193    # load chatrex model194    model = AutoModelForCausalLM.from_pretrained(195        "IDEA-Research/ChatRex-7B",196        trust_remote_code=True,197        use_safetensors=True,198    ).to("cuda")199 200    # load upn model201    print(f"loading upn model...")202    ckpt_path = "checkpoints/upn_checkpoints/upn_large.pth"203    model_upn = UPNWrapper(ckpt_path)204    test_image_path = "tests/images/test_chatrex_detection.jpg"205 206    # get upn predictions207    fine_grained_proposals = model_upn.inference(208        test_image_path, prompt_type="fine_grained_prompt"209    )210    fine_grained_filtered_proposals = model_upn.filter(211        fine_grained_proposals, min_score=0.3, nms_value=0.8212    )213 214    inputs = processor.process(215        image=Image.open(test_image_path),216        question="Please detect person; pigeon in this image. Answer the question with object indexes.",217        bbox=fine_grained_filtered_proposals["original_xyxy_boxes"][218            0219        ],  # box in xyxy format220    )221 222    inputs = {k: v.to("cuda") for k, v in inputs.items()}223 224    # perform inference225    gen_config = GenerationConfig(226        max_new_tokens=512,227        do_sample=False,228        eos_token_id=processor.tokenizer.eos_token_id,229        pad_token_id=(230            processor.tokenizer.pad_token_id231            if processor.tokenizer.pad_token_id is not None232            else processor.tokenizer.eos_token_id233        ),234    )235    with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):236        prediction = model.generate(237            inputs, gen_config=gen_config, tokenizer=processor.tokenizer238        )239    print(f"prediction:", prediction)240 241    # visualize the prediction242    vis_image = visualize_chatrex_output(243        Image.open(test_image_path),244        fine_grained_filtered_proposals["original_xyxy_boxes"][0],245        prediction,246        font_size=15,247        draw_width=5,248    )249    vis_image.save("tests/test_chatrex_detection.jpeg")250    print(f"prediction is saved at tests/test_chatrex_detection.jpeg")251```252 253The output from LLM is like:254```text255<ground>person</ground><objects><obj10><obj14><obj15><obj27><obj28><obj32><obj33><obj35><obj38><obj47><obj50></objects>256<ground>pigeon</ground><objects><obj0><obj1><obj2><obj3><obj4><obj5><obj6><obj7><obj8><obj9><obj11><obj12><obj13><obj16><obj17><obj18><obj19><obj20><obj21><obj22><obj23><obj24><obj25><obj26><obj29><obj31><obj37><obj39><obj40><obj41><obj44><obj49></objects>257```258 259The visualization of the output is like:260 261<div align=center>262  <img src="assets/vis_output/test_chatrex_detection.jpeg" width=600 >263</div>264 265</details>266 267----268 269### 3.2.2 ChatRex for Region Caption270Example Prompt for Region Caption tasks:271 272```text273# Single Object Detection274## caption in category name275What is the category name of <obji>? Answer the question with its category name in free format.276 277## caption in short phrase278Can you provide me with a short phrase to describe <obji>? Answer the question with a short phrase.279 280## caption in referring style281Can you provide me with a brief description of <obji>? Answer the question with brief description.282 283## caption in one sentence284Can you provide me with a one sentence of <obji>? Answer the question with one sentence description.285 286# multiple object detection, use ; to separate the objects287```288 289<details close>290<summary><strong>Example Code</strong></summary>291 292```python293import torch294from PIL import Image295from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig296 297from chatrex.tools.visualize import visualize_chatrex_output298from chatrex.upn import UPNWrapper299 300if __name__ == "__main__":301    # load the processor302    processor = AutoProcessor.from_pretrained(303        "IDEA-Research/ChatRex-7B",304        trust_remote_code=True,305        device_map="cuda",306    )307 308    print(f"loading chatrex model...")309    # load chatrex model310    model = AutoModelForCausalLM.from_pretrained(311        "IDEA-Research/ChatRex-7B",312        trust_remote_code=True,313        use_safetensors=True,314    ).to("cuda")315 316    test_image_path = "tests/images/test_chatrex_install.jpg"317 318    inputs = processor.process(319        image=Image.open(test_image_path),320        question="Can you provide a one sentence description of <obj0> in the image? Answer the question with a one sentence description.",321        bbox=[[73.88417, 56.62228, 227.69223, 216.34338]],322    )323 324    inputs = {k: v.to("cuda") for k, v in inputs.items()}325 326    # perform inference327    gen_config = GenerationConfig(328        max_new_tokens=512,329        do_sample=False,330        eos_token_id=processor.tokenizer.eos_token_id,331        pad_token_id=(332            processor.tokenizer.pad_token_id333            if processor.tokenizer.pad_token_id is not None334            else processor.tokenizer.eos_token_id335        ),336    )337    with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):338        prediction = model.generate(339            inputs, gen_config=gen_config, tokenizer=processor.tokenizer340        )341    print(f"prediction:", prediction)342 343    # visualize the prediction344    vis_image = visualize_chatrex_output(345        Image.open(test_image_path),346        [[73.88417, 56.62228, 227.69223, 216.34338]],347        prediction,348        font_size=15,349        draw_width=5,350    )351    vis_image.save("tests/test_chatrex_region_caption.jpeg")352    print(f"prediction is saved at tests/test_chatrex_region_caption.jpeg")353```354 355The output from LLM is like:356```text357<ground>A brown dog is lying on a bed, appearing relaxed and comfortable</ground><objects><obj0></objects>358```359 360The visualization of the output is like:361 362<div align=center>363  <img src="assets/vis_output/test_chatrex_region_caption.jpeg" width=600 >364</div>365 366</details>367 368----369 370### 3.2.3 ChatRex for Grounded Image Captioning371Example Prompt for Region Caption tasks:372 373```text374# Brief Grounded Imager Caption375Please breifly describe this image in one sentence and detect all the mentioned objects. Answer the question with grounded answer.376 377# Detailed Grounded Image Caption378Please provide a detailed description of the image and detect all the mentioned objects. Answer the question with grounded object indexes.379```380 381<details close>382<summary><strong>Example Code</strong></summary>383 384```python385import torch386from PIL import Image387from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig388 389from chatrex.tools.visualize import visualize_chatrex_output390from chatrex.upn import UPNWrapper391 392if __name__ == "__main__":393    # load the processor394    processor = AutoProcessor.from_pretrained(395        "IDEA-Research/ChatRex-7B",396        trust_remote_code=True,397        device_map="cuda",398    )399 400    print(f"loading chatrex model...")401    # load chatrex model402    model = AutoModelForCausalLM.from_pretrained(403        "IDEA-Research/ChatRex-7B",404        trust_remote_code=True,405        use_safetensors=True,406    ).to("cuda")407 408    # load upn model409    print(f"loading upn model...")410    ckpt_path = "checkpoints/upn_checkpoints/upn_large.pth"411    model_upn = UPNWrapper(ckpt_path)412    test_image_path = "tests/images/test_chatrex_grounded_caption.jpg"413 414    # get upn predictions415    fine_grained_proposals = model_upn.inference(416        test_image_path, prompt_type="fine_grained_prompt"417    )418    fine_grained_filtered_proposals = model_upn.filter(419        fine_grained_proposals, min_score=0.3, nms_value=0.8420    )421 422    inputs = processor.process(423        image=Image.open(test_image_path),424        question="Please breifly describe this image in one sentence and detect all the mentioned objects. Answer the question with grounded answer.",425        bbox=fine_grained_filtered_proposals["original_xyxy_boxes"][426            0427        ],  # box in xyxy format428    )429 430    inputs = {k: v.to("cuda") for k, v in inputs.items()}431 432    # perform inference433    gen_config = GenerationConfig(434        max_new_tokens=512,435        do_sample=False,436        eos_token_id=processor.tokenizer.eos_token_id,437        pad_token_id=(438            processor.tokenizer.pad_token_id439            if processor.tokenizer.pad_token_id is not None440            else processor.tokenizer.eos_token_id441        ),442    )443    with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):444        prediction = model.generate(445            inputs, gen_config=gen_config, tokenizer=processor.tokenizer446        )447    print(f"prediction:", prediction)448 449    # visualize the prediction450    vis_image = visualize_chatrex_output(451        Image.open(test_image_path),452        fine_grained_filtered_proposals["original_xyxy_boxes"][0],453        prediction,454        font_size=15,455        draw_width=5,456    )457    vis_image.save("tests/test_chatrex_grounded_image_caption.jpeg")458    print(f"prediction is saved at tests/test_chatrex_grounded_image_caption.jpeg")459```460 461The output from LLM is like:462```text463The image depicts a cozy living room with a <ground>plaid couch,</ground><objects><obj2></objects> a <ground>wooden TV stand</ground><objects><obj3></objects>holding a <ground>black television,</ground><objects><obj1></objects> a <ground>red armchair,</ground><objects><obj4></objects> and a <ground>whiteboard</ground><objects><obj0></objects>with writing on the wall, accompanied by a <ground>framed poster</ground><objects><obj6></objects>of a <ground>couple.</ground><objects><obj9><obj11></objects>464```465 466The visualization of the output is like:467 468<div align=center>469  <img src="assets/vis_output/test_chatrex_grounded_image_caption.jpeg" width=600 >470</div>471 472</details>473 474----475 476### 3.2.4 ChatRex for Grounded Conversation477Example Prompt for Region Caption tasks:478 479```text480Answer the question in Grounded format. Question481```482 483<details close>484<summary><strong>Example Code</strong></summary>485 486```python487import torch488from PIL import Image489from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig490 491from chatrex.tools.visualize import visualize_chatrex_output492from chatrex.upn import UPNWrapper493 494if __name__ == "__main__":495    # load the processor496    processor = AutoProcessor.from_pretrained(497        "IDEA-Research/ChatRex-7B",498        trust_remote_code=True,499        device_map="cuda",500    )501 502    print(f"loading chatrex model...")503    # load chatrex model504    model = AutoModelForCausalLM.from_pretrained(505        "IDEA-Research/ChatRex-7B",506        trust_remote_code=True,507        use_safetensors=True,508    ).to("cuda")509 510    # load upn model511    print(f"loading upn model...")512    ckpt_path = "checkpoints/upn_checkpoints/upn_large.pth"513    model_upn = UPNWrapper(ckpt_path)514    test_image_path = "tests/images/test_grounded_conversation.jpg"515 516    # get upn predictions517    fine_grained_proposals = model_upn.inference(518        test_image_path, prompt_type="coarse_grained_prompt"519    )520    fine_grained_filtered_proposals = model_upn.filter(521        fine_grained_proposals, min_score=0.3, nms_value=0.8522    )523 524    inputs = processor.process(525        image=Image.open(test_image_path),526        question="Answer the question in grounded format. This is a photo of my room, and can you tell me what kind of person I am?  ",527        bbox=fine_grained_filtered_proposals["original_xyxy_boxes"][528            0529        ],  # box in xyxy format530    )531 532    inputs = {k: v.to("cuda") for k, v in inputs.items()}533 534    # perform inference535    gen_config = GenerationConfig(536        max_new_tokens=512,537        do_sample=False,538        eos_token_id=processor.tokenizer.eos_token_id,539        pad_token_id=(540            processor.tokenizer.pad_token_id541            if processor.tokenizer.pad_token_id is not None542            else processor.tokenizer.eos_token_id543        ),544    )545    with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):546        prediction = model.generate(547            inputs, gen_config=gen_config, tokenizer=processor.tokenizer548        )549    print(f"prediction:", prediction)550 551    # visualize the prediction552    vis_image = visualize_chatrex_output(553        Image.open(test_image_path),554        fine_grained_filtered_proposals["original_xyxy_boxes"][0],555        prediction,556        font_size=30,557        draw_width=10,558    )559    vis_image.save("tests/test_chatrex_grounded_conversation.jpeg")560    print(f"prediction is saved at tests/test_chatrex_grounded_conversation.jpeg")561 562```563 564The output from LLM is like:565```text566Based on the items in the image, it can be inferred that the <ground>person</ground><objects><obj1></objects> who owns this room has an interest in fitness and possibly enjoys reading. The presence of the <ground>dumbbell</ground><objects><obj2></objects> suggests a commitment to physical activity, while the <ground>book</ground><objects><obj3></objects> indicates a liking for literature or reading. The <ground>sneaker</ground><objects><obj0></objects>s and the <ground>plush toy</ground><objects><obj1></objects> add a personal touch, suggesting that the <ground>person</ground><objects><obj1></objects> might also value comfort and perhaps has a playful or nostalgic side. However, without more context, it is not possible to accurately determine the individual's specific traits or <ground>person</ground><objects><obj1></objects>ality.567```568 569The visualization of the output is like:570 571<div align=center>572  <img src="assets/test_chatrex_grounded_conversation.jpeg" width=600 >573</div>574 575</details>576 577----578 579 580# 5. LICENSE581 582ChatRex is licensed under the IDEA License 1.0, Copyright (c) IDEA. All Rights Reserved. Note that this project utilizes certain datasets and checkpoints that are subject to their respective original licenses. Users must comply with all terms and conditions of these original licenses including but not limited to the:583- [OpenAI Terms of Use](https://openai.com/policies/terms-of-use) for the dataset. 584- For the LLM used in this project, the model is [lmsys/vicuna-7b-v1.5](https://huggingface.co/lmsys/vicuna-7b-v1.5/tree/main), which is licensed under [Llama 2 Community License Agreement](https://huggingface.co/lmsys/vicuna-7b-v1.5).585- For the high resolution vision encoder, we are using [laion/CLIP-convnext_large_d.laion2B-s26B-b102K-augreg](https://huggingface.co/laion/CLIP-convnext_large_d.laion2B-s26B-b102K-augreg) which is licensed under [MIT LICENSE](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/mit.md).586- For the low resolution vision encoder, we are using [openai/clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) which is licensed under [MIT LICENSE](https://github.com/openai/CLIP/blob/main/LICENSE)587# BibTeX 📚588```589@misc{jiang2024chatrextamingmultimodalllm,590      title={ChatRex: Taming Multimodal LLM for Joint Perception and Understanding}, 591      author={Qing Jiang and Gen Luo and Yuqin Yang and Yuda Xiong and Yihao Chen and Zhaoyang Zeng and Tianhe Ren and Lei Zhang},592      year={2024},593      eprint={2411.18363},594      archivePrefix={arXiv},595      primaryClass={cs.CV},596      url={https://arxiv.org/abs/2411.18363}, 597}598```599