Search code examples
pythonfits

FITS image copy and pasting, Python


So I have a FITS image file. I want to crop one particular area (it is a straight line in the fits image which I want to copy and rotate), rotate it at an angle and then create a new image with 12 such straight lines in a circle. Like one at 30 degrees other at 60 and so on. I have tried rotating it but I don't know how do I paste it and create a new image. This is on python.

I am using this code to open the fits image:

hdul = pyfits.open('img.fits')

img=hdul[0].data

This is the image


Solution

  • Update 1: I have edited the code based on your updated specifications.

    There are three sub problems to your question:
    Step 1) extracting the region you have specified.
    Step 2) rotating the extracted region.
    Step 3) copying it to an image(basically another numpy array).

    This is still a dirty code. Can be improved based on more inputs about what to expect in input and extraction logic for feature. As of now feature extraction is still hardcoded.

    import numpy as np
    import copy
    import math
    from scipy import ndimage
    from scipy.spatial import distance
    import matplotlib.pyplot as plt
    from astropy.io import fits
    
    
    # This is for Step 1.
    def extract_roi(fits_image: np.ndarray) -> np.ndarray:
        """Returns the extracted region of interest(ROI) as numpy array.
    
        ROI: Rectangular region around white vertical line. Currently hard-coding
        the bounds. They have to be set manually.
        Assumption: ROI will be a vertical feature.
        """
        ROW_START, ROW_END = 0, 50
        COL_START, COL_END = 70, 75
    
        extracted_region = fits_image[ROW_START: ROW_END, COL_START:COL_END]
        return extracted_region
    
    
    def square_pad_roi(roi_image, pivot):
        """ Pads the image with zeros and centers is about the pivot.
    
        This eases rotation about a pivot using scipy.
        Assumption: Image Top Left: Array (0, 0)
        Assumption: Image Bot Right: Array(Rows, Cols)
        :param roi_image: The image with the extracted ROI.
        :param pivot: [row, col] pivot position to rotate image about.
                      row, col should be w.r.t. roi_image co-ords.
        :return: square padded image centered about the provided pivot.
        """
        r, c = roi_image.shape[:2]
        piv_row, piv_col = pivot
        corners = [[0, 0], [0, c], [r, c], [r, 0]]  # [TL, TR, BR, BL]
        dist = math.ceil(max([distance.euclidean(x, pivot) for x in corners]))
    
        pad_row = [abs(piv_row - dist), abs(piv_row + dist - r)]
        pad_col = [abs(piv_col - dist), abs(piv_col + dist - c)]
    
        return np.pad(roi_image, [pad_row, pad_col], mode='constant',
                      constant_values=0)
    
    
    # This is for Step 2.
    def rotate_and_scale(image, angle, out_width) -> np.ndarray:
        """Rotates input image about center and scales it.
    
        :param image: ROI image input which has to be rotated.
                      Expected to be square in shaped and centered about pivot.
        :param angle: Angle in degrees to rotate image by.
        :param out_width: width in pixels to fit the image into.
    
        :return: rotated and scaled output image.
        """
        img = copy.deepcopy(image)
        rotated_roi = ndimage.rotate(img, angle, reshape=False)
    
        zoom_scale = out_width / img.shape[0]
        return ndimage.zoom(rotated_roi, zoom_scale)
    
    
    # This is for Step 3.
    def modify_output_image(output_image: np.ndarray,
                            rotated_roi: np.ndarray):
        """Overwrites inplace the rotated ROI image to the output image."""
        THRESHOLD = 5
        idx = np.where(rotated_roi > THRESHOLD)
        output_image[idx] = rotated_roi[idx]
    
    
    # This is what you should run.
    def problem_logic(input_fits_image: np.ndarray, output_width) -> np.ndarray:
        """Takes input FITS image, and returns a modified image based on usecase."""
    
        output_image = np.zeros((output_width, output_width))
    
        extracted_roi = extract_roi(input_fits_image)
        pivot = [0, extracted_roi.shape[1] // 2]  # change this based on your req.
    
        pivoted_roi = square_pad_roi(extracted_roi, pivot)
    
        # I am assuming below as your rotation logic.
        T0TAL_STEPS = 12  # For each hour orientation.
        ROTATION_STEP = float(360 / T0TAL_STEPS)  # In degrees.
    
        for i in range(T0TAL_STEPS):
            rotated_image = rotate_and_scale(pivoted_roi, ROTATION_STEP * i,
                                             output_width)
            modify_output_image(output_image, rotated_image)
    
        return output_image
    
    
    def write_as_fits(image: np.ndarray, path: str):
        """ Writes the numpy image to the provided path in fits format"""
        hdu = fits.PrimaryHDU(image)
        hdu.writeto(path, overwrite=True)
    
    
    def driver_main():
        sample_fits = np.zeros((50, 100))
        sample_fits[:, 70:73] = 255
    
        # This creates the rotated ROI.
        # Replace sample_fits with your fits image.
        output_img = problem_logic(
            input_fits_image=sample_fits, output_width=400
        )
    
        # This writes it to output file in fits format.
        out_path = "sample.fits"
        write_as_fits(output_img, out_path)
    
        # This is just to visualize output.
        plt.imshow(output_img, cmap='gray')
        plt.show()
    
    
    if __name__ == "__main__":
        driver_main()