Views
No views yet
1from huggingface_hub import hf_hub_download
2from ultralytics import YOLO
3
4# First, you'll need to download the model file from HuggingFace
5model_path = hf_hub_download(repo_id="SoyoKaze83/trashnet-clf",
6 filename="weights/yolov8.pt")
7
8# Then load the downloaded model
9model = YOLO(model_path)
10
11# Single image inference
12results = model("path/to/image.jpg") # replace with your image path
13
14# Get prediction for single image
15for r in results:
16 # Get the predicted class and confidence
17 probs = r.probs # probability for each class
18 cls_id = int(probs.top1) # index of top class
19 conf = float(probs.top1conf) # confidence of top class
20 cls_name = model.names[cls_id] # name of predicted class
21
22 print(f"Predicted class: {cls_name} with confidence: {conf:.2f}")
23
24 # If you want all class probabilities
25 all_probs = probs.data.tolist() # probabilities for all classes
26 for i, prob in enumerate(all_probs):
27 print(f"Class {model.names[i]}: {prob:.2f}")
28
29# Batch inference
30image_paths = [
31 "path/to/image1.jpg",
32 "path/to/image2.jpg",
33 "path/to/image3.jpg"
34]
35
36# Process batch of images
37results = model(image_paths)
38for i, r in enumerate(results):
39 probs = r.probs
40 cls_id = int(probs.top1)
41 conf = float(probs.top1conf)
42 cls_name = model.names[cls_id]
43 print(f"\nImage {i+1}:")
44 print(f"Predicted class: {cls_name} with confidence: {conf:.2f}")