| | --- |
| | license: apache-2.0 |
| | datasets: |
| | - HuggingFaceFW/fineweb-edu |
| | language: |
| | - en |
| | metrics: |
| | - accuracy |
| | library_name: flair |
| | pipeline_tag: image-classification |
| | tags: |
| | - art |
| | - code |
| | --- |
| | |
| | # Model Card for Custom CNN Model for Garbage Classification |
| |
|
| | This model card provides information about a custom Convolutional Neural Network (CNN) designed for classifying images of garbage items into predefined categories. |
| |
|
| | ## Model Details |
| |
|
| | ### Model Description |
| |
|
| | The CNN architecture (CNNModel) consists of: |
| |
|
| | - Four convolutional layers with batch normalization, ReLU activation, max pooling, and dropout for feature extraction. |
| | - Two fully connected layers for classification. |
| |
|
| | The model is trained using the Adam optimizer with cross-entropy loss and a ReduceLROnPlateau scheduler. |
| |
|
| | ### Model Source |
| |
|
| | - **Repository:** [https://www.kaggle.com/datasets/asdasdasasdas/garbage-classification] |
| | - **License:** Apache 2.0 |
| |
|
| | ## Uses |
| |
|
| | ### Direct Use |
| |
|
| | This model can be used to classify images of garbage items into specific categories. |
| |
|
| | ### Downstream Use |
| |
|
| | Fine-tuning the model on a specific garbage classification dataset or integrating it into an application for waste management. |
| |
|
| | ### Out-of-Scope Use |
| |
|
| | This model is not suitable for general image classification tasks outside of garbage classification. |
| |
|
| | ## Bias, Risks, and Limitations |
| |
|
| | The model's performance may be affected by biases in the training data, such as underrepresentation of certain garbage types. |
| |
|
| | ### Recommendations |
| |
|
| | Users should be aware of the model's limitations and consider domain-specific data augmentation to improve performance. |
| |
|
| | ## How to Get Started with the Model |
| |
|
| | ### Example Usage |
| |
|
| | ```python |
| | import torch |
| | from torchvision import transforms |
| | from PIL import Image |
| | |
| | # Define your CNN model |
| | class CNNModel(nn.Module): |
| | def __init__(self, num_classes): |
| | # Define layers as per your CNNModel definition |
| | |
| | def forward(self, x): |
| | # Define forward pass as per your CNNModel forward method |
| | |
| | # Set device to GPU if available |
| | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| | |
| | # Load the model |
| | model = CNNModel(num_classes=num_classes).to(device) |
| | |
| | # Load the best trained weights |
| | model.load_state_dict(torch.load('best_model.pth')) |
| | model.eval() |
| | |
| | # Preprocess an image |
| | transform = transforms.Compose([ |
| | transforms.Resize((224, 224)), |
| | transforms.ToTensor(), |
| | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), |
| | ]) |
| | |
| | img_path = 'path_to_your_image.jpg' # Replace with your image path |
| | img = Image.open(img_path) |
| | input_tensor = transform(img) |
| | input_batch = input_tensor.unsqueeze(0) |
| | |
| | # Use the model for prediction |
| | with torch.no_grad(): |
| | output = model(input_batch) |
| | |
| | # Get the predicted class |
| | _, predicted = torch.max(output, 1) |
| | predicted_class = train_dataset.classes[predicted.item()] |
| | print(f'Predicted class: {predicted_class}') |
