metadata

license: mit
datasets:
  - zh-plus/tiny-imagenet
pipeline_tag: image-to-image
tags:
  - image
  - autoencoder
  - Bueorm

Autoencoder Image Denoising

Model Description

AID is a convolutional autoencoder model designed to remove noise from images. The model has been trained on a custom dataset with various levels of Gaussian noise, as well as grayscale variations, making it versatile for image denoising tasks. The architecture leverages several convolutional and deconvolutional layers to encode and decode the input images, aiming to reconstruct the original image from noisy data.

This model can handle a variety of noise types, including Gaussian noise and salt-and-pepper noise, and is robust in restoring details in color and grayscale images.

Model Architecture

The model follows a typical autoencoder structure with an encoder and a decoder:

Encoder: Multiple convolutional layers with ReLU activations and down-sampling through max-pooling.
Bottleneck: A latent space representation that compresses the input image.
Decoder: Deconvolutional layers with ReLU activations, using up-sampling to reconstruct the original image.
Output Layer: A final convolutional layer with sigmoid activation to ensure the output is scaled between 0 and 1.

The model is trained using the Mean Squared Error (MSE) loss function, optimized with the Adam optimizer, and has been trained on a GPU for faster convergence.

Hyperparameters

Optimizer: Adam (learning_rate=0.001)
Loss Function: Mean Squared Error (MSE)
Batch Size: 16
Input Image Size: 256x256x3 (RGB format)
Noise Types: Gaussian noise (var=0.1), Salt-and-Pepper noise, and grayscale conversion.

How to Use the Model

You can load and use the model for image denoising directly in Hugging Face using the following code:

import tensorflow as tf
from huggingface_hub import hf_hub_download
from skimage.util import random_noise
import numpy as np
import matplotlib.pyplot as plt

model_path = hf_hub_download(repo_id="BueormLLC/AID", filename="model.h5")

autoencoder = tf.keras.models.load_model(model_path, custom_objects={'mse': tf.keras.losses.MeanSquaredError()})

def add_noise(img, noise_type="gaussian"):
    img_np = img.numpy()
    if noise_type == "gaussian":
        noisy_img = random_noise(img_np, mode='gaussian', var=0.1)
    elif noise_type == "salt_pepper":
        noisy_img = random_noise(img_np, mode='s&p', amount=0.1)
    return np.clip(noisy_img, 0., 1.)

def predict_denoised_image(autoencoder, image):
    img_resized = tf.image.resize(image, (256, 256)) / 255.0
    img_array = np.expand_dims(img_resized, axis=0)

    noisy_image = add_noise(img_resized)

    denoised_image = autoencoder.predict(np.expand_dims(noisy_image, axis=0))

    fig, ax = plt.subplots(1, 2, figsize=(10, 5))
    ax[0].imshow(noisy_image)
    ax[0].set_title("Noisy Image")
    ax[0].axis('off')

    ax[1].imshow(denoised_image[0])
    ax[1].set_title("Denoised Image")
    ax[1].axis('off')

    plt.show()

test_image = tf.keras.preprocessing.image.load_img('image.jpg', target_size=(256, 256))
test_image = tf.keras.preprocessing.image.img_to_array(test_image)

predict_denoised_image(autoencoder, test_image)

Performance

The model has been evaluated on a variety of noisy images and performs well in denoising tasks. While specific benchmarks were not run, it shows a significant reduction in noise for both color and grayscale images across a variety of noise levels.

Test Image

Donations

Subscription

Bueorm AI

Acknowledgements

We would like to thank the creators of TensorFlow and the Hugging Face team for providing the infrastructure and tools to develop and host this model. Special thanks to the community for the support in improving and fine-tuning this model.

Thanks

Feel free to use and test this model for your image denoising tasks, and don't hesitate to reach out with any feedback or suggestions.