Mask RCNN 간단한 Head Fine Tuning 예제¶
In [ ]:
import os
import sys
import random
import math
import numpy as np
import cv2
import matplotlib.pyplot as plt
from tqdm import tqdm
import pandas as pd
import glob
from sklearn.model_selection import KFold
from PIL import Image
In [113]:
from mrcnn.config import Config
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize
from mrcnn.model import log
In [114]:
import os
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import numpy as np
데이터 사이즈 변경¶
In [116]:
mask = os.listdir('../maskrcnn/mask/')
train = os.listdir('../maskrcnn/train/')
In [117]:
def save_image(from_path, target_path) :
images = os.listdir(from_path)
for i in images:
if i.lower().endswith(('.png', '.jpg', '.jpeg')) :
image = Image.open(from_path + i)
image = image.resize((128,128), Image.ANTIALIAS)
ext = '.jpg'
image.save(target_path + i + ext)
In [118]:
save_image('../maskrcnn/mask/','../maskrcnn/resized/mask/' )
In [119]:
save_image('../maskrcnn/train/','../maskrcnn/resized/train/' )
In [153]:
mask_resized = os.listdir('../maskrcnn/resized/mask/')
train_resized = os.listdir('../maskrcnn/resized/train/')
데이터셋 만들기¶
In [172]:
class DetectorDataset(utils.Dataset):
def load_dataset(self, image_filenames):
# Add classes
self.add_class("openimage", 1, "hi")
self.add_class("openimage", 2, "bi")
for i in image_filenames:
print(i)
self.add_image("openimage", image_id=i, path='../maskrcnn/resized/train/' + i,width=128, height=128)
# def load_mask(self, image_id):
# return np.zeros((128,128,1))
# return np.array(Image.open('../maskrcnn/resized/mask/' + image_id))
dataset = DetectorDataset()
dataset.load_dataset(train_resized)
dataset.prepare()
In [170]:
ROOT_DIR = os.path.abspath("../../")
# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
# Local path to trained weights file
COCO_MODEL_PATH = "mask_rcnn_coco.h5"
# Download COCO trained weights from Releases if needed
if not os.path.exists("mask_rcnn_coco.h5"):
utils.download_trained_weights(COCO_MODEL_PATH)
In [145]:
class ShapesConfig(Config):
"""Configuration for training on the toy shapes dataset.
Derives from the base Config class and overrides values specific
to the toy shapes dataset.
"""
# Give the configuration a recognizable name
NAME = "shapes"
# Train on 1 GPU and 8 images per GPU. We can put multiple images on each
# GPU because the images are small. Batch size is 8 (GPUs * images/GPU).
GPU_COUNT = 1
IMAGES_PER_GPU = 8
# Number of classes (including background)
NUM_CLASSES = 1 + 3 # background + 3 shapes
# Use small images for faster training. Set the limits of the small side
# the large side, and that determines the image shape.
IMAGE_MIN_DIM = 128
IMAGE_MAX_DIM = 128
# Use smaller anchors because our image and objects are small
RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128) # anchor side in pixels
# Reduce training ROIs per image because the images are small and have
# few objects. Aim to allow ROI sampling to pick 33% positive ROIs.
TRAIN_ROIS_PER_IMAGE = 32
# Use a small epoch since the data is simple
STEPS_PER_EPOCH = 100
# use small validation steps since the epoch is small
VALIDATION_STEPS = 5
config = ShapesConfig()
config.display()
In [146]:
# Directory to save logs and trained model
MODEL_DIR = os.path.join("./model/")
In [147]:
model = modellib.MaskRCNN(mode="training", config=config,model_dir=MODEL_DIR)
In [148]:
model.load_weights(COCO_MODEL_PATH, by_name=True,
exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",
"mrcnn_bbox", "mrcnn_mask"])
In [ ]:
model.train(dataset, dataset, learning_rate=config.LEARNING_RATE, epochs=1, layers='heads')
In [ ]: