Update inference.py

inference.py

@@ -5,21 +5,13 @@ import numpy as np
 from PIL import Image
 
 IMG_SIZE = 256
-
-# ------------------------------
-# Load the generator model from Hugging Face Hub
-# ------------------------------
 model_path = hf_hub_download(
-    repo_id="saraayum/cgan4colsar",   # your Hugging Face repo
+    repo_id="saraayum/cgan4colsar",  
     filename="generator_epoch_180.keras"
 )
 
-# Safe model loading
 generator = load_model(model_path, compile=False, safe_mode=False)
 
-# ------------------------------
-# Preprocess the SAR image (normalize [0, 1])
-# ------------------------------
 def preprocess_image(image_path):
     img = tf.io.read_file(image_path)
     img = tf.image.decode_png(img, channels=1)        # SAR = grayscale
@@ -28,9 +20,6 @@ def preprocess_image(image_path):
     img = tf.expand_dims(img, 0)                      # Add batch dimension
     return img
 
-# ------------------------------
-# Postprocess the generator output
-# ------------------------------
 def postprocess_output(output_tensor):
     # Convert output from [0,1] to [0,255]
     output_tensor = output_tensor[0] * 255.0
@@ -38,19 +27,13 @@ def postprocess_output(output_tensor):
     output_image = tf.cast(output_tensor, tf.uint8)
     return Image.fromarray(output_image.numpy())
 
-# ------------------------------
-# Predict function for inference
-# ------------------------------
 def predict(image_path, save_path="output.png"):
     sar_input = preprocess_image(image_path)
     gen_output = generator(sar_input, training=False)
     output_image = postprocess_output(gen_output)
     output_image.save(save_path)
-    print(f"✅ Colorized image saved as: {save_path}")
+    print(f"Colorized image saved as: {save_path}")
     return output_image
 
-# ------------------------------
-# Example usage (comment this on HF)
-# ------------------------------
 if __name__ == "__main__":
     predict("sample_sar.png", "predicted_colorized.png")