[시각 지능] Roboflow

Roboflow

 

https://docs.roboflow.com/

 

Roboflow는 개발자가 기술이나 경험에 관계없이 자신의 컴퓨터 비전 애플리케이션을 구축할 수 있도록 지원한다. 원시 이미지를 훈련된 맞춤형 컴퓨터 비전 모델로 변환하고 애플리케이션에서 사용하기 위해 배포하는 데 필요한 모든 도구를 제공한다. 현재 Roboflow는 객체 감지 및 분류 모델을 지원한다.

 

https://public.roboflow.com/object-detection/mask-wearing/4에서 포맷을 YOLO v5 Pytorch로 다운로드하고 파일 이름을 Mask_Data.zip으로 변경한다.

 

YOLOv5 설치

 

# Colab root dir
ROOT_DIR = '/content'

import os

YOLOv5_ROOT_DIR = os.path.join(ROOT_DIR, 'yolov5')

DATA_ROOT_DIR = os.path.join(ROOT_DIR, 'Mask_Data')  # Custom Data

print(YOLOv5_ROOT_DIR, DATA_ROOT_DIR, type(YOLOv5_ROOT_DIR), type(DATA_ROOT_DIR))

%rm -rf {YOLOv5_ROOT_DIR}

%cd {ROOT_DIR}

# Git clone
!git clone https://github.com/ultralytics/yolov5.git

# 필수 라이브러리 설치
!pip install -r {os.path.join(YOLOv5_ROOT_DIR, 'requirements.txt')}

# yolo5s.pt download
!wget -P {YOLOv5_ROOT_DIR}  https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5s.pt

 

데이터 전처리

 

import shutil

if os.path.exists(DATA_ROOT_DIR):
    shutil.rmtree(DATA_ROOT_DIR)
    print(DATA_ROOT_DIR + ' is removed.')

import zipfile

with zipfile.ZipFile(os.path.join(ROOT_DIR, 'Mask_Data.zip'), 'r') as target_fle:
    target_fle.extractall(DATA_ROOT_DIR)

TRAIN_IMAGE_DATA_DIR = os.path.join(DATA_ROOT_DIR, 'train/images/')
TRAIN_LABEL_DATA_DIR = os.path.join(DATA_ROOT_DIR, 'train/labels/')

NEW_TRAIN_IMAGE_DATA_DIR = os.path.join(DATA_ROOT_DIR, 'new_train/images/')
NEW_TRAIN_LABEL_DATA_DIR = os.path.join(DATA_ROOT_DIR, 'new_train/labels/')

VALID_IMAGE_DATA_DIR = os.path.join(DATA_ROOT_DIR, 'valid/images/')
VALID_LABEL_DATA_DIR = os.path.join(DATA_ROOT_DIR, 'valid/labels/')

NEW_VALID_IMAGE_DATA_DIR = os.path.join(DATA_ROOT_DIR, 'new_valid/images/')
NEW_VALID_LABEL_DATA_DIR = os.path.join(DATA_ROOT_DIR, 'new_valid/labels/')

# 임의의 train 이미지 파일 2개, 임의의 train 정답 파일 3개 삭제
import random
import glob

train_image_file_path_list = glob.glob(TRAIN_IMAGE_DATA_DIR+'*')
train_label_file_path_list = glob.glob(TRAIN_LABEL_DATA_DIR+'*')

print('[before os.remove] nums of image file path = ', len(train_image_file_path_list))
print('[before os.remove] nums of label file path = ', len(train_label_file_path_list))

random.shuffle(train_image_file_path_list)  # image shuffle
random.shuffle(train_label_file_path_list)  # label shuffle

# remove train 2 images
removed_image_nums = 2

for removed_image_file in train_image_file_path_list[:removed_image_nums]:
    print(removed_image_file+ ' is removing...')
    os.remove(removed_image_file)

# remove train 3 labels
removed_label_nums = 3

for removed_label_file in train_label_file_path_list[:removed_label_nums]:
    print(removed_label_file+' is removing...')
    os.remove(removed_label_file)

