app.py

import json

from KMVE_RG.models.SGF_model import SGF
from KMVE_RG.modules.tokenizers import Tokenizer
from KMVE_RG.modules.metrics import compute_scores
import numpy as np
from Demo.utils.thyroid_gen_config import config as thyroid_args
from Demo.utils.liver_gen_config import config as liver_args
from Demo.utils.breast_gen_config import config as breast_args

import gradio as gr
import torch
from PIL import Image
import os
from torchvision import transforms

np.random.seed(9233)
torch.manual_seed(9233)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

class Generator(object):
    def __init__(self, model_type):
        if model_type == '甲状腺':
            self.args = thyroid_args
        elif model_type == '乳腺':
            self.args = breast_args
        elif model_type == '肝脏':
            self.args = liver_args
        self.tokenizer = Tokenizer(self.args)
        self.model = SGF(self.args, self.tokenizer)
        sd = torch.load(self.args.models)['state_dict']
        msg = self.model.load_state_dict(sd)
        print(msg)
        self.model.eval()
        self.metrics = compute_scores
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406),
                                 (0.229, 0.224, 0.225))])
        with open(self.args.ann_path, 'r', encoding='utf-8-sig') as f:
            self.data = json.load(f)
        print('模型加载完成')

    def image_process(self, img_paths):
        image_1 = Image.open(os.path.join(self.args.image_dir, img_paths[0])).convert('RGB')
        image_2 = Image.open(os.path.join(self.args.image_dir, img_paths[1])).convert('RGB')
        if self.transform is not None:
            image_1 = self.transform(image_1)
            image_2 = self.transform(image_2)
        image = torch.stack((image_1, image_2), 0)
        return image

    def generate(self, uid):
        img_paths, report = self.data[uid]['img_paths'], self.data[uid]['report']
        imgs = self.image_process(img_paths)
        imgs = imgs.unsqueeze(0)
        with torch.no_grad():
            output, _ = self.model(imgs, mode='sample')
        pred = self.tokenizer.decode(output[0].cpu().numpy())
        gt = self.tokenizer.decode(self.tokenizer(report[:self.args.max_seq_length])[1:])
        scores = self.metrics({0: [gt]}, {0: [pred]})
        return pred, gt, scores

    def visualize_images(self, uid):
        image_1 = Image.open(os.path.join(self.args.image_dir, self.data[uid]['img_paths'][0])).convert('RGB')
        image_2 = Image.open(os.path.join(self.args.image_dir, self.data[uid]['img_paths'][1])).convert('RGB')
        return image_1, image_2

# 主应用程序
def demo():
    with gr.Blocks() as app:
        gr.Markdown("# 超声报告生成Demo")
        gr.Markdown('### SIAT认知与交互技术中心')
        gr.Markdown('### 项目主页：https://lijunrio.github.io/Ultrasound-Report-Generation/')

        # 选择模型
        with gr.Row():
            model_choice = gr.Radio(choices=["甲状腺", "乳腺", "肝脏"], label="请选择模型类型", interactive=True)

        model = gr.State()

        # 展示UID按钮
        uids = [f"uid_{i}" for i in range(20)]
        with gr.Row():
            uid_choice = gr.Radio(choices=[f"{uid}" for uid in uids], label="请选择uid", interactive=False)

        # 定义展示图片的组件
        with gr.Row():
            image1_display = gr.Image(label="图像1", visible=True)
            image2_display = gr.Image(label="图像2", visible=True)

        # 定义生成报告的按钮和文本框
        generate_button = gr.Button("生成报告", interactive=False)
        generated_report_display = gr.Textbox(label="生成的报告", visible=True)
        ground_truth_display = gr.Textbox(label="Ground Truth报告", visible=True)
        nlp_score_display = gr.Textbox(label="NLP得分", visible=True)

        # 加载模型的回调函数
        def load_model_and_uids(model_type):
            model = Generator(model_type)
            return model, gr.update(interactive=True)

        # 点击UID按钮后加载对应的图片
        def on_uid_click(model, uid):
            image1, image2 = model.visualize_images(uid)
            # 显示图片和生成按钮
            return image1, image2, gr.update(interactive=True)

        # 点击生成按钮生成报告
        def on_generate_click(model, uid):
            generated_report, ground_truth_report, nlp_score = model.generate(uid)
            # 展示生成的报告、Ground Truth 和 NLP 得分
            return generated_report, ground_truth_report, f"NLP得分: {nlp_score}"

        # 链接模型选择与UID按钮显示
        model_choice.change(load_model_and_uids, inputs=model_choice, outputs=[model, uid_choice])

        # 链接UID按钮点击与图片显示
        # for uid_button in uid_buttons:
        #     uid_button.click(on_uid_click, inputs=[model, uid_button], outputs=[image1_display, image2_display, uid_choice])
        uid_choice.change(on_uid_click, inputs=[model, uid_choice], outputs=[image1_display, image2_display, generate_button])

        generate_button.click(on_generate_click, inputs=[model, uid_choice], outputs=[generated_report_display, ground_truth_display, nlp_score_display])

    return app

