机器视觉在农业上的应用：基于Mask R-CNN的大豆豆荚期识别与豆荚计数的研究（一）

1.前期准备

 

 

 转换代码

1. 用labelme进行标注

2. 将生成的json文件转换成需要的dataset

import argparse
import json
import os
import os.path as osp
import warnings
 
import PIL.Image
import yaml
 
from labelme import utils
import base64
 
def main():
    warnings.warn("This script is aimed to demonstrate how to convert the\n"
                  "JSON file to a single image dataset, and not to handle\n"
                  "multiple JSON files to generate a real-use dataset.")
    parser = argparse.ArgumentParser()
    parser.add_argument('json_file')
    parser.add_argument('-o', '--out', default=None)
    args = parser.parse_args()
 
    json_file = args.json_file
    if args.out is None:
        out_dir = osp.basename(json_file).replace('.', '_')
        out_dir = osp.join(osp.dirname(json_file), out_dir)
    else:
        out_dir = args.out
    if not osp.exists(out_dir):
        os.mkdir(out_dir)
 
    count = os.listdir(json_file) 
    for i in range(0, len(count)):
        path = os.path.join(json_file, count[i])
        if os.path.isfile(path):
            data = json.load(open(path))
            
            if data['imageData']:
                imageData = data['imageData']
            else:
                imagePath = os.path.join(os.path.dirname(path), data['imagePath'])
                with open(imagePath, 'rb') as f:
                    imageData = f.read()
                    imageData = base64.b64encode(imageData).decode('utf-8')
            img = utils.img_b64_to_arr(imageData)
            label_name_to_value = {'_background_': 0}
            for shape in data['shapes']:
                label_name = shape['label']
                if label_name in label_name_to_value:
                    label_value = label_name_to_value[label_name]
                else:
                    label_value = len(label_name_to_value)
                    label_name_to_value[label_name] = label_value
            
            # label_values must be dense
            label_values, label_names = [], []
            for ln, lv in sorted(label_name_to_value.items(), key=lambda x: x[1]):
                label_values.append(lv)
                label_names.append(ln)
            assert label_values == list(range(len(label_values)))
            
            lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
            
            captions = ['{}: {}'.format(lv, ln)
                for ln, lv in label_name_to_value.items()]
            lbl_viz = utils.draw_label(lbl, img, captions)
            
            out_dir = osp.basename(count[i]).replace('.', '_')
            out_dir = osp.join(osp.dirname(count[i]), out_dir)
            if not osp.exists(out_dir):
                os.mkdir(out_dir)
 
            PIL.Image.fromarray(img).save(osp.join(out_dir, 'img.png'))
            #PIL.Image.fromarray(lbl).save(osp.join(out_dir, 'label.png'))
            utils.lblsave(osp.join(out_dir, 'label.png'), lbl)
            PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, 'label_viz.png'))
 
            with open(osp.join(out_dir, 'label_names.txt'), 'w') as f:
                for lbl_name in label_names:
                    f.write(lbl_name + '\n')
 
            warnings.warn('info.yaml is being replaced by label_names.txt')
            info = dict(label_names=label_names)
            with open(osp.join(out_dir, 'info.yaml'), 'w') as f:
                yaml.safe_dump(info, f, default_flow_style=False)
 
            print('Saved to: %s' % out_dir)
if __name__ == '__main__':
    main()

将cv16转为方便处理的cv8

 

转换代码

from PIL import Image
import numpy as np
import math
import os

path = 'K:/PicturesofBean/20180928-pod/90-100/load'
newpath = 'K:/PicturesofBean/20180928-pod/90-100/cv2_mask/'


def toeight():
    filelist = os.listdir(path)
    for file in filelist:
        whole_path = os.path.join(path, file)
        img = Image.open(whole_path)
        img = np.array(img)
        #img = Image.fromarray(np.uint8(img / float(math.pow(2, 16) - 1) * 255))
        img = Image.fromarray(np.uint8(img))
        img.save(newpath + file)


toeight()

2.训练自己的数据

"""
Created on Mon Nov 19 20:40:59 2018

@author: 34905
"""

import os
import sys
import random
import math
import re
import time
import numpy as np
import cv2
import matplotlib
import matplotlib.pyplot as plt
import tensorflow as tf

#ROOT_DIR = os.getcwd()
ROOT_DIR = os.path.abspath("G:/Mask_RCNN")
sys.path.append(ROOT_DIR)

