Upload 2 files

app.py

@@ -11,19 +11,20 @@ from monai.transforms import (
     LoadImage, Orientation, Compose, ToTensor, Activations, 
     FillHoles, KeepLargestConnectedComponent, AsDiscrete, ScaleIntensityRange
 )
+import llama_cpp
+import llama_cpp.llama_tokenizer
 
 
 # global params 
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
 examples_path = [
     os.path.join(THIS_DIR, 'examples', 'HCC_003.nrrd'),
-    os.path.join(THIS_DIR, 'examples', 'HCC_006.nrrd'),
     os.path.join(THIS_DIR, 'examples', 'HCC_007.nrrd'),
     os.path.join(THIS_DIR, 'examples', 'HCC_018.nrrd')
 ]
 models_path = {
     "liver": os.path.join(THIS_DIR, 'checkpoints', 'liver_3DSegResNetVAE.pth'),
-    "tumor": os.path.join(THIS_DIR, 'checkpoints', 'tumor_3DSegResNetVAE_weak_morp.pth') 
+    "tumor": os.path.join(THIS_DIR, 'checkpoints', 'tumor_3DSegResNetVAE.pth') 
 }
 cache_path = {
     "liver mask": "liver_mask.npy",
@@ -127,7 +128,7 @@ def load_image(image, slider, selected_slice):
     
     image, annotations = render(image_name, slider, selected_slice)
     
-    return f"Your image is successfully loaded! Please use the slider to view the image (zmin: 1, zmax: {input.shape[-1]}).", (image, annotations)
+    return f"😊 Your image is successfully loaded! Please use the slider to view the image (zmin: 1, zmax: {input.shape[-1]}).", (image, annotations)
     
 
 def segment_tumor(image_name):
@@ -245,44 +246,87 @@ def segment(image, selected_mask, slider, selected_slice):
     
     image, annotations = render(image, slider, selected_slice)
     
-    return f"Segmentation is completed! ", download_liver, download_tumor, (image, annotations)
+    return f"🥳 Segmentation is completed. You can use the slider to view slices or proceed with generating a summary report.", download_liver, download_tumor, (image, annotations)
 
 
 def generate_summary(image):
+    
     image_name = image.name.split('/')[-1].replace(".nrrd","")
+    
+    if "liver mask" not in mydict[image_name] or "tumor mask" not in mydict[image_name]:
+        return "⛔ You need to generate both liver and tumor masks before we can create a summary report.", "Not generated" 
+    
+    # extract tumor features from CT scan 
     features = generate_features(mydict[image_name]["img"], mydict[image_name]["liver mask"], mydict[image_name]["tumor mask"])
     print(features)
     
-    return ""
+    # initialize LLM pulling from hugging face 
+    llama = llama_cpp.Llama.from_pretrained(
+        repo_id="Qwen/Qwen1.5-0.5B-Chat-GGUF",
+        filename="*q8_0.gguf",
+        tokenizer=llama_cpp.llama_tokenizer.LlamaHFTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B"),
+        verbose=False
+    )
+
+    # openai.api_key = os.environ["OPENAI"]
+    system_msg = """
+    You are a radiologist. You use a segmentation model that extracts tumor characteristics from CT scans from which you generate a diagnosis report. 
+    The report should include recommendations for next steps, and a disclaimer that these results should be taken with a grain of salt.
+    """
+    
+    user_msg = f"""
+    The tumor characteristics are:
+    {str(features)}
+    Please provide your interpretation of the findings and a differential diagnosis, considering the possibility of liver cancer (hepatocellular carcinoma or metastatic liver lesions).
+    """
+    print(user_msg)
+
+    response = llama.create_chat_completion(
+        messages = [
+            {"role": "system", "content": system_msg},
+            {"role": "user", "content": user_msg}
+        ],
+        temperature=0.7
+    )
+    print(response)
+    
+    try:
+        report = response["choices"][0]["message"]["content"]
+        return "📝 Your AI diagnosis summary report is generated! Please review below. Thank you for trying this tool!", report 
+    except Exception as e:
+        return "Sorry. There was an error in report generation: " + e, "To be generated"
     
 
 with gr.Blocks() as app:
     with gr.Column():
         gr.Markdown(
         """
-        # Lung Tumor Segmentation App
+        # MedAssist-Liver: an AI-powered Liver Tumor Segmentation Tool
+        
+        Welcome to explore the power of AI for automated medical image analysis with our user-friendly app! 
         
-        This tool is designed to assist in the identification and segmentation of lung and tumor from medical images. By uploading a CT scan image, a pre-trained machine learning model will automatically segment the lung and tumor regions. Segmented tumor's characteristics such as shape, size, and location are then analyzed to produce an AI-generated diagnosis report of the lung cancer.
+        This tool is designed to assist in the identification and segmentation of liver and tumor from medical images. By uploading a CT scan image, a pre-trained machine learning model will automatically segment the liver and tumor regions. Segmented tumor's characteristics such as shape, size, and location are then analyzed to produce an AI-generated diagnosis report of the liver cancer.
         
-        ⚠️ Important disclaimer: these model outputs should NOT replace the medical diagnosis of healthcare professionals. For your reference, our model was trained on the [HCC-TACE-Seg dataset](https://www.cancerimagingarchive.net/collection/hcc-tace-seg/) and achieved 0.954 dice score for lung segmentation and 0.570 dice score for tumor segmentation. Improving tumor segmentation is still an active area of research! 
+        ⚠️ Important disclaimer: these model outputs should NOT replace the medical diagnosis of healthcare professionals. For your reference, our model was trained on the [HCC-TACE-Seg dataset](https://www.cancerimagingarchive.net/collection/hcc-tace-seg/) and achieved 0.954 dice score for liver segmentation and 0.570 dice score for tumor segmentation. Improving tumor segmentation is still an active area of research! 
         """)
     
     with gr.Row():
-        comment = gr.Textbox(label='Your tool guide:', value="👋 Hi there, welcome to explore the power of AI for automated medical image analysis with our user-friendly app! Start by uploading a CT scan image. Note that for now we accept .nrrd formats only.")
+        comment = gr.Textbox(label='🤖 Your tool guide:', value="👋 Hi there, I will be helping you use this tool. To get started, upload a CT scan image or select one from examples.")
         
     
     with gr.Row():
         
         with gr.Column(scale=2):
             image_file = gr.File(label="Step 1: Upload a CT image (.nrrd)", file_count='single', file_types=['.nrrd'], type='filepath')
+            gr.Examples(examples_path, [image_file])
             btn_upload = gr.Button("Upload")
             
         with gr.Column(scale=2):
             selected_mask = gr.CheckboxGroup(label='Step 2: Select mask to produce', choices=['liver mask', 'tumor mask'], value = ['liver mask'])
-            btn_segment = gr.Button("Segment")
+            btn_segment = gr.Button("Generate Segmentation")
                 
     with gr.Row():
-        slider = gr.Slider(1, 100, step=1, label="Slice (z)")
+        slider = gr.Slider(1, 100, step=1, label="Image slice: ")
         selected_slice = gr.State(value=1)
         
     with gr.Row():
@@ -295,17 +339,11 @@ with gr.Blocks() as app:
             btn_download_tumor = gr.DownloadButton("Download tumor mask", visible=False)
         
     with gr.Row():       
-        report = gr.Textbox(label='Step 4. Generate summary report using AI:')
+        report = gr.Textbox(label='Step 4. Generate summary report using AI:', value="To be generated. ")
             
     with gr.Row():
         btn_report = gr.Button("Generate summary")  
-    
 
-    gr.Examples(
-        examples_path,
-        [image_file],
-    )
-    
     btn_upload.click(fn=load_image, 
         inputs=[image_file, slider, selected_slice], 
         outputs=[comment, myimage],
@@ -323,7 +361,8 @@ with gr.Blocks() as app:
     )
     
     btn_report.click(fn=generate_summary,
-        outputs=report
+        inputs=[image_file],
+        outputs=[comment, report]
     )
     
     

requirements.txt

@@ -5,4 +5,6 @@ pynrrd==1.0.0
 nibabel==5.2.1
 scikit-image==0.23.2
 morphsnakes==2.0.1
-opencv-python==4.9.0.80
+opencv-python==4.9.0.80
+llama-cpp-python==0.2.63
+transformers==4.40.0