[python] Yolo v5 object detection 고추 병해 데이터 셋 학습해보기

하루하루 최종 프로젝트가 단단해지고 있다.

오늘은 기존 고추 잎으로 해본 것을 고추에 있는 탄저병을 구분해보려고 한다.

model은 yolo v5를 사용해서 학습을 실시해보자

(코랩에서는 사람들이 올려놓은 예제를 따라했을때 잘됬는데, 안되서 애좀 먹었다.

알고보니 폴더 구성을 잘못해서...ㄷㄷ;)

 

실행은 CPU환경으로 했다.

import torch
import utils
display = utils.notebook_init()  # checks

YOLOv5  v6.1-325-g3e85863 Python-3.9.13 torch-1.12.0+cpu CPU

Setup complete  (16 CPUs, 31.9 GB RAM, 295.3/466.1 GB disk)

 

PC의 Cuda 버전을 11.7 버전을 설치했는데, 텐서플로는 호환이 됬지만, 파이토치는 11.6버전까지 아직 지원이 안되는거 같다. 버전을 낮춰서 다시 실시해야 될거같다.

 

 

In [99]:

import torch
from IPython.display import Image

In [104]:

 

기존 AIHUB 데이터 셋이 구성이 되어 있었지만, 바운딩 박스가 JSON 파일로 되어있었다.

파일 이름들도 전부 일정하지 않아서 이것을 어떻게 처리해줄지 고민이 많았었다.

이미지 크기가 보통 5000X5000수준으로 되어있어서 Bounding Box도 같이 줄여주고 제대로 작동하는지 확인하는 부분이 힘들었다. 

전처리 하느라 고생해준 팀원의 노고에 감사를 보낸다 ..ㅠ 

 

# 리스트 정리
from glob import glob

train_img_list = glob('custom_data/images/train/*.jpg')
test_img_list = glob('custom_data/images/test/*.jpg')
valid_img_list = glob('custom_data/images/val/*.jpg')
print(len(train_img_list), len(test_img_list), len(valid_img_list)) #검증

 

 

2856 683 672

In [105]:

import yaml
# txt 파일로 저장한다 !
with open('C:/Users/User/Desktop/yolo_test/yolov5/custom_data/train.txt','w') as f:
  f.write('\n'.join(train_img_list) + '\n')
with open('C:/Users/User/Desktop/yolo_test/yolov5/custom_data/test.txt','w') as f:
  f.write('\n'.join(test_img_list) + '\n')
with open('C:/Users/User/Desktop/yolo_test/yolov5/custom_data/valid.txt','w') as f:
  f.write('\n'.join(valid_img_list) + '\n')

In [106]:

from IPython.core.magic import register_line_cell_magic

@register_line_cell_magic
def writetemplate(line, cell):
  with open(line, 'w') as f:
    f.write(cell.format(**globals()))

In [107]:

 

YOLO V5에서 제공하는 가중치를 설치하기위한 코드. YOLO V5 깃허브에서 따와서 사용했다.

가중치를 포함하고 안하고에 따라 성능도 많은 차이를 보인다고 한다. 학습시 기본 가중치를 꼭 넣어주고 학습하자

 

class Albumentations:
    # YOLOv5 Albumentations class (optional, used if package is installed)
    def __init__(self):
        self.transform = None
        try:
            import albumentations as A
            check_version(A.__version__, '1.0.3')  # version requirement

            self.transform = A.Compose([
                A.Blur(blur_limit=50, p=0.1),
                A.MedianBlur(blur_limit=51, p=0.1),
                A.ToGray(p=0.3)],
                bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))

            logging.info(colorstr('albumentations: ') + ', '.join(f'{x}' for x in self.transform.transforms))
        except ImportError:  # package not installed, skip
            pass
        except Exception as e:
            logging.info(colorstr('albumentations: ') + f'{e}')

    def __call__(self, im, labels, p=1.0):
        if self.transform and random.random() < p:
            new = self.transform(image=im, bboxes=labels[:, 1:], class_labels=labels[:, 0])  # transformed
            im, labels = new['image'], np.array([[c, *b] for c, b in zip(new['class_labels'], new['bboxes'])])
        return im, labels

In [108]:

import yaml

with open('data.yaml', 'r', encoding="utf-8") as f:
  data = yaml.safe_load(f)


data['train'] = 'C:\\Users\\User\\Desktop\\yolo_test\\yolov5\\custom_data\\train.txt'
data['val'] = 'C:\\Users\\User\\Desktop\\yolo_test\\yolov5\\custom_data\\valid.txt'

with open('data.yaml', 'w', encoding="utf-8") as f:
  yaml.dump(data, f)

In [109]:

print(data)

{'names': ['index0', 'index1', 'index2', 'index3', 'index4', 'gochu', 'index6', 'index7', 'index8', 'index9', 'index10', 'index11', 'index12', 'index13', 'index14', 'index15', 'index16', 'lisk1', 'lisk2'], 'nc': 19, 'train': 'C:\\Users\\User\\Desktop\\yolo_test\\yolov5\\custom_data\\train.txt', 'val': 'C:\\Users\\User\\Desktop\\yolo_test\\yolov5\\custom_data\\valid.txt'}

In [110]:

 

커널이 죽지않을까 걱정했지만, 다행히 살아서 잘돌려줬다. 에폭 10번도는데 대략 340분정도 걸렸다.

 

!python train.py --img 640 --batch 10 --epochs 10 --data data.yaml  --cfg C:/Users/User/Desktop/yolo_test/yolov5/models/yolov5s.yaml --weights yolov5s.pt --name gochu_results --cache

 

결과를 확인해보자

 

 

 

 

F1-Curve

 

 

Confusion_Matrix

 

 

2500여 장이고, 10번정도밖에 epohcs를 안돌렸는데, 생각보다 학습이 잘됬다. 

이미지 예측파일들을 보자

 

마지막으로 테스트 데이터도 예측해보자

 

!python detect.py --weights runs/train/gochu_results21/weights/best.pt --img 640 --conf 0.4 --source custom_data/images/test
# test 이미지에 대해서 디텍션을 모두 사용한다. -> 

# 다 출력할 수는 없으니, 랜덤으로 !
import glob
import random
from IPython.display import Image, display

image_name = random.choice(glob.glob('runs/detect/exp4/*.jpg'))

display(Image(filename=image_name)) # 랜덤으로 이미지가 나온다 ! 학습을 더 많이 하면 더 좋은 결과 !

TEST 데이터에서는 스레시홀드 값이 낮게 나온다 어느정도 병해를 예측하지만 뭔가 만족스럽진 않다. 

역시 데이터셋 부족이라고 생각이 든다. 데이터 정리를 다시해서 DETECTION 모델을 구성해 봐야 겠다.