Migrated to U2NETP and onnx

Author: code2k13

README.md

@@ -1,6 +1,6 @@
 # Star reduction in deep sky images
 
-Starrem2k13 is a simple tool for removing stars from astronomical images. Starrem2k13 uses a GAN trained on augmented data. It's code was inspired from a [sample at Tensorflow's website](https://www.tensorflow.org/tutorials/generative/pix2pix). The training data consists of only three base images.
+Starrem2k13 is a simple tool for removing stars from astronomical images. Starrem2k13 uses a U2NETP model trained on augmented data. The training data consists of only three base images.
 
 Below are examples of what it can do:
 
@@ -81,4 +81,11 @@ Url: [https://commons.wikimedia.org/wiki/File:The_star_cluster_NGC_3572_and_its_
 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)
 
 
-
+@InProceedings{Qin_2020_PR,
+title = {U2-Net: Going Deeper with Nested U-Structure for Salient Object Detection},
+author = {Qin, Xuebin and Zhang, Zichen and Huang, Chenyang and Dehghan, Masood and Zaiane, Osmar and Jagersand, Martin},
+journal = {Pattern Recognition},
+volume = {106},
+pages = {107404},
+year = {2020}
+}

export_to_onnx.py

@@ -0,0 +1,17 @@
+import tensorflow as tf
+import tf2onnx
+import model  # your model.py
+
+# Build and load weights
+G2 = model.Generator()
+G2.load_weights("weights/generator_epoch_1000.weights.h5")
+
+# Optional: Run once to build the model (especially if using Functional API or Subclassed model)
+G2.predict(tf.random.normal([1, 512, 512]))  # Adjust input shape as needed
+
+# Convert to ONNX
+spec = (tf.TensorSpec((1, 512, 512), tf.float32),)  # shape should match model input
+output_path = "generator.onnx"
+
+model_proto, _ = tf2onnx.convert.from_keras(G2, input_signature=spec, opset=13, output_path=output_path)
+print(f"ONNX model saved to {output_path}")

model.py

@@ -1,72 +1,62 @@
-import tensorflow as tf 
-
-IMG_SIZE = 512
-OUTPUT_CHANNELS = 1
-
-def downsample(filters, size, apply_batchnorm=True,strides = 2,name=''):
-  initializer = tf.random_normal_initializer(0., 0.02)
-  result = tf.keras.Sequential(name=name)
-  result.add(tf.keras.layers.Conv2D(filters, size, strides=strides, padding='same',kernel_initializer=initializer, use_bias=False))
-
-  if apply_batchnorm:
-    result.add(tf.keras.layers.BatchNormalization())
-
-  result.add(tf.keras.layers.LeakyReLU())
-  return result
-
-def upsample(filters, size, apply_dropout=False,strides = 2,name=''):
-  initializer = tf.random_normal_initializer(0., 0.02)
-  result = tf.keras.Sequential(name=name)
-  result.add(tf.keras.layers.Conv2DTranspose(filters, size, strides=strides,padding='same',
-                                    kernel_initializer=initializer,use_bias=False))
-
-  result.add(tf.keras.layers.BatchNormalization())
-  if apply_dropout:
-      result.add(tf.keras.layers.Dropout(0.5))        
-  result.add(tf.keras.layers.ReLU())
-  return result
-
-def Generator():
-  inputs = tf.keras.layers.Input(shape=[IMG_SIZE, IMG_SIZE, OUTPUT_CHANNELS])
-
-  down_stack = [
-    downsample(16, 5, apply_batchnorm=False,strides = 1, name='gd_1'),  # (batch_size, 128, 128, 64)
-    downsample(32, 5,name='gd_2'),  # (batch_size, 64, 64, 128)
-    downsample(64, 5,name='gd_3'),  # (batch_size, 32, 32, 256)
-    downsample(128, 5,name='gd_4'),  # (batch_size, 16, 16, 512)
-    downsample(256, 5,name='gd_5'),  # (batch_size, 8, 8, 512)
-    downsample(256, 5,name='gd_6'),  # (batch_size, 4, 4, 512)
-    downsample(512, 5,name='gd_7'),  # (batch_size, 2, 2, 512)
-    downsample(512, 5,name='gd_8'),  # (batch_size, 1, 1, 512)
-  ]
-
-  up_stack = [
-    upsample(512, 5, apply_dropout=True,name='gu_1'),  # (batch_size, 2, 2, 1024)
-    upsample(256, 5, apply_dropout=True,name='gu_2'),  # (batch_size, 4, 4, 1024)
-    upsample(256, 5, apply_dropout=True,name='gu_3'),  # (batch_size, 8, 8, 1024)      
-    upsample(128, 5,name='gu_4'),  # (batch_size, 16, 16, 1024)
-    upsample(64, 5,name='gu_5'),  # (batch_size, 32, 32, 512)
-    upsample(32, 5,name='gu_6'),  # (batch_size, 64, 64, 256)
-    upsample(16, 5,name='gu_7'),  # (batch_size, 128, 128, 128)
-  ]
-
-  initializer = tf.random_normal_initializer(0., 0.02)
-  last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,strides=1,padding='same',
-                                         kernel_initializer=initializer,activation='relu')  # (batch_size, 256, 256, 3)
-
-  x = inputs
-  # Downsampling through the model
-  skips = []
-  for down in down_stack:
-    x = down(x)
-    skips.append(x)
-
-  skips = reversed(skips[:-1])
-
-  for up, skip in zip(up_stack, skips):
-    x = up(x)
-    x = tf.keras.layers.Concatenate()([x, skip])
-
-  x = last(x)
-
-  return tf.keras.Model(inputs=inputs, outputs=x)
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras import layers
+
+def REBNCONV(x, filters, dilation_rate=1, name_prefix="rebn"):
+    x = layers.Conv2D(filters, 3, padding='same', dilation_rate=dilation_rate,
+                      kernel_initializer='he_normal', name=f"{name_prefix}_conv")(x)
+    x = layers.BatchNormalization(name=f"{name_prefix}_bn")(x)
+    x = layers.ReLU(name=f"{name_prefix}_relu")(x)
+    return x
+
+def RSU4(x_input, in_filters, mid_filters, out_filters, name_prefix="rsu4"):
+    hxin = REBNCONV(x_input, out_filters, name_prefix=f"{name_prefix}_in")
+
+    hx1 = REBNCONV(hxin, mid_filters, name_prefix=f"{name_prefix}_conv1")
+    pool1 = layers.MaxPooling2D(pool_size=2)(hx1)
+
+    hx2 = REBNCONV(pool1, mid_filters, name_prefix=f"{name_prefix}_conv2")
+    pool2 = layers.MaxPooling2D(pool_size=2)(hx2)
+
+    hx3 = REBNCONV(pool2, mid_filters, name_prefix=f"{name_prefix}_conv3")
+    hx4 = REBNCONV(hx3, mid_filters, dilation_rate=2, name_prefix=f"{name_prefix}_conv4")
+
+    hx3d = REBNCONV(layers.Concatenate()([hx4, hx3]), mid_filters, name_prefix=f"{name_prefix}_conv3d")
+    hx3d_up = layers.UpSampling2D(size=2, interpolation='bilinear')(hx3d)
+
+    hx2d = REBNCONV(layers.Concatenate()([hx3d_up, hx2]), mid_filters, name_prefix=f"{name_prefix}_conv2d")
+    hx2d_up = layers.UpSampling2D(size=2, interpolation='bilinear')(hx2d)
+
+    hx1d = REBNCONV(layers.Concatenate()([hx2d_up, hx1]), out_filters, name_prefix=f"{name_prefix}_conv1d")
+
+    return layers.Add(name=f"{name_prefix}_add")([hx1d, hxin])
+
+def Generator(input_shape=(512, 512, 1)):
+    inputs = keras.Input(shape=input_shape)
+
+    stage1 = RSU4(inputs, 1, 16, 64, name_prefix="stage1")
+    pool12 = layers.MaxPooling2D(pool_size=2)(stage1)
+
+    stage2 = RSU4(pool12, 64, 16, 64, name_prefix="stage2")
+    pool23 = layers.MaxPooling2D(pool_size=2)(stage2)
+
+    stage3 = RSU4(pool23, 64, 16, 64, name_prefix="stage3")
+    pool34 = layers.MaxPooling2D(pool_size=2)(stage3)
+
+    stage4 = RSU4(pool34, 64, 16, 64, name_prefix="stage4")
+
+    stage3d = RSU4(layers.Concatenate()([
+        layers.UpSampling2D(size=2, interpolation='bilinear')(stage4), stage3
+    ]), 128, 16, 64, name_prefix="stage3d")
+
+    stage2d = RSU4(layers.Concatenate()([
+        layers.UpSampling2D(size=2, interpolation='bilinear')(stage3d), stage2
+    ]), 128, 16, 64, name_prefix="stage2d")
+
+    stage1d = RSU4(layers.Concatenate()([
+        layers.UpSampling2D(size=2, interpolation='bilinear')(stage2d), stage1
+    ]), 128, 16, 64, name_prefix="stage1d")
+
+    output = layers.Conv2D(1, 1, padding='same', activation='sigmoid', name="final_output")(stage1d)
+
+    return keras.Model(inputs, output, name="U2NETP_Gray2Gray")

