[OpenCV] 코랩으로 이미지 Crop(자르기), Mask(가리기)하기

정기용·2023년 1월 20일
1

lunalabs

목록 보기
2/5
post-thumbnail

#목표

  • 라벨링한 이미지 1000장으로 Crop과 Mask를 한다.
  • 클래스 0은 차량 번호판으로, Mask할 것이다.
  • 클래스 1, 2는 각각 차량 앞면부와 차량 전체로, Crop할 것이다.

환경은 코랩으로 해보았다.

이미지 업로드와 다운로드하는데 시간이 오래걸리니 로컬(주피터 노트북 등등)에서 돌리는게 낫다.

나는 로컬에 cv2 설정하다가 안돼서 그냥 코랩에 돌렸다.
코랩에서는 imshow()도 안되니 더더욱 로컬에서 돌리자.

1. 구글 드라이브 연결하기

구글 드라이브에 이미지 압축파일을 업로드 한다.
압축 안해서 올리면 엄청 오래걸린다. 압축파일을 올린 후 밑에서 코드로 압축을 해제해준다.

그리고 코랩에 연결해준다. (버튼 클릭 후 생성된 코드 실행)

2. 이미지 리스트 불러오기 & 정렬

import os

base_path = '/content/drive/MyDrive/yolov5/data/car+plate+carfront/'
img_path = base_path + 'images'
label_path = base_path + 'labels'

img_file_list = os.listdir(img_path)
label_file_list = os.listdir(label_path)

나는 하나의 디렉토리에 images와 labels를 분리해서 넣었다.
os.listdir로 각각 폴더 위치를 지정해 폴더의 하위 파일 이름을 리스트로 가져온다.

img_file_list.sort()
label_file_list.sort()

print('img len: {}, label len: {}\n'.format(len(img_file_list), len(label_file_list)))
print(img_file_list[:100])
print(label_file_list[:100])

그리고 리스트를 정렬해준다.

정렬 안하면 밑에서 이미지와 라벨 짝이 안맞아 이상 부분을 crop & mask한다.

2_1. 검수코드

## 이미지 파일과 라벨 파일 이름 똑같음
cnt=0

for a, b in zip (img_file_list, label_file_list):
  a = a.rstrip('jpgtxt')
  a = a.rstrip('.')
  b = b.rstrip('jpgtxt')
  b = b.rstrip('.')
  if a != b: print(cnt)
  cnt+=1
  • 이미지와 라벨 짝이 맞는지 검사하는 코드다.

아무것도 출력이 안되면 잘 짝이 맞는 것이다.
짝이 안맞을 경우 해당 인덱스를 출력한다.

3. 이미지 좌표 구하기

import cv2

cnt = 0

for img_file, label_file in zip(img_file_list, label_file_list):

    cnt += 1
    print(cnt)


    img = cv2.imread(img_path + '/' + img_file)

    with open(label_path + '/' + label_file, 'r') as f:
        labels = f.readlines()
        car_whole = list(map(float, labels[2].split()[1:]))  ## 차 전체 이미지 바운딩박스
        car_front = list(map(float, labels[1].split()[1:]))  ## 차 앞부분 이미지 바운딩박스
        carplate_masking = list(map(float, labels[0].split()[1:]))  ## 차 앞부분 이미지 바운딩박스
        h, w, _ = img.shape  ## 사진 전체 크기 (heigt, width)

우선 for문에서 이미지 파일과 라벨 파일을 정렬한 순서대로 하나씩 선택한다.
cv2.imread로 img에 이미지를 저장한다.
with open~ 으로 f에 라벨을 저장한다.

## 라벨 예시
0 0.281785 0.303789 0.193897 0.092345
1 0.425525 0.287692 0.581912 0.378795
2 0.455954 0.285394 0.730389 0.509354
  • 라벨 설명
    • 라벨 클래스는 0(자동차 번호판), 1(자동차 앞면 그릴), 2(자동차 전체)이다.
    • 뒤에 소수점으로 표시된 좌표는 각각 (바운딩박스 중점 x좌표, 바운딩박스 중점 y좌표, 바운딩박스 가로 넓이, 바운딩박스 세로 넓이) 이다.
    • 전부 소수점으로, 사진 전체 중 비율을 기준으로 라벨링이 되어있다.

각각 list로 변수에 좌표를 저장한다.

3_1. 좌표 진짜로 구하기

원하는 부분을 사각형모양으로 crop하거나 mask하려면 왼쪽 위 꼭짓점 좌표와 오른쪽 아래 꼭짓점 좌표를 알아야한다.

다음은 이를 계산한 좌표다.

starty = car_whole[1] - (car_whole[3] / 2)
        startx = car_whole[0] - (car_whole[2] / 2)
        endy = car_whole[1] + (car_whole[3] / 2)
        endx = car_whole[0] + (car_whole[2] / 2)

