PixelForge

KO | EN

gitlite — search

#generative-ai #proxy #pytorch #list #windows #agents #cloud-native #ai-agents #containers #redis

(★ 9)

No description available.

파일 탐색기

ZIP 다운로드

- best_unet.pth
- __init__.py
- dataset.py
- generate_synthetic_data.py
- transforms.py
- __init__.py
- charbonnier.py
- combined_loss.py
- perceptual.py
- ssim_loss.py
- __init__.py
- builder.py
- esrgan.py
- unet.py
- demo_visualization.ipynb
- demo_visualization.py
- restoration_comparison.png
.gitignore
eval.py
package-lock.json
README.md
requirements.txt
train.py
utils.py

# CDN으로 사용하기

jsDelivr

jsDelivr는 공개 GitHub 리포지토리를 별도 설정 없이 CDN으로 즉시 서빙합니다. 버전과 파일을 고르면 웹페이지에 바로 붙일 수 있는 링크와 예시 코드가 만들어집니다.

버전

파일

링크

예시

# 프로젝트 배지

PyTorch 바로가기
Python 바로가기
License 바로가기

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

**SEMICON India Hackathon 2026 — Problem Statement 1 (KLA): AI-Based Restoration of Degraded Images for Semiconductor Inspection**

An end-to-end PyTorch deep learning pipeline designed to solve complex image degradation problems by learning the inverse mapping $f^{-1}: y \to x$.

---

## 📌 Problem Overview & Formulations

The task requires learning the inverse mapping $f^{-1}: y \to x$, where:
- **Ground Truth ($x$)**: $256 \times 256$ clean image, normalized to range $[0, 1]$.
- **Corrupt Image ($y$)**: $128 \times 128$ image degraded through a combination of:
  1. Multiplicative Speckle Noise ($y_{\text{speckle}} = x + x \cdot \eta_1$)
  2. $2\times$ Spatial Downsampling ($256 \times 256 \to 128 \times 128$)
  3. Additive Gaussian Noise ($y = y_{\text{speckle, down}} + \eta_2$)

> [!IMPORTANT]
> **Out-of-Bounds Pixel Range ($y \notin [0, 1]$)**
> Because of speckle noise and additive Gaussian noise, pixel values in $y$ can drop below $0.0$ or exceed $1.0$. Clipping $y$ directly to $[0, 1]$ during preprocessing would destroy crucial phase/noise distribution information. Our dataset loader preserves raw float values or applies z-score standardization.

---

## 🏗️ Architecture Design Rationale

PixelForge provides two interchangeable model architectures:

### 1. UNet-SR (Primary Model — used for final results below)
- **Residual Encoder-Decoder**: Multi-scale feature extraction with skip connections.
- **Channel Attention (SE-style)**: Squeeze-and-Excitation attention blocks focus capacity on noise-free feature channels.
- **Sub-Pixel Convolution (PixelShuffle 2x Head)**: Restores spatial dimensions from $128 \times 128 \to 256 \times 256$ cleanly, eliminating checkerboard artifacts inherent to standard transposed convolutions.
- **Global Residual Skip**: The network predicts the residual correction $\Delta = f^{-1}(y) - \text{Bilinear}(y)$, stabilizing gradient flow.

### 2. Lightweight ESRGAN Generator (Alternative Model)
- **Residual-in-Residual Dense Blocks (RRDB)**: Dense feature reuse across deep layers allows high-frequency edge restoration.
- **Sub-Pixel Upsampler**: $2\times$ upsampling module via sub-pixel convolution for comparative benchmark evaluation against UNet-SR.

---

## 🎯 Loss Function Rationale

We employ a composite loss function balancing pixel-level accuracy, structural similarity, and perceptual quality:

$$\mathcal{L}_{\text{total}} = \lambda_{\text{pixel}} \mathcal{L}_{\text{Charbonnier}} + \lambda_{\text{ssim}} \mathcal{L}_{\text{SSIM}} + \lambda_{\text{perceptual}} \mathcal{L}_{\text{Perceptual}}$$

