diff --git a/rfplot.py b/rfplot.py index aa57fc4..d44e795 100644 --- a/rfplot.py +++ b/rfplot.py @@ -8,6 +8,10 @@ import matplotlib.dates as mdates from matplotlib.backend_bases import MouseButton from matplotlib.widgets import RectangleSelector import matplotlib as mpl +from astropy.visualization import (ZScaleInterval, ImageNormalize,SqrtStretch) +import imageio +from skimage.morphology import binary_dilation, remove_small_objects + mpl.rcParams['keymap.save'].remove('s') mpl.rcParams['keymap.fullscreen'].remove('f') @@ -42,11 +46,43 @@ if __name__ == "__main__": def line_select_callback(eclick, erelease): x1, y1 = eclick.xdata, eclick.ydata x2, y2 = erelease.xdata, erelease.ydata - array = mark.get_offsets() - maskx = np.logical_and(array[:,0] >= min(x1,x2), array[:,0] <= max(x1,x2)) - masky = np.logical_and(array[:,1] >= min(y1,y2), array[:,1] <= max(y1,y2)) - mask = np.logical_and(maskx, masky) - mark.set_offsets(array[np.logical_not(mask),:]) + + t1_ind = round(len(s.t) * (x1 - tmin) / (tmax - tmin)) + t2_ind = round(len(s.t) * (x2 - tmin) / (tmax - tmin)) + f1_ind = round(len(s.freq) * (y1 - fmin) / (fmax - fmin)) + f2_ind = round(len(s.freq) * (y2 - fmin) / (fmax - fmin)) + + + submat = s.z[f1_ind:f2_ind,t1_ind:t2_ind] + + signal = submat - np.median(submat, axis=0) + background = np.copy(signal) + filter = np.ones(50)/50 + for i in range(signal.shape[1]): + background[:,i] = np.convolve(signal[:,i], filter, mode="same") + + sig_without_background = signal - background + mask = sig_without_background > 3 * np.std(sig_without_background, axis=0) + sig_without_background[mask] = background[mask] + sig_without_background[np.logical_not(mask)] = signal[np.logical_not(mask)] + + for i in range(signal.shape[1]): + background[:,i] = np.convolve(sig_without_background[:,i], filter, mode="same") + + sig_without_background = signal - background + mask = (sig_without_background > 3 * np.std(sig_without_background, axis=0)).astype(np.uint8) + mask = binary_dilation(mask) + remove_small_objects(mask, min_size=16, in_place=True) + mask = np.flipud(mask) + + imageio.imwrite(f'test3.png', 255 * mask) + + print(s.z[f1_ind:f2_ind,t1_ind:t2_ind].shape) + # array = mark.get_offsets() + # maskx = np.logical_and(array[:,0] >= min(x1,x2), array[:,0] <= max(x1,x2)) + # masky = np.logical_and(array[:,1] >= min(y1,y2), array[:,1] <= max(y1,y2)) + # mask = np.logical_and(maskx, masky) + # mark.set_offsets(array[np.logical_not(mask),:]) fig.canvas.draw() print(f"select over {x1},{y1},{x2},{y2}")