Updated app.py

app.py

@@ -184,16 +184,32 @@ def run_3D_registration(user_section, ):
                                               moving=affined_fixed_atlas,
                                               transformlist=mytx_non_rigid['fwdtransforms'],
                                               interpolator='nearestNeighbor')
-  gallery_images = load_gallery_images()
+  gallery_images = natsorted(load_gallery_images())
   transformed_images = []
   if not(os.path.exists("Overlaped_registered")):
     os.mkdir("Overlaped_registered")
-  registered = nonrigid_fixed_atlas.numpy()/255
-  for id in list(range((registered.shape[0]//2)-15, (registered.shape[0]//2)+15, 2)):
-    print(id)
-    plt.imsave(f'Overlaped_registered/{id}.png',registered[id, :, :], cmap = 'gray' )
-    transformed_images.append(f'Overlaped_registered/{id}.png')
-
+  # registered = nonrigid_fixed_atlas.numpy()/255
+  # for id in list(range((registered.shape[0]//2)-15, (registered.shape[0]//2)+15, 2)):
+  #   print(id)
+  #   plt.imsave(f'Overlaped_registered/{id}.png',registered[id, :, :], cmap = 'gray' )
+  #   transformed_images.append(f'Overlaped_registered/{id}.png')
+  for i in range(len(gallery_images)-10):
+    im = plt.imread(gallery_images[i])
+    fname = os.path.split(gallery_images[i])[-1]
+    moving_image_slice = ants.from_numpy(square_padding(gray_scale(im)))
+    affined_fixed_atlas = ants.apply_transforms(fixed=fixed_image,
+                                                moving=moving_image,
+                                                transformlist=mytx['fwdtransforms'],
+                                                interpolator='nearestNeighbor')
+    nonrigid_fixed_atlas = ants.apply_transforms(fixed=fixed_image,
+                                                moving=affined_fixed_atlas,
+                                                transformlist=mytx_non_rigid['fwdtransforms'],
+                                                interpolator='nearestNeighbor')
+    # print(im.shape, nonrigid_fixed_atlas.numpy().shape)
+    reconverted_img = reconvert_to_rgb(im[:,:,:3], nonrigid_fixed_atlas.numpy()[i,:,:])
+    plt.imsave(f'Overlaped_registered/{fname}',(reconverted_img * 255).astype(np.uint8))
+    transformed_images.append(f'Overlaped_registered/{fname}')
+  transformed_images = natsorted(load_gallery_images())
   return transformed_images
 
 
@@ -219,9 +235,11 @@ def run_2D_registration(user_section, slice_idx):
   mytx = affine_reg(fixed_image,moving_image)
   mytx_non_rigid = nonrigid_reg(fixed_image,mytx)
   gallery_imgs = natsorted(load_gallery_images())
-  moving_gallery_img = ants.from_numpy(square_padding(plt.imread(gallery_imgs[gallery_selected_data])))
+  im = plt.imread(gallery_imgs[gallery_selected_data])
+  print(im.shape)
+  moving_gallery_img = ants.from_numpy(square_padding(gray_scale(im)))
   affined_fixed_atlas = ants.apply_transforms(fixed=fixed_image,
-                                                moving=moving_image,
+                                                moving=moving_gallery_img,
                                                 transformlist=mytx['fwdtransforms'],
                                                 interpolator='nearestNeighbor')
   nonrigid_fixed_atlas = ants.apply_transforms(fixed=fixed_image,
@@ -232,10 +250,29 @@ def run_2D_registration(user_section, slice_idx):
   transformed_images = []
   if not(os.path.exists("Overlaped_registered")):
     os.mkdir("Overlaped_registered")
-  plt.imsave(f'Overlaped_registered/registered_slice.png',nonrigid_fixed_atlas.numpy()/255, cmap = 'gray')
+  print("Reconverting Image")
+  reconverted_img = reconvert_to_rgb(im[:,:,:3], nonrigid_fixed_atlas.numpy())
+  plt.imsave(f'Overlaped_registered/registered_slice_reconverted_1.png',(reconverted_img * 255).astype(np.uint8))
   
-  return ['Overlaped_registered/registered_slice.png']
+  return ['Overlaped_registered/registered_slice_reconverted_1.png']
 
+def reconvert_to_rgb(img_rgb, img_gray_processed):
+
+  # 3. Resize original RGB to match processed grayscale shape
+  original_shape = img_gray_processed.shape
+  img_rgb_resized = resize(img_rgb, (original_shape[0], original_shape[1]), preserve_range=True)
+
+  # 4. Convert resized RGB to grayscale
+  gray_resized = np.mean(img_rgb_resized, axis=2) + 1e-8  # avoid divide-by-zero
+
+  # 5. Compute ratio of new_gray / old_gray, apply to RGB channels
+  ratio = img_gray_processed / gray_resized
+  img_recolored = img_rgb_resized * ratio[..., np.newaxis]
+
+  # 6. Clip values to [0, 1] if image is in float format (common for imread)
+  # img_recolored = np.clip(img_recolored, 0, 1)
+
+  return img_recolored
 
 def embeddings_classifier(user_section, atlas_embeddings,atlas_labels):
     class SliceEncoder(nn.Module):
@@ -307,6 +344,8 @@ def atlas_slice_prediction(user_section, axis = 'coronal'):
 
 
 
+
+
 example_files = [
     ["./resampled_green_25.nii.gz", "CCF registered Sample", "3D"],
     ["./Brain_1.png", "Custom Sample", "2D"],
@@ -467,6 +506,9 @@ def handle_data_type_change(dt):
         return gr.update(visible=True, minimum=0, maximum=len(actual_ids)-1, value=len(actual_ids)//2)
 
 def on_select(evt: gr.SelectData):
+  print("Selected index:", evt)
+  print("Selected value:", evt.value)
+  print("Selected coordinates:", evt.selected)
   gallery_selected_data = evt.index
 
 gallery_images = natsorted(load_gallery_images())
@@ -474,7 +516,7 @@ with gr.Blocks() as demo:
     gr.Markdown("# Map My Sections\n### This GUI is part of the submission to the Allen Institute's Map My Sections tool by Tibbling Technologies.")
 
     with gr.Row():
-      gr.Markdown("### Step 1: Upload your sample, currently only .nii.gz and .png supported")
+      gr.Markdown("### Step 1: Upload your sample, currently only .nii.gz (3D) and .png (2D) supported")
       gr.Markdown("### Step 2: Select your sample and data type.")
     with gr.Row():
         nifti_file = gr.File(label="File Upload")
@@ -492,10 +534,10 @@ with gr.Blocks() as demo:
         with gr.Column(scale=1):
             gr.Markdown("### Step 4: Visualizing Allen Brain Cell Types Atlas")
             gallery = gr.Gallery(label="ABC Atlas", value = gallery_images,columns = 5, height = 450)
-            gr.Markdown("### Step 5: Run cell type mappin and/or registeration. ")
+            gr.Markdown("### Step 5: Run cell type mappinp and/or registeration. ")
             with gr.Row():
-                run_button = gr.Button("Run Mapping")
-                reg_button = gr.Button("Run Registration", visible=False)
+                run_button = gr.Button("Map My Sections")
+                reg_button = gr.Button("Run Registration (Optional)", visible=False)
 
     with gr.Column(visible=False) as plot_row:
         gr.Markdown("### Step 6: Quantitative results of the mapping model.")
@@ -522,7 +564,7 @@ with gr.Blocks() as demo:
           outputs=slice_slider
       )
     
-    gallery.select(on_select, None, None)
+    gallery.select(on_select, inputs=None, outputs=None)
 
     slice_slider.change(update_slice, inputs=slice_slider, outputs=[image_display, prob_plot, gallery])
     run_button.click(run_mapping, outputs=[prob_plot, plot_row, download_button, download_step])