lenscorrection/DistortionCorrectionTask.java

   1 /**
   2  *
   3  */
   4 package lenscorrection;
   5
   6 import java.awt.Rectangle;
   7 import java.awt.geom.AffineTransform;
   8 import java.util.ArrayList;
   9 import java.util.Collection;
  10 import java.util.List;
  11 import java.util.Set;
  12 import java.util.concurrent.atomic.AtomicInteger;
  13
  14 import lenscorrection.Distortion_Correction.BasicParam;
  15 import mpicbg.models.PointMatch;
  16 import mpicbg.models.Tile;
  17 import mpicbg.trakem2.align.AbstractAffineTile2D;
  18 import mpicbg.trakem2.align.Align;
  19 import mpicbg.trakem2.transform.CoordinateTransform;
  20
  21 import ij.IJ;
  22 import ij.gui.GenericDialog;
  23 import ini.trakem2.display.Display;
  24 import ini.trakem2.display.Displayable;
  25 import ini.trakem2.display.Layer;
  26 import ini.trakem2.display.Patch;
  27 import ini.trakem2.display.Selection;
  28 import ini.trakem2.utils.Worker;
  29 import ini.trakem2.utils.Bureaucrat;
  30 import ini.trakem2.utils.IJError;
  31 import ini.trakem2.utils.Utils;
  32
  33 /**
  34  * Methods collection to be called from the GUI for alignment tasks.
  35  *
  36  */
  37 final public class DistortionCorrectionTask
  38 {
  39         static public class CorrectDistortionFromSelectionParam extends BasicParam
  40         {
  41                 public int firstLayerIndex;
  42                 public int lastLayerIndex;
  43                 public boolean clearTransform = false;
  44                 public boolean visualize = false;
  45
  46                 public void addFields( final GenericDialog gd, final Selection selection )
  47                 {
  48                         addFields( gd );
  49                         gd.addMessage( "Apply Distortion Correction :" );
  50
  51                         Utils.addLayerRangeChoices( selection.getLayer(), gd );
  52                         gd.addCheckbox( "clear_present_transforms", clearTransform );
  53                         gd.addCheckbox( "visualize_distortion_model", visualize );
  54                 }
  55
  56                 @Override
  57                 public boolean readFields( final GenericDialog gd )
  58                 {
  59                         super.readFields( gd );
  60                         firstLayerIndex = gd.getNextChoiceIndex();
  61                         lastLayerIndex = gd.getNextChoiceIndex();
  62                         clearTransform = gd.getNextBoolean();
  63                         visualize = gd.getNextBoolean();
  64                         return !gd.invalidNumber();
  65                 }
  66
  67                 public boolean setup( final Selection selection )
  68                 {
  69                         final GenericDialog gd = new GenericDialog( "Distortion Correction" );
  70                         addFields( gd, selection );
  71                         do
  72                         {
  73                                 gd.showDialog();
  74                                 if ( gd.wasCanceled() ) return false;
  75                         }
  76                         while ( !readFields( gd ) );
  77
  78                         return true;
  79                 }
  80
  81                 @Override
  82                 public CorrectDistortionFromSelectionParam clone()
  83                 {
  84                         final CorrectDistortionFromSelectionParam p = new CorrectDistortionFromSelectionParam();
  85                         p.sift.set( sift );
  86                         p.dimension = dimension;
  87                         p.expectedModelIndex = expectedModelIndex;
  88                         p.lambda = lambda;
  89                         p.maxEpsilon = maxEpsilon;
  90                         p.minInlierRatio = minInlierRatio;
  91                         p.rod = rod;
  92                         p.firstLayerIndex = firstLayerIndex;
  93                         p.lastLayerIndex = lastLayerIndex;
  94                         p.clearTransform = clearTransform;
  95                         p.visualize = visualize;
  96
  97                         return p;
  98                 }
  99         }
 100
 101
 102         /**
 103          * Sets a {@link CoordinateTransform} to a list of patches.
 104          */
 105         final static protected class SetCoordinateTransformThread extends Thread
 106         {
 107                 final protected List< Patch > patches;
 108                 final protected CoordinateTransform transform;
 109                 final protected AtomicInteger ai;
 110
 111                 public SetCoordinateTransformThread(
 112                                 final List< Patch > patches,
 113                                 final CoordinateTransform transform,
 114                                 final AtomicInteger ai )
 115                 {
 116                         this.patches = patches;
 117                         this.transform = transform;
 118                         this.ai = ai;
 119                 }
 120
 121                 @Override
 122                 final public void run()
 123                 {
 124                         for ( int i = ai.getAndIncrement(); i < patches.size() && !isInterrupted(); i = ai.getAndIncrement() )
 125                         {
 126                                 final Patch patch = patches.get( i );
 127                                 //IJ.log( "Setting transform \"" + transform + "\" for patch \"" + patch.getTitle() + "\"." );
 128                                 patch.setCoordinateTransform( transform );
 129                                 patch.updateMipmaps();
 130
 131                                 IJ.showProgress( i, patches.size() );
 132                         }
 133                 }
 134         }
 135
 136         final static protected void setCoordinateTransform(
 137                         final List< Patch > patches,
 138                         final CoordinateTransform transform,
 139                         final int numThreads )
 140         {
 141                 final AtomicInteger ai = new AtomicInteger( 0 );
 142                 final List< SetCoordinateTransformThread > threads = new ArrayList< SetCoordinateTransformThread >();
 143
 144                 for ( int i = 0; i < numThreads; ++i )
 145                 {
 146                         final SetCoordinateTransformThread thread = new SetCoordinateTransformThread( patches, transform, ai );
 147                         threads.add( thread );
 148                         thread.start();
 149                 }
 150                 try
 151                 {
 152                         for ( final Thread thread : threads )
 153                                 thread.join();
 154                 }
 155                 catch ( InterruptedException e )
 156                 {
 157                         IJ.log( "Setting CoordinateTransform failed.\n" + e.getMessage() + "\n" + e.getStackTrace() );
 158                 }
 159         }
 160
 161
 162         /**
 163          * Appends a {@link CoordinateTransform} to a list of patches.
 164          */
 165         final static protected class AppendCoordinateTransformThread extends Thread
 166         {
 167                 final protected List< Patch > patches;
 168                 final protected CoordinateTransform transform;
 169                 final protected AtomicInteger ai;
 170
 171                 public AppendCoordinateTransformThread(
 172                                 final List< Patch > patches,
 173                                 final CoordinateTransform transform,
 174                                 final AtomicInteger ai )
 175                 {
 176                         this.patches = patches;
 177                         this.transform = transform;
 178                         this.ai = ai;
 179                 }
 180
 181                 @Override
 182                 final public void run()
 183                 {
 184                         for ( int i = ai.getAndIncrement(); i < patches.size() && !isInterrupted(); i = ai.getAndIncrement() )
 185                         {
 186                                 final Patch patch = patches.get( i );
 187                                 patch.appendCoordinateTransform( transform );
 188                                 patch.updateMipmaps();
 189
 190                                 IJ.showProgress( i, patches.size() );
 191                         }
 192                 }
 193         }
 194
 195         final static protected void appendCoordinateTransform(
 196                         final List< Patch > patches,
 197                         final CoordinateTransform transform,
 198                         final int numThreads )
 199         {
 200                 final AtomicInteger ai = new AtomicInteger( 0 );
 201                 final List< AppendCoordinateTransformThread > threads = new ArrayList< AppendCoordinateTransformThread >();
 202
 203                 for ( int i = 0; i < numThreads; ++i )
 204                 {
 205                         final AppendCoordinateTransformThread thread = new AppendCoordinateTransformThread( patches, transform, ai );
 206                         threads.add( thread );
 207                         thread.start();
 208                 }
 209                 try
 210                 {
 211                         for ( final Thread thread : threads )
 212                                 thread.join();
 213                 }
 214                 catch ( InterruptedException e )
 215                 {
 216                         IJ.log( "Appending CoordinateTransform failed.\n" + e.getMessage() + "\n" + e.getStackTrace() );
 217                 }
 218         }
 219
 220         final static public CorrectDistortionFromSelectionParam correctDistortionFromSelectionParam = new CorrectDistortionFromSelectionParam();
 221
 222         final static public Bureaucrat correctDistortionFromSelectionTask ( final Selection selection )
 223         {
 224                 Worker worker = new Worker("Distortion Correction", false, true) {
 225                         public void run() {
 226                                 startedWorking();
 227                                 try {
 228                                         correctDistortionFromSelection( selection );
 229                                         Display.repaint(selection.getLayer());
 230                                 } catch (Throwable e) {
 231                                         IJError.print(e);
 232                                 } finally {
 233                                         finishedWorking();
 234                                 }
 235                         }
 236                         public void cleanup() {
 237                                 if (!selection.isEmpty())
 238                                         selection.getLayer().getParent().undoOneStep();
 239                         }
 240                 };
 241                 return Bureaucrat.createAndStart( worker, selection.getProject() );
 242         }
 243
 244         final static public Bureaucrat correctDistortionFromSelection( final Selection selection )
 245         {
 246                 final List< Patch > patches = new ArrayList< Patch >();
 247                 for ( Displayable d : Display.getFront().getSelection().getSelected() )
 248                         if ( d instanceof Patch ) patches.add( ( Patch )d );
 249
 250                 if ( patches.size() < 2 )
 251                 {
 252                         Utils.log("No images in the selection.");
 253                         return null;
 254                 }
 255
 256                 final Worker worker = new Worker( "Lens correction" )
 257                 {
 258                         final public void run()
 259                         {
 260                                 try
 261                                 {
 262                                         startedWorking();
 263
 264                                         if ( !correctDistortionFromSelectionParam.setup( selection ) ) return;
 265
 266                                         final CorrectDistortionFromSelectionParam p = correctDistortionFromSelectionParam.clone();
 267                                         final Align.ParamOptimize ap = Align.paramOptimize.clone();
 268                                         ap.sift.set( p.sift );
 269                                         ap.desiredModelIndex = ap.expectedModelIndex = p.expectedModelIndex;
 270                                         ap.maxEpsilon = p.maxEpsilon;
 271                                         ap.minInlierRatio = p.minInlierRatio;
 272                                         ap.rod = p.rod;
 273
 274                                         /** Get all patches that will be affected. */
 275                                         final List< Patch > allPatches = new ArrayList< Patch >();
 276                                         for ( final Layer l : selection.getLayer().getParent().getLayers().subList( p.firstLayerIndex, p.lastLayerIndex + 1 ) )
 277                                                 for ( Displayable d : l.getDisplayables( Patch.class ) )
 278                                                         allPatches.add( ( Patch )d );
 279
 280                                         /** Unset the coordinate transforms of all patches if desired. */
 281                                         if ( p.clearTransform )
 282                                         {
 283                                                 setTaskName( "Clearing present transforms" );
 284                                                 setCoordinateTransform( allPatches, null, Runtime.getRuntime().availableProcessors() );
 285                                                 Display.repaint();
 286                                         }
 287
 288                                         setTaskName( "Establishing SIFT correspondences" );
 289
 290                                         List< AbstractAffineTile2D< ? > > tiles = new ArrayList< AbstractAffineTile2D< ? > >();
 291                                         List< AbstractAffineTile2D< ? > > fixedTiles = new ArrayList< AbstractAffineTile2D< ? > > ();
 292                                         List< Patch > fixedPatches = new ArrayList< Patch >();
 293                                         final Displayable active = selection.getActive();
 294                                         if ( active != null && active instanceof Patch )
 295                                                 fixedPatches.add( ( Patch )active );
 296                                         Align.tilesFromPatches( ap, patches, fixedPatches, tiles, fixedTiles );
 297
 298                                         final List< AbstractAffineTile2D< ? >[] > tilePairs = new ArrayList< AbstractAffineTile2D< ? >[] >();
 299
 300                                         /** TODO Verena: Why not each to each? */
 301 //                                      AbstractAffineTile2D.pairTiles( tiles, tilePairs );
 302                                         final AbstractAffineTile2D< ? > fixedTile = fixedTiles.iterator().next();
 303                                         for ( final AbstractAffineTile2D< ? > t : tiles )
 304                                                  if ( t != fixedTile )
 305                                                          tilePairs.add( new AbstractAffineTile2D< ? >[]{ t, fixedTile } );
 306
 307                                         Align.connectTilePairs( ap, tiles, tilePairs, Runtime.getRuntime().availableProcessors() );
 308
 309
 310                                         /** Shift all local coordinates into the original image frame */
 311                                         for ( final AbstractAffineTile2D< ? > tile : tiles )
 312                                         {
 313                                                 final Rectangle box = tile.getPatch().getCoordinateTransformBoundingBox();
 314                                                 for ( final PointMatch m : tile.getMatches() )
 315                                                 {
 316                                                         final float[] l = m.getP1().getL();
 317                                                         final float[] w = m.getP1().getW();
 318                                                         l[ 0 ] += box.x;
 319                                                         l[ 1 ] += box.y;
 320                                                         w[ 0 ] = l[ 0 ];
 321                                                         w[ 1 ] = l[ 1 ];
 322                                                 }
 323                                         }
 324
 325                                         if ( Thread.currentThread().isInterrupted() ) return;
 326
 327                                         List< Set< Tile< ? > > > graphs = AbstractAffineTile2D.identifyConnectedGraphs( tiles );
 328                                         if ( graphs.size() > 1 )
 329                                                 IJ.log( "Could not interconnect all images with correspondences.  " );
 330
 331                                         final List< AbstractAffineTile2D< ? > > interestingTiles;
 332
 333                                         /** Find largest graph. */
 334                                         Set< Tile< ? > > largestGraph = null;
 335                                         for ( Set< Tile< ? > > graph : graphs )
 336                                                 if ( largestGraph == null || largestGraph.size() < graph.size() )
 337                                                         largestGraph = graph;
 338
 339                                         interestingTiles = new ArrayList< AbstractAffineTile2D< ? > >();
 340                                         for ( Tile< ? > t : largestGraph )
 341                                                 interestingTiles.add( ( AbstractAffineTile2D< ? > )t );
 342
 343                                         if ( Thread.currentThread().isInterrupted() ) return;
 344
 345                                         Align.optimizeTileConfiguration( ap, interestingTiles, fixedTiles );
 346
 347                                         /** Some data shuffling for the lens correction interface */
 348                                         final List< Collection< PointMatch > > matches = new ArrayList< Collection< PointMatch > >();
 349                                         final List< AffineTransform > affines = new ArrayList< AffineTransform >();
 350                                         for ( AbstractAffineTile2D< ? >[] tilePair : tilePairs )
 351                                         {
 352                                                 matches.add( tilePair[ 0 ].getMatches() );
 353                                                 final AffineTransform a = tilePair[ 0 ].createAffine();
 354                                                 a.preConcatenate( fixedTile.getModel().createInverseAffine() );
 355                                                 affines.add( a );
 356                                         }
 357
 358                                         setTaskName( "Estimating lens distortion correction" );
 359
 360                                         final NonLinearTransform lensModel = Distortion_Correction.createInverseDistortionModel(
 361                                                 matches,
 362                                                 affines,
 363                                                 p.dimension,
 364                                                 p.lambda,
 365                                                 ( int )fixedTile.getWidth(),
 366                                                 ( int )fixedTile.getHeight() );
 367
 368                                         if ( p.visualize )
 369                                         {
 370                                                 setTaskName( "Visualizing lens distortion correction" );
 371                                                 lensModel.visualizeSmall( p.lambda );
 372                                         }
 373
 374                                         setTaskName( "Applying lens distortion correction" );
 375
 376                                         appendCoordinateTransform( allPatches, lensModel, Runtime.getRuntime().availableProcessors() );
 377
 378                                         IJ.log( "Done." );
 379
 380                                         Display.repaint();
 381                                 }
 382                                 catch ( Exception e ) { IJError.print( e ); }
 383                                 finally { finishedWorking(); }
 384                         }
 385                 };
 386
 387                 return Bureaucrat.createAndStart( worker, selection.getProject() );
 388         }
 389 }