Upload 6 files

app.py

@@ -0,0 +1,71 @@
+import gradio as gr
+from PIL import Image
+import os
+import tensorflow as tf
+import requests
+
+from .models.auto_encoder_gray2color import SpatialAttention
+
+# Load the saved model once at startup
+load_model_path = "./ckpts/best_model.h5"
+if not os.path.exists(load_model_path):
+    os.makedirs(os.path.dirname(load_model_path), exist_ok=True)
+    url = "https://huggingface.co/danhtran2mind/autoencoder-grayscale-to-color-landscape/resolve/main/ckpts/best_model.h5"
+    print(f"Downloading model from {url}...")
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status()
+        with open(load_model_path, "wb") as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                f.write(chunk)
+    print("Download complete.")
+
+print(f"Loading model from {load_model_path}...")
+loaded_autoencoder = tf.keras.models.load_model(
+    load_model_path,
+    custom_objects={'SpatialAttention': SpatialAttention}
+)
+
+def process_image(input_img):
+    # Convert PIL Image to numpy array and normalize
+    img = input_img.convert("RGB")
+    img = img.resize((256, 256))  # adjust size as needed
+    img_array = tf.keras.preprocessing.image.img_to_array(img) / 255.0
+    img_array = img_array[None, ...]  # add batch dimension
+
+    # Run inference
+    output_array = loaded_autoencoder.predict(img_array)
+    output_img = tf.keras.preprocessing.image.array_to_img(output_array[0])
+
+    return output_img
+
+custom_css = """
+body {background: linear-gradient(135deg, #232526 0%, #414345 100%) !important;}
+.gradio-container {background: transparent !important;}
+h1, .gr-title {color: #00e6d3 !important; font-family: 'Segoe UI', sans-serif;}
+.gr-description {color: #e0e0e0 !important; font-size: 1.1em;}
+.gr-input, .gr-output {border-radius: 18px !important; box-shadow: 0 4px 24px rgba(0,0,0,0.18);}
+.gr-button {background: linear-gradient(90deg, #00e6d3 0%, #0072ff 100%) !important; color: #fff !important; border: none !important; border-radius: 12px !important;}
+"""
+
+demo = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil", label="Upload Grayscale Landscape", image_mode="L", shape=(256, 256)),
+    outputs=gr.Image(type="pil", label="Colorized Output"),
+    title="🌄 Gray2Color Landscape Autoencoder",
+    description=(
+        "<div style='font-size:1.15em;line-height:1.6em;'>"
+        "Transform your <b>grayscale landscape</b> photos into vivid color with a state-of-the-art autoencoder.<br>"
+        "Simply upload a grayscale image and see the magic happen!"
+        "</div>"
+    ),
+    theme="soft",
+    css=custom_css,
+    allow_flagging="never",
+    examples=[
+        ["examples/grayscale_landscape1.jpg"],
+        ["examples/grayscale_landscape2.jpg"]
+    ]
+)
+
+if __name__ == "__main__":
+    demo.launch()

ckpts/best_model.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6d0361fa140c1dc3b279bcce8107c28b6e10a4e1bc31f770e5b071a44f5f76d
+size 20800096

dataset/README.md

@@ -0,0 +1,42 @@
+```markdown
+# Landscape Pictures Dataset Processing
+
+[![Kaggle](https://img.shields.io/badge/Dataset-Kaggle-blue.svg)](https://www.kaggle.com/datasets/arnaud58/landscape-pictures)
+
+This README provides instructions for downloading, extracting, and processing the landscape pictures dataset from Kaggle.
+
+## Dataset Source
+
+The dataset is sourced from Kaggle: Landscape Pictures by Arnaud58. Follow this link: [Kaggle Dataset](https://www.kaggle.com/datasets/arnaud58/landscape-pictures)
+
+## Setup
+
+1. **Create a Dataset Directory**: Create a directory to store the dataset:
+
+```python
+import os
+
+ds_path = "./dataset/landscape-pictures"
+os.makedirs(ds_path, exist_ok=True)
+```
+
+2. **Download the Dataset**: Use the following command to download the dataset from Kaggle:
+
+```bash
+curl -L https://www.kaggle.com/api/v1/datasets/download/arnaud58/landscape-pictures -o ./dataset/landscape-pictures.zip
+```
+
+Note: You may need a Kaggle API token for authentication. Ensure you have the `kaggle.json` file configured in `~/.kaggle/` or set up the Kaggle API as per Kaggle's API documentation.
+
+3. **Extract the Dataset**: Run the following Python code to extract the downloaded zip file:
+
+```python
+import zipfile
+import os
+
+with zipfile.ZipFile('dataset/landscape-pictures.zip', 'r') as zip_ref:
+    zip_ref.extractall(ds_path)
+```
+
+This will extract the dataset into the `./dataset` directory.
+```

