train/star-removal-from-astronomical-images-with-pix2pix.ipynb

{"cells":[{"cell_type":"markdown","metadata":{},"source":["![](https://github.com/code2k13/starreduction/raw/main/images/example.jpg)"]},{"cell_type":"markdown","metadata":{"execution":{"iopub.execute_input":"2021-09-13T04:37:26.761387Z","iopub.status.busy":"2021-09-13T04:37:26.760844Z","iopub.status.idle":"2021-09-13T04:37:26.853225Z","shell.execute_reply":"2021-09-13T04:37:26.851922Z","shell.execute_reply.started":"2021-09-13T04:37:26.761294Z"}},"source":["This notebook trains a GAN to remove stars from deep sky images. The code is inspired by [sample from tensorflow's website](https://www.tensorflow.org/tutorials/generative/pix2pix). The training data consists of only two images. One image of the Antenna Galaxy and another is a starmap that was created from a star cluster image.\n","                                                                                          \n","                                                                                        "]},{"cell_type":"markdown","metadata":{},"source":["![](https://github.com/code2k13/starreduction/raw/main/images/star_reduction_title.png)"]},{"cell_type":"markdown","metadata":{},"source":["## Attributions\n","\n","The training images used in this code were sourced from Wikimedia Commons and processed using GIMP.\n","### Antennae Galaxy Image\n","This image was downloaded from Wikimedia Commons and converted to grayscale using GIMP by me for purpose of training the model.\n","\n","Link to processed image: https://github.com/code2k13/starreduction/blob/main/training_data/Antennae_galaxies_xl.png\n",">[NASA, ESA, and the Hubble Heritage Team (STScI/AURA)-ESA/Hubble Collaboration](https://commons.wikimedia.org/wiki/File:Antennae_galaxies_xl.jpg), Public domain, via Wikimedia Commons\n","\n","Url : [https://commons.wikimedia.org/wiki/File:Antennae_galaxies_xl.jpg](https://commons.wikimedia.org/wiki/File:Antennae_galaxies_xl.jpg)\n","\n","Direct Link: [https://upload.wikimedia.org/wikipedia/commons/f/f6/Antennae_galaxies_xl.jpg](https://upload.wikimedia.org/wikipedia/commons/f/f6/Antennae_galaxies_xl.jpg)\n","\n","### Star cluster NGC 3572 and its surroundings\n","This image was downloaded from Wikimedia Commons and star mask was created by me using GIMP.\n","\n","Link to the processed image: https://github.com/code2k13/starreduction/blob/main/training_data/star_map_base.png\n","\n",">[ESO/G. Beccari](https://commons.wikimedia.org/wiki/File:The_star_cluster_NGC_3572_and_its_dramatic_surroundings.jpg\"), [https://creativecommons.org/licenses/by/4.0] (via Wikimedia Commons) \n","\n","Url: [https://commons.wikimedia.org/wiki/File:Antennae_galaxies_xl.jpg](https://commons.wikimedia.org/wiki/File:Antennae_galaxies_xl.jpg) \n","\n","Direct Link: [https://upload.wikimedia.org/wikipedia/commons/9/95/The_star_cluster_NGC_3572_and_its_dramatic_surroundings.jpg](https://upload.wikimedia.org/wikipedia/commons/9/95/The_star_cluster_NGC_3572_and_its_dramatic_surroundings.jpg)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:43.321921Z","iopub.status.busy":"2022-10-27T17:26:43.321492Z","iopub.status.idle":"2022-10-27T17:26:49.507719Z","shell.execute_reply":"2022-10-27T17:26:49.506396Z","shell.execute_reply.started":"2022-10-27T17:26:43.321797Z"},"trusted":true},"outputs":[],"source":["import tensorflow as tf\n","import os\n","import time\n","import matplotlib.pyplot as plt\n","from IPython.display import clear_output\n","from PIL import Image, ImageFilter,ImageEnhance\n","import cv2\n","import random\n","from matplotlib import pyplot as plt\n","import numpy as np\n","%matplotlib inline"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:49.510342Z","iopub.status.busy":"2022-10-27T17:26:49.509989Z","iopub.status.idle":"2022-10-27T17:26:49.515522Z","shell.execute_reply":"2022-10-27T17:26:49.514468Z","shell.execute_reply.started":"2022-10-27T17:26:49.510301Z"},"trusted":true},"outputs":[],"source":["IMG_SIZE = 1024\n","OUTPUT_CHANNELS = 1\n","LAMBDA = 10 "]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:49.517335Z","iopub.status.busy":"2022-10-27T17:26:49.516851Z","iopub.status.idle":"2022-10-27T17:26:49.526622Z","shell.execute_reply":"2022-10-27T17:26:49.525823Z","shell.execute_reply.started":"2022-10-27T17:26:49.517295Z"},"trusted":true},"outputs":[],"source":["def downsample(filters, size, apply_batchnorm=True,strides = 2,name=''):\n","  initializer = tf.random_normal_initializer(0., 0.02)\n","  result = tf.keras.Sequential(name=name)\n","  result.add(tf.keras.layers.Conv2D(filters, size, strides=strides, padding='same',kernel_initializer=initializer, use_bias=False))\n","\n","  if apply_batchnorm:\n","    result.add(tf.keras.layers.BatchNormalization())\n","\n","  result.add(tf.keras.layers.LeakyReLU())\n","  return result"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:49.529374Z","iopub.status.busy":"2022-10-27T17:26:49.528786Z","iopub.status.idle":"2022-10-27T17:26:49.537892Z","shell.execute_reply":"2022-10-27T17:26:49.537091Z","shell.execute_reply.started":"2022-10-27T17:26:49.529331Z"},"trusted":true},"outputs":[],"source":["def upsample(filters, size, apply_dropout=False,strides = 2,name=''):\n","  initializer = tf.random_normal_initializer(0., 0.02)\n","  result = tf.keras.Sequential(name=name)\n","  result.add(tf.keras.layers.Conv2DTranspose(filters, size, strides=strides,padding='same',\n","                                    kernel_initializer=initializer,use_bias=False))\n","\n","  result.add(tf.keras.layers.BatchNormalization())\n","  if apply_dropout:\n","      result.add(tf.keras.layers.Dropout(0.5))        \n","  result.add(tf.keras.layers.ReLU())\n","  return result"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:49.539829Z","iopub.status.busy":"2022-10-27T17:26:49.539265Z","iopub.status.idle":"2022-10-27T17:26:49.552664Z","shell.execute_reply":"2022-10-27T17:26:49.551952Z","shell.execute_reply.started":"2022-10-27T17:26:49.539788Z"},"trusted":true},"outputs":[],"source":["def Generator():\n","  inputs = tf.keras.layers.Input(shape=[IMG_SIZE, IMG_SIZE, OUTPUT_CHANNELS])\n","\n","  down_stack = [\n","    downsample(16, 5, apply_batchnorm=False,strides = 1, name='gd_1'),  # (batch_size, 128, 128, 64)\n","    downsample(32, 5,name='gd_2'),  # (batch_size, 64, 64, 128)\n","    downsample(64, 5,name='gd_3'),  # (batch_size, 32, 32, 256)\n","    downsample(128, 5,name='gd_4'),  # (batch_size, 16, 16, 512)\n","    downsample(256, 5,name='gd_5'),  # (batch_size, 8, 8, 512)\n","    downsample(256, 5,name='gd_6'),  # (batch_size, 4, 4, 512)\n","    downsample(512, 5,name='gd_7'),  # (batch_size, 2, 2, 512)\n","    downsample(512, 5,name='gd_8'),  # (batch_size, 1, 1, 512)\n","  ]\n","\n","  up_stack = [\n","    upsample(512, 5, apply_dropout=True,name='gu_1'),  # (batch_size, 2, 2, 1024)\n","    upsample(256, 5, apply_dropout=True,name='gu_2'),  # (batch_size, 4, 4, 1024)\n","    upsample(256, 5, apply_dropout=True,name='gu_3'),  # (batch_size, 8, 8, 1024)      \n","    upsample(128, 5,name='gu_4'),  # (batch_size, 16, 16, 1024)\n","    upsample(64, 5,name='gu_5'),  # (batch_size, 32, 32, 512)\n","    upsample(32, 5,name='gu_6'),  # (batch_size, 64, 64, 256)\n","    upsample(16, 5,name='gu_7'),  # (batch_size, 128, 128, 128)\n","  ]\n","\n","  initializer = tf.random_normal_initializer(0., 0.02)\n","  last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,strides=1,padding='same',\n","                                         kernel_initializer=initializer,activation='relu')  # (batch_size, 256, 256, 3)\n","\n","  x = inputs\n","  # Downsampling through the model\n","  skips = []\n","  for down in down_stack:\n","    x = down(x)\n","    skips.append(x)\n","\n","  skips = reversed(skips[:-1])\n","\n","  for up, skip in zip(up_stack, skips):\n","    x = up(x)\n","    x = tf.keras.layers.Concatenate()([x, skip])\n","\n","  x = last(x)\n","\n","  return tf.keras.Model(inputs=inputs, outputs=x)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:49.554837Z","iopub.status.busy":"2022-10-27T17:26:49.553948Z","iopub.status.idle":"2022-10-27T17:26:49.569269Z","shell.execute_reply":"2022-10-27T17:26:49.568588Z","shell.execute_reply.started":"2022-10-27T17:26:49.554792Z"},"trusted":true},"outputs":[],"source":["loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:49.571033Z","iopub.status.busy":"2022-10-27T17:26:49.570702Z","iopub.status.idle":"2022-10-27T17:26:49.581230Z","shell.execute_reply":"2022-10-27T17:26:49.580345Z","shell.execute_reply.started":"2022-10-27T17:26:49.570994Z"},"trusted":true},"outputs":[],"source":["def generator_loss(disc_generated_output, gen_output, target,disc_target_int_op,disc_gen_int_op):\n","  gan_loss = loss_object(tf.ones_like(disc_generated_output), disc_generated_output)\n","  l1_loss = tf.reduce_mean(tf.abs(target - gen_output))  \n","  gd1_loss = tf.reduce_mean(tf.abs(disc_target_int_op[0] - disc_gen_int_op[0]))\n","  gd2_loss = tf.reduce_mean(tf.abs(disc_target_int_op[1] - disc_gen_int_op[1]))\n","  gd3_loss = tf.reduce_mean(tf.abs(disc_target_int_op[2] - disc_gen_int_op[2]))\n","  gd4_loss = tf.reduce_mean(tf.abs(disc_target_int_op[3] - disc_gen_int_op[3]))\n","    \n","\n","  total_gen_loss = gan_loss + (LAMBDA * l1_loss) + gd1_loss + gd2_loss + gd3_loss + gd4_loss\n","  return total_gen_loss, gan_loss, l1_loss"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:49.583029Z","iopub.status.busy":"2022-10-27T17:26:49.582691Z","iopub.status.idle":"2022-10-27T17:26:49.597024Z","shell.execute_reply":"2022-10-27T17:26:49.596248Z","shell.execute_reply.started":"2022-10-27T17:26:49.582988Z"},"trusted":true},"outputs":[],"source":["\n","def Discriminator(): \n","  initializer = tf.random_normal_initializer(0., 0.02)\n","  inp = tf.keras.layers.Input(shape=[IMG_SIZE, IMG_SIZE, OUTPUT_CHANNELS], name='input_image')\n","  tar = tf.keras.layers.Input(shape=[IMG_SIZE, IMG_SIZE, OUTPUT_CHANNELS], name='target_image')\n","  x = tf.keras.layers.concatenate([inp, tar])  # (batch_size, 256, 256, channels*2)\n","  down0 = downsample(16, 4, False,name='dd_1')(x) \n","  down1 = downsample(32, 4,name='dd_2')(down0)  # (batch_size, 128, 128, 64)\n","  down2 = downsample(64, 4,name='dd_3')(down1)  # (batch_size, 64, 64, 128)\n","  down3 = downsample(128, 4,name='dd_4')(down2)  # (batch_size, 32, 32, 256)\n","\n","  zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3)  # (batch_size, 34, 34, 256)\n","  conv = tf.keras.layers.Conv2D(256, 4, strides=1,kernel_initializer=initializer,use_bias=False)(zero_pad1)  # (batch_size, 31, 31, 512)\n","\n","  batchnorm1 = tf.keras.layers.BatchNormalization()(conv)\n","  leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1)\n","  zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu)  # (batch_size, 33, 33, 512)\n","  last = tf.keras.layers.Conv2D(1, 4, strides=1,kernel_initializer=initializer)(zero_pad2)  # (batch_size, 30, 30, 1)\n","  return tf.keras.Model(inputs=[inp, tar], outputs=[last,down0,down1,down2,down3])"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:49.599684Z","iopub.status.busy":"2022-10-27T17:26:49.599163Z","iopub.status.idle":"2022-10-27T17:26:49.609855Z","shell.execute_reply":"2022-10-27T17:26:49.609220Z","shell.execute_reply.started":"2022-10-27T17:26:49.599641Z"},"trusted":true},"outputs":[],"source":["def discriminator_loss(disc_real_output, disc_generated_output):\n","  real_loss = loss_object(tf.ones_like(disc_real_output), disc_real_output)\n","  generated_loss = loss_object(tf.zeros_like(disc_generated_output), disc_generated_output)\n","  total_disc_loss = real_loss + generated_loss\n","  return total_disc_loss"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:49.613867Z","iopub.status.busy":"2022-10-27T17:26:49.613349Z","iopub.status.idle":"2022-10-27T17:26:49.631415Z","shell.execute_reply":"2022-10-27T17:26:49.630737Z","shell.execute_reply.started":"2022-10-27T17:26:49.613823Z"},"trusted":true},"outputs":[],"source":["generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n","discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:49.633014Z","iopub.status.busy":"2022-10-27T17:26:49.632408Z","iopub.status.idle":"2022-10-27T17:26:50.847320Z","shell.execute_reply":"2022-10-27T17:26:50.846347Z","shell.execute_reply.started":"2022-10-27T17:26:49.632965Z"},"trusted":true},"outputs":[],"source":["generator = Generator()\n","discriminator = Discriminator()"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:50.848897Z","iopub.status.busy":"2022-10-27T17:26:50.848633Z","iopub.status.idle":"2022-10-27T17:26:50.869076Z","shell.execute_reply":"2022-10-27T17:26:50.868233Z","shell.execute_reply.started":"2022-10-27T17:26:50.848855Z"},"trusted":true},"outputs":[],"source":["generator.summary()"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:50.870489Z","iopub.status.busy":"2022-10-27T17:26:50.870252Z","iopub.status.idle":"2022-10-27T17:26:50.882682Z","shell.execute_reply":"2022-10-27T17:26:50.881660Z","shell.execute_reply.started":"2022-10-27T17:26:50.870461Z"},"trusted":true},"outputs":[],"source":["discriminator.summary()"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:50.884682Z","iopub.status.busy":"2022-10-27T17:26:50.884242Z","iopub.status.idle":"2022-10-27T17:26:50.891374Z","shell.execute_reply":"2022-10-27T17:26:50.890342Z","shell.execute_reply.started":"2022-10-27T17:26:50.884643Z"},"trusted":true},"outputs":[],"source":["checkpoint_dir = './training_checkpoints'\n","checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n","checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,\n","                                 discriminator_optimizer=discriminator_optimizer,\n","                                 generator=generator,\n","                                 discriminator=discriminator)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:50.893217Z","iopub.status.busy":"2022-10-27T17:26:50.892633Z","iopub.status.idle":"2022-10-27T17:26:50.903505Z","shell.execute_reply":"2022-10-27T17:26:50.902472Z","shell.execute_reply.started":"2022-10-27T17:26:50.893185Z"},"trusted":true},"outputs":[],"source":["def train_step(input_image, target, step):\n","    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n","        gen_output = generator(input_image, training=True)\n","        \n","        disc_real_output,a,b,c,d = discriminator([input_image, target], training=True)\n","        dr1 = [a,b,c,d]\n","        \n","        disc_generated_output,e,f,g,h = discriminator([input_image, gen_output], training=True)\n","        dr2 = [e,f,g,h]\n","        \n","        gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(disc_generated_output, gen_output, target,\n","                                                                  dr1,dr2)\n","        disc_loss = discriminator_loss(disc_real_output, disc_generated_output)\n","\n","    generator_gradients = gen_tape.gradient(gen_total_loss,generator.trainable_variables)\n","    discriminator_gradients = disc_tape.gradient(disc_loss,discriminator.trainable_variables)\n","    generator_optimizer.apply_gradients(zip(generator_gradients,generator.trainable_variables))\n","    discriminator_optimizer.apply_gradients(zip(discriminator_gradients,discriminator.trainable_variables))\n","\n","    with summary_writer.as_default():\n","        tf.summary.scalar('gen_total_loss', gen_total_loss, step=step//1000)\n","        tf.summary.scalar('gen_gan_loss', gen_gan_loss, step=step//1000)\n","        tf.summary.scalar('gen_l1_loss', gen_l1_loss, step=step//1000)\n","        tf.summary.scalar('disc_loss', disc_loss, step=step//1000)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:50.905265Z","iopub.status.busy":"2022-10-27T17:26:50.905047Z","iopub.status.idle":"2022-10-27T17:26:52.664903Z","shell.execute_reply":"2022-10-27T17:26:52.664033Z","shell.execute_reply.started":"2022-10-27T17:26:50.905241Z"},"trusted":true},"outputs":[],"source":["import albumentations as A\n","transform = A.Compose([\n","        A.RandomRotate90(),\n","        A.Flip(),\n","        A.Transpose(),\n","        A.OneOf([\n","            A.IAAAdditiveGaussianNoise(),\n","            A.GaussNoise(),\n","        ], p=0.2),\n","        A.OneOf([\n","            A.MotionBlur(p=.2),\n","            A.MedianBlur(blur_limit=3, p=0.1),\n","            A.Blur(blur_limit=3, p=0.1),\n","        ], p=0.2),\n","        A.ShiftScaleRotate(shift_limit=0.5625, scale_limit=0.8, rotate_limit=145, p=0.2),\n","        A.OneOf([\n","            A.OpticalDistortion(p=0.3),\n","            A.GridDistortion(p=.1),\n","            A.IAAPiecewiseAffine(p=0.3),\n","        ], p=0.2),\n","        A.OneOf([\n","            A.CLAHE(clip_limit=2),\n","            A.IAASharpen(),\n","            A.IAAEmboss(),\n","            A.RandomBrightnessContrast(),            \n","        ], p=0.3),\n","        A.HueSaturationValue(p=0.3),\n","    ])\n","\n"," "]},{"cell_type":"markdown","metadata":{},"source":["We are generating images on the fly. "]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:26:52.666409Z","iopub.status.busy":"2022-10-27T17:26:52.666190Z","iopub.status.idle":"2022-10-27T17:27:00.311566Z","shell.execute_reply":"2022-10-27T17:27:00.310349Z","shell.execute_reply.started":"2022-10-27T17:26:52.666384Z"},"trusted":true},"outputs":[],"source":["!wget https://raw.githubusercontent.com/code2k13/starreduction/main/images/helix_bw_base.jpg\n","!wget https://github.com/code2k13/starreduction/raw/main/training_data/Antennae_galaxies_xl.png\n","!wget https://github.com/code2k13/starreduction/raw/main/training_data/star_map_base.png"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:27:00.313767Z","iopub.status.busy":"2022-10-27T17:27:00.313493Z","iopub.status.idle":"2022-10-27T17:27:03.026401Z","shell.execute_reply":"2022-10-27T17:27:03.025435Z","shell.execute_reply.started":"2022-10-27T17:27:00.313733Z"},"trusted":true},"outputs":[],"source":["!wget https://github.com/code2k13/starreduction/raw/main/training_data/turbulence.jpg"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:27:03.028321Z","iopub.status.busy":"2022-10-27T17:27:03.028032Z","iopub.status.idle":"2022-10-27T17:27:03.033184Z","shell.execute_reply":"2022-10-27T17:27:03.032273Z","shell.execute_reply.started":"2022-10-27T17:27:03.028269Z"},"trusted":true},"outputs":[],"source":["flip = False"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:45:24.001254Z","iopub.status.busy":"2022-10-27T17:45:24.000922Z","iopub.status.idle":"2022-10-27T17:45:24.018540Z","shell.execute_reply":"2022-10-27T17:45:24.017739Z","shell.execute_reply.started":"2022-10-27T17:45:24.001222Z"},"trusted":true},"outputs":[],"source":["def get_dataset_batch(batch_size = 2):\n","    flip = False    \n","    im2 = Image.open(\"star_map_base.png\")\n","    trainA = []\n","    trainB = []\n","    for i in range(0,batch_size):\n","        if flip == True:\n","            turbulence_size = random.choice([1,2])\n","            im = Image.open(\"turbulence.jpg\")\n","            x = random.randint(0,4096-IMG_SIZE*turbulence_size)\n","            y =  random.randint(0,2136-IMG_SIZE*turbulence_size)\n","            corp_actual_rect = (x, y, x+IMG_SIZE*turbulence_size, y+IMG_SIZE*turbulence_size)    \n","            corped_actual = im.crop(corp_actual_rect)\n","            corped_actual = corped_actual.resize((IMG_SIZE,IMG_SIZE))      \n","        else:\n","            helix_size = random.choice([1,2,3])\n","            im = Image.open(\"helix_bw_base.jpg\")\n","            x = random.randint(0,4096-IMG_SIZE*helix_size)\n","            y = random.randint(0,4096-IMG_SIZE*helix_size)\n","            corp_actual_rect = (x, y, x+IMG_SIZE*helix_size, y+IMG_SIZE*helix_size)\n","            corped_actual = im.crop(corp_actual_rect)\n","            corped_actual = corped_actual.resize((IMG_SIZE,IMG_SIZE))\n","        flip = not flip\n","        corped_actual = corped_actual.convert('LA')\n","        #enhancer1 = ImageEnhance.Brightness(corped_actual)\n","        #factor =   0.15 + random.random()*1\n","        star_overlayed = corped_actual #enhancer1.enhance(factor)  \n","        star_overlayed = star_overlayed.rotate(random.randint(1,360), expand=False) \n","        star_overlayed = star_overlayed.convert('L')\n","        star_overlayed = Image.fromarray(transform(image = np.asarray(star_overlayed))[\"image\"]/1)\n","        star_overlayed = star_overlayed.convert('LA')\n","        ca  = star_overlayed.copy()          \n","        ca = ca.convert('L') \n","        star_size = random.choice([1])\n","        x = random.randint(0,1024-IMG_SIZE*star_size)\n","        y =  random.randint(0,1024-IMG_SIZE*star_size)\n","        crop_rectangle = (x, y, x+IMG_SIZE*star_size, y+IMG_SIZE*star_size)\n","        scale_factor = random.choice([1,2,3,4,5,6])\n","        im2 = im2.resize((IMG_SIZE*scale_factor,IMG_SIZE*scale_factor))\n","        star_corped = im2.crop(crop_rectangle) \n","        star_corped = star_corped.resize((IMG_SIZE,IMG_SIZE))\n","        star_corped = star_corped.filter(ImageFilter.GaussianBlur(0.2))\n","        star_corped = star_corped.rotate(random.randint(1,360), expand=False)  \n","       \n","        #enhancer2 = ImageEnhance.Brightness(star_corped)\n","        #factor =  0.4  + random.random() * 0.2\n","        #star_enhanced= enhancer2.enhance(factor)      \n","        star_enhanced = star_corped.convert('RGBA')\n","        star_overlayed = star_overlayed.convert('RGBA')  \n","        op = overlay_images(star_overlayed,star_enhanced)\n","        op = Image.fromarray(op)\n","        star_overlayed = op.convert('L')\n","        a = np.asarray(ca,dtype=\"float32\").reshape(1,IMG_SIZE,IMG_SIZE,1)/512          \n","        b = np.asarray(star_overlayed,dtype=\"float32\").reshape(1,IMG_SIZE,IMG_SIZE,1)/512               \n","        trainA.append(a)\n","        trainB.append(b)    \n","    return trainA,trainB"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:45:24.473019Z","iopub.status.busy":"2022-10-27T17:45:24.472481Z","iopub.status.idle":"2022-10-27T17:45:24.479592Z","shell.execute_reply":"2022-10-27T17:45:24.478670Z","shell.execute_reply.started":"2022-10-27T17:45:24.472984Z"},"trusted":true},"outputs":[],"source":["\n","def overlay_images(background, foreground):\n","    b_img = np.array(background)\n","    f_img = np.array(foreground)    \n","    y1, y2 = 0,b_img.shape[0]\n","    x1, x2 = 0,b_img.shape[1]    \n","    m = np.max(b_img)    \n","    f_alpha = f_img[:, :, 3] / 255.0\n","    b_alpha = 1.0 - f_alpha\n","    for c in range(0, 3):\n","        b_img[y1:y2, x1:x2, c] = (f_alpha * f_img[:, :, c] +\n","                                  b_alpha * b_img[y1:y2, x1:x2, c]*2)\n","    return b_img"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:45:25.349798Z","iopub.status.busy":"2022-10-27T17:45:25.349519Z","iopub.status.idle":"2022-10-27T17:45:40.307228Z","shell.execute_reply":"2022-10-27T17:45:40.306292Z","shell.execute_reply.started":"2022-10-27T17:45:25.349771Z"},"trusted":true},"outputs":[],"source":["plt.rcParams['figure.figsize'] = (30, 10)\n","random.seed(10)\n","test_starless, test_stars = get_dataset_batch(20)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:45:40.309248Z","iopub.status.busy":"2022-10-27T17:45:40.308940Z","iopub.status.idle":"2022-10-27T17:45:41.064882Z","shell.execute_reply":"2022-10-27T17:45:41.064134Z","shell.execute_reply.started":"2022-10-27T17:45:40.309209Z"},"trusted":true},"outputs":[],"source":["\n","plt.imshow(test_stars[2].reshape((IMG_SIZE,IMG_SIZE,1)),cmap='gray')\n","plt.show()\n","plt.imshow(test_starless[2].reshape((IMG_SIZE,IMG_SIZE,1)),cmap='gray')\n","plt.show()"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:27:11.709971Z","iopub.status.busy":"2022-10-27T17:27:11.709475Z","iopub.status.idle":"2022-10-27T17:27:11.725396Z","shell.execute_reply":"2022-10-27T17:27:11.724340Z","shell.execute_reply.started":"2022-10-27T17:27:11.709922Z"},"trusted":true},"outputs":[],"source":["import datetime\n","log_dir=\"logs/\"\n","summary_writer = tf.summary.create_file_writer(log_dir + \"fit/\" + datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))"]},{"cell_type":"markdown","metadata":{},"source":["There is no real notion of 'epochs' here. Everytime random image is being created for training. We have some control over repetability using the 'n' variable which is reset to 'total_samples' everytime it exceeds 'total_samples'. This 'n' variable is used as random seed before creation of images."]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2022-10-27T17:27:11.727609Z","iopub.status.busy":"2022-10-27T17:27:11.726852Z","iopub.status.idle":"2022-10-27T17:39:38.117729Z","shell.execute_reply":"2022-10-27T17:39:38.116261Z","shell.execute_reply.started":"2022-10-27T17:27:11.727566Z"},"trusted":true},"outputs":[],"source":["from matplotlib.pyplot import figure\n","EPOCHS = 12000\n","SAVE_MODEL_AFTER = 400\n","n = 0\n","total_samples = 500\n","BATCH_SIZE = 2\n","plt.rcParams['figure.figsize'] = (15, 5)\n","\n","for epoch in range(EPOCHS):\n","    #random.seed(n)\n","    #print(epoch)\n","    train_starless, train_stars = get_dataset_batch(BATCH_SIZE)\n","    start = time.time()\n","    train_step(np.asarray(train_stars).reshape((-1,IMG_SIZE,IMG_SIZE,OUTPUT_CHANNELS)),\n","               np.asarray(train_starless).reshape((-1,IMG_SIZE,IMG_SIZE,OUTPUT_CHANNELS)),n)\n","    \n","    if (epoch + 1) % SAVE_MODEL_AFTER == 0:\n","        generator.save_weights(\"generator\" + str(epoch))\n","        #ckpt_save_path = checkpoint.save('test_')\n","        print ('Saving checkpoint for epoch {} '.format(epoch+1))\n","        sample_prediction = generator(test_stars[8], training=False)\n","        f, axarr = plt.subplots(1,3)\n","        axarr[0].imshow(test_starless[8].reshape((IMG_SIZE,IMG_SIZE,OUTPUT_CHANNELS)),cmap=\"gray\")\n","        axarr[1].imshow(test_stars[8].reshape((IMG_SIZE,IMG_SIZE,OUTPUT_CHANNELS)),cmap=\"gray\")\n","        axarr[2].imshow(sample_prediction.numpy().reshape((IMG_SIZE,IMG_SIZE,OUTPUT_CHANNELS)),cmap=\"gray\")\n","        #axarr.imshow()\n","        plt.show()\n","        sample_prediction = generator(test_stars[5], training=False)\n","        f, axarr = plt.subplots(1,3)\n","        axarr[0].imshow(test_starless[5].reshape((IMG_SIZE,IMG_SIZE,OUTPUT_CHANNELS)),cmap=\"gray\")\n","        axarr[1].imshow(test_stars[5].reshape((IMG_SIZE,IMG_SIZE,OUTPUT_CHANNELS)),cmap=\"gray\")\n","        axarr[2].imshow(sample_prediction.numpy().reshape((IMG_SIZE,IMG_SIZE,OUTPUT_CHANNELS)),cmap=\"gray\")\n","        #axarr.imshow()\n","        plt.show()\n","    n = n + BATCH_SIZE\n","    if n > total_samples:\n","        n = 0"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.status.busy":"2022-10-27T17:39:38.119258Z","iopub.status.idle":"2022-10-27T17:39:38.119728Z","shell.execute_reply":"2022-10-27T17:39:38.119521Z","shell.execute_reply.started":"2022-10-27T17:39:38.119498Z"},"trusted":true},"outputs":[],"source":["!rm test*"]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.8.13 (default, Sep  6 2022, 16:33:03) \n[GCC 9.4.0]"},"vscode":{"interpreter":{"hash":"9ac03a0a6051494cc606d484d27d20fce22fb7b4d169f583271e11d5ba46a56e"}}},"nbformat":4,"nbformat_minor":4}