move custom library to code

Dockerfile

@@ -19,8 +19,6 @@ COPY --chown=user . .
 
 RUN pip install -r ./app/requirements.txt
 
-COPY ./libs/image.py /home/user/.local/lib/python3.8/site-packages/mmcv/visualization/image.py
-
 EXPOSE 3000
 
 CMD [ "uvicorn", "app.main:app", "--host" ,"0.0.0.0" ,"--port", "3000"]

app/custom_mmcv/color.py

@@ -0,0 +1,52 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from enum import Enum
+from typing import Union
+
+import numpy as np
+from mmengine.utils import is_str
+
+
+class Color(Enum):
+    """An enum that defines common colors.
+
+    Contains red, green, blue, cyan, yellow, magenta, white and black.
+    """
+
+    red = (0, 0, 255)
+    green = (0, 255, 0)
+    blue = (255, 0, 0)
+    cyan = (255, 255, 0)
+    yellow = (0, 255, 255)
+    magenta = (255, 0, 255)
+    white = (255, 255, 255)
+    black = (0, 0, 0)
+
+
+def color_val(color: Union[Color, str, tuple, int, np.ndarray]) -> tuple:
+    """Convert various input to color tuples.
+
+    Args:
+        color (:obj:`Color`/str/tuple/int/ndarray): Color inputs
+
+    Returns:
+        tuple[int]: A tuple of 3 integers indicating BGR channels.
+    """
+    if is_str(color):
+        return Color[color].value  # type: ignore
+    elif isinstance(color, Color):
+        return color.value
+    elif isinstance(color, tuple):
+        assert len(color) == 3
+        for channel in color:
+            assert 0 <= channel <= 255
+        return color
+    elif isinstance(color, int):
+        assert 0 <= color <= 255
+        return color, color, color
+    elif isinstance(color, np.ndarray):
+        assert color.ndim == 1 and color.size == 3
+        assert np.all((color >= 0) & (color <= 255))
+        color = color.astype(np.uint8)
+        return tuple(color)
+    else:
+        raise TypeError(f"Invalid type for color: {type(color)}")

libs/image.py → app/custom_mmcv/main.py

@@ -1,4 +1,6 @@
+# THE ORIGINAL mmcv.imshow_det_bboxes
 # Copyright (c) OpenMMLab. All rights reserved.
+
 from typing import List, Optional, Union
 
 import cv2
@@ -32,61 +34,6 @@ def imshow(img: Union[str, np.ndarray], win_name: str = "", wait_time: int = 0):
         ret = cv2.waitKey(wait_time)
 
 
-def imshow_bboxes(
-    img: Union[str, np.ndarray],
-    bboxes: Union[list, np.ndarray],
-    colors: ColorType = "green",
-    top_k: int = -1,
-    thickness: int = 1,
-    show: bool = True,
-    win_name: str = "",
-    wait_time: int = 0,
-    out_file: Optional[str] = None,
-):
-    """Draw bboxes on an image.
-
-    Args:
-        img (str or ndarray): The image to be displayed.
-        bboxes (list or ndarray): A list of ndarray of shape (k, 4).
-        colors (Color or str or tuple or int or ndarray): A list of colors.
-        top_k (int): Plot the first k bboxes only if set positive.
-        thickness (int): Thickness of lines.
-        show (bool): Whether to show the image.
-        win_name (str): The window name.
-        wait_time (int): Value of waitKey param.
-        out_file (str, optional): The filename to write the image.
-
-    Returns:
-        ndarray: The image with bboxes drawn on it.
-    """
-    img = imread(img)
-    img = np.ascontiguousarray(img)
-
-    if isinstance(bboxes, np.ndarray):
-        bboxes = [bboxes]
-    if not isinstance(colors, list):
-        colors = [colors for _ in range(len(bboxes))]
-    colors = [color_val(c) for c in colors]
-    assert len(bboxes) == len(colors)
-
-    for i, _bboxes in enumerate(bboxes):
-        _bboxes = _bboxes.astype(np.int32)
-        if top_k <= 0:
-            _top_k = _bboxes.shape[0]
-        else:
-            _top_k = min(top_k, _bboxes.shape[0])
-        for j in range(_top_k):
-            left_top = (_bboxes[j, 0], _bboxes[j, 1])
-            right_bottom = (_bboxes[j, 2], _bboxes[j, 3])
-            cv2.rectangle(img, left_top, right_bottom, colors[i], thickness=thickness)
-
-    if show:
-        imshow(img, win_name, wait_time)
-    if out_file is not None:
-        imwrite(img, out_file)
-    return img
-
-
 def imshow_det_bboxes(
     img: Union[str, np.ndarray],
     bboxes: np.ndarray,
@@ -132,7 +79,7 @@ def imshow_det_bboxes(
     assert bboxes.shape[0] == labels.shape[0]
     assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5
     img = imread(img)
-    
+
     if score_thr > 0:
         assert bboxes.shape[1] == 5
         scores = bboxes[:, -1]
@@ -160,7 +107,7 @@ def imshow_det_bboxes(
             cv2.FONT_HERSHEY_TRIPLEX,
             font_scale,
             text_color,
-            4
+            4,
         )
 
     if show:

app/routers/image.py

@@ -1,10 +1,11 @@
 from fastapi import APIRouter, File, Response, WebSocket, WebSocketDisconnect
-import mmcv
 import cv2
 import numpy as np
 from app.constants import classNames, colors
 from app import detector
 
+from mmcv import imfrombytes
+from app.custom_mmcv.main import imshow_det_bboxes
 
 router = APIRouter(prefix="/image")
 
@@ -15,7 +16,7 @@ async def handleImageRequest(
     threshold: float = 0.3,
     raw: bool = False,
 ):
-    img = mmcv.imfrombytes(file, cv2.IMREAD_COLOR)
+    img = imfrombytes(file, cv2.IMREAD_COLOR)
     if raw:
         bboxes, labels = inferenceImage(img, threshold, raw)
         return {"bboxes": bboxes.tolist(), "labels": labels.tolist()}
@@ -42,7 +43,7 @@ def inferenceImage(img, threshold: float, isRaw: bool = False):
         labels = np.delete(labels, removeIndexs)
 
         return bboxes, labels
-    return mmcv.imshow_det_bboxes(
+    return imshow_det_bboxes(
         img=img,
         bboxes=bboxes,
         labels=labels,
@@ -59,7 +60,7 @@ async def websocketEndpoint(websocket: WebSocket, threshold: float = 0.3):
     try:
         while True:
             data = await websocket.receive_bytes()
-            img = mmcv.imfrombytes(data, cv2.IMREAD_COLOR)
+            img = imfrombytes(data, cv2.IMREAD_COLOR)
             bboxes, labels = inferenceImage(img, threshold, True)
             await websocket.send_json(
                 {"bboxes": bboxes.tolist(), "labels": labels.tolist()}