---
tags:
- image-classification
- pytorch
- vision
- vit
- transformers
license: apache-2.0
library_name: transformers
pipeline_tag: image-classification
model-index:
- name: Image Quality Regression Model
  results: []
---

# Image Quality Regression Model

This model is trained on the dataset [yigagilbert/image_quality_dataset](https://huggingface.co/datasets/yigagilbert/image_quality_dataset) and performs regression tasks to predict image quality scores. 

## Model Details

- **Dataset**: yigagilbert/image_quality_dataset
- **Target Column**: quality_score
- **Test Split**: 20% test data
- **Training Epochs**: 3
- **Learning Rate**: 5e-5
- **Max Value in Dataset**: 54.02

This model fine-tunes the [google/vit-base-patch16-224](https://huggingface.co/google/vit-base-patch16-224) Vision Transformer using PyTorch and Hugging Face's 🤗 Transformers library. It predicts a numerical score based on the quality of the input image.


# Image Regression Model

This repository contains a model for **image regression** tasks, where the goal is to predict a numerical value from an input image. The model fine-tunes the [google/vit-base-patch16-224](https://huggingface.co/google/vit-base-patch16-224) Vision Transformer using PyTorch and 🤗 Hugging Face tools. You can train the model, upload it to the 🤗 Model Hub, and perform inference using a simple API.

## Installation
Install the required packages by running:
```bash
pip install -r requirements.txt
```

## Usage

### Import Functions
```python
from ImageRegression import train_model, upload_model, predict
```

### Train Model
Train the model using the `train_model()` function. Below are the key parameters:

- **dataset_id**: Hugging Face dataset identifier or path to your local dataset.
- **value_column_name**: Column in the dataset containing the target regression values.
- **test_split**: Proportion of data to use for testing (e.g., `0.2` for 20% test data).
- **output_dir**: Directory where model checkpoints will be saved.
- **num_train_epochs**: Number of training epochs.
- **learning_rate**: Learning rate for the optimizer.

```python
train_model(dataset_id='yigagilbert/image_quality_dataset',
            value_column_name='quality_score',
            test_split=0.2,
            output_dir='./model_output',
            num_train_epochs=10,
            learning_rate=1e-4)
```

Training progress will be logged, and checkpoints will be saved in `output_dir`. These checkpoints can be used for model inference and uploaded to the 🤗 Hub.

### Upload Model to Hugging Face Hub
To upload your trained model to the 🤗 Hub, use the `upload_model()` function:

- **model_id**: The name of the model repository on the 🤗 Hub.
- **token**: Authentication token (create one [here](https://huggingface.co/settings/tokens)).
- **checkpoint_dir**: Directory where the trained model checkpoints are located.

```python
upload_model(model_id='yigagilbert/image-qaulity-model',
             token='your_HF_token',
             checkpoint_dir='./model_output/checkpoint-940')
```
Once uploaded, the model can be used for inference directly from the Hub.

### Model Inference (Prediction)
You can perform inference using the `predict()` function. 

- **repo_id**: The repository identifier of the uploaded model.
- **image_path**: Path to the image file you want to run predictions on.

```python
predict(repo_id='yigagilbert/image-qaulity-model',
        image_path='path_to_image.jpg')
```
The first time you run inference, the model will be downloaded from the Hugging Face Hub. Subsequent inferences will run faster as the model is cached locally.