Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

15 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Aftershoot AI White Balance Prediction

📸 Aftershoot AI Hiring Challenge

Hybrid Deep Learning for Automatic White Balance Prediction

A computer vision and deep learning solution that predicts Temperature and Tint adjustments by combining image features with camera metadata.


📖 Overview

This project was developed as part of the Aftershoot AI Hiring Challenge, where the objective was to build an intelligent system capable of predicting white balance corrections for photographs.

Unlike traditional approaches that rely solely on image pixels, this solution combines:

  • 📷 Visual features extracted using a pretrained CNN
  • 📊 Camera metadata (EXIF information)
  • 🤖 Deep regression network for final prediction

The hybrid architecture improves prediction accuracy by leveraging both visual content and contextual camera information.


🎯 Problem Statement

Professional photo editing software such as Lightroom and Aftershoot automatically estimate white balance adjustments to improve image quality.

The challenge is to predict two continuous values:

  • Temperature
  • Tint

using

  • Input image
  • Camera metadata

This project formulates the task as a supervised regression problem and implements a hybrid deep learning architecture for accurate prediction.


🚀 Features

✅ Hybrid CNN + Metadata Architecture

✅ ResNet-18 Image Feature Extractor

✅ Metadata Encoding & Normalization

✅ Feature Fusion Network

✅ End-to-End Training Pipeline

✅ Automatic Validation

✅ Submission File Generation

✅ PyTorch Implementation


🛠 Tech Stack

Category Technologies
Language Python
Deep Learning PyTorch
Computer Vision TorchVision, Pillow
Machine Learning Scikit-learn, LightGBM
Data Processing Pandas, NumPy
Progress Tracking tqdm

📂 Project Structure

MachineLearning_aftershoot_project/

│

├── assets/
│   └── banner.png

├── dataset/
│   ├── Train/
│   │   ├── images/
│   │   └── sliders.csv
│   │
│   └── Validation/
│       ├── images/
│       └── sliders_input.csv

├── submissions/
│   └── submission_hybrid.csv

├── hybrid_model_final.pt

├── main.py

├── requirements.txt

└── README.md

🧠 Model Architecture

Image Feature Extraction

A pretrained ResNet-18 network is used to extract high-level visual representations from each image.

Image
   │
   ▼
ResNet-18
   │
512-D Feature Vector

Metadata Processing

Camera metadata is processed separately.

Features include

  • Camera Make
  • Camera Model
  • ISO
  • Aperture
  • Exposure
  • Focal Length
  • Other EXIF information

Processing pipeline

Metadata

↓

Label Encoding

↓

Standard Scaling

↓

MLP

↓

64-D Feature Vector

Feature Fusion

Both feature vectors are concatenated before being passed through fully connected layers.

Image Features (512)

        +

Metadata Features (64)

        │

        ▼

Feature Fusion

        │

Regression Head

        │

        ▼

Temperature

Tint

⚙️ Training Pipeline

The complete workflow consists of

Dataset

↓

Preprocessing

↓

Feature Extraction

↓

Metadata Encoding

↓

Hybrid Network Training

↓

Validation

↓

Inference

↓

submission_hybrid.csv

📊 Output

The trained model predicts

ID Temperature Tint
EB5BEE31... 6248 11
DE666E1F... 5996 7

Output file

submission_hybrid.csv

📈 Evaluation

The model is evaluated using regression metrics including

  • Mean Absolute Error (MAE)
  • Validation Loss
  • Separate prediction errors for
    • Temperature
    • Tint

▶️ Installation

Clone the repository

git clone https://github.com/jagadeesh-0/MachineLearning_aftershoot_project.git

cd MachineLearning_aftershoot_project

Install dependencies

pip install -r requirements.txt

▶️ Dataset

Place the dataset in the following structure

dataset/

├── Train/

│   ├── images/

│   └── sliders.csv

│

└── Validation/

    ├── images/

    └── sliders_input.csv

▶️ Run

Train the model and generate predictions

python main.py

Generated files

hybrid_model_final.pt

submission_hybrid.csv

💡 Future Improvements

Potential enhancements include

  • EfficientNet / ConvNeXt backbone
  • Vision Transformer (ViT)
  • Hyperparameter Optimization (Optuna)
  • Mixed Precision Training
  • Cross Validation
  • ONNX Export
  • Docker Support
  • Web API Deployment using FastAPI

📌 Key Learnings

Through this project I gained practical experience in

  • Hybrid Deep Learning Architectures
  • Computer Vision
  • Image Feature Extraction
  • Metadata Engineering
  • Regression Models
  • PyTorch Model Development
  • Data Preprocessing
  • Model Evaluation
  • End-to-End ML Pipelines

🤝 Contributing

Contributions, suggestions and improvements are always welcome.

Feel free to open an issue or submit a pull request.


📄 License

This project is licensed under the MIT License.


👨‍💻 Author

Jagadeesh Kumar

Machine Learning & AI Enthusiast

GitHub: https://github.com/jagadeesh-0

LinkedIn: (Add your LinkedIn URL here if available.)


⭐ If you found this project useful, consider giving it a star!

About

Computer Vision project built with PyTorch for image white balance prediction using a modular deep learning pipeline.

Topics

Resources

Stars

1 star

Watchers

0 watching

Forks

Contributors

Languages