用TensorFlowObjectDetectionAPIModel 采用ssd_mobilenet_v1模型训练自己的数据:

1.下载 https://github.com/tensorflow/models.git 安装请自行百度,官网上也提供说明 我的路径 /home/sy/models/

2.制作数据集:

由于我之前用yolo训练过,识别率也可以,但是在移植转化的时候出里点问题,所哟用ssd尝试下,我现有的数据:

  • 将标注文件.txt->.xml
import os
import xml.etree.ElementTree as ET
from PIL import Image
import numpy as np

def change_txt_xml(img_dir, labels_dir, annotations_dir, class_file):
    labels = os.listdir(labels_dir)
    with open(class_file, 'r')as f:
        classes = f.readlines()
        classes = [cls.strip('\n') for cls in classes]

    def write_xml(imgname, filepath, labeldicts):
        root = ET.Element('Annotation')
        ET.SubElement(root, 'folder').text = 'image'
        ET.SubElement(root, 'filename').text = str(imgname)
        ET.SubElement(root,'path').text=str(img_dir)+str(imgname)
        source = ET.SubElement(root, 'source')
        ET.SubElement(source, 'databashe').text = 'Unknown'
        size = ET.SubElement(root, 'size')
        ET.SubElement(size, 'width').text = '1280'
        ET.SubElement(size, 'height').text = '1024'
        ET.SubElement(size, 'depth').text = '3'
        # segmented = ET.SubElement(size, 'segmented')
        ET.SubElement(root,'segmented').text = '0'

        for labeldict in labeldicts:
            objects = ET.SubElement(root, 'object')
            ET.SubElement(objects, 'name').text = labeldict['name']
            ET.SubElement(objects, 'pose').text = 'Unspecified'
            ET.SubElement(objects, 'truncated').text = '0'
            ET.SubElement(objects, 'difficult').text = '0'
            bndbox = ET.SubElement(objects, 'bndbox')
            ET.SubElement(bndbox, 'xmin').text = str(int(labeldict['xmin']))
            ET.SubElement(bndbox, 'ymin').text = str(int(labeldict['ymin']))
            ET.SubElement(bndbox, 'xmax').text = str(int(labeldict['xmax']))
            ET.SubElement(bndbox, 'ymax').text = str(int(labeldict['ymax']))
        tree = ET.ElementTree(root)
        tree.write(filepath, encoding='utf-8')

    for label in labels:
        with open(labels_dir + label, 'r')as f:
            img_id = os.path.splitext(label)[0]+'.jpg'
            contents = f.readlines()
            labeldicts = []
            for content in contents:
                img = np.array(Image.open(img_dir + label.strip('.txt') + '.jpg'))
                sh, sw = img.shape[0], img.shape[1]
                content = content.strip('\n').split()
                x = float(content[1]) * sw
                y = float(content[2]) * sh
                w = float(content[3]) * sw
                h = float(content[4]) * sh
                new_dict = {'name': classes[int(content[0])],
                            'difficult': '0',
                            'xmin': x + 1 - w / 2,
                            'ymin': y + 1 - h / 2,
                            'xmax': x + 1 + w / 2,
                            'ymax': y + 1 + h / 2}
                labeldicts.append(new_dict)
            write_xml(img_id, annotations_dir + label.strip('.txt') + '.xml', labeldicts)


if __name__ == '__main__':
    img_dir = '/home/sy/data/image_jpg/'
    labels_dir = '/home/sy/data/label-txt/'
    annotations_dir = '/home/sy/data/xml_jpg/'
    class_file = '/home/sy/data/class/classes.txt'
    change_txt_xml(img_dir,labels_dir,annotations_dir,class_file)
  • 将.xml->.csv
import os
import glob
import pandas as pd
import xml.etree.ElementTree as ET

def xml_to_csv(path):
    xml_list = []
    for xml_file in glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for member in root.findall('object'):
            value = (root.find('filename').text,
                     int(root.find('size')[0].text),
                     int(root.find('size')[1].text),
                     member[0].text,
                     int(member[4][0].text),
                     int(member[4][1].text),
                     int(member[4][2].text),
                     int(member[4][3].text)
                     )
            xml_list.append(value)
    column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
    xml_df = pd.DataFrame(xml_list, columns=column_name)
    return xml_df


if __name__=="__main__":
    image_path = '/home/sy/data/test_xml_jpg'
    xml_df = xml_to_csv(image_path)
    xml_df.to_csv('/home/sy/data/pupils/data/test.csv', index=None)
    print('Successfully converted xml to csv.')
  • 将.csv->.record(注意train 和 test 分开) 把生成的.record放在一个文件夹下   

/home/sy/data/pupils/data/


import os
import io
import pandas as pd
import tensorflow as tf

from PIL import Image
from object_detection.utils import dataset_util
from collections import namedtuple, OrderedDict

