src/controllers/sack_detector.py

import cv2
from ultralytics import YOLO

from src.handlers.image_processing import draw_box, resize_image

import os
from dotenv import load_dotenv
load_dotenv()

class SackDetector:
    def __init__(self):
        self.model_path = os.getenv("YOLO_MODEL_PATH", None)  # Default path
        if self.model_path is None:
            raise ValueError("YOLO_MODEL_PATH environment variable is not set.")
        self.model = YOLO(self.model_path)
        self.class_list = self.model.model.names

    def detect_objects(self, image_path):
        results = self.model(image_path)
        result = results[0]
        labeled_img = draw_box(result.orig_img, result, self.class_list)
        # display_img = resize_image(labeled_img, 100)
        return labeled_img, result

    def display_image(self, image):
        cv2.imshow('outputimage', image)
        cv2.waitKey(0)  # Wait for user input to close the window