[코드] 이미지 상자 치기 코드draw_boxs_in_images

코드 검색을 위해, 다음을 이용하자. 

케글 - 이미 개발된 코드가 있을 수 있고, discussion에 data augmentation에 대한 정보도 많다

깃 - 패키지 안에, 검색을 함으로써 이미 만들어진 코드를 검색할 수 있다.

구글 - 어떤 코드, 함수를 찾으려면 위의 방법을 이용하는게 낫다. [~하는 코드]는 구글에서 찾기 힘들다

---

draw_boxs_in_images

<실행 방법>

python drawRbox_in_Images.py --image_path=./images --save_path=./images_with_boxs --csv_path=./baseline.csv

---

<위 실행을 위한 data structure>

• drawRbox_in_Images.py
• images
  • 0.png
  • 1.png
  • 2.png
• images_with_box
  • (empty)
• baseline.csv

---

<세부 사항>

1. images에 있는 사진만 박스 친다.
2. test사진 1240장 모두가 꼭 있어야 하는 것은 아니다.

from matplotlib import pyplot as plt
import numpy as np
import argparse
from PIL import Image
from matplotlib import pyplot as plt
import cv2
import pandas as pd
from PIL import ImageFont
from PIL import ImageDraw



def save_image_with_box(box_info, imageName , start_index, last_index , 
                        image_path, save_path):
    """
    box_info : panda로 받아드린 csv파일
    image_Name : image_Num.png 파일을 봐야한다. 
    start, last : box_info['file_name'][start ~ last]가  image_Num.png에 대한 정보를 담고 있다. 
    img_path : 이미지가 있는 directory name
    save_path : 이미지가 저장될 directory name
    """
    try:
        im = Image.open(image_path + imageName)

    except :
        print("no such file in directory : ", imageName)
        return
    plt.figure(figsize = (30,30))
    plt.imshow(im)
    color_set = ['r','b','y','g']
    for i in range(start_index, last_index+1):
        point1 = (box_info["point1_x"][i],box_info["point1_y"][i])
        point2 = (box_info["point2_x"][i],box_info["point2_y"][i])
        point3 = (box_info["point3_x"][i],box_info["point3_y"][i])
        point4 = (box_info["point4_x"][i],box_info["point4_y"][i])
        plt.plot([point1[0],point2[0]],[point1[1],point2[1]], linewidth=3, color = color_set[box_info['class_id'][i] - 1])
        plt.plot([point2[0],point3[0]],[point2[1],point3[1]], linewidth=3, color = color_set[box_info['class_id'][i] - 1])
        plt.plot([point3[0],point4[0]],[point3[1],point4[1]], linewidth=3, color = color_set[box_info['class_id'][i] - 1])
        plt.plot([point4[0],point1[0]],[point4[1],point1[1]], linewidth=3, color = color_set[box_info['class_id'][i] - 1])

    plt.savefig(save_path  + imageName)
    plt.close()
    print("saved : ", imageName)


if __name__ == "__main__":
    # 기본 설정
    # image_path = './Alpha-project/images/'
    # save_path = './Alpha-project/images_with_boxs/'
    # csv_path = './Alpha-project/baseline.csv'
    parser = argparse.ArgumentParser(description='draw_rbox_in_images')
    parser.add_argument('--image_path', type=str, default='images')
    parser.add_argument('--save_path', type=str, default='images_with_boxs')
    parser.add_argument('--csv_path', type=str, default='baseline.csv')

    args = parser.parse_args()
    image_path = args.image_path + '/'  # imagepath == images/
    save_path = args.save_path + '/'    # save_path == images_with_boxs/
    csv_path = args.csv_path

    #csv file load
    box_info = pd.read_csv(csv_path)

    #run
    start_index = 0
    for i in range(len(box_info)+1):
        try:
            # i+1번째 파일이 다른 이미지라면, i번째 파일에 대해서 박스가 처진 그림을 그린다. 
            if box_info['file_name'][i][0:-4] != box_info['file_name'][i+1][0:-4]:  
                save_image_with_box(box_info, box_info['file_name'][i] , start_index, i,
                                    image_path ,save_path )
                start_index = i+1
        except:
            # box_info['file_name'][i+1]가 존재하지 않으면 , 즉 999.png 이후를 바라본다면
            save_image_with_box(box_info,box_info['file_name'][i],start_index,i,
                                image_path ,save_path )