Commit 013c915d authored by Adrian Nievergelt's avatar Adrian Nievergelt

First draft of k2 line correction script for dose fractionated tomography stacks

parent 9d5b3a24
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri May 10 11:45:03 2019
Correct line streaks in broken K2 image files
@author: adrian
"""
import mrcfile
import numpy as np
import scipy.stats, scipy.signal
import argparse
import os,sys
parser = argparse.ArgumentParser(description = "Removes streak outliers generated by aging K2 camera from a stack of frames by variational cutoff analysis.")
# parser.add_argument("-f", "--file", help="create overlay for a single file",
parser.add_argument('-i', '--input', nargs='?', default='.', help='Input file to process')
parser.add_argument('-o', '--output', nargs='?', default='.', help='Output file to write to. If no argument is provided, a _destreak suffix will be added to the original file name, unless directory output mode is enabled')
parser.add_argument('-d', '--outDir', nargs='?', default='.', help='Optional output directory to write to. Saves the file under the same name into specified output directory.')
parser.add_argument('-c', '--cutoff', nargs='?', type=float, default='1', help='The threshold of how many standard deviations above or below the a line must deviate before it is corrected')
parser.add_argument('-g', '--graph', action='store_true', help='Store a graph of the corrected and uncorrected line intensities for debug purposes.')
args = parser.parse_args()
if args.graph:
import matplotlib.pyplot as plt
# input checking and sanitizing
args.input = args.input.strip()
args.output = args.output.strip()
args.outDir = args.outDir.strip()
if args.input == '.':
print("Missing input file");
print(parser.print_help())
sys.exit()
if args.output == '.' and args.outDir == '.':
inPath = os.path.splitext(args.input)
args.output = inPath[0] + "_destripe" + inPath[1]
if not args.outDir == '.':
if not os.path.isdir(args.outDir):
try:
print(args.outDir)
os.mkdir(args.outDir)
except:
print("Could not create output directory!")
sys.exit()
if args.output == '.':
args.output = os.path.basename(args.input)
args.output = os.path.join(args.outDir, os.path.basename(args.output))
inFile = args.input.strip()
mrc = mrcfile.mmap(inFile)
testIm = np.sum(mrc.data,axis=0)
# dispIm = (mrc.data)[1,1000:2200, 1500:3500].astype('uint8')
lineSums = np.sum(testIm, axis=1)
lineSumsAll = np.sum(mrc.data, axis=(0,2))
lineNo = len(lineSumsAll)
# Find sharp spikes in the image sum vector
envLines = 30
environKernel = envLines*[-1/(envLines*2)]
detectionKernel = np.array(environKernel+[1]+environKernel)
spikeDetect = abs(np.convolve(lineSumsAll, detectionKernel, mode='same'))
spikeSelect = spikeDetect > args.cutoff*np.std(spikeDetect)
spikeSelect[0:envLines] = 0
spikeSelect[lineNo-envLines:lineNo] = 0
# print(np.std(spikeDetect))
# Construct a correction vector
# First generate an environment vector with a clipped mean
envirDensity = scipy.signal.savgol_filter(lineSumsAll, 61, 3)
correctionVector = np.array([x*y for x,y in zip((envirDensity - lineSumsAll), spikeSelect)])
lineSumsAllCorr = lineSumsAll + correctionVector
correctionVector = np.round((correctionVector/mrc.data.shape[0]/mrc.data.shape[2])).astype(np.int16)
correctionVector = np.reshape(correctionVector, [1, len(correctionVector),1]).astype(np.int16)
correctionMatrix = correctionVector*np.ones([mrc.data.shape[0],1,mrc.data.shape[2]])
# Write the output file
with mrcfile.new(args.output, overwrite=True) as outmrc:
# outmrc.header = mrc.header
outmrc.set_data((mrc.data+correctionMatrix).astype(np.int16))
outmrc.set_extended_header(mrc.extended_header)
if args.graph:
plt.plot(lineSumsAll)
plt.plot(lineSumsAllCorr)
outPath = os.path.splitext(args.output)
plotPath = outPath[0] + "_plot" + outPath[1] + ".png"
plt.savefig(plotPath)
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment