CelebA数据库转换为VOC、YOLO格式

CelebA数据库转换为VOC、YOLO格式

原文：https://blog.youkuaiyun.com/minstyrain/article/details/77888176

 

CelebA是香港中文大学发布的20多万的名人人脸数据库，被很多算法用来训练，取得了不错的效果，其主页为http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html

它提供了人脸包围框和5个关键点的标注，都放置在一个txt文件中，这和faster-rcnn、YOLO等的要求还是有一些差距的，因此需要做些变换才能使用。

下面这个脚本实现了讲CelebA自带标注转换为VOC和YOLO格式的功能.

 

1. import cv2,h5py,os

2. import numpy as np

3. from xml.dom.minidom import Document

4. import progressbar

5. rootdir="../"

6. imgdir=rootdir+"Img/img_celeba"

7.  

8. landmarkpath=rootdir+"Anno/list_landmarks_celeba.txt"

9. bboxpath=rootdir+"Anno/list_bbox_celeba.txt"

10. vocannotationdir=rootdir+"/"+"Annotations"

11. labelsdir=rootdir+"/"+"labels"

12.  

13. convet2yoloformat=True

14. convert2vocformat=True

15.  

16. resized_dim=(48,48)

17.  

18. datasetprefix="/home/yanhe/data/CelebA/images/"

19. progress = progressbar.ProgressBar(widgets=[

20. progressbar.Percentage(),

21. ' (', progressbar.SimpleProgress(), ') ',

22. ' (', progressbar.Timer(), ') ',

23. ' (', progressbar.ETA(), ') ',])

24. def drawbboxandlandmarks(img,bbox,landmark):

25. cv2.rectangle(img,(bbox[0],bbox[1]),(bbox[0]+bbox[2],bbox[1]+bbox[3]),(0,255,0))

26. for i in range(int(len(landmark)/2)):

27. cv2.circle(img,(int(landmark[2*i]),int(landmark[2*i+1])),2,(0,0,255))

28.  

29. def loadgt():

30. imgpaths=[]

31. landmarks=[]

32. bboxes=[]

33. with open(landmarkpath) as landmarkfile:

34. lines=landmarkfile.readlines()

35. lines=lines[2:]

36. for line in lines:

37. landmarkline=line.split()

38. imgpath=landmarkline[0]

39. imgpaths.append(imgpath)

40. landmarkline=landmarkline[1:]

41. landmark=[int(str) for str in landmarkline]

42. landmarks.append(landmark)

43. with open(bboxpath) as bboxfile:

44. lines=bboxfile.readlines()

45. lines=lines[2:]

46. for line in lines:

47. bboxline=line.split()

48. imgpath=bboxline[0]

49. bboxline=bboxline[1:]

50. bbox=[int(bb) for bb in bboxline]

51. bboxes.append(bbox)

52. return imgpaths,bboxes,landmarks

53.  

54. def generate_hdf5():

55. imgpaths,bboxes,landmarks=loadgt()

56. numofimg=len(imgpaths)

57. faces=[]

58. labels=[]

59. #numofimg=2

60. for i in range(numofimg):

61. imgpath=imgdir+"/"+imgpaths[i]

62. print(i)#,imgpath)

63. bbox=bboxes[i]

64. landmark=landmarks[i]

65. img=cv2.imread(imgpath)

66. if bbox[2]<=0 or bbox[3]<=0:

67. continue

68. face=img[bbox[1]:bbox[1]+bbox[3],bbox[0]:bbox[0]+bbox[2]]

69. face=cv2.resize(face,resized_dim)

70. faces.append(face)

71. label=[]

72. label.append(1)

73. for i in range(len(bbox)):

74. label.append(bbox[i])

75. for i in range(len(landmark)):

76. lm=landmark[i]

77. if i%2==0:

78. lm=(lm-bbox[0])*1.0/(bbox[2])

79. else:

80. lm=(lm-bbox[1])*1.0/(bbox[3])

81. label.append(lm)

82. labels.append(label)

83. faces=np.asarray(faces)

84. labels=np.asarray(labels)

85. f=h5py.File('train.h5','w')

86. f['data']=faces.astype(np.float32)

87. f['labels']=labels.astype(np.float32)

88. f.close()

89. def viewginhdf5():

90. f = h5py.File('train.h5','r')

91. f.keys()

92. faces=f['data'][:]

93. labels=f['labels'][:]

94. for i in range(len(faces)):

95. print(i)

96. face=faces[i].astype(np.uint8)

97. label=labels[i]

98. bbox=label[1:4]

99. landmark=label[5:]

100. for i in range(int(len(landmark)/2)):

101. cv2.circle(face,(int(landmark[2*i]*resized_dim[0]),int(landmark[2*i+1]*resized_dim[1])),1,(0,0,255))

102. cv2.imshow("img",face)

103. cv2.waitKey()

104. f.close()

105. def showgt():

106. landmarkfile=open(landmarkpath)

107. bboxfile=open(bboxpath)

108. numofimgs=int(landmarkfile.readline())

109. _=landmarkfile.readline()

110. _=bboxfile.readline()

111. _=bboxfile.readline()

112. index=0

113. pbar = progress.start()

114. if convet2yoloformat:

115. if not os.path.exists(labelsdir):

116. os.mkdir(labelsdir)

117. if convert2vocformat:

118. if not os.path.exists(vocannotationdir):

119. os.mkdir(vocannotationdir)

120. # while(index<numofimgs):

121. for i in pbar(range(numofimgs)):

122. #pbar.update(int((index/(numofimgs-1))*10000))

123. landmarkline=landmarkfile.readline().split()

124. filename=landmarkline[0]

125. #sys.stdout.write("\r"+str(index)+":"+filename)

126. #sys.stdout.flush()

127. imgpath=imgdir+"/"+filename

128. img=cv2.imread(imgpath)

129. landmarkline=landmarkline[1:]

130. landmark=[int(pt) for pt in landmarkline]

131. bboxline=bboxfile.readline().split()

132. imgpath2=imgdir+"/"+bboxline[0]

133. bboxline=bboxline[1:]

134. bbox=[int(bb) for bb in bboxline]

135. drawbboxandlandmarks(img,bbox,landmark)

136. if convet2yoloformat:

137. height=img.shape[0]

138. width=img.shape[1]

139. txtpath=labelsdir+"/"+filename

140. txtpath=txtpath[:-3]+"txt"

141. ftxt=open(txtpath,'w')

142. xcenter=(bbox[0]+bbox[2]*0.5)/width

143. ycenter=(bbox[1]+bbox[3]*0.5)/height

144. wr=bbox[2]*1.0/width

145. hr=bbox[3]*1.0/height

146. line="0 "+str(xcenter)+" "+str(ycenter)+" "+str(wr)+" "+str(hr)+"\n"

147. ftxt.write(line)

148. ftxt.close()

149. if convert2vocformat:

150. xmlpath=vocannotationdir+"/"+filename

151. xmlpath=xmlpath[:-3]+"xml"

152. doc = Document()

153. annotation = doc.createElement('annotation')

154. doc.appendChild(annotation)

155. folder = doc.createElement('folder')

156. folder_name = doc.createTextNode('CelebA')

157. folder.appendChild(folder_name)

158. annotation.appendChild(folder)

159. filenamenode = doc.createElement('filename')

160. filename_name = doc.createTextNode(filename)

161. filenamenode.appendChild(filename_name)

162. annotation.appendChild(filenamenode)

163. source = doc.createElement('source')

164. annotation.appendChild(source)

165. database = doc.createElement('database')

166. database.appendChild(doc.createTextNode('CelebA Database'))

167. source.appendChild(database)

168. annotation_s = doc.createElement('annotation')

169. annotation_s.appendChild(doc.createTextNode('PASCAL VOC2007'))

170. source.appendChild(annotation_s)

171. image = doc.createElement('image')

172. image.appendChild(doc.createTextNode('flickr'))

173. source.appendChild(image)

174. flickrid = doc.createElement('flickrid')

175. flickrid.appendChild(doc.createTextNode('-1'))

176. source.appendChild(flickrid)

177. owner = doc.createElement('owner')

178. annotation.appendChild(owner)

179. flickrid_o = doc.createElement('flickrid')

180. flickrid_o.appendChild(doc.createTextNode('tdr'))

181. owner.appendChild(flickrid_o)

182. name_o = doc.createElement('name')

183. name_o.appendChild(doc.createTextNode('yanyu'))

184. owner.appendChild(name_o)

185. size = doc.createElement('size')

186. annotation.appendChild(size)

187. width = doc.createElement('width')

188. width.appendChild(doc.createTextNode(str(img.shape[1])))

189. height = doc.createElement('height')

190. height.appendChild(doc.createTextNode(str(img.shape[0])))