from mrcnn.config import Config
#import utils
from mrcnn import model as modellib,utils
from mrcnn import visualize
import yaml
from mrcnn.model import log
from PIL import Image


MODEL_DIR = os.path.join(ROOT_DIR, "logs/pods20181209T2207")

iter_num=0

COCO_MODEL_PATH = os.path.join(MODEL_DIR, "mask_rcnn_coco.h5")

if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)
    
class ShapesConfig(Config):
     NAME = "pods"
    
     GPU_COUNT = 1
     
     IMAGES_PER_GPU = 1
    
     NUM_CLASSES = 1 + 1
     
     IMAGE_MIN_DIM = 320
     IMAGE_MAX_DIM = 384
     
     RPN_ANCHOR_SCALES = (8 * 4, 16 * 4, 32 * 4)
#     RPN_ANCHOR_SCALES = (8 * 6, 16 * 6, 32 * 6, 64 * 6, 128 * 6)
#     RPN_ANCHOR_SCALES = (32, 64, 128, 256, 512)
     
     RPN_ANCHOR_RATIOS = [0.5, 1, 2]
     
     RPN_ANCHOR_STRIDE = 1
     RPN_NMS_THRESHOLD = 0.7
     RPN_TRAIN_ANCHORS_PER_IMAGE = 256
     
     TRAIN_ROIS_PER_IMAGE = 100
     
     STEPS_PER_EPOCH = 5
     
     MAX_GT_INSTANCES = 100
     
     VALIDATION_STEPS = 1
     
     DETECTION_MIN_CONFIDENCE = 0.9
     
     DETECTION_NMS_THRESHOLD = 0.1
config = ShapesConfig()
config.display()

class DrugDataset(utils.Dataset):
    def get_obj_index(self, image):
        n = np.max(image)
        return n
    
    def from_yaml_get_class(self, image_id):
        info = self.image_info[image_id]
        with open(info['yaml_path']) as f:
            temp = yaml.load(f.read())
            labels = temp['label_names']
            del labels[0]
        return labels
    
    def draw_mask(self, num_obj, mask, image,image_id):
        info = self.image_info[image_id]
        for index in range(num_obj):
            for i in range(info['width']):
                for j in range(info['height']):
                    at_pixel = image.getpixel((i, j))
                    if at_pixel == index + 1:
                        mask[j, i, index] = 1
        return mask
    
    def load_shapes(self, count, img_folder, mask_folder, imglist, dataset_root_path):
        self.add_class("shapes", 1, "pod")
        for i in range(count):
            filestr = imglist[i].split(".")[0]
            mask_path = mask_folder + "/" + filestr + ".png"
            yaml_path = dataset_root_path + "labelme_json/" + filestr + "_json/info.yaml"
            print(dataset_root_path + "labelme_json/" + filestr + "_json/img.png")
            cv_img = cv2.imread(dataset_root_path + "labelme_json/" + filestr + "_json/img.png")
            
            self.add_image("shapes", image_id=i, path=img_folder + "/" + imglist[i],
                           width=cv_img.shape[1], height=cv_img.shape[0], mask_path=mask_path, yaml_path=yaml_path)

    def load_mask(self, image_id):
        global iter_num
        print("image_id",image_id)
        info = self.image_info[image_id]
        count = 1  # number of object
        img = Image.open(info['mask_path'])
        num_obj = self.get_obj_index(img)
        mask = np.zeros([info['height'], info['width'], num_obj], dtype=np.uint8)
        mask = self.draw_mask(num_obj, mask, img,image_id)
        occlusion = np.logical_not(mask[:, :, -1]).astype(np.uint8)
        for i in range(count - 2, -1, -1):
            mask[:, :, i] = mask[:, :, i] * occlusion
            
            occlusion = np.logical_and(occlusion, np.logical_not(mask[:, :, i]))
        labels = []
        labels = self.from_yaml_get_class(image_id)
        labels_form = []
        for i in range(len(labels)):
            if labels[i].find("pod") != -1:
                # print "box"
                labels_form.append("pod")
        class_ids = np.array([self.class_names.index(s) for s in labels_form])
        return mask, class_ids.astype(np.int32)

def get_ax(rows=1, cols=1, size=8):
    _, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
    return ax

