machine learning

Deep Learning - 4

안범진 2020. 12. 10. 13:53

Mask-RCNN custom dataset 돌리기 위한 준비

requirements.txt에 있는 것을 모두 설치해준다음에

tensorflow-gpu==1.15.0

gpu를 사용하기 위해 위와 같이 입력.

keras == 2.0.8, scikit-image = 0.16.2 로 설치해주어야 정상 훈련이 가능하다.

# Root directory of the project

ROOT_DIR = os.path.abspath("/home/ubuntu/..../.../...")

프로젝트의 ROOT_DIR을 자신의 경로와 맞게 설정한다.

그 다음에 configuration을 지정해준다. 여기서 Name은 나중에 .h5 로그 파일이 지정되는 폴더의 이름의 일부가 된다.

############################################################
#  Configurations
############################################################


class CustomConfig(Config):
    """Configuration for training on the toy  dataset.
    Derives from the base Config class and overrides some values.
    """
    # Give the configuration a recognizable name
    NAME = "object"

    # We use a GPU with 12GB memory, which can fit two images.
    # Adjust down if you use a smaller GPU.
    IMAGES_PER_GPU = 2

    # Number of classes (including background)
    NUM_CLASSES = 1 + 3  # Background + toy

    # Number of training steps per epoch
    STEPS_PER_EPOCH = 100

    # Skip detections with < 90% confidence
    DETECTION_MIN_CONFIDENCE = 0.9

반드시 via_region_data.json파일에 있는 regional_attributes가 올바른지 확인할 것

아래에서 self.add_class의 Dent, Scratch, Others가 실제로 위의 json파일의 regional attributes에 없을 경우 dimension 관련 에러가 발생하니 꼭 유의 할 것

제대로 지정하지 않은 경우 아래와 같은 에러가 발생한다.

def_load_custom만 잘 설정해주면 나머지는 알아서 된다.

############################################################
#  Dataset
############################################################

class CustomDataset(utils.Dataset):

    def load_custom(self, dataset_dir, subset):
        """Load a subset of the bottle dataset.
        dataset_dir: Root directory of the dataset.
        subset: Subset to load: train or val
        """
        # Add classes. We have only one class to add.
        self.add_class("object", 1, "Dent")
        self.add_class("object", 2, "Scratch")
        self.add_class("object", 3, "Others")
        

        # Train or validation dataset?
        assert subset in ["train", "val"]
        dataset_dir = os.path.join(dataset_dir, subset)

        # Load annotations
        # VGG Image Annotator saves each image in the form:
        # { 'filename': '28503151_5b5b7ec140_b.jpg',
        #   'regions': {
        #       '0': {
        #           'region_attributes': {},
        #           'shape_attributes': {
        #               'all_points_x': [...],
        #               'all_points_y': [...],
        #               'name': 'polygon'}},
        #       ... more regions ...
        #   },
        #   'size': 100202
        # }
        # We mostly care about the x and y coordinates of each region
        annotations = json.load(open(os.path.join(dataset_dir, "via_region_data.json")))
        annotations = list(annotations.values())  # don't need the dict keys


        annotations = [a for a in annotations if a['regions']]
        
        # Add images
        for a in annotations:
            # Get the x, y coordinaets of points of the polygons that make up
            # the outline of each object instance. There are stores in the
            # shape_attributes (see json format above)
            # for b in a['regions'].values():
            #    polygons = [{**b['shape_attributes'], **b['region_attributes']}]
            # print("string=", polygons)
            # for r in a['regions'].values():
            #    polygons = [r['shape_attributes']]
            #    # print("polygons=", polygons)
            #    multi_numbers = [r['region_attributes']]
                # print("multi_numbers=", multi_numbers)
            polygons = [r['shape_attributes'] for r in a['regions'].values()]
            objects = [s['region_attributes'] for s in a['regions'].values()]
            # print("multi_numbers=", multi_numbers)
            # num_ids = [n for n in multi_numbers['number'].values()]
            # for n in multi_numbers:
            num_ids = []
            for n in objects:
        	#        print(n)
        	#print(type(n))
                try:
                    print(n)
                    if 'Dent' == n['name']:
                        
                        num_ids.append(1)
                    elif 'Scratch' == n['name']:
                        num_ids.append(2)
                    elif 'Others' == n['name']:
                        num_ids.append(3)
                except:
                    pass

            # print("num_ids=", num_ids)
            # print("num_ids_new=", num_ids_new)
            # categories = [s['region_attributes'] for s in a['regions'].values()]
            # load_mask() needs the image size to convert polygons to masks.
            # Unfortunately, VIA doesn't include it in JSON, so we must read
            # the image. This is only managable since the dataset is tiny.
            image_path = os.path.join(dataset_dir, a['filename'])
            image = skimage.io.imread(image_path)
            height, width = image.shape[:2]

            self.add_image(
                "object",
                image_id=a['filename'],  # use file name as a unique image id
                path=image_path,
                width=width, height=height,
                polygons=polygons,
                num_ids=num_ids)

훈련!!!

mask rcnn 파일이 있는 곳 까지 가서 파이썬 파일을 실행시켜준다.

—dataset 과 —weight이 모두 존재해야 훈련이 가능하다

—weights = 'coco' 는 인터넷에서 코코를 받으면 훈련이 되어 편리하므로 이렇게 사용했음.

python3 final.py train --dataset=/home/ubuntu/projects/mask_rcnn_damage_detection/new --weights='coco'

segmentation 해보기

여기서 weights는 아까 train에서 저장된 h5파일을 지정하면된다.

그리고 image는 segmentation하고자 하는 이미지를 넣어주면 된다.

python3 final.py splash --weights=/home/ubuntu/projects/mask_rcnn_damage_detection/logs/balloon20200901T1339/mask_rcnn_balloon_0008.h5 \
--image=/home/ubuntu/projects/mask_rcnn_damage_detection/customImages/val/38305_2.jpg