From 4b57251d5e7a421b440e9d4faa2ea394759bd4ce Mon Sep 17 00:00:00 2001
From: fcdl94 <S203813s3>
Date: Tue, 10 Dec 2024 09:56:41 +0100
Subject: [PATCH] Fix misalignement between Local and Remote model infer input
 types.

---
 focoos/__init__.py     | 1 +
 focoos/remote_model.py | 8 ++++++--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/focoos/__init__.py b/focoos/__init__.py
index 69049cd..b8ab9d6 100644
--- a/focoos/__init__.py
+++ b/focoos/__init__.py
@@ -1,3 +1,4 @@
 from .focoos import Focoos
+from .local_model import LocalModel
 from .ports import *
 from .remote_model import RemoteModel
diff --git a/focoos/remote_model.py b/focoos/remote_model.py
index a5b08f4..7ccb337 100644
--- a/focoos/remote_model.py
+++ b/focoos/remote_model.py
@@ -4,6 +4,7 @@
 from time import sleep
 from typing import Optional, Tuple, Union
 
+import cv2
 import numpy as np
 from supervision import BoxAnnotator, Detections, LabelAnnotator, MaskAnnotator
 
@@ -120,16 +121,19 @@ def _annotate(self, im: np.ndarray, detections: Detections) -> np.ndarray:
 
     def infer(
         self,
-        image: Union[str, Path, bytes],
+        image: Union[str, Path, np.ndarray, bytes],
         threshold: float = 0.5,
         annotate: bool = False,
     ) -> Tuple[FocoosDetections, Optional[np.ndarray]]:
         image_bytes = None
-        if not isinstance(image, bytes):
+        if isinstance(image, str) or isinstance(image, Path):
             if not os.path.exists(image):
                 logger.error(f"Image file not found: {image}")
                 raise FileNotFoundError(f"Image file not found: {image}")
             image_bytes = open(image, "rb").read()
+        elif isinstance(image, np.ndarray):
+            _, buffer = cv2.imencode(".jpg", image)
+            image_bytes = buffer.tobytes()
         else:
             image_bytes = image
         files = {"file": image_bytes}