#dataset_root_path="train_data/"
dataset_root_path = os.path.join(ROOT_DIR, "train_data/")
img_folder = dataset_root_path + "pic"
mask_folder = dataset_root_path + "cv2_mask"
imglist = os.listdir(img_folder)
count = len(imglist)

dataset_train = DrugDataset()
dataset_train.load_shapes(count, img_folder, mask_folder, imglist,dataset_root_path)
dataset_train.prepare()

dataset_val = DrugDataset()
dataset_val.load_shapes(5, img_folder, mask_folder, imglist,dataset_root_path)
dataset_val.prepare()

model = modellib.MaskRCNN(mode="training", config=config,model_dir=MODEL_DIR)

init_with = "coco"

if init_with == "imagenet":
    model.load_weights(model.get_imagenet_weights(), by_name=True)
elif init_with == "coco":
    model.load_weights(COCO_MODEL_PATH, by_name=True,exclude=["mrcnn_class_logits", "mrcnn_bbox_fc","mrcnn_bbox", "mrcnn_mask"])
elif init_with == "last":
    model.load_weights(model.find_last()[1], by_name=True)
    
#model.train(dataset_train, dataset_val,learning_rate=config.LEARNING_RATE,epochs=1,layers='heads')
model.train(dataset_train, dataset_val,learning_rate=config.LEARNING_RATE,epochs=5,layers='heads')

model.train(dataset_train, dataset_val,learning_rate=config.LEARNING_RATE / 10,epochs=2,layers="all")
                                                                        #model.train(dataset_train, dataset_val,learning_rate=config.LEARNING_RATE / 10,epochs=1,layers="all")

3.测试自己的数据

# -*- coding: utf-8 -*-
"""
Created on Fri Nov 23 21:15:41 2018

@author: 34905
"""

import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt
import cv2
import time

ROOT_DIR = os.path.abspath("G:/Mask_RCNN")

sys.path.append(ROOT_DIR)

from mrcnn.config import Config
from datetime import datetime
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize

sys.path.append(os.path.join(ROOT_DIR, "samples/coco/"))
from samples.coco import coco

MODEL_DIR = os.path.join(ROOT_DIR, "logs")
#MODEL_DIR = os.path.join(ROOT_DIR, "logs")

COCO_MODEL_PATH = os.path.join(MODEL_DIR, "mask_rcnn_pods_0200.h5")
#COCO_MODEL_PATH = os.path.join(MODEL_DIR, "mask_rcnn_coco.h5")

IMAGE_DIR = os.path.join(ROOT_DIR, "train_data/load")
#IMAGE_DIR = os.path.join(ROOT_DIR, "images")

class ShapesConfig(Config):
    NAME = "shapes"
    
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
    
    NUM_CLASSES = 1 + 1
    
    IMAGE_MIN_DIM = 320
    IMAGE_MAX_DIM = 384
    
#    RPN_ANCHOR_SCALES = (4 * 8, 4 * 16, 4 * 32, 4 * 64)
    RPN_ANCHOR_SCALES = (32, 64, 128, 256, 512)
    
    TRAIN_ROIS_PER_IMAGE =100
    
    STEPS_PER_EPOCH = 10
    
    VALIDATION_STEPS = 1
    
    DETECTION_MIN_CONFIDENCE = 0.9
     
    DETECTION_NMS_THRESHOLD = 0.2
    
class InferenceConfig(ShapesConfig):
    
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
    
config = InferenceConfig()

model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

model.load_weights(COCO_MODEL_PATH, by_name=True)

class_names = ['BG', 'pod']
count = os.listdir(IMAGE_DIR)
for i in range(len(count)-1,len(count)):
    path=os.path.join(IMAGE_DIR,count[i])
#    if os.path.isfile(path):
#        file_names = next(os.walk(IMAGE_DIR))[2]
#        image = skimage.io.imread(os.path.join(IMAGE_DIR, random.choice(file_names)))
#        image = skimage.io.imread(os.path.join(IMAGE_DIR, count[i]))
#        image = skimage.io.imread(os.path.join(IMAGE_DIR, "1108-02a-h.JPG"))
    image = skimage.io.imread(sys.argv[1])
    a=datetime.now()

    results = model.detect([image], verbose=1)
    b=datetime.now()

    print("time:",(b-a).seconds)
    r = results[0]
    print(r['class_ids'].shape[0])
    visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores'])