diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..afc0d01 --- /dev/null +++ b/.DS_Store Binary files differ diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..5ada641 --- /dev/null +++ b/.gitignore @@ -0,0 +1,162 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class +OAI_CONFIG_LIST + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/latest/usage/project/#working-with-version-control +.pdm.toml +.pdm-python +.pdm-build/ + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/OAI_CONFIG_LIST.example b/OAI_CONFIG_LIST.example new file mode 100644 index 0000000..57d2ea0 --- /dev/null +++ b/OAI_CONFIG_LIST.example @@ -0,0 +1,14 @@ +[ + { + "model": "Qwen/Qwen2.5-32B-Instruct", + "api_key": "", + "base_url": "https://api.siliconflow.cn/v1", + "tags": ["silicon"] + }, + { + "model": "gpt-4o-2024-08-06", + "api_key": "", + "base_url": "", + "tags": ["openai"] + } +] diff --git a/README.md b/README.md index ec21a73..b2c8af5 100644 --- a/README.md +++ b/README.md @@ -1 +1,13 @@ -# MedSegAgent \ No newline at end of file +# Med-SegAgent: Universal & Scalable Multi-Agent System for Instructive Medical Image Segmentation + + + +## Key Features: +* **Universal & Scalable:** Tackle diverse medical image segmentation tasks using natural language instructions. +* **Precise Automation:** Automatically selects the most suitable segmentation models. +* **Enhanced Robustness:** Improves reliability through multi-model integration and ensemble capabilities. + +## Quick Start: +[WIP] diff --git a/agents/__init__.py b/agents/__init__.py new file mode 100644 index 0000000..37a20cb --- /dev/null +++ b/agents/__init__.py @@ -0,0 +1,12 @@ +from agents.native import chat_with_llm_native +from agents.cot import chat_with_llm_cot +from agents.react import chat_with_llm_react +from agents.react_rag import chat_with_llm_react_rag +from agents.react_selfverify import chat_with_llm_react_selfverify + +from agents.react_multiexpert import chat_with_llm_react_multiexpert +from agents.react_multiexpert_experience import chat_with_llm_react_multiexpert_experience +from agents.react_multiexpert_selfverify import chat_with_llm_react_multiexpert_selfverify +from agents.react_multiexpert_selfverify_experience import chat_with_llm_react_multiexpert_selfverify_experience + +from agents.c2f_tree_search import chat_with_llm_c2f \ No newline at end of file diff --git a/agents/assigner_experts.py b/agents/assigner_experts.py new file mode 100644 index 0000000..e0e5cac --- /dev/null +++ b/agents/assigner_experts.py @@ -0,0 +1,98 @@ +from prompts.assigner_experts import DEPARTMENT_EXPERTS_INFO, generate_expert_sys_prompt_dict, generate_doctor_expert_sys_prompt_dict +from prompts.data import DATASET_NAME_TO_PATH +from agents.single_turn import chat_with_llm_single_turn +from agents.doctor_expert import chat_with_llm_multi_turn_doctor_expert +from utils.reply_utils import extract_llm_answer +''' +The task is decouple into this workflow: +1. query - -> chosen s +2. for each , reuse the `doctor_expert` pattern to generate results +3. a general check the result with selected information (extract only related information from dataset.json) + and decided +''' + + +def chat_with_llm_assigner_expert_single_turn(query, model, autogen_config_path): + profile_dataset_mapping = DEPARTMENT_EXPERTS_INFO + ''' step1: assigner choose experts ''' + # skip now + ''' step2: experts answer the question ''' + general_reply_list = [] + general_chat_history = [] + general_answer = dict() + + for profile, dataset_list in profile_dataset_mapping.items(): + sys_prompt_dict = generate_expert_sys_prompt_dict(profile, dataset_list) + expert_res = chat_with_llm_single_turn(query, + model, + autogen_config_path, + sys_prompt_dict=sys_prompt_dict) + general_reply_list.append(expert_res["reply"]) + expert_res['chat_history'][-1]['role'] = profile.replace(".", "") + general_chat_history.extend(expert_res['chat_history']) + + expert_answer = extract_llm_answer(expert_res["reply"]) + for dataset, labels in expert_answer.items(): + if (str(dataset) not in DATASET_NAME_TO_PATH): + continue + general_answer[dataset] = labels + + dataset_reply = ", ".join(list(general_answer.keys())) + labels_reply = ", ".join([str(general_answer[gak]) for gak in list(general_answer.keys())]) + general_reply = "{\"dataset\": \"" + dataset_reply + "\",\"labels\": \"" + labels_reply + "\"}" + general_chat_history.extend([ + { + "role": "Summarizer", + "content": general_reply + }, + ]) + + return { + "reply": general_reply, + "chat_history": general_chat_history, + } + + +def chat_with_llm_assigner_expert_multi_turn_verifer(query, model, autogen_config_path): + profile_dataset_mapping = DEPARTMENT_EXPERTS_INFO + ''' step1: assigner choose experts ''' + # skip now + ''' step2: experts answer the question ''' + general_reply_list = [] + general_chat_history = [] + general_answer = dict() + + from tenacity import retry, stop_after_attempt + retry_decorator = retry(stop=stop_after_attempt(3)) + chat_func_with_retry = retry_decorator(chat_with_llm_multi_turn_doctor_expert) + + for profile, dataset_list in profile_dataset_mapping.items(): + sys_prompt_dict = generate_doctor_expert_sys_prompt_dict(profile, dataset_list) + expert_res = chat_func_with_retry(query, + model, + autogen_config_path, + system_prompt_dict=sys_prompt_dict) + general_reply_list.append(expert_res["reply"]) + expert_res['chat_history'][-1]['role'] = profile.replace(".", "") + general_chat_history.extend(expert_res['chat_history']) + + expert_answer = extract_llm_answer(expert_res["reply"]) + for dataset, labels in expert_answer.items(): + if (str(dataset) not in DATASET_NAME_TO_PATH): + continue + general_answer[dataset] = labels + + dataset_reply = ", ".join(list(general_answer.keys())) + labels_reply = ", ".join([str(general_answer[gak]) for gak in list(general_answer.keys())]) + general_reply = "{\"dataset\": \"" + dataset_reply + "\",\"labels\": \"" + labels_reply + "\"}" + general_chat_history.extend([ + { + "role": "Summarizer", + "content": general_reply + }, + ]) + + return { + "reply": general_reply, + "chat_history": general_chat_history, + } diff --git a/agents/basic.py b/agents/basic.py new file mode 100644 index 0000000..98884c9 --- /dev/null +++ b/agents/basic.py @@ -0,0 +1,90 @@ +from openai import OpenAI +from autogen import config_list_from_json + + +def chat_with_agent_single_turn(query, model, autogen_config_path, sys_prompt): + config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model, "default"]})[0] + client = OpenAI(api_key=config['api_key'], base_url=config['base_url']) + chat_history = [{"role": "user", "content": query}] + completion = client.chat.completions.create( + model=model, + messages=[{ + "role": "system", + "content": sys_prompt + }, { + "role": "user", + "content": query + }], + ) + assistant_reply = completion.choices[0].message.content + chat_history.append({"role": "assistant", "content": assistant_reply}) + return { + "reply": assistant_reply, + "chat_history": chat_history, + } + + +class Agent: + + def __init__(self, name, system_message, llm_config, *args): + """ + Initialize the Agent with a name, system message, and LLM configuration. + + :param name: The name of the agent. + :param system_message: The system message or prompt for the agent. + :param llm_config: Configuration for the language model. + """ + self.name = name + self.system_message = system_message + self.llm_config = llm_config + self.model = self.llm_config["model"] + self.client = OpenAI(api_key=self.llm_config['api_key'], + base_url=self.llm_config['base_url']) + + def send(self, recipient, message): + """ + Send a query to another agent. + + :param query: The query to send. + :param recipient: The target agent to send the query to. + :return: The generated response. + """ + return recipient.process_query(message) + + def process_query(self, query): + """ + Process a query received from another agent. + + :param query: The query to process. + :return: The generated response. + """ + response = self._generate_response(query) + return response + + def _generate_response(self, query): + """ + Generate a response using the LLM. + + :param query: The query to generate a response for. + :return: The generated response. + """ + chat_history = [{"role": "user", "content": query}] + import pdb + pdb.set_trace() + completion = self.client.chat.completions.create( + model=self.model, + messages=[{ + "role": "system", + "content": self.system_message + }, { + "role": "user", + "content": query + }], + ) + assistant_reply = completion.choices[0].message.content + chat_history.append({"role": "assistant", "content": assistant_reply}) + return { + "reply": assistant_reply, + "chat_history": chat_history, + } diff --git a/agents/c2f_tree_search.py b/agents/c2f_tree_search.py new file mode 100644 index 0000000..cdf49a8 --- /dev/null +++ b/agents/c2f_tree_search.py @@ -0,0 +1,518 @@ +import json +import os.path as osp +import logging +from glob import glob + +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent + + +# Function to load JSON dataset information from a directory +def load_dataset_info(json_dir, + keys_to_keep=["name", "description", "input_modalities", "labels"], + file_list=None): + """ + Complete list of keys: ["name", "description", "input_modalities", "labels", "models"] + """ + json_info_list = [] + # Append .json extension to the filenames in the list + if file_list is not None: + file_list = {name + ".json" for name in file_list} + + for json_path in glob(osp.join(json_dir, "*.json")): + file_name = osp.basename(json_path) + + # Check if the file is in the specified file list + if file_list is None or file_name in file_list: + with open(json_path, encoding='utf-8') as f: + json_content = json.load(f) + # Keep only specified keys + filtered_content = {key: json_content.get(key) for key in keys_to_keep} + json_info_list.append(json.dumps(filtered_content)) + + return "\n".join(json_info_list) + + +# Function to extract JSON data from LLM output +def extract_json_from_output(llm_output, expected_type='list'): + import re + if expected_type == 'list': + pattern = r"Final Answer:\s*(\[[\s\S]*?\])" + elif expected_type == 'dict': + pattern = r"Final Answer:\s*(\{[\s\S]*?\})" + else: + raise ValueError("expected_type must be 'list' or 'dict'") + + match = re.search(pattern, llm_output) + if match: + json_str = match.group(1) + # Remove comments from JSON string + json_str = re.sub(r'//.*', '', json_str) # Remove single-line comments + json_str = re.sub(r'/\*[\s\S]*?\*/', '', json_str) # Remove multi-line comments + try: + json_data = json.loads(json_str) + return json_data + except json.JSONDecodeError as e: + logging.error("JSON decoding error:", e) + logging.info("Attempting to fix JSON string...") + json_str_fixed = json_str.replace("'", '"') + try: + json_data = json.loads(json_str_fixed) + return json_data + except json.JSONDecodeError as e: + logging.error("Failed to fix JSON string.") + return None + else: + logging.error("No 'Final Answer' found in LLM output.") + return None + return [] if expected_type == 'list' else {} + + +PROJECT_DIR = osp.dirname(osp.dirname(__file__)) +DATASET_DIR = osp.join(PROJECT_DIR, "dataset") + +SYSTEM_PROMPT_QUERY_PARSER = """ +# Role: +You are an AI assistant specialized in parsing queries related to medical image segmentation. Your task is to split the query into its imaging modality and segmentation target with as much specificity as possible. + +## Skills: +1. Query Parsing: +- Identify and extract the imaging modality, including specific subtypes (e.g., CT contrast, CT non-contrast, MR ADC, MR FLAIR, etc.) from the query. +- If specific subtypes are not mentioned, extract only the main modality (e.g., CT, MR). +- Identify and extract the segmentation target (e.g., brain, liver, tumor, etc.). + +## Constraints: +1. Response Format +- MUST use the following JSON format for the Final Answer: +Final Answer: {"modality": "modality_name", "target": "segmentation_target"} +Here are some examples: +{"modality": "CT", "target": "liver"} +{"modality": "MR", "target": "spleen"} +{"modality": "MR ADC", "target": "brain tumor"} +""" + +FORMATTED_OUTPUT_DATASET_LIST = '["AbdomenAtlasMini", "AMOS22_Task2", "FLARE23"]' +FORMATTED_OUTPUT_DATASET_VALUE_DICT = '{"AbdomenAtlasMini": [1], "AMOS22_Task2": [2, 3], "FLARE23": [15]}' + +SYSTEM_PROMPT_MODALITY_FILTER = f""" +# Role: +You are an AI assistant specializing in medical image segmentation datasets. Your task is to filter relevant datasets based on the specified imaging modalities. + +## Skills: +1. Dataset Selection by Modality: +- Accurately identify datasets that match the specified imaging modalities. +- Handle composite modalities like "MR or CT", treating them as both "MR" and "CT". +- Include variants such as "head CT", "non-contrast CT", or "contrast CT" under "CT". + +## Constraints: +1. Response Format +- MUST use the following JSON format for the Final Answer: +Final Answer: ["Dataset1", "Dataset2"] + + +## Edge Cases: +- If no datasets match the specified imaging modalities are found, return an empty list: [] + +Here is the dataset information (only name and input modalities are provided): +""" + +SYSTEM_PROMPT_RELEVANCE_FILTER = f""" +# Role: +You are an AI assistant specializing in medical image segmentation datasets. Your task is to filter datasets based on their labels for relevance to the specified segmentation target. + +## Skills: +1. Dataset Relevance Filtering: +- Identify datasets where a single label or a **combination of labels** semantically match the segmentation target, regardless of word order. +- Consider combinations where the label names contain all necessary components to form the target, even if the order differs. +- **Do not** include datasets where labels only partially match or where combinations do not fully represent the exact target. + +2. Thought Process: +You MUST use this format to think which datasets and their labels should be chosen: +``` +Question: the input question you must answer +Thought: you should always think about what to do +Action: the action to take +Action Input: the input to the action +Observation: the result of the action +... (this process can repeat multiple times) +Thought: I now know the final answer +Final Answer: ["Dataset1", "Dataset2"] +``` + +## Constraints: +1. Response Format +- MUST use the following JSON format for the Final Answer: +Final Answer: ["Dataset1", "Dataset2"] + + +## Edge Cases: +- If no relevant datasets are found, return an empty list: [] + +Here is the dataset information (only name and labels are provided): +""" + +SYSTEM_PROMPT_LABEL_SELECTION = """ +# Role: +You are an AI expert in medical image segmentation. Your task is to match datasets and their relevant labels based on the specified segmentation target. + +## Skills: +1. Dataset & Label Selection +- Choose datasets and corresponding labels based on the input question. +- Use the provided dataset information to ensure compatibility. +- Validate the selected labels can cover the need to segment the requested targets. + +2. Thought Process: +You MUST use this format to think which datasets and their labels should be chosen: +``` +Question: the input question you must answer +Thought: you should always think about what to do +Action: the action to take +Action Input: the input to the action +Observation: the result of the action +... (this process can repeat multiple times) +Thought: I now know the final answer +Final Answer: {"Dataset1": [1, 2], "Dataset2": [3]} +``` + +## Constraints: +1. Response Format +- MUST use the following JSON format for the Final Answer: +Final Answer: {"Dataset1": [1, 2], "Dataset2": [3]} + +## Edge Cases: +- If no relevant datasets or labels are found, return an empty dictionary: {} + +Here is the dataset information: +""" + +AUTOGEN_CONFIG_PATH = osp.join(osp.dirname(osp.dirname(__file__)), "OAI_CONFIG_LIST") + + +def chat_with_llm_c2f(query, model, autogen_config_path): + """ + query: A natural language string specifying the desired segmentation task. + model: Name of the LLM model to use. + autogen_config_path: Path to the configuration file for the autogen agent. + """ + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + logging.info("步骤 1: 查询解析代理") + # Step 1: Query Parsing Agent + query_parser_agent = ConversableAgent( + name="query_parser_agent", + system_message=SYSTEM_PROMPT_QUERY_PARSER, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + # 定义不同的查询前缀 + query_prefixes = [ + "My query is: ", + "I have a question: ", + "Can you help with: " + ] + max_attempts = 3 + attempt = 0 + parsed_query = None # 初始化 parsed_query + + while attempt < max_attempts: + prefix = query_prefixes[attempt % len(query_prefixes)] + message = prefix + query + query_parsing_result = user_proxy.initiate_chat( + query_parser_agent, + message=message, + max_turns=1 + ) + + # 从查询解析器的输出中提取 JSON + temp_parsed_query = extract_json_from_output(query_parsing_result.summary, expected_type='dict') + + if temp_parsed_query: + modality = temp_parsed_query.get("modality") + target = temp_parsed_query.get("target") or temp_parsed_query.get("targets") + if modality and target: + parsed_query = temp_parsed_query + logging.info(f"Query parsed successfully: {parsed_query}") + break + else: + logging.warning(f"Attempt {attempt + 1}: Parsed query missing 'modality' or 'target'.") + else: + logging.warning(f"Attempt {attempt + 1} failed to parse query.") + + attempt += 1 + + if not parsed_query: + logging.error("Failed to parse query. Exiting.") + return {} + + # 确保 target 是字符串 + modality = parsed_query.get("modality") + target = parsed_query.get("target") or parsed_query.get("targets") + target = str(target) if target is not None and not isinstance(target, str) else target + if not modality or not target: + plogging.error("Parsed query is incomplete. Exiting.") + return {} + + logging.info("步骤 2: 对数据集针对模态进行过滤") + # Step 2: Modality Filtering Agent + + # 获取所有数据集名称作为集合 + all_dataset_names = set( + [osp.splitext(osp.basename(path))[0] for path in glob(osp.join(DATASET_DIR, "*.json"))] + ) + # 初始化已选择和未选择的数据集 + modality_filtered_datasets = set() + unselected_datasets = all_dataset_names.copy() + + max_iterations = 3 + iteration = 0 + + # 定义不同的模态过滤消息 + modality_messages = [ + f"I would like to segment an image with the modality {modality}.", + f"The target modality of the image is {modality}.", + f"Please segment the image with modality {modality}." + ] + + while iteration < max_iterations and unselected_datasets: + + logging.info(f"Modality filtering iteration {iteration+1}") + + # 加载未选择数据集的模态信息 + unselected_dataset_info_modality = load_dataset_info( + DATASET_DIR, keys_to_keep=["name", "input_modalities"], file_list=unselected_datasets + ) + + # 创建包含更新后的数据集信息的系统提示 + current_system_prompt = SYSTEM_PROMPT_MODALITY_FILTER + unselected_dataset_info_modality + + # 创建新的模态过滤代理 + modality_filter_agent = ConversableAgent( + name="modality_filter_agent", + system_message=current_system_prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + max_attempts = 3 + attempt = 0 + newly_selected_datasets_raw = None + while attempt < max_attempts: + modality_message = modality_messages[(iteration+attempt) % len(modality_messages)] + modality_filter_result = user_proxy.initiate_chat(modality_filter_agent, + message=modality_message, + max_turns=1) + newly_selected_datasets_raw = extract_json_from_output(modality_filter_result.summary, + expected_type='list') + if isinstance(newly_selected_datasets_raw, list) and all(isinstance(item, str) for item in newly_selected_datasets_raw): + break + + attempt += 1 + + if not (isinstance(newly_selected_datasets_raw, list) and all(isinstance(item, str) for item in newly_selected_datasets_raw)): + return {} + + newly_selected_datasets = set(newly_selected_datasets_raw) + + if not newly_selected_datasets: + break + + modality_filtered_datasets.update(newly_selected_datasets) + unselected_datasets -= newly_selected_datasets + + iteration += 1 + + # 检查是否有任何数据集被选中 + if not modality_filtered_datasets: + logging.error("No datasets match the modality filter after multiple iterations. Exiting.") + return {} + + # 循环结束后，打印所有已选择的数据集 + logging.info("Final selected datasets after modality filtering:") + logging.info(sorted(modality_filtered_datasets)) + + logging.info("步骤 3: 对数据集针对类别进行过滤") + # Step 3: Relevance Filtering Agent + + # 初始化已选择和未选择的数据集 + relevance_filtered_datasets = set() + unselected_datasets = set(modality_filtered_datasets) + + max_iterations = 2 + iteration = 0 + + # 定义不同的相关性过滤消息 + target_messages = [ + f"I'm interested in segmenting the {target}.", + f"The segmentation target is {target}.", + f"I want to segment the {target} in the image." + ] + while iteration < max_iterations and unselected_datasets: + + logging.info(f"Relevance filtering iteration {iteration+1}") + + # 加载未选择数据集的相关信息 + unselected_dataset_info = load_dataset_info( + DATASET_DIR, + keys_to_keep=["name", "labels"], + file_list=unselected_datasets + ) + + # 创建包含更新后的数据集信息的系统提示 + current_system_prompt = SYSTEM_PROMPT_RELEVANCE_FILTER + unselected_dataset_info + + relevance_filter_agent = ConversableAgent( + name="relevance_filter_agent", + system_message=current_system_prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + max_attempts = 3 + attempt = 0 + newly_selected_datasets_raw = None + + while attempt < max_attempts: + target_message = target_messages[(iteration + attempt) % len(target_messages)] + relevance_filter_result = user_proxy.initiate_chat( + relevance_filter_agent, + message=target_message, + max_turns=1 + ) + + newly_selected_datasets_raw = extract_json_from_output( + relevance_filter_result.summary, + expected_type='list' + ) + + if isinstance(newly_selected_datasets_raw, list) and all(isinstance(item, str) for item in newly_selected_datasets_raw): + break + + attempt += 1 + + if not (isinstance(newly_selected_datasets_raw, list) and all(isinstance(item, str) for item in newly_selected_datasets_raw)): + return {} + + newly_selected_datasets = set(newly_selected_datasets_raw) + + if not newly_selected_datasets and iteration > 0: + break + + # 更新已选择和未选择的数据集 + relevance_filtered_datasets.update(newly_selected_datasets) + unselected_datasets -= newly_selected_datasets + + iteration += 1 + + # 检查是否有任何数据集被选中 + if not relevance_filtered_datasets: + logging.error("No datasets match the relevance filter after multiple iterations. Exiting.") + return {} + + # 循环结束后，打印所有已选择的数据集 + logging.info("Final selected datasets after relevance filtering:") + logging.info(sorted(relevance_filtered_datasets)) + + logging.info("步骤 4: 选择数据集和类别") + # Step 4: Label Selection Agent + + # 修改标签选择部分的逻辑，模仿相关性过滤的迭代过程 + label_selected_labels = {} + unselected_datasets = relevance_filtered_datasets + max_iterations = 1 # 一轮选择 + 一轮检查 + iteration = 0 + + # 定义标签选择的消息 + label_selection_messages = [ + f"I want to segment the {target} in the {modality} image", + f"The segmentation is {target} in {modality}", + f"Please segment the {target} in the {modality} images" + ] + + while iteration < max_iterations and unselected_datasets: + + logging.info(f"Label selection iteration {iteration+1}") + + # 加载当前未选数据集的详细信息 + current_dataset_info = load_dataset_info( + DATASET_DIR, + keys_to_keep=["name", "input_modalities", "labels"], + file_list=unselected_datasets + ) + + # 创建新的系统提示，包含当前数据集信息 + current_system_prompt = SYSTEM_PROMPT_LABEL_SELECTION + current_dataset_info + + # 创建新的标签选择代理 + label_selection_agent = ConversableAgent( + name="label_selection_agent", + system_message=current_system_prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + max_attempts = 3 + attempt = 0 + dataset_label_dict = None + + while attempt < max_attempts: + + label_selection_message = label_selection_messages[(iteration+attempt)%len(label_selection_messages)] + + label_selection_result = user_proxy.initiate_chat( + label_selection_agent, + message=label_selection_message, + max_turns=1 + ) + + # 从标签选择代理的输出中提取 JSON + dataset_label_dict = extract_json_from_output( + label_selection_result.summary, + expected_type='dict' + ) + + # 判断是否成功提取，并且是包含整数或浮点数列表的字典 + if isinstance(dataset_label_dict, dict) and dataset_label_dict and all( + isinstance(dataset, str) and isinstance(labels, list) and all(isinstance(x, (int, float)) for x in labels) + for dataset, labels in dataset_label_dict.items() + ): + break + attempt += 1 + + if dataset_label_dict is None or not all( + isinstance(labels, list) + and all(isinstance(x, (int, float)) for x in labels) + for labels in dataset_label_dict.values() + ): + dataset_label_dict = {} + + # 合并已选择的标签 + for dataset, labels in dataset_label_dict.items(): + if dataset in label_selected_labels: + # 合并标签列表，避免重复 + existing_labels = set(label_selected_labels[dataset]) + new_labels = set(labels) + combined_labels = existing_labels.union(new_labels) + label_selected_labels[dataset] = list(combined_labels) + else: + label_selected_labels[dataset] = labels + + # 移除已选择的标签的数据集 + unselected_datasets -= set(dataset_label_dict.keys()) + + iteration += 1 + + logging.info("最终选择的标签:") + logging.info(label_selected_labels) + + return label_selected_labels + + +def main(): + result = chat_with_llm_c2f("I want to segment the liver vessel in this CT image", "claude-3-5-haiku-20241022", AUTOGEN_CONFIG_PATH) + print(result) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/agents/c2f_tree_search1.py b/agents/c2f_tree_search1.py new file mode 100644 index 0000000..db2d944 --- /dev/null +++ b/agents/c2f_tree_search1.py @@ -0,0 +1,504 @@ +import json +import os.path as osp +from glob import glob + +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent + + +# Function to load JSON dataset information from a directory +def load_dataset_info(json_dir, + keys_to_keep=["name", "description", "input_modalities", "labels"], + file_list=None): + """ + Complete list of keys: ["name", "description", "input_modalities", "labels", "models"] + """ + json_info_list = [] + # Append .json extension to the filenames in the list + if file_list is not None: + file_list = {name + ".json" for name in file_list} + + for json_path in glob(osp.join(json_dir, "*.json")): + file_name = osp.basename(json_path) + + # Check if the file is in the specified file list + if file_list is None or file_name in file_list: + with open(json_path, encoding='utf-8') as f: + json_content = json.load(f) + # Keep only specified keys + filtered_content = {key: json_content.get(key) for key in keys_to_keep} + json_info_list.append(json.dumps(filtered_content)) + + return "\n".join(json_info_list) + + +# Function to extract JSON data from LLM output +def extract_json_from_output(llm_output, expected_type='list'): + import re + if expected_type == 'list': + pattern = r"Final Answer:\s*(\[[\s\S]*?\])" + elif expected_type == 'dict': + pattern = r"Final Answer:\s*(\{[\s\S]*?\})" + else: + raise ValueError("expected_type must be 'list' or 'dict'") + + match = re.search(pattern, llm_output) + if match: + json_str = match.group(1) + # Remove comments from JSON string + json_str = re.sub(r'//.*', '', json_str) # Remove single-line comments + json_str = re.sub(r'/\*[\s\S]*?\*/', '', json_str) # Remove multi-line comments + try: + json_data = json.loads(json_str) + return json_data + except json.JSONDecodeError as e: + print("JSON decoding error:", e) + print("Attempting to fix JSON string...") + json_str_fixed = json_str.replace("'", '"') + try: + json_data = json.loads(json_str_fixed) + return json_data + except json.JSONDecodeError as e: + print("Failed to fix JSON string.") + return None + else: + print("No 'Final Answer' found in LLM output.") + return None + return [] if expected_type == 'list' else {} + + +PROJECT_DIR = osp.dirname(osp.dirname(__file__)) +DATASET_DIR = osp.join(PROJECT_DIR, "dataset") + +SYSTEM_PROMPT_QUERY_PARSER = """ +# Role: +You are an AI assistant specialized in parsing queries related to medical image segmentation. Your task is to split the query into its imaging modality and segmentation target with as much specificity as possible. + +## Skills: +1. Query Parsing: +- Identify and extract the imaging modality, including specific subtypes (e.g., CT contrast, CT non-contrast, MR ADC, MR FLAIR, etc.) from the query. +- If specific subtypes are not mentioned, extract only the main modality (e.g., CT, MR). +- Identify and extract the segmentation target (e.g., brain, liver, tumor, etc.). + +2. Thought Process: +Use this format for reasoning: +Question: the input question +Thought: analyze the input query to extract "modality" and "target". +Action: return a dictionary with the parsed components. +Observation: the parsed components. + +Final Answer: a dictionary in the format: +{"modality": "modality_name", "target": "segmentation_target"} +""" + +FORMATTED_OUTPUT_DATASET_LIST = '["AbdomenAtlasMini", "AMOS22_Task2", "FLARE23"]' +FORMATTED_OUTPUT_DATASET_VALUE_DICT = '{"AbdomenAtlasMini": [1], "AMOS22_Task2": [2, 3], "FLARE23": [15]}' + +SYSTEM_PROMPT_MODALITY_FILTER = f""" +# Role: +You are an AI assistant specializing in medical image segmentation datasets. Your task is to filter relevant datasets based on the specified imaging modalities. + +## Skills: +1. Dataset Selection by Modality: +- Match datasets using synonyms, abbreviations, and partial matches. +- Handle composite modalities like "MR or CT", treating them as both "MR" and "CT". +- Include variants such as "head CT", "non-contrast CT", or "contrast CT" under "CT". +- Perform case-insensitive comparisons. + +2. Thought Process: +Use this format for reasoning: +Modality Filter: the input filter condition specifying imaging modalities. +Thought: analyze the input modalities and match them with the datasets' modalities. +Action: list datasets that match the specified modalities. +Observation: datasets that match the filter. + +Final Answer: a list of matching datasets in valid JSON format without using code blocks. + +## Response Format: +- Return the answer as a JSON list of dataset names, e.g., ["Dataset1", "Dataset2"] +- Answer after **Final Answer** + +## Edge Cases: +- If no relevant datasets are found, return an empty list: [] + +Here is the dataset information (only name and input modalities are provided): +""" + +SYSTEM_PROMPT_RELEVANCE_FILTER = f""" +# Role: +You are an AI assistant specializing in medical image segmentation datasets. Your task is to filter datasets based on their labels for relevance to the specified segmentation target. + +## Skills: +1. Dataset Relevance Filtering: +- Analyze the labels of the datasets. +- Identify datasets where a single label or a **combination of labels** semantically match the segmentation target, regardless of word order. +- Consider combinations where the label names contain all necessary components to form the target, even if the order differs. +- **Do not** include datasets where labels only partially match or where combinations do not fully represent the exact target. + +2. Thought Process: +You MUST use this format to think which datasets and their labels should be chosen: +``` +Question: the input question you must answer +Thought: you should always think about what to do +Action: the action to take +Action Input: the input to the action +Observation: the result of the action +... (this process can repeat multiple times) +Thought: I now know the final answer +Final Answer: the final answer to the original input question in JSON format (Looking for examples below) +``` + +Final Answer: a list of relevant datasets in valid JSON format without using code blocks. + +## Response Format: +- Return the answer as a JSON list: +{FORMATTED_OUTPUT_DATASET_LIST} + +## Edge Cases: +- If no relevant datasets within provided datasets are found, return an empty list: [] + +Here is the dataset information: +""" + +SYSTEM_PROMPT_LABEL_SELECTION = f""" +# Role: +You are an AI expert in medical image segmentation. Your task is to match datasets and their relevant labels based on the specified segmentation target. + +## Skills: +1. Dataset & Label Selection +- Choose datasets and corresponding labels based on the input question. +- Use the provided dataset information to ensure compatibility. +- Validate the selected labels can cover the need to segment the requested targets. + +2. Thought Process: +You MUST use this format to think which datasets and their labels should be chosen: +``` +Question: the input question you must answer +Thought: you should always think about what to do +Action: the action to take +Action Input: the input to the action +Observation: the result of the action +... (this process can repeat multiple times) +Thought: I now know the final answer +Final Answer: the final answer to the original input question in JSON format (Looking for examples below) +``` + +## Response Format: +- Return the answer as a dictionary: +{FORMATTED_OUTPUT_DATASET_VALUE_DICT} + +## Edge Cases: +- If no relevant datasets or labels are found, return an empty dictionary. + +Here is the dataset information: +""" + +AUTOGEN_CONFIG_PATH = osp.join(osp.dirname(osp.dirname(__file__)), "OAI_CONFIG_LIST") + + +def chat_with_llm_c2f1(query, model, autogen_config_path): + """ + query: A natural language string specifying the desired segmentation task. + model: Name of the LLM model to use. + autogen_config_path: Path to the configuration file for the autogen agent. + """ + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + # Step 1: Query Parsing Agent + query_parser_agent = ConversableAgent( + name="query_parser_agent", + system_message=SYSTEM_PROMPT_QUERY_PARSER, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + max_attempts = 3 + attempt = 0 + while attempt < max_attempts: + query_parsing_result = user_proxy.initiate_chat(query_parser_agent, + message=query, + max_turns=1) + + # 从查询解析器的输出中提取 JSON + parsed_query = extract_json_from_output(query_parsing_result.summary, expected_type='dict') + + if parsed_query is not None: + # 成功解析，退出循环 + break + + attempt += 1 + + if not parsed_query: + print("Failed to parse query. Exiting.") + return {} + + modality = parsed_query.get("modality") + target = parsed_query.get("target") or parsed_query.get("targets") + target = str(target) if target is not None and not isinstance(target, str) else target + if not modality or not target: + print("Parsed query is incomplete. Exiting.") + return {} + + # Step 2: Modality Filtering Agent + + # Get all dataset names as a set + all_dataset_names = set( + [osp.splitext(osp.basename(path))[0] for path in glob(osp.join(DATASET_DIR, "*.json"))]) + + # Initialize selected and unselected datasets + modality_filtered_datasets = set() + unselected_datasets = all_dataset_names.copy() + + max_iterations = 3 + iteration = 0 + + modality_messages = [ + f"I want to segment that a image of the modality {modality} image", f"The image target modality is {modality}", + f"Please segment the {modality} modality image" + ] + + while iteration < max_iterations and unselected_datasets: + iteration += 1 + print(f"Modality filtering iteration {iteration}") + + # Load modality info for unselected datasets + unselected_dataset_info_modality = load_dataset_info( + DATASET_DIR, keys_to_keep=["name", "input_modalities"], file_list=unselected_datasets) + + # Create a new system prompt with updated dataset info + current_system_prompt = SYSTEM_PROMPT_MODALITY_FILTER + unselected_dataset_info_modality + + # Create a new agent with the updated system prompt + modality_filter_agent = ConversableAgent( + name="modality_filter_agent", + system_message=current_system_prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + max_attempts = 3 + + attempt = 0 + while attempt < max_attempts: + modality_message = modality_messages[attempt+iteration-1] if attempt+iteration-1 < len( + modality_messages) else modality_messages[-1] + # Use the agent to filter datasets based on modality + modality_filter_result = user_proxy.initiate_chat(modality_filter_agent, + message=modality_message, + max_turns=1) + # Extract the newly selected datasets + newly_selected_datasets_raw = extract_json_from_output(modality_filter_result.summary, + expected_type='list') + if newly_selected_datasets_raw is not None: + break + attempt += 1 + + if newly_selected_datasets_raw is None: + continue + + # Process the list to extract dataset names + newly_selected_datasets = set() + for item in newly_selected_datasets_raw: + if isinstance(item, dict): + if 'name' in item: + newly_selected_datasets.add(item['name']) + else: + print("Warning: Dictionary item without 'name' key:", item) + elif isinstance(item, str): + newly_selected_datasets.add(item) + else: + print("Warning: Unexpected item type:", type(item), item) + + # If no new datasets are selected, exit the loop + if not newly_selected_datasets: + break + + # Update selected and unselected datasets + modality_filtered_datasets.update(newly_selected_datasets) + unselected_datasets -= newly_selected_datasets + + # 循环结束后，打印所有已选择的数据集 + print("\nFinal selected datasets after modality filtering:") + print(sorted(modality_filtered_datasets)) + print() + + # Check if any datasets have been selected + if not modality_filtered_datasets: + print("No datasets match the modality filter after multiple iterations. Exiting.") + return {} + + # Convert the set to a list for subsequent steps + modality_filtered_datasets = list(modality_filtered_datasets) + + # Step 3: Relevance Filtering Agent + + # Initialize selected and unselected datasets for relevance filtering + relevance_filtered_datasets = set() + unselected_datasets = set(modality_filtered_datasets) + + max_iterations = 3 + iteration = 0 + target_messages = [ + f"I want to segment the {target} in the image", f"The segmentation target is {target}", f"Please segment {target}", + f"What is the anatomical region of the {target} in the body, and can you provide datasets related to that region?" + ] + while iteration < max_iterations and unselected_datasets: + + target_message = target_messages[iteration] if iteration < len( + target_messages) else target_messages[-1] + + iteration += 1 + print(f"Relevance filtering iteration {iteration}") + + # Load dataset info for unselected datasets + unselected_dataset_info = load_dataset_info(DATASET_DIR, + keys_to_keep=["name", "labels"], + file_list=unselected_datasets) + + # Create a new system prompt with updated dataset info + current_system_prompt = SYSTEM_PROMPT_RELEVANCE_FILTER + unselected_dataset_info + + # Create a new agent with the updated system prompt + relevance_filter_agent = ConversableAgent( + name="relevance_filter_agent", + system_message=current_system_prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + max_attempts = 3 + attempt = 0 + while attempt < max_attempts: + target_message = target_messages[(iteration + attempt) % len(target_messages)] + # Use the agent to filter datasets based on relevance + relevance_filter_result = user_proxy.initiate_chat(relevance_filter_agent, + message=target_message, + max_turns=1) + + # Extract the newly selected datasets + newly_selected_datasets_raw = extract_json_from_output(relevance_filter_result.summary, + expected_type='list') + + if newly_selected_datasets_raw is not None: + break + attempt += 1 + + # Process the list to extract dataset names + newly_selected_datasets = set() + for item in newly_selected_datasets_raw: + if isinstance(item, dict): + if 'name' in item: + newly_selected_datasets.add(item['name']) + else: + print("Warning: Dictionary item without 'name' key:", item) + elif isinstance(item, str): + newly_selected_datasets.add(item) + else: + print("Warning: Unexpected item type:", type(item), item) + + # If no new datasets are selected, exit the loop + if not newly_selected_datasets and iteration > 1: + break + + # Update selected and unselected datasets + relevance_filtered_datasets.update(newly_selected_datasets) + unselected_datasets -= newly_selected_datasets + + # Check if any datasets have been selected + if not relevance_filtered_datasets: + print("No datasets match the relevance filter after multiple iterations. Exiting.") + return {} + + # 循环结束后，打印所有已选择的数据集 + print("\nFinal selected datasets after target filtering:") + print(sorted(relevance_filtered_datasets)) + print() + + # Convert the set to a list for subsequent steps + relevance_filtered_datasets = list(relevance_filtered_datasets) + + # Step 4: Label Selection Agent + + # 修改标签选择部分的逻辑，模仿相关性过滤的迭代过程 + label_selected_labels = {} + unselected_datasets = set(relevance_filtered_datasets) + max_iterations = 2 # 一轮选择 + 一轮检查 + iteration = 0 + + # 定义标签选择的消息 + label_selection_messages = [ + f"I want to segment the {target} in the {modality} image", + f"The segmentation is {target} in {modality}", + f"Please segment the {target} in the {modality} images" + ] + + while iteration < max_iterations and unselected_datasets: + iteration += 1 + print(f"Label selection iteration {iteration}") + + # 加载当前未选数据集的详细信息 + current_dataset_info = load_dataset_info(DATASET_DIR, + keys_to_keep=["name", "input_modalities", "labels"], + file_list=unselected_datasets) + + # 创建新的系统提示，包含当前数据集信息 + current_system_prompt = SYSTEM_PROMPT_LABEL_SELECTION + current_dataset_info + + # 创建新的标签选择代理 + label_selection_agent = ConversableAgent( + name="label_selection_agent", + system_message=current_system_prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + # 发送消息并尝试获取标签选择结果 + attempt_inner = 0 + while attempt_inner < max_attempts: + # 确保索引不会越界 + label_message_index = attempt_inner + iteration - 1 if (attempt_inner + iteration - 1) < len(label_selection_messages) else -1 + label_message = label_selection_messages[label_message_index] if label_message_index >= 0 else label_selection_messages[-1] + label_selection_result = user_proxy.initiate_chat(label_selection_agent, + message=label_message, + max_turns=1) + + # 从标签选择代理的输出中提取 JSON + dataset_label_dict = extract_json_from_output(label_selection_result.summary, + expected_type='dict') + + if dataset_label_dict is not None: + break + attempt_inner += 1 + + if dataset_label_dict is None or not all( + isinstance(labels, list) + and all(isinstance(x, (int, float)) for x in labels) + for labels in dataset_label_dict.values() + ): + continue + + # 合并已选择的标签 + for dataset, labels in dataset_label_dict.items(): + if dataset in label_selected_labels: + # 合并标签列表，避免重复 + existing_labels = set(label_selected_labels[dataset]) + new_labels = set(labels) + combined_labels = existing_labels.union(new_labels) + label_selected_labels[dataset] = list(combined_labels) + else: + label_selected_labels[dataset] = labels + + # 移除已选择的标签的数据集 + unselected_datasets -= set(dataset_label_dict.keys()) + + return label_selected_labels + + +def main(): + result = chat_with_llm_c2f("I want to segment the liver vessel in this CT image", "claude-3-5-haiku-20241022", AUTOGEN_CONFIG_PATH) + print(result) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/agents/c2f_tree_search2.py b/agents/c2f_tree_search2.py new file mode 100644 index 0000000..9a15392 --- /dev/null +++ b/agents/c2f_tree_search2.py @@ -0,0 +1,518 @@ +import json +import os.path as osp +import logging +from glob import glob + +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent + + +# Function to load JSON dataset information from a directory +def load_dataset_info(json_dir, + keys_to_keep=["name", "description", "input_modalities", "labels"], + file_list=None): + """ + Complete list of keys: ["name", "description", "input_modalities", "labels", "models"] + """ + json_info_list = [] + # Append .json extension to the filenames in the list + if file_list is not None: + file_list = {name + ".json" for name in file_list} + + for json_path in glob(osp.join(json_dir, "*.json")): + file_name = osp.basename(json_path) + + # Check if the file is in the specified file list + if file_list is None or file_name in file_list: + with open(json_path, encoding='utf-8') as f: + json_content = json.load(f) + # Keep only specified keys + filtered_content = {key: json_content.get(key) for key in keys_to_keep} + json_info_list.append(json.dumps(filtered_content)) + + return "\n".join(json_info_list) + + +# Function to extract JSON data from LLM output +def extract_json_from_output(llm_output, expected_type='list'): + import re + if expected_type == 'list': + pattern = r"Final Answer:\s*(\[[\s\S]*?\])" + elif expected_type == 'dict': + pattern = r"Final Answer:\s*(\{[\s\S]*?\})" + else: + raise ValueError("expected_type must be 'list' or 'dict'") + + match = re.search(pattern, llm_output) + if match: + json_str = match.group(1) + # Remove comments from JSON string + json_str = re.sub(r'//.*', '', json_str) # Remove single-line comments + json_str = re.sub(r'/\*[\s\S]*?\*/', '', json_str) # Remove multi-line comments + try: + json_data = json.loads(json_str) + return json_data + except json.JSONDecodeError as e: + logging.error("JSON decoding error:", e) + logging.info("Attempting to fix JSON string...") + json_str_fixed = json_str.replace("'", '"') + try: + json_data = json.loads(json_str_fixed) + return json_data + except json.JSONDecodeError as e: + logging.error("Failed to fix JSON string.") + return None + else: + logging.error("No 'Final Answer' found in LLM output.") + return None + return [] if expected_type == 'list' else {} + + +PROJECT_DIR = osp.dirname(osp.dirname(__file__)) +DATASET_DIR = osp.join(PROJECT_DIR, "dataset") + +SYSTEM_PROMPT_QUERY_PARSER = """ +# Role: +You are an AI assistant specialized in parsing queries related to medical image segmentation. Your task is to split the query into its imaging modality and segmentation target with as much specificity as possible. + +## Skills: +1. Query Parsing: +- Identify and extract the imaging modality, including specific subtypes (e.g., CT contrast, CT non-contrast, MR ADC, MR FLAIR, etc.) from the query. +- If specific subtypes are not mentioned, extract only the main modality (e.g., CT, MR). +- Identify and extract the segmentation target (e.g., brain, liver, tumor, etc.). + +## Constraints: +1. Response Format +- MUST use the following JSON format for the Final Answer: +Final Answer: {"modality": "modality_name", "target": "segmentation_target"} +Here are some examples: +{"modality": "CT", "target": "liver"} +{"modality": "MR", "target": "spleen"} +{"modality": "MR ADC", "target": "brain tumor"} +""" + +FORMATTED_OUTPUT_DATASET_LIST = '["AbdomenAtlasMini", "AMOS22_Task2", "FLARE23"]' +FORMATTED_OUTPUT_DATASET_VALUE_DICT = '{"AbdomenAtlasMini": [1], "AMOS22_Task2": [2, 3], "FLARE23": [15]}' + +SYSTEM_PROMPT_MODALITY_FILTER = f""" +# Role: +You are an AI assistant specializing in medical image segmentation datasets. Your task is to filter relevant datasets based on the specified imaging modalities. + +## Skills: +1. Dataset Selection by Modality: +- Accurately identify datasets that match the specified imaging modalities. +- Handle composite modalities like "MR or CT", treating them as both "MR" and "CT". +- Include variants such as "head CT", "non-contrast CT", or "contrast CT" under "CT". + +## Constraints: +1. Response Format +- MUST use the following JSON format for the Final Answer: +Final Answer: ["Dataset1", "Dataset2"] + + +## Edge Cases: +- If no datasets match the specified imaging modalities are found, return an empty list: [] + +Here is the dataset information (only name and input modalities are provided): +""" + +SYSTEM_PROMPT_RELEVANCE_FILTER = f""" +# Role: +You are an AI assistant specializing in the task of medical image segmentation datasets. Your task is to filter datasets based on their labels for relevance to the specified segmentation target. + +## Skills: +1. Dataset Relevance Filtering: +- Identify datasets where a single label or a **combination of labels** semantically match the segmentation target, regardless of word order. +- Consider combinations where the label names contain all necessary components to form the target, even if the order differs. +- **Do not** include datasets where labels only partially match or where combinations do not fully represent the exact target. + +2. Thought Process: +You MUST use this format to think which datasets and their labels should be chosen: +``` +Question: the input question you must answer +Thought: you should always think about what to do +Action: the action to take +Action Input: the input to the action +Observation: the result of the action +... (this process can repeat multiple times) +Thought: I now know the final answer +Final Answer: ["Dataset1", "Dataset2"] +``` + +## Constraints: +1. Response Format +- MUST use the following JSON format for the Final Answer: +Final Answer: ["Dataset1", "Dataset2"] + + +## Edge Cases: +- If no relevant datasets are found, return an empty list: [] + +Here is the dataset information (only name and labels are provided): +""" + +SYSTEM_PROMPT_LABEL_SELECTION = """ +# Role: +You are an AI expert in the task of medical image segmentation. Your task is to match datasets and their relevant labels based on the specified segmentation target. + +## Skills: +1. Dataset & Label Selection +- Choose datasets and corresponding labels based on the input question. +- Use the provided dataset information to ensure compatibility. +- Validate the selected labels can cover the need to segment the requested targets. + +2. Thought Process: +You MUST use this format to think which datasets and their labels should be chosen: +``` +Question: the input question you must answer +Thought: you should always think about what to do +Action: the action to take +Action Input: the input to the action +Observation: the result of the action +... (this process can repeat multiple times) +Thought: I now know the final answer +Final Answer: {"Dataset1": [1, 2], "Dataset2": [3]} +``` + +## Constraints: +1. Response Format +- MUST use the following JSON format for the Final Answer: +Final Answer: {"Dataset1": [1, 2], "Dataset2": [3]} + +## Edge Cases: +- If no relevant datasets or labels are found, return an empty dictionary: {} + +Here is the dataset information: +""" + +AUTOGEN_CONFIG_PATH = osp.join(osp.dirname(osp.dirname(__file__)), "OAI_CONFIG_LIST") + + +def chat_with_llm_c2f2(query, model, autogen_config_path): + """ + query: A natural language string specifying the desired segmentation task. + model: Name of the LLM model to use. + autogen_config_path: Path to the configuration file for the autogen agent. + """ + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + logging.info("步骤 1: 查询解析代理") + # Step 1: Query Parsing Agent + query_parser_agent = ConversableAgent( + name="query_parser_agent", + system_message=SYSTEM_PROMPT_QUERY_PARSER, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + # 定义不同的查询前缀 + query_prefixes = [ + "My query is: ", + "I have a question: ", + "Can you help with: " + ] + max_attempts = 3 + attempt = 0 + parsed_query = None # 初始化 parsed_query + + while attempt < max_attempts: + prefix = query_prefixes[attempt % len(query_prefixes)] + message = prefix + query + query_parsing_result = user_proxy.initiate_chat( + query_parser_agent, + message=message, + max_turns=1 + ) + + # 从查询解析器的输出中提取 JSON + temp_parsed_query = extract_json_from_output(query_parsing_result.summary, expected_type='dict') + + if temp_parsed_query: + modality = temp_parsed_query.get("modality") + target = temp_parsed_query.get("target") or temp_parsed_query.get("targets") + if modality and target: + parsed_query = temp_parsed_query + logging.info(f"Query parsed successfully: {parsed_query}") + break + else: + logging.warning(f"Attempt {attempt + 1}: Parsed query missing 'modality' or 'target'.") + else: + logging.warning(f"Attempt {attempt + 1} failed to parse query.") + + attempt += 1 + + if not parsed_query: + logging.error("Failed to parse query. Exiting.") + return {} + + # 确保 target 是字符串 + modality = parsed_query.get("modality") + target = parsed_query.get("target") or parsed_query.get("targets") + target = str(target) if target is not None and not isinstance(target, str) else target + if not modality or not target: + plogging.error("Parsed query is incomplete. Exiting.") + return {} + + logging.info("步骤 2: 对数据集针对模态进行过滤") + # Step 2: Modality Filtering Agent + + # 获取所有数据集名称作为集合 + all_dataset_names = set( + [osp.splitext(osp.basename(path))[0] for path in glob(osp.join(DATASET_DIR, "*.json"))] + ) + # 初始化已选择和未选择的数据集 + modality_filtered_datasets = set() + unselected_datasets = all_dataset_names.copy() + + max_iterations = 3 + iteration = 0 + + # 定义不同的模态过滤消息 + modality_messages = [ + f"I would like to segment an image with the modality {modality}.", + f"The target modality of the image is {modality}.", + f"Please segment the image with modality {modality}." + ] + + while iteration < max_iterations and unselected_datasets: + + logging.info(f"Modality filtering iteration {iteration+1}") + + # 加载未选择数据集的模态信息 + unselected_dataset_info_modality = load_dataset_info( + DATASET_DIR, keys_to_keep=["name", "input_modalities"], file_list=unselected_datasets + ) + + # 创建包含更新后的数据集信息的系统提示 + current_system_prompt = SYSTEM_PROMPT_MODALITY_FILTER + unselected_dataset_info_modality + + # 创建新的模态过滤代理 + modality_filter_agent = ConversableAgent( + name="modality_filter_agent", + system_message=current_system_prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + max_attempts = 3 + attempt = 0 + newly_selected_datasets_raw = None + while attempt < max_attempts: + modality_message = modality_messages[(iteration+attempt) % len(modality_messages)] + modality_filter_result = user_proxy.initiate_chat(modality_filter_agent, + message=modality_message, + max_turns=1) + newly_selected_datasets_raw = extract_json_from_output(modality_filter_result.summary, + expected_type='list') + if isinstance(newly_selected_datasets_raw, list) and all(isinstance(item, str) for item in newly_selected_datasets_raw): + break + + attempt += 1 + + if not (isinstance(newly_selected_datasets_raw, list) and all(isinstance(item, str) for item in newly_selected_datasets_raw)): + return {} + + newly_selected_datasets = set(newly_selected_datasets_raw) + + if not newly_selected_datasets: + break + + modality_filtered_datasets.update(newly_selected_datasets) + unselected_datasets -= newly_selected_datasets + + iteration += 1 + + # 检查是否有任何数据集被选中 + if not modality_filtered_datasets: + logging.error("No datasets match the modality filter after multiple iterations. Exiting.") + return {} + + # 循环结束后，打印所有已选择的数据集 + logging.info("Final selected datasets after modality filtering:") + logging.info(sorted(modality_filtered_datasets)) + + logging.info("步骤 3: 对数据集针对类别进行过滤") + # Step 3: Relevance Filtering Agent + + # 初始化已选择和未选择的数据集 + relevance_filtered_datasets = set() + unselected_datasets = set(modality_filtered_datasets) + + max_iterations = 2 + iteration = 0 + + # 定义不同的相关性过滤消息 + target_messages = [ + f"I'm interested in segmenting the {target}.", + f"The segmentation target is {target}.", + f"I want to segment the {target} in the image." + ] + while iteration < max_iterations and unselected_datasets: + + logging.info(f"Relevance filtering iteration {iteration+1}") + + # 加载未选择数据集的相关信息 + unselected_dataset_info = load_dataset_info( + DATASET_DIR, + keys_to_keep=["name", "labels"], + file_list=unselected_datasets + ) + + # 创建包含更新后的数据集信息的系统提示 + current_system_prompt = SYSTEM_PROMPT_RELEVANCE_FILTER + unselected_dataset_info + + relevance_filter_agent = ConversableAgent( + name="relevance_filter_agent", + system_message=current_system_prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + max_attempts = 3 + attempt = 0 + newly_selected_datasets_raw = None + + while attempt < max_attempts: + target_message = target_messages[(iteration + attempt) % len(target_messages)] + relevance_filter_result = user_proxy.initiate_chat( + relevance_filter_agent, + message=target_message, + max_turns=1 + ) + + newly_selected_datasets_raw = extract_json_from_output( + relevance_filter_result.summary, + expected_type='list' + ) + + if isinstance(newly_selected_datasets_raw, list) and all(isinstance(item, str) for item in newly_selected_datasets_raw): + break + + attempt += 1 + + if not (isinstance(newly_selected_datasets_raw, list) and all(isinstance(item, str) for item in newly_selected_datasets_raw)): + return {} + + newly_selected_datasets = set(newly_selected_datasets_raw) + + if not newly_selected_datasets and iteration > 0: + break + + # 更新已选择和未选择的数据集 + relevance_filtered_datasets.update(newly_selected_datasets) + unselected_datasets -= newly_selected_datasets + + iteration += 1 + + # 检查是否有任何数据集被选中 + if not relevance_filtered_datasets: + logging.error("No datasets match the relevance filter after multiple iterations. Exiting.") + return {} + + # 循环结束后，打印所有已选择的数据集 + logging.info("Final selected datasets after relevance filtering:") + logging.info(sorted(relevance_filtered_datasets)) + + logging.info("步骤 4: 选择数据集和类别") + # Step 4: Label Selection Agent + + # 修改标签选择部分的逻辑，模仿相关性过滤的迭代过程 + label_selected_labels = {} + unselected_datasets = relevance_filtered_datasets + max_iterations = 1 # 一轮选择 + 一轮检查 + iteration = 0 + + # 定义标签选择的消息 + label_selection_messages = [ + f"I want to segment the {target} in the {modality} image", + f"The segmentation is {target} in {modality}", + f"Please segment the {target} in the {modality} images" + ] + + while iteration < max_iterations and unselected_datasets: + + logging.info(f"Label selection iteration {iteration+1}") + + # 加载当前未选数据集的详细信息 + current_dataset_info = load_dataset_info( + DATASET_DIR, + keys_to_keep=["name", "input_modalities", "labels"], + file_list=unselected_datasets + ) + + # 创建新的系统提示，包含当前数据集信息 + current_system_prompt = SYSTEM_PROMPT_LABEL_SELECTION + current_dataset_info + + # 创建新的标签选择代理 + label_selection_agent = ConversableAgent( + name="label_selection_agent", + system_message=current_system_prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + max_attempts = 3 + attempt = 0 + dataset_label_dict = None + + while attempt < max_attempts: + + label_selection_message = label_selection_messages[(iteration+attempt)%len(label_selection_messages)] + + label_selection_result = user_proxy.initiate_chat( + label_selection_agent, + message=label_selection_message, + max_turns=1 + ) + + # 从标签选择代理的输出中提取 JSON + dataset_label_dict = extract_json_from_output( + label_selection_result.summary, + expected_type='dict' + ) + + # 判断是否成功提取，并且是包含整数或浮点数列表的字典 + if isinstance(dataset_label_dict, dict) and dataset_label_dict and all( + isinstance(dataset, str) and isinstance(labels, list) and all(isinstance(x, (int, float)) for x in labels) + for dataset, labels in dataset_label_dict.items() + ): + break + attempt += 1 + + if dataset_label_dict is None or not all( + isinstance(labels, list) + and all(isinstance(x, (int, float)) for x in labels) + for labels in dataset_label_dict.values() + ): + dataset_label_dict = {} + + # 合并已选择的标签 + for dataset, labels in dataset_label_dict.items(): + if dataset in label_selected_labels: + # 合并标签列表，避免重复 + existing_labels = set(label_selected_labels[dataset]) + new_labels = set(labels) + combined_labels = existing_labels.union(new_labels) + label_selected_labels[dataset] = list(combined_labels) + else: + label_selected_labels[dataset] = labels + + # 移除已选择的标签的数据集 + unselected_datasets -= set(dataset_label_dict.keys()) + + iteration += 1 + + logging.info("最终选择的标签:") + logging.info(label_selected_labels) + + return label_selected_labels + + +def main(): + result = chat_with_llm_c2f2("I want to segment the liver vessel in this CT image", "claude-3-5-haiku-20241022", AUTOGEN_CONFIG_PATH) + print(result) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/agents/cot.py b/agents/cot.py new file mode 100644 index 0000000..5e7b314 --- /dev/null +++ b/agents/cot.py @@ -0,0 +1,68 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from prompts.react_2shot import SYSTEM_PROMPT_REACT +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent +from utils.reply_utils import extract_llm_answer + +from prompts.data import DATASET_JSON_INFO + + +FORMATTED_OUTPUT = '{ "thought":"", "answer":{"AbdomenAtlasMini": [1], "AMOS22_Task2": [2,3], "FLARE23": [15]} }' + +FORMATTED_EXAMPLES = '{ "thought":"I need to find all CT datasets containing duodenum. I should first check the modality, then the labels.", "answer":{"AbdomenAtlasMini":[13], "AMOS22_Task2":[12], "FLARE23":[19]} }\n{ "thought":"To segment adrenal gland in CT, I need to first find all dataset with adrenal gland.", "answer":{"AMOS22_Task2":[11,12], "FLARE23":[8,9], "TotalSegmentator_v2":[7,8]} }' + +MISSING_MODALITY = '{ "thought":"No suitable modality", "answer": {} }' +NO_DATASET_FOUND = '{ "thought":"No suitable dataset", "answer": {} }' + +SYSTEM_PROMPT_COT = f""" +# Role: +You are an AI expert on medical image segmentation, specializing in addressing clinical tasks efficiently. You will answer the input question as best you can. + +## Skills: +1. Dataset & Label Selection +- Choose datasets and corresponding labels based on the input question. +- Use the provided dataset information to ensure compatibility. +- Validate the selected labels can cover the need to segment the requested targets. + +## Constraints: +1. Response Format +- MUST use the following JSON format for Final Answer: +Final Answer: {FORMATTED_OUTPUT} +Here are some examples: +{FORMATTED_EXAMPLES} + +2. Handling Missing Modalities +- If fail to find any needed modalities, reply in JSON: {MISSING_MODALITY} +- If no suitable dataset is available, reply in JSON: {NO_DATASET_FOUND} + +## Reasoning (CoT): +To ensure accuracy and transparency, provide a step-by-step reasoning process for selecting datasets, labels, and handling constraints. This reasoning should reflect logical steps leading to the final JSON response. + +Here's the dataset information: +{DATASET_JSON_INFO} + +Let’s Begin! +""" + + + +def chat_with_llm_cot(query, model, autogen_config_path, sys_prompt_dict=SYSTEM_PROMPT_COT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + general_chat_history = [] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + expert_agent = ConversableAgent( + name=f"single_expert", + system_message=SYSTEM_PROMPT_COT, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + expert_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + + final_llm_answer = extract_llm_answer(expert_result.summary) + + return { + "reply": final_llm_answer, + "chat_history": general_chat_history, + } diff --git a/agents/doctor_expert.py b/agents/doctor_expert.py new file mode 100644 index 0000000..34106cd --- /dev/null +++ b/agents/doctor_expert.py @@ -0,0 +1,48 @@ +from autogen import ConversableAgent, config_list_from_json +from prompts.basic_v2 import FORMATTED_OUTPUT, FORMATTED_EXAMPLES +from prompts.doctor_expert import DOCTOR_EXPERT_SYS_PROMPT + + +def chat_with_llm_multi_turn_doctor_expert(query, + model, + autogen_config_path, + system_prompt_dict=DOCTOR_EXPERT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + for agent_name in ["Doctor", "Expert"]: + if agent_name not in system_prompt_dict: + raise ValueError(f"Missing {agent_name} in system_prompt_dict") + + doctor_agent = ConversableAgent( + name="Doctor", + system_message=system_prompt_dict["Doctor"], + llm_config=autogen_config, + ) + expert_agent = ConversableAgent( + name="Expert", + system_message=system_prompt_dict["Expert"] + query, + llm_config=autogen_config, + ) + + SUMMARY_PROMPT = f""" Summarize the final results of dataset and label selctions from the conversation. + Please output using the following JSON format: ``` {FORMATTED_OUTPUT} ``` + Make sure that all suitable datasets and corresponding labels are chosen. + Here are some examples: + {FORMATTED_EXAMPLES} + """ + # expert_agent (recipient) will summarize with the prompt + chat_result = doctor_agent.initiate_chat(expert_agent, + message=query, + max_turns=2, + summary_method="reflection_with_llm", + summary_args={"summary_prompt": SUMMARY_PROMPT}) + + return { + "reply": + chat_result.summary, + "chat_history": + chat_result.chat_history + [{ + "role": "Summarizer", + "content": chat_result.summary + }], + } diff --git a/agents/evo_react_2shot.py b/agents/evo_react_2shot.py new file mode 100644 index 0000000..997a93c --- /dev/null +++ b/agents/evo_react_2shot.py @@ -0,0 +1,52 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from autogen import config_list_from_json, UserProxyAgent +from agents.evo_utils import EvolvableAgent + +import json + + +def factual_exp_filter(json_content): + # json_content.pop("llm_raw_answer", None) + json_content.pop("checklist", None) + return json_content + + +factual_exp_file = "C:/Sync_Code/MedSegAgent/ckpt/ReAct_2shot_train_set_100_gpt-4o-2024-08-06_factual_exp_suggested.jsonl" +factual_exp_info = [factual_exp_filter(json.loads(line)) for line in open(factual_exp_file)] + + +# print(factual_exp_info) +def working_exp_filter(json_content): + return json_content.get("checklist", None) + + +working_exp_info = [working_exp_filter(json.loads(line)) for line in open(factual_exp_file)] +working_exp_info = [we for we in working_exp_info if (we)] +# working_exp_info = working_exp_info[:3] +# print(working_exp_info) +# print(len(working_exp_info)) + + +def chat_with_llm_evo_react(query, model, autogen_config_path, sys_prompt_dict=REACT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + expert_agent = EvolvableAgent(name="expert", + system_message=sys_prompt_dict["Expert"], + llm_config=autogen_config, + human_input_mode="NEVER", + max_exp_size=1000) + + for i, fexp in enumerate(factual_exp_info): + expert_agent.update_factual_exp(f"case-{i}", str(fexp)) + for i, wexp in enumerate(working_exp_info): + expert_agent.update_working_exp(f"exp-{i}", str(wexp)) + + expert_agent.update_profile_with_exp(silent=True) + + chat_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + + return { + "reply": chat_result.summary, + "chat_history": chat_result.chat_history, + } diff --git a/agents/evo_react_verifier_2shot.py b/agents/evo_react_verifier_2shot.py new file mode 100644 index 0000000..0db0928 --- /dev/null +++ b/agents/evo_react_verifier_2shot.py @@ -0,0 +1,93 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from prompts.verifier import VERIFIER_SYS_PROMPT +from autogen import config_list_from_json, UserProxyAgent +from agents.evo_utils import EvolvableAgent +from prompts.data import retrieve_dataset_label_info +from utils.reply_utils import extract_llm_answer + +import json + + +def factual_exp_filter(json_content): + json_content.pop("checklist", None) + return json_content + + +factual_exp_file = "C:/Sync_Code/MedSegAgent/ckpt/ReAct_2shot_train_set_100_gpt-4o-2024-08-06_factual_exp_suggested.jsonl" +factual_exp_info = [factual_exp_filter(json.loads(line)) for line in open(factual_exp_file)] + + +def working_exp_filter(json_content): + return json_content.get("checklist", None) + + +working_exp_info = [working_exp_filter(json.loads(line)) for line in open(factual_exp_file)] +working_exp_info = [we for we in working_exp_info if (we)] + +# def chat_with_llm_evo_react(query, model, autogen_config_path, sys_prompt_dict=REACT_SYS_PROMPT): +# autogen_config = config_list_from_json(env_or_file=autogen_config_path, +# filter_dict={"model": [model]})[0] +# user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) +# expert_agent = EvolvableAgent(name="expert", +# system_message=sys_prompt_dict["Expert"], +# llm_config=autogen_config, +# human_input_mode="NEVER", +# max_exp_size=1000) + +# for i, fexp in enumerate(factual_exp_info): +# expert_agent.update_factual_exp(f"case-{i}", str(fexp)) +# expert_agent.update_profile_with_exp(silent=True) + +# chat_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + +# return { +# "reply": chat_result.summary, +# "chat_history": chat_result.chat_history, +# } + + +def chat_with_llm_evo_react_verifier(query, + model, + autogen_config_path, + sys_prompt_dict=REACT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + general_chat_history = [] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + expert_agent = EvolvableAgent(name="expert", + system_message=sys_prompt_dict["Expert"], + llm_config=autogen_config, + human_input_mode="NEVER", + max_exp_size=1000) + verifier_agent = EvolvableAgent(name="verifier", + system_message=VERIFIER_SYS_PROMPT, + llm_config=autogen_config, + human_input_mode="NEVER", + max_exp_size=1000) + for i, fexp in enumerate(factual_exp_info): + expert_agent.update_factual_exp(f"case-{i}", str(fexp)) + expert_agent.update_profile_with_exp(silent=True) + + # user -> expert + expert_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + general_chat_history.extend(expert_result.chat_history) + llm_answer = extract_llm_answer(expert_result.summary) + + # user -> verifier + verf_query = f"LLM Answer:`{llm_answer}`\nReference Information:```\n" + for dataset, labels in llm_answer.items(): + label_info = retrieve_dataset_label_info(dataset, labels) + verf_query += f"{dataset}:{label_info}\n" + verf_query += "```\nPlease verify the answer and return in the constrainted format." + verf_result = user_proxy.initiate_chat(verifier_agent, message=verf_query, max_turns=1) + general_chat_history.extend(verf_result.chat_history) + + # verifier -> expert + final_ans_query = f"Verifier Comments:`{verf_result.summary}`\n Question:`{query}`\nPlease provide your new answer." + ans_result = user_proxy.initiate_chat(expert_agent, message=final_ans_query, max_turns=1) + general_chat_history.extend(ans_result.chat_history) + + return { + "reply": ans_result.summary, + "chat_history": general_chat_history, + } diff --git a/agents/native.py b/agents/native.py new file mode 100644 index 0000000..3adc044 --- /dev/null +++ b/agents/native.py @@ -0,0 +1,31 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from prompts.react_2shot import SYSTEM_PROMPT_REACT +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent +from utils.reply_utils import extract_llm_answer + +from prompts.native_2shot import SYSTEM_PROMPT_NATIVE + + +def chat_with_llm_native(query, model, autogen_config_path, sys_prompt_dict=REACT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + general_chat_history = [] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + print(REACT_SYS_PROMPT) + + expert_agent = ConversableAgent( + name=f"single_expert", + system_message=SYSTEM_PROMPT_NATIVE, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + expert_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + + final_llm_answer = extract_llm_answer(expert_result.summary) + + return { + "reply": final_llm_answer, + "chat_history": general_chat_history, + } diff --git a/agents/native_2shot.py b/agents/native_2shot.py new file mode 100644 index 0000000..854d522 --- /dev/null +++ b/agents/native_2shot.py @@ -0,0 +1,26 @@ +from prompts.native_2shot import NATIVE_SYS_PROMPT +from openai import OpenAI +from autogen import config_list_from_json + + +def chat_with_llm_native_old(query, model, autogen_config_path, sys_prompt_dict=NATIVE_SYS_PROMPT): + config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model, "default"]})[0] + client = OpenAI(api_key=config['api_key'], base_url=config['base_url']) + chat_history = [{"role": "user", "content": query}] + completion = client.chat.completions.create( + model=model, + messages=[{ + "role": "system", + "content": sys_prompt_dict["Expert"] + }, { + "role": "user", + "content": query + }], + ) + assistant_reply = completion.choices[0].message.content + chat_history.append({"role": "assistant", "content": assistant_reply}) + return { + "reply": assistant_reply, + "chat_history": chat_history, + } diff --git a/agents/react.py b/agents/react.py new file mode 100644 index 0000000..057ca17 --- /dev/null +++ b/agents/react.py @@ -0,0 +1,27 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from prompts.react_2shot import SYSTEM_PROMPT_REACT +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent +from utils.reply_utils import extract_llm_answer + + +def chat_with_llm_react(query, model, autogen_config_path, sys_prompt_dict=REACT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + general_chat_history = [] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + expert_agent = ConversableAgent( + name=f"single_expert", + system_message=SYSTEM_PROMPT_REACT, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + expert_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + + final_llm_answer = extract_llm_answer(expert_result.summary) + + return { + "reply": final_llm_answer, + "chat_history": general_chat_history, + } diff --git a/agents/react_2shot.py b/agents/react_2shot.py new file mode 100644 index 0000000..5465f23 --- /dev/null +++ b/agents/react_2shot.py @@ -0,0 +1,26 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from openai import OpenAI +from autogen import config_list_from_json + + +def chat_with_llm_react_old(query, model, autogen_config_path, sys_prompt_dict=REACT_SYS_PROMPT): + config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model, "default"]})[0] + client = OpenAI(api_key=config['api_key'], base_url=config['base_url']) + chat_history = [{"role": "user", "content": query}] + completion = client.chat.completions.create( + model=model, + messages=[{ + "role": "system", + "content": sys_prompt_dict["Expert"] + }, { + "role": "user", + "content": query + }], + ) + assistant_reply = completion.choices[0].message.content + chat_history.append({"role": "assistant", "content": assistant_reply}) + return { + "reply": assistant_reply, + "chat_history": chat_history, + } diff --git a/agents/react_multiexpert.py b/agents/react_multiexpert.py new file mode 100644 index 0000000..8ec34b8 --- /dev/null +++ b/agents/react_multiexpert.py @@ -0,0 +1,79 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from prompts.assigner_experts_v2 import ( + ASSIGNER_PROMPT, + WHOLEBODY_EXPERT_PROMPT_NO_EXPERIENCE, + HEAD_EXPERT_PROMPT_NO_EXPERIENCE, + CARDIAC_EXPERT_PROMPT_NO_EXPERIENCE, + THORACIC_EXPERT_PROMPT_NO_EXPERIENCE, + ABDOMINAL_EXPERT_PROMPT_NO_EXPERIENCE, +) +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent +from utils.reply_utils import extract_llm_answer + + +def chat_with_llm_react_multiexpert(query, + model, + autogen_config_path, + sys_prompt_dict=REACT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + general_chat_history = [] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + # Assigner agent + assigner_agent = ConversableAgent( + name="assigner", + system_message=ASSIGNER_PROMPT, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + # Mapping of expert names to their corresponding agents and no-experience prompts + expert_prompts_no_experience = { + "WHOLEBODY": WHOLEBODY_EXPERT_PROMPT_NO_EXPERIENCE, + "ABDOMINAL": ABDOMINAL_EXPERT_PROMPT_NO_EXPERIENCE, + "HEAD": HEAD_EXPERT_PROMPT_NO_EXPERIENCE, + "CARDIAC": CARDIAC_EXPERT_PROMPT_NO_EXPERIENCE, + "THORACIC": THORACIC_EXPERT_PROMPT_NO_EXPERIENCE, + } + + # Create expert agents dynamically using no-experience prompts + experts = { + name: + ConversableAgent( + name=f"{name.lower()}_expert", + system_message=prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + for name, prompt in expert_prompts_no_experience.items() + } + + # User -> Assigner + assigner_result = user_proxy.initiate_chat(assigner_agent, message=query, max_turns=1) + general_chat_history.extend(assigner_result.chat_history) + expert_names = extract_llm_answer(assigner_result.summary) + + # Assigner -> Expert(s) + expert_answers = {} + for expert_name in expert_names: + expert_agent = experts.get(expert_name) + if expert_agent: + # Expert provides initial answer + expert_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + llm_answer = extract_llm_answer(expert_result.summary) + general_chat_history.extend(expert_result.chat_history) + + expert_answers[expert_name] = llm_answer + else: + print(f"Expert '{expert_name}' not found.") + + # Combine answers + final_llm_answer = {} + for ans in expert_answers.values(): + final_llm_answer.update(ans) + + return { + "reply": final_llm_answer, + "chat_history": general_chat_history, + } diff --git a/agents/react_multiexpert_experience.py b/agents/react_multiexpert_experience.py new file mode 100644 index 0000000..9c96965 --- /dev/null +++ b/agents/react_multiexpert_experience.py @@ -0,0 +1,79 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from prompts.assigner_experts_v2 import ( + ASSIGNER_PROMPT, + WHOLEBODY_EXPERT_PROMPT, + HEAD_EXPERT_PROMPT, + CARDIAC_EXPERT_PROMPT, + THORACIC_EXPERT_PROMPT, + ABDOMINAL_EXPERT_PROMPT, +) +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent +from utils.reply_utils import extract_llm_answer + + +def chat_with_llm_react_multiexpert_experience(query, + model, + autogen_config_path, + sys_prompt_dict=REACT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + general_chat_history = [] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + # Assigner agent + assigner_agent = ConversableAgent( + name="assigner", + system_message=ASSIGNER_PROMPT, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + # Mapping of expert names to their corresponding agents and prompts + expert_prompts = { + "WHOLEBODY": WHOLEBODY_EXPERT_PROMPT, + "ABDOMINAL": ABDOMINAL_EXPERT_PROMPT, + "HEAD": HEAD_EXPERT_PROMPT, + "CARDIAC": CARDIAC_EXPERT_PROMPT, + "THORACIC": THORACIC_EXPERT_PROMPT, + } + + # Create expert agents dynamically + experts = { + name: + ConversableAgent( + name=f"{name.lower()}_expert", + system_message=prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + for name, prompt in expert_prompts.items() + } + + # User -> Assigner + assigner_result = user_proxy.initiate_chat(assigner_agent, message=query, max_turns=1) + general_chat_history.extend(assigner_result.chat_history) + expert_names = extract_llm_answer(assigner_result.summary) + + # Assigner -> Expert(s) + expert_answers = {} + for expert_name in expert_names: + expert_agent = experts.get(expert_name) + if expert_agent: + # Expert provides initial answer + expert_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + llm_answer = extract_llm_answer(expert_result.summary) + general_chat_history.extend(expert_result.chat_history) + + expert_answers[expert_name] = llm_answer + else: + print(f"Expert '{expert_name}' not found.") + + # Combine refined answers + final_llm_answer = {} + for ans in expert_answers.values(): + final_llm_answer.update(ans) + + return { + "reply": final_llm_answer, + "chat_history": general_chat_history, + } diff --git a/agents/react_multiexpert_selfverify.py b/agents/react_multiexpert_selfverify.py new file mode 100644 index 0000000..7548ffb --- /dev/null +++ b/agents/react_multiexpert_selfverify.py @@ -0,0 +1,92 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from prompts.assigner_experts_v2 import ( + ASSIGNER_PROMPT, + WHOLEBODY_EXPERT_PROMPT_NO_EXPERIENCE, + HEAD_EXPERT_PROMPT_NO_EXPERIENCE, + CARDIAC_EXPERT_PROMPT_NO_EXPERIENCE, + THORACIC_EXPERT_PROMPT_NO_EXPERIENCE, + ABDOMINAL_EXPERT_PROMPT_NO_EXPERIENCE, +) +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent +from utils.reply_utils import extract_llm_answer + + +def chat_with_llm_react_multiexpert_selfverify(query, + model, + autogen_config_path, + sys_prompt_dict=REACT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + general_chat_history = [] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + # Assigner agent + assigner_agent = ConversableAgent( + name="assigner", + system_message=ASSIGNER_PROMPT, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + # Mapping of expert names to their corresponding agents and no-experience prompts + expert_prompts_no_experience = { + "WHOLEBODY": WHOLEBODY_EXPERT_PROMPT_NO_EXPERIENCE, + "ABDOMINAL": ABDOMINAL_EXPERT_PROMPT_NO_EXPERIENCE, + "HEAD": HEAD_EXPERT_PROMPT_NO_EXPERIENCE, + "CARDIAC": CARDIAC_EXPERT_PROMPT_NO_EXPERIENCE, + "THORACIC": THORACIC_EXPERT_PROMPT_NO_EXPERIENCE, + } + + # Create expert agents dynamically using no-experience prompts + experts = { + name: + ConversableAgent( + name=f"{name.lower()}_expert", + system_message=prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + for name, prompt in expert_prompts_no_experience.items() + } + + # User -> Assigner + assigner_result = user_proxy.initiate_chat(assigner_agent, message=query, max_turns=1) + general_chat_history.extend(assigner_result.chat_history) + expert_names = extract_llm_answer(assigner_result.summary) + + # Assigner -> Expert(s) -> Self-Verify + expert_answers = {} + for expert_name in expert_names: + expert_agent = experts.get(expert_name) + if expert_agent: + # Expert provides initial answer + expert_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + llm_answer = extract_llm_answer(expert_result.summary) + general_chat_history.extend(expert_result.chat_history) + + # Expert self-verifies their answer + self_verify_query = ( + f"The question was: {query}\n" + f"Your previous answer was: {llm_answer}\n" + "Please review your previous answer. Does it align with your knowledge and expertise?\n" + "If your answer is correct, please confirm it. If there are any errors or omissions, please correct them in your reply." + ) + self_verify_result = user_proxy.initiate_chat(expert_agent, + message=self_verify_query, + max_turns=1) + refined_llm_answer = extract_llm_answer(self_verify_result.summary) + general_chat_history.extend(self_verify_result.chat_history) + + expert_answers[expert_name] = refined_llm_answer + else: + print(f"Expert '{expert_name}' not found.") + + # Combine refined answers + final_llm_answer = {} + for ans in expert_answers.values(): + final_llm_answer.update(ans) + + return { + "reply": final_llm_answer, + "chat_history": general_chat_history, + } diff --git a/agents/react_multiexpert_selfverify_experience.py b/agents/react_multiexpert_selfverify_experience.py new file mode 100644 index 0000000..370579e --- /dev/null +++ b/agents/react_multiexpert_selfverify_experience.py @@ -0,0 +1,92 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from prompts.assigner_experts_v2 import ( + ASSIGNER_PROMPT, + WHOLEBODY_EXPERT_PROMPT, + HEAD_EXPERT_PROMPT, + CARDIAC_EXPERT_PROMPT, + THORACIC_EXPERT_PROMPT, + ABDOMINAL_EXPERT_PROMPT, +) +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent +from utils.reply_utils import extract_llm_answer + + +def chat_with_llm_react_multiexpert_selfverify_experience(query, + model, + autogen_config_path, + sys_prompt_dict=REACT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + general_chat_history = [] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + # Assigner agent + assigner_agent = ConversableAgent( + name="assigner", + system_message=ASSIGNER_PROMPT, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + # Mapping of expert names to their corresponding agents and prompts + expert_prompts = { + "WHOLEBODY": WHOLEBODY_EXPERT_PROMPT, + "ABDOMINAL": ABDOMINAL_EXPERT_PROMPT, + "HEAD": HEAD_EXPERT_PROMPT, + "CARDIAC": CARDIAC_EXPERT_PROMPT, + "THORACIC": THORACIC_EXPERT_PROMPT, + } + + # Create expert agents dynamically + experts = { + name: + ConversableAgent( + name=f"{name.lower()}_expert", + system_message=prompt, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + for name, prompt in expert_prompts.items() + } + + # User -> Assigner + assigner_result = user_proxy.initiate_chat(assigner_agent, message=query, max_turns=1) + general_chat_history.extend(assigner_result.chat_history) + expert_names = extract_llm_answer(assigner_result.summary) + + # Assigner -> Expert(s) -> Self-Verify + expert_answers = {} + for expert_name in expert_names: + expert_agent = experts.get(expert_name) + if expert_agent: + # Expert provides initial answer + expert_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + llm_answer = extract_llm_answer(expert_result.summary) + general_chat_history.extend(expert_result.chat_history) + + # Expert self-verifies their answer + self_verify_query = ( + f"The question was: {query}\n" + f"Your previous answer was: {llm_answer}\n" + "Please review your previous answer. Does it align with your Common Experiences and Expert Experiences? \n" + "If your answer is correct, please confirm it. If there are any errors or omissions, please correct them in your reply." + ) + self_verify_result = user_proxy.initiate_chat(expert_agent, + message=self_verify_query, + max_turns=1) + refined_llm_answer = extract_llm_answer(self_verify_result.summary) + general_chat_history.extend(self_verify_result.chat_history) + + expert_answers[expert_name] = refined_llm_answer + else: + print(f"Expert '{expert_name}' not found.") + + # Combine refined answers + final_llm_answer = {} + for ans in expert_answers.values(): + final_llm_answer.update(ans) + + return { + "reply": final_llm_answer, + "chat_history": general_chat_history, + } diff --git a/agents/react_rag.py b/agents/react_rag.py new file mode 100644 index 0000000..0061726 --- /dev/null +++ b/agents/react_rag.py @@ -0,0 +1,45 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from prompts.react_2shot import SYSTEM_PROMPT_REACT +from prompts.rag import RAG_HELPER_SYS_PROMPT +from prompts.data import retrieve_dataset_from_label_name +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent +from utils.reply_utils import extract_llm_answer + + +def chat_with_llm_react_rag(query, model, autogen_config_path, sys_prompt_dict=REACT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + general_chat_history = [] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + rag_helper = ConversableAgent( + name=f"rag_helper", + system_message=RAG_HELPER_SYS_PROMPT, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + expert_agent = ConversableAgent( + name=f"single_expert", + system_message=SYSTEM_PROMPT_REACT, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + rag_helper_result = user_proxy.initiate_chat( + rag_helper, message=f"Please give one keyword for this query: {query}", max_turns=1) + rag_keyword = rag_helper_result.summary.strip() + retrieved_info = retrieve_dataset_from_label_name(rag_keyword, backend="spacy") + + expert_result = user_proxy.initiate_chat( + expert_agent, + message= + f"Here are some more relevant datasets, please pay special attention to them:{retrieved_info}\n{query}", + max_turns=1) + + final_llm_answer = extract_llm_answer(expert_result.summary) + + return { + "reply": final_llm_answer, + "chat_history": general_chat_history, + } diff --git a/agents/react_selfverify.py b/agents/react_selfverify.py new file mode 100644 index 0000000..9411f16 --- /dev/null +++ b/agents/react_selfverify.py @@ -0,0 +1,42 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from prompts.react_2shot import SYSTEM_PROMPT_REACT +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent +from utils.reply_utils import extract_llm_answer + + +def chat_with_llm_react_selfverify(query, model, autogen_config_path, sys_prompt_dict=REACT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + general_chat_history = [] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + + expert_agent = ConversableAgent( + name=f"single_expert", + system_message=SYSTEM_PROMPT_REACT, + llm_config=autogen_config, + human_input_mode="NEVER", + ) + + # Get initial answer from expert + expert_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + initial_llm_answer = extract_llm_answer(expert_result.summary) + general_chat_history.extend(expert_result.chat_history) + + # Expert self-verifies their answer + self_verify_query = ( + f"The question was: {query}\n" + f"Your previous answer was: {initial_llm_answer}\n" + "Please review your previous answer. Does it align with your knowledge and expertise?\n" + "If your answer is correct, please confirm it. If there are any errors or omissions, please correct them in your reply." + ) + + self_verify_result = user_proxy.initiate_chat(expert_agent, + message=self_verify_query, + max_turns=1) + final_llm_answer = extract_llm_answer(self_verify_result.summary) + general_chat_history.extend(self_verify_result.chat_history) + + return { + "reply": final_llm_answer, + "chat_history": general_chat_history, + } \ No newline at end of file diff --git a/agents/react_verifier_2shot.py b/agents/react_verifier_2shot.py new file mode 100644 index 0000000..a6377e6 --- /dev/null +++ b/agents/react_verifier_2shot.py @@ -0,0 +1,49 @@ +from prompts.react_2shot import REACT_SYS_PROMPT +from prompts.verifier import VERIFIER_SYS_PROMPT +from autogen import config_list_from_json, UserProxyAgent, ConversableAgent +from prompts.data import retrieve_dataset_label_info +from utils.reply_utils import extract_llm_answer + + +def chat_with_llm_react_verifier(query, + model, + autogen_config_path, + sys_prompt_dict=REACT_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + general_chat_history = [] + user_proxy = UserProxyAgent("user_proxy", code_execution_config=False) + expert_agent = ConversableAgent(name="expert", + system_message=sys_prompt_dict["Expert"], + llm_config=autogen_config, + human_input_mode="NEVER") + verifier_agent = ConversableAgent(name="verifier", + system_message=VERIFIER_SYS_PROMPT, + llm_config=autogen_config, + human_input_mode="NEVER") + + # user - + + # user -> expert + expert_result = user_proxy.initiate_chat(expert_agent, message=query, max_turns=1) + general_chat_history.extend(expert_result.chat_history) + llm_answer = extract_llm_answer(expert_result.summary) + + # user -> verifier + verf_query = f"Question: {query}\nLLM Answer:{llm_answer}\nReference Information:\n" + for dataset, labels in llm_answer.items(): + label_info = retrieve_dataset_label_info(dataset, labels) + verf_query += f"{dataset}:{label_info}\n" + verf_query += "Please verify the answer and return in the constrainted format." + verf_result = user_proxy.initiate_chat(verifier_agent, message=verf_query, max_turns=1) + general_chat_history.extend(verf_result.chat_history) + + # verifier -> expert + final_ans_query = f"Verifier Comments:{verf_result.summary}\n Question:{query}\nPlease provide your new answer." + ans_result = user_proxy.initiate_chat(expert_agent, message=final_ans_query, max_turns=1) + general_chat_history.extend(ans_result.chat_history) + + return { + "reply": ans_result.summary, + "chat_history": general_chat_history, + } diff --git a/agents/reflexion_2shot.py b/agents/reflexion_2shot.py new file mode 100644 index 0000000..50b9fbe --- /dev/null +++ b/agents/reflexion_2shot.py @@ -0,0 +1,80 @@ +from autogen import config_list_from_json, ConversableAgent +from prompts.reflexion_2shot import REFLECTXION_SYS_PROMPT +from prompts.react_2shot import REACT_FORMATTED_EXAMPLES +from prompts.data import retrieve_dataset_label_info +from utils.reply_utils import extract_llm_answer + +MAX_TRY = 2 + + +def chat_with_llm_reflexion(query, + model, + autogen_config_path, + sys_prompt_dict=REFLECTXION_SYS_PROMPT): + autogen_config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model]})[0] + expert_agent = ConversableAgent(name="actor", + system_message=sys_prompt_dict["Expert"], + llm_config=autogen_config, + human_input_mode="NEVER") + + eval_agent = ConversableAgent(name="evaluator", + system_message=sys_prompt_dict["Evaluator"], + llm_config=autogen_config, + human_input_mode="NEVER") + + reflection_agent = ConversableAgent(name="self_reflection", + system_message=sys_prompt_dict["Reflection"], + llm_config=autogen_config, + human_input_mode="NEVER") + + general_chat_history = [] + curr_query = query + for i in range(MAX_TRY): + # ask for first result + actor_result = eval_agent.initiate_chat(expert_agent, message=curr_query, max_turns=1) + general_chat_history += actor_result.chat_history + actor_final_answer = actor_result.summary + if ("Final Answer:" in actor_final_answer): + actor_final_answer = actor_final_answer.split("Final Answer:")[-1].strip() + + # get eval + llm_answer = extract_llm_answer(actor_final_answer) + eval_query = f"Previous Answer:{actor_final_answer}\nVerified Labels:\n" + for dataset, labels in llm_answer.items(): + eval_query += f"dataset:{dataset}," + "labels:{" + retrieve_dataset_label_info( + dataset, labels) + "}\n" + eval_query += "\nPlease evaluate." + + eval_result = expert_agent.initiate_chat(eval_agent, message=eval_query, max_turns=1) + general_chat_history += eval_result.chat_history[1:] + eval_final_result = eval_result.summary + if ("Final Result: True" in eval_final_result): + eval_final_result = "True" + + # eval think we get the right answer or upon the max try time + if (eval_final_result == "True" or i == MAX_TRY - 1): + return { + "reply": actor_result.summary, + "chat_history": general_chat_history, + } + + # if not find the right answer + # self-reflection for suggestion + reflect_result = eval_agent.initiate_chat(reflection_agent, + message=f"result:{actor_result.summary}\n" + f"evaluation:{eval_result.summary}\n" + "Please give the suggestion.", + max_turns=1) + general_chat_history += reflect_result.chat_history + + # generate new query + suggestion = "No suggestion." + if ("Final Suggestion:" in reflect_result.summary): + suggestion = reflect_result.summary.split("Final Suggestion:")[-1].strip() + curr_query = f"Task:{query}\nWrong Answer:{actor_final_answer}\nSuggestion for this task:'{suggestion}'\nHere's some examples: {REACT_FORMATTED_EXAMPLES}" + + return { + "reply": actor_result.summary, + "chat_history": general_chat_history, + } diff --git a/agents/single_turn.py b/agents/single_turn.py new file mode 100644 index 0000000..a68ec26 --- /dev/null +++ b/agents/single_turn.py @@ -0,0 +1,29 @@ +from prompts.basic_v2 import BASIC_V2_SYS_PROMPT +from openai import OpenAI +from autogen import config_list_from_json + + +def chat_with_llm_single_turn(query, + model, + autogen_config_path, + sys_prompt_dict=BASIC_V2_SYS_PROMPT): + config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model, "default"]})[0] + client = OpenAI(api_key=config['api_key'], base_url=config['base_url']) + chat_history = [{"role": "user", "content": query}] + completion = client.chat.completions.create( + model=model, + messages=[{ + "role": "system", + "content": sys_prompt_dict["Expert"] + }, { + "role": "user", + "content": query + }], + ) + assistant_reply = completion.choices[0].message.content + chat_history.append({"role": "assistant", "content": assistant_reply}) + return { + "reply": assistant_reply, + "chat_history": chat_history, + } diff --git a/assets/method.png b/assets/method.png new file mode 100644 index 0000000..4ff1688 --- /dev/null +++ b/assets/method.png Binary files differ diff --git a/assets/motivation.png b/assets/motivation.png new file mode 100644 index 0000000..d9e20d2 --- /dev/null +++ b/assets/motivation.png Binary files differ diff --git a/assets/result1.png b/assets/result1.png new file mode 100644 index 0000000..0e0a14a --- /dev/null +++ b/assets/result1.png Binary files differ diff --git a/dataset/AMOS22_Task2.json b/dataset/AMOS22_Task2.json new file mode 100644 index 0000000..c504e58 --- /dev/null +++ b/dataset/AMOS22_Task2.json @@ -0,0 +1,45 @@ +{ + "name": "AMOS22_Task2", + "description": "MICCAI2022 Multi-Modality Abdominal Multi-Organ Segmentation Task 2", + "input_modalities": { + "0": "MR or CT" + }, + "labels": { + "background": 0, + "spleen": 1, + "kidney_right": 2, + "kidney_left": 3, + "gallbladder": 4, + "esophagus": 5, + "liver": 6, + "stomach": 7, + "aorta": 8, + "inferior_vena_cava": 9, + "pancreas": 10, + "adrenal_gland_right": 11, + "adrenal_gland_left": 12, + "duodenum": 13, + "bladder": 14, + "prostate_or_uterus": 15 + }, + "numTraining": 240, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/ATM22.json b/dataset/ATM22.json new file mode 100644 index 0000000..26c3656 --- /dev/null +++ b/dataset/ATM22.json @@ -0,0 +1,31 @@ +{ + "name": "ATM22", + "description": "Airway segmentation is a crucial step for the analysis of pulmonary diseases including asthma, bronchiectasis, and emphysema. The accurate segmentation based on X-Ray computed tomography (CT) enables the quantitative measurements of airway dimensions and wall thickness, which can reveal the abnormality of patients with chronic obstructive pulmonary disease (COPD). Besides, the extraction of patient-specific airway models from CT images is required for navigation in bronchoscopic-assisted surgery.", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "pulmonary airway": 1 + }, + "numTraining": 299, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/AbdomenAtlasMini.json b/dataset/AbdomenAtlasMini.json new file mode 100644 index 0000000..cdf1e97 --- /dev/null +++ b/dataset/AbdomenAtlasMini.json @@ -0,0 +1,39 @@ +{ + "name": "AbdomenAtlasMini", + "description": "The largest, fully-annotated CT dataset to date, including 5,195 annotated CT volumes (with spleen, liver, kidneys, stomach, gallbladder, pancreas, aorta, and IVC annotations).", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "aorta": 1, + "gallbladder": 2, + "kidney_left": 3, + "kidney_right": 4, + "liver": 5, + "pancreas": 6, + "postcava": 7, + "spleen": 8, + "stomach": 9 + }, + "numTraining": 5170, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} diff --git a/dataset/Adrenal-ACC-Ki67-Seg.json b/dataset/Adrenal-ACC-Ki67-Seg.json new file mode 100644 index 0000000..bf75d73 --- /dev/null +++ b/dataset/Adrenal-ACC-Ki67-Seg.json @@ -0,0 +1,31 @@ +{ + "name": "Adrenal-ACC-Ki67-Seg", + "description": "This retrospectively acquired data includes contrast enhanced CT imaging studies of 53 confirmed ACC patients between 2006 to 2018 with available clinical and pathological data", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "Adrenocortical carcinoma": 1 + }, + "numTraining": 57, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/AutoPET.json b/dataset/AutoPET.json new file mode 100644 index 0000000..10c8bc6 --- /dev/null +++ b/dataset/AutoPET.json @@ -0,0 +1,32 @@ +{ + "name": "AutoPET", + "description": "The challenge cohort consists of patients with histologically proven malignant melanoma, lymphoma or lung cancer as well as negative control patients who were examined by FDG-PET/CT in two large medical centers (University Hospital Tübingen, Germany & University Hospital of the LMU in Munich, Germany). All PET/CT data within this challenge have been acquired on state-of-the-art PET/CT scanners (Siemens Biograph mCT, mCT Flow and Biograph 64, GE Discovery 690) using standardized protocols following international guidelines. CT as well as PET data are provided as 3D volumes consisting of stacks of axial slices. Data provided as part of this challenge consists of whole-body examinations. Usually, the scan range of these examinations extends from the skull base to the mid-thigh level. If clinically relevant, scans can be extended to cover the entire body including the entire head and legs/feet.", + "input_modalities": { + "0": "PET", + "1": "CT" + }, + "labels": { + "background": 0, + "tumor": 1 + }, + "numTraining": 1014, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/BraTS21.json b/dataset/BraTS21.json new file mode 100644 index 0000000..4fc7bb0 --- /dev/null +++ b/dataset/BraTS21.json @@ -0,0 +1,36 @@ +{ + "name": "BraTS21", + "description": "Gliomas segmentation tumour in on brain images", + "input_modalities": { + "0": "MR T2 FLAIR", + "1": "MR T1 Weighted", + "2": "MR T1 weighted imaging with gadolinium", + "3": "MR T2 Weighted" + }, + "labels": { + "background": 0, + "whole_tumor": [1, 2, 3], + "tumor_core": [1, 3], + "enhancing_tumor": 3 + }, + "numTraining": 1251, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/CT-ORG.json b/dataset/CT-ORG.json new file mode 100644 index 0000000..b013f60 --- /dev/null +++ b/dataset/CT-ORG.json @@ -0,0 +1,36 @@ +{ + "name": "CT-ORG", + "description": "This dataset consists of 140 computed tomography (CT) scans, each with five organs labeled in 3D: lung, bones, liver, kidneys and bladder. The brain is also labeled on the minority of scans which show it.Patients were included based on the presence of lesions in one or more of the labeled organs. Most of the images exhibit liver lesions, both benign and malignant. Some also exhibit metastatic disease in other organs such as bones and lungs. The images come from a wide variety of sources, including abdominal and full-body; contrast and non-contrast; low-dose and high-dose CT scans. 131 images are dedicated CTs, the remaining 9 are the CT component taken from PET-CT exams. This makes the dataset ideal for training and evaluating organ segmentation algorithms, which ought to perform well in a wide variety of imaging conditions.", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "Liver": 1, + "Bladder": 2, + "Lungs": 3, + "Kidneys": 4, + "Bone": 5, + "Brain": 6 + }, + "numTraining": 140, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/CrossMoDA2023.json b/dataset/CrossMoDA2023.json new file mode 100644 index 0000000..61657b9 --- /dev/null +++ b/dataset/CrossMoDA2023.json @@ -0,0 +1,15 @@ +{ + "name": "CrossMoDA2023", + "description": "", + "input_modalities": { + "0": "ceT1 and T2" + }, + "labels": { + "background": 0, + "intra-meatal tumor": 1, + "extra-meatal tumor": 2, + "cochlea": 3 + }, + "numTraining": 227, + "file_ending": ".nii.gz" +} \ No newline at end of file diff --git a/dataset/FLARE22.json b/dataset/FLARE22.json new file mode 100644 index 0000000..fc1e74e --- /dev/null +++ b/dataset/FLARE22.json @@ -0,0 +1,43 @@ +{ + "name": "FLARE22", + "description": "Abdominal multi-organs segmentation datasets.", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "liver": 1, + "right kidney": 2, + "spleen": 3, + "pancreas": 4, + "aorta": 5, + "inferior vena cava": 6, + "right adrenal gland": 7, + "left adrenal gland": 8, + "gallbladder": 9, + "esophagus": 10, + "stomach": 11, + "duodenum": 12, + "left kidney": 13 + }, + "numTraining": 50, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/FUMPE.json b/dataset/FUMPE.json new file mode 100644 index 0000000..d499731 --- /dev/null +++ b/dataset/FUMPE.json @@ -0,0 +1,13 @@ +{ + "name": "FUMPE", + "description": "The FUMPE (Ferdowsi University of Mashhad's Pulmonary Embolism) dataset is a dataset for segmenting pulmonary embolism (PE) regions in computed tomography angiography (CTA) images.", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "Pulmonary Embolism": 1 + }, + "numTraining": 35, + "file_ending": ".nii.gz" +} \ No newline at end of file diff --git a/dataset/HECKTOR2022.json b/dataset/HECKTOR2022.json new file mode 100644 index 0000000..f0af38a --- /dev/null +++ b/dataset/HECKTOR2022.json @@ -0,0 +1,32 @@ +{ + "name": "HECKTOR2022", + "description": "The automatic segmentation of Head and Neck (H&N) primary tumors and lymph nodes (new!) in FDG-PET/CT images, focusing on oropharyngeal cancers", + "input_modalities": { + "0": "PET", + "1": "CT" + }, + "labels": { + "GTVp": 1, + "GTVnd": 2 + }, + "numTraining": 524, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} diff --git a/dataset/ISLES22.json b/dataset/ISLES22.json new file mode 100644 index 0000000..2e10836 --- /dev/null +++ b/dataset/ISLES22.json @@ -0,0 +1,34 @@ +{ + "name": "ISLES22", + "website": "https://www.isles-challenge.org/", + "description": "Multimodal MRI infarct segmentation in acute and sub-acute stroke. Participants are requested to Segment brain infarct lesions from acute and sub-acute stroke scans using DWI, ADC and FLAIR images.", + "input_modalities": { + "0": "MR adc", + "1": "MR dwi", + "2": "MR flair" + }, + "labels": { + "background": 0, + "stroke lesion": 1 + }, + "numTraining": 250, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/ISLES22_ATLAS.json b/dataset/ISLES22_ATLAS.json new file mode 100644 index 0000000..a28b678 --- /dev/null +++ b/dataset/ISLES22_ATLAS.json @@ -0,0 +1,32 @@ +{ + "name": "ISLES22_ATLAS", + "website": "https://atlas.grand-challenge.org/", + "description": "Single channel T1-weighted lesion segmentation in acute, sub-acute and chronic stroke (ATLAS challenge). Participants are requested to Segment brain lesions from acute, sub-acute and chronic stroke in single channel weighted T1 images.", + "input_modalities": { + "0": "MR T1 weighted" + }, + "labels": { + "background": 0, + "stroke lesion": 1 + }, + "numTraining": 655, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/ImageCAS.json b/dataset/ImageCAS.json new file mode 100644 index 0000000..76724ec --- /dev/null +++ b/dataset/ImageCAS.json @@ -0,0 +1,31 @@ +{ + "name": "ImageCAS", + "description": "A Large Scale Dataset and Benchmark for Coronary Artery Segmentation based on CT", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "Coronary Artery": 1 + }, + "numTraining": 1000, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/Instance22.json b/dataset/Instance22.json new file mode 100644 index 0000000..e0d891a --- /dev/null +++ b/dataset/Instance22.json @@ -0,0 +1,32 @@ +{ + "name": "Instance22", + "website": "https://instance.grand-challenge.org/", + "description": "non-contrast head CT volumes of clinically diagnosed patients with different kinds of ICH, including subdural hemorrhage, epidural hemorrhage, intraventricular hemorrhage, intraparenchymal hemorrhage, and subarachnoid hemorrhage.", + "input_modalities": { + "0": "Non-Contrast head CT" + }, + "labels": { + "background": 0, + "Intracranial Hemorrhage": 1 + }, + "numTraining": 100, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/KiTS23.json b/dataset/KiTS23.json new file mode 100644 index 0000000..59770c8 --- /dev/null +++ b/dataset/KiTS23.json @@ -0,0 +1,33 @@ +{ + "name": "KiTS23", + "description": "The 2023 Kidney and Kidney Tumor Segmentation challenge (abbreviated KiTS23) is a competition in which teams compete to develop the best system for automatic semantic segmentation of kidneys, renal tumors, and renal cysts.", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "kidney and masses": [1,2,3], + "kidney masses": [2,3], + "kidney tumor": [3] + }, + "numTraining": 489, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/LNQ2023.json b/dataset/LNQ2023.json new file mode 100644 index 0000000..af03485 --- /dev/null +++ b/dataset/LNQ2023.json @@ -0,0 +1,13 @@ +{ + "name": "LNQ2023", + "description": "Lymph Node Quantification (LNQ2023) is a CT dataset for automated systems to segment mediastinal lymph nodes.", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "mediastinal lymph nodes": 1 + }, + "numTraining": 393, + "file_ending": ".nii.gz" +} \ No newline at end of file diff --git a/dataset/LiTS.json b/dataset/LiTS.json new file mode 100644 index 0000000..6c56c11 --- /dev/null +++ b/dataset/LiTS.json @@ -0,0 +1,32 @@ +{ + "name": "LiTS", + "description": "Liver and liver tumor segmentation", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "liver": 1, + "liver_tumour": 2 + }, + "numTraining": 131, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/MM-WHS.json b/dataset/MM-WHS.json new file mode 100644 index 0000000..ab12ab4 --- /dev/null +++ b/dataset/MM-WHS.json @@ -0,0 +1,37 @@ +{ + "name": "MM-WHS", + "description": "MM-WHS: Multi-Modality Whole Heart Segmentation. The purpose of disseminating the Data is to perform a multi-institutional analysis of a database of anonymized clinical MRI and CT scans for whole heart segmentation.", + "input_modalities": { + "0": "MR or CT" + }, + "labels": { + "background": 0, + "heart_left_ventricular_myocardium": 1, + "heart_left_atrium_blood_cavity": 2, + "heart_left_ventricle_blood_cavity": 3, + "heart_right_atrium_blood_cavity": 4, + "heart_right_ventricle_blood_cavity": 5, + "heart_ascending_aorta": 6, + "pulmonary_artery": 7 + }, + "numTraining": 40, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/MMs.json b/dataset/MMs.json new file mode 100644 index 0000000..7dc6327 --- /dev/null +++ b/dataset/MMs.json @@ -0,0 +1,33 @@ +{ + "name": "MMs", + "description": "Multi-Centre, Multi-Vendor & Multi-Disease Cardiac Image Segmentation Challenge", + "input_modalities": { + "0": "MR" + }, + "labels": { + "background": 0, + "Left Ventricle": 1, + "Right Ventricle": 2, + "Myocardium": 3 + }, + "numTraining": 150, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/Parse22.json b/dataset/Parse22.json new file mode 100644 index 0000000..0f9e980 --- /dev/null +++ b/dataset/Parse22.json @@ -0,0 +1,31 @@ +{ + "name": "Parse22", + "description": "3D volumes with refined pulmonary artery label, these Contrast Enhanced CT Pulmonary Angiography (CTPA) data are obtained from dual-source 64-slice CT scanner", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "pulmonary artery": 1 + }, + "numTraining": 100, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} \ No newline at end of file diff --git a/dataset/README.md b/dataset/README.md new file mode 100644 index 0000000..e20591c --- /dev/null +++ b/dataset/README.md @@ -0,0 +1,101 @@ +# 数据集 JSON 模板说明 + +在添加一个新的数据集时，需要按照以下 JSON 模板进行准备。请确保每个字段都正确填写，并遵循模板中的格式。 + +```json +{ + "name": "dataset name", + "description": "", + "input_modalities": { + "0": "modality1", + "1": "modality2" + }, + "labels": { + "background": 0, + "class1": 1, + "class2": 2 + }, + "numTraining": 100, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} +``` +## 字段说明 + +- **name**: 数据集的名称。 + - 类型: 字符串 + - 示例: `"dataset name"` + +- **description**: 数据集的描述。 + - 类型: 字符串 + - 示例: `""`（可为空） + +- **input_modalities**: 输入模态的列表，每个模态对应一个键值对。 + - 类型: 对象 + - 示例: + ```json + { + "0": "modality1", + "1": "modality2" + } + ``` + +- **labels**: 标签列表，每个标签对应一个键值对，键为标签名称，值为标签的数字表示。 + - 类型: 对象 + - 示例: + ```json + { + "background": 0, + "class1": 1, + "class2": 2 + } + ``` + +- **numTraining**: 训练样本的数量。 + - 类型: 整数 + - 示例: `100` + +- **file_ending**: 数据文件的后缀名。 + - 类型: 字符串 + - 示例: `".nii.gz"` + +- **models**: 模型列表，每个模型框架包含多个模型名称。 + - 类型: 对象 + - 示例: + ```json + { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } + ``` diff --git a/dataset/SPPIN.json b/dataset/SPPIN.json new file mode 100644 index 0000000..66dfabd --- /dev/null +++ b/dataset/SPPIN.json @@ -0,0 +1,13 @@ +{ + "name": "SPPIN", + "description": "The SPPIN dataset includes T2-weighted MRI, contrast-enhanced T1-weighted MRI, and DWI (b0 and b100) modalities, targeting the segmentation of pediatric neuroblastoma tumors for surgical planning.", + "input_modalities": { + "0": "MR" + }, + "labels": { + "background": 0, + "Pediatric Neuroblastoma": 1 + }, + "numTraining": 767, + "file_ending": ".nii.gz" +} \ No newline at end of file diff --git a/dataset/SegRap2023_Task1.json b/dataset/SegRap2023_Task1.json new file mode 100644 index 0000000..7217e35 --- /dev/null +++ b/dataset/SegRap2023_Task1.json @@ -0,0 +1,76 @@ +{ + "name": "SegRap2023_Task1", + "description": "Segmentation of Organs-at-Risk and Gross Tumor Volume of NPC for Radiotherapy Planning (SegRap2023) both contrast and non-contrast CT", + "input_modalities": { + "0": "non_contrast_CT", + "1": "contrast_CT" + }, + "labels": { + "background": 0, + "Brain": 1, + "BrainStem": 2, + "Chiasm": 3, + "TemporalLobe_L": [4, 6], + "TemporalLobe_R": [5, 7], + "Hippocampus_L": [6, 8], + "Hippocampus_R": [7, 9], + "Eye_L": 10, + "Eye_R": 11, + "Lens_L": 12, + "Lens_R": 13, + "OpticNerve_L": 14, + "OpticNerve_R": 15, + "MiddleEar_L": [16, 20, 24, 30], + "MiddleEar_R": [17, 21, 25, 31], + "IAC_L": 18, + "IAC_R": 19, + "TympanicCavity_L": [20, 22], + "TympanicCavity_R": [21, 23], + "VestibulSemi_L": [24, 26], + "VestibulSemi_R": [25, 27], + "Cochlea_L": 28, + "Cochlea_R": 29, + "ETbone_L": [30, 32], + "ETbone_R": [31, 33], + "Pituitary": 34, + "OralCavity": 35, + "Mandible_L": 36, + "Mandible_R": 37, + "Submandibular_L": 38, + "Submandibular_R": 39, + "Parotid_L": 40, + "Parotid_R": 41, + "Mastoid_L": 42, + "Mastoid_R": 43, + "TMjoint_L": 44, + "TMjoint_R": 45, + "SpinalCord": 46, + "Esophagus": 47, + "Larynx": [48, 51], + "Larynx_Glottic": 49, + "Larynx_Supraglot": 50, + "PharynxConst": [51, 52], + "Thyroid": 53, + "Trachea": 54 + }, + "numTraining": 120, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} diff --git a/dataset/SegRap2023_Task2.json b/dataset/SegRap2023_Task2.json new file mode 100644 index 0000000..36ad33a --- /dev/null +++ b/dataset/SegRap2023_Task2.json @@ -0,0 +1,32 @@ +{ + "name": "SegRap2023_Task2", + "description": "Segmentation of Organs-at-Risk and Gross Tumor Volume of NPC for Radiotherapy Planning (SegRap2023)", + "input_modalities": { + "0": "non_contrast_CT", + "1": "contrast_CT" + }, + "labels": { + "GTVp": 1, + "GTVnd": 2 + }, + "numTraining": 120, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} diff --git a/dataset/TDSC-ABUS2023.json b/dataset/TDSC-ABUS2023.json new file mode 100644 index 0000000..b4c656f --- /dev/null +++ b/dataset/TDSC-ABUS2023.json @@ -0,0 +1,13 @@ +{ + "name": "TDSC-ABUS2023", + "description": "tumor segmentation in automated 3D breast ultrasound (ABUS) images.", + "input_modalities": { + "0": "US" + }, + "labels": { + "background": 0, + "breast tumor": 1 + }, + "numTraining": 100, + "file_ending": ".nii.gz" +} \ No newline at end of file diff --git a/dataset/TotalSegmentator_MRI.json b/dataset/TotalSegmentator_MRI.json new file mode 100644 index 0000000..015c17f --- /dev/null +++ b/dataset/TotalSegmentator_MRI.json @@ -0,0 +1,73 @@ +{ + "name": "TotalSegmentator_MRI", + "description": "Segmentation of TotalSegmentator MRI classes", + "reference": "", + "licence": "Apache 2.0", + "release": "2.0", + "input_modalities": { + "0": "MR" + }, + "labels": { + "background": 0, + "spleen": 1, + "kidney_right": 2, + "kidney_left": 3, + "gallbladder": 4, + "liver": 5, + "stomach": 6, + "pancreas": 7, + "suprarenal_gland_right": 8, + "suprarenal_gland_left": 9, + "lung_left": 10, + "lung_right": 11, + "esophagus": 12, + "small_intestine": 13, + "duodenum": 14, + "colon": 15, + "urinary_bladder": 16, + "prostate": 17, + "sacrum": 18, + "vertebrae": 19, + "intervertebral_discs": 20, + "spinal_cord": 21, + "heart": 22, + "aorta": 23, + "inferior_vena_cava": 24, + "hepatic_portal_vein": 25, + "common_iliac_artery_left": 26, + "common_iliac_artery_right": 27, + "common_iliac_vein_left": 28, + "common_iliac_vein_right": 29, + "humerus_left": 30, + "humerus_right": 31, + "fibula": 32, + "tibia": 33, + "femur_left": 34, + "femur_right": 35, + "hip_left": 36, + "hip_right": 37, + "gluteus_maximus_muscle_left": 38, + "gluteus_maximus_muscle_right": 39, + "gluteus_medius_muscle_left": 40, + "gluteus_medius_muscle_right": 41, + "gluteus_minimus_muscle_left": 42, + "gluteus_minimus_muscle_right": 43, + "autochthon_left": 44, + "autochthon_right": 45, + "iliopsoas_muscle_left": 46, + "iliopsoas_muscle_right": 47, + "quadriceps_femoris_left": 48, + "quadriceps_femoris_right": 49, + "thigh_medial_compartment_left": 50, + "thigh_medial_compartment_right": 51, + "thigh_posterior_compartment_left": 52, + "thigh_posterior_compartment_right": 53, + "sartorius_left": 54, + "sartorius_right": 55, + "brain": 56 + }, + + "numTraining": 298, + "file_ending": ".nii.gz", + "overwrite_image_reader_writer": "NibabelIOWithReorient" +} \ No newline at end of file diff --git a/dataset/TotalSegmentator_v2.json b/dataset/TotalSegmentator_v2.json new file mode 100644 index 0000000..b8e8e67 --- /dev/null +++ b/dataset/TotalSegmentator_v2.json @@ -0,0 +1,147 @@ +{ + "name": "TotalSegmentator_v2", + "description": "Tool for segmentation of over 117 classes in CT images. It was trained on a wide range of different CT images (different scanners, institutions, protocols,...) and therefore should work well on most images.", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "spleen": 1, + "kidney_right": 2, + "kidney_left": 3, + "gallbladder": 4, + "liver": 5, + "stomach": 6, + "pancreas": 7, + "adrenal_gland_right": 8, + "adrenal_gland_left": 9, + "lung_upper_lobe_left": 10, + "lung_lower_lobe_left": 11, + "lung_upper_lobe_right": 12, + "lung_middle_lobe_right": 13, + "lung_lower_lobe_right": 14, + "esophagus": 15, + "trachea": 16, + "thyroid_gland": 17, + "small_bowel": 18, + "duodenum": 19, + "colon": 20, + "urinary_bladder": 21, + "prostate": 22, + "kidney_cyst_left": 23, + "kidney_cyst_right": 24, + "sacrum": 25, + "vertebrae_S1": 26, + "vertebrae_L5": 27, + "vertebrae_L4": 28, + "vertebrae_L3": 29, + "vertebrae_L2": 30, + "vertebrae_L1": 31, + "vertebrae_T12": 32, + "vertebrae_T11": 33, + "vertebrae_T10": 34, + "vertebrae_T9": 35, + "vertebrae_T8": 36, + "vertebrae_T7": 37, + "vertebrae_T6": 38, + "vertebrae_T5": 39, + "vertebrae_T4": 40, + "vertebrae_T3": 41, + "vertebrae_T2": 42, + "vertebrae_T1": 43, + "vertebrae_C7": 44, + "vertebrae_C6": 45, + "vertebrae_C5": 46, + "vertebrae_C4": 47, + "vertebrae_C3": 48, + "vertebrae_C2": 49, + "vertebrae_C1": 50, + "heart": 51, + "aorta": 52, + "pulmonary_vein": 53, + "brachiocephalic_trunk": 54, + "subclavian_artery_right": 55, + "subclavian_artery_left": 56, + "common_carotid_artery_right": 57, + "common_carotid_artery_left": 58, + "brachiocephalic_vein_left": 59, + "brachiocephalic_vein_right": 60, + "atrial_appendage_left": 61, + "superior_vena_cava": 62, + "inferior_vena_cava": 63, + "portal_vein_and_splenic_vein": 64, + "iliac_artery_left": 65, + "iliac_artery_right": 66, + "iliac_vena_left": 67, + "iliac_vena_right": 68, + "humerus_left": 69, + "humerus_right": 70, + "scapula_left": 71, + "scapula_right": 72, + "clavicula_left": 73, + "clavicula_right": 74, + "femur_left": 75, + "femur_right": 76, + "hip_left": 77, + "hip_right": 78, + "spinal_cord": 79, + "gluteus_maximus_left": 80, + "gluteus_maximus_right": 81, + "gluteus_medius_left": 82, + "gluteus_medius_right": 83, + "gluteus_minimus_left": 84, + "gluteus_minimus_right": 85, + "autochthon_left": 86, + "autochthon_right": 87, + "iliopsoas_left": 88, + "iliopsoas_right": 89, + "brain": 90, + "skull": 91, + "rib_right_4": 92, + "rib_right_3": 93, + "rib_left_1": 94, + "rib_left_2": 95, + "rib_left_3": 96, + "rib_left_4": 97, + "rib_left_5": 98, + "rib_left_6": 99, + "rib_left_7": 100, + "rib_left_8": 101, + "rib_left_9": 102, + "rib_left_10": 103, + "rib_left_11": 104, + "rib_left_12": 105, + "rib_right_1": 106, + "rib_right_2": 107, + "rib_right_5": 108, + "rib_right_6": 109, + "rib_right_7": 110, + "rib_right_8": 111, + "rib_right_9": 112, + "rib_right_10": 113, + "rib_right_11": 114, + "rib_right_12": 115, + "sternum": 116, + "costal_cartilages": 117 + }, + "numTraining": 1139, + "file_ending": ".nii.gz", + "models": { + "nnV1": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "nnV2": [ + "nnUNet", + "STU-Net-S", + "STU-Net-B", + "STU-Net-L" + ], + "monai": [ + "SwinUNETR", + "UNETR" + ] + } +} diff --git a/dataset/WORD.json b/dataset/WORD.json new file mode 100644 index 0000000..099b2fc --- /dev/null +++ b/dataset/WORD.json @@ -0,0 +1,28 @@ +{ + "name": "WORD", + "description": "Segmentation of 16 Abdominal Organs", + "input_modalities": { + "0": "CT" + }, + "labels": { + "background": 0, + "liver": 1, + "spleen": 2, + "kidney left": 3, + "kidney right": 4, + "stomach": 5, + "gallbladder": 6, + "esophagus": 7, + "pancreas": 8, + "duodenum": 9, + "colon": 10, + "intestine": 11, + "adrenal": 12, + "rectum": 13, + "bladder": 14, + "femur left": 15, + "femur right": 16 + }, + "numTraining": 100, + "file_ending": ".nii.gz" +} \ No newline at end of file diff --git a/eval_diff_llms.sh b/eval_diff_llms.sh new file mode 100644 index 0000000..d5425e2 --- /dev/null +++ b/eval_diff_llms.sh @@ -0,0 +1,4 @@ +python ./evaluate.py --test_file_path "model_selection_bench.jsonl" --test_pattern "ReAct" --model "Qwen/Qwen2.5-7B-Instruct" --log_to_file +python ./evaluate.py --test_file_path "model_selection_bench.jsonl" --test_pattern "C2F" --model "Qwen/Qwen2.5-7B-Instruct" --log_to_file +python ./evaluate.py --test_file_path "model_selection_bench.jsonl" --test_pattern "ReAct" --model "meta-llama/Meta-Llama-3.1-8B-Instruct" --log_to_file +python ./evaluate.py --test_file_path "model_selection_bench.jsonl" --test_pattern "C2F" --model "meta-llama/Meta-Llama-3.1-8B-Instruct" --log_to_file \ No newline at end of file diff --git a/eval_example.sh b/eval_example.sh new file mode 100644 index 0000000..1e625e3 --- /dev/null +++ b/eval_example.sh @@ -0,0 +1,5 @@ +python ./evaluate.py \ + --test_file_path "model_selection_test_case.jsonl" \ + --test_pattern "C2F" \ + --model "gpt-4o-2024-08-06" \ + --log_to_file diff --git a/evaluate.py b/evaluate.py new file mode 100644 index 0000000..5ac201f --- /dev/null +++ b/evaluate.py @@ -0,0 +1,292 @@ +import argparse +import json +import logging +import colorlog +import os +import os.path as osp +from datetime import datetime + +from utils.chat_utils import chat_with_llm +from utils.metric import dataset_hit_rate, all_dataset_found_rate, any_dataset_found_rate +from utils.reply_utils import * + +import time + + +def score_the_answer(gt_answer, llm_answer): + return { + "hit_rate": dataset_hit_rate(gt_answer, llm_answer), + "strict_hit_rate": all_dataset_found_rate(gt_answer, llm_answer), + "soft_hit_rate": any_dataset_found_rate(gt_answer, llm_answer), + } + + +def run_tests_from_jsonl(file_path, + llm_func, + result_file, + resume=False, + state_filename='task_state.json'): + final_score, total_score = 0, 0 + final_strict_score, final_soft_score = 0, 0 + start_line_id = 0 + evaluation_results = dict() + + # try to resume state + if resume: + logging.info("[Results] resume from state") + state = load_state(state_filename=state_filename) + if state: + start_line_id = state['current_line'] + final_score = state['final_score'] + final_strict_score = state['final_strict_score'] + final_soft_score = state['final_soft_score'] + total_score = state['total_score'] + evaluation_results = state['evaluation_results'] + + with open(file_path, 'r') as file: + for line_id, line in enumerate(file): + line_id += 1 + if line_id <= start_line_id: + continue + + test_case = json.loads(line) + query = test_case['query'] + gt_answer = extract_ground_truth(test_case) + + res_dict = llm_func(query) + evaluation_results[line_id] = res_dict + try: + res = res_dict['reply'] + except KeyError: + res = res_dict + try: + chat_history = res_dict['chat_history'] + except KeyError: + chat_history = "" + formatted_history = '\n'.join([ + chat['role'] + '(' + chat.get('name', "unamed") + "): " + chat['content'] + for chat in chat_history + ]) + logging.info(f"[{line_id}] Chat History:\n{formatted_history}") + # llm_answer = extract_llm_answer(res) + llm_answer = res + case_score = score_the_answer(gt_answer, llm_answer) + + hit_rate = case_score['hit_rate'] + strict_hit_rate = case_score['strict_hit_rate'] + soft_hit_rate = case_score['soft_hit_rate'] + + final_score += hit_rate + final_strict_score += strict_hit_rate + final_soft_score += soft_hit_rate + total_score += 1.0 + + if hit_rate == 1.0: # answer correctly + logging.info( + f"[{line_id}] Correct! (hit_rate: {hit_rate}, final: {final_score})\nExpected answer:\n{gt_answer}\nLLM:\n{llm_answer}" + ) + print( + f"[{line_id}] Correct! (hit_rate: {hit_rate}, final: {final_score})\nExpected answer:\n{gt_answer}\nLLM:\n{llm_answer}" + ) + else: + logging.error( + f"[{line_id}] Wrong! (hit_rate: {hit_rate}, final: {final_score})\nExpected answer:\n{gt_answer}\nLLM:\n{llm_answer}" + ) + print( + f"[{line_id}] Wrong! (hit_rate: {hit_rate}, final: {final_score})\nExpected answer:\n{gt_answer}\nLLM:\n{llm_answer}" + ) + + # time.sleep(5) + + # save the state + state = { + 'current_line': line_id, + 'final_score': final_score, + 'final_strict_score': final_strict_score, + 'final_soft_score': final_soft_score, + 'total_score': total_score, + 'evaluation_results': evaluation_results, + } + save_state(state, state_filename=state_filename) + + logging.info(f"[Summary] [Datasets Hit Rate] final score {final_score:.2f} / {total_score}") + logging.info( + f"[Summary] [All Dataset Found] final score {final_strict_score:.2f} / {total_score}") + logging.info( + f"[Summary] [Any Dataset Found] final score {final_soft_score:.2f} / {total_score}") + + if not args.skip_results_saving: + logging.info(f"[Results] result saved to {result_file}") + with open(result_file, 'w') as f: + json.dump(evaluation_results, f, indent=4) + else: + logging.info("[Results] skip results saving") + + if os.path.exists(state_filename): + logging.info("[Results] clean state") + os.remove(state_filename) + + return { + 'final_score': final_score, + 'final_strict_score': final_strict_score, + 'final_soft_score': final_soft_score, + 'total_score': total_score, + } + + +def recompute_metrics_from_file(question_file, result_file): + final_score, total_score = 0, 0 + final_strict_score, final_soft_score = 0, 0 + evaluation_results = json.load(open(result_file, "r")) + + question_numbers = 0 + with open(question_file, 'r') as file: + for line_id, line in enumerate(file): + question_numbers += 1 + + if (len(evaluation_results) != question_numbers): + logging.error("") + return -1 + + with open(question_file, 'r') as file: + for line_id, line in enumerate(file): + line_id += 1 + + test_case = json.loads(line) + gt_answer = extract_ground_truth(test_case) + + res_dict = evaluation_results[str(line_id)] + res = res_dict['reply'] + chat_history = res_dict['chat_history'] + formatted_history = '\n'.join( + [chat['role'] + ": " + chat['content'] for chat in chat_history]) + logging.info(f"[{line_id}] Chat History:\n{formatted_history}") + llm_answer = extract_llm_answer(res) + + case_score = score_the_answer(gt_answer, llm_answer) + + hit_rate = case_score['hit_rate'] + strict_hit_rate = case_score['strict_hit_rate'] + soft_hit_rate = case_score['soft_hit_rate'] + + final_score += hit_rate + final_strict_score += strict_hit_rate + final_soft_score += soft_hit_rate + total_score += 1.0 + + if hit_rate == 1.0: # answer correctly + logging.info( + f"[{line_id}] Correct! (hit_rate: {hit_rate}, final: {final_score})\nExpected answer:\n{gt_answer}\nLLM:\n{llm_answer}" + ) + else: + logging.error( + f"[{line_id}] Wrong! (hit_rate: {hit_rate}, final: {final_score})\nExpected answer:\n{gt_answer}\nLLM:\n{llm_answer}" + ) + + logging.warning(f"[Summary] [Dataset Hit Rate] final score {final_score:.2f} / {total_score}") + logging.warning( + f"[Summary] [Dataset All Found] final score {final_strict_score:.2f} / {total_score}") + logging.warning( + f"[Summary] [Dataset Any Found] final score {final_soft_score:.2f} / {total_score}") + + return { + 'final_score': final_score, + 'final_strict_score': final_strict_score, + 'final_soft_score': final_soft_score, + 'total_score': total_score, + } + + +def get_timestamped_log_file(pattern, file_path, tag): + timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') + base_name = osp.basename(file_path).replace('.jsonl', '') + os.makedirs("results", exist_ok=True) + return osp.join("results", f'{pattern}_{base_name}_{tag}_{timestamp}.log') + + +def main(test_file_path, test_pattern, model, tag, resume, log_to_file): + tag = osp.basename(model) + log_outfile = get_timestamped_log_file(test_pattern, test_file_path, tag) + if (args.recompute_metrics_from_file): + log_outfile = args.recompute_metrics_from_file.replace(".json", "_recompute.log") + if resume: + state = load_state() + if state: + log_outfile = log_outfile.replace(".log", "_resume.log") + + if log_to_file: + print(f"saving file to {log_outfile}") + logging.basicConfig(filename=log_outfile, + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s') + else: + handler = colorlog.StreamHandler() + log_colors = { + 'WARNING': 'yellow', + 'ERROR': 'red', + 'CRITICAL': 'bold_red', + } + formatter = colorlog.ColoredFormatter( + '%(log_color)s%(asctime)s - %(levelname)s - %(message)s', log_colors=log_colors) + handler.setFormatter(formatter) + logging.basicConfig(level=logging.INFO) + logger = logging.getLogger() + logger.handlers.clear() # Clear existing handlers + logger.addHandler(handler) + + if (args.recompute_metrics_from_file): + pass + final_score_dict = recompute_metrics_from_file( + question_file=test_file_path, + result_file=args.recompute_metrics_from_file, + ) + print( + f"recomputed hit rate: {final_score_dict['final_score']:.2f} / {final_score_dict['total_score']}" + ) + print( + f"recomputed all found: {final_score_dict['final_strict_score']:.2f} / {final_score_dict['total_score']}" + ) + print( + f"recomputed any found: {final_score_dict['final_soft_score']:.2f} / {final_score_dict['total_score']}" + ) + else: + final_score_dict = run_tests_from_jsonl( + test_file_path, + llm_func=lambda x: chat_with_llm(test_pattern, x, model=model), + result_file=log_outfile.replace(".log", ".json"), + resume=resume) + print( + f"hit rate: {final_score_dict['final_score']:.2f} / {final_score_dict['total_score']}") + print( + f"all found: {final_score_dict['final_strict_score']:.2f} / {final_score_dict['total_score']}" + ) + print( + f"any found: {final_score_dict['final_soft_score']:.2f} / {final_score_dict['total_score']}" + ) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Process some parameters.') + parser.add_argument('--recompute_metrics_from_file', + type=str, + default=None, + help='Recompute the metrics without query') + parser.add_argument('--test_file_path', type=str, required=True, help='Path to the test file') + parser.add_argument('--test_pattern', + type=str, + required=True, + default="Single_Turn", + help='Test pattern') + parser.add_argument('--model', + type=str, + required=True, + default='gpt-4o-mini-2024-07-18', + help='Model name') + parser.add_argument('--tag', type=str, default=None, required=False, help='Tag name') + parser.add_argument('--resume', action='store_true', help='Resume flag') + parser.add_argument('--log_to_file', action='store_true', help='Log to file flag') + parser.add_argument('--skip_results_saving', action='store_true', help='Log to file flag') + args = parser.parse_args() + + main(args.test_file_path, args.test_pattern, args.model, args.tag, args.resume, + args.log_to_file) diff --git a/model_selection_bench.jsonl b/model_selection_bench.jsonl new file mode 100644 index 0000000..c3e51e5 --- /dev/null +++ b/model_selection_bench.jsonl @@ -0,0 +1,343 @@ +{"query": "Please help me segment the spleen in CT images.", "expected_datasets": {"TotalSegmentator_v2": [1], "FLARE22": [3], "AbdomenAtlasMini": [8], "AMOS22_Task2": [1], "WORD": [2]}} +{"query": "Please help me segment the right kidney in CT images.", "expected_datasets": {"TotalSegmentator_v2": [2], "FLARE22": [2], "AbdomenAtlasMini": [4], "AMOS22_Task2": [2], "WORD": [4]}} +{"query": "Please help me segment the left kidney in CT images.", "expected_datasets": {"TotalSegmentator_v2": [3], "FLARE22": [13], "AbdomenAtlasMini": [3], "AMOS22_Task2": [3], "WORD": [3]}} +{"query": "Please help me segment the kidneys in CT images.", "expected_datasets": {"TotalSegmentator_v2": [2, 3], "FLARE22": [2, 13], "AbdomenAtlasMini": [3, 4], "AMOS22_Task2": [2, 3], "KiTS23": [1], "CT-ORG": [4], "WORD": [3, 4]}} +{"query": "Please help me segment the gallbladder in CT images.", "expected_datasets": {"TotalSegmentator_v2": [4], "FLARE22": [9], "AbdomenAtlasMini": [2], "AMOS22_Task2": [4], "WORD": [6]}} +{"query": "Please help me segment the liver in CT images.", "expected_datasets": {"TotalSegmentator_v2": [5], "FLARE22": [1], "AbdomenAtlasMini": [5], "AMOS22_Task2": [6], "CT-ORG": [1], "LiTS": [1], "WORD": [1]}} +{"query": "Please help me segment the stomach in CT images.", "expected_datasets": {"TotalSegmentator_v2": [6], "FLARE22": [11], "AbdomenAtlasMini": [9], "AMOS22_Task2": [7], "WORD": [5]}} +{"query": "Please help me segment the pancreas in CT images.", "expected_datasets": {"TotalSegmentator_v2": [7], "FLARE22": [4], "AbdomenAtlasMini": [6], "AMOS22_Task2": [10], "WORD": [8]}} +{"query": "Please help me segment the right adrenal gland in CT images.", "expected_datasets": {"TotalSegmentator_v2": [8], "FLARE22": [7], "AMOS22_Task2": [11]}} +{"query": "Please help me segment the left adrenal gland in CT images.", "expected_datasets": {"TotalSegmentator_v2": [9], "FLARE22": [8], "AMOS22_Task2": [12]}} +{"query": "Please help me segment the adrenal glands in CT images.", "expected_datasets": {"TotalSegmentator_v2": [8, 9], "FLARE22": [7, 8], "AMOS22_Task2": [11, 12], "WORD": [12]}} +{"query": "Please help me segment the upper lobe of the left lung in CT images.", "expected_datasets": {"TotalSegmentator_v2": [10]}} +{"query": "Please help me segment the lower lobe of the left lung in CT images.", "expected_datasets": {"TotalSegmentator_v2": [11]}} +{"query": "Please help me segment the upper lobe of the right lung in CT images.", "expected_datasets": {"TotalSegmentator_v2": [12]}} +{"query": "Please help me segment the middle lobe of the right lung in CT images.", "expected_datasets": {"TotalSegmentator_v2": [13]}} +{"query": "Please help me segment the lower lobe of the right lung in CT images.", "expected_datasets": {"TotalSegmentator_v2": [14]}} +{"query": "Please help me segment the left lung in CT images.", "expected_datasets": {"TotalSegmentator_v2": [10, 11]}} +{"query": "Please help me segment the right lung in CT images.", "expected_datasets": {"TotalSegmentator_v2": [12, 13, 14]}} +{"query": "Please help me segment the lungs in CT images.", "expected_datasets": {"CT-ORG": [3], "TotalSegmentator_v2": [10,11,12,13,14]}} +{"query": "Please help me segment the esophagus in CT images.", "expected_datasets": {"TotalSegmentator_v2": [15], "SegRap2023_Task1": [47], "FLARE22": [10], "AMOS22_Task2": [5], "WORD": [7]}} +{"query": "Please help me segment the trachea in CT images.", "expected_datasets": {"TotalSegmentator_v2": [16], "SegRap2023_Task1": [54]}} +{"query": "Please help me segment the thyroid gland in CT images.", "expected_datasets": {"TotalSegmentator_v2": [17], "SegRap2023_Task1": [53]}} +{"query": "Please help me segment the small bowel in CT images.", "expected_datasets": {"TotalSegmentator_v2": [18]}} +{"query": "Please help me segment the duodenum in CT images.", "expected_datasets": {"TotalSegmentator_v2": [19], "FLARE22": [12], "AMOS22_Task2": [13], "WORD": [9]}} +{"query": "Please help me segment the colon in CT images.", "expected_datasets": {"TotalSegmentator_v2": [20], "WORD": [10]}} +{"query": "Please help me segment the urinary bladder in CT images.", "expected_datasets": {"TotalSegmentator_v2": [21], "CT-ORG": [2], "AMOS22_Task2": [14], "WORD": [14]}} +{"query": "Please help me segment the prostate in CT images.", "expected_datasets": {"TotalSegmentator_v2": [22], "AMOS22_Task2": [15]}} +{"query": "Please help me segment the left kidney cyst in CT images.", "expected_datasets": {"TotalSegmentator_v2": [23]}} +{"query": "Please help me segment the right kidney cyst in CT images.", "expected_datasets": {"TotalSegmentator_v2": [24]}} +{"query": "Please help me segment the kidney cyst in CT images.", "expected_datasets": {"TotalSegmentator_v2": [23, 24]}} +{"query": "Please help me segment the sacrum in CT images.", "expected_datasets": {"TotalSegmentator_v2": [25]}} +{"query": "Please help me segment the C1 cervical vertebra (the first cervical vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [50]}} +{"query": "Please help me segment the C2 cervical vertebra (the second cervical vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [49]}} +{"query": "Please help me segment the C3 cervical vertebra (the third cervical vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [48]}} +{"query": "Please help me segment the C4 cervical vertebra (the fourth cervical vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [47]}} +{"query": "Please help me segment the C5 cervical vertebra (the fifth cervical vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [46]}} +{"query": "Please help me segment the C6 cervical vertebra (the sixth cervical vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [45]}} +{"query": "Please help me segment the C7 cervical vertebra (the seventh cervical vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [44]}} +{"query": "Please help me segment the cervical spine (C1-C7 cervical vertebrae) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [44, 45, 46, 47, 48, 49, 50]}} +{"query": "Please help me segment the T1 thoracic vertebra (the first thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [43]}} +{"query": "Please help me segment the T2 thoracic vertebra (the second thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [42]}} +{"query": "Please help me segment the T3 thoracic vertebra (the third thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [41]}} +{"query": "Please help me segment the T4 thoracic vertebra (the fourth thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [40]}} +{"query": "Please help me segment the T5 thoracic vertebra (the fifth thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [39]}} +{"query": "Please help me segment the T6 thoracic vertebra (the sixth thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [38]}} +{"query": "Please help me segment the T7 thoracic vertebra (the seventh thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [37]}} +{"query": "Please help me segment the T8 thoracic vertebra (the eighth thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [36]}} +{"query": "Please help me segment the T9 thoracic vertebra (the ninth thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [35]}} +{"query": "Please help me segment the T10 thoracic vertebra (the tenth thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [34]}} +{"query": "Please help me segment the T11 thoracic vertebra (the eleventh thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [33]}} +{"query": "Please help me segment the T12 thoracic vertebra (the twelfth thoracic vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [32]}} +{"query": "Please help me segment the thoracic spine (T1-T12 thoracic vertebrae) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43]}} +{"query": "Please help me segment the L1 lumbar vertebra (the first lumbar vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [31]}} +{"query": "Please help me segment the L2 lumbar vertebra (the second lumbar vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [30]}} +{"query": "Please help me segment the L3 lumbar vertebra (the third lumbar vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [29]}} +{"query": "Please help me segment the L4 lumbar vertebra (the fourth lumbar vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [28]}} +{"query": "Please help me segment the L5 lumbar vertebra (the fifth lumbar vertebra) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [27]}} +{"query": "Please help me segment the lumbar spine (L1-L5 lumbar vertebrae) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [27, 28, 29, 30, 31]}} +{"query": "Please help me segment the first sacral vertebra (S1) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [26]}} +{"query": "Please help me segment the lumbosacral spine (L1-S1 lumbar and sacral vertebrae) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [26, 27, 28, 29, 30, 31]}} +{"query": "Please help me segment the entire spine (C1-S1 cervical, thoracic, lumbar, and sacral vertebrae) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50]}} +{"query": "Please help me segment the heart in CT images.", "expected_datasets": {"TotalSegmentator_v2": [51], "MM-WHS": [1, 2, 3, 4, 5, 6, 7]}} +{"query": "Please help me segment the aorta in CT images.", "expected_datasets": {"TotalSegmentator_v2": [52], "FLARE22": [5], "AbdomenAtlasMini": [1], "AMOS22_Task2": [8]}} +{"query": "Please help me segment the pulmonary vein in CT images.", "expected_datasets": {"TotalSegmentator_v2": [53]}} +{"query": "Please help me segment the brachiocephalic trunk in CT images.", "expected_datasets": {"TotalSegmentator_v2": [54]}} +{"query": "Please help me segment the right subclavian artery in CT images.", "expected_datasets": {"TotalSegmentator_v2": [55]}} +{"query": "Please help me segment the left subclavian artery in CT images.", "expected_datasets": {"TotalSegmentator_v2": [56]}} +{"query": "Please help me segment the subclavian arteries in CT images.", "expected_datasets": {"TotalSegmentator_v2": [55, 56]}} +{"query": "Please help me segment the right common carotid artery in CT images.", "expected_datasets": {"TotalSegmentator_v2": [57]}} +{"query": "Please help me segment the left common carotid artery in CT images.", "expected_datasets": {"TotalSegmentator_v2": [58]}} +{"query": "Please help me segment the common carotid arteries in CT images.", "expected_datasets": {"TotalSegmentator_v2": [57, 58]}} +{"query": "Please help me segment the left brachiocephalic vein in CT images.", "expected_datasets": {"TotalSegmentator_v2": [59]}} +{"query": "Please help me segment the right brachiocephalic vein in CT images.", "expected_datasets": {"TotalSegmentator_v2": [60]}} +{"query": "Please help me segment the brachiocephalic veins in CT images.", "expected_datasets": {"TotalSegmentator_v2": [59, 60]}} +{"query": "Please help me segment the left atrial appendage in CT images.", "expected_datasets": {"TotalSegmentator_v2": [61]}} +{"query": "Please help me segment the superior vena cava in CT images.", "expected_datasets": {"TotalSegmentator_v2": [62]}} +{"query": "Please help me segment the inferior vena cava in CT images.", "expected_datasets": {"TotalSegmentator_v2": [63], "FLARE22": [6], "AMOS22_Task2": [9], "AbdomenAtlasMini": [7]}} +{"query": "Please help me segment the portal vein and splenic vein in CT images.", "expected_datasets": {"TotalSegmentator_v2": [64]}} +{"query": "Please help me segment the left iliac artery in CT images.", "expected_datasets": {"TotalSegmentator_v2": [65]}} +{"query": "Please help me segment the right iliac artery in CT images.", "expected_datasets": {"TotalSegmentator_v2": [66]}} +{"query": "Please help me segment the iliac arteries in CT images.", "expected_datasets": {"TotalSegmentator_v2": [65, 66]}} +{"query": "Please help me segment the left iliac vein in CT images.", "expected_datasets": {"TotalSegmentator_v2": [67]}} +{"query": "Please help me segment the right iliac vein in CT images.", "expected_datasets": {"TotalSegmentator_v2": [68]}} +{"query": "Please help me segment the iliac veins in CT images.", "expected_datasets": {"TotalSegmentator_v2": [67, 68]}} +{"query": "Please help me segment the iliac vessels in CT images.", "expected_datasets": {"TotalSegmentator_v2": [65, 66, 67, 68]}} +{"query": "Please help me segment the left humerus in CT images.", "expected_datasets": {"TotalSegmentator_v2": [69]}} +{"query": "Please help me segment the right humerus in CT images.", "expected_datasets": {"TotalSegmentator_v2": [70]}} +{"query": "Please help me segment the humeri in CT images.", "expected_datasets": {"TotalSegmentator_v2": [69, 70]}} +{"query": "Please help me segment the left scapula in CT images.", "expected_datasets": {"TotalSegmentator_v2": [71]}} +{"query": "Please help me segment the right scapula in CT images.", "expected_datasets": {"TotalSegmentator_v2": [72]}} +{"query": "Please help me segment the scapulae in CT images.", "expected_datasets": {"TotalSegmentator_v2": [71, 72]}} +{"query": "Please help me segment the left clavicle in CT images.", "expected_datasets": {"TotalSegmentator_v2": [73]}} +{"query": "Please help me segment the right clavicle in CT images.", "expected_datasets": {"TotalSegmentator_v2": [74]}} +{"query": "Please help me segment the clavicles in CT images.", "expected_datasets": {"TotalSegmentator_v2": [73, 74]}} +{"query": "Please help me segment the left femur in CT images.", "expected_datasets": {"TotalSegmentator_v2": [75], "WORD": [15]}} +{"query": "Please help me segment the right femur in CT images.", "expected_datasets": {"TotalSegmentator_v2": [76], "WORD": [16]}} +{"query": "Please help me segment the femurs in CT images.", "expected_datasets": {"TotalSegmentator_v2": [75, 76], "WORD": [15, 16]}} +{"query": "Please help me segment the left hip in CT images.", "expected_datasets": {"TotalSegmentator_v2": [77]}} +{"query": "Please help me segment the right hip in CT images.", "expected_datasets": {"TotalSegmentator_v2": [78]}} +{"query": "Please help me segment the hips in CT images.", "expected_datasets": {"TotalSegmentator_v2": [77, 78]}} +{"query": "Please help me segment the spinal cord in CT images.", "expected_datasets": {"TotalSegmentator_v2": [79], "SegRap2023_Task1": [46]}} +{"query": "Please help me segment the left gluteus maximus in CT images.", "expected_datasets": {"TotalSegmentator_v2": [80]}} +{"query": "Please help me segment the right gluteus maximus in CT images.", "expected_datasets": {"TotalSegmentator_v2": [81]}} +{"query": "Please help me segment the gluteus maximus muscles in CT images.", "expected_datasets": {"TotalSegmentator_v2": [80, 81]}} +{"query": "Please help me segment the left gluteus medius in CT images.", "expected_datasets": {"TotalSegmentator_v2": [82]}} +{"query": "Please help me segment the right gluteus medius in CT images.", "expected_datasets": {"TotalSegmentator_v2": [83]}} +{"query": "Please help me segment the gluteus medius muscles in CT images.", "expected_datasets": {"TotalSegmentator_v2": [82, 83]}} +{"query": "Please help me segment the left gluteus minimus in CT images.", "expected_datasets": {"TotalSegmentator_v2": [84]}} +{"query": "Please help me segment the right gluteus minimus in CT images.", "expected_datasets": {"TotalSegmentator_v2": [85]}} +{"query": "Please help me segment the gluteus minimus muscles in CT images.", "expected_datasets": {"TotalSegmentator_v2": [84, 85]}} +{"query": "Please help me segment the gluteus muscles in CT images.", "expected_datasets": {"TotalSegmentator_v2": [80, 81, 82, 83, 84, 85]}} +{"query": "Please help me segment the left autochthon muscle group in CT images.", "expected_datasets": {"TotalSegmentator_v2": [86]}} +{"query": "Please help me segment the right autochthon muscle group in CT images.", "expected_datasets": {"TotalSegmentator_v2": [87]}} +{"query": "Please help me segment the autochthon muscle groups in CT images.", "expected_datasets": {"TotalSegmentator_v2": [86, 87]}} +{"query": "Please help me segment the left iliopsoas muscle in CT images.", "expected_datasets": {"TotalSegmentator_v2": [88]}} +{"query": "Please help me segment the right iliopsoas muscle in CT images.", "expected_datasets": {"TotalSegmentator_v2": [89]}} +{"query": "Please help me segment the iliopsoas muscles in CT images.", "expected_datasets": {"TotalSegmentator_v2": [88, 89]}} +{"query": "Please help me segment the brain in CT images.", "expected_datasets": {"TotalSegmentator_v2": [90], "SegRap2023_Task1": [1], "CT-ORG": [6]}} +{"query": "Please help me segment the skull in CT images.", "expected_datasets": {"TotalSegmentator_v2": [91]}} +{"query": "Please help me segment the left rib 1 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [94]}} +{"query": "Please help me segment the left rib 2 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [95]}} +{"query": "Please help me segment the left rib 3 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [96]}} +{"query": "Please help me segment the left rib 4 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [97]}} +{"query": "Please help me segment the left rib 5 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [98]}} +{"query": "Please help me segment the left rib 6 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [99]}} +{"query": "Please help me segment the left rib 7 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [100]}} +{"query": "Please help me segment the left rib 8 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [101]}} +{"query": "Please help me segment the left rib 9 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [102]}} +{"query": "Please help me segment the left rib 10 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [103]}} +{"query": "Please help me segment the left rib 11 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [104]}} +{"query": "Please help me segment the left rib 12 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [105]}} +{"query": "Please help me segment the left ribs from 1 to 12 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105]}} +{"query": "Please help me segment the right rib 1 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [106]}} +{"query": "Please help me segment the right rib 2 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [107]}} +{"query": "Please help me segment the right rib 3 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [93]}} +{"query": "Please help me segment the right rib 4 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [92]}} +{"query": "Please help me segment the right rib 5 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [108]}} +{"query": "Please help me segment the right rib 6 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [109]}} +{"query": "Please help me segment the right rib 7 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [110]}} +{"query": "Please help me segment the right rib 8 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [111]}} +{"query": "Please help me segment the right rib 9 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [112]}} +{"query": "Please help me segment the right rib 10 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [113]}} +{"query": "Please help me segment the right rib 11 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [114]}} +{"query": "Please help me segment the right rib 12 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [115]}} +{"query": "Please help me segment the right ribs from 1 to 12 in CT images.", "expected_datasets": {"TotalSegmentator_v2": [106, 107, 93, 92, 108, 109, 110, 111, 112, 113, 114, 115]}} +{"query": "Please help me segment all ribs in CT images.", "expected_datasets": {"TotalSegmentator_v2": [94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 93, 92, 108, 109, 110, 111, 112, 113, 114, 115]}} +{"query": "Please help me segment the left true ribs (1-7) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [94, 95, 96, 97, 98, 99, 100]}} +{"query": "Please help me segment the right true ribs (1-7) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [106, 107, 93, 92, 108, 109, 110]}} +{"query": "Please help me segment all true ribs (1-7) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [94, 95, 96, 97, 98, 99, 100, 106, 107, 93, 92, 108, 109, 110]}} +{"query": "Please help me segment the left false ribs (8-12) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [101, 102, 103, 104, 105]}} +{"query": "Please help me segment the right false ribs (8-12) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [111, 112, 113, 114, 115]}} +{"query": "Please help me segment all false ribs (8-12) in CT images.", "expected_datasets": {"TotalSegmentator_v2": [101, 102, 103, 104, 105, 111, 112, 113, 114, 115]}} +{"query": "Please help me segment the sternum in CT images.", "expected_datasets": {"TotalSegmentator_v2": [116]}} +{"query": "Please help me segment the costal cartilages in CT images.", "expected_datasets": {"TotalSegmentator_v2": [117]}} +{"query": "Please help me segment the spleen in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [1], "AMOS22_Task2": [1]}} +{"query": "Please help me segment the right kidney in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [2], "AMOS22_Task2": [2]}} +{"query": "Please help me segment the left kidney in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [3], "AMOS22_Task2": [3]}} +{"query": "Please help me segment the kidneys in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [2, 3], "AMOS22_Task2": [2, 3]}} +{"query": "Please help me segment the gallbladder in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [4], "AMOS22_Task2": [4]}} +{"query": "Please help me segment the liver in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [5], "AMOS22_Task2": [6]}} +{"query": "Please help me segment the stomach in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [6], "AMOS22_Task2": [7]}} +{"query": "Please help me segment the pancreas in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [7], "AMOS22_Task2": [10]}} +{"query": "Please help me segment the right suprarenal gland in MR images.", "expected_datasets": {"AMOS22_Task2": [11],"TotalSegmentator_MRI": [8]}} +{"query": "Please help me segment the left suprarenal gland in MR images.", "expected_datasets": {"AMOS22_Task2": [12], "TotalSegmentator_MRI": [9]}} +{"query": "Please help me segment the suprarenal glands in MR images.", "expected_datasets": {"AMOS22_Task2": [11, 12],"TotalSegmentator_MRI": [8, 9]}} +{"query": "Please help me segment the left lung in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [10]}} +{"query": "Please help me segment the right lung in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [11]}} +{"query": "Please help me segment the lungs in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [10, 11]}} +{"query": "Please help me segment the esophagus in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [12], "AMOS22_Task2": [5]}} +{"query": "Please help me segment the small intestine in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [13]}} +{"query": "Please help me segment the duodenum in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [14], "AMOS22_Task2": [13]}} +{"query": "Please help me segment the colon in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [15]}} +{"query": "Please help me segment the urinary bladder in MR images.", "expected_datasets": {"AMOS22_Task2": [14], "TotalSegmentator_MRI": [16]}} +{"query": "Please help me segment the prostate in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [17], "AMOS22_Task2": [15]}} +{"query": "Please help me segment the sacrum in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [18]}} +{"query": "Please help me segment the vertebrae in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [19]}} +{"query": "Please help me segment the intervertebral discs in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [20]}} +{"query": "Please help me segment the spinal cord in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [21]}} +{"query": "Please help me segment the heart in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [22], "MM-WHS": [1, 2, 3, 4, 5, 6, 7], "MMs": [1,2,3]}} +{"query": "Please help me segment the aorta in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [23], "AMOS22_Task2": [8]}} +{"query": "Please help me segment the inferior vena cava in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [24], "AMOS22_Task2": [9]}} +{"query": "Please help me segment the hepatic portal vein in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [25]}} +{"query": "Please help me segment the left common iliac artery in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [26]}} +{"query": "Please help me segment the right common iliac artery in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [27]}} +{"query": "Please help me segment the common iliac arteries in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [26, 27]}} +{"query": "Please help me segment the left common iliac vein in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [28]}} +{"query": "Please help me segment the right common iliac vein in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [29]}} +{"query": "Please help me segment the common iliac veins in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [28, 29]}} +{"query": "Please help me segment the common iliac vessels in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [26, 27, 28, 29]}} +{"query": "Please help me segment the left humerus in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [30]}} +{"query": "Please help me segment the right humerus in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [31]}} +{"query": "Please help me segment the humeri in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [30, 31]}} +{"query": "Please help me segment the fibula in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [32]}} +{"query": "Please help me segment the tibia in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [33]}} +{"query": "Please help me segment the left femur in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [34]}} +{"query": "Please help me segment the right femur in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [35]}} +{"query": "Please help me segment the femurs in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [34, 35]}} +{"query": "Please help me segment the left hip in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [36]}} +{"query": "Please help me segment the right hip in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [37]}} +{"query": "Please help me segment the hips in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [36, 37]}} +{"query": "Please help me segment the left gluteus maximus muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [38]}} +{"query": "Please help me segment the right gluteus maximus muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [39]}} +{"query": "Please help me segment the gluteus maximus muscles in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [38, 39]}} +{"query": "Please help me segment the left gluteus medius muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [40]}} +{"query": "Please help me segment the right gluteus medius muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [41]}} +{"query": "Please help me segment the gluteus medius muscles in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [40, 41]}} +{"query": "Please help me segment the left gluteus minimus muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [42]}} +{"query": "Please help me segment the right gluteus minimus muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [43]}} +{"query": "Please help me segment the gluteus minimus muscles in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [42, 43]}} +{"query": "Please help me segment the left autochthon muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [44]}} +{"query": "Please help me segment the right autochthon muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [45]}} +{"query": "Please help me segment the autochthon muscles in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [44, 45]}} +{"query": "Please help me segment the left iliopsoas muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [46]}} +{"query": "Please help me segment the right iliopsoas muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [47]}} +{"query": "Please help me segment the iliopsoas muscles in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [46, 47]}} +{"query": "Please help me segment the left quadriceps femoris muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [48]}} +{"query": "Please help me segment the right quadriceps femoris muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [49]}} +{"query": "Please help me segment the quadriceps femoris muscles in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [48, 49]}} +{"query": "Please help me segment the left medial compartment of the thigh in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [50]}} +{"query": "Please help me segment the right medial compartment of the thigh in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [51]}} +{"query": "Please help me segment the medial compartments of the thighs in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [50, 51]}} +{"query": "Please help me segment the left posterior compartment of the thigh in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [52]}} +{"query": "Please help me segment the right posterior compartment of the thigh in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [53]}} +{"query": "Please help me segment the posterior compartments of the thighs in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [52, 53]}} +{"query": "Please help me segment the left sartorius muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [54]}} +{"query": "Please help me segment the right sartorius muscle in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [55]}} +{"query": "Please help me segment the sartorius muscles in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [54, 55]}} +{"query": "Please help me segment the brain in MR images.", "expected_datasets": {"TotalSegmentator_MRI": [56]}} +{"query": "Please help me segment the uterus in MR images.", "expected_datasets": {"AMOS22_Task2": [15]}} +{"query": "Please help me segment the brainstem in CT images.", "expected_datasets": {"SegRap2023_Task1": [2]}} +{"query": "Please help me segment the chiasm in CT images.", "expected_datasets": {"SegRap2023_Task1": [3]}} +{"query": "Please help me segment the left temporal lobe in CT images.", "expected_datasets": {"SegRap2023_Task1": [4, 6]}} +{"query": "Please help me segment the right temporal lobe in CT images.", "expected_datasets": {"SegRap2023_Task1": [5, 7]}} +{"query": "Please help me segment the left hippocampus in CT images.", "expected_datasets": {"SegRap2023_Task1": [6, 8]}} +{"query": "Please help me segment the right hippocampus in CT images.", "expected_datasets": {"SegRap2023_Task1": [7, 9]}} +{"query": "Please help me segment the left eye in CT images.", "expected_datasets": {"SegRap2023_Task1": [10]}} +{"query": "Please help me segment the right eye in CT images.", "expected_datasets": {"SegRap2023_Task1": [11]}} +{"query": "Please help me segment the eyes in CT images.", "expected_datasets": {"SegRap2023_Task1": [10, 11]}} +{"query": "Please help me segment the left lens in CT images.", "expected_datasets": {"SegRap2023_Task1": [12]}} +{"query": "Please help me segment the right lens in CT images.", "expected_datasets": {"SegRap2023_Task1": [13]}} +{"query": "Please help me segment the lens in CT images.", "expected_datasets": {"SegRap2023_Task1": [12, 13]}} +{"query": "Please help me segment the left optic nerve in CT images.", "expected_datasets": {"SegRap2023_Task1": [14]}} +{"query": "Please help me segment the right optic nerve in CT images.", "expected_datasets": {"SegRap2023_Task1": [15]}} +{"query": "Please help me segment the optic nerves in CT images.", "expected_datasets": {"SegRap2023_Task1": [14, 15]}} +{"query": "Please help me segment the left middle ear in CT images.", "expected_datasets": {"SegRap2023_Task1": [16, 20, 24, 30]}} +{"query": "Please help me segment the right middle ear in CT images.", "expected_datasets": {"SegRap2023_Task1": [17, 21, 25, 31]}} +{"query": "Please help me segment the middle ears in CT images.", "expected_datasets": {"SegRap2023_Task1": [16, 17, 20, 21, 24, 25, 30, 31]}} +{"query": "Please help me segment the left internal auditory canal in CT images.", "expected_datasets": {"SegRap2023_Task1": [18]}} +{"query": "Please help me segment the right internal auditory canal in CT images.", "expected_datasets": {"SegRap2023_Task1": [19]}} +{"query": "Please help me segment the internal auditory canals in CT images.", "expected_datasets": {"SegRap2023_Task1": [18, 19]}} +{"query": "Please help me segment the left tympanic cavity in CT images.", "expected_datasets": {"SegRap2023_Task1": [20, 22]}} +{"query": "Please help me segment the right tympanic cavity in CT images.", "expected_datasets": {"SegRap2023_Task1": [21, 23]}} +{"query": "Please help me segment the tympanic cavities in CT images.", "expected_datasets": {"SegRap2023_Task1": [20, 21, 22, 23]}} +{"query": "Please help me segment the left vestibular semi-circular canals in CT images.", "expected_datasets": {"SegRap2023_Task1": [24, 26]}} +{"query": "Please help me segment the right vestibular semi-circular canals in CT images.", "expected_datasets": {"SegRap2023_Task1": [25, 27]}} +{"query": "Please help me segment the vestibular semi-circular canals in CT images.", "expected_datasets": {"SegRap2023_Task1": [24, 25, 26, 27]}} +{"query": "Please help me segment the left cochlea in CT images.", "expected_datasets": {"SegRap2023_Task1": [28]}} +{"query": "Please help me segment the right cochlea in CT images.", "expected_datasets": {"SegRap2023_Task1": [29]}} +{"query": "Please help me segment the cochlea in CT images.", "expected_datasets": {"SegRap2023_Task1": [28, 29]}} +{"query": "Please help me segment the left eustachian tube bone in CT images.", "expected_datasets": {"SegRap2023_Task1": [30, 32]}} +{"query": "Please help me segment the right eustachian tube bone in CT images.", "expected_datasets": {"SegRap2023_Task1": [31, 33]}} +{"query": "Please help me segment the eustachian tube bone in CT images.", "expected_datasets": {"SegRap2023_Task1": [30, 31, 32, 33]}} +{"query": "Please help me segment the pituitary gland in CT images.", "expected_datasets": {"SegRap2023_Task1": [34]}} +{"query": "Please help me segment the oral cavity in CT images.", "expected_datasets": {"SegRap2023_Task1": [35]}} +{"query": "Please help me segment the left mandible in CT images.", "expected_datasets": {"SegRap2023_Task1": [36]}} +{"query": "Please help me segment the right mandible in CT images.", "expected_datasets": {"SegRap2023_Task1": [37]}} +{"query": "Please help me segment the mandible in CT images.", "expected_datasets": {"SegRap2023_Task1": [36, 37]}} +{"query": "Please help me segment the left submandibular gland in CT images.", "expected_datasets": {"SegRap2023_Task1": [38]}} +{"query": "Please help me segment the right submandibular gland in CT images.", "expected_datasets": {"SegRap2023_Task1": [39]}} +{"query": "Please help me segment the submandibular gland in CT images.", "expected_datasets": {"SegRap2023_Task1": [38, 39]}} +{"query": "Please help me segment the left parotid gland in CT images.", "expected_datasets": {"SegRap2023_Task1": [40]}} +{"query": "Please help me segment the right parotid gland in CT images.", "expected_datasets": {"SegRap2023_Task1": [41]}} +{"query": "Please help me segment the parotid gland in CT images.", "expected_datasets": {"SegRap2023_Task1": [40, 41]}} +{"query": "Please help me segment the left mastoid in CT images.", "expected_datasets": {"SegRap2023_Task1": [42]}} +{"query": "Please help me segment the right mastoid in CT images.", "expected_datasets": {"SegRap2023_Task1": [43]}} +{"query": "Please help me segment the mastoid in CT images.", "expected_datasets": {"SegRap2023_Task1": [42, 43]}} +{"query": "Please help me segment the left temporomandibular joint in CT images.", "expected_datasets": {"SegRap2023_Task1": [44]}} +{"query": "Please help me segment the right temporomandibular joint in CT images.", "expected_datasets": {"SegRap2023_Task1": [45]}} +{"query": "Please help me segment the temporomandibular joint in CT images.", "expected_datasets": {"SegRap2023_Task1": [44, 45]}} +{"query": "Please help me segment the larynx in CT images.", "expected_datasets": {"SegRap2023_Task1": [48, 49, 50, 51]}} +{"query": "Please help me segment the glottic part of the larynx in CT images.", "expected_datasets": {"SegRap2023_Task1": [49]}} +{"query": "Please help me segment the supraglottic part of the larynx in CT images.", "expected_datasets": {"SegRap2023_Task1": [50]}} +{"query": "Please help me segment the pharyngeal constrictor muscles in CT images.", "expected_datasets": {"SegRap2023_Task1": [51, 52]}} +{"query": "Please help me segment the bones in the CT image", "expected_datasets": {"CT-ORG": [5], "TotalSegmentator_v2": [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117]}} +{"query": "Please help me segment the abdominal tumors in CT scans.", "expected_datasets": {"Adrenal-ACC-Ki67-Seg": [1], "LiTS": [2], "KiTS23": [3]}} +{"query": "Please help me segment the coronary arteries in CT images.", "expected_datasets": {"ImageCAS": [1]}} +{"query": "Please help me segment the adrenal tumor in CT images.", "expected_datasets": {"Adrenal-ACC-Ki67-Seg": [1]}} +{"query": "Please help me segment the tumors in whole-body PET images.", "expected_datasets": {"AutoPET": [1]}} +{"query": "Please help me segment the tumors in whole-body CT images.", "expected_datasets": {"Adrenal-ACC-Ki67-Seg": [1], "KiTS23": [3], "AutoPET": [1], "LiTS": [2]}} +{"query": "Please help me segment the tumors in whole-body combined PET/CT images.", "expected_datasets": {"AutoPET": [1]}} +{"query": "Please help me segment the whole tumor in brain MRI using T2 FLAIR images.", "expected_datasets": {"BraTS21": [1, 2, 3]}} +{"query": "Please help me segment the whole tumor in brain MRI using T1 Weighted images.", "expected_datasets": {"BraTS21": [1, 2, 3]}} +{"query": "Please help me segment the whole tumor in brain MRI using T1 weighted imaging with gadolinium.", "expected_datasets": {"BraTS21": [1, 2, 3]}} +{"query": "Please help me segment the whole tumor in brain MRI using T2 Weighted images.", "expected_datasets": {"BraTS21": [1, 2, 3]}} +{"query": "Please help me segment the whole tumor in brain MRI using a combination of T2 FLAIR, T1 Weighted, T1 weighted imaging with gadolinium, and T2 Weighted images.", "expected_datasets": {"BraTS21": [1, 2, 3]}} +{"query": "Please help me segment the tumor core in brain MRI using T2 FLAIR images.", "expected_datasets": {"BraTS21": [1, 3]}} +{"query": "Please help me segment the tumor core in brain MRI using T1 Weighted images.", "expected_datasets": {"BraTS21": [1, 3]}} +{"query": "Please help me segment the tumor core in brain MRI using T1 weighted imaging with gadolinium.", "expected_datasets": {"BraTS21": [1, 3]}} +{"query": "Please help me segment the tumor core in brain MRI using T2 Weighted images.", "expected_datasets": {"BraTS21": [1, 3]}} +{"query": "Please help me segment the tumor core in brain MRI using a combination of T2 FLAIR, T1 Weighted, T1 weighted imaging with gadolinium, and T2 Weighted images.", "expected_datasets": {"BraTS21": [1, 3]}} +{"query": "Please help me segment the enhancing tumor in brain MRI using T2 FLAIR images.", "expected_datasets": {"BraTS21": [3]}} +{"query": "Please help me segment the enhancing tumor in brain MRI using T1 Weighted images.", "expected_datasets": {"BraTS21": [3]}} +{"query": "Please help me segment the enhancing tumor in brain MRI using T1 weighted imaging with gadolinium.", "expected_datasets": {"BraTS21": [3]}} +{"query": "Please help me segment the enhancing tumor in brain MRI using T2 Weighted images.", "expected_datasets": {"BraTS21": [3]}} +{"query": "Please help me segment the enhancing tumor in brain MRI using a combination of T2 FLAIR, T1 Weighted, T1 weighted imaging with gadolinium, and T2 Weighted images.", "expected_datasets": {"BraTS21": [3]}} +{"query": "Please help me segment the myocardium in MR images.", "expected_datasets": {"MM-WHS": [1], "MMs": [3]}} +{"query": "Please help me segment the myocardium in CT images.", "expected_datasets": {"MM-WHS": [1]}} +{"query": "Please help me segment the left atrium cavity in MR images.", "expected_datasets": {"MM-WHS": [2]}} +{"query": "Please help me segment the left atrium cavity in CT images.", "expected_datasets": {"MM-WHS": [2]}} +{"query": "Please help me segment the left ventricle cavity in MR images.", "expected_datasets": {"MM-WHS": [3], "MMs": [1]}} +{"query": "Please help me segment the left ventricle cavity in CT images.", "expected_datasets": {"MM-WHS": [3]}} +{"query": "Please help me segment the right atrium cavity in MR images.", "expected_datasets": {"MM-WHS": [4]}} +{"query": "Please help me segment the right atrium cavity in CT images.", "expected_datasets": {"MM-WHS": [4]}} +{"query": "Please help me segment the right ventricle cavity in MR images.", "expected_datasets": {"MM-WHS": [5], "MMs": [2]}} +{"query": "Please help me segment the right ventricle cavity in CT images.", "expected_datasets": {"MM-WHS": [5]}} +{"query": "Please help me segment the ascending aorta in MR images.", "expected_datasets": {"MM-WHS": [6]}} +{"query": "Please help me segment the ascending aorta in CT images.", "expected_datasets": {"MM-WHS": [6]}} +{"query": "Please help me segment the pulmonary artery in MR images.", "expected_datasets": {"MM-WHS": [7]}} +{"query": "Please help me segment the pulmonary artery in CT images.", "expected_datasets": {"MM-WHS": [7], "Parse22": [1]}} +{"query": "Please help me segment the left ventricle and myocardium in MR images.", "expected_datasets": {"MM-WHS": [1, 3], "MMs": [1, 3]}} +{"query": "Please help me segment the left ventricle and myocardium in CT images.", "expected_datasets": {"MM-WHS": [1, 3]}} +{"query": "Please help me segment the stroke lesions in T1-weighted MR images.", "expected_datasets": {"ISLES22_ATLAS": [1]}} +{"query": "Please help me segment the stroke lesions using ADC MRI images.", "expected_datasets": {"ISLES22": [1]}} +{"query": "Please help me segment the stroke lesions using DWI MRI images.", "expected_datasets": {"ISLES22": [1]}} +{"query": "Please help me segment the stroke lesions using FLAIR MRI images.", "expected_datasets": {"ISLES22": [1]}} +{"query": "Please help me segment the stroke lesions using a combination of ADC, DWI, and FLAIR MRI images.", "expected_datasets": {"ISLES22": [1]}} +{"query": "Please help me segment the primary tumor (GTVp) in CT images.", "expected_datasets": {"HECKTOR2022": [1], "SegRap2023_Task2": [1]}} +{"query": "Please help me segment the lymph nodes (GTVnd) in CT images.", "expected_datasets": {"HECKTOR2022": [2], "SegRap2023_Task2": [2]}} +{"query": "Please help me segment the primary tumor (GTVp) in PET images.", "expected_datasets": {"HECKTOR2022": [1]}} +{"query": "Please help me segment the lymph nodes (GTVnd) in PET images.", "expected_datasets": {"HECKTOR2022": [2]}} +{"query": "Please help me segment the primary tumor (GTVp) in combined PET/CT images.", "expected_datasets": {"HECKTOR2022": [1]}} +{"query": "Please help me segment the lymph nodes (GTVnd) in combined PET/CT images.", "expected_datasets": {"HECKTOR2022": [2]}} +{"query": "Please help me segment both the primary tumor (GTVp) and lymph nodes (GTVnd) in CT images.", "expected_datasets": {"HECKTOR2022": [1, 2], "SegRap2023_Task2": [1, 2]}} +{"query": "Please help me segment both the primary tumor (GTVp) and lymph nodes (GTVnd) in PET images.", "expected_datasets": {"HECKTOR2022": [1, 2]}} +{"query": "Please help me segment both the primary tumor (GTVp) and lymph nodes (GTVnd) in combined PET/CT images.", "expected_datasets": {"HECKTOR2022": [1, 2]}} +{"query": "Please help me segment the intracranial hemorrhage in non-contrast head CT images.", "expected_datasets": {"Instance22": [1]}} +{"query": "Please help me segment the kidney and masses in CT images.", "expected_datasets": {"KiTS23": [1, 2, 3]}} +{"query": "Please help me segment the kidney masses in CT images.", "expected_datasets": {"KiTS23": [2, 3]}} +{"query": "Please help me segment the kidney tumors in CT images.", "expected_datasets": {"KiTS23": [3]}} +{"query": "Please help me segment the liver tumors in CT images.", "expected_datasets": {"LiTS": [2]}} +{"query": "Please help me segment the liver and liver tumors in CT images.", "expected_datasets": {"LiTS": [1, 2]}} +{"query": "Please help me segment the pulmonary airway in CT images.", "expected_datasets": {"ATM22": [1]}} +{"query": "Please help me segment the intestine in CT images.", "expected_datasets": {"WORD": [11]}} \ No newline at end of file diff --git a/model_selection_test_case.jsonl b/model_selection_test_case.jsonl new file mode 100644 index 0000000..1e8e812 --- /dev/null +++ b/model_selection_test_case.jsonl @@ -0,0 +1,2 @@ +{"query": "Please help me segment the spleen in CT images.", "expected_datasets": {"TotalSegmentator_v2": [1], "FLARE22": [3], "AbdomenAtlasMini": [8], "AMOS22_Task2": [1], "WORD": [2]}} +{"query": "Please help me segment the primary tumor (GTVp) in PET images.", "expected_datasets": {"HECKTOR2022": [1]}} \ No newline at end of file diff --git a/prompts/__init__.py b/prompts/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/prompts/__init__.py diff --git a/prompts/agent_optimizer.py b/prompts/agent_optimizer.py new file mode 100644 index 0000000..b951b08 --- /dev/null +++ b/prompts/agent_optimizer.py @@ -0,0 +1,64 @@ +from prompts.data import DATASET_JSON_INFO + +REACT_FORMATTED_OUTPUT = "{ \"thought\": \"\", \"dataset\": \"\", \"labels\": \"\" }" + +REACT_FORMATTED_EXAMPLES = """ +``` +{\"thought\": \"I need to find all CT datasets containing duodenum. I should first check the modality, then the labels.\",\"dataset\": \"AMOS22_Task2, FLARE23, TotalSegmentator_v2\",\"labels\": \"[13], [12], [19]\",} +``` +``` +{\"thought\": \"To segment adrenal gland in CT, I need to first find all dataset with adrenal gland. Then I need to find \"\"dataset\": \"AMOS22_Task2, FLARE23, TotalSegmentator_v2\"\"labels\": \"[11,12], [8,9], [7,8]\"} +``` +""" + +REACT_MISSING_MODALITY = '{"thought": "No suitable modality", "dataset": "None", "labels": "None"}' +REACT_NO_DATASET_FOUND = '{"thought": "No suitable dataset", "dataset": "None", "labels": "None"}' + +SYSTEM_PROMPT_REACT = f""" +# Role: +You are an AI expert on medical image segmentation, specializing in addressing clinical tasks efficiently. You will answer the input question as best you can. + +## Skills: +1. Dataset & Label Selection +- Choose datasets and corresponding labels based on the input question. +- Use the provided dataset information to ensure compatibility. +- Validate the selected labels can cover the need to segment the requested targets. + +2. Thought Process Optimization +You MUST use this format to think which datasets and their labels should be chosen: +``` +Question: the input question you must answer +Thought: you should always think about what to do +Action: the action to take +Action Input: the input to the action +Observation: the result of the action +... (this process can repeat multiple times) +Thought: I now know the final answer +Final Answer: the final answer to the original input question in JSON format +``` +## Constraints: +1. Thought Process Constraints +- Thoughts must be concise and short (within 3 sentences). +- All suitable datasets and corresponding labels should be selected. + +2. Response Format +- Use the following JSON format for Final Answer: +``` +{REACT_FORMATTED_OUTPUT} +``` +Here are some examples: +{REACT_FORMATTED_EXAMPLES} + +3. Handling Missing Modalities +- If fail to find any needed modalities, reply in JSON: {REACT_MISSING_MODALITY} ('MR' is equal to 'MRI') +- If no suitable dataset is available, reply in JSON: {REACT_NO_DATASET_FOUND} + +Here's the dataset information: +{DATASET_JSON_INFO} + +Let’s Begin! +""" + +REACT_SYS_PROMPT = { + "Expert": SYSTEM_PROMPT_REACT, +} diff --git a/prompts/assigner_experts.py b/prompts/assigner_experts.py new file mode 100644 index 0000000..f014eef --- /dev/null +++ b/prompts/assigner_experts.py @@ -0,0 +1,142 @@ +from prompts.basic_v2 import FORMATTED_OUTPUT +from prompts.data import DATASET_NAME_TO_PATH, get_info_from_json_path_list + +DEFAULT_EXPERT_PROFILE = "You are an AI expert on medical image segmentation, specializing in addressing clinical tasks efficiently." + + +def get_selected_dataset_info(dataset_name_list): + json_path_list = [DATASET_NAME_TO_PATH.get(dn, None) for dn in dataset_name_list] + if (None in json_path_list): + raise ValueError(f"non-exist dataset name found: {dataset_name_list}") + return get_info_from_json_path_list(json_path_list) + + +def generate_expert_sys_prompt_dict(expert_profile, dataset_name_list): + selected_dataset_info = get_selected_dataset_info(dataset_name_list) + ai_expert_sys_prompt = f""" + # Role: + {expert_profile}. + + ## Skills: + 1. Thought Process Optimization + - Quickly analyze the specific clinical need and the segmentation targets. + - Choose all suitable datasets and corresponding labels. + + 2. Dataset & Label Selection + - Use the provided dataset information to ensure compatibility. + - Validate the selected labels can cover the need to segment the requested targets. + + ## Constraints: + 1. Thought Process Constraints + - Thoughts must be concise and short (within 3 sentences). + - All suitable datasets and corresponding labels should be selected. + + 2. Response Format + - Use the following JSON format: + ``` + {FORMATTED_OUTPUT} + ``` + + 3. Handling Missing Modalities + - If fail to find any needed modalities, reply in JSON with: 'Missing modalities.' ('MR' is equal to 'MRI') + - If no suitable dataset is available, reply in JSON with: 'Dataset: None\nLabels: None' + + Here's the dataset information: + {selected_dataset_info} + + Let’s Begin! + """ + return {"Expert": ai_expert_sys_prompt} + + +def generate_doctor_expert_sys_prompt_dict(expert_profile, dataset_name_list): + selected_dataset_info = get_selected_dataset_info(dataset_name_list) + doctor_sys_prompt = f"""" + You are a doctor who request a AI expert to select the correct tools (datasets and labels) to the segment requested targets on the medical image with specific modality. + After the expert's response, you will verify if the labels of the selected datasets meets your request, and give your feedback. + + The dataset informations are here: + {selected_dataset_info}. + + Your request: + """ + ai_expert_sys_prompt = generate_expert_sys_prompt_dict(expert_profile, dataset_name_list) + ai_expert_sys_prompt = ai_expert_sys_prompt["Expert"].replace( + "Let’s Begin!", + "When you are given the feedback, you will adjust your response and return it with the correct format. Remember, ALL suitable datasets and corresponding labels should be selected. \n Let’s Begin!" + ) + return { + "Doctor": doctor_sys_prompt, + "Expert": ai_expert_sys_prompt, + } + + +def verify_dataset_coverage(profile_to_datasets_mapping): + dataset_list = [] + for assigned_datasets in profile_to_datasets_mapping.values(): + dataset_list.extend(list(assigned_datasets)) + unique_values = set(dataset_list) + target_set = set(DATASET_NAME_TO_PATH.keys()) + + if unique_values == target_set: + print("The sets are the same.") + else: + missing_items = target_set - unique_values + extra_items = unique_values - target_set + + if missing_items: + print(f"Missing items: {missing_items}") + if extra_items: + print(f"Extra items: {extra_items}") + + for dataset_name in unique_values: + dataset_count = dataset_list.count(dataset_name) + if (dataset_count > 1): + print(f"The dataset '{dataset_name}' is included {dataset_count} times.") + + +DEPARTMENT_EXPERTS_INFO_OVERLAPPED = { + "An Expert in Abdominal Medical Image Segmentation in Gastroenterology Department.": + ["AMOS22_Task2", "FLARE23", "AbdomenAtlasMini"], + "An Expert in Kidney and Urological Medical Image Segmentation in Nephrology/Urology Department.": + ["KiTS23", "TotalSegmentator_v2", "CT-ORG"], + "An Expert in Thoracic Medical Image Segmentation in Pulmonology Department.": + ["ATM22", "Parse22", "TotalSegmentator_v2"], + "An Expert in Head and Neck Medical Image Segmentation in Otolaryngology Department.": + ["HECKTOR2022", "SegRap2023_Task1", "SegRap2023_Task2"], + "An Expert in Brain and Neurological Medical Image Segmentation in Neurology Department.": + ["BraTS21", "ISLES22", "TotalSegmentator_MRI"], + "An Expert in Cardiac Medical Image Segmentation in Cardiology Department.": + ["MM-WHS", "MMs", "TotalSegmentator_v2", "ImageCAS"], + "An Expert in Hepatic and Oncological Medical Image Segmentation in Hepatology/Oncology Department.": + ["LiTS", "AutoPET", "FLARE23", "Adrenal-ACC-Ki67-Seg"], + "An Expert in Stroke Lesion Segmentation in Emergency/Radiology Department.": + ["ISLES22", "ISLES22_ATLAS", "Instance22"], +} + +DEPARTMENT_EXPERTS_INFO = { + "An Expert in Full-Body CT Medical Image Segmentation in Radiology Department.": + ["CT-ORG", "TotalSegmentator_v2"], + "An Expert in Full-Body MRI Medical Image Segmentation in Radiology Department.": + ["TotalSegmentator_MRI"], + "An Expert in Abdominal Medical Image Segmentation in Gastroenterology Department.": + ["AMOS22_Task2", "FLARE23", "AbdomenAtlasMini", "KiTS23"], + "An Expert in Thoracic Medical Image Segmentation in Pulmonology Department.": + ["ATM22", "Parse22"], + "An Expert in Head and Neck Medical Image Segmentation in Otolaryngology Department.": + ["HECKTOR2022", "SegRap2023_Task1", "SegRap2023_Task2"], + "An Expert in Brain and Neurological Medical Image Segmentation in Neurology Department.": + ["BraTS21", "ISLES22", "ISLES22_ATLAS", "Instance22"], + "An Expert in Cardiac Medical Image Segmentation in Cardiology Department.": + ["MM-WHS", "MMs", "ImageCAS"], + "An Expert in Oncological Medical Image Segmentation in Oncology Department.": + ["LiTS", "AutoPET", "Adrenal-ACC-Ki67-Seg"], +} + +if __name__ == "__main__": + print( + generate_expert_sys_prompt_dict( + "An Expert in Abdominal Medical Image Segmentation in Gastroenterology Department.", + ["AMOS22_Task2"])) + + verify_dataset_coverage(DEPARTMENT_EXPERTS_INFO) diff --git a/prompts/assigner_experts_v2.py b/prompts/assigner_experts_v2.py new file mode 100644 index 0000000..940002a --- /dev/null +++ b/prompts/assigner_experts_v2.py @@ -0,0 +1,127 @@ +from prompts.data import DATASET_NAME_TO_PATH, get_info_from_json_path_list, DATASET_JSON_INFO +from prompts.react_2shot import SYSTEM_PROMPT_REACT +import json + + +# 加载专家配置文件 +def load_expert_config(config_path): + with open(config_path, 'r') as config_file: + return json.load(config_file) + + +# 加载配置数据 +config_data = load_expert_config('./prompts/expert_config.json') + +# 定义 ASSIGNER_PROMPT 所需的格式化输出和示例 +ASSIGNER_FORMATTED_OUTPUTS = '{"thought": "", "answer": ["WHOLEBODY", "EXPERT1", "EXPERT2"]}' + +ASSIGNER_FORMATTED_EXAMPLES = ''' +{"thought": "Liver segmentation requires both ABDOMINAL and WHOLEBODY expertise because the liver is an abdominal organ that may also appear in whole-body segmentation tasks.", "answer": ["WHOLEBODY", "ABDOMINAL"]}, +{"thought": "Heart segmentation requires CARDIAC expertise because the heart is the primary focus of cardiac specialists. THORACIC expertise is necessary because the heart is located within the thoracic cavity and may involve surrounding thoracic structures. WHOLEBODY expertise is relevant since the heart can appear in whole-body scans.", "answer": ["WHOLEBODY", "CARDIAC", "THORACIC"]}, +''' + + +# 获取选定数据集的信息 +def get_selected_dataset_info(dataset_name_list): + json_path_list = [DATASET_NAME_TO_PATH.get(dn, None) for dn in dataset_name_list] + if None in json_path_list: + missing_datasets = [ + dn for dn, path in zip(dataset_name_list, json_path_list) if path is None + ] + raise ValueError(f"Non-existent dataset names found: {missing_datasets}") + return get_info_from_json_path_list(json_path_list) + + +# 创建一个格式化的共同经验字符串 +common_experiences_section = "\n**Common Experiences**:\n" + "\n".join( + f"- {exp}" for exp in config_data['common_experiences']) + +# 更新 COMMON_EXPERT_PROMPT 以包含格式化的经验 +COMMON_EXPERT_PROMPT = SYSTEM_PROMPT_REACT + common_experiences_section + + +# 创建每个专家的 prompt，包括他们的经验 +def create_expert_prompt(dataset_name_list, experience_list, include_experiences=True): + dataset_info = get_selected_dataset_info(dataset_name_list) + if include_experiences: + # 包含共同经验和专家独有经验 + complete_prompt = COMMON_EXPERT_PROMPT.replace(DATASET_JSON_INFO, dataset_info) + if experience_list: + experience_section = "\n**Expert Experience**:\n" + "\n".join( + f"- {exp}" for exp in experience_list) + complete_prompt += experience_section + else: + # 不包含任何经验 + complete_prompt = SYSTEM_PROMPT_REACT.replace(DATASET_JSON_INFO, dataset_info) + return complete_prompt + + +# 根据配置数据生成每个专家的 prompt +expert_prompts_with_experience = {} +expert_prompts_no_experience = {} +for expert in config_data['experts']: + expert_name = expert['name'].upper() + dataset_list = expert['datasets'] + experience_list = expert.get('expert_experiences', []) + # 生成包含经验的 prompt + expert_prompts_with_experience[expert_name] = create_expert_prompt(dataset_list, + experience_list, + include_experiences=True) + # 生成不包含经验的 prompt + expert_prompts_no_experience[expert_name] = create_expert_prompt(dataset_list, + experience_list, + include_experiences=False) + +# 为 ASSIGNER_PROMPT 准备专家的 profile +expert_profiles = [] +for expert in config_data['experts']: + expert_name = expert['name'].upper() + profile = expert.get('profile', '') + if profile: + profile_text = f"**{expert_name}**: {profile}" + else: + profile_text = f"**{expert_name}**: No specific profile provided." + expert_profiles.append(profile_text) + +expert_profiles_text = "\n\n".join(expert_profiles) + +# 更新 ASSIGNER_PROMPT，引用配置中的专家 +expert_names = ', '.join([expert['name'].upper() for expert in config_data['experts']]) + +ASSIGNER_PROMPT = f""" +You are an AI expert assigner specializing in medical image segmentation tasks. Your responsibility is to select the most suitable specialist(s) based on anatomical regions and technical requirements. The specialists' areas of expertise include: {expert_names}. + +**Experts' Profiles**: +{expert_profiles_text} + +**Instructions**: + +1. **Analyze Task Requirements**: Identify the specified anatomical region(s) (e.g., HEAD, ABDOMINAL) and required segmentation techniques or imaging modalities (e.g., MRI, CT). **If the anatomical region is not specified, assume that the task could involve any relevant regions where the structures might appear, and select all applicable specialists.** + +2. **Match Expert Profiles**: Based on each expert's anatomical specialization, experience, and preferred modalities, select the most appropriate expert(s) for the task. **Always include WHOLEBODY expertise unless the task is exclusively limited to a specific region without any whole-body context.** + +3. **Provide Brief Justification**: Briefly explain your choice, referencing each expert's specialization and limitations. + +**Response Format**: +Use the following JSON format for Final Answer: +{ASSIGNER_FORMATTED_OUTPUTS} +Replace with your reasoning for selecting the expert(s). +Replace "WHOLEBODY", "EXPERT1", "EXPERT2" with the names of the assigned specialists. + +Here are some examples: +{ASSIGNER_FORMATTED_EXAMPLES} +""" + +# 使用包含经验的 expert_prompts_with_experience 字典 +HEAD_EXPERT_PROMPT = expert_prompts_with_experience.get("HEAD", "") +CARDIAC_EXPERT_PROMPT = expert_prompts_with_experience.get("CARDIAC", "") +ABDOMINAL_EXPERT_PROMPT = expert_prompts_with_experience.get("ABDOMINAL", "") +THORACIC_EXPERT_PROMPT = expert_prompts_with_experience.get("THORACIC", "") +WHOLEBODY_EXPERT_PROMPT = expert_prompts_with_experience.get("WHOLEBODY", "") + +# 使用不包含经验的 expert_prompts_no_experience 字典 +HEAD_EXPERT_PROMPT_NO_EXPERIENCE = expert_prompts_no_experience.get("HEAD", "") +CARDIAC_EXPERT_PROMPT_NO_EXPERIENCE = expert_prompts_no_experience.get("CARDIAC", "") +ABDOMINAL_EXPERT_PROMPT_NO_EXPERIENCE = expert_prompts_no_experience.get("ABDOMINAL", "") +THORACIC_EXPERT_PROMPT_NO_EXPERIENCE = expert_prompts_no_experience.get("THORACIC", "") +WHOLEBODY_EXPERT_PROMPT_NO_EXPERIENCE = expert_prompts_no_experience.get("WHOLEBODY", "") diff --git a/prompts/basic.py b/prompts/basic.py new file mode 100644 index 0000000..b950776 --- /dev/null +++ b/prompts/basic.py @@ -0,0 +1,46 @@ +from prompts.data import DATASET_JSON_INFO, MODEL_JSON_INFO + +SYSTEM_PROMPT_AI_EXPERT = f""" +You are an AI expert of medical image segmentation and good at dealing with clinical tasks. + +Remember: +1. You should ALWAYS think about what to do, but all the thought is short, at most in 3 sentences. +2. Only one dataset and one model should be chosen. Make sure the model trained on the chosen dataset exists. +3. Make sure all modalities in the selected dataset is provided. If some modalities are missing, respond with 'Missing modalities.' +4. If no suitable dataset is found, respond with 'Dataset: No suitable dataset is found.' + +Task description: +You should choose suitable AI model trained on suitable dataset for the real-world clinical needs. +Specifically, you have access to the following dataset: +{DATASET_JSON_INFO} + +Model introduction: +{MODEL_JSON_INFO} + +Desired format: +Thought: +Dataset: +Labels: + +You should reply in the format of this examples: +Thought: Since the user ask for brain tumor segmentation with low latency +Dataset: BraTS21 +Labels: [1, 2, 3] + +Let’s Begin! +""" + +SYSTEM_PROMPT_DATASET_ONLY = f""" +You are an AI expert of medical image segmentation and good at dealing with clinical tasks. + +Task description: +You have access to the following dataset: +{DATASET_JSON_INFO} +""" + +BASIC_SYS_PROMPT = { + "Expert": SYSTEM_PROMPT_AI_EXPERT, +} + +if __name__ == "__main__": + print(BASIC_SYS_PROMPT["Expert"]) diff --git a/prompts/basic_v2.py b/prompts/basic_v2.py new file mode 100644 index 0000000..1335147 --- /dev/null +++ b/prompts/basic_v2.py @@ -0,0 +1,65 @@ +from prompts.data import DATASET_JSON_INFO + +FORMATTED_OUTPUT = "{ \"thought\": \"\", \"dataset\": \"\", \"labels\": \"\" }" + +NATRUAL_EXAMPLES = """ +## Usage Example: +``` +thought: Please segment duodenum in CT images. +dataset: AMOS22_Task2, FLARE23, TotalSegmentator_v2 +labels: [13], [12], [19] +``` +``` +thought: Please segment adrenal gland in CT images. +dataset: AMOS22_Task2, TotalSegmentator_v2, FLARE23 +labels: [11,12], [8,9], [7,8] +``` +""" + +FORMATTED_EXAMPLES = """ +## Usage Example: +``` +{\"thought\": \"I need to find all CT datasets containing duodenum. I should first check the modality, then the labels.\",\"dataset\": \"AMOS22_Task2, FLARE23, TotalSegmentator_v2\",\"labels\": \"[13], [12], [19]\",} +``` +``` +{\"thought\": \"To segment adrenal gland in CT, I need to first find all dataset with adrenal gland. Then I need to find \"\"dataset\": \"AMOS22_Task2, FLARE23, TotalSegmentator_v2\"\"labels\": \"[11,12], [8,9], [7,8]\"} +``` +""" + +SYSTEM_PROMPT_AI_EXPERT_V2 = f""" +# Role: +You are an AI expert on medical image segmentation, specializing in addressing clinical tasks efficiently. + +## Skills: +1. Thought Process Optimization +- Quickly analyze the specific clinical need and the segmentation targets. +- Choose all suitable datasets and corresponding labels. + +2. Dataset & Label Selection +- Use the provided dataset information to ensure compatibility. +- Validate the selected labels can cover the need to segment the requested targets. + +## Constraints: +1. Thought Process Constraints +- Thoughts must be concise and short (within 3 sentences). +- All suitable datasets and corresponding labels should be selected. + +2. Response Format +- Use the following JSON format: +``` +{FORMATTED_OUTPUT} +``` + +3. Handling Missing Modalities +- If fail to find any needed modalities, reply in JSON with: 'Missing modalities.' ('MR' is equal to 'MRI') +- If no suitable dataset is available, reply in JSON: 'Dataset: None\nLabels: None' + +Here's the dataset information: +{DATASET_JSON_INFO} + +Let’s Begin! +""" + +BASIC_V2_SYS_PROMPT = { + "Expert": SYSTEM_PROMPT_AI_EXPERT_V2, +} diff --git a/prompts/data.py b/prompts/data.py new file mode 100644 index 0000000..14b58c7 --- /dev/null +++ b/prompts/data.py @@ -0,0 +1,122 @@ +import json +import os.path as osp +from glob import glob +from collections import defaultdict + + +def remove_useless_keys(json_content): + json_content.pop("models", None) + json_content["labels"].pop("background", None) + return json_content + + +def get_info_from_json_dir(json_dir, json_filter=lambda x: x): + json_path_list = glob(osp.join(json_dir, "*.json")) + + json_info_list = [ + json.dumps(json_filter(json.load(open(json_path, encoding='utf-8')))) + for json_path in json_path_list + ] + json_info = "\n".join(json_info_list) + return json_info + + +def get_info_from_json_path_list(json_path_list, json_filter=remove_useless_keys): + json_info_list = [ + json.dumps(json_filter(json.load(open(json_path, encoding='utf-8')))) + for json_path in json_path_list + ] + json_info = "\n".join(json_info_list) + return json_info + + +PROJECT_DIR = osp.dirname(osp.dirname(__file__)) +DATASET_DIR = osp.join(PROJECT_DIR, "dataset") +DATASET_JSON_INFO = get_info_from_json_dir(DATASET_DIR, json_filter=remove_useless_keys) +MODEL_JSON_INFO = get_info_from_json_dir(osp.join(PROJECT_DIR, "dataset", "model")) + +DATASET_NAME_TO_PATH = dict() +DATASET_NAME_TO_LABEL_INFO = dict() +DATASET_NAME_TO_FILTERED_INFO = dict() +for json_path in glob(osp.join(DATASET_DIR, "*.json")): + dataset_info = json.load(open(json_path, encoding='utf-8')) + dataset_name = dataset_info['name'] + DATASET_NAME_TO_PATH[dataset_name] = json_path + # dataset_label_to_name = dict() + # for label_name, label in dataset_info["labels"].items(): + # dataset_label_to_name[str(label)] = label_name + # DATASET_NAME_TO_INFO[dataset_name] = dataset_label_to_name + DATASET_NAME_TO_LABEL_INFO[dataset_name] = dataset_info["labels"] + DATASET_NAME_TO_FILTERED_INFO[dataset_name] = remove_useless_keys(dataset_info) + +LABEL_NAME_TO_DATASET = defaultdict(list) +ALL_LABEL_LIST = set() +ALL_DATASET_LIST = set() +for dataset, info in DATASET_NAME_TO_FILTERED_INFO.items(): + ALL_DATASET_LIST.add(dataset) + label_list = list(info["labels"].keys()) + for label in label_list: + LABEL_NAME_TO_DATASET[label].append(dataset) + ALL_LABEL_LIST.add(label) +ALL_LABEL_LIST = sorted(list(ALL_LABEL_LIST)) +ALL_DATASET_LIST = sorted(list(ALL_DATASET_LIST)) + + +def retrieve_dataset_label_info(dataset_name, label_list): + result = "" + dataset_info = DATASET_NAME_TO_LABEL_INFO.get(dataset_name, dict()) + + for label in label_list: + label_found = False + for label_name, label_values in dataset_info.items(): + if (isinstance(label_values, int) and label == label_values): + result += f"{label}({label_name})," + label_found = True + break + if (isinstance(label_values, list) and label_list == label_values): + result += f"{label}({label_name})," + label_found = True + break + if not label_found: + result += f"{label}(Not Found)," + + if result == "": + return "No Label Found " + + return result[:-1] # Remove the trailing comma + + +def retrieve_dataset_from_label_name(target_label, mode="label", top_k=5, backend="openai"): + from utils.embed_utils import calculate_similarities + retrieved_datasets = set() + if mode == "label": + top_k_labels = calculate_similarities(ALL_LABEL_LIST, target_label, + backend=backend)[:top_k] + for label, _ in top_k_labels: + for dataset in LABEL_NAME_TO_DATASET[label]: + retrieved_datasets.add(dataset) + elif mode == "dataset" and backend == "openai": + top_k_datasets = calculate_similarities(ALL_DATASET_LIST, target_label, + backend=backend)[:top_k] + for dataset, _ in top_k_datasets: + retrieved_datasets.add(dataset) + else: + raise NotImplementedError() + + if not retrieved_datasets: + return "No Dataset Found" + + results = "\n".join( + [str(DATASET_NAME_TO_FILTERED_INFO[dataset]) for dataset in retrieved_datasets]) + return results + + +if __name__ == "__main__": + # print(retrieve_dataset_label_info("KiTS23", [2, 3])) + # print(retrieve_dataset_label_info("AMOS22_Task2", [1])) + # print(retrieve_dataset_label_info("FLARE23", [8])) + # print(retrieve_dataset_label_info("TotalSegmentator_v2", [9])) + # print(retrieve_dataset_from_label_name("brain", mode="label", backend="openai", top_k=1000)) + print(retrieve_dataset_from_label_name("brain", mode="label", backend="spacy", top_k=1000)) + # print(retrieve_dataset_from_label_name("brain", mode="dataset", backend="openai", top_k=1000)) + # print(retrieve_dataset_from_label_name("brain", top_k=1000)) diff --git a/prompts/doctor_expert.py b/prompts/doctor_expert.py new file mode 100644 index 0000000..cd623c9 --- /dev/null +++ b/prompts/doctor_expert.py @@ -0,0 +1,25 @@ +from prompts.data import DATASET_JSON_INFO +from prompts.basic_v2 import SYSTEM_PROMPT_AI_EXPERT_V2 + +DOCTOR_SYS_PROMPT = f"""" +You are a doctor who request a AI expert to select the correct tools (datasets and labels) to the segment requested targets on the medical image with specific modality. +After the expert's response, you will verify if the labels of the selected datasets meets your request, and give your feedback. + +The dataset informations are here: +{DATASET_JSON_INFO}. + +Your request: +""" + +EXPERT_SYS_PROMPT = SYSTEM_PROMPT_AI_EXPERT_V2.replace( + "Let’s Begin!", + "When you are given the feedback, you will adjust your response and return it with the correct format. Remember, ALL suitable datasets and corresponding labels should be selected. \n Let’s Begin!" +) + +DOCTOR_EXPERT_SYS_PROMPT = { + "Doctor": DOCTOR_SYS_PROMPT, + "Expert": EXPERT_SYS_PROMPT, +} + +if __name__ == "__main__": + print(DOCTOR_EXPERT_SYS_PROMPT) diff --git a/prompts/expert_config.json b/prompts/expert_config.json new file mode 100644 index 0000000..81f4016 --- /dev/null +++ b/prompts/expert_config.json @@ -0,0 +1,49 @@ +{ + "common_experiences": [ + "Notice that when asked to segment the left or right part (e.g., kidney_right), do not provide the entire target (kidney) but rather the specific part (kidney_right).", + "Ensure that only labels mentioned in the task are included, excluding any unnecessary or unrelated labels even similar or related." + ], + + "experts": [ + { + "name": "HEAD", + "datasets": ["BraTS21", "ISLES22", "ISLES22_ATLAS", "Instance22", "HECKTOR2022", "SegRap2023_Task1", "SegRap2023_Task2"], + "profile": "Specializes in head and neck imaging, including organs, lesions, tumors, and lymph nodes in this region.", + "expert_experiences": [ + ] + }, + { + "name": "CARDIAC", + "datasets": ["MM-WHS", "MMs", "ImageCAS"], + "profile": "Expert in cardiac imaging, including cardiac chambers, vessels, and related structures like pulmonary artery", + "expert_experiences": [ + "When segmenting the heart, leverage all available related labels in the dataset to provide comprehensive context." + ] + }, + { + "name": "ABDOMINAL", + "datasets": ["AMOS22_Task2", "FLARE22", "AbdomenAtlasMini", "KiTS23", "Adrenal-ACC-Ki67-Seg", "LiTS", "WORD"], + "profile": "Specializes in abdominal organ segmentation, such as liver, kidney, and adrenal glands. ", + "expert_experiences": [ + "The KiTS23 dataset contain kidneys segmentation but you should choose the right label value", + "Masses and cysts are distinct entities." + ] + }, + { + "name": "THORACIC", + "datasets": ["ATM22", "Parse22"], + "profile": "Focuses on thoracic imaging, including lungs and mediastinal organs. ", + "expert_experiences": [ + + ] + }, + { + "name": "WHOLEBODY", + "datasets": ["CT-ORG", "TotalSegmentator_v2", "TotalSegmentator_MRI", "AutoPET"], + "profile": "Proficient in full-body scan analysis, covering multi-organ segmentation and whole-body tumor segmentation.", + "expert_experiences": [ + "Masses and cysts are distinct entities." + ] + } + ] +} \ No newline at end of file diff --git a/prompts/native_2shot.py b/prompts/native_2shot.py new file mode 100644 index 0000000..029bf2d --- /dev/null +++ b/prompts/native_2shot.py @@ -0,0 +1,39 @@ +from prompts.data import DATASET_JSON_INFO + +FORMATTED_OUTPUT = '{ "thought":"", "answer":{"AbdomenAtlasMini": [1], "AMOS22_Task2": [2,3], "FLARE23": [15]} }' + +FORMATTED_EXAMPLES = '{ "thought":"I need to find all CT datasets containing duodenum. I should first check the modality, then the labels.", "answer":{"AbdomenAtlasMini":[13], "AMOS22_Task2":[12], "FLARE23":[19]} }\n{ "thought":"To segment adrenal gland in CT, I need to first find all dataset with adrenal gland.", "answer":{"AMOS22_Task2":[11,12], "FLARE23":[8,9], "TotalSegmentator_v2":[7,8]} }' + +MISSING_MODALITY = '{ "thought":"No suitable modality", "answer": {} }' +NO_DATASET_FOUND = '{ "thought":"No suitable dataset", "answer": {} }' + +SYSTEM_PROMPT_NATIVE = f""" +# Role: +You are an AI expert on medical image segmentation, specializing in addressing clinical tasks efficiently. You will answer the input question as best you can. + +## Skills: +1. Dataset & Label Selection +- Choose datasets and corresponding labels based on the input question. +- Use the provided dataset information to ensure compatibility. +- Validate the selected labels can cover the need to segment the requested targets. + +## Constraints: +1. Response Format +- MUST use the following JSON format for Final Answer: +Final Answer: {FORMATTED_OUTPUT} +Here are some examples: +{FORMATTED_EXAMPLES} + +2. Handling Missing Modalities +- If fail to find any needed modalities, reply in JSON: {MISSING_MODALITY} +- If no suitable dataset is available, reply in JSON: {NO_DATASET_FOUND} + +Here's the dataset information: +{DATASET_JSON_INFO} + +Let’s Begin! +""" + +NATIVE_SYS_PROMPT = { + "Expert": SYSTEM_PROMPT_NATIVE, +} diff --git a/prompts/rag.py b/prompts/rag.py new file mode 100644 index 0000000..cf69af3 --- /dev/null +++ b/prompts/rag.py @@ -0,0 +1,12 @@ +RAG_HELPER_SYS_PROMPT = ''' +You are a helpful assitant. Please only give me one keyword (can be more than 1 word) to retrieve for related target. + +For example: +Question: Please segment the kidney masses in CT images. +Keyword: kidney masses + +Question: Please segment the spleen in MRI images. +Keyword: spleen + +Please follow the above exmaples to give out clear keyword. +''' \ No newline at end of file diff --git a/prompts/react_2shot.py b/prompts/react_2shot.py new file mode 100644 index 0000000..0e8247e --- /dev/null +++ b/prompts/react_2shot.py @@ -0,0 +1,56 @@ +from prompts.data import DATASET_JSON_INFO + +FORMATTED_OUTPUT = '{ "thought":"", "answer":{"AbdomenAtlasMini": [1], "AMOS22_Task2": [2,3], "FLARE23": [15]} }' + +FORMATTED_EXAMPLES = '{ "thought":"I need to find all CT datasets containing duodenum. I should first check the modality, then the labels.", "answer":{"AbdomenAtlasMini":[13], "AMOS22_Task2":[12], "FLARE23":[19]} }\n{ "thought":"To segment adrenal gland in CT, I need to first find all dataset with adrenal gland.", "answer":{"AMOS22_Task2":[11,12], "FLARE23":[8,9], "TotalSegmentator_v2":[7,8]} }' + +MISSING_MODALITY = '{ "thought":"No suitable modality", "answer": {} }' +NO_DATASET_FOUND = '{ "thought":"No suitable dataset", "answer": {} }' + +SYSTEM_PROMPT_REACT = f""" +# Role: +You are an AI expert on medical image segmentation, specializing in addressing clinical tasks efficiently. You will answer the input question as best you can. + +## Skills: +1. Dataset & Label Selection +- Choose datasets and corresponding labels based on the input question. +- Use the provided dataset information to ensure compatibility. +- Validate the selected labels can cover the need to segment the requested targets. + +2. Thought Process Optimization +You MUST use this format to think which datasets and their labels should be chosen: +``` +Question: the input question you must answer +Thought: you should always think about what to do +Action: the action to take +Action Input: the input to the action +Observation: the result of the action +... (this process can repeat multiple times) +Thought: I now know the final answer +Final Answer: the final answer to the original input question in JSON format (Looking for examples below) +``` + +## Constraints: +1. Thought Process Constraints +- Thoughts must be concise and short (within 3 sentences). +- All suitable datasets and corresponding labels should be selected. + +2. Response Format +- Use the following JSON format for Final Answer: +{FORMATTED_OUTPUT} +Here are some examples: +{FORMATTED_EXAMPLES} + +3. Handling Missing Modalities +- If fail to find any needed modalities, reply in JSON: {MISSING_MODALITY} +- If no suitable dataset is available, reply in JSON: {NO_DATASET_FOUND} + +Here's the dataset information: +{DATASET_JSON_INFO} + +Let’s Begin! +""" + +REACT_SYS_PROMPT = { + "Expert": SYSTEM_PROMPT_REACT, +} diff --git a/prompts/reflexion_2shot.py b/prompts/reflexion_2shot.py new file mode 100644 index 0000000..c2d4f96 --- /dev/null +++ b/prompts/reflexion_2shot.py @@ -0,0 +1,34 @@ +from prompts.react_2shot import SYSTEM_PROMPT_REACT + +EVALUATOR_SYS_PROMPT = """ You are an AI expert specializing in evaluating answers related to the use of AI agents in medical tools. +You will be provided with an answer to a question. You will then think and provide an evaluation of the answer. +Provide the evaluation in the following format: +``` +Question: The input question you must answer +Previous Answer: The input previous answer for the question (only Final Answer) +Thought: You should always reflect and think about how to better answer the question +Final Evaluation: The final evaluation of the previous answer for the input question +Final Result: True or False (should be a bool) +``` +""" + +REFLECTION_SYS_PROMPT = """You are an AI expert specializing in self-reflection and give suggestion. +You will be given the previous answer and its evaluation results. +You will then provide one short suggestion on how to change the answer. +If there's no need to change, you should suggest "not need to change". +DON'T SUGGEST TO EXPLAIN ANY THING. ONLY SUGGEST ON HOW TO CHANGE. +Provide the suggestion in the following format: +``` +Question: the input question you must answer +Previous Answer: the input answer for the question (only Final Answer) +Previous Evaluation: the input evaluation for the answer (Final Evaluation & Final Result) +Thought: you should always reflect and think about how to change the answer +Final Suggestion: the final suggestion to change the previous answer for the input question (e.g. `Final Suggestion: The dataset 'CT-ORG' and its labels should be removed.`, `Final Suggestion: The label [4] of dataset 'CT-ORG' should be changed to [5].`) +``` +""" + +REFLECTXION_SYS_PROMPT = { + "Expert": SYSTEM_PROMPT_REACT, + "Evaluator": EVALUATOR_SYS_PROMPT, + "Reflection": REFLECTION_SYS_PROMPT, +} diff --git a/prompts/verifier.py b/prompts/verifier.py new file mode 100644 index 0000000..e59a12c --- /dev/null +++ b/prompts/verifier.py @@ -0,0 +1,21 @@ +from prompts.data import DATASET_JSON_INFO + +VERIFIER_SYS_PROMPT = f""" +# Role: +You are a careful verifier for LLM answer, specializing in checking the correctness of the answer based on the reference information. You will verify every detail of the answer and give a concise conclusion as best you can. + +## Skills: +- Carefully read the question, LLM answer, and reference information. Then you will compare the question and verify whether the LLM answer is correct based on the reference information. + +## Constraints: +- Notice that when asked to segment the left or right part (e.g., kidney_right), do not provide the entire target (kidney) but rather the specific part (kidney_right). +- You MUST verify in the following format: +Question: The input question you must answer +LLM Answer: The input previous answer for the question (only Final Answer) +Reference Information: The input reference information for the question and answer +Thought: You should always reflect and think about how to better answer the question +Final Conclusion: True or False (should be a bool) +Final Suggestion: (if the answer is False, you should provide a better answer based on the reference information, otherwise, you should return the original answer.) + +Let’s Begin! +""" \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..ba82164 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +pyautogen==0.2.35 +spacy==3.8.7 +tenacity==9.1.2 +colorlog +nnunetv2 \ No newline at end of file diff --git a/results/C2F_model_selection_test_case_gpt-4o-2024-08-06_20250604_103650.json b/results/C2F_model_selection_test_case_gpt-4o-2024-08-06_20250604_103650.json new file mode 100644 index 0000000..437512f --- /dev/null +++ b/results/C2F_model_selection_test_case_gpt-4o-2024-08-06_20250604_103650.json @@ -0,0 +1,24 @@ +{ + "1": { + "AbdomenAtlasMini": [ + 8 + ], + "WORD": [ + 2 + ], + "FLARE22": [ + 3 + ], + "TotalSegmentator_v2": [ + 1 + ], + "AMOS22_Task2": [ + 1 + ] + }, + "2": { + "HECKTOR2022": [ + 1 + ] + } +} \ No newline at end of file diff --git a/results/C2F_model_selection_test_case_gpt-4o-2024-08-06_20250604_103650.log b/results/C2F_model_selection_test_case_gpt-4o-2024-08-06_20250604_103650.log new file mode 100644 index 0000000..3c45052 --- /dev/null +++ b/results/C2F_model_selection_test_case_gpt-4o-2024-08-06_20250604_103650.log @@ -0,0 +1,64 @@ +2025-06-04 10:36:50,690 - INFO - 步骤 1: 查询解析代理 +2025-06-04 10:36:53,094 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:36:53,103 - INFO - Query parsed successfully: {'modality': 'CT', 'target': 'spleen'} +2025-06-04 10:36:53,103 - INFO - 步骤 2: 对数据集针对模态进行过滤 +2025-06-04 10:36:53,103 - INFO - Modality filtering iteration 1 +2025-06-04 10:36:57,977 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:36:57,982 - INFO - Modality filtering iteration 2 +2025-06-04 10:37:00,519 - INFO - Retrying request to /chat/completions in 0.378578 seconds +2025-06-04 10:37:05,826 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:37:05,834 - INFO - Final selected datasets after modality filtering: +2025-06-04 10:37:05,834 - INFO - ['AMOS22_Task2', 'ATM22', 'AbdomenAtlasMini', 'Adrenal-ACC-Ki67-Seg', 'AutoPET', 'CT-ORG', 'FLARE22', 'FUMPE', 'HECKTOR2022', 'ImageCAS', 'Instance22', 'KiTS23', 'LNQ2023', 'LiTS', 'MM-WHS', 'Parse22', 'SegRap2023_Task1', 'SegRap2023_Task2', 'TotalSegmentator_v2', 'WORD'] +2025-06-04 10:37:05,834 - INFO - 步骤 3: 对数据集针对类别进行过滤 +2025-06-04 10:37:05,834 - INFO - Relevance filtering iteration 1 +2025-06-04 10:37:11,087 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:37:11,093 - INFO - Relevance filtering iteration 2 +2025-06-04 10:37:14,889 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:37:14,895 - INFO - Final selected datasets after relevance filtering: +2025-06-04 10:37:14,895 - INFO - ['AMOS22_Task2', 'AbdomenAtlasMini', 'FLARE22', 'TotalSegmentator_v2', 'WORD'] +2025-06-04 10:37:14,895 - INFO - 步骤 4: 选择数据集和类别 +2025-06-04 10:37:14,895 - INFO - Label selection iteration 1 +2025-06-04 10:37:18,120 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:37:18,126 - INFO - 最终选择的标签: +2025-06-04 10:37:18,126 - INFO - {'AbdomenAtlasMini': [8], 'WORD': [2], 'FLARE22': [3], 'TotalSegmentator_v2': [1], 'AMOS22_Task2': [1]} +2025-06-04 10:37:18,126 - INFO - [1] Chat History: + +2025-06-04 10:37:18,126 - INFO - [1] Correct! (hit_rate: 1.0, final: 1.0) +Expected answer: +{'TotalSegmentator_v2': [1], 'FLARE22': [3], 'AbdomenAtlasMini': [8], 'AMOS22_Task2': [1], 'WORD': [2]} +LLM: +{'AbdomenAtlasMini': [8], 'WORD': [2], 'FLARE22': [3], 'TotalSegmentator_v2': [1], 'AMOS22_Task2': [1]} +2025-06-04 10:37:18,127 - INFO - 步骤 1: 查询解析代理 +2025-06-04 10:37:19,670 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:37:19,677 - INFO - Query parsed successfully: {'modality': 'PET', 'target': 'primary tumor (GTVp)'} +2025-06-04 10:37:19,677 - INFO - 步骤 2: 对数据集针对模态进行过滤 +2025-06-04 10:37:19,677 - INFO - Modality filtering iteration 1 +2025-06-04 10:37:21,135 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:37:21,144 - INFO - Modality filtering iteration 2 +2025-06-04 10:37:23,130 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:37:23,140 - INFO - Final selected datasets after modality filtering: +2025-06-04 10:37:23,140 - INFO - ['AutoPET', 'HECKTOR2022'] +2025-06-04 10:37:23,140 - INFO - 步骤 3: 对数据集针对类别进行过滤 +2025-06-04 10:37:23,140 - INFO - Relevance filtering iteration 1 +2025-06-04 10:37:27,489 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:37:27,495 - INFO - Relevance filtering iteration 2 +2025-06-04 10:37:31,924 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:37:31,930 - INFO - Final selected datasets after relevance filtering: +2025-06-04 10:37:31,930 - INFO - ['HECKTOR2022'] +2025-06-04 10:37:31,930 - INFO - 步骤 4: 选择数据集和类别 +2025-06-04 10:37:31,930 - INFO - Label selection iteration 1 +2025-06-04 10:37:35,821 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" +2025-06-04 10:37:35,826 - INFO - 最终选择的标签: +2025-06-04 10:37:35,827 - INFO - {'HECKTOR2022': [1]} +2025-06-04 10:37:35,827 - INFO - [2] Chat History: + +2025-06-04 10:37:35,827 - INFO - [2] Correct! (hit_rate: 1.0, final: 2.0) +Expected answer: +{'HECKTOR2022': [1]} +LLM: +{'HECKTOR2022': [1]} +2025-06-04 10:37:35,827 - INFO - [Summary] [Datasets Hit Rate] final score 2.00 / 2.0 +2025-06-04 10:37:35,827 - INFO - [Summary] [All Dataset Found] final score 2.00 / 2.0 +2025-06-04 10:37:35,827 - INFO - [Summary] [Any Dataset Found] final score 2.00 / 2.0 +2025-06-04 10:37:35,827 - INFO - [Results] result saved to results/C2F_model_selection_test_case_gpt-4o-2024-08-06_20250604_103650.json +2025-06-04 10:37:35,827 - INFO - [Results] clean state diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/utils/__init__.py diff --git a/utils/chat_utils.py b/utils/chat_utils.py new file mode 100644 index 0000000..556cabc --- /dev/null +++ b/utils/chat_utils.py @@ -0,0 +1,74 @@ +import os.path as osp +from openai import OpenAI +from autogen import config_list_from_json +from tenacity import ( + retry, + stop_after_attempt, + wait_fixed, +) +from agents import * + + +AUTOGEN_CONFIG_PATH = osp.join(osp.dirname(osp.dirname(__file__)), "OAI_CONFIG_LIST") + + +def get_OAI_client_for_model(model, autogen_config_path=AUTOGEN_CONFIG_PATH): + config = config_list_from_json(env_or_file=autogen_config_path, + filter_dict={"model": [model, "default"]})[0] + client = OpenAI(api_key=config['api_key'], base_url=config['base_url']) + return client + + +def get_model_list(): + all_available_models = [] + for possible_model in ["default", "gpt-4o"]: + client = get_OAI_client_for_model(possible_model) + model_list = client.models.list() + model_list = [model.id for model in model_list] + all_available_models += model_list + return all_available_models + + +def chat_with_llm(pattern, query, model="gpt-4o"): + if pattern == "Native": + res = chat_with_llm_native(query=query, + model=model, + autogen_config_path=AUTOGEN_CONFIG_PATH) + elif pattern == "CoT": + res = chat_with_llm_cot(query=query, model=model, autogen_config_path=AUTOGEN_CONFIG_PATH) + elif pattern == "ReAct": + res = chat_with_llm_react(query=query, + model=model, + autogen_config_path=AUTOGEN_CONFIG_PATH) + elif pattern == "ReAct_RAG": + res = chat_with_llm_react_rag(query=query, + model=model, + autogen_config_path=AUTOGEN_CONFIG_PATH) + + elif pattern == "ReAct_SelfVerify": + res = chat_with_llm_react_selfverify(query=query, + model=model, + autogen_config_path=AUTOGEN_CONFIG_PATH) + + elif pattern == "ReAct_MultiExpert": + res = chat_with_llm_react_multiexpert(query=query, + model=model, + autogen_config_path=AUTOGEN_CONFIG_PATH) + elif pattern == "ReAct_MultiExpert_Experience": + res = chat_with_llm_react_multiexpert_experience(query=query, + model=model, + autogen_config_path=AUTOGEN_CONFIG_PATH) + elif pattern == "ReAct_MultiExpert_SelfVerify": + res = chat_with_llm_react_multiexpert_selfverify(query=query, + model=model, + autogen_config_path=AUTOGEN_CONFIG_PATH) + elif pattern == "ReAct_MultiExpert_SelfVerify_Experience": + res = chat_with_llm_react_multiexpert_selfverify_experience( + query=query, model=model, autogen_config_path=AUTOGEN_CONFIG_PATH) + + elif pattern == "C2F": + res = chat_with_llm_c2f(query=query, model=model, autogen_config_path=AUTOGEN_CONFIG_PATH) + + else: + raise NotImplementedError(f"No pattern named '{pattern}'") + return res diff --git a/utils/embed_utils.py b/utils/embed_utils.py new file mode 100644 index 0000000..1710ebe --- /dev/null +++ b/utils/embed_utils.py @@ -0,0 +1,72 @@ +import spacy +import numpy as np +from openai import OpenAI +from autogen import config_list_from_json +import os.path as osp +import os +import json + +nlp = spacy.load("en_core_web_sm") +PROJECT_DIR = osp.dirname(osp.dirname(__file__)) +AUTOGEN_CONFIG_PATH = osp.join(PROJECT_DIR, "OAI_CONFIG_LIST") +openai_embed_store_file = osp.join(PROJECT_DIR, "ckpt", "openai_embed_store.json") +OPENAI_EMBED_STORE = dict() +if (osp.exists(openai_embed_store_file)): + OPENAI_EMBED_STORE = json.load(open(openai_embed_store_file, "r", encoding='utf-8')) + +config = config_list_from_json(env_or_file=AUTOGEN_CONFIG_PATH, filter_dict={"model": + ["gpt-4"]})[0] +client = OpenAI(api_key=config['api_key'], base_url=config['base_url']) + + +def get_text_embedding(text, model="text-embedding-ada-002"): + if (text in OPENAI_EMBED_STORE): + return OPENAI_EMBED_STORE[text] + print(f"recompute since no cache found for: {text}") + text = text.replace("\n", " ") + response = client.embeddings.create(input=text, model=model) + embed = response.data[0].embedding + return np.array(embed) + + +def calculate_similarities(label_list, target_label, backend="openai"): + """ + 计算目标标签与标签列表中每个标签的语义相似度，并返回相似度列表。 + + :param label_list: 包含标签的字符串列表 + :param target_label: 目标标签字符串 + :return: 包含 (标签, 相似度) 元组的列表 + """ + if backend == "spacy": + embedding_func = lambda x: nlp(x).vector + elif backend == "openai": + embedding_func = get_text_embedding + else: + raise NotImplementedError(f"{backend} is not supported for embeddon") + + similarities = [] + + target_vector = embedding_func(target_label) + for label in label_list: + label_vector = embedding_func(label) + + similarity = np.dot(target_vector, label_vector) / (np.linalg.norm(target_vector) * + np.linalg.norm(label_vector)) + + similarities.append((label, similarity)) + + return sorted(similarities, key=lambda x: x[1], reverse=True) + + +if __name__ == "__main__": + categories = [ + "kidney", "kidney_right", "kidney_left", "right kidney", "liver", "heart", "lung" + ] + target = "left kidney" + + for backend in ["spacy", "openai"]: + print("backend:", backend) + similarities = calculate_similarities(categories, target, backend=backend) + for category, similarity in similarities: + print(f"Category: {category}, Similarity: {similarity}") + print("----------------------------------------------------------") diff --git a/utils/metric.py b/utils/metric.py new file mode 100644 index 0000000..f951da2 --- /dev/null +++ b/utils/metric.py @@ -0,0 +1,64 @@ +import numpy as np + +# expected data format: +# {"dataset1": [1,2,3], "dataset2": [1]} +''' How many percentage of datasets are found''' +def dataset_hit_rate(gt_answer, llm_answer): + try: + hit_count = 0 + total_datasets = len(set(gt_answer.keys()) | set(llm_answer.keys())) + + for dataset, gt_labels in gt_answer.items(): + if dataset in llm_answer: + val = llm_answer[dataset] + # 若存在子列表，将其一层展平 + if isinstance(val, list) and any(isinstance(x, list) for x in val): + val = [elem for sub in val for elem in (sub if isinstance(sub, list) else [sub])] + + # 进行集合比较 + if set(gt_labels) == set(val): + hit_count += 1 + + return hit_count / total_datasets + except Exception as e: + print(f"Error occurred: {e}") + return 0 + + +''' If any dataset is found successfully, return 1, else 0''' + + +def any_dataset_found_rate(gt_answer, llm_answer): + hr = dataset_hit_rate(gt_answer, llm_answer) + return 1 if (hr > 0) else 0 + + +''' If all datasets are found, return 1, else 0''' + + +def all_dataset_found_rate(gt_answer, llm_answer): + hr = dataset_hit_rate(gt_answer, llm_answer) + return 1 if (hr == 1) else 0 + + +if __name__ == "__main__": + gt = { + 'TotalSegmentator_v2': [2, 3], + 'FLARE23': [2, 13], + 'AbdomenAtlasMini': [3, 4], + 'AMOS22_Task2': [2, 3], + 'KiTS23': [1], + 'CT-ORG': [4] + } + pred = { + 'AbdomenAtlasMini': [3, 4], + 'CT-ORG': [4], + 'KiTS23': [1, 2, 3], + 'LiTS': [1, 2], + 'TotalSegmentator_v2': [2, 3] + } + + gt = {'SegRap2023 Task1': [40, 41]} + pred = {'SegRap2023_Task1': [40, 41]} + res = dataset_hit_rate(gt, pred) + print(res) diff --git a/utils/reply_utils.py b/utils/reply_utils.py new file mode 100644 index 0000000..05efc77 --- /dev/null +++ b/utils/reply_utils.py @@ -0,0 +1,106 @@ +import json +import logging +import os +import re + + +def save_state(state, state_filename='task_state.json'): + with open(state_filename, 'w') as f: + json.dump(state, f) + + +def load_state(state_filename='task_state.json'): + if os.path.exists(state_filename): + with open(state_filename, 'r') as f: + return json.load(f) + return None + + +def clean_dataset_format(input_str): + input_str = input_str.replace("'", "").replace("\"", "") + if input_str.startswith("[") and input_str.endswith("]"): + input_str = input_str[1:-1] + return input_str + + +def clean_label_format(input_str): + input_str = input_str.replace("'", "").replace("\"", "").strip() + + if input_str.startswith("'") and input_str.endswith("'"): + input_str = input_str[1:-1] + + if input_str.startswith('"') and input_str.endswith('"'): + input_str = input_str[1:-1] + + if input_str.startswith("[[") and input_str.endswith("]]"): + input_str = input_str[1:-1] + + return input_str + + +def extract_llm_answer(input_str, key="Final Answer:"): + # 1. 移除代码块标记 + input_str = input_str.replace("```json\n", "").replace("```", "").strip() + + if key in input_str: + # 2. 分割字符串并获取关键字后的部分 + input_str = input_str.split(key)[-1].strip() + + # 替换 JSON 键的单引号为双引号 + input_str = re.sub(r"(?<=\{)\s*'(\w+)'\s*:", r'"\1":', input_str) + input_str = re.sub(r",\s*'(\w+)'\s*:", r', "\1":', input_str) + + # 替换第一个值中的首尾单引号为双引号 + input_str = re.sub(r'"\s*:\s*\'(.*?)\'\s*(,|\})', r'": "\1"\2', input_str) + + # 替换 "answer" 字段中的所有单引号为双引号 + if '"answer":' in input_str: + answer_match = re.search(r'("answer"\s*:\s*{.*?})', input_str, re.DOTALL) + if answer_match: + answer_content = answer_match.group(1) + # 将内容中的单引号替换为双引号 + answer_content = answer_content.replace("'", '"') + input_str = input_str.replace(answer_match.group(1), answer_content) + + # 3. 使用正则表达式提取符合 {"thought": "", "answer": {}} 格式的 JSON 字符串 + json_pattern = r'\{\s*"thought"\s*:\s*"[^"]*",\s*"answer"\s*:\s*(\{.*?\}|\[.*?\]|".*?"|\d+)\s*\}' + match = re.search(json_pattern, input_str, re.DOTALL) + + if match: + json_str = match.group(0) + else: + # 如果未匹配到预期的 JSON 格式，使用整个 input_str + json_str = input_str + + result = {} + try: + result = json.loads(json_str) + except json.JSONDecodeError: + logging.warning(f"无法将以下内容解析为 JSON: {json_str}") + return {} + + # 4. 获取 'answer' 部分并返回整个字典 + answer = result.get('answer', {}) + + return answer + + +def extract_ground_truth(test_case): + expected_dataset_info = test_case.get('expected_datasets', {"None": []}) + return expected_dataset_info + + +if __name__ == "__main__": + input_str = """ + ``` + Question: Please segment the left adrenal gland in CT images. + Thought: I need to find CT datasets with annotations for the left adrenal gland. + Action: Check datasets with CT modality for left adrenal gland labels. + Action Input: Dataset information provided. + Observation: Found datasets with left adrenal gland labels in CT. + Thought: I now know the final answer. + Final Answer: { "thought": "The left adrenal gland can be segmented using CT datasets. I found suitable annotations.", "dataset": "AMOS22_Task2, FLARE23, TotalSegmentator_v2", "labels": "[12], [8], [9]" } + ``` + """ + res = extract_llm_answer(input_str) + print(res) diff --git a/utils/spacy_model.md b/utils/spacy_model.md new file mode 100644 index 0000000..d481a08 --- /dev/null +++ b/utils/spacy_model.md @@ -0,0 +1,4 @@ +use this to download: +``` +python -m spacy download en_core_web_sm +``` \ No newline at end of file