This repository contains a fine-tuned ResNet18 model for classifying trash into multiple categories using the TrashNet dataset. The model was trained using PyTorch and is ready to use for further experimentation or deployment.
Model Overview
This model is based on ResNet18, a pre-trained convolutional neural network (CNN) from PyTorch. The model has been fine-tuned on the TrashNet dataset, which consists of images of trash items categorized into several classes.
Key Features
Base Model: ResNet18 (pre-trained on ImageNet).
Input Size: Images are resized to 224x224 pixels.
Normalization: Mean and standard deviation normalized to [0.5, 0.5, 0.5].
Number of Classes: 6 (based on TrashNet categories).
How to Use the Model
You can use this model with the Hugging Face Transformers library or directly with PyTorch. Below are the steps to use it.
Installation
Make sure you have the necessary libraries installed:
pip install torch torchvision transformers
Load the Model
python
1import torch
2from torchvision import transforms
3from PIL import Image
45# Load the model from Hugging Face Hub6model = torch.hub.load('huggingface/pytorch','resnet18_trash_classification')7model.eval()# Set the model to evaluation mode89# Define the transformation10transform = transforms.Compose([11 transforms.Resize((224,224)),12 transforms.ToTensor(),13 transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])14])1516# Load an image17image = Image.open('path_to_image.jpg')18input_tensor = transform(image).unsqueeze(0)# Add batch dimension1920# Perform inference21with torch.no_grad():22 outputs = model(input_tensor)23 _, predicted = torch.max(outputs,1)24print(f'Predicted class: {predicted.item()}')
Training Details
The model was fine-tuned on the TrashNet dataset using the following configuration:
Optimizer: Adam
Learning Rate: 0.001
Loss Function: CrossEntropyLoss
Batch Size: 32
Epochs: 10
Metrics
Training Loss: Decreasing steadily, as seen in the logs.
Validation Accuracy: Reached over 85% after 10 epochs.
Dataset
The model was trained on the TrashNet dataset. You can find the dataset here. The dataset contains labeled images of trash items for classification.
How to Reproduce
Requirements
Ensure the following dependencies are installed:
pip install -r requirements.txt
Training the Model
Use the included training script to train the model:
python train.py
Inference Script
Run the inference script to test the model on new images:
python inference.py --image path_to_image.jpg
Model Card
Model Name: ResNet18 Trash Classification
Framework: PyTorch
License: Apache 2.0
Dataset: TrashNet
Fine-Tuning: Yes, on TrashNet dataset
License
This model is licensed under the Apache 2.0 License. See the LICENSE file for more details.
Citation
If you use this model, please consider citing:
@article{resnet,
title={Deep Residual Learning for Image Recognition},
author={He, Kaiming and Zhang, Xiangyu and Ren, Shaoqing and Sun, Jian},
journal={arXiv preprint arXiv:1512.03385},
year={2015}
}
Acknowledgements
ResNet18: Pre-trained ResNet18 model from PyTorch.
Dataset: TrashNet dataset.
Frameworks: PyTorch, Hugging Face Transformers.
Feel free to raise an issue or submit a pull request if you encounter any problems or have suggestions for improvement.
