SoyoKaze83/trashnet-clf

YOLOv8 on Trashnet

This model comes from the YOLO family, specifically YOLOv8, fine-tuned on the Trashnet dataset which can be viewed and downloaded from this link. The dataset provides us the task to classify images of trash categories including paper, glass, cardboard, plastic, metal, and trash.

Training Details

The model is trained on 2527 images resized to 512x512 with a split of 80% for the training set and 20% for the validation set. The main metric for evaluating the model is Top 1 Accuracy which shows us how well the model classifies an image based on the highest probability from the output probability distribution.

Running Inference

from huggingface_hub import hf_hub_download
from ultralytics import YOLO

# First, you'll need to download the model file from HuggingFace
model_path = hf_hub_download(repo_id="SoyoKaze83/trashnet-clf", 
                            filename="weights/yolov8.pt")

# Then load the downloaded model
model = YOLO(model_path)

# Single image inference
results = model("path/to/image.jpg")  # replace with your image path

# Get prediction for single image
for r in results:
    # Get the predicted class and confidence
    probs = r.probs  # probability for each class
    cls_id = int(probs.top1)  # index of top class
    conf = float(probs.top1conf)  # confidence of top class
    cls_name = model.names[cls_id]  # name of predicted class
    
    print(f"Predicted class: {cls_name} with confidence: {conf:.2f}")
    
    # If you want all class probabilities
    all_probs = probs.data.tolist()  # probabilities for all classes
    for i, prob in enumerate(all_probs):
        print(f"Class {model.names[i]}: {prob:.2f}")

# Batch inference
image_paths = [
    "path/to/image1.jpg",
    "path/to/image2.jpg",
    "path/to/image3.jpg"
]

# Process batch of images
results = model(image_paths)
for i, r in enumerate(results):
    probs = r.probs
    cls_id = int(probs.top1)
    conf = float(probs.top1conf)
    cls_name = model.names[cls_id]
    print(f"\nImage {i+1}:")
    print(f"Predicted class: {cls_name} with confidence: {conf:.2f}")