Update README.md

README.md

@@ -74,38 +74,63 @@ pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
 ).to("cuda")
 generator = torch.Generator("cuda").manual_seed(0)
 
-prompt = "Convert this non-contrast CT slice to mimic a venous-phase contrast-enhanced CT. Brighten and enhance the portal and hepatic veins, and emphasize organ boundaries."
+prompt_art = "Convert this non-contrast CT slice to mimic an arterial-phase contrast-enhanced CT. Brighten and enhance the aorta, major arteries, and adjacent organ boundaries to emphasize arterial flow, focusing on clarity and contrast in these areas while maintaining other features unchanged."
+prompt_ven = "Convert this non-contrast CT slice to mimic a venous-phase contrast-enhanced CT. Brighten and enhance the veins, especially the portal and hepatic veins, and emphasize organ boundaries to mimic venous flow, focusing on brightness and contrast in these areas while maintaining other features unchanged."
 
+# read all dicoms
 def load_dicom_folder(dicom_folder):
     dicom_folder = os.path.join(dicom_folder, 'DICOM')
-    files = sorted(glob(os.path.join(dicom_folder, "*")))
-    slices = [dcmread(f).pixel_array.astype(np.float32) for f in files]
-    volume = np.stack(slices, axis=0)
-    volume += dcmread(files[0]).RescaleIntercept
-    volume = np.clip(volume, -1000, 1000)
-    return (volume + 1000) / 2000.0
-
-def process(volume):
-    results = []
-    for i in tqdm(range(volume.shape[0])):
-        img = Image.fromarray((volume[i] * 255).astype(np.uint8)).convert("RGB")
-        out = pipe(prompt, image=img, num_inference_steps=20,
-                   image_guidance_scale=1.5, guidance_scale=10,
-                   generator=generator).images[0]
-        gray = np.array(out.convert("L")).astype(np.float32) / 255.0
-        gray = gray * 2000 - 1000
-        results.append(gray)
-    return np.stack(results, axis=0)
-
-def save_nifti(volume, path):
-    nib.save(nib.Nifti1Image(volume, np.eye(4)), path)
-
-input_path = "/path/to/dicom_folder"
-output_path = "/path/to/output.nii.gz"
-
-vol = load_dicom_folder(input_path)
-out_vol = process(vol)
-save_nifti(out_vol, output_path)
+    dicom_files = sorted(glob(os.path.join(dicom_folder, "*")))
+
+    slices = []
+    for dicom_file in dicom_files:
+        ds = dcmread(dicom_file)
+        slices.append(ds.pixel_array.astype(np.float32))
+    
+    dicom_array = np.stack(slices, axis=0)
+    dicom_array += ds.RescaleIntercept
+    dicom_array = np.clip(dicom_array, -1000, 1000)
+    dicom_array = (dicom_array + 1000) / 2000.0
+    return dicom_array
+
+# transfer to RGB and send to diffusion
+def process_slices(dicom_array):
+    outputs = []
+    for i in tqdm(range(dicom_array.shape[0])):
+        slice_img = (dicom_array[i] * 255).astype(np.uint8)
+        pil_img = Image.fromarray(slice_img).convert("RGB")
+
+        edited_image = pipe(
+            prompt, #### chose prompt_art or prompt_ven
+            image=pil_img,
+            num_inference_steps=20,
+            image_guidance_scale=1.5,
+            guidance_scale=10,
+            generator=generator,
+        ).images[0]
+
+        gray = edited_image.convert("L")
+        gray_np = np.array(gray).astype(np.float32) / 255.0
+        gray_np = gray_np * 2000 - 1000  # scale back to [-1000, 1000]
+        outputs.append(gray_np)
+
+    volume = np.stack(outputs, axis=0)
+    return volume
+
+# save to nii.gz
+def save_nifti(volume, output_path):
+    affine = np.eye(4)
+    nii = nib.Nifti1Image(volume, affine)
+    nib.save(nii, output_path)
+
+# main function
+def main(input_dicom_path, output_nifti_path):
+    dicom_array = load_dicom_folder(input_dicom_path)
+    edited_volume = process_slices(dicom_array)
+    save_nifti(edited_volume, output_nifti_path)
+
+# DMEO 
+# main("/path/to/dicom_folder", "/path/to/output.nii.gz")
 ```
 
 ---