train_image_file_path_list = glob.glob(TRAIN_IMAGE_DATA_DIR+'*')
train_label_file_path_list = glob.glob(TRAIN_LABEL_DATA_DIR+'*')

print('[after os.remove] nums of image file path = ', len(train_image_file_path_list))
print('[after os.remove] nums of label file path = ', len(train_label_file_path_list))

# TRAIN_IMAGE_DATA_DIR, TRAIN_LABEL_DATA_DIR 파일이름 비교 후 train -> new_train 복사
train_image_file_list = []
train_label_file_list = []

train_image_file_name_list = []
train_label_file_name_list = []

train_image_file_name_set = set()
train_label_file_name_set = set()

train_image_label_difference_set = set()
train_image_label_intersection_set = set()

# train data
print('====================================================================')
print('[train] compare image file nums with label file nums')
print('====================================================================')

train_image_file_list = os.listdir(TRAIN_IMAGE_DATA_DIR)
train_label_file_list = os.listdir(TRAIN_LABEL_DATA_DIR)

train_image_file_list.sort()
train_label_file_list.sort()

# image file name 추출
for index in range(len(train_image_file_list)):
    image_name = train_image_file_list[index].split('.jpg')[0] # image file name
    train_image_file_name_list.append(image_name)
    
train_image_file_name_set = set(train_image_file_name_list)  # 이미지 파일 이름 집합생성

# label file name 추출
for index in range(len(train_label_file_list)):
    label_name = train_label_file_list[index].split('.txt')[0]  # label file name
    train_label_file_name_list.append(label_name)

train_label_file_name_set = set(train_label_file_name_list)  # 정답 파일 이름 집합생성

# image file / label file name 비교
train_image_label_difference_set = train_image_file_name_set - train_label_file_name_set  # 차집합

train_image_label_intersection_set = train_image_file_name_set & train_label_file_name_set  # 교집합

# image file / label file 개수가 같은 경우
if len(train_image_label_difference_set) == 0:
    print('>>> image is exactly same as label  <<<')

# image file / label file 개수가 다른 경우   
else:
    print('--------------------------------------------------------------------')
    print('[train] image file != label file')
    print('--------------------------------------------------------------------')

    # image file name 추출
    for index in range(len(train_image_file_list)):
        image_name = train_image_file_list[index].split('.jpg')[0] # image file name
        train_image_file_name_list.append(image_name)
    
    train_image_file_name_set = set(train_image_file_name_list)  # 이미지 파일 이름 집합생성

    # label file name 추출
    for index in range(len(train_label_file_list)):
        label_name = train_label_file_list[index].split('.txt')[0]  # label file name
        train_label_file_name_list.append(label_name)

    train_label_file_name_set = set(train_label_file_name_list)  # 정답 파일 이름 집합생성

    # new_train/images, new_train/labels 디렉토리 생성
    if not os.path.exists(os.path.join(DATA_ROOT_DIR, 'new_train')):
        os.mkdir(os.path.join(DATA_ROOT_DIR, 'new_train'))
        print(os.path.join(DATA_ROOT_DIR, 'new_train') + ' is created!')

    if not os.path.exists(NEW_TRAIN_IMAGE_DATA_DIR):
        os.mkdir(NEW_TRAIN_IMAGE_DATA_DIR)
        print(NEW_TRAIN_IMAGE_DATA_DIR + ' is created!')

    if not os.path.exists(NEW_TRAIN_LABEL_DATA_DIR):
        os.mkdir(NEW_TRAIN_LABEL_DATA_DIR)
        print(NEW_TRAIN_LABEL_DATA_DIR + ' is created!')

    # image / label 교집합 생성
    train_image_label_intersection_set = train_image_file_name_set & train_label_file_name_set  # 교집합

    print('nums of image_label intersection set => ', len(train_image_label_intersection_set))

    # src / dst path 설정
    src_image_path = TRAIN_IMAGE_DATA_DIR
    dst_image_path = NEW_TRAIN_IMAGE_DATA_DIR

    src_label_path = TRAIN_LABEL_DATA_DIR
    dst_label_path = NEW_TRAIN_LABEL_DATA_DIR

    # src -> dst 교집합 파일 복사
    copy_nums = 0

    for file_name in train_image_label_intersection_set:
        # image file copy
        shutil.copy(os.path.join(src_image_path, file_name+'.jpg'),
                    os.path.join(dst_image_path, file_name+'.jpg'))
        
        # label file copy
        shutil.copy(os.path.join(src_label_path, file_name+'.txt'),
                    os.path.join(dst_label_path, file_name+'.txt'))

        copy_nums += 1

    print('total copy nums (src->dst) = ', copy_nums)