os.chdir('/home/sy/models/research/object_detection/')

flags = tf.app.flags
flags.DEFINE_string('csv_input', '', 'Path to the CSV input')
flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
FLAGS = flags.FLAGS

# TO-DO replace this with label map
def class_text_to_int(row_label):
    if row_label == 'pupils':
        return 1
    else:
        None


def split(df, group):
    data = namedtuple('data', ['filename', 'object'])
    gb = df.groupby(group)
    return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)]


def create_tf_example(group, path):
    with tf.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid:
        encoded_jpg = fid.read()
    encoded_jpg_io = io.BytesIO(encoded_jpg)
    image = Image.open(encoded_jpg_io)
    width, height = image.size

    filename = group.filename.encode('utf8')
    image_format = b'jpg'
    xmins = []
    xmaxs = []
    ymins = []
    ymaxs = []
    classes_text = []
    classes = []

    for index, row in group.object.iterrows():
        xmins.append(row['xmin'] / width)
        xmaxs.append(row['xmax'] / width)
        ymins.append(row['ymin'] / height)
        ymaxs.append(row['ymax'] / height)
        classes_text.append(row['class'].encode('utf8'))
        classes.append(class_text_to_int(row['class']))

    tf_example = tf.train.Example(features=tf.train.Features(feature={
        'image/height': dataset_util.int64_feature(height),
        'image/width': dataset_util.int64_feature(width),
        'image/filename': dataset_util.bytes_feature(filename),
        'image/source_id': dataset_util.bytes_feature(filename),
        'image/encoded': dataset_util.bytes_feature(encoded_jpg),
        'image/format': dataset_util.bytes_feature(image_format),
        'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
        'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
        'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
        'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
        'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
        'image/object/class/label': dataset_util.int64_list_feature(classes),
    }))
    return tf_example


def main(_):
    writer = tf.python_io.TFRecordWriter(FLAGS.output_path)
    path = os.path.join(os.getcwd(), '/home/sy/data/image_jpg')
    examples = pd.read_csv(FLAGS.csv_input)
    grouped = split(examples, 'filename')
    for group in grouped:
        tf_example = create_tf_example(group, path)
        writer.write(tf_example.SerializeToString())

    writer.close()
    output_path = os.path.join(os.getcwd(), FLAGS.output_path)
    print('Successfully created the TFRecords: {}'.format(output_path))


if __name__ == '__main__':
    tf.app.run()
  • 生成.pbtxt文件 也放在/home/sy/data/pupils/data
item{
  id:1
  name:'pupils'
}
  • 修改config 文件

把/home/sy/models/research/object_detection/samples/configs/ssd_mobilenet_v1_coco.config文件 复制到 /home/sy/data/pupils/data 做如下修改

  1. num_classes: 1
  2. fine_tune_checkpoint: "/home/sy/data/pupils/data/ssd_mobilenet_v1_coco_2017_11_17/model.ckpt"

     (需要下载预训练模型,如果你想从头开始,可以注释掉这句和下句)

     3. 修改传入路径和传出路径

(官网上提供一条很长的指令,

sed -i "s|PATH_TO_BE_CONFIGURED|/home/user/model/reaearch/objectdetect/data|g" ssd_mobilenet_v2_coco.config)

也是在做这个工作,只不过你的数据需要在它指定的路径下,还是手动比较保险)

 4.如果电脑显存不够的话,把batch_size设置的小点

batch_size: 2(我的GTX2080,大的也许可以)

3.开始训练

Pythonobject_detection/model_main.py --pipeline_config_path=/home/sy/data/pupils/data/ssd_mobilenet_v1_coco.config --model_dir=/home/sy/data/pupils/result  

--num_train_steps=50000

 --num_eval_steps=2000

--alsologtostderr

(在这个阶段遇到的问题)

TypeError: can't pickle dict_values objects

Ans:

If you're using python3 , add list() to category_index.values() in model_lib.py about line 381 as this list(category_index.values()).

https://github.com/tensorflow/models/issues/4780

4.用官网工具转换成.pb 文件

Ckpt->.pb

Python object_detection/export_inference_graph.py

--input_type=image_tensor

--pipeline_config_path=/home/sy/data/pupils/data/ssd_mobilenet_v1_coco.config

--trained_checkpoint_prefix=/home/sy/data/pupils/result/model.ckpt-17402

--output_directory=/home/sy/data/pupils/result/log/pb_model

 

5. 替换官网demo中模型路径和label文件(需要下载github上的tensorflow,命令行安装的不包含example)

 

6.最后运行效果

当我启动截屏键出现了右图,很神奇有没有 很酷有没有 exceting(层,图片,时间什么都有了)

 

  • 0
    点赞
  • 5
    收藏
    觉得还不错? 一键收藏
  • 3
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值