requirements.txt

@@ -1,6 +1,7 @@
-tensorflow
 numpy
 pillow
 tqdm
 tensorflowjs
-pyinstaller
+pyinstaller
+tf2onnx
+onnxruntime

samples_starless/app.py

@@ -1,16 +1,11 @@
 #!/usr/bin/python3
 
-import os
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
-import tensorflow as tf
-tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
-
-import model
 from PIL import Image
 import numpy as np
 import sys
 from tqdm import tqdm
 import math
+import onnxruntime as ort
 
 
 IMG_SIZE = 512
@@ -27,38 +22,46 @@
   print ('Argument List:', str(sys.argv)) 
   exit(1)
 
-def process_tile(channel,i,j,pad_width,model,output_image):
-  corp_rect = (i*IMG_SIZE,j*IMG_SIZE,i*IMG_SIZE+IMG_SIZE,j*IMG_SIZE+IMG_SIZE)
-  current_tile = channel.crop(corp_rect)
-  current_tile = current_tile.convert('L')
-  blank_image =  current_tile.copy()  
-  current_tile = current_tile.resize((IMG_SIZE-pad_width*2,IMG_SIZE-pad_width*2))
-  blank_image.paste(current_tile,(pad_width,pad_width))
-  blank_image = blank_image.resize((MODEL_SIZE,MODEL_SIZE))
-  blank_image  = np.asarray(blank_image,dtype="float32").reshape(1,MODEL_SIZE,MODEL_SIZE)/512
-  predicted_section = model.predict(blank_image,verbose=0)      
-  predicted_section = predicted_section.reshape(MODEL_SIZE,MODEL_SIZE)*512
-  predicted_section = Image.fromarray(predicted_section).convert('L')
-  predicted_section =  predicted_section.resize((IMG_SIZE,IMG_SIZE))            
-  predicted_section = predicted_section.crop((pad_width,pad_width,IMG_SIZE-pad_width,IMG_SIZE-pad_width))
-  predicted_section = predicted_section.resize((IMG_SIZE,IMG_SIZE))  
-  output_image.paste(predicted_section, (i*IMG_SIZE,j*IMG_SIZE),  mask=None)
+def process_tile(channel,i,j,pad_width,output_image):
+    corp_rect = (i * IMG_SIZE, j * IMG_SIZE, i * IMG_SIZE + IMG_SIZE, j * IMG_SIZE + IMG_SIZE)
+    current_tile = channel.crop(corp_rect).convert('L')
+    blank_image = current_tile.copy()
+
+    # Resize to remove padding, then paste into padded blank
+    current_tile = current_tile.resize((IMG_SIZE - pad_width * 2, IMG_SIZE - pad_width * 2))
+    blank_image.paste(current_tile, (pad_width, pad_width))
+
+    # Resize to match model input
+    blank_image = blank_image.resize((MODEL_SIZE, MODEL_SIZE))
+    input_array = np.asarray(blank_image, dtype="float32").reshape(1, MODEL_SIZE, MODEL_SIZE)
+    input_array = input_array /382  # Normalize same as training
+
+    # ONNX Inference
+    output_array = onnx_session.run(None, {input_name: input_array})[0]
+    output_array = output_array.reshape(MODEL_SIZE, MODEL_SIZE) * 382  # De-normalize
+
+    # Convert to image and paste
+    predicted_section = Image.fromarray(output_array.astype(np.uint8)).convert('L')
+    predicted_section = predicted_section.resize((IMG_SIZE, IMG_SIZE))
+    predicted_section = predicted_section.crop((pad_width, pad_width, IMG_SIZE - pad_width, IMG_SIZE - pad_width))
+    predicted_section = predicted_section.resize((IMG_SIZE, IMG_SIZE))
+
+    output_image.paste(predicted_section, (i * IMG_SIZE, j * IMG_SIZE), mask=None)
 
 def process_channel(channel,pad_width,input_image_size):    
     global progress_bar,step_size,current_progress
     output_image =  Image.new('L', input_image_size)
     for i in range(0,int(channel.size[0]/IMG_SIZE)):
         for j in range(0,int(channel.size[1]/IMG_SIZE)):                   
-            process_tile(channel,i,j,pad_width,G2,output_image)
+            process_tile(channel,i,j,pad_width,output_image)
             current_progress = current_progress + step_size
             if current_progress <= 100:
                 progress_bar.update(round(step_size,2))
     return output_image
 
 try:
-  G2 = model.Generator()
-  G2.load_weights("weights/weights")
-
+  onnx_session = ort.InferenceSession("weights/model.onnx", providers=["CPUExecutionProvider"])
+  input_name = onnx_session.get_inputs()[0].name
   source_image = Image.open(args[1])
   mode = source_image.mode
   size = source_image.size