# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %load_ext autoreload
# %autoreload 2

# +
from __future__ import print_function

import numpy as np

from astropy.io import fits
# -

np.random.seed(42)

# +

generator = rotate_angle_generator()

def image_augementation_generator(zoom_factor):
    def _rotate_angle_generator():
        while True:
            angle = np.random.randint(0, 360)
            yield angle

    def _zoom_generator():
        while True:
            zoom = np.round(zoom_factor * (np.random.random() - 0.5), 3) + 1.0
            yield zoom

    def _boolean_generator():
        while True:
            prob = np.random.random()
            yield prob >= 0.5

    rotate_angle_generator = _rotate_angle_generator()
    zoom_generator = _zoom_generator()
    horiztonal_flip_generator = _boolean_generator()
    vertical_flip_generator = _boolean_generator()
    
    while True:
        rotate_angle = next(rotate_angle_generator)
        zoom = next(zoom_generator)
        horiztonal_flip = next(horiztonal_flip_generator)
        vertical_flip = next(vertical_flip_generator)
        
        yield rotate_angle, zoom, horiztonal_flip, vertical_flip

for idx, value in enumerate(image_augementation_generator(0.1)):
    print(value)
    if idx == 20:
        break
# -

np.random.random()

# %ls -l

fits_file = fits.open('obj-1237658491746254860.fits.bz2')

# +
from astropy.wcs import WCS

WCS(fits_file[0].header)


# -

def plot_images(imgs, size):
    num_rows = int(np.ceil(size/3.0))
    figsize_y = 5 * num_rows

    fig = plt.figure(figsize=(20,figsize_y))
    for idx in range(0, size):
        img = imgs[idx]
        # make scale between 0 and 1.0 plotting
        img_min = img.min()
        img_max = img.max()
        img = (img - img_min) / (img_max - img_min)

        fig.add_subplot(num_rows, 3, idx + 1)
        plt.imshow(img, cmap=plt.get_cmap('gray'), interpolation='bicubic')

    plt.show()


# +
from PIL import Image

def flip_horizontal(data):
    return np.flip(data, 2)

def flip_vertical(data):
    return np.flip(data, 1)

def rotate(data, deg):
    for idx in range(data.shape[0]):
        img = Image.fromarray(data[idx])
        img = img.rotate(deg, resample=Image.BICUBIC, expand=True)
        
        if (idx == 0):
            out = np.empty((3,) + img.size)

        out[idx] = np.array(img)

    return out

def resize(data, factor):
    new_size = int(np.round(factor * data.shape[1]))
    
    out = np.empty((3, new_size, new_size))
    
    for idx in range(data.shape[0]):
        img = Image.fromarray(data[idx])
        img = img.resize((new_size, new_size), resample=Image.BICUBIC)
        out[idx] = np.array(img)
    return out

def crop(data, crop_size=42):
    curr_size = data.shape[1]
    out = np.empty((3, 42, 42))

    top = int((curr_size - crop_size)/2)
    bttm = top + crop_size

    for idx in range(3):
        out[idx] = data[idx, top:bttm, top:bttm]

    return out

img_data = fits_file[0].data

outputs = np.empty((9, 3, 42, 42))
outputs[0] = sgd.scale_rgb(crop(img_data), sigma=1/6)
generator = image_augementation_generator(0.5)
for idx in range(1,9):
    angle, zoom, flip_h, flip_v = next(generator)
    print(f'Rotate: {angle}, Zoom: {zoom}, Flip Horizon.: {flip_h}, Flip Vertical: {flip_v}')
    data = rotate(img_data, angle)
    data = resize(data, zoom)

    if flip_h:
        data = flip_horizontal(data)
    
    if flip_v:
        data = flip_vertical(data)
        
    data = crop(data)
    outputs[idx] = sgd.scale_rgb(data, sigma=1/6)

plot_images(np.moveaxis(outputs, 1, -1), 9)


# +
import sdss_gz_data as sgd
import matplotlib.pyplot as plt

#Image.fromarray(np.moveaxis(sgd.scale_rgb(img_data), 0, -1))
plt.imshow(np.moveaxis(sgd.scale_rgb(crop(img_data)), 0, -1), interpolation='bicubic')
# -

plt.imshow(np.moveaxis(sgd.scale_rgb(crop(rotated_image)), 0, -1), interpolation='bicubic')

plt.imshow(np.moveaxis(sgd.scale_rgb(crop(resized)), 0, -1), interpolation='bicubic')

from astropy.io import fits

fits_file = fits.open('data/obj-1237650369407221944.fits.bz2')
data = fits_file[0].data

data.shape

nans = np.isnan(data)

data[nans] = 0

np.min(data)

data[nans] = np.min(data)

data



StarcoderdataJupyterScriptsDedupFiltered