change to yolo

app/__init__.py

@@ -3,7 +3,6 @@ import json
 import logging
 
 from dotenv import load_dotenv
-from mmdeploy_runtime import Detector
 from supabase import create_client, Client
 from firebase_admin import credentials, initialize_app, firestore
 from neo4j import GraphDatabase
@@ -15,10 +14,6 @@ logger = logging.getLogger(__name__)
 
 load_dotenv()
 
-# LOAD MODEL
-model_path = "./model"
-detector = Detector(model_path=model_path, device_name="cpu", device_id=0)
-
 # LOAD SUPABASE
 url: str = os.environ.get("SUPABASE_URL")
 key: str = os.environ.get("SUPABASE_KEY")

app/constants.py

@@ -84,87 +84,17 @@ colors = [
     (253, 245, 230),
 ]
 
-classNames = [
-    "person",
-    "bicycle",
-    "car",
-    "motorcycle",
-    "airplane",
-    "bus",
-    "train",
-    "truck",
-    "boat",
-    "traffic light",
-    "fire hydrant",
-    "stop sign",
-    "parking meter",
-    "bench",
-    "bird",
-    "cat",
-    "dog",
-    "horse",
-    "sheep",
-    "cow",
-    "elephant",
-    "bear",
-    "zebra",
-    "giraffe",
-    "backpack",
-    "umbrella",
-    "handbag",
-    "tie",
-    "suitcase",
-    "frisbee",
-    "skis",
-    "snowboard",
-    "sports ball",
-    "kite",
-    "baseball bat",
-    "baseball glove",
-    "skateboard",
-    "surfboard",
-    "tennis racket",
-    "bottle",
-    "wine glass",
-    "cup",
-    "fork",
-    "knife",
-    "spoon",
-    "bowl",
-    "banana",
-    "apple",
-    "sandwich",
-    "orange",
-    "broccoli",
-    "carrot",
-    "hot dog",
-    "pizza",
-    "donut",
-    "cake",
-    "chair",
-    "couch",
-    "potted plant",
-    "bed",
-    "dining table",
-    "toilet",
-    "tv",
-    "laptop",
-    "mouse",
-    "remote",
-    "keyboard",
-    "cell phone",
-    "microwave",
-    "oven",
-    "toaster",
-    "sink",
-    "refrigerator",
-    "book",
-    "clock",
-    "vase",
-    "scissors",
-    "teddy bear",
-    "hair drier",
-    "toothbrush",
+class_names = [
+    "Bicycle",
+    "Bus",
+    "Car",
+    "Dog",
+    "Electric pole",
+    "Motorcycle",
+    "Person",
+    "Traffic signs",
+    "Tree",
+    "Uncovered manhole",
 ]
 
-deviceId = 'dvg7OhkORV6mYFlUuNKNaE:APA91bHp-SGj4uLyyI9zpEFGrQDrklonQff0bJjtQWhGrJ-tyd_1u5NZSenxEkJzDZ-lyYsZCvmIrZnwNeu9CetsuJNGAU0lkNePNngR-76ytR8ZLvqTOMyf5RHWexsxaPa9Rop05WMi'
+deviceId = "dvg7OhkORV6mYFlUuNKNaE:APA91bHp-SGj4uLyyI9zpEFGrQDrklonQff0bJjtQWhGrJ-tyd_1u5NZSenxEkJzDZ-lyYsZCvmIrZnwNeu9CetsuJNGAU0lkNePNngR-76ytR8ZLvqTOMyf5RHWexsxaPa9Rop05WMi"

app/custom_mmcv/color.py

@@ -1,52 +0,0 @@
-# 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)}")

app/custom_mmcv/main.py

@@ -1,88 +0,0 @@
-# THE ORIGINAL mmcv.imshow_det_bboxes
-# Copyright (c) OpenMMLab. All rights reserved.
-
-from typing import List, Optional, Union
-
-import cv2
-import numpy as np
-
-from mmcv.image import imread, imwrite
-from .color import Color, color_val
-
-# a type alias declares the optional types of color argument
-ColorType = Union[Color, str, tuple, int, np.ndarray]
-
-
-def imshow_det_bboxes(
-    img: Union[str, np.ndarray],
-    bboxes: np.ndarray,
-    labels: np.ndarray,
-    class_names: List[str] = [],
-    score_thr: float = 0,
-    bbox_color: ColorType = "green",
-    text_color: ColorType = "green",
-    thickness: int = 1,
-    font_scale: float = 1,
-    out_file: Optional[str] = None,
-    colors: np.ndarray = None,
-):
-    """Draw bboxes and class labels (with scores) on an image.
-
-    Args:
-        img (str or ndarray): The image to be displayed.
-        bboxes (ndarray): Bounding boxes (with scores), shaped (n, 4) or
-            (n, 5).
-        labels (ndarray): Labels of bboxes.
-        class_names (list[str]): Names of each classes.
-        score_thr (float): Minimum score of bboxes to be shown.
-        bbox_color (Color or str or tuple or int or ndarray): Color
-            of bbox lines.
-        text_color (Color or str or tuple or int or ndarray): Color
-            of texts.
-        thickness (int): Thickness of lines.
-        font_scale (float): Font scales of texts.
-        out_file (str or None): The filename to write the image.
-        colors (array of tuple RGB int): the color of bbox and label of each class
-
-    Returns:
-        ndarray: The image with bboxes drawn on it.
-    """
-    assert bboxes.ndim == 2
-    assert labels.ndim == 1
-    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]
-        inds = scores > score_thr
-        bboxes = bboxes[inds, :]
-        labels = labels[inds]
-
-    bbox_color = color_val(bbox_color)
-    text_color = color_val(text_color)
-
-    for bbox, label in zip(bboxes, labels):
-        bbox_int = bbox.astype(np.int32)
-        left_top = (bbox_int[0], bbox_int[1])
-        right_bottom = (bbox_int[2], bbox_int[3])
-        if colors is not None and len(colors) > 0:
-            bbox_color = text_color = color_val(colors[label])
-        cv2.rectangle(img, left_top, right_bottom, bbox_color, thickness=thickness)
-        label_text = class_names[label] if class_names is not None else f"cls {label}"
-        if len(bbox) > 4:
-            label_text += f"|{bbox[-1]:.02f}"
-        cv2.putText(
-            img,
-            label_text,
-            (bbox_int[0], bbox_int[1] - 2),
-            cv2.FONT_HERSHEY_TRIPLEX,
-            font_scale,
-            text_color,
-            4,
-        )
-
-    if out_file is not None:
-        imwrite(img, out_file)
-    return img