191. depth = doc.createElement('depth')

192. depth.appendChild(doc.createTextNode(str(img.shape[2])))

193. size.appendChild(width)

194. size.appendChild(height)

195. size.appendChild(depth)

196. segmented = doc.createElement('segmented')

197. segmented.appendChild(doc.createTextNode('0'))

198. annotation.appendChild(segmented)

199. for i in range(1):

200. objects = doc.createElement('object')

201. annotation.appendChild(objects)

202. object_name = doc.createElement('name')

203. object_name.appendChild(doc.createTextNode('face'))

204. objects.appendChild(object_name)

205. pose = doc.createElement('pose')

206. pose.appendChild(doc.createTextNode('Unspecified'))

207. objects.appendChild(pose)

208. truncated = doc.createElement('truncated')

209. truncated.appendChild(doc.createTextNode('1'))

210. objects.appendChild(truncated)

211. difficult = doc.createElement('difficult')

212. difficult.appendChild(doc.createTextNode('0'))

213. objects.appendChild(difficult)

214. bndbox = doc.createElement('bndbox')

215. objects.appendChild(bndbox)

216. xmin = doc.createElement('xmin')

217. xmin.appendChild(doc.createTextNode(str(bbox[0])))

218. bndbox.appendChild(xmin)

219. ymin = doc.createElement('ymin')

220. ymin.appendChild(doc.createTextNode(str(bbox[1])))

221. bndbox.appendChild(ymin)

222. xmax = doc.createElement('xmax')

223. xmax.appendChild(doc.createTextNode(str(bbox[0]+bbox[2])))

224. bndbox.appendChild(xmax)

225. ymax = doc.createElement('ymax')

226. ymax.appendChild(doc.createTextNode(str(bbox[1]+bbox[3])))

227. bndbox.appendChild(ymax)

228. f=open(xmlpath,"w")

229. f.write(doc.toprettyxml(indent = ''))

230. f.close()

231. cv2.imshow("img",img)

232. cv2.waitKey(1)

233. index=index+1

234. pbar.finish()

235.  

236. def generatetxt(trainratio=0.7,valratio=0.2,testratio=0.1):

237. files=os.listdir(labelsdir)

238. ftrain=open(rootdir+"/"+"train.txt","w")

239. fval=open(rootdir+"/"+"val.txt","w")

240. ftrainval=open(rootdir+"/"+"trainval.txt","w")

241. ftest=open(rootdir+"/"+"test.txt","w")

242. index=0

243. for i in range(len(files)):

244. filename=files[i]

245. filename=datasetprefix+filename[:-3]+"jpg"+"\n"

246. if i<trainratio*len(files):

247. ftrain.write(filename)

248. ftrainval.write(filename)

249. elif i<(trainratio+valratio)*len(files):

250. fval.write(filename)

251. ftrainval.write(filename)

252. elif i<(trainratio+valratio+testratio)*len(files):

253. ftest.write(filename)

254. ftrain.close()

255. fval.close()

256. ftrainval.close()

257. ftest.close()

258.  

259. def generatevocsets(trainratio=0.7,valratio=0.2,testratio=0.1):

260. if not os.path.exists(rootdir+"/ImageSets"):

261. os.mkdir(rootdir+"/ImageSets")

262. if not os.path.exists(rootdir+"/ImageSets/Main"):

263. os.mkdir(rootdir+"/ImageSets/Main")

264. ftrain=open(rootdir+"/ImageSets/Main/train.txt",'w')

265. fval=open(rootdir+"/ImageSets/Main/val.txt",'w')

266. ftrainval=open(rootdir+"/ImageSets/Main/trainval.txt",'w')

267. ftest=open(rootdir+"/ImageSets/Main/test.txt",'w')

268. files=os.listdir(labelsdir)

269. for i in range(len(files)):

270. imgfilename=files[i][:-4]

271. ftrainval.write(imgfilename+"\n")

272. if i<int(len(files)*trainratio):

273. ftrain.write(imgfilename+"\n")

274. elif i<int(len(files)*(trainratio+valratio)):

275. fval.write(imgfilename+"\n")

276. else:

277. ftest.write(imgfilename+"\n")

278. ftrain.close()

279. fval.close()

280. ftrainval.close()

281. ftest.close()

282.  

283. if __name__=="__main__":

284. showgt()

285. generatevocsets()

286. generatetxt()

287. #generate_hdf5()

288.