# 임의의 valid 이미지 파일 3개, 임의의 train 정답 파일 1개 삭제
valid_image_file_path_list = glob.glob(VALID_IMAGE_DATA_DIR+'*')
valid_label_file_path_list = glob.glob(VALID_LABEL_DATA_DIR+'*')

print('[before os.remove] nums of image file path = ', len(valid_image_file_path_list))
print('[before os.remove] nums of label file path = ', len(valid_label_file_path_list))

random.shuffle(valid_image_file_path_list)  # image shuffle
random.shuffle(valid_label_file_path_list)  # label shuffle

# remove valid 3 images
removed_image_nums = 3

for removed_image_file in valid_image_file_path_list[:removed_image_nums]:

    print(removed_image_file+ ' is removing...')
    os.remove(removed_image_file)

# remove train 1 labels
removed_label_nums = 1

for removed_label_file in valid_label_file_path_list[:removed_label_nums]:
    
    print(removed_label_file+' is removing...')
    os.remove(removed_label_file)

valid_image_file_path_list = glob.glob(VALID_IMAGE_DATA_DIR+'*')
valid_label_file_path_list = glob.glob(VALID_LABEL_DATA_DIR+'*')

print('[after os.remove] nums of image file path = ', len(valid_image_file_path_list))
print('[after os.remove] nums of label file path = ', len(valid_label_file_path_list))

# VALID_IMAGE_DATA_DIR, VALID_LABEL_DATA_DIR 파일이름 비교 후 valid -> new_valid 복사
valid_image_file_list = []
valid_label_file_list = []

valid_image_file_name_list = []
valid_label_file_name_list = []

valid_image_file_name_set = set()
valid_label_file_name_set = set()

valid_image_label_difference_set = set()
valid_image_label_intersection_set = set()

# valid data

print('====================================================================')
print('[valid] compare image file nums with label file nums')
print('====================================================================')
valid_image_file_list = os.listdir(VALID_IMAGE_DATA_DIR)
valid_label_file_list = os.listdir(VALID_LABEL_DATA_DIR)

valid_image_file_list.sort()
valid_label_file_list.sort()

# image file name 추출
for index in range(len(valid_image_file_list)):
    image_name = valid_image_file_list[index].split('.jpg')[0] # image file name
    valid_image_file_name_list.append(image_name)
    
valid_image_file_name_set = set(valid_image_file_name_list)  # 이미지 파일 이름 집합생성

# label file name 추출
for index in range(len(valid_label_file_list)):
    label_name = valid_label_file_list[index].split('.txt')[0]  # label file name
    valid_label_file_name_list.append(label_name)

valid_label_file_name_set = set(valid_label_file_name_list)  # 정답 파일 이름 집합생성

# image file / label file name 비교
valid_image_label_difference_set = valid_image_file_name_set - valid_label_file_name_set  # 차집합

valid_image_label_intersection_set = valid_image_file_name_set & valid_label_file_name_set  # 교집합

# image file / label file 개수가 같은 경우
if len(valid_image_label_difference_set) == 0:
    print('>>> image is exactly same as label  <<<')

