Skip to content

Instantly share code, notes, and snippets.

@biochem-fan
Last active September 16, 2023 23:37
Show Gist options
  • Star 3 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save biochem-fan/76ff5934495585da56ab2c0af4fe2563 to your computer and use it in GitHub Desktop.
Save biochem-fan/76ff5934495585da56ab2c0af4fe2563 to your computer and use it in GitHub Desktop.
WarpConsole
/* BoxNet2D runner
* Based on Warp 1.0.7 by Dimitry Tegunov https://github.com/cramerlab/warp/
* Command line interface and Linux port by @biochem_fan
* Licensed under GPLv3
*/
using System;
using System.IO;
using System.Linq; // for toList()
using System.Diagnostics; // for Stopwatch
using System.Collections.Generic;
using Warp;
using Warp.Headers;
using Warp.Tools;
using CommandLine;
using CommandLine.Text;
class Options
{
[Option('i', Required = true, HelpText = "Movie file names to process")]
public IEnumerable<string> InputFileNames { get; set; }
[Option("out", Required = true, HelpText = "Output directory for picked coordinates or a trained model")]
public string OutDir { get; set; }
[Option("no_xml", Default=false, HelpText = "Don't use XML file")]
public bool NoXML { get; set; }
[Option("model", Required = true, HelpText = "BoxNet2D model directory (READ)")]
public string ModelDir { get; set; }
[Option("min_distance", Default = 0, HelpText = "Minimum mask distance")]
public double MinimumMaskDistance { get; set; }
[Option("min_score", Default = 0.95, HelpText = "Minimum score for picking")]
public double MinimumScore { get; set; }
[Option("particle_diameter", Default = 200.0, HelpText = "Particle diameter in Angstrom")]
public double ExpectedDiameter { get; set; }
[Option("train", Default=false, HelpText = "Do training, instead of picking")]
public bool DoTrain { get; set; }
[Option("filament", Default=false, HelpText = "This is filament")]
public bool IsFilament { get; set; }
[Option("corpus", Default="", HelpText = "BoxNet2D corpus directory (READ)")]
public string CorpusDir { get; set; }
[Option("folder_positive", Default = "positive/", HelpText = "Positive example directory")]
public string PositiveDir { get; set; }
[Option("suffix_positive", Default = "_positive", HelpText = "Suffix for positive examples")]
public string SuffixPositive { get; set; }
[Option("folder_false_positive", Default = "false_positive", HelpText = "False positive example directory")]
public string FalsePositiveDir { get; set; }
[Option("suffix_false_positive", Default = "false_positive", HelpText = "Suffix for false positive examples")]
public string SuffixFalsePositive { get; set; }
[Option("folder_uncertain", Default = "uncertain", HelpText = "Uncertain example directory")]
public string UncertainDir { get; set; }
[Option("suffix_uncertain", Default = "uncertain", HelpText = "Suffix for uncertain examples")]
public string SuffixUncertain { get; set; }
[Option("negative", Default=false, HelpText = "Negative stain?")]
public bool IsNegative { get; set; }
[Option("train_masking", Default=false, HelpText = "Train masking?")]
public bool TrainMasking { get; set; }
}
class Program
{
static bool UseCorpus, TrainMasking, IsNegative;
static float Diameter;
static string PotentialNewName;
static string FolderPositive, SuffixPositive;
static string FolderFalsePositive, SuffixFalsePositive;
static string FolderUncertain, SuffixUncertain;
static void Main(string[] args)
{
Console.WriteLine("============================================================");
Console.WriteLine();
Console.WriteLine(" BoxNet2D runner");
Console.WriteLine(" Based on Warp 1.0.7 by Dimitry Tegunov");
Console.WriteLine(" Command line interface and Linux port by @biochem_fan");
Console.WriteLine();
Options Options = new Options();
var errors = new List<CommandLine.Error>();
CommandLine.Parser.Default.ParseArguments<Options>(args)
.WithNotParsed(x => errors = x.ToList())
.WithParsed( x => Options = x);
Console.WriteLine("============================================================");
Console.WriteLine(" Input parameters:");
Console.WriteLine($" Input file name: {Options.InputFileNames.First()} ...");
Console.WriteLine($" Inputs are without XML: {Options.NoXML}");
Console.WriteLine($" Pre-trained Model (READ): {Options.ModelDir}");
Console.WriteLine($" Minimum score: {Options.MinimumScore}");
Console.WriteLine($" Particle diameter: {Options.ExpectedDiameter}");
Console.WriteLine($" Minimum distance from bad area: {Options.MinimumMaskDistance}");
Console.WriteLine();
Console.WriteLine($" Train, instead of pick: {Options.DoTrain}");
Console.WriteLine($" Model output (WRITTEN): {Options.OutDir}");
Console.WriteLine($" Train filaments: {Options.IsFilament}");
Console.WriteLine($" Positive example directory: {Options.PositiveDir}");
Console.WriteLine($" Positive example suffix: {Options.SuffixPositive}");
Console.WriteLine($" False positive example directory: {Options.FalsePositiveDir}");
Console.WriteLine($" False positive example suffix: {Options.SuffixFalsePositive}");
Console.WriteLine($" Uncertain example directory: {Options.UncertainDir}");
Console.WriteLine($" Uncertain example suffix: {Options.SuffixUncertain}");
Console.WriteLine($" Is negative stain: {Options.IsNegative}");
Console.WriteLine($" Train masking: {Options.TrainMasking}");
Console.WriteLine("============================================================");
Console.WriteLine();
// TODO: Need a STAR file for visuallisation.
//
// DONE: Needs support from RELION; scale CoordinateX/Y. => tentatively implemented --cord_scale
// relion_manualpick --i denoised.star --coord_scale 0.138 --odir warp/matching/ --pickname BoxNet2_20180602 --sigma_contrast 3 --particle_diameter 140
// Remove trailing "/" from ModelName because this becomes the suffix for coordinate STAR files
Options.ModelDir = Options.ModelDir.TrimEnd('/');
Options.OutDir = Options.OutDir.TrimEnd('/');
ProcessingOptionsBoxNet optionsBoxNet = new ProcessingOptionsBoxNet();
optionsBoxNet.PickingInvert = false;
optionsBoxNet.ExpectedDiameter = (decimal)Options.ExpectedDiameter;
optionsBoxNet.MinimumMaskDistance = (decimal)Options.MinimumMaskDistance;
optionsBoxNet.MinimumScore = (decimal)Options.MinimumScore;
optionsBoxNet.ModelName = Options.ModelDir;
List<Movie> ValidInputs = new List<Movie>();
if (Options.InputFileNames.Count() == 1 && Options.InputFileNames.First().EndsWith("lst"))
{
Console.WriteLine($"Input is a list file.");
StreamReader sr = new StreamReader(Options.InputFileNames.First());
while (!sr.EndOfStream)
{
string line = sr.ReadLine().Trim();
Movie movie = new Movie(Path.GetFullPath(line));
string ImagePath = (Options.NoXML) ? movie.Path : movie.AveragePath;
Console.WriteLine($"File {ImagePath}");
if (File.Exists(ImagePath))
ValidInputs.Add(movie);
}
}
else
{
foreach (string MovieFileName in Options.InputFileNames)
{
Console.WriteLine($"Input {MovieFileName}");
Movie movie = new Movie(Path.GetFullPath(MovieFileName));
string ImagePath = (Options.NoXML) ? movie.Path : movie.AveragePath;
if (File.Exists(ImagePath))
ValidInputs.Add(movie);
}
}
const int gpu = 0;
BoxNet2 BoxNetwork = new BoxNet2(Options.ModelDir, gpu, 2, 1, false);
Uri UriPwd = new Uri(Environment.CurrentDirectory + "/");
foreach (Movie movie in ValidInputs)
{
string ImagePath = (Options.NoXML) ? movie.Path : movie.AveragePath;
string RelativeImagePath = UriPwd.MakeRelativeUri(new Uri(ImagePath)).ToString();
Console.WriteLine($"Processing {RelativeImagePath}");
Image stack = Image.FromFile(ImagePath);
movie.MatchBoxNet2(new[] { BoxNetwork }, stack, optionsBoxNet, null);
stack.Dispose();
// For RELION, the outputs must be in OutDir/Movies/average/XXX or OutDir/Movies/denoised/XXX
// Unfortunately MatchBoxNet2 writes to Movies/matching/XXX and this is hard-coded in Movie.cs
string CoordinateName = movie.RootName + "_" + Helper.PathToNameWithExtension(Options.ModelDir) + ".star";
string WrittenPath = movie.MatchingDir + CoordinateName;
WrittenPath = UriPwd.MakeRelativeUri(new Uri(WrittenPath)).ToString();
string DestinationPath = Options.OutDir + "/" + Helper.PathToFolder(RelativeImagePath) + CoordinateName;
Directory.CreateDirectory(Helper.PathToFolder(DestinationPath));
if (File.Exists(DestinationPath))
File.Delete(DestinationPath);
File.Move(WrittenPath, DestinationPath);
// Link OutDir/Movies/denoised to OutDir/Movies/average (for visuallisation)
if (Options.NoXML) continue;
string RelativeMoviePath = UriPwd.MakeRelativeUri(new Uri(movie.DirectoryName)).ToString();
string DenoisedPath = RelativeMoviePath + "denoised";
System.Diagnostics.Process Process = new System.Diagnostics.Process();
Process.StartInfo.FileName = "/bin/ln";
Process.StartInfo.Arguments = $"-sf average {Options.OutDir + "/" + DenoisedPath}";
Process.Start();
Process.WaitForExit();
// TODO: Write a STAR file
}
}
// Adapted from Warp/Controls/BoxNet/BoxNetTrain.xaml.cs
static void Train(Options Options, List<Movie> Movies)
{
int gpu = 0;
UseCorpus = (Options.CorpusDir != "");
TrainMasking = (bool)Options.TrainMasking;
Diameter = (float)Options.ExpectedDiameter;
IsNegative = Options.IsNegative;
PotentialNewName = Options.OutDir;
FolderPositive = Options.PositiveDir;
SuffixPositive = Options.SuffixPositive;
FolderFalsePositive = Options.FalsePositiveDir;
SuffixFalsePositive = Options.SuffixFalsePositive;
FolderUncertain = Options.UncertainDir;
SuffixUncertain = Options.SuffixUncertain;
GPU.SetDevice(0);
#region Prepare new examples
Console.WriteLine("+ Preparing new examples...");
Directory.CreateDirectory(Movies[0].DirectoryName + PotentialNewName + "_training/");
string NewExamplesPath = Movies[0].DirectoryName + PotentialNewName + "_training/" + PotentialNewName + ".tif";
//try
{
PrepareData(NewExamplesPath, Movies, Options.NoXML, Options.IsFilament);
}
/*catch (Exception exc)
{
Console.WriteLine("Failed to prepare data");
return;
}*/
#endregion
#region Load background and new examples
Console.WriteLine("Loading examples...");
int2 DimsLargest = new int2(1);
List<float[]>[] AllMicrographs = { new List<float[]>(), new List<float[]>() };
List<float[]>[] AllLabels = { new List<float[]>(), new List<float[]>() };
List<float[]>[] AllUncertains = { new List<float[]>(), new List<float[]>() };
List<int2>[] AllDims = { new List<int2>(), new List<int2>() };
List<float3>[] AllLabelWeights = { new List<float3>(), new List<float3>() };
string[][] AllPaths = UseCorpus
? new[]
{
new[] { NewExamplesPath },
Directory.EnumerateFiles(Options.CorpusDir, "*.tif").ToArray()
}
: new[] { new[] { NewExamplesPath } };
long[] ClassHist = new long[3];
for (int icorpus = 0; icorpus < AllPaths.Length; icorpus++)
{
foreach (var examplePath in AllPaths[icorpus])
{
Image ExampleImage = Image.FromFile(examplePath);
int N = ExampleImage.Dims.Z / 3;
for (int n = 0; n < N; n++)
{
float[] Mic = ExampleImage.GetHost(Intent.Read)[n * 3 + 0];
MathHelper.FitAndSubtractGrid(Mic, new int2(ExampleImage.Dims), new int2(4));
AllMicrographs[icorpus].Add(Mic);
AllLabels[icorpus].Add(ExampleImage.GetHost(Intent.Read)[n * 3 + 1]);
AllUncertains[icorpus].Add(ExampleImage.GetHost(Intent.Read)[n * 3 + 2]);
AllDims[icorpus].Add(new int2(ExampleImage.Dims));
float[] Labels = ExampleImage.GetHost(Intent.Read)[n * 3 + 1];
float[] Uncertains = ExampleImage.GetHost(Intent.Read)[n * 3 + 2];
for (int i = 0; i < Labels.Length; i++)
{
int Label = (int)Labels[i];
if (!TrainMasking && Label == 2)
{
Label = 0;
Labels[i] = 0;
}
ClassHist[Label]++;
}
}
DimsLargest.X = Math.Max(DimsLargest.X, ExampleImage.Dims.X);
DimsLargest.Y = Math.Max(DimsLargest.Y, ExampleImage.Dims.Y);
ExampleImage.Dispose();
}
}
{
float[] LabelWeights = { 1f, 1f, 1f };
LabelWeights[0] = (float)Math.Pow((float)ClassHist[1] / ClassHist[0], 1 / 3.0);
LabelWeights[2] = 1;//(float)Math.Sqrt((float)ClassHist[1] / ClassHist[2]);
for (int icorpus = 0; icorpus < AllPaths.Length; icorpus++)
for (int i = 0; i < AllMicrographs[icorpus].Count; i++)
AllLabelWeights[icorpus].Add(new float3(LabelWeights[0], LabelWeights[1], LabelWeights[2]));
}
int NNewExamples = AllMicrographs[0].Count;
int NOldExamples = UseCorpus ? AllMicrographs[1].Count : 0;
#endregion
#region Load models
Console.WriteLine("+ Loading old BoxNet model...");
int NThreads = 2;
BoxNet2 NetworkTrain = null;
try
{
NetworkTrain = new BoxNet2(Options.ModelDir, GPU.GetDeviceCount() - 1, NThreads, 8, true);
}
catch // It might be an old version of BoxNet that doesn't support batch size != 1
{
Console.WriteLine("Old BoxNet2D");
NetworkTrain = new BoxNet2(Options.ModelDir, GPU.GetDeviceCount() - 1, NThreads, 1, true);
}
//BoxNet2 NetworkOld = new BoxNet2(Options.MainWindow.LocatePickingModel(ModelName), (GPU.GetDeviceCount() * 2 - 2) % GPU.GetDeviceCount(), NThreads, 1, false);
#endregion
#region Training
Console.WriteLine("+ Training...");
int2 DimsAugmented = BoxNet2.BoxDimensionsTrain;
int Border = (BoxNet2.BoxDimensionsTrain.X - BoxNet2.BoxDimensionsValidTrain.X) / 2;
int BatchSize = NetworkTrain.BatchSize;
int PlotEveryN = 10;
int SmoothN = 30;
//List<ObservablePoint>[] AccuracyPoints = Helper.ArrayOfFunction(i => new List<ObservablePoint>(), 4);
Queue<float>[] LastAccuracies = { new Queue<float>(SmoothN), new Queue<float>(SmoothN) };
List<float>[] LastBaselines = { new List<float>(), new List<float>() };
GPU.SetDevice(0);
IntPtr d_MaskUncertain;
{
float[] DataUncertain = new float[DimsAugmented.Elements()];
for (int y = 0; y < DimsAugmented.Y; y++)
{
for (int x = 0; x < DimsAugmented.X; x++)
{
if (x >= Border &&
y >= Border &&
x < DimsAugmented.X - Border &&
y < DimsAugmented.Y - Border)
DataUncertain[y * DimsAugmented.X + x] = 1;
else
DataUncertain[y * DimsAugmented.X + x] = 0.1f;
}
}
d_MaskUncertain = GPU.MallocDeviceFromHost(DataUncertain, DataUncertain.Length);
}
IntPtr[] d_OriData = Helper.ArrayOfFunction(i => GPU.MallocDevice(DimsLargest.Elements()), NetworkTrain.MaxThreads);
IntPtr[] d_OriLabels = Helper.ArrayOfFunction(i => GPU.MallocDevice(DimsLargest.Elements()), NetworkTrain.MaxThreads);
IntPtr[] d_OriUncertains = Helper.ArrayOfFunction(i => GPU.MallocDevice(DimsLargest.Elements()), NetworkTrain.MaxThreads);
IntPtr[] d_AugmentedData = Helper.ArrayOfFunction(i => GPU.MallocDevice(DimsAugmented.Elements() * BatchSize), NetworkTrain.MaxThreads);
IntPtr[] d_AugmentedLabels = Helper.ArrayOfFunction(i => GPU.MallocDevice(DimsAugmented.Elements() * BatchSize * 3), NetworkTrain.MaxThreads);
IntPtr[] d_AugmentedWeights = Helper.ArrayOfFunction(i => GPU.MallocDevice(DimsAugmented.Elements() * BatchSize), NetworkTrain.MaxThreads);
Stopwatch Watch = new Stopwatch();
Watch.Start();
Random[] RG = Helper.ArrayOfFunction(i => new Random(i), NetworkTrain.MaxThreads);
RandomNormal[] RGN = Helper.ArrayOfFunction(i => new RandomNormal(i), NetworkTrain.MaxThreads);
//float[][] h_AugmentedData = Helper.ArrayOfFunction(i => new float[DimsAugmented.Elements()], NetworkTrain.MaxThreads);
//float[][] h_AugmentedLabels = Helper.ArrayOfFunction(i => new float[DimsAugmented.Elements()], NetworkTrain.MaxThreads);
//float[][] h_AugmentedWeights = Helper.ArrayOfFunction(i => new float[DimsAugmented.Elements()], NetworkTrain.MaxThreads);
//float[][] LabelsOneHot = Helper.ArrayOfFunction(i => new float[DimsAugmented.Elements() * 3], NetworkTrain.MaxThreads);
int NIterations = NNewExamples * 100 * AllMicrographs.Length;
int NDone = 0;
Helper.ForCPU(0, NIterations, NetworkTrain.MaxThreads,
threadID => GPU.SetDevice(0),
(b, threadID) =>
{
int icorpus;
lock (Watch)
icorpus = NDone % AllPaths.Length;
float2[] PositionsGround;
for (int ib = 0; ib < BatchSize; ib++)
{
int ExampleID = RG[threadID].Next(AllMicrographs[icorpus].Count);
int2 Dims = AllDims[icorpus][ExampleID];
float2[] Translations = Helper.ArrayOfFunction(x => new float2(RG[threadID].Next(Dims.X - Border * 2) + Border - DimsAugmented.X / 2,
RG[threadID].Next(Dims.Y - Border * 2) + Border - DimsAugmented.Y / 2), 1);
float[] Rotations = Helper.ArrayOfFunction(i => (float)(RG[threadID].NextDouble() * Math.PI * 2), 1);
float3[] Scales = Helper.ArrayOfFunction(i => new float3(0.8f + (float)RG[threadID].NextDouble() * 0.4f,
0.8f + (float)RG[threadID].NextDouble() * 0.4f,
(float)(RG[threadID].NextDouble() * Math.PI * 2)), 1);
float StdDev = (float)Math.Abs(RGN[threadID].NextSingle(0, 0.3f));
float[] DataMicrograph = AllMicrographs[icorpus][ExampleID];
float[] DataLabels = AllLabels[icorpus][ExampleID];
float[] DataUncertains = AllUncertains[icorpus][ExampleID];
GPU.CopyHostToDevice(DataMicrograph, d_OriData[threadID], Dims.Elements());
GPU.CopyHostToDevice(DataLabels, d_OriLabels[threadID], Dims.Elements());
GPU.CopyHostToDevice(DataUncertains, d_OriUncertains[threadID], Dims.Elements());
//GPU.ValueFill(d_OriUncertains[threadID], Dims.Elements(), 1f);
GPU.BoxNet2Augment(d_OriData[threadID],
d_OriLabels[threadID],
d_OriUncertains[threadID],
Dims,
new IntPtr((long)d_AugmentedData[threadID] + ib * DimsAugmented.Elements() * sizeof(float)),
new IntPtr((long)d_AugmentedLabels[threadID] + ib * DimsAugmented.Elements() * 3 * sizeof(float)),
new IntPtr((long)d_AugmentedWeights[threadID] + ib * DimsAugmented.Elements() * sizeof(float)),
DimsAugmented,
AllLabelWeights[icorpus][ExampleID],
Helper.ToInterleaved(Translations),
Rotations,
Helper.ToInterleaved(Scales),
StdDev,
RG[threadID].Next(99999),
(uint)1);
}
GPU.MultiplySlices(d_AugmentedWeights[threadID], d_MaskUncertain,
d_AugmentedWeights[threadID], DimsAugmented.Elements(), (uint)BatchSize);
float LearningRate = 0.00002f;
long[][] ResultLabels = new long[2][];
float[][] ResultProbabilities = new float[2][];
float Loss = 0;
lock (NetworkTrain)
Loss = NetworkTrain.Train(d_AugmentedData[threadID], d_AugmentedLabels[threadID],
d_AugmentedWeights[threadID], LearningRate,
threadID,
out ResultLabels[0], out ResultProbabilities[0]);
lock (Watch)
{
NDone++;
//if (!float.IsNaN(AccuracyParticles[0]))
{
LastAccuracies[icorpus].Enqueue(Loss);
if (LastAccuracies[icorpus].Count > SmoothN)
LastAccuracies[icorpus].Dequeue();
}
//if (!float.IsNaN(AccuracyParticles[1]))
// LastBaselines[icorpus].Add(AccuracyParticles[1]);
if (NDone % PlotEveryN == 0)
{
for (int iicorpus = 0; iicorpus < AllMicrographs.Length; iicorpus++)
{
Console.WriteLine($"{NDone} / {NIterations}, iicorpus = {iicorpus}, Acc = {MathHelper.Mean(LastAccuracies[iicorpus].Where(v => !float.IsNaN(v)))}");
//AccuracyPoints[iicorpus * 2 + 1].Clear();
//AccuracyPoints[iicorpus * 2 + 1].Add(new ObservablePoint(0,
// MathHelper.Mean(LastBaselines[iicorpus].Where(v => !float.IsNaN(v)))));
//AccuracyPoints[iicorpus * 2 + 1].Add(new ObservablePoint((float)NDone / NIterations * 100,
// MathHelper.Mean(LastBaselines[iicorpus].Where(v => !float.IsNaN(v)))));
}
long Elapsed = Watch.ElapsedMilliseconds;
double Estimated = (double)Elapsed / NDone * NIterations;
int Remaining = (int)(Estimated - Elapsed);
TimeSpan SpanRemaining = new TimeSpan(0, 0, 0, 0, Remaining);
Console.WriteLine(SpanRemaining.ToString((int)SpanRemaining.TotalHours > 0 ? @"hh\:mm\:ss" : @"mm\:ss"));
}
}
}, null);
foreach (var ptr in d_OriData)
GPU.FreeDevice(ptr);
foreach (var ptr in d_OriLabels)
GPU.FreeDevice(ptr);
foreach (var ptr in d_OriUncertains)
GPU.FreeDevice(ptr);
foreach (var ptr in d_AugmentedData)
GPU.FreeDevice(ptr);
foreach (var ptr in d_AugmentedLabels)
GPU.FreeDevice(ptr);
foreach (var ptr in d_AugmentedWeights)
GPU.FreeDevice(ptr);
GPU.FreeDevice(d_MaskUncertain);
#endregion
Console.WriteLine("+ Saving new BoxNet model...");
string BoxNetDir = System.IO.Path.Combine(Environment.CurrentDirectory, Options.OutDir);
Directory.CreateDirectory(BoxNetDir);
NetworkTrain.Export(BoxNetDir);
NetworkTrain.Dispose();
//NetworkOld.Dispose();
}
// Adapted from Warp/Controls/BoxNet/BoxNetTrain.xaml.cs
static void PrepareData(string savePath, List<Movie> Movies, bool noXML, bool IsFilament)
{
string ErrorString = "";
if (string.IsNullOrEmpty(SuffixPositive))
ErrorString += "No positive examples selected.\n";
if (!string.IsNullOrEmpty(ErrorString))
throw new Exception(ErrorString);
List<Movie> ValidMovies = new List<Movie>();
foreach (var movie in Movies)
{
if (movie.UnselectManual != null && (bool)movie.UnselectManual)
continue;
string ImagePath = (noXML) ? movie.Path : movie.AveragePath;
if (!File.Exists(ImagePath))
continue;
Console.WriteLine($"Positive Examples should be {FolderPositive + movie.RootName + SuffixPositive + ".star"}");
bool HasExamples = File.Exists(FolderPositive + movie.RootName + SuffixPositive + ".star");
if (!HasExamples)
continue;
ValidMovies.Add(movie);
}
if (ValidMovies.Count == 0)
throw new Exception("No movie averages could be found to create training examples. Please process the movies first to create the averages.");
List<Image> AllAveragesBN = new List<Image>();
List<Image> AllLabelsBN = new List<Image>();
List<Image> AllCertainBN = new List<Image>();
int MoviesDone = 0;
foreach (var movie in ValidMovies)
{
string ImagePath = (noXML) ? movie.Path : movie.AveragePath;
MapHeader Header = MapHeader.ReadFromFile(ImagePath);
float PixelSize = Header.PixelSize.X;
Console.WriteLine($"Image {ImagePath}: Pixel Size {PixelSize}");
#region Load positions, and possibly move on to next movie
List<float2> PosPositive = new List<float2>();
List<float2> PosFalse = new List<float2>();
List<float2> PosUncertain = new List<float2>();
if (File.Exists(FolderPositive + movie.RootName + SuffixPositive + ".star"))
PosPositive.AddRange(Star.LoadFloat2(FolderPositive + movie.RootName + SuffixPositive + ".star",
"rlnCoordinateX", "rlnCoordinateY")
.Select(v => v * PixelSize / BoxNet2.PixelSize));
if (PosPositive.Count == 0)
continue;
if (File.Exists(FolderFalsePositive + movie.RootName + SuffixFalsePositive + ".star"))
PosFalse.AddRange(Star.LoadFloat2(FolderFalsePositive + movie.RootName + SuffixFalsePositive + ".star",
"rlnCoordinateX", "rlnCoordinateY")
.Select(v => v * PixelSize / BoxNet2.PixelSize));
if (File.Exists(FolderUncertain + movie.RootName + SuffixUncertain + ".star"))
PosUncertain.AddRange(Star.LoadFloat2(FolderUncertain + movie.RootName + SuffixUncertain + ".star",
"rlnCoordinateX", "rlnCoordinateY")
.Select(v => v * PixelSize / BoxNet2.PixelSize));
#endregion
Image Average = Image.FromFile(ImagePath);
int2 Dims = new int2(Average.Dims);
Image Mask = null;
if (File.Exists(movie.MaskPath))
Mask = Image.FromFile(movie.MaskPath);
float RadiusParticle = Math.Max(1, Diameter / 2 / BoxNet2.PixelSize);
float RadiusPeak = Math.Max(1.5f, Diameter / 2 / BoxNet2.PixelSize / 4);
float RadiusFalse = Math.Max(1, Diameter / 2 / BoxNet2.PixelSize);
float RadiusUncertain = Math.Max(1, Diameter / 2 / BoxNet2.PixelSize);
#region Rescale everything and allocate memory
int2 DimsBN = new int2(new float2(Dims) * PixelSize / BoxNet2.PixelSize + 0.5f) / 2 * 2;
Image AverageBN = Average.AsScaled(DimsBN);
Average.Dispose();
if (IsNegative)
AverageBN.Multiply(-1f);
GPU.Normalize(AverageBN.GetDevice(Intent.Read),
AverageBN.GetDevice(Intent.Write),
(uint)AverageBN.ElementsSliceReal,
1);
Image MaskBN = null;
if (Mask != null)
{
MaskBN = Mask.AsScaled(DimsBN);
Mask.Dispose();
}
Image LabelsBN = new Image(new int3(DimsBN));
Image CertainBN = new Image(new int3(DimsBN));
CertainBN.Fill(1f);
#endregion
#region Paint all positive and uncertain peaks
for (int i = 0; i < 3; i++)
{
var ori_positions = (new[] { PosPositive, PosFalse, PosUncertain })[i];
float R = (new[] { RadiusPeak, RadiusFalse, RadiusUncertain })[i];
float R2 = R * R;
float Label = (new[] { 1, 4, 0 })[i];
float[] ImageData = (new[] { LabelsBN.GetHost(Intent.ReadWrite)[0], CertainBN.GetHost(Intent.ReadWrite)[0], CertainBN.GetHost(Intent.ReadWrite)[0] })[i];
var positions = ori_positions;
if (IsFilament)
{
positions = new List<float2>();
for (int ipos = 0, imax = ori_positions.Count; ipos + 1 < imax; ipos += 2)
{
float2 delta = ori_positions[ipos + 1] - ori_positions[ipos];
float dist = (float)Math.Sqrt(delta.X * delta.X + delta.Y * delta.Y);
delta /= dist;
// Console.WriteLine($"ipos {ipos}, {ori_positions[ipos]} - {ori_positions[ipos + 1]} delta {delta} dist {dist}");
const float step = 3;
for (int j = 0; step * j < dist; j++)
{
positions.Add(ori_positions[ipos] + delta * step * j);
}
}
}
if (i == 0) Console.WriteLine($" Number of positions: {positions.Count}");
foreach (var pos in positions)
{
int2 Min = new int2(Math.Max(0, (int)(pos.X - R)), Math.Max(0, (int)(pos.Y - R)));
int2 Max = new int2(Math.Min(DimsBN.X - 1, (int)(pos.X + R)), Math.Min(DimsBN.Y - 1, (int)(pos.Y + R)));
for (int y = Min.Y; y <= Max.Y; y++)
{
float yy = y - pos.Y;
yy *= yy;
for (int x = Min.X; x <= Max.X; x++)
{
float xx = x - pos.X;
xx *= xx;
float r2 = xx + yy;
if (r2 <= R2)
ImageData[y * DimsBN.X + x] = Label;
}
}
}
}
#endregion
#region Add junk mask if there is one
if (MaskBN != null)
{
float[] LabelsBNData = LabelsBN.GetHost(Intent.ReadWrite)[0];
float[] MaskBNData = MaskBN.GetHost(Intent.Read)[0];
for (int i = 0; i < LabelsBNData.Length; i++)
if (MaskBNData[i] > 0.5f)
LabelsBNData[i] = 2;
}
#endregion
#region Clean up
MaskBN?.Dispose();
AllAveragesBN.Add(AverageBN);
AverageBN.FreeDevice();
AllLabelsBN.Add(LabelsBN);
LabelsBN.FreeDevice();
AllCertainBN.Add(CertainBN);
CertainBN.FreeDevice();
#endregion
}
#region Figure out smallest dimensions that contain everything
int2 DimsCommon = new int2(1);
foreach (var image in AllAveragesBN)
{
DimsCommon.X = Math.Max(DimsCommon.X, image.Dims.X);
DimsCommon.Y = Math.Max(DimsCommon.Y, image.Dims.Y);
}
#endregion
#region Put everything in one stack and save
Image Everything = new Image(new int3(DimsCommon.X, DimsCommon.Y, AllAveragesBN.Count * 3));
float[][] EverythingData = Everything.GetHost(Intent.ReadWrite);
for (int i = 0; i < AllAveragesBN.Count; i++)
{
if (AllAveragesBN[i].Dims == new int3(DimsCommon)) // No padding needed
{
EverythingData[i * 3 + 0] = AllAveragesBN[i].GetHost(Intent.Read)[0];
EverythingData[i * 3 + 1] = AllLabelsBN[i].GetHost(Intent.Read)[0];
EverythingData[i * 3 + 2] = AllCertainBN[i].GetHost(Intent.Read)[0];
}
else // Padding needed
{
{
Image Padded = AllAveragesBN[i].AsPadded(DimsCommon);
AllAveragesBN[i].Dispose();
EverythingData[i * 3 + 0] = Padded.GetHost(Intent.Read)[0];
Padded.FreeDevice();
}
{
Image Padded = AllLabelsBN[i].AsPadded(DimsCommon);
AllLabelsBN[i].Dispose();
EverythingData[i * 3 + 1] = Padded.GetHost(Intent.Read)[0];
Padded.FreeDevice();
}
{
Image Padded = AllCertainBN[i].AsPadded(DimsCommon);
AllCertainBN[i].Dispose();
EverythingData[i * 3 + 2] = Padded.GetHost(Intent.Read)[0];
Padded.FreeDevice();
}
}
}
Everything.WriteTIFF(savePath, BoxNet2.PixelSize, typeof(float));
#endregion
}
}
/*
* Warp console
* Based on Warp 1.0.9 by Dimitry Tegunov
* https://github.com/cramerlab/warp/
* Command line interface and Linux port by @biochem_fan
* Licensed under GPLv3
*/
using System;
using System.IO;
using System.Linq; // for toList()
using System.Collections.Generic; // for CultureInfo
using System.Globalization;
using Warp;
using Warp.Headers;
using Warp.Tools;
using CommandLine;
class Options
{
[Option('i', Required = true, HelpText = "Movie file name to process")]
public IEnumerable<string> InputFileNames { get; set; }
[Option("update_defocus", Default = "", HelpText = "Update defocus on this STAR file")]
public string StarIn { get; set; }
[Option("angpix", Default = 1.0, HelpText = "Pixel size of the input (unbinned) micrograph")]
public double PixelSize { get; set; }
[Option("motion_range_min", Default = 0.05, HelpText = "Minimum wavenumber (Nyquist) to consider for motion")]
public double MotionRangeMin { get; set; }
[Option("motion_range_max", Default = 0.25, HelpText = "Maximum wavenumber (Nyquist) to consider for motion")]
public double MotionRangeMax { get; set; }
[Option("bfactor", Default = -500.0, HelpText = "B factor for motion estimation")]
public double Bfactor { get; set; }
[Option("bin", Default = 0, HelpText = "Binning factor (power of 2)")]
public int BinTimes { get; set; }
[Option("gain", Default = "", HelpText = "Gain reference")]
public string GainPath { get; set; }
[Option("gain_flipX", Default = false, HelpText = "Flip the gain reference in X")]
public bool GainFlipX { get; set; }
[Option("gain_flipY", Default = false, HelpText = "Flip the gain reference in Y")]
public bool GainFlipY { get; set; }
[Option("gain_transpose", Default = false, HelpText = "Transpose the gain reference")]
public bool GainTranspose { get; set; }
[Option("motion_grid_x", Default = 5, HelpText = "Number of grid points along X for motion")]
public int MotionGridX { get; set; }
[Option("motion_grid_y", Default = 5, HelpText = "Number of grid points along Y for motion")]
public int MotionGridY { get; set; }
[Option("motion_grid_t", Default = -1, HelpText = "Number of grid points along time for motion")]
public int MotionGridT { get; set; }
[Option("dose", Default = 1.0, HelpText = "Dose per angstrom^2 per frame")]
public double DosePerAngstromFrame { get; set; }
[Option("ctf_window", Default = 512, HelpText = "CTF box size (pixel)")]
public int CtfWindow { get; set; }
[Option("use_movie_sum", Default = true, HelpText = "Use movie sum for CTF")]
public bool UseMovieSum { get; set; }
[Option("ctf_grid_x", Default = 4, HelpText = "Number of grid points along X for CTF")]
public int CtfGridX { get; set; }
[Option("ctf_grid_y", Default = 4, HelpText = "Number of grid points along Y for CTF")]
public int CtfGridY { get; set; }
[Option("ctf_grid_t", Default = 1, HelpText = "Number of grid points along time for CTF")]
public int CtfGridT { get; set; }
[Option("ctf_range_min", Default = 0.10, HelpText = "Minimum wavenumber (Nyquist) to consider for CTF")]
public double CtfRangeMin { get; set; }
[Option("ctf_range_max", Default = 0.50, HelpText = "Maximum wavenumber (Nyquist) to consider for CTF")]
public double CtfRangeMax { get; set; }
[Option("defocus_min", Default = 0.30, HelpText = "Minimum defocus to search (um)")]
public double ZMin { get; set; }
[Option("defocus_max", Default = 4.00, HelpText = "Maximum defocus to search (um)")]
public double ZMax { get; set; }
[Option("cs", Default = 2.7, HelpText = "Spherical aberration (mm)")]
public double Cs { get; set; }
[Option("ac", Default = 0.10, HelpText = "Amplitude contrast")]
public double Amplitude { get; set; }
[Option("voltage", Default = 300, HelpText = "Electron voltage (kV)")]
public int Voltage { get; set; }
[Option("do_phase", Default = false, HelpText = "Estimate phase shift")]
public bool DoPhase { get; set; }
// TODO:
//
// - redo only CTF
// - test gain orientation
// - Specify EER Supersample
// - Treat EER gain properly
// - multiGPU: read gain to each GPU, then use: Helper.ForEachGPU(items, (item, gpuID) =>
}
class OptionsImport
{
public int HeaderlessWidth, HeaderlessHeight;
public string HeaderlessType;
public long HeaderlessOffset;
public bool GainFlipX, GainFlipY, GainTranspose;
public string GainPath;
}
class Program
{
public const int MaxDenoisingTrainData = 128;
static void Main(string[] args)
{
Console.WriteLine("============================================================");
Console.WriteLine();
Console.WriteLine(" Warp Console");
Console.WriteLine(" Based on Warp 1.0.9 by Dimitry Tegunov");
Console.WriteLine(" Command line interface and Linux port by @biochem_fan");
Console.WriteLine();
Options Options = new Options();
Parser.Default.ParseArguments<Options>(args).WithParsed<Options>(opts => Options = opts);
Console.WriteLine("============================================================");
Console.WriteLine(" Input parameters:");
Console.WriteLine($" Input file names: {Options.InputFileNames.First()} ...");
Console.WriteLine($" Pixel size (A/px): {Options.PixelSize}");
Console.WriteLine($" Gain reference: {Options.GainPath}");
Console.WriteLine($" Gain flip X: {Options.GainFlipX}");
Console.WriteLine($" Gain flip X: {Options.GainFlipX}");
Console.WriteLine($" Gain transpose: {Options.GainTranspose}");
Console.WriteLine($" Voltage (kV) {Options.Voltage}");
Console.WriteLine($" Spherical aberration (mm): {Options.Cs}");
Console.WriteLine($" Amplitude contrast: {Options.Amplitude}");
Console.WriteLine();
Console.WriteLine($" Binning: {Options.BinTimes}");
Console.WriteLine($" B factor: {Options.Bfactor}");
Console.WriteLine($" Motion Range (Nyquist): {Options.MotionRangeMin} to {Options.MotionRangeMax}");
Console.WriteLine($" Motion Grid: X={Options.MotionGridX} Y={Options.MotionGridY} T={Options.MotionGridT}");
Console.WriteLine();
Console.WriteLine($" CTF Range (Nyquist): {Options.CtfRangeMin} to {Options.CtfRangeMax}");
Console.WriteLine($" CTF Grid: X={Options.CtfGridX} Y={Options.CtfGridY} T={Options.CtfGridT}");
Console.WriteLine($" Defocus search range (um): {Options.ZMin} to {Options.ZMax}");
Console.WriteLine($" Estimate PhaseShift: {Options.DoPhase}");
Console.WriteLine("============================================================");
Console.WriteLine();
int gpuCount = GPU.GetDeviceCount();
Console.WriteLine($"Detected {gpuCount} GPU(s).");
if (Options.StarIn != "")
{
UpdateLocalDefocus(Options.InputFileNames, Options.StarIn, Helper.PathToName(Options.StarIn) + "_updated.star");
return;
}
Image imageGain = null;
OptionsImport optionsImport = new OptionsImport();
if (Options.GainPath != "")
{
if (!File.Exists(Options.GainPath))
{
Console.WriteLine("Input gain reference not found.");
return;
}
optionsImport.GainPath = Options.GainPath;
optionsImport.GainFlipX = Options.GainFlipX;
optionsImport.GainFlipY = Options.GainFlipY;
optionsImport.GainTranspose = Options.GainTranspose;
imageGain = LoadAndPrepareGainReference(optionsImport);
Console.WriteLine("Read the gain reference.");
// Save the transformed gain
string TransformedGainPath = Helper.PathToFolder(Options.GainPath) + "transformed-gain.mrc";
Directory.CreateDirectory(Helper.PathToFolder(TransformedGainPath));
imageGain.WriteMRC(TransformedGainPath, (float)Options.PixelSize, true);
Options.GainPath = TransformedGainPath;
optionsImport.GainFlipX = false;
optionsImport.GainFlipY = false;
optionsImport.GainTranspose = false;
}
ProcessingOptionsMovieMovement movieOptions = new ProcessingOptionsMovieMovement();
movieOptions.GainPath = Options.GainPath;
movieOptions.BinTimes = Options.BinTimes;
movieOptions.PixelSizeX = (decimal)Options.PixelSize;
movieOptions.PixelSizeY = (decimal)Options.PixelSize;
movieOptions.RangeMin = (decimal)Options.MotionRangeMin;
movieOptions.RangeMax = (decimal)Options.MotionRangeMax;;
movieOptions.Bfactor = (decimal)Options.Bfactor;
ProcessingOptionsMovieCTF options = new ProcessingOptionsMovieCTF();
options.Window = Options.CtfWindow;
options.UseMovieSum = Options.UseMovieSum;
options.GridDims = new int3(Options.CtfGridX, Options.CtfGridY, Options.CtfGridT);
options.RangeMin = (decimal)Options.CtfRangeMin;
options.RangeMax = (decimal)Options.CtfRangeMax;
options.Voltage = Options.Voltage;
options.ZMin = (decimal)Options.ZMin;
options.ZMax = (decimal)Options.ZMax;
options.Cs = (decimal)Options.Cs;
options.Amplitude = (decimal)Options.Amplitude;
options.DoPhase = Options.DoPhase;
options.PixelSizeX = (decimal)Options.PixelSize;
options.PixelSizeY = (decimal)Options.PixelSize;
ProcessingOptionsMovieExport exportOptions = new ProcessingOptionsMovieExport();
exportOptions.DoAverage = true;
exportOptions.DoStack = false;
exportOptions.DoDenoise = true;
exportOptions.Voltage = Options.Voltage;
exportOptions.DosePerAngstromFrame = (decimal)Options.DosePerAngstromFrame;
exportOptions.PixelSizeX = (decimal)Options.PixelSize;
exportOptions.PixelSizeY = (decimal)Options.PixelSize;
exportOptions.BinTimes = Options.BinTimes;
Star TableOut = new Star(new[] { "rlnMicrographName" });
TableOut.AddColumn("rlnMagnification");
TableOut.AddColumn("rlnDetectorPixelSize");
TableOut.AddColumn("rlnVoltage");
TableOut.AddColumn("rlnSphericalAberration");
TableOut.AddColumn("rlnAmplitudeContrast");
TableOut.AddColumn("rlnPhaseShift");
TableOut.AddColumn("rlnDefocusU");
TableOut.AddColumn("rlnDefocusV");
TableOut.AddColumn("rlnDefocusAngle");
TableOut.AddColumn("rlnCtfMaxResolution");
TableOut.AddColumn("rlnBeamTiltX");
TableOut.AddColumn("rlnBeamTiltY");
foreach(string inputFileName in Options.InputFileNames)
{
Console.WriteLine($"Processing {inputFileName}");
if (!File.Exists(inputFileName))
{
Console.WriteLine("Input file not found.");
continue;
}
Movie movie = new Movie(Path.GetFullPath(inputFileName));
decimal ScaleFactor = 1M / (decimal)Math.Pow(2, (double)Options.BinTimes);
MapHeader header = null;
Image stack = null;
if (movie.OptionsCTF == null || movie.OptionsMovement == null) // this still reads movies when Z = 1...
{
Console.Write(" Loading the movie...");
LoadAndPrepareHeaderAndMap(inputFileName, imageGain, ScaleFactor, optionsImport, out header, out stack);
Console.WriteLine(" Done");
Console.WriteLine($" Input movie size: X={stack.Dims.X} Y={stack.Dims.Y} N={stack.Dims.Z}");
}
if (movie.OptionsCTF == null)
{
options.Dimensions = new float3(stack.Dims);
movie.ProcessCTF(stack, options);
Console.WriteLine($" CTF estimation done.");
}
if (movie.OptionsMovement == null && stack.Dims.Z > 1)
{
movieOptions.Dimensions = new float3(stack.Dims);
if (Options.MotionGridT <= 0) Options.MotionGridT = stack.Dims.Z;
movieOptions.GridDims = new int3(Options.MotionGridX, Options.MotionGridY, Options.MotionGridT);
movie.ProcessShift(stack, movieOptions);
ExportMotionSTARFile(header, movie, Options);
Console.WriteLine(" Motion estimation done.");
bool NeedsMoreDenoisingExamples = !Directory.Exists(movie.DenoiseTrainingDirOdd) ||
Directory.EnumerateFiles(movie.DenoiseTrainingDirOdd, "*.mrc").Count() < MaxDenoisingTrainData;
exportOptions.DoDenoise = NeedsMoreDenoisingExamples;
exportOptions.Dimensions = new float3(stack.Dims);
}
else
{
exportOptions.DoAverage = false;
}
if (stack != null)
{
movie.ExportMovie(stack, exportOptions);
Console.WriteLine($" Export done.");
}
if (stack != null && !exportOptions.DoAverage)
{
// Make symbolic link to average directory; use for denoising etc
System.Diagnostics.Process Process = new System.Diagnostics.Process();
Process.StartInfo.FileName = "/bin/ln";
Process.StartInfo.Arguments = $"-sf ../{movie.RootName}.mrc {movie.AveragePath}";
Process.Start();
Process.WaitForExit();
}
List<string> Row = new List<string>();
Row.AddRange(new[]
{
RelativePathFromPWD(movie.AveragePath), // rlnMicrographName
(1e4M / movie.OptionsCTF.BinnedPixelSizeMean).ToString("F1", CultureInfo.InvariantCulture), // rlnMagnification
"1.0", // rlnDetectorPixelSize
movie.CTF.Voltage.ToString("F1", CultureInfo.InvariantCulture), // rlnVoltage
movie.CTF.Cs.ToString("F4", CultureInfo.InvariantCulture), // rlnSphericalAberration
movie.CTF.Amplitude.ToString("F3", CultureInfo.InvariantCulture), // rlnAmplitudeContrast
(movie.CTF.PhaseShift * 180M).ToString("F1", CultureInfo.InvariantCulture), // rlnPhaseShift
((movie.CTF.Defocus + movie.CTF.DefocusDelta / 2) * 1e4M).ToString("F1", CultureInfo.InvariantCulture), // rlnDefocusU
((movie.CTF.Defocus - movie.CTF.DefocusDelta / 2) * 1e4M).ToString("F1", CultureInfo.InvariantCulture), // rlnDefocusV
movie.CTF.DefocusAngle.ToString("F1", CultureInfo.InvariantCulture), // rlnDefocusAngle
movie.CTFResolutionEstimate.ToString("F1", CultureInfo.InvariantCulture), // rlnCtfMaxResolution
movie.CTF.BeamTilt.X.ToString("F2", CultureInfo.InvariantCulture), // rlnBeamTiltX
movie.CTF.BeamTilt.Y.ToString("F2", CultureInfo.InvariantCulture) // rlnBeamTiltY
});
if (movie.OptionsMovement != null)
{
if (!TableOut.HasColumn("rlnMicrographMetadata"))
TableOut.AddColumn("rlnMicrographMetadata");
Row.Add(RelativePathFromPWD(movie.DirectoryName + "motion/" + movie.RootName + ".star")); // rlnMicrographMetadata
}
TableOut.AddRow(Row);
}
TableOut.Save("exported.star");
return;
}
public static string RelativePathFromPWD(string Path)
{
Uri UriPwd = new Uri(Environment.CurrentDirectory + "/");
return UriPwd.MakeRelativeUri(new Uri(Path)).ToString();
}
// Adapted from ButtonWrite_OnClick() in Warp/Controls/TaskDialogs/2D/Dialog2DList.xaml.cs
public static void ExportMotionSTARFile(MapHeader Header, Movie Movie, Options Options)
{
List<string> HeaderNames = new List<string>()
{
"rlnImageSizeX",
"rlnImageSizeY",
"rlnImageSizeZ",
"rlnMicrographMovieName",
"rlnMicrographBinning",
"rlnMicrographOriginalPixelSize",
"rlnMicrographDoseRate",
"rlnMicrographPreExposure",
"rlnVoltage",
"rlnMicrographStartFrame",
"rlnMotionModelVersion"
};
List<string> HeaderValues = new List<string>()
{
Header.Dimensions.X.ToString(),
Header.Dimensions.Y.ToString(),
Header.Dimensions.Z.ToString(),
RelativePathFromPWD(Movie.Path),
Math.Pow(2.0, (double)(Options.BinTimes)).ToString("F5"),
Options.PixelSize.ToString("F5"),
Options.DosePerAngstromFrame.ToString("F5"),
"0",
Options.Voltage.ToString(),
"1",
"0"
};
if (!string.IsNullOrEmpty(Options.GainPath))
{
HeaderNames.Add("rlnMicrographGainName");
HeaderValues.Add(Options.GainPath);
}
StarParameters ParamsTable = new StarParameters(HeaderNames.ToArray(), HeaderValues.ToArray());
float2[] MotionTrack = Movie.GetMotionTrack(new float2(0.5f, 0.5f), 1).Select(v => v / (float)Options.PixelSize).ToArray();
Star TrackTable = new Star(new[]
{
Helper.ArrayOfSequence(1, MotionTrack.Length + 1, 1).Select(v => v.ToString()).ToArray(),
MotionTrack.Select(v => (-v.X).ToString("F5")).ToArray(),
MotionTrack.Select(v => (-v.Y).ToString("F5")).ToArray()
},
new[]
{
"rlnMicrographFrameNumber",
"rlnMicrographShiftX",
"rlnMicrographShiftY"
});
Directory.CreateDirectory(Movie.DirectoryName + "motion/");
Star.SaveMultitable(Movie.DirectoryName + "motion/" + Movie.RootName + ".star", new Dictionary<string, Star>()
{
{ "general", ParamsTable },
{ "global_shift", TrackTable },
});
}
// Adapted from Warp/MainWindow.xaml.cs
public static Image LoadAndPrepareGainReference(OptionsImport optionsImport)
{
Image Gain = Image.FromFilePatient(50, 500, optionsImport.GainPath,
new int2(optionsImport.HeaderlessWidth, optionsImport.HeaderlessHeight),
(int)optionsImport.HeaderlessOffset,
ImageFormatsHelper.StringToType(optionsImport.HeaderlessType));
if (optionsImport.GainFlipX)
Gain = Gain.AsFlippedX();
if (optionsImport.GainFlipY)
Gain = Gain.AsFlippedY();
if (optionsImport.GainTranspose)
Gain = Gain.AsTransposed();
return Gain;
}
// Adapted from Warp/MainWindow.xaml.cs
public static void LoadAndPrepareHeaderAndMap(string path, Image imageGain, decimal scaleFactor, OptionsImport optionsImport, out MapHeader header, out Image stack, bool needStack = true, int maxThreads = 8)
{
header = MapHeader.ReadFromFilePatient(50, 500, path,
new int2(optionsImport.HeaderlessWidth, optionsImport.HeaderlessHeight),
optionsImport.HeaderlessOffset,
ImageFormatsHelper.StringToType(optionsImport.HeaderlessType));
string Extension = Helper.PathToExtension(path).ToLower();
bool IsTiff = header.GetType() == typeof(HeaderTiff);
bool IsEER = header.GetType() == typeof(HeaderEER);
if (imageGain != null)
if (!IsEER)
if (header.Dimensions.X != imageGain.Dims.X || header.Dimensions.Y != imageGain.Dims.Y)
throw new Exception("Gain reference dimensions do not match image.");
int EERSupersample = 1;
if (imageGain != null && IsEER)
{
if (header.Dimensions.X == imageGain.Dims.X)
EERSupersample = 1;
else if (header.Dimensions.X * 2 == imageGain.Dims.X)
EERSupersample = 2;
else if (header.Dimensions.X * 4 == imageGain.Dims.X)
EERSupersample = 3;
else
throw new Exception("Invalid supersampling factor requested for EER based on gain reference dimensions");
}
HeaderEER.SuperResolution = EERSupersample;
if (IsEER && imageGain != null)
{
header.Dimensions.X = imageGain.Dims.X;
header.Dimensions.Y = imageGain.Dims.Y;
}
MapHeader Header = header;
int NThreads = (IsTiff || IsEER) ? 6 : 2;
int CurrentDevice = GPU.GetDevice();
if (needStack)
{
byte[] TiffBytes = null;
if (IsTiff)
{
MemoryStream Stream = new MemoryStream();
using (Stream BigBufferStream = IOHelper.OpenWithBigBuffer(path))
BigBufferStream.CopyTo(Stream);
TiffBytes = Stream.GetBuffer();
}
if (scaleFactor == 1M)
{
stack = new Image(header.Dimensions);
float[][] OriginalStackData = stack.GetHost(Intent.Write);
Helper.ForCPU(0, header.Dimensions.Z, NThreads, threadID => GPU.SetDevice(CurrentDevice), (z, threadID) =>
{
Image Layer = null;
MemoryStream TiffStream = TiffBytes != null ? new MemoryStream(TiffBytes) : null;
if (!IsEER)
Layer = Image.FromFilePatient(50, 500, path,
new int2(optionsImport.HeaderlessWidth, optionsImport.HeaderlessHeight),
(int)optionsImport.HeaderlessOffset,
ImageFormatsHelper.StringToType(optionsImport.HeaderlessType),
z,
TiffStream);
else
{
Layer = new Image(Header.Dimensions.Slice());
EERNative.ReadEERPatient(50, 500,
path, z * 10, (z + 1) * 10, EERSupersample, Layer.GetHost(Intent.Write)[0]);
}
lock (OriginalStackData)
{
if (imageGain != null)
Layer.MultiplySlices(imageGain);
Layer.Xray(20f);
OriginalStackData[z] = Layer.GetHost(Intent.Read)[0];
Layer.Dispose();
}
}, null);
}
else // binning
{
int3 ScaledDims = new int3((int)Math.Round(header.Dimensions.X * scaleFactor) / 2 * 2,
(int)Math.Round(header.Dimensions.Y * scaleFactor) / 2 * 2,
header.Dimensions.Z);
stack = new Image(ScaledDims);
float[][] OriginalStackData = stack.GetHost(Intent.Write);
int PlanForw = GPU.CreateFFTPlan(header.Dimensions.Slice(), 1);
int PlanBack = GPU.CreateIFFTPlan(ScaledDims.Slice(), 1);
Helper.ForCPU(0, ScaledDims.Z, NThreads, threadID => GPU.SetDevice(CurrentDevice), (z, threadID) =>
{
Image Layer = null;
MemoryStream TiffStream = TiffBytes != null ? new MemoryStream(TiffBytes) : null;
if (!IsEER)
Layer = Image.FromFilePatient(50, 500, path,
new int2(optionsImport.HeaderlessWidth, optionsImport.HeaderlessHeight),
(int)optionsImport.HeaderlessOffset,
ImageFormatsHelper.StringToType(optionsImport.HeaderlessType),
z,
TiffStream);
else
{
Layer = new Image(Header.Dimensions.Slice());
EERNative.ReadEERPatient(50, 500,
path, z * 10, (z + 1) * 10, EERSupersample, Layer.GetHost(Intent.Write)[0]);
}
Image ScaledLayer = null;
lock (OriginalStackData)
{
if (imageGain != null)
Layer.MultiplySlices(imageGain);
Layer.Xray(20f);
ScaledLayer = Layer.AsScaled(new int2(ScaledDims), PlanForw, PlanBack);
Layer.Dispose();
}
OriginalStackData[z] = ScaledLayer.GetHost(Intent.Read)[0];
ScaledLayer.Dispose();
}, null);
GPU.DestroyFFTPlan(PlanForw);
GPU.DestroyFFTPlan(PlanBack);
}
}
else // !needStack
{
stack = null;
}
}
// Adapted from ButtonWrite_OnClick() in Warp/Controls/TaskDialogs/2D/Dialog2DList.xaml.cs
public static void UpdateLocalDefocus(IEnumerable<string> InputFileNames, string ImportPath, string ExportPath)
{
#region Get all movies that can potentially be used
List<Movie> ValidMovies = new List<Movie>();
if (InputFileNames.Count() == 1 && InputFileNames.First().EndsWith("lst"))
{
Console.WriteLine($"Input is a list file.");
StreamReader sr = new StreamReader(InputFileNames.First());
while (!sr.EndOfStream)
{
string line = sr.ReadLine().Trim();
Movie Movie = new Movie(Path.GetFullPath(line));
if (Movie.OptionsCTF != null)
ValidMovies.Add(Movie);
else
Console.WriteLine($"CTF not available for {line}");
}
}
else
{
foreach(string inputFileName in InputFileNames)
{
if (!File.Exists(inputFileName))
{
Console.WriteLine($"Input file {inputFileName} not found.");
continue;
}
Movie Movie = new Movie(Path.GetFullPath(inputFileName));
if (Movie.OptionsCTF != null)
ValidMovies.Add(Movie);
else
Console.WriteLine($"CTF not available for {inputFileName}");
}
}
List<string> ValidMovieNames = ValidMovies.Select(m => m.RootName).ToList();
Console.WriteLine($"Accepted {ValidMovies.Count()} movies.");
#endregion
#region Read table and intersect its micrograph set with valid movies
Star TableIn, OpticsTable = null;
bool HasOpticsGroup = Star.ContainsTable(ImportPath, "optics");
if (HasOpticsGroup)
{
TableIn = new Star(ImportPath, "particles");
OpticsTable = new Star(ImportPath, "optics");
}
else
{
TableIn = new Star(ImportPath);
}
if (!TableIn.HasColumn("rlnMicrographName"))
throw new Exception("Couldn't find rlnMicrographName column.");
if (!TableIn.HasColumn("rlnCoordinateX"))
throw new Exception("Couldn't find rlnCoordinateX column.");
if (!TableIn.HasColumn("rlnCoordinateY"))
throw new Exception("Couldn't find rlnCoordinateY column.");
Dictionary<string, List<int>> Groups = new Dictionary<string, List<int>>();
{
string[] ColumnMicNames = TableIn.GetColumn("rlnMicrographName");
for (int r = 0; r < ColumnMicNames.Length; r++)
{
if (!Groups.ContainsKey(ColumnMicNames[r]))
Groups.Add(ColumnMicNames[r], new List<int>());
Groups[ColumnMicNames[r]].Add(r);
}
// Match only file name
Groups = Groups.ToDictionary(group => Helper.PathToName(group.Key), group => group.Value);
// Where returns Enumerable; thus have to make Dictionary again
// @ allows the use of reserved word ('group')
Groups = Groups.Where(@group => ValidMovieNames.Contains(@group.Key)).ToDictionary(@group => @group.Key, @group => @group.Value);
}
bool[] RowsIncluded = new bool[TableIn.RowCount];
foreach (var @group in Groups)
foreach (var r in group.Value)
RowsIncluded[r] = true;
List<int> RowsNotIncluded = new List<int>();
for (int r = 0; r < RowsIncluded.Length; r++)
if (!RowsIncluded[r])
RowsNotIncluded.Add(r);
Console.WriteLine($"Movies for {RowsNotIncluded.Count} particles not found.");
#endregion
#region Make sure all columns are there
if (!HasOpticsGroup)
{
if (!TableIn.HasColumn("rlnVoltage"))
TableIn.AddColumn("rlnVoltage", "300.0");
if (!TableIn.HasColumn("rlnSphericalAberration"))
TableIn.AddColumn("rlnSphericalAberration", "2.7");
if (!TableIn.HasColumn("rlnAmplitudeContrast"))
TableIn.AddColumn("rlnAmplitudeContrast", "0.07");
}
if (!TableIn.HasColumn("rlnPhaseShift"))
TableIn.AddColumn("rlnPhaseShift", "0.0");
if (!TableIn.HasColumn("rlnDefocusU"))
TableIn.AddColumn("rlnDefocusU", "0.0");
if (!TableIn.HasColumn("rlnDefocusV"))
TableIn.AddColumn("rlnDefocusV", "0.0");
if (!TableIn.HasColumn("rlnDefocusAngle"))
TableIn.AddColumn("rlnDefocusAngle", "0.0");
if (!TableIn.HasColumn("rlnCtfMaxResolution"))
TableIn.AddColumn("rlnCtfMaxResolution", "999.0");
#endregion
float[] PosX = TableIn.GetColumn("rlnCoordinateX").Select(v => float.Parse(v, CultureInfo.InvariantCulture)).ToArray();
float[] PosY = TableIn.GetColumn("rlnCoordinateY").Select(v => float.Parse(v, CultureInfo.InvariantCulture)).ToArray();
foreach (var movie in ValidMovies)
{
if (!Groups.ContainsKey(movie.RootName))
continue;
List<int> GroupRows = Groups[movie.RootName];
float Astigmatism = (float)movie.CTF.DefocusDelta / 2;
float PhaseShift = movie.OptionsCTF.DoPhase ? movie.GridCTFPhase.GetInterpolated(new float3(0.5f)) * 180 : 0;
Console.WriteLine($"Working on {GroupRows.Count()} particles from {movie.RootName}");
foreach (var r in GroupRows)
{
// TODO: What happens if binned?
float3 Position = new float3(PosX[r] * (float)movie.OptionsCTF.PixelSizeMean / movie.OptionsCTF.Dimensions.X,
PosY[r] * (float)movie.OptionsCTF.PixelSizeMean / movie.OptionsCTF.Dimensions.Y,
0.5f);
//Console.WriteLine($"Pos = {PosX[r]}, {PosY[r]}, AngPix = {movie.OptionsCTF.PixelSizeMean}, Dim = {movie.OptionsCTF.Dimensions}");
float LocalDefocus = movie.GridCTFDefocus.GetInterpolated(Position);
TableIn.SetRowValue(r, "rlnDefocusU", ((LocalDefocus + Astigmatism) * 1e4f).ToString("F1", CultureInfo.InvariantCulture));
TableIn.SetRowValue(r, "rlnDefocusV", ((LocalDefocus - Astigmatism) * 1e4f).ToString("F1", CultureInfo.InvariantCulture));
TableIn.SetRowValue(r, "rlnDefocusAngle", movie.CTF.DefocusAngle.ToString("F1", CultureInfo.InvariantCulture));
if (!HasOpticsGroup)
{
TableIn.SetRowValue(r, "rlnVoltage", movie.CTF.Voltage.ToString("F1", CultureInfo.InvariantCulture));
TableIn.SetRowValue(r, "rlnSphericalAberration", movie.CTF.Cs.ToString("F4", CultureInfo.InvariantCulture));
TableIn.SetRowValue(r, "rlnAmplitudeContrast", movie.CTF.Amplitude.ToString("F3", CultureInfo.InvariantCulture));
}
else
{
// TODO: Update the optics group table entry
}
TableIn.SetRowValue(r, "rlnPhaseShift", PhaseShift.ToString("F1", CultureInfo.InvariantCulture));
TableIn.SetRowValue(r, "rlnCtfMaxResolution", movie.CTFResolutionEstimate.ToString("F1", CultureInfo.InvariantCulture));
} // particles
} // movies
if (HasOpticsGroup)
{
Star.SaveMultitable(ExportPath, new Dictionary<string, Star>()
{
{ "optics", OpticsTable },
{ "particles", TableIn },
});
}
else
{
TableIn.Save(ExportPath);
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment