개발 프로세스/Machine Learning

[Machine Learning] Colab에서 YOLO 테스트 하기

slowreem 2021. 4. 5. 22:41

1. YOLO weights 파일 다운로드

YOLO 공식 홈페이지 첫 화면에서 스크롤바를 내리다 보면 yolov3.weights 파일을 다운로드하는 링크와 명령어를 확인할 수 있다.

 

'here'을 클릭하거나 wget 명령어를 통해 weights 파일을 다운로드한다.

 

 

2. 구글 드라이브에 업로드

기존에 생성하였던 darknet 폴더에 weights 폴더를 새로 생성하였다.

 

 

3. Colab 가상환경으로 복사

!cp -r "/content/gdrive/MyDrive/darknet/weights" ./weights

 

darknet 폴더 위치에서 %ls 명령어를 실행하여 weights 폴더가 복사된 것을 확인하였다

 

 

4. Colab 가상환경을 위한 함수 선언

def imShow(path):
  import cv2
  import matplotlib.pyplot as plt
  %matplotlib inline

  image = cv2.imread(path)
  height, width = image.shape[:2]
  resized_image = cv2.resize(image,(3*width, 3*height), interpolation = cv2.INTER_CUBIC)

  fig = plt.gcf()
  fig.set_size_inches(18, 10)
  plt.axis("off")
  #plt.rcParams['figure.figsize'] = [10, 5]
  plt.imshow(cv2.cvtColor(resized_image, cv2.COLOR_BGR2RGB))
  plt.show()
  
  
def upload():
  from google.colab import files
  uploaded = files.upload() 
  for name, data in uploaded.items():
    with open(name, 'wb') as f:
      f.write(data)
      print ('saved file', name)
def download(path):
  from google.colab import files
  files.download(path)

 

 

 

5. Finally YOLO test !!

predictions.jpg 는 detect 명령어를 실행했을 때 생성하는 결과 이미지 파일이다.

!./darknet detect cfg/yolov3.cfg weights/yolov3.weights data/dog.jpg -dont_show
imShow('predictions.jpg')

 

test 이미지는 'data/' 항목에서 변경해볼 수 있다.