From a90a15db060c0963801ff65cad4e8cc1eaf632b0 Mon Sep 17 00:00:00 2001
From: Ali Khan <alik@robarts.ca>
Date: Sun, 26 Jan 2025 23:10:41 -0500
Subject: [PATCH] have ground truth fusion working now

still need to sort out optimization..
---
 dask-stitch/Snakefile                         |  1 -
 dask-stitch/scripts/assign_translation.py     | 22 +++++---------
 .../scripts/create_test_dataset_singletile.py | 30 ++-----------------
 3 files changed, 11 insertions(+), 42 deletions(-)

diff --git a/dask-stitch/Snakefile b/dask-stitch/Snakefile
index 0790cdd..51d4a83 100644
--- a/dask-stitch/Snakefile
+++ b/dask-stitch/Snakefile
@@ -25,7 +25,6 @@ rule create_test_dataset_single_ome_zarr:
     output: 
         ome_zarr=directory('results/tile-{tile}_SPIM.ome.zarr'),
         nifti='results/tile-{tile}_SPIM.nii',
-        nifti_fromzarr='results/tile-{tile}_fromzarr_SPIM.nii',
         true_offset='results/tile-{tile}_desc-groundtruth_offset.txt'
     script: 'scripts/create_test_dataset_singletile.py'
 
diff --git a/dask-stitch/scripts/assign_translation.py b/dask-stitch/scripts/assign_translation.py
index bd05abf..2e99fd5 100644
--- a/dask-stitch/scripts/assign_translation.py
+++ b/dask-stitch/scripts/assign_translation.py
@@ -1,6 +1,6 @@
 
 import numpy as np
-from zarrnii import ZarrNii
+from zarrnii import ZarrNii, AffineTransform
 
 def assign_translations(ome_zarr_paths, translations, output_niftis, output_ome_zarrs):
     """
@@ -20,21 +20,15 @@ def assign_translations(ome_zarr_paths, translations, output_niftis, output_ome_
         print(translation)
         # Load the OME-Zarr dataset
         znimg = ZarrNii.from_ome_zarr(path)
-        # Update the affine matrix
-        updated_affine = znimg.affine
-        print(f'original affine: {updated_affine}')
-        
-        #HACK FIX:
-        #affine has translation negated and flipped, for whatever reason..  need to fix this in zarrnii
-        #updated_affine[:3,3] = -1 * np.flip(updated_affine[:3,3])
 
+        affine = znimg.affine
+   
+        if znimg.axes_order == 'ZYX':
+            affine[:3, 3] += translation[::-1]  # Add the optimized translation
+        else:
+            affine[:3, 3] += translation  # Add the optimized translation
 
-
-
-        updated_affine[:3, 3] += translation  # Add the optimized translation
-        print(f'updated affine: {updated_affine}')
-
-        znimg.affine = updated_affine
+        znimg.affine = affine
 
         # Save the updated ZarrNii
         znimg.to_nifti(out_nii)
diff --git a/dask-stitch/scripts/create_test_dataset_singletile.py b/dask-stitch/scripts/create_test_dataset_singletile.py
index 21542c5..1f75577 100644
--- a/dask-stitch/scripts/create_test_dataset_singletile.py
+++ b/dask-stitch/scripts/create_test_dataset_singletile.py
@@ -68,13 +68,13 @@ def create_test_dataset_single(tile_index, template="MNI152NLin2009cAsym", res=2
 
     # TODO: Simulate error by applying a transformation to the image before 
     
-    random_offset_low=-0
-    random_offset_high=0
+    random_offset_low=-10
+    random_offset_high=10
     # initially lets just do a random jitter 
     offset = np.random.uniform(random_offset_low, random_offset_high, size=(grid_shape[0],grid_shape[1],3))  # Random 3D offsets for each tile
 
     #save this offset to a text file so we know the ground truth
-    np.savetxt(snakemake.output.true_offset, -offset[x,y,:].reshape((1,3)), fmt="%.6f")
+    np.savetxt(snakemake.output.true_offset, offset[x,y,:].reshape((1,3)), fmt="%.6f")
 
 
     xfm_img_data = affine_transform(img_data,matrix=np.eye(3,3),offset=offset[x,y,:],order=3,mode='nearest')
@@ -124,27 +124,3 @@ def create_test_dataset_single(tile_index, template="MNI152NLin2009cAsym", res=2
 
 
 
-print("LOOP0")
-print(test_znimg)
-
-test_znimg2 = ZarrNii.from_ome_zarr(snakemake.output.ome_zarr)
-print("LOOP1")
-print(test_znimg2) #we want this to be identical to test_znimg
-
-
-
-test_znimg2.to_nifti(snakemake.output.nifti_fromzarr)
-
-"""
-
-test_znimg2.to_ome_zarr('mytest_writeafterread.ome.zarr')
-test_znimg3 = ZarrNii.from_path('mytest_writeafterread.ome.zarr')
-print("LOOP2")
-print(test_znimg3)
-
-
-test_znimg3.to_ome_zarr('mytest_writeafterread2.ome.zarr')
-test_znimg4 = ZarrNii.from_path('mytest_writeafterread2.ome.zarr')
-print("LOOP3")
-print(test_znimg4)
-"""