det_demo.py

Go to the documentation of this file.

#system path 설정
import sys
sys_path = ['/workspace/smlab', '/workspace']
for path in sys_path:
    if path not in sys.path:
        sys.path.append(path)
 
#import package
import cv2
import smrunner
from smutils.utils_vis import vis_instance_segmentation, draw_bboxes_list
 
#경로 설정
# cfg_path = '/workspace/smlab/smdetect/configs/mmdet/rtmdet/rtmdet_trt.py' # rtmdet
cfg_path = '/workspace/smlab/smdetect/configs/yolo/yolov8.py' # yolov8
image_path = '/workspace/smlab/smdetect/demo/sample.bmp'
save_path = '/workspace/smlab/smdetect/demo/det_demo.jpg'
 
# 모델 빌드
det_model = smrunner.build_model(cfg_path)
 
# 이미지 로드 : 이미지 경로를 설정해야함
image = cv2.imread(image_path)
 
# 모델 inference
result = det_model.run_detector(image)
 
#시각화 영상 만들기 및 저장
vis_img = image.copy()
if 'masks' in result:
    vis_img = vis_instance_segmentation(vis_img, result['masks']) #마스크 그리기
 
vis_img = draw_bboxes_list(vis_img, result['det_bboxes'], (255, 0, 0), 2, 'xyxy') #박스그리기
cv2.imwrite(save_path, vis_img) #저장
 
#모델 결과 출력
print(result)