1. **Charbonnier Loss ($\mathcal{L}_{\text{Charbonnier}}$)**:
   $$\mathcal{L}_{\text{Charbonnier}} = \sqrt{||\hat{x} - x||^2 + \epsilon^2} \quad (\epsilon = 10^{-3})$$
   *Rationale*: A smooth differentiable variant of $L_1$ loss. Unlike $L_2$ (which produces blurry predictions) or standard $L_1$ (non-smooth at 0), Charbonnier handles sparse noise outliers gracefully while preventing gradient saturation.
2. **Structural Similarity Loss ($\mathcal{L}_{\text{SSIM}}$)**:
   $$\mathcal{L}_{\text{SSIM}} = 1 - \text{SSIM}(\hat{x}, x)$$
   *Rationale*: Enforces structural coherence, local luminance contrast, and edge sharpness.
3. **VGG-19 Perceptual Loss ($\mathcal{L}_{\text{Perceptual}}$)**:
   *Rationale*: Evaluates Euclidean distance in VGG feature space (`relu2_2`, `relu3_3`), preserving high-frequency textures. Implemented but not enabled for the final run below (see Known Limitations).

---

## 🧪 Training Hygiene & Experiments Summary

| Experiment | Hypothesis / Modification | Result / Findings | Status |
| :--- | :--- | :--- | :--- |
| **Exp 1: Hard Clipping [0,1] on $y$** | Clamp noisy LR input $y$ to $[0,1]$ before feeding to model. | Truncated negative Gaussian tails and speckle peaks, causing loss of variance detail in restoration. | ❌ Rejected |
| **Exp 2: Transposed Convolution for 2x SR** | Use `ConvTranspose2d` for upscaling $128 \to 256$. | Produced visible checkerboard grid patterns in high-frequency regions. | ❌ Replaced |
| **Exp 3: Sub-Pixel Convolution (PixelShuffle) + ResSkip** | Use `PixelShuffle(scale_factor=2)` + residual base addition. | Completely eliminated checkerboard artifacts; improved Val PSNR by +2.1 dB. | ✅ Adopted |
| **Exp 4: Charbonnier + SSIM Composite Loss** | Combine smooth $L_1$ with differentiable SSIM ($1 - \text{SSIM}$). | Significantly sharper structural edges and cleaner noise removal. | ✅ Adopted |

### Data Verification (before training)

Before any training run, the full pipeline was verified end-to-end on the real
KLA-provided dataset:
- Confirmed 3200 real GT/NoisyLR pairs, correctly matched by filename (no
  mismatches after removing macOS resource-fork artifacts from the zip).
- Confirmed tensor shapes: NoisyLR `(1, 128, 128)`, GT `(1, 256, 256)`.
- Confirmed NoisyLR values genuinely exceed `[0, 1]` (observed range in
  samples: as low as -0.30, as high as 1.61) and are preserved, not clipped.
- Ran a 1-epoch CPU smoke test first to confirm the full data → model → loss
  → checkpoint pipeline worked before committing to a longer GPU run.

---

## 📊 Final Results (UNet-SR, 40 epochs, real KLA dataset)

Trained on Google Colab (free T4 GPU), 2880 train / 320 validation split
(90/10, fixed seed), batch size 16, learning rate 1e-4.

| Metric | Epoch 1 (CPU smoke test) | Epoch 40 (final, GPU) |
|---|---|---|
| Validation PSNR | 21.67 dB | **28.36 dB** |
| Validation SSIM | 0.6385 | **0.7787** |

Validation PSNR climbed steadily across training and plateaued around
epoch 33–40 (28.32 → 28.36 dB), indicating the model reached a stable
optimum for this architecture/loss configuration rather than stopping
early or training unstably.

See [`results/restoration_comparison.png`](results/restoration_comparison.png)
for side-by-side comparisons (Noisy Input | Restored | Ground Truth) across
several validation samples.

