머신러닝 | 쪽지시험 - YOLOv5를 이용해 이미지처리

---

1. YOLOv5를 사용하기 위해 패키지 설치 

pip install -qr https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

➡️ 패키지 설치가 완료되면 yolov5s가 설치된 것을 확인할 수 있다. 

 

 

2. 모델 생성

import ssl
ssl._create_default_https_context = ssl._create_unverified_context

import torch
import cv2
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
imgs = ['https://teamsparta.notion.site/image/https%3A%2F%2Fs3-us-west-2.amazonaws.com%2Fsecure.notion-static.com%2F396579ac-a81f-4073-9409-248ba51bc5ea%2FUntitled.jpeg?table=block&id=e70a6c23-18cd-4bc6-8c46-b585967dfd56&spaceId=83c75a39-3aba-4ba4-a792-7aefe4b07895&width=2000&userId=&cache=v2']  # batch of images

results = model(imgs)

results.save()

➡️ 사용할 이미지의 경로를 imgs에 입력해준다. 이렇게 모델 저장을 하게 되면 runs/detect라는 폴더가 생기고 이미지 안에서 모델을 탐지한 결과이미지가 저장된다.

 

 

3. 데이터 확보

import ssl
ssl._create_default_https_context = ssl._create_unverified_context
import torch
import cv2
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
imgs = ['https://teamsparta.notion.site/image/https%3A%2F%2Fs3-us-west-2.amazonaws.com%2Fsecure.notion-static.com%2F396579ac-a81f-4073-9409-248ba51bc5ea%2FUntitled.jpeg?table=block&id=e70a6c23-18cd-4bc6-8c46-b585967dfd56&spaceId=83c75a39-3aba-4ba4-a792-7aefe4b07895&width=2000&userId=&cache=v2']  # batch of images

results = model(imgs)

print(results.xyxy[0], results.xyxy[0][0][0].item())  # img1 predictions (tensor)
print(results.pandas().xyxy[0])  # img1 predictions (pandas)


tmp_img = cv2.imread('picture.jpg')
cv2.rectangle(tmp_img, (int(results.xyxy[0][0][0].item()), int(results.xyxy[0][0][1].item())), (int(results.xyxy[0][0][2].item()), int(results.xyxy[0][0][3].item())), (255,255,255))
cv2.imwrite('result1.png', tmp_img)

tmp_img = cv2.imread('picture.jpg')
cv2.rectangle(tmp_img, (int(results.xyxy[0][1][0].item()), int(results.xyxy[0][1][1].item())), (int(results.xyxy[0][1][2].item()), int(results.xyxy[0][1][3].item())), (255,255,255))
cv2.imwrite('result2.png', tmp_img)

tmp_img = cv2.imread('picture.jpg')
cv2.rectangle(tmp_img, (int(results.xyxy[0][2][0].item()), int(results.xyxy[0][2][1].item())), (int(results.xyxy[0][2][2].item()), int(results.xyxy[0][2][3].item())), (255,255,255))
cv2.imwrite('result3.png', tmp_img)

tmp_img = cv2.imread('picture.jpg')
cv2.rectangle(tmp_img, (int(results.xyxy[0][3][0].item()), int(results.xyxy[0][3][1].item())), (int(results.xyxy[0][3][2].item()), int(results.xyxy[0][3][3].item())), (255,255,255))
cv2.imwrite('result4.png', tmp_img)

tmp_img = cv2.imread('picture.jpg')
cv2.rectangle(tmp_img, (int(results.xyxy[0][4][0].item()), int(results.xyxy[0][4][1].item())), (int(results.xyxy[0][4][2].item()), int(results.xyxy[0][4][3].item())), (255,255,255))
cv2.imwrite('result5.png', tmp_img)

➡️ 위 파일을 실행해주면 사람 5명을 인식해서 사람 1명씩 사각형이 그려지게 된다. tensor에 저장된 배열은 pandas를 이용해 우리가 알아보기 쉽게 출력된다. xmin, ymin, xmax, ymax를 이용해서 사각형을 그리는 것!

 

 

 

4. 데이터 라벨링

import ssl
ssl._create_default_https_context = ssl._create_unverified_context
import torch
import cv2
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

img = cv2.imread('picture.jpg')
results = model(img)
results.save()

result = results.pandas().xyxy[0].to_numpy()
result = [item for item in result if item[6]=='person']

tmp_img = cv2.imread('picture.jpg')
print(tmp_img.shape)

for i in range(len(result)):
    cropped = tmp_img[int(result[i][1]):int(result[i][3]), int(result[i][0]):int(result[i][2])]
    print(cropped.shape)
    cv2.imwrite(f'people{i}.png', cropped)

➡️ 위 코드를 실행시키면 바로 위에 있는 사진처럼 사각형으로 그려진 부분을 따로 잘라내서 파일로 저장해준다.