Update app.py

app.py

@@ -7,10 +7,15 @@ import torch
 import gradio as gr
 
 from huggingface_hub import hf_hub_download
+
+
 def cached_download(*args, **kwargs):
     print("Warning: cached_download is deprecated, using hf_hub_download instead.")
     return hf_hub_download(*args, **kwargs)
+
+
 import sys
+
 sys.modules["huggingface_hub.cached_download"] = cached_download
 
 from diffusers import AutoencoderKL, DDPMScheduler
@@ -20,20 +25,6 @@ from medical_pipeline import MedicalPipeline
 from diffusers import DDIMScheduler
 from StableDiffusion.Our_Pipe import StableDiffusionPipeline
 
-def get_random_values(my_dict):
-    values_list = list(my_dict.values())
-    num_choices = random.randint(1, len(values_list))
-    kinds = random.sample(values_list, num_choices)
-    kind = ''
-
-    for k in kinds:
-        if kind == '':
-            kind = k
-        else:
-            kind = kind + ',' + k
-
-    return kind
-
 model_repo_id = "runwayml/stable-diffusion-v1-5"
 medsegfactory_id = "JohnWeck/StableDiffusion"
 filename = 'checkpoint-300.pth'
@@ -75,42 +66,72 @@ pipe = StableDiffusionPipeline.from_pretrained(
 
 pipeline = MedicalPipeline(pipe, device)
 
+# 定义 keys 与 organ 及 kind 的映射
+keys_to_organ_kind = {
+    "CVC-ClinicDB": {
+        "organs": ["polyp colonoscopy"],
+        "kinds": {"polyp colonoscopy": ["polyp"]}
+    },
+    "BUSI": {
+        "organs": ["breast ultrasound"],
+        "kinds": {"breast ultrasound": ["normal", "breast tumor"]}
+    },
+    "LiTS2017": {
+        "organs": ["abdomen CT scans"],
+        "kinds": {"abdomen CT scans": ["liver","liver tumor"]}
+    },
+    "KiTS2019": {
+            "organs": ["abdomen CT scans"],
+            "kinds": {"abdomen CT scans": ["kidney","kidney tumor"]}
+        },
+    "ACDC": {
+            "organs": ["cardiovascular ventricle mri"],
+            "kinds": {"cardiovascular ventricle mri": ["right ventricle", "myocardium","left ventricle"]}
+        },
+    "AMOS2022": {
+            "organs": ["abdomen CT scans"],
+            "kinds": {"abdomen CT scans": ["liver", "right kidney", "spleen", "pancreas", "aorta", "inferior vena cava",
+                         "right adrenal gland", "left adrenal gland", "gall bladder", "esophagus", "stomach", "duodenum", "left kidney",
+                         "bladder", "prostate"]}
+        }
+}
+
+
+def update_organ_and_kind(selected_key):
+    organs = keys_to_organ_kind[selected_key]["organs"]
+    first_organ = organs[0]  # 默认选第一个 organ
+    kinds = keys_to_organ_kind[selected_key]["kinds"][first_organ]
+    return gr.update(choices=organs, value=first_organ), gr.update(choices=kinds, value=kinds)
+
+
+def update_kind(selected_key, selected_organ):
+    kinds = keys_to_organ_kind[selected_key]["kinds"][selected_organ]
+    return gr.update(choices=kinds, value=kinds)
+
+
 def generate_image(organ, kind, keys):
-    # 使用 MedSegFactory 生成图像
     image, label = pipeline.generate(organ=organ, kind=kind, keys=keys)
-    plt.subplot(1,2,1)
+    plt.subplot(1, 2, 1)
     plt.imshow(image)
     plt.axis('off')
-    plt.subplot(1,2,2)
+    plt.subplot(1, 2, 2)
     plt.imshow(label)
     plt.axis('off')
-    plt.savefig('pred.png',bbox_inches='tight', pad_inches = 0)
+    plt.savefig('pred.png', bbox_inches='tight', pad_inches=0)
     return "pred.png"
 
-AMOS2022 = {1: 'liver', 2: 'right kidney', 3: 'spleen', 4: 'pancreas', 5: 'aorta', 6: 'inferior vena cava',
-            7: 'right adrenal gland', 8: 'left adrenal gland',
-            9: 'gall bladder', 10: 'esophagus', 11: 'stomach', 12: 'duodenum', 13: 'left kidney',
-            14: 'bladder', 15: 'prostate'}
-ACDC = {1: 'right ventricle', 2: 'myocardium', 3: 'left ventricle'}
-LiTS2017 = {1: 'liver', 2: 'liver tumor'}
-KiTS2019 = {1: 'kidney', 2: 'kidney tumor'}
-
-# 预定义的输入案例
-examples = [
-    ['polyp colonoscopy', 'polyp', 'CVC-ClinicDB'],
-    ['breast ultrasound', 'normal', 'BUSI'],
-    ['breast ultrasound', 'breast tumor', 'BUSI'],
-]
-
-# 创建 Gradio 接口
-interface = gr.Interface(
-    fn=generate_image,
-    inputs=[gr.Dropdown(["polyp colonoscopy"], label="organ"),
-            gr.Dropdown(["polyp"], label="kind"),
-            gr.Dropdown(['CVC-ClinicDB'], label="keys")],
-    outputs=gr.Image(label="Visualization"),  # 返回 Matplotlib 渲染的图片
-    examples=examples  # 添加输入案例
-)
 
-# 启动 Gradio 应用
-interface.launch()
+with gr.Blocks() as demo:
+    keys_dropdown = gr.Dropdown(list(keys_to_organ_kind.keys()), label="Keys", value="CVC-ClinicDB")
+    organ_dropdown = gr.Dropdown(keys_to_organ_kind["CVC-ClinicDB"]["organs"], label="Organ")
+    kind_checkbox = gr.CheckboxGroup(keys_to_organ_kind["CVC-ClinicDB"]["kinds"]["polyp colonoscopy"], label="Kind")
+
+    keys_dropdown.change(update_organ_and_kind, inputs=keys_dropdown, outputs=[organ_dropdown, kind_checkbox])
+    organ_dropdown.change(update_kind, inputs=[keys_dropdown, organ_dropdown], outputs=kind_checkbox)
+
+    generate_button = gr.Button("Generate Image")
+    output_image = gr.Image(label="Visualization")
+
+    generate_button.click(generate_image, inputs=[organ_dropdown, kind_checkbox, keys_dropdown], outputs=output_image)
+
+demo.launch()