models/auto_encoder_gray2color.py

@@ -0,0 +1,92 @@
+
+import tensorflow as tf
+from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, BatchNormalization, Add, Concatenate, Multiply
+from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import Adam
+
+# Spatial Attention Layer
+# Define SpatialAttention layer
+class SpatialAttention(tf.keras.layers.Layer):
+    def __init__(self, kernel_size=7, **kwargs):
+        super(SpatialAttention, self).__init__(**kwargs)
+        self.kernel_size = kernel_size
+        self.conv = Conv2D(filters=1, kernel_size=kernel_size, padding='same', activation='sigmoid')
+
+    def call(self, inputs):
+        avg_pool = tf.reduce_mean(inputs, axis=-1, keepdims=True)
+        max_pool = tf.reduce_max(inputs, axis=-1, keepdims=True)
+        concat = Concatenate()([avg_pool, max_pool])
+        attention = self.conv(concat)
+        return Multiply()([inputs, attention])
+
+    def get_config(self):
+        config = super(SpatialAttention, self).get_config()
+        config.update({'kernel_size': self.kernel_size})
+        return config
+        
+# Build Autoencoder
+def build_autoencoder(height, width,):
+    input_img = Input(shape=(height, width, 1))
+    
+    # Encoder
+    x = Conv2D(96, (3, 3), activation='relu', padding='same')(input_img)
+    x = BatchNormalization()(x)
+    x = SpatialAttention()(x)
+    x = MaxPooling2D((2, 2), padding='same')(x)
+    
+    # Residual Block 1
+    residual = Conv2D(192, (1, 1), padding='same')(x)
+    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Add()([x, residual])
+    x = SpatialAttention()(x)
+    x = MaxPooling2D((2, 2), padding='same')(x)
+    
+    # Residual Block 2
+    residual = Conv2D(384, (1, 1), padding='same')(x)
+    x = Conv2D(384, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Conv2D(384, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Add()([x, residual])
+    x = SpatialAttention()(x)
+    encoded = MaxPooling2D((2, 2), padding='same')(x)
+    
+    # Decoder
+    x = Conv2D(384, (3, 3), activation='relu', padding='same')(encoded)
+    x = BatchNormalization()(x)
+    x = SpatialAttention()(x)
+    x = UpSampling2D((2, 2))(x)
+    
+    # Residual Block 3
+    residual = Conv2D(192, (1, 1), padding='same')(x)
+    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Add()([x, residual])
+    x = SpatialAttention()(x)
+    x = UpSampling2D((2, 2))(x)
+    
+    x = Conv2D(96, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = SpatialAttention()(x)
+    x = UpSampling2D((2, 2))(x)
+    
+    decoded = Conv2D(2, (3, 3), activation=None, padding='same')(x)
+    
+    return Model(input_img, decoded)
+
+
+
+
+
+if __name__ == "__main__":
+    # Define constants
+    HEIGHT, WIDTH = 512, 512
+    # Compile model
+    autoencoder = build_autoencoder()
+    autoencoder.summary()
+    autoencoder.compile(optimizer=Adam(learning_rate=7e-5), loss=tf.keras.losses.MeanSquaredError())

requirements.txt

@@ -0,0 +1,5 @@
+numpy==1.26.4
+tensorflow==2.18.0
+opencv-python==4.11.0.86
+scikit-image==0.25.2
+matplotlib==3.7.2