| # HuggingFace Hub | |
| from huggingface_hub import from_pretrained_keras | |
| model = from_pretrained_keras("Plsek/CADET-v1") | |
| # Basic libraries | |
| import os | |
| import shutil | |
| import numpy as np | |
| from scipy.ndimage import center_of_mass | |
| import matplotlib.pyplot as plt | |
| from matplotlib.colors import LogNorm | |
| from matplotlib.patches import Rectangle | |
| # Astropy | |
| from astropy.io import fits | |
| from astropy.wcs import WCS | |
| from astropy.nddata import Cutout2D, CCDData | |
| from astropy.convolution import Gaussian2DKernel as Gauss | |
| from astropy.convolution import convolve | |
| # Scikit-learn | |
| from sklearn.cluster import DBSCAN | |
| # Streamlit | |
| import streamlit as st | |
| st.set_option('deprecation.showPyplotGlobalUse', False) | |
| # # Define function to plot the uploaded image | |
| # def plot_image(image, scale): | |
| # plt.figure(figsize=(4, 4)) | |
| # x0 = image.shape[0] // 2 - scale * 128 / 2 | |
| # plt.imshow(image, origin="lower") | |
| # plt.gca().add_patch(Rectangle((x0-0.5, x0-0.5), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none')) | |
| # plt.axis('off') | |
| # plt.tight_layout() | |
| # with colA: st.pyplot() | |
| # # Define function to plot the prediction | |
| # def plot_prediction(pred): | |
| # plt.figure(figsize=(4, 4)) | |
| # plt.imshow(pred, origin="lower") | |
| # plt.axis('off') | |
| # with colB: st.pyplot() | |
| # # Define function to plot the decomposed prediction | |
| # def plot_decomposed(decomposed): | |
| # plt.figure(figsize=(4, 4)) | |
| # plt.imshow(decomposed, origin="lower") #, norm=LogNorm()) | |
| # N = int(np.max(decomposed)) | |
| # for i in range(N): | |
| # new = np.where(decomposed == i+1, 1, 0) | |
| # x0, y0 = center_of_mass(new) | |
| # color = "white" if i < N//2 else "black" | |
| # plt.text(y0, x0, f"{i+1}", ha="center", va="center", fontsize=15, color=color) | |
| # plt.axis('off') | |
| # with colC: st.pyplot() | |
| # # Define function to cut input image and rebin it to 128x128 pixels | |
| # def cut(data0, wcs0, scale=1): | |
| # shape = data0.shape[0] | |
| # x0 = shape / 2 | |
| # size = 128 * scale | |
| # cutout = Cutout2D(data0, (x0, x0), (size, size), wcs=wcs0) | |
| # data, wcs = cutout.data, cutout.wcs | |
| # # Regrid data | |
| # factor = size // 128 | |
| # data = data.reshape(128, factor, 128, factor).mean(-1).mean(1) | |
| # # Regrid wcs | |
| # ra, dec = wcs.wcs_pix2world(np.array([[63, 63]]),0)[0] | |
| # wcs.wcs.cdelt[0] = wcs.wcs.cdelt[0] * factor | |
| # wcs.wcs.cdelt[1] = wcs.wcs.cdelt[1] * factor | |
| # wcs.wcs.crval[0] = ra | |
| # wcs.wcs.crval[1] = dec | |
| # wcs.wcs.crpix[0] = 64 / factor | |
| # wcs.wcs.crpix[1] = 64 / factor | |
| # return data, wcs | |
| # # Define function to apply cutting and produce a prediction | |
| # # @st.cache | |
| # def cut_n_predict(data, wcs, scale): | |
| # data, wcs = cut(data, wcs, scale=scale) | |
| # image = np.log10(data+1) | |
| # y_pred = 0 | |
| # for j in [0,1,2,3]: | |
| # rotated = np.rot90(image, j) | |
| # pred = model.predict(rotated.reshape(1, 128, 128, 1)).reshape(128 ,128) | |
| # pred = np.rot90(pred, -j) | |
| # y_pred += pred / 4 | |
| # return y_pred, wcs | |
| # # Define function to decompose prediction into individual cavities | |
| # # @st.cache | |
| # def decompose_cavity(pred, th2=0.7, amin=6): | |
| # X, Y = pred.nonzero() | |
| # data = np.array([X,Y]).reshape(2, -1) | |
| # # DBSCAN clustering | |
| # try: clusters = DBSCAN(eps=1.0, min_samples=3).fit(data.T).labels_ | |
| # except: clusters = [] | |
| # N = len(set(clusters)) | |
| # cavities = [] | |
| # for i in range(N): | |
| # img = np.zeros((128,128)) | |
| # b = clusters == i | |
| # xi, yi = X[b], Y[b] | |
| # img[xi, yi] = pred[xi, yi] | |
| # # # Thresholding #2 | |
| # # if not (img > th2).any(): continue | |
| # # Minimal area | |
| # if np.sum(img) <= amin: continue | |
| # cavities.append(img) | |
| # # Save raw and decomposed predictions to predictions folder | |
| # ccd = CCDData(pred, unit="adu", wcs=wcs) | |
| # ccd.write(f"predictions/predicted.fits", overwrite=True) | |
| # image_decomposed = np.zeros((128,128)) | |
| # for i, cav in enumerate(cavities): | |
| # ccd = CCDData(cav, unit="adu", wcs=wcs) | |
| # ccd.write(f"predictions/predicted_{i+1}.fits", overwrite=True) | |
| # image_decomposed += (i+1) * np.where(cav > 0, 1, 0) | |
| # # shutil.make_archive("predictions", 'zip', "predictions") | |
| # return image_decomposed | |
| # # @st.cache | |
| # def load_file(fname): | |
| # with fits.open(fname) as hdul: | |
| # data = hdul[0].data | |
| # wcs = WCS(hdul[0].header) | |
| # return data, wcs | |
| # def change_scale(): | |
| # del st.session_state["threshold"] | |
| # # Use wide layout and create columns | |
| # st.set_page_config(page_title="Cavity Detection Tool", layout="wide") | |
| # bordersize = 0.45 | |
| # _, col, _ = st.columns([bordersize, 3, bordersize]) | |
| # os.system("mkdir -p predictions") | |
| # with col: | |
| # # Create heading and description | |
| # st.markdown("<h1 align='center'>Cavity Detection Tool</h1>", unsafe_allow_html=True) | |
| # st.markdown("Cavity Detection Tool (CADET) is a machine learning pipeline trained to detect X-ray cavities from noisy Chandra images of early-type galaxies.") | |
| # st.markdown("To use this tool: upload your image, select the scale of interest, make a prediction, and decompose it into individual cavities!") | |
| # st.markdown("Input images should be in units of counts, centred at the galaxy center, and point sources should be filled with surrounding background ([dmfilth](https://cxc.cfa.harvard.edu/ciao/ahelp/dmfilth.html)).") | |
| # st.markdown("If you use this tool for your research, please cite [Plšek et al. 2023](https://arxiv.org/abs/2304.05457)") | |
| # # _, col_1, col_2, col_3, _ = st.columns([bordersize, 2.0, 0.5, 0.5, bordersize]) | |
| # # with col: | |
| # uploaded_file = st.file_uploader("Choose a FITS file", type=['fits']) | |
| # # with col_2: | |
| # # st.markdown("### Examples") | |
| # # NGC4649 = st.button("NGC4649") | |
| # # with col_3: | |
| # # st.markdown("""<style>[data-baseweb="select"] {margin-top: 26px;}</style>""", unsafe_allow_html=True) | |
| # # NGC5813 = st.button("NGC5813") | |
| # # if NGC4649: | |
| # # uploaded_file = "NGC4649_example.fits" | |
| # # elif NGC5813: | |
| # # uploaded_file = "NGC5813_example.fits" | |
| # # # If file is uploaded, read in the data and plot it | |
| # # if uploaded_file is not None: | |
| # # data, wcs = load_file(uploaded_file) | |
| # # # if "data" not in locals(): | |
| # # # data = np.zeros((128,128)) | |
| # # # Make six columns for buttons | |
| # # _, col1, col2, col3, col4, col5, col6, _ = st.columns([bordersize,0.5,0.5,0.5,0.5,0.5,0.5,bordersize]) | |
| # # col1.subheader("Input image") | |
| # # col3.subheader("Prediction") | |
| # # col5.subheader("Decomposed") | |
| # # col6.subheader("") | |
| # # with col1: | |
| # # st.markdown("""<style>[data-baseweb="select"] {margin-top: -46px;}</style>""", unsafe_allow_html=True) | |
| # # max_scale = int(data.shape[0] // 128) | |
| # # scale = st.selectbox('Scale:',[f"{(i+1)*128}x{(i+1)*128}" for i in range(max_scale)], label_visibility="hidden", on_change=change_scale) | |
| # # scale = int(scale.split("x")[0]) // 128 | |
| # # # Detect button | |
| # # with col3: detect = st.button('Detect', key="detect") | |
| # # # Threshold slider | |
| # # with col4: | |
| # # st.markdown("") | |
| # # # st.markdown("""<style>[data-baseweb="select"] {margin-top: -36px;}</style>""", unsafe_allow_html=True) | |
| # # threshold = st.slider("Threshold", 0.0, 1.0, 0.0, 0.05, key="threshold") #, label_visibility="hidden") | |
| # # # Decompose button | |
| # # with col5: decompose = st.button('Decompose', key="decompose") | |
| # # # Make two columns for plots | |
| # # _, colA, colB, colC, _ = st.columns([bordersize,1,1,1,bordersize]) | |
| # # image = np.log10(data+1) | |
| # # plot_image(image, scale) | |
| # # if detect or threshold: | |
| # # # if st.session_state.get("detect", True): | |
| # # y_pred, wcs = cut_n_predict(data, wcs, scale) | |
| # # y_pred_th = np.where(y_pred > threshold, y_pred, 0) | |
| # # plot_prediction(y_pred_th) | |
| # # if decompose or st.session_state.get("download", False): | |
| # # image_decomposed = decompose_cavity(y_pred_th) | |
| # # plot_decomposed(image_decomposed) | |
| # # with col6: | |
| # # st.markdown("<br style='margin:4px 0'>", unsafe_allow_html=True) | |
| # # # st.markdown("""<style>[data-baseweb="select"] {margin-top: 16px;}</style>""", unsafe_allow_html=True) | |
| # # fname = uploaded_file.name.strip(".fits") | |
| # # # if st.session_state.get("download", False): | |
| # # shutil.make_archive("predictions", 'zip', "predictions") | |
| # # with open('predictions.zip', 'rb') as f: | |
| # # res = f.read() | |
| # # download = st.download_button(label="Download", data=res, key="download", | |
| # # file_name=f'{fname}_{int(scale*128)}.zip', | |
| # # # disabled=st.session_state.get("disabled", True), | |
| # # mime="application/octet-stream") |