# image file / label file 개수가 다른 경우   
else:
    print('--------------------------------------------------------------------')
    print('[valid] image file != label file')
    print('--------------------------------------------------------------------')

    # image file name 추출
    for index in range(len(valid_image_file_list)):
        image_name = valid_image_file_list[index].split('.jpg')[0] # image file name
        valid_image_file_name_list.append(image_name)
    
    valid_image_file_name_set = set(valid_image_file_name_list)  # 이미지 파일 이름 집합생성

    # label file name 추출
    for index in range(len(valid_label_file_list)):
        label_name = valid_label_file_list[index].split('.txt')[0]  # label file name
        valid_label_file_name_list.append(label_name)

    valid_label_file_name_set = set(valid_label_file_name_list)  # 정답 파일 이름 집합생성

    # new_valid/images, new_valid/labels 디렉토리 생성
    if not os.path.exists(os.path.join(DATA_ROOT_DIR, 'new_valid')):
        os.mkdir(os.path.join(DATA_ROOT_DIR, 'new_valid'))
        print(os.path.join(DATA_ROOT_DIR, 'new_valid') + ' is created !!!')

    if not os.path.exists(NEW_VALID_IMAGE_DATA_DIR):
        os.mkdir(NEW_VALID_IMAGE_DATA_DIR)
        print(NEW_VALID_IMAGE_DATA_DIR + ' is created !!!')

    if not os.path.exists(NEW_VALID_LABEL_DATA_DIR):
        os.mkdir(NEW_VALID_LABEL_DATA_DIR)
        print(NEW_VALID_LABEL_DATA_DIR + ' is created !!!')

    # image / label 교집합 생성
    valid_image_label_intersection_set = valid_image_file_name_set & valid_label_file_name_set  # 교집합

    print('nums of image_label intersection set => ', len(valid_image_label_intersection_set))

    # src / dst path 설정
    src_image_path = VALID_IMAGE_DATA_DIR
    dst_image_path = NEW_VALID_IMAGE_DATA_DIR

    src_label_path = VALID_LABEL_DATA_DIR
    dst_label_path = NEW_VALID_LABEL_DATA_DIR

    # src -> dst 교집합 파일 복사
    copy_nums = 0

    for file_name in valid_image_label_intersection_set:

        # image file copy
        shutil.copy(os.path.join(src_image_path, file_name+'.jpg'),
                    os.path.join(dst_image_path, file_name+'.jpg'))
        
        # label file copy
        shutil.copy(os.path.join(src_label_path, file_name+'.txt'),
                    os.path.join(dst_label_path, file_name+'.txt'))

        copy_nums += 1

    print('total copy nums (src->dst) = ', copy_nums)

# new_train/images 파일과 new_train/labels 파일 비교
# new_train data
print('====================================================================')
print('[new_train] compare image file nums with label file nums')
print('====================================================================')
new_train_image_file_list = os.listdir(NEW_TRAIN_IMAGE_DATA_DIR)
new_train_label_file_list = os.listdir(NEW_TRAIN_LABEL_DATA_DIR)

new_train_image_file_list.sort()
new_train_label_file_list.sort()

print('new image nums = ', len(new_train_image_file_list), ' , new label nums = ',len(new_train_label_file_list))
print(new_train_image_file_list[:2], new_train_image_file_list[-2:])
print(new_train_label_file_list[:2], new_train_label_file_list[-2:])

# new_valid/images 파일과 new_valid/labels 파일 비교
# new_valid data
print('====================================================================')
print('[new_valid] compare image file nums with label file nums')
print('====================================================================')
new_valid_image_file_list = os.listdir(NEW_VALID_IMAGE_DATA_DIR)
new_valid_label_file_list = os.listdir(NEW_VALID_LABEL_DATA_DIR)

new_valid_image_file_list.sort()
new_valid_label_file_list.sort()

print('new image nums = ', len(new_valid_image_file_list), ' , new label nums = ',len(new_valid_label_file_list))
print(new_valid_image_file_list[:2], new_valid_image_file_list[-2:])
print(new_valid_label_file_list[:2], new_valid_label_file_list[-2:])

 

