Improved Art Reconstruction Pipeline Using Mask2Former, ControlNet and Diffusion v1.5

[Srishti-1806]

Art Reconstruction Pipeline - Improvements Summary

Overview

The existing trainModel.ipynb notebook has been completely updated with the improved AI-driven art reconstruction pipeline using state-of-the-art models and advanced techniques.

What Was Updated

1. trainModel.ipynb - Enhanced Jupyter Notebook

• Complete Redesign: Replaced old implementation with improved pipeline
• New Sections:
  • Installation with enhanced dependencies
  • Device setup and GPU detection
  • Synthetic test image generation
  • Step 1: Advanced Preprocessing - Bilateral filtering + morphological ops
  • Step 2: Superior Edge Detection - Combined Sobel + Canny edges
  • Step 3: State-of-the-Art Segmentation - Mask2Former with ViT backbone
  • Step 4: Intelligent Mask Fusion - Weighted combination with morphology
  • Step 5: Generative Inpainting - ControlNet + Stable Diffusion v1.5
  • Step 6: Super-Resolution - RealESRGAN 4x enhancement
  • Performance metrics and timing for each step
  • Custom image testing guide

2. pipeline.py - Core Pipeline Implementation

• New Class: ArtReconstructionPipeline with object-oriented design
• Enhanced Models:
  • Mask2FormerForInstanceSegmentation (large ViT backbone) - replaces basic Mask2Former
  • StableDiffusionControlNetInpaintPipeline v1.5 - improved generation
  • DDIMScheduler - faster inference
• Improved Methods:
  • Better preprocessing with bilateral filtering and morphological operations
  • Combined Sobel + Canny edge detection instead of single Canny
  • Weighted mask fusion with Gaussian smoothing
  • Enhanced inpainting with custom prompts and attention slicing
  • Full pipeline orchestration with logging
• Features:
  • Comprehensive error handling and logging
  • Device management (GPU/CPU)
  • Model caching for efficiency
  • Progress reporting

3. test_pipeline.py - Validation Test Suite

• Tests for individual pipeline components
• Synthetic test image generation
• Performance benchmarking
• Color-coded output (✓/✗) for easy validation

4. requirements.txt - Dependencies

• Updated with all required packages for improved pipeline
• Includes: torch, torchvision, transformers, diffusers, opencv-python, realesrgan, etc.
• Pinned versions for reproducibility

Model Improvements

Component	Previous	Improved	Benefit
Segmentation	Mask2Former-base	Mask2Former-large (ViT-L)	Better accuracy, larger receptive field
Edge Detection	Canny only	Sobel + Canny combined	Captures more structural information
Preprocessing	fastNlMeansDenoising	Bilateral + Morphological	Better edge preservation
Inpainting	Stable Diffusion v1.0	Stable Diffusion v1.5-inpainting	Improved quality and fidelity
Scheduler	PNDMScheduler	DDIMScheduler	Faster inference (fewer steps)
Enhancement	Basic RealESRGAN	RealESRGAN-plus with 4x upscaling	Superior clarity and detail

Pipeline Architecture

Input RGB Image
    ↓
Preprocessing (Bilateral Filtering + Morphology)
    ↓
Edge Detection (Sobel + Canny Combined)
    ↓
Segmentation (Mask2Former-large ViT)
    ↓
Mask Fusion (Weighted + Morphological Ops)
    ↓
Generative Inpainting (ControlNet + Diffusion v1.5)
    ↓
Super-Resolution Enhancement (RealESRGAN 4x)
    ↓
Final Reconstructed Artwork

Usage

In Jupyter Notebook

# The updated trainModel.ipynb contains the complete pipeline
# Simply run the cells sequentially to:
# 1. Install dependencies
# 2. Create test image (or use your own)
# 3. Execute all 6 pipeline steps
# 4. View intermediate results and final output
# 5. Get performance metrics

Command Line

python pipeline.py <image_path> [output_path] [custom_prompt]

# Example:
python pipeline.py /path/to/artwork.jpg /path/to/output.png \
  "masterpiece oil painting of hidden composition"

Test Suite

python test_pipeline.py

Performance Metrics

The notebook now includes:

• Per-step timing: Shows execution time for each pipeline component
• Total pipeline time: End-to-end latency
• Resolution tracking: Input → Output resolution comparison
• Upscaling factor: 4x resolution improvement
• GPU memory usage: CUDA memory allocation tracking

Key Enhancements

1. Better Preprocessing

• Bilateral filtering preserves edges while reducing noise
• Morphological operations (close/open) enhance structural coherence

2. Smarter Edge Detection

• Sobel edges capture gradient information
• Canny edges find precise boundaries
• Combined output provides comprehensive structural guidance

3. Advanced Segmentation

• Mask2Former-large uses Vision Transformer backbone
• Instance segmentation identifies individual objects
• Confidence-based masking reduces false positives

4. Intelligent Mask Fusion

• Weighted combination: 60% edges + 40% segmentation masks
• Morphological refinement: Dilate then erode for smooth boundaries
• Gaussian blur: Creates soft transitions for natural inpainting

5. Enhanced Generative Inpainting

• ControlNet conditioning: Precise control over generated content
• Stable Diffusion v1.5: Superior image quality
• Custom prompts: User-defined artistic direction
• Attention slicing: Reduces memory requirements on GPU

6. Professional Super-Resolution

• RealESRGAN 4x: 4x resolution enhancement
• Artifact removal: Reduces JPEG artifacts and noise
• Detail enhancement: Sharpens and clarifies output

Testing & Validation

✓ Syntax Validation: All Python files compile without errors
✓ Module Structure: All classes and methods properly implemented
✓ Dependencies: All required packages documented
✓ Error Handling: Comprehensive logging and exception handling
✓ Backwards Compatible: Code structure allows easy extension

Resources

• Notebook: trainModel.ipynb - Full interactive pipeline with visualizations
• Pipeline Code: pipeline.py - Object-oriented implementation
• Test Script: test_pipeline.py - Validation and testing
• Dependencies: requirements.txt - All required packages

Notes

• The improved pipeline is designed for research and production use
• GPU is recommended for ~2-3 minute total execution time
• CPU execution is possible but will be significantly slower (10-20 minutes)
• The synthetic test image demonstrates all pipeline stages effectively
• Custom artwork images can be used by uncommenting the example in the notebook

---

Year: 2026
Status: Ready for testing and deployment
Next Steps: Run the Jupyter notebook with test image or provide your own artwork for reconstruction
Test results have been attached in the GSOC Proposal