Skip to main content

​
Overview

The ModelLoader class handles the initialization and management of YOLO segmentation models. It loads the pre-trained trash classification model and assigns it to the specified PyTorch device for optimal performance.

​
Class Definition

from trash_classificator.segmentation.model_loader import ModelLoader
import torch

loader = ModelLoader(device=torch.device("cuda"))
model = loader.get_model()

​
Constructor

​
__init__(device: torch.device)

Initializes the ModelLoader with a specific device and loads the YOLO model.
​
device
torch.device
required
The PyTorch device to load the model onto. Can be CPU, CUDA (GPU), or MPS (Apple Silicon).

​
Behavior

  • Loads the YOLO model from the path specified in trash_model_path
  • Transfers the model to the specified device
  • The model path resolves to: trash_classificator/segmentation/models/trash_segmentation_model_v2.pt
The model is automatically loaded from the trash_model_path constant defined in trash_classificator.segmentation.models.trash_model.

​
Methods

​
get_model() -> YOLO

Returns the loaded YOLO model instance.
​
model
YOLO
The Ultralytics YOLO model instance configured for trash segmentation. This model can classify three categories: cardboard and paper, metal, and plastic.

​
Model Path Handling

The model path is dynamically constructed based on the package installation location: 
# From trash_model.py
current_path = os.path.dirname(os.path.abspath(__file__))
trash_model_path: str = current_path + "/trash_segmentation_model_v2.pt"
This ensures the model file is always located relative to the package installation directory.

​
Usage Example

​
Basic Usage

import torch
from trash_classificator.segmentation.model_loader import ModelLoader
from trash_classificator.segmentation.device_manager import DeviceManager

# Get optimal device
device = DeviceManager.get_device()

# Load model
loader = ModelLoader(device=device)
model = loader.get_model()

# Run inference
results = model("image.jpg")

​
Device-Specific Loading

import torch
from trash_classificator.segmentation.model_loader import ModelLoader

# Force CPU usage
loader = ModelLoader(device=torch.device("cpu"))
model = loader.get_model()

# Use CUDA if available
if torch.cuda.is_available():
    loader = ModelLoader(device=torch.device("cuda"))
    model = loader.get_model()
Use DeviceManager.get_device() to automatically select the best available device instead of manually specifying it.

​
Parameters and Return Types

MethodParametersReturn TypeDescription
__init__device: torch.deviceNoneInitializes and loads the model
get_modelNoneYOLOReturns the loaded model instance

​
Model Specifications

  • Model Type: YOLOv8 Segmentation
  • Model File: trash_segmentation_model_v2.pt
  • Classes Detected:
    • 0: cardboard and paper
    • 1: metal
    • 2: plastic

​
Source Reference

Implementation: trash_classificator/segmentation/model_loader.py:6-11

Build docs developers (and LLMs) love

Get started for freeTalk to us