I'm trying to detect 5 classes of weed species with mask rcnn, but I only get loss: nan - val_loss: nan in mask_rcnn during training, I already tried to change the value of learning hate but I always got the same result. Below I will leave my code, i'll also leave an example of how the .xml is organized
class WeedsDataset(Dataset):
def load_dataset(self, dataset_dir, is_train=True):
# define classes
self.add_class("dataset", 1, "vassourinha")
self.add_class("dataset", 2, "agriaozinho")
self.add_class("dataset", 3, "capim_amargoso")
self.add_class("dataset", 4, "mastruco")
self.add_class("dataset", 5, "coqueiro")
# define data locations
images_dir = dataset_dir + '/images/'
annotations_dir = dataset_dir + '/anotations/'
# find all images
for filename in listdir(images_dir):
#print(filename)
# extract image id
image_id = filename[:-4]
#print('IMAGE ID: ',image_id)
if is_train and int(image_id) >= 500:
continue
# skip all images before 115 if we are building the test/val set
if not is_train and int(image_id) < 500:
continue
img_path = images_dir + filename
ann_path = annotations_dir + image_id + '.xml'
# add to dataset
self.add_image('dataset', image_id=image_id, path=img_path, annotation=ann_path, class_ids = [0,1,2,3,4,5])
# extract bounding boxes from an annotation file
def extract_boxes(self, filename):
# load and parse the file
tree = ElementTree.parse(filename)
# get the root of the document
root = tree.getroot()
# extract each bounding box
boxes = list()
for box in root.findall('.//object'):
name = box.find('name').text #Add label name to the box list
xmin = int(box.find('./bndbox/xmin').text)
xmax = int(box.find('./bndbox/xmax').text)
ymin = int(box.find('./bndbox/ymin').text)
ymax = int(box.find('./bndbox/ymax').text)
coors = [xmin, ymin, xmax, ymax, name]
boxes.append(coors)
# extract image dimensions
width = int(root.find('.//size/width').text)
height = int(root.find('.//size/height').text)
return boxes, width, height
# load the masks for an image
def load_mask(self, image_id):
# get details of image
info = self.image_info[image_id]
# define box file location
path = info['annotation']
#return info, path
# load XML
boxes, w, h = self.extract_boxes(path)
# create one array for all masks, each on a different channel
masks = zeros([h, w, len(boxes)], dtype='uint8')
# create masks
class_ids = list()
for i in range(len(boxes)):
box = boxes[i]
row_s, row_e = box[1], box[3]
col_s, col_e = box[0], box[2]
# box[4] will have the name of the class
if (box[4] == 'vassourinha'):
masks[row_s:row_e, col_s:col_e, i] = 1
class_ids.append(self.class_names.index('vassourinha'))
elif(box[4] == 'agriaozinho'):
masks[row_s:row_e, col_s:col_e, i] = 2
class_ids.append(self.class_names.index('agriaozinho'))
elif(box[4] == 'capim_amargoso'):
masks[row_s:row_e, col_s:col_e, i] = 3
class_ids.append(self.class_names.index('capim_amargoso'))
elif(box[4] == 'mastruco'):
masks[row_s:row_e, col_s:col_e, i] = 4
class_ids.append(self.class_names.index('mastruco'))
elif(box[4] == 'coqueiro'):
masks[row_s:row_e, col_s:col_e, i] = 5
class_ids.append(self.class_names.index('coqueiro'))
return masks, asarray(class_ids, dtype='int32')
# load an image reference
def image_reference(self, image_id):
info = self.image_info[image_id]
return info['path']
# train set
dataset_dir='C:/TCC Gelson/xml_proj.v1i.voc/train_backup'
train_set = WeedsDataset()
train_set.load_dataset(dataset_dir, is_train=True)
train_set.prepare()
print('Train: %d' % len(train_set.image_ids))
# test/val set
test_set = WeedsDataset()
test_set.load_dataset(dataset_dir, is_train=False)
test_set.prepare()
print('Test: %d' % len(test_set.image_ids))
import random
num=random.randint(0, len(train_set.image_ids))
# define image id
image_id = num
# load the image
image = train_set.load_image(image_id)
# load the masks and the class ids
mask, class_ids = train_set.load_mask(image_id)
# extract bounding boxes from the masks
bbox = extract_bboxes(mask)
# display image with masks and bounding boxes
display_instances(image, bbox, mask, class_ids, train_set.class_names)
# define a configuration for the model
class WeedsConfig(Config):
NAME = "weeds_cfg"
GPU_COUNT = 1
IMAGES_PER_GPU = 1
NUM_CLASSES = 1 + 5
LEARNING_RATE = 0.0025
# prepare config
config = WeedsConfig()
config.display()
# define the model
model = MaskRCNN(mode='training', model_dir=MODEL_DIR, config=config)
# load weights (mscoco) and exclude the output layers
model.load_weights("C:/TCC Gelson/xml_proj.v1i.voc/train_backup/weights/mask_rcnn_coco.h5", by_name=True, exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"])
# train weights (output layers or 'heads')
model.train(train_set, test_set, learning_rate=config.LEARNING_RATE, epochs=5, layers='heads')
# change directory to logs directory
os.chdir(DEFAULT_LOGS_DIR)
# save weights
model_path = 'Weeds_mask_rcnn.h5'
model.keras_model.save_weights(model_path)