### Known Limitations

- The model restores structured, high-contrast content (sharp edges,
  textured surfaces, distinct shapes) very well, but is noticeably weaker
  on smooth, low-contrast, heavily noise-dominated regions, where output
  is visibly blurrier than ground truth. This is a common failure mode for
  pixel/SSIM-loss-trained restoration networks, which favor "safe" smooth
  predictions in low-signal areas — enabling the VGG perceptual loss
  (already implemented, see `losses/perceptual.py`) is the planned next
  experiment to address this.
- Results above are on our internal validation split; we have not yet run
  inference on KLA's separate `Test_NoisyLR` set (which has no ground
  truth), so generalization beyond our validation distribution is not yet
  independently quantified.

---

## 📁 Repository Structure

```
PixelForge/
├── data/
│   ├── dataset.py                  # Paired GT & NoisyLR PyTorch Dataset
│   ├── transforms.py               # Paired spatial augmentations (flips, rotations)
│   └── generate_synthetic_data.py  # Synthetic benchmark dataset generator script
├── models/
│   ├── unet.py                     # UNet-SR architecture with PixelShuffle 2x upsampler
│   ├── esrgan.py                   # Lightweight ESRGAN-RRDB generator
│   └── builder.py                  # Model factory helper
├── losses/
│   ├── charbonnier.py              # Charbonnier smooth pixel loss
│   ├── ssim_loss.py                # Differentiable SSIM loss
│   ├── perceptual.py              # VGG-19 Perceptual loss
│   └── combined_loss.py            # Composite weighted loss manager
├── notebooks/
│   ├── demo_visualization.py       # Standalone visualization script
│   └── demo_visualization.ipynb    # Interactive visual & histogram analysis notebook
├── checkpoints/
│   └── best_unet.pth               # Trained model weights (Val PSNR 28.36 dB)
├── results/
│   └── restoration_comparison.png  # Visual proof: Noisy | Restored | Ground Truth
├── utils.py                        # Metric routines (PSNR, SSIM) & plotting helpers
├── train.py                        # Training script with logging and checkpointing
├── eval.py                         # Evaluation script (metrics, comparison grids, histograms)
├── requirements.txt                # Python package dependencies
└── README.md                       # Complete technical documentation
```

---

## 🚀 Quick Start & Reproduction Guide

### 1. Installation
Clone the repository and install dependencies:
```bash
git clone https://github.com/sabeeshvar/PixelForge.git
cd PixelForge
pip install -r requirements.txt
```

### 2. Generate Synthetic Benchmark Dataset (Optional)
If you don't have custom data ready, generate a synthetic paired dataset for testing:
```bash
python data/generate_synthetic_data.py --output_dir data/sample_dataset --num_train 32 --num_val 8
```

### 3. Train Model
Train the primary UNet-SR architecture (as used for the final results above):
```bash
python train.py --gt_dir data/train/GT --noisy_dir data/train/NoisyLR \
  --model_name unet --channels 1 --epochs 40 --batch_size 16 --lr 1e-4 \
  --val_split 0.1 --checkpoint_dir checkpoints --log_dir runs
```

To train the alternative Lightweight ESRGAN architecture for comparison:
```bash
python train.py --model_name esrgan --epochs 20 --batch_size 4 --lr 2e-4
```

### 4. Evaluate & Generate Visual Artifacts
Run evaluation on held-out validation images to compute PSNR/SSIM metrics, generate side-by-side comparison grids (`Ground Truth` | `Corrupt` | `Restored`), and pixel intensity histograms:
```bash
python eval.py --model_name unet --checkpoint_path checkpoints/best_unet.pth --output_dir eval_results
```

### 5. Interactive Notebook
Launch Jupyter Notebook to explore image outputs and intensity distributions:
```bash
jupyter notebook notebooks/demo_visualization.ipynb
```

PixelForge

파일 탐색기

# CDN으로 사용하기

# 프로젝트 배지

Was this content helpful?

Recent Feedback

Download README