# 主应用程序
# def demo():
#     with gr.Blocks() as app:
#         gr.Markdown("# 医学报告生成 Demo")
#
#         # 选择模型类型
#         model_choice = gr.Radio(choices=["甲状腺", "乳腺", "肝脏"], label="请选择模型类型", interactive=True)
#
#         # 创建空的 Generator 实例（将稍后初始化）
#         generator_instance = gr.State()
#
#         # 展示 UID 按钮
#         uids = [f"uid_{i}" for i in range(20)]
#         selected_uid = gr.State()  # 用于存储当前选择的 UID
#         uid_buttons = [gr.Button(f"{uid}") for uid in uids]
#
#         # 定义展示图片的组件
#         with gr.Row():
#             image1_display = gr.Image(label="图像 1", visible=False)
#             image2_display = gr.Image(label="图像 2", visible=False)
#
#         # 定义生成报告的按钮和文本框
#         generate_button = gr.Button("生成报告", visible=False)
#         generated_report_display = gr.Textbox(label="生成的报告", visible=False)
#         ground_truth_display = gr.Textbox(label="Ground Truth 报告", visible=False)
#         nlp_score_display = gr.Textbox(label="NLP 得分", visible=False)
#
#         # 模型选择后初始化 Generator 类
#         def initialize_generator(model_type):
#             generator = Generator(model_type)  # 初始化 Generator
#             return generator, True  # 返回生成器实例，显示 UID 按钮
#
#         # 点击 UID 按钮后可视化对应图片
#         def on_uid_click(uid, generator):
#             image1, image2 = generator.visual_images(uid)
#             return image1, image2, uid, True  # 返回图片、UID，显示生成按钮
#
#         # 点击生成按钮生成 Ground Truth 报告、预测结果和 NLP 分数
#         def on_generate_click(generator, uid):
#             ground_truth, predict, nlp_score = generator.generate(uid)
#             return ground_truth, predict, nlp_score
#
#         # 链接模型选择与生成器初始化
#         model_choice.change(initialize_generator, inputs=model_choice, outputs=[generator_instance, uid_buttons[0]])
#
#         # 链接 UID 按钮点击与图片显示
#         for i, uid_button in enumerate(uid_buttons):
#             uid_button.click(on_uid_click, inputs=[selected_uid, generator_instance],
#                              outputs=[image1_display, image2_display, selected_uid, generate_button],
#                              fn=lambda uid=uids[i]: uid)
#
#         # 点击生成按钮时生成报告
#         generate_button.click(on_generate_click, inputs=[generator_instance, selected_uid],
#                               outputs=[ground_truth_display, generated_report_display, nlp_score_display])
#
#     return app

if __name__ == '__main__':
    # 启动应用程序
    demo().launch()