update: log

src/parcellation.py

@@ -93,39 +93,57 @@ def run_inference(input_file, only_face_cropping, only_skull_stripping):
     print(f"Using device: {device}")
 
     # Load the pre-trained models from the fixed "model/" folder
+    print("Loading models...")
     cnet, ssnet, pnet_c, pnet_s, pnet_a, hnet_c, hnet_a = load_model("model/", device=device)
+    print("Models loaded successfully.")
     # cnet, ssnet, pnet_a, hnet_c, hnet_a = load_model("model/", device=device)
 
     # --- Processing Flow (based on the original parcellation.py) ---
     # 1. Load the input image, convert to canonical orientation, and remove extra dimensions
+    print("Loading and preprocessing the input image...")
     odata = nib.squeeze_image(nib.as_closest_canonical(nib.load(input_file.name)))
     nii = nib.Nifti1Image(odata.get_fdata().astype(np.float32), affine=odata.affine)
     os.makedirs(os.path.join(opt.o, "original"), exist_ok=True)
     original_nii_path = os.path.join(opt.o, f"original/{basename}.nii")
     nib.save(nii, original_nii_path)
+    print(f"Input image saved to: {original_nii_path}")
 
     # 2. Preprocess the image
+    print("Preprocessing the input image...")
     odata, data = preprocessing(input_file.name, opt.o, basename)
+    print("Preprocessing completed.")
 
     # 3. Cropping
+    print("Cropping the input image...")
     cropped, out_filename = cropping(opt.o, basename, odata, data, cnet, device)
+    print("Cropping completed.")
     if only_face_cropping:
         pass
 
     else:
         # 4. Skull stripping
+        print("Performing skull stripping...")
         stripped, shift, out_filename = stripping(opt.o, basename, cropped, odata, data, ssnet, device)
+        print("Skull stripping completed.")
         if only_skull_stripping:
             pass
         else:
             # 5. Parcellation
+            print("Starting parcellation...")
             parcellated = parcellation(stripped, pnet_c, pnet_s, pnet_a, device)
+            print("Parcellation completed.")
             # 6. Separate into hemispheres
+            print("Separating hemispheres...")
             separated = hemisphere(stripped, hnet_c, hnet_a, device)
+            print("Hemispheres separated.")
             # 7. Postprocessing
+            print("Postprocessing the parcellated data...")
             output = postprocessing(parcellated, separated, shift, device)
+            print("Postprocessing completed.")
             # 8. Create CSV with volume information, etc.
+            print("Creating CSV with volume information...")
             df = make_csv(output, opt.o, basename)
+            print("CSV created successfully.")
             # 9. Create and save the parcellation result NIfTI file
             nii_out = nib.Nifti1Image(output.astype(np.uint16), affine=data.affine)
             header = odata.header
@@ -140,6 +158,7 @@ def run_inference(input_file, only_face_cropping, only_skull_stripping):
             out_filename = os.path.join(out_parcellated_dir, f"{basename}_Type1_Level5.nii")
             nib.save(nii_out, out_filename)
             create_parcellated_images(output, opt.o, basename, odata, data)
+            print(f"Parcellation result saved to: {out_filename}")
 
     # Zip the entire output directory into a ZIP file
     zip_path = os.path.join(os.path.dirname(opt.o), f"{basename}_results.zip")
@@ -156,7 +175,7 @@ def run_inference(input_file, only_face_cropping, only_skull_stripping):
 
     # *** Cleanup: Remove the temporary output directory ***
     # Note: This is performed before returning. It is not possible to execute code after the return statement.
-    shutil.rmtree(opt.o)
+    # shutil.rmtree(opt.o)
 
     # Return the ZIP file path and the two visualization images
     return zip_path, Image.open(input_png_path), Image.open(parcellation_png_path)