Update app.py

app.py

@@ -1,189 +1,130 @@
-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()
-
+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 utils.thyroid_gen_config import config as thyroid_args
+from utils.liver_gen_config import config as liver_args
+from 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按钮点击与图片显示
+        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
+
+if __name__ == '__main__':
+    # 启动应用程序
+    demo().launch()
+