services/app_service/model.py

from abc import ABC, abstractmethod
from pathlib import Path

from PIL import Image
from ultralytics import YOLO


class BaseModel(ABC):
    @abstractmethod
    def __init__(self):
        pass

    @abstractmethod
    def predict_image(self, image):
        pass

    @abstractmethod
    def predict_video(self, video):
        pass


class YOLOModelv1(BaseModel):
    """Model: modelYOLOv8n_datasetDOTAv2_epochs5_batch1.pt"""

    def __init__(self):
        repo_root = Path(__file__).resolve().parents[1]
        weights_path = (
            repo_root / "models" / "modelYOLOv8n_datasetDOTAv2_epochs5_batch1.pt"
        )
        self.model = YOLO(str(weights_path), task="detect")

    def predict_image(self, image, min_confidence):
        results = self.model.predict(image, save=False, imgsz=640, conf=min_confidence)
        annotated_image_filename = "annotated_image.png"
        last_im = None
        for result in results:
            im_array = result.plot()
            last_im = Image.fromarray(im_array[..., ::-1])  # RGB PIL image
            last_im.save(annotated_image_filename)
        # Return PIL Image for robust display in Streamlit
        return last_im if last_im is not None else Image.open(annotated_image_filename)

    def predict_video(self, video, min_confidence, target_dir_name="annotated_video"):
        self.model.predict(
            video,
            save=True,
            project=".",
            name=target_dir_name,
            exist_ok=True,
            imgsz=640,
            conf=min_confidence,
        )


class YOLOModelv2(BaseModel):
    """Model: modelYOLOv8n_datasetDIOR_epochs50_batch16.pt"""

    def __init__(self):
        repo_root = Path(__file__).resolve().parents[1]
        weights_path = (
            repo_root / "models" / "modelYOLOv8n_datasetDIOR_epochs50_batch16.pt"
        )
        self.model = YOLO(str(weights_path), task="detect")

    def predict_image(self, image, min_confidence, classes=None):
        results = self.model.predict(
            image, save=False, imgsz=800, conf=min_confidence, classes=classes
        )
        annotated_image_filename = "annotated_image.png"
        last_im = None
        for result in results:
            im_array = result.plot()
            last_im = Image.fromarray(im_array[..., ::-1])  # RGB PIL image
            last_im.save(annotated_image_filename)
        # Return PIL Image for robust display in Streamlit
        return last_im if last_im is not None else Image.open(annotated_image_filename)

    def predict_video(
        self, video, min_confidence, target_dir_name="annotated_video", classes=None
    ):
        self.model.predict(
            video,
            save=True,
            project=".",
            name=target_dir_name,
            exist_ok=True,
            imgsz=800,
            conf=min_confidence,
            classes=classes,
        )