SIAT-RZJS/Ultrasound_Report_Generation

1import json2 3from KMVE_RG.models.SGF_model import SGF4from KMVE_RG.modules.tokenizers import Tokenizer5from KMVE_RG.modules.metrics import compute_scores6import numpy as np7from utils.thyroid_gen_config import config as thyroid_args8from utils.liver_gen_config import config as liver_args9from utils.breast_gen_config import config as breast_args10 11import gradio as gr12import torch13from PIL import Image14import os15from torchvision import transforms16 17np.random.seed(9233)18torch.manual_seed(9233)19torch.backends.cudnn.deterministic = True20torch.backends.cudnn.benchmark = False21 22class Generator(object):23    def __init__(self, model_type):24        if model_type == '甲状腺':25            self.args = thyroid_args26        elif model_type == '乳腺':27            self.args = breast_args28        elif model_type == '肝脏':29            self.args = liver_args30        self.tokenizer = Tokenizer(self.args)31        self.model = SGF(self.args, self.tokenizer)32        sd = torch.load(self.args.models)['state_dict']33        msg = self.model.load_state_dict(sd)34        print(msg)35        self.model.eval()36        self.metrics = compute_scores37        self.transform = transforms.Compose([38            transforms.Resize((224, 224)),39            transforms.ToTensor(),40            transforms.Normalize((0.485, 0.456, 0.406),41                                 (0.229, 0.224, 0.225))])42        with open(self.args.ann_path, 'r', encoding='utf-8-sig') as f:43            self.data = json.load(f)44        print('模型加载完成')45 46    def image_process(self, img_paths):47        image_1 = Image.open(os.path.join(self.args.image_dir, img_paths[0])).convert('RGB')48        image_2 = Image.open(os.path.join(self.args.image_dir, img_paths[1])).convert('RGB')49        if self.transform is not None:50            image_1 = self.transform(image_1)51            image_2 = self.transform(image_2)52        image = torch.stack((image_1, image_2), 0)53        return image54 55    def generate(self, uid):56        img_paths, report = self.data[uid]['img_paths'], self.data[uid]['report']57        imgs = self.image_process(img_paths)58        imgs = imgs.unsqueeze(0)59        with torch.no_grad():60            output, _ = self.model(imgs, mode='sample')61        pred = self.tokenizer.decode(output[0].cpu().numpy())62        gt = self.tokenizer.decode(self.tokenizer(report[:self.args.max_seq_length])[1:])63        scores = self.metrics({0: [gt]}, {0: [pred]})64        return pred, gt, scores65 66    def visualize_images(self, uid):67        image_1 = Image.open(os.path.join(self.args.image_dir, self.data[uid]['img_paths'][0])).convert('RGB')68        image_2 = Image.open(os.path.join(self.args.image_dir, self.data[uid]['img_paths'][1])).convert('RGB')69        return image_1, image_270 71# 主应用程序72def demo():73    with gr.Blocks() as app:74        gr.Markdown("# 超声报告生成Demo")75        gr.Markdown('### SIAT认知与交互技术中心')76        gr.Markdown('### 项目主页：https://lijunrio.github.io/Ultrasound-Report-Generation/')77 78        # 选择模型79        with gr.Row():80            model_choice = gr.Radio(choices=["甲状腺", "乳腺", "肝脏"], label="请选择模型类型", interactive=True)81 82        model = gr.State()83 84        # 展示UID按钮85        uids = [f"uid_{i}" for i in range(20)]86        with gr.Row():87            uid_choice = gr.Radio(choices=[f"{uid}" for uid in uids], label="请选择uid", interactive=False)88 89        # 定义展示图片的组件90        with gr.Row():91            image1_display = gr.Image(label="图像1", visible=True)92            image2_display = gr.Image(label="图像2", visible=True)93 94        # 定义生成报告的按钮和文本框95        generate_button = gr.Button("生成报告", interactive=False)96        generated_report_display = gr.Textbox(label="生成的报告", visible=True)97        ground_truth_display = gr.Textbox(label="Ground Truth报告", visible=True)98        nlp_score_display = gr.Textbox(label="NLP得分", visible=True)99 100        # 加载模型的回调函数101        def load_model_and_uids(model_type):102            model = Generator(model_type)103            return model, gr.update(interactive=True)104 105        # 点击UID按钮后加载对应的图片106        def on_uid_click(model, uid):107            image1, image2 = model.visualize_images(uid)108            # 显示图片和生成按钮109            return image1, image2, gr.update(interactive=True)110 111        # 点击生成按钮生成报告112        def on_generate_click(model, uid):113            generated_report, ground_truth_report, nlp_score = model.generate(uid)114            # 展示生成的报告、Ground Truth 和 NLP 得分115            return generated_report, ground_truth_report, f"NLP得分: {nlp_score}"116 117        # 链接模型选择与UID按钮显示118        model_choice.change(load_model_and_uids, inputs=model_choice, outputs=[model, uid_choice])119 120        # 链接UID按钮点击与图片显示121        uid_choice.change(on_uid_click, inputs=[model, uid_choice], outputs=[image1_display, image2_display, generate_button])122 123        generate_button.click(on_generate_click, inputs=[model, uid_choice], outputs=[generated_report_display, ground_truth_display, nlp_score_display])124 125    return app126 127if __name__ == '__main__':128    # 启动应用程序129    demo().launch()130 131