YOLO-NAS 예제

Smart Parking Project 프로젝트를 진행 중
나중에 유용하다 싶은 내용을 기록합니다.

 

환경 세팅

pip install super-gradients
pip install opencv-python

 

라이브러리 및 GPU 설정 - (GPU 사용이 불가할 경우 CPU로 우회)

import cv2
import torch
from super_gradients.training import models

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
use_cuda = torch.cuda.is_available()
print(use_cuda)
if use_cuda:
  print(torch.cuda.get_device_name(0))

 
처음 위와 같이 설정을 하지 않고 NAS버전을 그냥 실행하니 아래와 같은 오류가 발생.
 

RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx

 
 
해결 방안을 찾아보니
위와 같이 작업을 gpu가 아닌 cpu로 지정하여 수행하게 해주는 코드가 있더라
 
 
 

학습 모델 지정

model1 = models.get("yolo_nas_l", pretrained_weights ="coco").to(device)
# model2 = models.get("yolo_nas_m", pretrained_weights ="coco").to(device)
# model3 = models.get("yolo_nas_s", pretrained_weights ="coco").to(device)

yolo_nas_l, yolo_nas_m, yolo_nas_s 이렇게 3가지 모델이 있다.

• yolo_nas_l 이 느리지만 가장 정확도가 좋고
• yolo_nas_s는 가장 빠르지만 l에 비해 정확도가 좋지 않다. 

 

객체 인식 및 차량 개수 파악

# 이미지 호출
img = cv2.imread("Test_video/test_photo1.jpg")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (0, 0), fx=0.5, fy=0.5)

# 지정한 모델을 이용한 검출
results = model1.predict(img, conf=0.25, fuse_model=False)

# 인식된 객체의 라밸 리스트를 가져옴
for result in results :
    labels = result.prediction.labels

# 차량의 라밸 번호는 자옹차:2, 트럭:7로 위 번호의 개수를 더함
labels = list(labels)
cnt = labels.count(2) + labels.count(7)

# 차량 개수 파악 및 시각화
print("인식된 차량의 개수: " + str(cnt))
results.show(box_thickness=2, show_confidence=True)

이미지를 불러와 지정한 모델로 객체 검출을 수행
이후 각 객체의 라벨 값을 확인하여 차량인지 판단하여 총 대수를 파악하였음.

 
yolo_nas_l 모델을 사용하면 정확도가 상당히 높다는 것을 볼 수 있다.
 
 

전체 코드

import cv2
import torch
from super_gradients.training import models

#-- GPU 설정
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
use_cuda = torch.cuda.is_available()
print(use_cuda)
if use_cuda:
  print(torch.cuda.get_device_name(0))

#-- 사전학습된 Yolo_nas_small 모델 불러오기(빠르지만 정확도가 낮음)
model1 = models.get("yolo_nas_l", pretrained_weights ="coco").to(device)
# model2 = models.get("yolo_nas_m", pretrained_weights ="coco").to(device)
# model3 = models.get("yolo_nas_s", pretrained_weights ="coco").to(device)

# 이미지 호출
img = cv2.imread("Test_video/test_photo1.jpg")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (0, 0), fx=0.5, fy=0.5)

# 지정한 모델을 이용한 검출
results = model1.predict(img, conf=0.25, fuse_model=False)

# 인식된 객체의 라밸 리스트를 가져옴
for result in results :
    labels = result.prediction.labels

# 차량의 라밸 번호는 자옹차:2, 트럭:7로 위 번호의 개수를 더함
labels = list(labels)
cnt = labels.count(2) + labels.count(7)

# 차량 개수 파악 및 시각화
print("인식된 차량의 개수: " + str(cnt))
results.show(box_thickness=2, show_confidence=True)