File size: 3,636 Bytes
c77dd3c 83e808f c77dd3c 2187529 c77dd3c 2187529 f1cac73 080aab9 f1cac73 080aab9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
---
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.
|