caffe 利用VoC2012制作图片数据集

代码

#coding='utf-8'
import scipy.misc
import caffe
import numpy as np
import os.path as osp
from xml.dom import minidom
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
import h5py

classes = ('__background__',  # always index 0
		   'aeroplane', 'bicycle', 'bird', 'boat',
		   'bottle', 'bus', 'car', 'cat', 'chair',
			'cow', 'diningtable', 'dog', 'horse',
			'motorbike', 'person', 'pottedplant',
			'sheep', 'sofa', 'train', 'tvmonitor')

def load_pascal_annotation(index, pascal_root):

	class_to_ind = dict(zip(classes, range(21)))
	filename = osp.join(pascal_root, 'Annotations', index + '.xml')
	# print 'Loading: {}'.format(filename)
	def get_data_from_tag(node, tag):
		return node.getElementsByTagName(tag)[0].childNodes[0].data

	with open(filename) as f:
		data = minidom.parseString(f.read())

	objs = data.getElementsByTagName('object')
	num_objs = len(objs)

	boxes = np.zeros((num_objs, 4), dtype=np.uint16)
	gt_classes = np.zeros((num_objs), dtype=np.int32)
	overlaps = np.zeros((num_objs, 21), dtype=np.float32)

	# Load object bounding boxes into a data frame.
	for ix, obj in enumerate(objs):
		# Make pixel indexes 0-based
		x1 = float(get_data_from_tag(obj, 'xmin')) - 1
		y1 = float(get_data_from_tag(obj, 'ymin')) - 1
		x2 = float(get_data_from_tag(obj, 'xmax')) - 1
		y2 = float(get_data_from_tag(obj, 'ymax')) - 1
		cls = class_to_ind[
			str(get_data_from_tag(obj, "name")).lower().strip()]
		boxes[ix, :] = [x1, y1, x2, y2]
		gt_classes[ix] = cls
		overlaps[ix, cls] = 1.0

	overlaps = scipy.sparse.csr_matrix(overlaps)

	return {'boxes': boxes,
			'gt_classes': gt_classes,
			'gt_overlaps': overlaps,
			'flipped': False,
			'index': index}

def val():

	train_num = 100
	pascal_root = 'D://learn/目标检测/VOCdevkit/VOC2012/'
	list_file = 'trainval.txt'
	indexlist = [line.rstrip('\n') for line in open(
	osp.join(pascal_root, 'ImageSets/Main', list_file))]

	f = h5py.File('train.h5', 'w')
	f.create_dataset('data', (train_num, 3,227,227), dtype='f8')
	f.create_dataset('label', (train_num, 20), dtype='i')

	f1 = h5py.File('test.h5', 'w')
	f1.create_dataset('data', (int(train_num*0.25), 3,227,227), dtype='f8')
	f1.create_dataset('label', (train_num, 20), dtype='i')

	# 然后输出前100个图像的label
	for i, name in enumerate(indexlist[0:train_num]):
		multilabel = np.zeros(20).astype(np.float32)
		img = Image.open(pascal_root+'JPEGImages/'+name+'.jpg')
		img = scipy.misc.imresize(img, (227,227)).transpose((2,0,1))
		anns = load_pascal_annotation(name, pascal_root)
		for label in anns['gt_classes']:
			multilabel[label - 1] = 1
		print (i,multilabel)
		f['data'][i] = img
		f['label'][i] = multilabel

	for i, name in enumerate(indexlist[train_num:int(train_num +train_num*0.25)]):
		multilabel1 = np.zeros(20).astype(np.float32)
		img1 = Image.open(pascal_root+'JPEGImages/'+name+'.jpg')
		img1 = scipy.misc.imresize(img, (227,227)).transpose((2,0,1))
		anns1 = load_pascal_annotation(name, pascal_root)
		for label in anns['gt_classes']:
			multilabel1[label - 1] = 1
		print (i,multilabel1)
		f1['data'][i] = img1
		f1['label'][i] = multilabel1


	with open("train.h5list","w")as f:
		f.write("train.h5\n")

	with open("test.h5list","w")as f1:
		f1.write("test.h5\n")



def main():
	pascal_root = 'D://learn/目标检测/VOCdevkit/VOC2012/'

	filename = '2007_000783' 

	plt.figure()
	img = Image.open(pascal_root+'JPEGImages/'+filename+'.jpg')
	gt = load_pascal_annotation(filename, pascal_root)
	draw = ImageDraw.Draw(img)

	for box in gt['boxes']:
		draw.rectangle(list(box), outline = 'red')

	plt.imshow(img)
	plt.title('GT: {}'.format([classes[g] for g in gt['gt_classes']]))
	plt.axis('off');
	plt.show() 

val()

main(显示标记框)

利用xml文件

val（制作数据集）