Skip to content

Commit

Permalink
add deploy predict test
Browse files Browse the repository at this point in the history
  • Loading branch information
Ash-Lee233 committed Aug 3, 2023
1 parent b80bb44 commit 3018ccc
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 0 deletions.
3 changes: 3 additions & 0 deletions deploy/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from .predict import detect

__all__ = ['detect']
51 changes: 51 additions & 0 deletions tests/modules/test_deploy_predict.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
import sys

sys.path.append(".")
import os
import numpy as np
from download import download
import cv2

from mindspore import nn

from mindyolo.utils.metrics import non_max_suppression, scale_coords, xyxy2xywh
from mindyolo.data import COCO80_TO_COCO91_CLASS
from deploy.infer_engine.lite import LiteModel
from mindyolo.utils.utils import draw_result
from deploy.predict import detect

def test_deploy_predict():
image_url = 'https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/image_cat.zip'
path = download(image_url, './', kind="zip", replace=True)
image_path = ('./image_cat/jpg/000000039769.jpg')

model_url = 'https://download.mindspore.cn/toolkits/mindyolo/yolov5/yolov5n_300e_mAP273-9b16bd7b-bd03027b.mindir'
model_path = download(model_url, './yolov5n.mindir', kind="file", replace=True)

network = LiteModel(model_path)
img = cv2.imread(image_path)
result_dict = detect(
network=network,
img=img,
conf_thres=0.1,
iou_thres=0.65,
conf_free=False,
nms_time_limit=20,
img_size=640,
is_coco_dataset=True,
)
save_path = os.path.join('./', "detect_results")
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' ]
draw_result(image_path, result_dict, names, save_path=save_path)


if __name__ == '__main__':
test_deploy_predict()

0 comments on commit 3018ccc

Please sign in to comment.