Spaces:
Sleeping
Sleeping
Rajeev-86
commited on
Commit
·
370e492
1
Parent(s):
cb52939
new models added
Browse files- .github/workflows/sync_to_hf.yml +0 -1
- Dockerfile +3 -3
- README.md +129 -2
- image_handler.py → handler.py +5 -3
- showcasing-few-results.py +90 -0
.github/workflows/sync_to_hf.yml
CHANGED
|
@@ -6,7 +6,6 @@ on:
|
|
| 6 |
paths:
|
| 7 |
- 'Dockerfile'
|
| 8 |
- 'requirements.txt'
|
| 9 |
-
- 'README.md'
|
| 10 |
- '.github/workflows/sync_to_hf.yml'
|
| 11 |
workflow_dispatch: # Allows manual trigger from the Actions tab
|
| 12 |
|
|
|
|
| 6 |
paths:
|
| 7 |
- 'Dockerfile'
|
| 8 |
- 'requirements.txt'
|
|
|
|
| 9 |
- '.github/workflows/sync_to_hf.yml'
|
| 10 |
workflow_dispatch: # Allows manual trigger from the Actions tab
|
| 11 |
|
Dockerfile
CHANGED
|
@@ -14,9 +14,9 @@ EXPOSE 8080
|
|
| 14 |
|
| 15 |
# 3. Installing gdown, downloading models and immediately uninstalling it for a smaller image
|
| 16 |
RUN pip install --no-cache-dir gdown && \
|
| 17 |
-
gdown
|
| 18 |
-
gdown
|
| 19 |
-
gdown
|
| 20 |
pip uninstall gdown -y && \
|
| 21 |
rm -rf /root/.cache/pip && \
|
| 22 |
rm -rf /var/lib/apt/lists/*
|
|
|
|
| 14 |
|
| 15 |
# 3. Installing gdown, downloading models and immediately uninstalling it for a smaller image
|
| 16 |
RUN pip install --no-cache-dir gdown && \
|
| 17 |
+
gdown h1vzykr8eH77ypl4R7_AsSb56W9p7LDlKo -O model_store/UNET.mar && \
|
| 18 |
+
gdown 1GdQyPrRn8yH5y7F1jC5Yuu2vWNLUs4K2 -O model_store/R-UNET.mar && \
|
| 19 |
+
gdown 1mw7VbQt9Lrv0TG2LX0Ke3hKp0kyukElA -O model_store/A-R-UNET.mar && \
|
| 20 |
pip uninstall gdown -y && \
|
| 21 |
rm -rf /root/.cache/pip && \
|
| 22 |
rm -rf /var/lib/apt/lists/*
|
README.md
CHANGED
|
@@ -9,6 +9,133 @@ app_file: Dockerfile
|
|
| 9 |
pinned: false
|
| 10 |
---
|
| 11 |
|
| 12 |
-
#
|
|
|
|
| 13 |
|
| 14 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 9 |
pinned: false
|
| 10 |
---
|
| 11 |
|
| 12 |
+
#Image Denoising using UNET and it's variants
|
| 13 |
+
_A Deep Learning approach to remove noise from images._
|
| 14 |
|
| 15 |
+
##1. Project Problem Statement: The Critical Need for Image Denoising
|
| 16 |
+
Digital images are indispensable data sources across numerous high-stakes industries, yet they are universally susceptible to noise corruption introduced during acquisition, transmission, or processing. This noise, whether Gaussian, Poisson, or Speckle, degrades images in two critical ways:
|
| 17 |
+
|
| 18 |
+
Impairing Human Perception: Noise obscures subtle features and textures, significantly lowering the visual fidelity required for accurate human interpretation.
|
| 19 |
+
|
| 20 |
+
Compromising Machine Reliability: Noise introduces spurious data points that confuse downstream Computer Vision tasks, drastically reducing the accuracy of algorithms used for analysis and automation.
|
| 21 |
+
|
| 22 |
+
The challenge is magnified across essential fields:
|
| 23 |
+
|
| 24 |
+
In Medical Imaging (e.g., MRI, CT), noise threatens the ability to identify critical, life-saving diagnostic features.
|
| 25 |
+
|
| 26 |
+
In Industrial Quality Control, noise leads to costly false positives or false negatives during automated inspection.
|
| 27 |
+
|
| 28 |
+
In Remote Sensing and Astronomy, noise prevents the reliable extraction of scientific data from satellite and telescopic imagery.
|
| 29 |
+
|
| 30 |
+
The objective of this project is to develop and evaluate a robust image denoising solution capable of effectively suppressing varied noise types while preserving crucial structural details, thereby elevating the reliability and precision of visual data for both human experts and advanced Machine Learning systems.
|
| 31 |
+
|
| 32 |
+
##2. Few Results
|
| 33 |
+
|
| 34 |
+

|
| 35 |
+
|
| 36 |
+
##3. Project Structure
|
| 37 |
+
|
| 38 |
+
├── .gitignore
|
| 39 |
+
├── .github/workflows
|
| 40 |
+
│ ├── sync_to_hf.yml
|
| 41 |
+
├── Dockerfile
|
| 42 |
+
├── api-test.py
|
| 43 |
+
├── handler.py
|
| 44 |
+
├── requirements.txt
|
| 45 |
+
├── images
|
| 46 |
+
│
|
| 47 |
+
├── README.md
|
| 48 |
+
└── (maybe some sample outputs/denoised images)
|
| 49 |
+
|
| 50 |
+
[Dataset Preparation](https://drive.google.com/file/d/1hY0OBv0TI8dsP5Y2Le6IT9kFwPM_t8_V/view?usp=sharing)
|
| 51 |
+
[UNET_training](https://www.kaggle.com/code/rajeev86/training-unet-for-image-denoising)
|
| 52 |
+
[Residual-UNET_training](https://www.kaggle.com/code/rajeev86/training-residual-unet-for-image-denoising)
|
| 53 |
+
[CBAM-Residual-UNET_training](https://www.kaggle.com/code/rajeev86/training-unet-with-residuals-and-cbam-layers)
|
| 54 |
+
[TorchScript_comparison](https://drive.google.com/file/d/1JC6WIi59fppT78v5kl26VSD4tX73ikgg/view?usp=sharing)
|
| 55 |
+
[Model Archiving](https://drive.google.com/file/d/1X4lMJYiC8ps3170X-Jj5-YvnaIDDvnbx/view?usp=sharing)
|
| 56 |
+
|
| 57 |
+
##4. Dataset Used
|
| 58 |
+
|
| 59 |
+
The model was trained on an augmented dataset of 32,000 clean/noisy patch pairs derived from the BSD500 dataset, utilizing a 128×128 patch size with dynamic D4 geometric augmentation. To ensure robustness against real-world degradation, we employed a hybrid noise model incorporating four components:
|
| 60 |
+
|
| 61 |
+
Mixed Sensor Noise: A combination of Additive White Gaussian Noise (σstd ∈[0,30]) and Signal-Dependent Poisson Noise (a∈[0,0.05]).
|
| 62 |
+
|
| 63 |
+
Impulse Noise: Sparse Salt-and-Pepper noise (∈[0.001,0.005]).
|
| 64 |
+
|
| 65 |
+
Structured Artifacts: JPEG compression with randomized quality (∈[70,95]).
|
| 66 |
+
|
| 67 |
+
Due to the complex, non-linear nature of this hybrid noise model, we quantified the overall degradation using the Effective Noise Level (σ eff), defined as the Standard Deviation of the entire noise residual (y−x) across the validation set. The measured effective noise level for the challenging dataset was σeff =79.32 (scaled to 0-255). All performance metrics (PSNR, SSIM) presented below are reported against this highly degraded baseline.
|
| 68 |
+
|
| 69 |
+
[Original Berkeley Segmentation Dataset 500 (BSDS500)](https://data.vision.ee.ethz.ch/cvl/DIV2K/)
|
| 70 |
+
[GDrive link for our modified Dataset](https://drive.google.com/drive/folders/1AObLCZGTHvtcv-lZFGPBA8k8xgC1k4_w?usp=sharing)
|
| 71 |
+
|
| 72 |
+
##5. Model Architectures
|
| 73 |
+
|
| 74 |
+
| Model | Description | Key Features |
|
| 75 |
+
| --------------------- | ---------------------------------------------------- | --------------------------------------- |
|
| 76 |
+
| U-Net | Baseline architecture for image-to-image restoration | Encoder-decoder skip connections |
|
| 77 |
+
| Residual U-Net | Adds residual blocks to improve feature flow | Residual connections within U-net blocks|
|
| 78 |
+
| Residual U-Net + CBAM | Incorporates Convolutional Block Attention Module | focuses noise removal on key locations |
|
| 79 |
+
|
| 80 |
+
##6. Training Setup
|
| 81 |
+
|
| 82 |
+
| Platform | Purpose | Notes |
|
| 83 |
+
| ------------ | ------------------------------ | ------------------------------ |
|
| 84 |
+
| Google Colab | Dataset prep + initial testing | Limited GPU runtime |
|
| 85 |
+
| Kaggle | Model training | Used for high-performance GPUs |
|
| 86 |
+
| Google Drive | Model & dataset storage | For cross-platform access |
|
| 87 |
+
|
| 88 |
+
##7. Optimization
|
| 89 |
+
Comparing ordinary serlialization vs TorchScript inference time
|
| 90 |
+
|
| 91 |
+
| Model | Avg. Inference Time (ms) |
|
| 92 |
+
| -------------------------| ------------------------ |
|
| 93 |
+
| U-Net | 39.18 % |
|
| 94 |
+
| Residual U-Net | 43.77 % |
|
| 95 |
+
| Attention Residual U-Net | 30.72 % |
|
| 96 |
+
|
| 97 |
+
##8. Deployment (Backend)
|
| 98 |
+
|
| 99 |
+
Backend Framework: TorchServe
|
| 100 |
+
|
| 101 |
+
Containerization: Docker
|
| 102 |
+
|
| 103 |
+
Deployment Platform: Hugging Face Spaces
|
| 104 |
+
|
| 105 |
+
Artifacts: .mar model files stored in [Gdrive link](https://drive.google.com/drive/folders/1Arnlrjdxqd0zBaIC4ECigDxxSrgqyAHX?usp=sharing)
|
| 106 |
+
|
| 107 |
+
[HuggingFace Space link](https://huggingface.co/spaces/Rexy-3d/Denoiser-Server)
|
| 108 |
+
|
| 109 |
+
##8. Frontend (Next.js)
|
| 110 |
+
|
| 111 |
+
Repo: [Frontend Repo Link](https://github.com/Rajeev-86/Denoiser_-Frontend-)
|
| 112 |
+
|
| 113 |
+
Platform: Vercel
|
| 114 |
+
|
| 115 |
+
Provides a simple web interface for uploading noisy images and visualizing denoised outputs.
|
| 116 |
+
|
| 117 |
+
[Vercel webpage Link](https://denoiserbyrajeev.vercel.app/)
|
| 118 |
+
|
| 119 |
+
##9. Results:
|
| 120 |
+
|
| 121 |
+
| Model | PSNR | SSIM | Notes |
|
| 122 |
+
| --------------------- | ------------ | ------------| ----------------------- |
|
| 123 |
+
| U-Net | 28.7583 | 0.8444 | Baseline |
|
| 124 |
+
| Residual U-Net | 28.7630 | 0.8415 | Better texture recovery |
|
| 125 |
+
| Residual U-Net + CBAM | **29.0086** | **0.8485** | Best performance |
|
| 126 |
+
|
| 127 |
+
##10. References
|
| 128 |
+
|
| 129 |
+
[U-Net: Convolutional Networks for Biomedical Image Segmentation](https://arxiv.org/abs/1505.04597)
|
| 130 |
+
[Recurrent Residual Convolutional Neural Network based on U-Net (R2U-Net) for Medical Image Segmentation](https://arxiv.org/pdf/1802.06955)
|
| 131 |
+
[Layer Normalization](https://arxiv.org/abs/1607.06450)
|
| 132 |
+
[CBAM: Convolutional Block Attention Module](https://arxiv.org/abs/1807.06521)
|
| 133 |
+
[Attention-based UNet enabled Lightweight Image Semantic Communication System over Internet of Things](https://arxiv.org/html/2401.07329v1)
|
| 134 |
+
[Application of ResUNet-CBAM in Thin-Section Image Segmentation of Rocks](https://www.mdpi.com/2078-2489/15/12/788)
|
| 135 |
+
|
| 136 |
+
##11. Author
|
| 137 |
+
|
| 138 |
+
Rajeev Ahirwar
|
| 139 |
+
|
| 140 |
+
[Linkedin](https://www.linkedin.com/in/86thrajeev/)
|
| 141 |
+
[GitHub](https://github.com/Rajeev-86)
|
image_handler.py → handler.py
RENAMED
|
@@ -6,6 +6,8 @@ import io
|
|
| 6 |
import time
|
| 7 |
import logging
|
| 8 |
import torch.nn.functional as F
|
|
|
|
|
|
|
| 9 |
|
| 10 |
logger = logging.getLogger(__name__)
|
| 11 |
|
|
@@ -41,12 +43,12 @@ class ImageHandler(BaseHandler):
|
|
| 41 |
output_tensor_resized = output_tensor
|
| 42 |
if output_tensor.shape != input_tensor.shape:
|
| 43 |
output_tensor_resized = F.interpolate(
|
| 44 |
-
output_tensor.unsqueeze(0),
|
| 45 |
size=input_tensor.shape[-2:],
|
| 46 |
-
mode='bilinear',
|
| 47 |
align_corners=False
|
| 48 |
).squeeze(0)
|
| 49 |
-
|
| 50 |
pixel_difference = torch.mean(torch.abs(input_tensor - output_tensor_resized)).item()
|
| 51 |
logger.info(f"OUTPUT_QUALITY: denoising_intensity={pixel_difference:.4f}")
|
| 52 |
|
|
|
|
| 6 |
import time
|
| 7 |
import logging
|
| 8 |
import torch.nn.functional as F
|
| 9 |
+
import numpy as np
|
| 10 |
+
from skimage.metrics import peak_signal_noise_ratio, structural_similarity
|
| 11 |
|
| 12 |
logger = logging.getLogger(__name__)
|
| 13 |
|
|
|
|
| 43 |
output_tensor_resized = output_tensor
|
| 44 |
if output_tensor.shape != input_tensor.shape:
|
| 45 |
output_tensor_resized = F.interpolate(
|
| 46 |
+
output_tensor.unsqueeze(0),
|
| 47 |
size=input_tensor.shape[-2:],
|
| 48 |
+
mode='bilinear',
|
| 49 |
align_corners=False
|
| 50 |
).squeeze(0)
|
| 51 |
+
|
| 52 |
pixel_difference = torch.mean(torch.abs(input_tensor - output_tensor_resized)).item()
|
| 53 |
logger.info(f"OUTPUT_QUALITY: denoising_intensity={pixel_difference:.4f}")
|
| 54 |
|
showcasing-few-results.py
ADDED
|
@@ -0,0 +1,90 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
import torchvision.transforms as transforms
|
| 3 |
+
from PIL import Image
|
| 4 |
+
import matplotlib.pyplot as plt
|
| 5 |
+
|
| 6 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 7 |
+
|
| 8 |
+
'''This script relies on predownloaded models, if you do'nt have them then run the following commands to download them-
|
| 9 |
+
|
| 10 |
+
!kaggle kernels output rajeev86/training-unet-for-image-denoising
|
| 11 |
+
!kaggle kernels output rajeev86/training-residual-unet-for-image-denoising
|
| 12 |
+
!kaggle kernels output rajeev86/training-unet-with-residuals-and-cbam-layers
|
| 13 |
+
|
| 14 |
+
Note that you may require kaggle credentials for successfully downloading the models
|
| 15 |
+
'''
|
| 16 |
+
|
| 17 |
+
model1_path = 'models/Script_Unet.pt'
|
| 18 |
+
model2_path = 'models/Script_Res-Unet.pt'
|
| 19 |
+
model3_path = 'models/Script_Att-Unet.pt'
|
| 20 |
+
|
| 21 |
+
try:
|
| 22 |
+
model1 = torch.jit.load(model1_path, map_location=device)
|
| 23 |
+
model1.eval()
|
| 24 |
+
model2 = torch.jit.load(model2_path, map_location=device)
|
| 25 |
+
model2.eval()
|
| 26 |
+
model3 = torch.jit.load(model3_path, map_location=device)
|
| 27 |
+
model3.eval()
|
| 28 |
+
except Exception as e:
|
| 29 |
+
print(f"Error loading model: {e}")
|
| 30 |
+
model1, model2, model3 = None, None, None
|
| 31 |
+
|
| 32 |
+
def denoise_image_and_show(image_paths):
|
| 33 |
+
for image_path in image_paths:
|
| 34 |
+
noisy_image = Image.open(image_path).convert('RGB')
|
| 35 |
+
|
| 36 |
+
transform = transforms.Compose([transforms.ToTensor()])
|
| 37 |
+
noisy_tensor = transform(noisy_image).unsqueeze(0).to(device)
|
| 38 |
+
|
| 39 |
+
denoised_tensor1 = None
|
| 40 |
+
denoised_tensor2 = None
|
| 41 |
+
denoised_tensor3 = None
|
| 42 |
+
|
| 43 |
+
with torch.no_grad():
|
| 44 |
+
if model1:
|
| 45 |
+
denoised_tensor1 = model1(noisy_tensor)
|
| 46 |
+
if model2:
|
| 47 |
+
denoised_tensor2 = model2(noisy_tensor)
|
| 48 |
+
if model3:
|
| 49 |
+
denoised_tensor3 = model3(noisy_tensor)
|
| 50 |
+
|
| 51 |
+
images_to_show = [noisy_image]
|
| 52 |
+
titles = ['Noisy Image']
|
| 53 |
+
|
| 54 |
+
if denoised_tensor1 is not None:
|
| 55 |
+
denoised1_image = transforms.ToPILImage()(denoised_tensor1.squeeze(0).cpu())
|
| 56 |
+
images_to_show.append(denoised1_image)
|
| 57 |
+
titles.append('Unet model')
|
| 58 |
+
|
| 59 |
+
if denoised_tensor2 is not None:
|
| 60 |
+
denoised2_image = transforms.ToPILImage()(denoised_tensor2.squeeze(0).cpu())
|
| 61 |
+
images_to_show.append(denoised2_image)
|
| 62 |
+
titles.append('res Unet model')
|
| 63 |
+
|
| 64 |
+
if denoised_tensor3 is not None:
|
| 65 |
+
denoised3_image = transforms.ToPILImage()(denoised_tensor3.squeeze(0).cpu())
|
| 66 |
+
images_to_show.append(denoised3_image)
|
| 67 |
+
titles.append('Att model')
|
| 68 |
+
|
| 69 |
+
fig, axes = plt.subplots(1, len(images_to_show), figsize=(5 * len(images_to_show), 5))
|
| 70 |
+
|
| 71 |
+
if len(images_to_show) == 1:
|
| 72 |
+
axes.imshow(images_to_show[0])
|
| 73 |
+
axes.set_title(titles[0])
|
| 74 |
+
axes.axis('off')
|
| 75 |
+
else:
|
| 76 |
+
for i, img in enumerate(images_to_show):
|
| 77 |
+
axes[i].imshow(img)
|
| 78 |
+
axes[i].set_title(titles[i])
|
| 79 |
+
axes[i].axis('off')
|
| 80 |
+
|
| 81 |
+
plt.tight_layout()
|
| 82 |
+
plt.show()
|
| 83 |
+
|
| 84 |
+
image_list = [
|
| 85 |
+
'images/145079.jpg',
|
| 86 |
+
'images/258089.jpg',
|
| 87 |
+
'images/29030.jpg',
|
| 88 |
+
'images/228076.jpg'
|
| 89 |
+
]
|
| 90 |
+
denoise_image_and_show(image_list)
|