app/detector/__init__.py

@@ -1,4 +1,4 @@
 from .yolov8.YOLOv8 import YOLOv8
 
-model_path = "./model"
+model_path = "./model/best20231112.onnx"
 detector = YOLOv8(model_path)

app/detector/yolov8/YOLOv8.py

@@ -3,7 +3,7 @@ import cv2
 import numpy as np
 import onnxruntime
 
-from yolov8.utils import xywh2xyxy, draw_detections, multiclass_nms
+from app.detector.yolov8.utils import xywh2xyxy, draw_detections, multiclass_nms
 
 
 class YOLOv8:
@@ -17,6 +17,9 @@ class YOLOv8:
     def __call__(self, image):
         return self.detect_objects(image)
 
+    def set_conf_threshold(self, conf_thres):
+        self.conf_threshold = conf_thres
+
     def initialize_model(self, path):
         self.session = onnxruntime.InferenceSession(
             path, providers=onnxruntime.get_available_providers()

app/detector/yolov8/utils.py

@@ -1,89 +1,9 @@
 from typing import Tuple
 import numpy as np
 import cv2
+from app.constants import class_names
 
-class_names = [
-    "person",
-    "bicycle",
-    "car",
-    "motorcycle",
-    "airplane",
-    "bus",
-    "train",
-    "truck",
-    "boat",
-    "traffic light",
-    "fire hydrant",
-    "stop sign",
-    "parking meter",
-    "bench",
-    "bird",
-    "cat",
-    "dog",
-    "horse",
-    "sheep",
-    "cow",
-    "elephant",
-    "bear",
-    "zebra",
-    "giraffe",
-    "backpack",
-    "umbrella",
-    "handbag",
-    "tie",
-    "suitcase",
-    "frisbee",
-    "skis",
-    "snowboard",
-    "sports ball",
-    "kite",
-    "baseball bat",
-    "baseball glove",
-    "skateboard",
-    "surfboard",
-    "tennis racket",
-    "bottle",
-    "wine glass",
-    "cup",
-    "fork",
-    "knife",
-    "spoon",
-    "bowl",
-    "banana",
-    "apple",
-    "sandwich",
-    "orange",
-    "broccoli",
-    "carrot",
-    "hot dog",
-    "pizza",
-    "donut",
-    "cake",
-    "chair",
-    "couch",
-    "potted plant",
-    "bed",
-    "dining table",
-    "toilet",
-    "tv",
-    "laptop",
-    "mouse",
-    "remote",
-    "keyboard",
-    "cell phone",
-    "microwave",
-    "oven",
-    "toaster",
-    "sink",
-    "refrigerator",
-    "book",
-    "clock",
-    "vase",
-    "scissors",
-    "teddy bear",
-    "hair drier",
-    "toothbrush",
-]
+class_names = class_names
 
 # Create a list of colors for each class where each color is a tuple of 3 integer values
 rng = np.random.default_rng(3)

app/routers/image.py

@@ -1,11 +1,8 @@
 import cv2
-import numpy as np
 
-from fastapi import APIRouter, File, Response, WebSocket, WebSocketDisconnect
-from app.constants import classNames, colors
-from app import detector
+from fastapi import APIRouter, File, Response
+from app.detector import detector
 from mmcv import imfrombytes
-from app.custom_mmcv.main import imshow_det_bboxes
 from app import logger
 
 router = APIRouter(prefix="/image", tags=["Image"])
@@ -15,14 +12,11 @@ router = APIRouter(prefix="/image", tags=["Image"])
 async def handleImageRequest(
     file: bytes = File(...),
     threshold: float = 0.3,
-    raw: bool = False,
 ):
     try:
         img = imfrombytes(file, cv2.IMREAD_COLOR)
-        if raw:
-            bboxes, labels = inferenceImage(img, threshold, raw)
-            return {"bboxes": bboxes.tolist(), "labels": labels.tolist()}
-        img = inferenceImage(img, threshold, raw)
+
+        img = inference_image(img, threshold)
     except Exception as e:
         logger.error(e)
         return Response(content="Failed to read image", status_code=400)
@@ -36,24 +30,7 @@ async def handleImageRequest(
     return Response(content=jpeg_bytes, media_type="image/jpeg")
 
 
-def inferenceImage(img, threshold: float, isRaw: bool):
-    bboxes, labels, _ = detector(img)
-    if isRaw:
-        removeIndexs = []
-        for i, bbox in enumerate(bboxes):
-            if bbox[4] < threshold:
-                removeIndexs.append(i)
-
-        bboxes = np.delete(bboxes, removeIndexs, axis=0)
-        labels = np.delete(labels, removeIndexs)
-
-        return bboxes, labels
-    return imshow_det_bboxes(
-        img=img,
-        bboxes=bboxes,
-        labels=labels,
-        class_names=classNames,
-        colors=colors,
-        score_thr=threshold,
-    )
-
+def inference_image(img, threshold):
+    detector.set_conf_threshold(threshold)
+    detector(img)
+    return detector.draw_detections(img)

app/routers/video.py

@@ -9,7 +9,6 @@ import cv2
 from multiprocessing import Process
 from fastapi import (
     APIRouter,
-    Depends,
     HTTPException,
     UploadFile,
     BackgroundTasks,
@@ -19,7 +18,7 @@ from firebase_admin import messaging
 from app import db
 from app import supabase
 from app.dependencies import get_current_user
-from app.routers.image import inferenceImage
+from app.routers.image import inference_image
 from google.cloud.firestore_v1.base_query import FieldFilter
 from app import logger
 
@@ -31,7 +30,6 @@ async def handleVideoRequest(
     file: UploadFile,
     background_tasks: BackgroundTasks,
     threshold: float = 0.3,
-    user=Depends(get_current_user),
 ):
     if re.search("^video\/", file.content_type) is None:
         raise HTTPException(
@@ -48,7 +46,7 @@ async def handleVideoRequest(
         async with aiofiles.open(os.path.join(id, "input.mp4"), "wb") as out_file:
             while content := await file.read(1024):
                 await out_file.write(content)
-        background_tasks.add_task(inferenceVideo, artifact_ref.id, id, threshold)
+        await inference_video(artifact_ref.id, id, threshold)
         return id + ".mp4"
     except ValueError as err:
         logger.error(err)
@@ -93,7 +91,7 @@ def inference_frame(inputDir, threshold: float = 0.3):
         if res == False:
             break
 
-        resFram = inferenceImage(frame, threshold, False)
+        resFram = inference_image(frame, threshold)
         result.write(resFram)
     cap.release()
     result.release()
@@ -102,9 +100,9 @@ def inference_frame(inputDir, threshold: float = 0.3):
     return thumbnail
 
 
-async def inferenceVideo(artifactId: str, inputDir: str, threshold: float):
+async def inference_video(artifactId: str, inputDir: str, threshold: float):
     try:
-        Process(updateArtifact(artifactId, {"status": "processing"})).start()
+        Process(update_artifact(artifactId, {"status": "processing"})).start()
         thumbnail = inference_frame(inputDir, threshold=threshold)
         createThumbnail(thumbnail, inputDir)
 
@@ -129,7 +127,7 @@ async def inferenceVideo(artifactId: str, inputDir: str, threshold: float):
         except Exception as e:
             print(e)
 
-        updateArtifact(
+        update_artifact(
             artifactId,
             {
                 "status": "success",
@@ -143,7 +141,7 @@ async def inferenceVideo(artifactId: str, inputDir: str, threshold: float):
         )
     except:
         Process(
-            updateArtifact(
+            update_artifact(
                 artifactId,
                 {
                     "status": "fail",
@@ -157,7 +155,7 @@ async def inferenceVideo(artifactId: str, inputDir: str, threshold: float):
             print(e)
 
 
-def updateArtifact(artifactId: str, body):
+def update_artifact(artifactId: str, body):
     artifact_ref = db.collection("artifacts").document(artifactId)
     artifact_snapshot = artifact_ref.get()
     if artifact_snapshot.exists:

model/{end2end.onnx → best20231112.onnx}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e0efff65be17a13541a3b93d6b0cde60ed37f996c4f2d57aa38071e48117f217
-size 22245735
+oid sha256:21f0a8a6f0398cef7e73d74919d8a50fabbf5c1f5dd9e4192b129ce53153a317
+size 12119302

model/deploy.json

@@ -1,16 +0,0 @@
-{
-    "version": "1.2.0",
-    "task": "Detector",
-    "models": [
-        {
-            "name": "rtmdet",
-            "net": "end2end.onnx",
-            "weights": "",
-            "backend": "onnxruntime",
-            "precision": "FP32",
-            "batch_size": 1,
-            "dynamic_shape": true
-        }
-    ],
-    "customs": []
-}

model/detail.json

@@ -1,68 +0,0 @@
-{
-    "version": "1.2.0",
-    "codebase": {
-        "task": "ObjectDetection",
-        "codebase": "mmdet",
-        "version": "3.1.0",
-        "pth": "rtmdet_tiny_8xb32-300e_coco_20220902_112414-78e30dcc.pth",
-        "config": "rtmdet_tiny_8xb32-300e_coco.py"
-    },
-    "codebase_config": {
-        "type": "mmdet",
-        "task": "ObjectDetection",
-        "model_type": "end2end",
-        "post_processing": {
-            "score_threshold": 0.05,
-            "confidence_threshold": 0.005,
-            "iou_threshold": 0.5,
-            "max_output_boxes_per_class": 200,
-            "pre_top_k": 5000,
-            "keep_top_k": 100,
-            "background_label_id": -1
-        }
-    },
-    "onnx_config": {
-        "type": "onnx",
-        "export_params": true,
-        "keep_initializers_as_inputs": false,
-        "opset_version": 11,
-        "save_file": "end2end.onnx",
-        "input_names": [
-            "input"
-        ],
-        "output_names": [
-            "dets",
-            "labels"
-        ],
-        "input_shape": null,
-        "optimize": true,
-        "dynamic_axes": {
-            "input": {
-                "0": "batch",
-                "2": "height",
-                "3": "width"
-            },
-            "dets": {
-                "0": "batch",
-                "1": "num_dets"
-            },
-            "labels": {
-                "0": "batch",
-                "1": "num_dets"
-            }
-        }
-    },
-    "backend_config": {
-        "type": "onnxruntime",
-        "precision": "fp16",
-        "common_config": {
-            "min_positive_val": 1e-07,
-            "max_finite_val": 10000.0,
-            "keep_io_types": false,
-            "disable_shape_infer": false,
-            "op_block_list": null,
-            "node_block_list": null
-        }
-    },
-    "calib_config": {}
-}

model/pipeline.json

@@ -1,132 +0,0 @@
-{
-    "pipeline": {
-        "input": [
-            "img"
-        ],
-        "output": [
-            "post_output"
-        ],
-        "tasks": [
-            {
-                "type": "Task",
-                "module": "Transform",
-                "name": "Preprocess",
-                "input": [
-                    "img"
-                ],
-                "output": [
-                    "prep_output"
-                ],
-                "transforms": [
-                    {
-                        "type": "LoadImageFromFile",
-                        "backend_args": null
-                    },
-                    {
-                        "type": "Resize",
-                        "keep_ratio": true,
-                        "size": [
-                            640,
-                            640
-                        ]
-                    },
-                    {
-                        "type": "Pad",
-                        "size": [
-                            640,
-                            640
-                        ],
-                        "pad_val": {
-                            "img": [
-                                114,
-                                114,
-                                114
-                            ]
-                        }
-                    },
-                    {
-                        "type": "Normalize",
-                        "to_rgb": false,
-                        "mean": [
-                            103.53,
-                            116.28,
-                            123.675
-                        ],
-                        "std": [
-                            57.375,
-                            57.12,
-                            58.395
-                        ]
-                    },
-                    {
-                        "type": "Pad",
-                        "size_divisor": 1
-                    },
-                    {
-                        "type": "DefaultFormatBundle"
-                    },
-                    {
-                        "type": "Collect",
-                        "meta_keys": [
-                            "img_path",
-                            "flip",
-                            "ori_filename",
-                            "pad_param",
-                            "pad_shape",
-                            "scale_factor",
-                            "valid_ratio",
-                            "flip_direction",
-                            "img_norm_cfg",
-                            "img_id",
-                            "img_shape",
-                            "ori_shape",
-                            "filename"
-                        ],
-                        "keys": [
-                            "img"
-                        ]
-                    }
-                ]
-            },
-            {
-                "name": "rtmdet",
-                "type": "Task",
-                "module": "Net",
-                "is_batched": true,
-                "input": [
-                    "prep_output"
-                ],
-                "output": [
-                    "infer_output"
-                ],
-                "input_map": {
-                    "img": "input"
-                },
-                "output_map": {}
-            },
-            {
-                "type": "Task",
-                "module": "mmdet",
-                "name": "postprocess",
-                "component": "ResizeBBox",
-                "params": {
-                    "nms_pre": 1000,
-                    "min_bbox_size": 0,
-                    "score_thr": 0.3,
-                    "nms": {
-                        "type": "nms",
-                        "iou_threshold": 0.65
-                    },
-                    "max_per_img": 100
-                },
-                "output": [
-                    "post_output"
-                ],
-                "input": [
-                    "prep_output",
-                    "infer_output"
-                ]
-            }
-        ]
-    }
-}

tests/test_custom_mmcv.py

@@ -1,6 +1,6 @@
 from app.custom_mmcv.color import color_val
 from app.custom_mmcv.main import imshow_det_bboxes
-from app.constants import classNames
+from app.constants import class_names
 import mmcv
 import cv2
 import numpy as np
@@ -30,22 +30,22 @@ class TestCustomMMCV():
         image = mmcv.imread('demo.jpg')
         bboxes = np.ones((1,5))
         labels = np.zeros(1, dtype=np.int32)
-        result = imshow_det_bboxes(image, bboxes, labels, class_names=classNames, bbox_color="red", text_color='red')
+        result = imshow_det_bboxes(image, bboxes, labels, class_names=class_names, bbox_color="red", text_color='red')
         assert (result[1,1,:] == (0,0,255)).all()
         with pytest.raises(AssertionError):
             bboxes = np.ones((1,3))
             labels = np.zeros(1, dtype=np.int32)
-            result = imshow_det_bboxes(image, bboxes, labels, class_names=classNames, bbox_color="red",text_color="red")
+            result = imshow_det_bboxes(image, bboxes, labels, class_names=class_names, bbox_color="red",text_color="red")
         with pytest.raises(AssertionError):
             bboxes = np.ones((1,7))
             labels = np.zeros(1, dtype=np.int32)
-            result = imshow_det_bboxes(image, bboxes, labels, class_names=classNames, bbox_color="red",text_color="red")
+            result = imshow_det_bboxes(image, bboxes, labels, class_names=class_names, bbox_color="red",text_color="red")
         with pytest.raises(AssertionError):
             bboxes = np.ones((1,5))
             labels = np.zeros(4, dtype=np.int32)
-            result = imshow_det_bboxes(image, bboxes, labels, class_names=classNames, bbox_color="red",text_color="red")
+            result = imshow_det_bboxes(image, bboxes, labels, class_names=class_names, bbox_color="red",text_color="red")
         with pytest.raises(AssertionError):
             bboxes = np.ones((2,5))
             labels = np.zeros(1, dtype=np.int32)
-            result = imshow_det_bboxes(image, bboxes, labels, class_names=classNames, bbox_color="red",text_color="red")
+            result = imshow_det_bboxes(image, bboxes, labels, class_names=class_names, bbox_color="red",text_color="red")
         

tests/test_image.py

@@ -1,7 +1,7 @@
 import os
 import mmcv
 import numpy as np
-from app.routers.image import inferenceImage
+from app.routers.image import inference_image
 import pytest
 import pytest
 from fastapi.testclient import TestClient
@@ -22,9 +22,9 @@ class TestImageRoute():
     img = mmcv.imread('demo.jpg')
     url = "http://0.0.0.0:3000/image"
     def test_inferenceImage(self):
-        bboxes, labels = inferenceImage(mmcv.imread('demo.jpg'), 0.3, True)
+        bboxes, labels = inference_image(mmcv.imread('demo.jpg'), 0.3, True)
         assert len(bboxes.tolist()) > 0 and len(labels.tolist()) > 0 and len(bboxes.tolist()) == len(labels.tolist())
-        result = inferenceImage(self.img, 0.3, False)
+        result = inference_image(self.img, 0.3, False)
         assert type(result) is np.ndarray and result.shape == self.img.shape
     def test_ImageAPI(self, client):
         payload = {}

tests/test_video.py

@@ -1,6 +1,6 @@
 from fastapi.testclient import TestClient
 from fastapi.routing import APIRoute
-from app.routers.video import updateArtifact, createThumbnail, inference_frame
+from app.routers.video import update_artifact, createThumbnail, inference_frame
 from app.main import app
 from app.constants import deviceId
 from app import db
@@ -107,7 +107,7 @@ class TestVideoAPI:
         else:
             test_artifact.update({"status": "testing", 'path': '', "thumbnailURL":""})
         # Testing update on each field
-        updateArtifact(test_artifact.id,{"status": "test_done"})
+        update_artifact(test_artifact.id,{"status": "test_done"})
         assert db.collection("artifacts").document('test').get().to_dict()['status'] == 'test_done'
         #Delete data for next time test
         test_artifact.delete()