그리고 비율이 0과 1을 넘지 않도록 다음 함수를 통과시켜준다.

def size_control(sx, sy, ex, ey):
  if sx < 0: sx = 0
  if sy < 0: sy = 0
  if ex > 1: ex = 1
  if ey > 1: ey = 1
  return sx, sy, ex, ey

합치면 다음과 같다.

def size_control(sx, sy, ex, ey):
  if sx < 0: sx = 0
  if sy < 0: sy = 0
  if ex > 1: ex = 1
  if ey > 1: ey = 1
  return sx, sy, ex, ey


starty = car_whole[1] - (car_whole[3] / 2)
        startx = car_whole[0] - (car_whole[2] / 2)
        endy = car_whole[1] + (car_whole[3] / 2)
        endx = car_whole[0] + (car_whole[2] / 2)
        
        startx, starty, endx, endy = size_control(startx, starty, endx, endy)

4. 이미지 Crop하기 & 저장

이미지 크롭은 슬라이싱으로 간단하게 할 수 있다.
이미지는 numpy array로 되어있어 슬라이싱이 가능하다.


        whole_crop_img = img[int(starty * h):int(endy * h), int(startx * w):int(endx * w)]

int 안붙이면 슬라이싱이 안먹힌다.
라벨링이 비율로 되어있으므로 h와 w를 곱해준다.

cv2.imwrite('/content/drive/MyDrive/crop_img/car_whole/' + 'whole_crop_' + img_file, whole_crop_img)

그리고 저장해준다.

5. 이미지 mask하기

이미지 마스크는 원하는 영역에 속이 꽉찬 사각형을 그려서 할 수 있다.

masking_img = img
        cv2.rectangle(masking_img, (int(startx * w), int(starty * h)), (int(endx * w), int(endy * h)),(255,255,255), -1)

뒤에 (255,255,255)는 하얀색을 가리킨다. (0,0,0)으로 하면 까만색 사각형을 그린다.

cv2.imwrite('/content/drive/MyDrive/crop_img/carplate_masking/' + 'carplate_masking_' + img_file, masking_img)

그리고 저장

crop & mask 전체 코드

import cv2

cnt = 0

for img_file, label_file in zip(img_file_list, label_file_list):

    cnt += 1
    print(cnt)


    img = cv2.imread(img_path + '/' + img_file)

    with open(label_path + '/' + label_file, 'r') as f:
        labels = f.readlines()
        car_whole = list(map(float, labels[2].split()[1:]))  ## 차 전체 이미지 바운딩박스
        car_front = list(map(float, labels[1].split()[1:]))  ## 차 앞부분 이미지 바운딩박스
        carplate_masking = list(map(float, labels[0].split()[1:]))  ## 차 앞부분 이미지 바운딩박스
        h, w, _ = img.shape  ## 사진 전체 크기 (heigt, width)

        ### car_whole
        starty = car_whole[1] - (car_whole[3] / 2)
        startx = car_whole[0] - (car_whole[2] / 2)
        endy = car_whole[1] + (car_whole[3] / 2)
        endx = car_whole[0] + (car_whole[2] / 2)
        startx, starty, endx, endy = size_control(startx, starty, endx, endy)


        whole_crop_img = img[int(starty * h):int(endy * h), int(startx * w):int(endx * w)]
        try:
          cv2.imwrite('/content/drive/MyDrive/crop_img/car_whole/' + 'whole_crop_' + img_file, whole_crop_img)
        except:
          pass


        ### car_front
        starty = car_front[1] - (car_front[3] / 2)
        startx = car_front[0] - (car_front[2] / 2)
        endy = car_front[1] + (car_front[3] / 2)
        endx = car_front[0] + (car_front[2] / 2)
        startx, starty, endx, endy = size_control(startx, starty, endx, endy)

        carfront_crop_img = img[int(starty * h):int(endy * h), int(startx * w):int(endx * w)]
        try:
          cv2.imwrite('/content/drive/MyDrive/crop_img/car_front/' + 'carfront_crop_' + img_file, carfront_crop_img)
        except:
          pass
        

        ## carplate_masking
        starty = carplate_masking[1] - (carplate_masking[3] / 2)
        startx = carplate_masking[0] - (carplate_masking[2] / 2)
        endy = carplate_masking[1] + (carplate_masking[3] / 2)
        endx = carplate_masking[0] + (carplate_masking[2] / 2)

        masking_img = img
        cv2.rectangle(masking_img, (int(startx * w), int(starty * h)), (int(endx * w), int(endy * h)),(255,255,255), -1)
        cv2.imwrite('/content/drive/MyDrive/crop_img/carplate_masking/' + 'carplate_masking_' + img_file, masking_img)

        f.close()

욜로 사용해보는 post는 다음 포스트에서

profile
SKKU CS 23

0개의 댓글