Upload app.py

app.py

@@ -17,20 +17,28 @@ import llama_cpp.llama_tokenizer
 
 # global params 
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+SW_OVERLAP = 0.50
 examples_path = [
-    os.path.join(THIS_DIR, 'examples', 'HCC_003.nrrd'),
-    os.path.join(THIS_DIR, 'examples', 'HCC_007.nrrd'),
-    os.path.join(THIS_DIR, 'examples', 'HCC_018.nrrd')
+    #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'),
+    #os.path.join(THIS_DIR, 'examples', 'HCC_020.nrrd'), # bad
+    os.path.join(THIS_DIR, 'examples', 'HCC_036.nrrd'), # 
+    os.path.join(THIS_DIR, 'examples', 'HCC_041.nrrd'), # good
+    os.path.join(THIS_DIR, 'examples', 'HCC_051.nrrd'), # ok, rerun with 0.3 
+    #os.path.join(THIS_DIR, 'examples', 'HCC_066.nrrd'), # very bad 
+    #os.path.join(THIS_DIR, 'examples', 'HCC_099.nrrd'), # bad 
 ]
 models_path = {
     "liver": os.path.join(THIS_DIR, 'checkpoints', 'liver_3DSegResNetVAE.pth'),
-    "tumor": os.path.join(THIS_DIR, 'checkpoints', 'tumor_3DSegResNetVAE.pth') 
+    "tumor": os.path.join(THIS_DIR, 'checkpoints', 'tumor_3DSegResNetVAE_weak_morp.pth') 
 }
 cache_path = {
     "liver mask": "liver_mask.npy",
     "tumor mask": "tumor_mask.npy"
 }
-device = "cpu"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 mydict = {}
 
 
@@ -149,22 +157,20 @@ def segment_tumor(image_name):
     
     postprocessing_tumor = Compose([
         Activations(sigmoid=True),
-        # Convert to binary predictions 
         AsDiscrete(argmax=True, to_onehot=3),
-        # Remove small connected components for 1=liver and 2=tumor
-        KeepLargestConnectedComponent(applied_labels=[2]),
-        # Fill holes in the binary mask for 1=liver and 2=tumor
-        FillHoles(applied_labels=[2]),
+        KeepLargestConnectedComponent(applied_labels=[1,2]),
+        FillHoles(applied_labels=[1,2]),
         ToTensor()
     ])
     
     # Preprocessing
     input = preprocessing_tumor(input)
-    input = torch.multiply(input, torch.from_numpy(mydict[image_name]['liver mask'])) # mask non-liver regions 
+    # mask non-liver regions 
+    input = torch.multiply(input, torch.from_numpy(mydict[image_name]['liver mask'])) 
         
     # Generate segmentation 
     with torch.no_grad():
-        segmented_mask = sw_inference(tumor_model, input[None, None, :], (256,256,32), False, discard_second_output=True, overlap=0.2)[0] # input dimensions [B,C,H,W,Z]
+        segmented_mask = sw_inference(tumor_model, input[None, None, :], (256,256,32), False, discard_second_output=True, overlap=SW_OVERLAP)[0] # input dimensions [B,C,H,W,Z]
 
     # Postprocess image 
     segmented_mask = postprocessing_tumor(segmented_mask)[-1].numpy() # background, liver, tumor
@@ -198,13 +204,9 @@ def segment_liver(image_name):
     ])
     
     postprocessing_liver = Compose([
-          # Apply softmax activation to convert logits to probabilities
           Activations(sigmoid=True),
-          # Convert predicted probabilities to discrete values (0 or 1)
           AsDiscrete(argmax=True, to_onehot=None),
-          # Remove small connected components for 1=liver and 2=tumor
           KeepLargestConnectedComponent(applied_labels=[1]),
-          # Fill holes in the binary mask for 1=liver and 2=tumor
           FillHoles(applied_labels=[1]),
           ToTensor()
     ])
@@ -214,7 +216,7 @@ def segment_liver(image_name):
  
     # Generate segmentation 
     with torch.no_grad():
-        segmented_mask = sw_inference(liver_model, input[None, None, :], (512,512,16), False, discard_second_output=True, overlap=0.2)[0] # input dimensions [B,C,H,W,Z]
+        segmented_mask = sw_inference(liver_model, input[None, None, :], (512,512,16), False, discard_second_output=True, overlap=0.25)[0] # input dimensions [B,C,H,W,Z]
 
     # Postprocess image 
     segmented_mask = postprocessing_liver(segmented_mask)[0].numpy() # first channel 
@@ -254,7 +256,7 @@ 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" 
+        return "⛔ You need to generate both liver and tumor masks before we can create a summary report.", "You need to generate both liver and tumor masks before we can create a summary report." 
     
     # extract tumor features from CT scan 
     features = generate_features(mydict[image_name]["img"], mydict[image_name]["liver mask"], mydict[image_name]["tumor mask"])
@@ -270,14 +272,12 @@ def generate_summary(image):
 
     # 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.
+    You are a radiologist. You need to write a diagnosis summary (1-2 paragraphs) given tumor characteristics observed from CT scans. 
+    The report should include your diagnosis, considering the possibility of liver cancer (hepatocellular carcinoma or metastatic liver lesions), 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).
+    The characteristics of this tumor are: {str(features)}. Please provide your diagnosis summary. 
     """
     print(user_msg)
 
@@ -294,7 +294,7 @@ def generate_summary(image):
         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"
+        return "Sorry. There was an error in report generation: " + e, "Sorry. There was an error in report generation: " + e
     
 
 with gr.Blocks() as app:
@@ -322,7 +322,7 @@ with gr.Blocks() as app:
             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'])
+            selected_mask = gr.CheckboxGroup(label='Step 2: Select mask to produce', choices=['liver mask', 'tumor mask'], value = ['liver mask', 'tumor mask'])
             btn_segment = gr.Button("Generate Segmentation")
                 
     with gr.Row():