yaml 파일 설정 (데이터셋 위치 알려주는 config file)

 

%cat {os.path.join(DATA_ROOT_DIR, 'data.yaml')}

import yaml

with open(os.path.join(DATA_ROOT_DIR, 'data.yaml'), 'r') as f:
    data = yaml.safe_load(f)

print(data)

data['train'] = os.path.join(DATA_ROOT_DIR, 'new_train/images/')  # train -> new_train
data['val'] = os.path.join(DATA_ROOT_DIR, 'new_valid/images/')    # valid -> new_valid

data['nc'] = 2
data['names'] = ['mask','no-mask']

with open(os.path.join(DATA_ROOT_DIR, 'data.yaml'), 'w') as f:
  yaml.dump(data, f)

print(data)

 

하이퍼 파라미터 설정

 

img_size = 416
bacth_size = 32
epochs = 100

data_path = os.path.join(DATA_ROOT_DIR, 'data.yaml')
yalm_path = os.path.join(YOLOv5_ROOT_DIR, 'models/yolov5s.yaml')
weights_path = os.path.join(YOLOv5_ROOT_DIR, 'yolov5s.pt')

# train
!python3 {os.path.join(YOLOv5_ROOT_DIR, 'train.py')} --img {img_size} --batch {bacth_size} --epochs {epochs} --data {data_path} \
--cfg {yalm_path} --weights {weights_path}

 

텐서보드 출력

 

%load_ext tensorboard

%tensorboard --logdir {os.path.join(YOLOv5_ROOT_DIR, 'runs/train/exp/')}

 

테스트 이미지

 

test_image_list = glob(os.path.join(DATA_ROOT_DIR, 'test/images/*.jpg'))

print(len(test_image_list))

test_image_list.sort()

for i in range(len(test_image_list)):
    print('i = ',i, test_image_list[i])

 

추론

 

# inference (> 0.4)
test_data_path = os.path.join(DATA_ROOT_DIR, 'test/images/')
best_weights_path = os.path.join(YOLOv5_ROOT_DIR, 'runs/train/exp/weights/best.pt')

!python3 {os.path.join(YOLOv5_ROOT_DIR, 'detect.py')} --weights {best_weights_path} --img 416 --conf 0.4 --source {test_data_path}

 

결과 확인

 

detetced_image_list = glob.glob((os.path.join(YOLOv5_ROOT_DIR, 'runs/detect/exp/*.jpg')))

detected_image_nums = len(detetced_image_list)

print(detected_image_nums)

print(detetced_image_list)

# 다운로드를 위한 inference image 압축
import zipfile
import os

if not os.path.exists(os.path.join(ROOT_DIR, 'detected_result')):
    os.mkdir(os.path.join(ROOT_DIR, 'detected_result'))
    print('detected_result dir is created !!!')

with zipfile.ZipFile(os.path.join(ROOT_DIR, 'detected_result/Mask_detected_images.zip'), 'w') as detected_images:
    for idx in range(detected_image_nums):
        detected_images.write(detetced_image_list[idx])

# 파일 다운로드
import shutil
from google.colab import files

shutil.copy(os.path.join(YOLOv5_ROOT_DIR, 'runs/train/exp/weights/best.pt'), os.path.join(YOLOv5_ROOT_DIR, 'Mask_best.pt'))

files.download(os.path.join(ROOT_DIR, 'detected_result/Mask_detected_images.zip'))

files.download(os.path.join(YOLOv5_ROOT_DIR, 'Mask_best.pt'))

from IPython.display import Image

val_img_path = img_list[1]

weights_path = os.path.join(YOLOv5_ROOT_DIR, 'runs/train/exp/weights/best.pt')

!python detect.py --weights "{weights_path}" --img 416 --conf 0.5 --source "{val_img_path}"

detect_img_path = os.path.join(YOLOv5_ROOT_DIR, 'runs/detect/exp')

Image(os.path.join(detect_img_path, os.path.basename(val_img_path)))