将分割数据集转为目标检测数据

１．通过代码找到分割图片里的颜色

from PIL import Image
image = Image.open("/image_segmentation/8256600123032.bmp")
print len(image.getcolors())
print image.getcolors()

２．把颜色显示出来，方便判断和选择

#!/usr/bin/python
# -*- coding: UTF-8 -*-

import numpy as np
import cv2 as cv
 
img = np.zeros((320, 320, 3), np.uint8) #生成一个空灰度图像
print img.shape # 输出：(480, 480, 3)

point_size = 1
point_color = (0, 220, 220) # BGR
thickness = 4 # 可以为 0 、4、8

# 要画的点的坐标
points_list = [(160, 160), (136, 160), (150, 200), (200, 180), (120, 150), (145, 180)]

for point in points_list:
	cv.circle(img, point, point_size, point_color, thickness)

# 画圆，圆心为：(160, 160)，半径为：60，颜色为：point_color，实心线
cv.circle(img, (160, 160), 60, point_color, 0)

cv.namedWindow("image")
cv.imshow('image', img)
cv.waitKey (10000) # 显示 10000 ms 即 10s 后消失
cv.destroyAllWindows()

３．二值化

from PIL import Image
import numpy as np
import cv2
import os
path="image_segmentation/"
new_path="/new_segmentation/"
for f in os.listdir(path):
    img = Image.open(path+f)
    img = img.convert('RGBA')
    pixdata = img.load()
    for y in range(img.size[1]):
        for x in range(img.size[0]):
            if pixdata[x,y][0]!=220 and pixdata[x,y][1]!=220 and pixdata[x,y][2]!=0:
                   pixdata[x, y] = (0, 0, 0,255)
            if pixdata[x,y][0]==220 and pixdata[x,y][1]==220 and pixdata[x,y][2]==0:
                   pixdata[x, y] = (255, 255, 255,255)

    img.save(new_path+f)

４．找连通阈，对每个连通域存label（这一步有点慢）

# -*- coding:utf-8 -*-
import cv2
import numpy as np
import matplotlib.pyplot as plt
import time
import os
color_path="/image_color/"
path="/new_segmentation/"
label_txt="/ros/label_txt/"
new_image="/ros/new/"
def find(img):
    img_flag = 255*np.ones(img.shape,np.int8)
    count = 0
    findpoint = []

    for x in range(img.shape[0]):
        for y in range(img.shape[1]):
            if(img[x][y] == 255 and img_flag[x][y] == 255):
                count += 1
                img_flag[x][y] = count
                findpoint.append((x,y))
            while len(findpoint) > 0:
                xx,yy = findpoint.pop()
                if xx > 0 :#上
                    if img[xx-1][yy] == 255 and img_flag[xx-1][yy] == 255:
                        findpoint.append((xx-1,yy))
                        img_flag[xx-1][yy] = count
                if xx < img.shape[0]-1:#下
                    if img[xx + 1][yy] == 255 and img_flag[xx + 1][yy] == 255:
                        findpoint.append((xx + 1, yy))
                        img_flag[xx+1][yy] = count
                if yy > 0:#左
                    if img[xx][yy-1] == 255 and img_flag[xx][yy-1] == 255:
                        findpoint.append((xx, yy-1))
                        img_flag[xx][yy-1] = count
                if yy < img.shape[1]-1:#右
             
                    if img[xx][yy+1] == 255 and img_flag[xx][yy+1] == 255:
                        findpoint.append((xx, yy+1))
                        img_flag[xx][yy+1] = count
    coutours = []
    for num in range(1,count+1):
        coutours.append([])
        for x in range(img_flag.shape[0]):
            for y in range(img_flag.shape[1]):
                if img_flag[x][y] == num:
                    coutours[num-1].append([x,y,img_flag[x][y]])
    desCoutous={}
    ii=0
    for num in range(len(coutours)):
        tmp = np.mat(coutours[num])
        minX = np.min(tmp[:,0])
        maxX = np.max(tmp[:,0])
        minY = np.min(tmp[:,1])
        maxY = np.max(tmp[:,1])
        dd=np.zeros((2,2))
        dd[0][0]=minX
        dd[0][1]=maxX
        dd[1][0]=minY
        dd[1][1]=maxY
        if maxX-minX>0 and maxY-minY>0 and (maxX-minX)*(maxY-minY)>=60:
            desCoutous.update({ii:dd})
            ii=ii+1
    return desCoutous



for f in os.listdir(path):
    img = cv2.imread(path+f,0)
    find_results = find(img)
    if find_results:

        img_color=cv2.imread(color_path+f)
        label_txt_all=label_txt+f[0:-4]+'.txt'
        for i in range(len(find_results)):
            str1='traffic_light'+ ' ' + str(int(find_results[i][1][0])) + ' ' + str(int(find_results[i][0][0])) + ' ' + str(int(find_results[i][1][1])) + ' ' + str(int(find_results[i][0][1])) + '\n' 
            with open(label_txt_all, 'a+') as txt:
                txt.write(str1)
            cv2.rectangle(img_color, (int(find_results[i][1][0]),int(find_results[i][0][0])), (int(find_results[i][1][1]), int(find_results[i][0][1])), (0, 255, 0), 1)
        cv2.imwrite(new_image+f,img_color)
#cv2.imshow("test",img_color)

#cv2.waitKey(0)

５．txt转xml
 

import xml.dom
import xml.dom.minidom
import os
import cv2
 
_TXT_PATH= 'label_txt/'
_IMAGE_PATH= 'image_color/'
 
_INDENT= ''*4
_NEW_LINE= '\n'
_FOLDER_NODE= ''
_ROOT_NODE= 'annotation'
_DATABASE_NAME= ''
_ANNOTATION= 'PASCAL VOC2007'
_AUTHOR= ''
_SEGMENTED= '0'
_DIFFICULT= '0'
_TRUNCATED= '0'
_POSE= 'Unspecified'
 

_IMAGE_COPY_PATH= 'JPEGImages'
_ANNOTATION_SAVE_PATH= 'Annotations'
 
 

def createElementNode(doc,tag, attr):  
    element_node = doc.createElement(tag)
 

    text_node = doc.createTextNode(attr)
 

    element_node.appendChild(text_node)
 
    return element_node
 

def createChildNode(doc,tag, attr,parent_node):
 
 
 
    child_node = createElementNode(doc, tag, attr)
 
    parent_node.appendChild(child_node)
 

 
def createObjectNode(doc,attrs):
 
    object_node = doc.createElement('object')
 
    createChildNode(doc, 'name', attrs['classification'],
                    object_node)
 
    createChildNode(doc, 'pose',
                    _POSE, object_node)
 
    createChildNode(doc, 'truncated',
                    _TRUNCATED, object_node)
 
    createChildNode(doc, 'difficult',
                    _DIFFICULT, object_node)
 
    bndbox_node = doc.createElement('bndbox')
 
    createChildNode(doc, 'xmin', attrs['x1'],
                    bndbox_node)
 
    createChildNode(doc, 'ymin', attrs['y1'],
                    bndbox_node)
 
    createChildNode(doc, 'xmax', attrs['x2'],
                    bndbox_node)
 
    createChildNode(doc, 'ymax', attrs['y2'],
                    bndbox_node)
 
 
    object_node.appendChild(bndbox_node)
 
    return object_node
 

def writeXMLFile(doc,filename):
 
    tmpfile =open('tmp.xml','w')
 
    doc.writexml(tmpfile, addindent=''*4,newl = '\n',encoding = 'utf-8')
 
    tmpfile.close()
 

    fin =open('tmp.xml')
 
    fout =open(filename, 'w')
 
    lines = fin.readlines()
 
    for line in lines[1:]:
 
        if line.split():
 
         fout.writelines(line)
 
        # new_lines = ''.join(lines[1:])
 
        # fout.write(new_lines)
 
    fin.close()
 
    fout.close()
 
def getFileList(path):
 
    fileList = []
    files = os.listdir(path)
    for f in files:
        if (os.path.isfile(path + '/' + f)):
            fileList.append(f)
    # print len(fileList)
    return fileList
 
 
if __name__ == "__main__":
 
    fileList = getFileList(_TXT_PATH)
    if fileList == 0:
        os._exit(-1)
 
    current_dirpath = os.path.dirname(os.path.abspath('__file__'))
 
    if not os.path.exists(_ANNOTATION_SAVE_PATH):
        os.mkdir(_ANNOTATION_SAVE_PATH)
 
    if not os.path.exists(_IMAGE_COPY_PATH):
        os.mkdir(_IMAGE_COPY_PATH)
 
    for xText in range(len(fileList)):
 
        saveName= "%06d" %(xText+1)
        pos = fileList[xText].rfind(".")
        textName = fileList[xText][:pos]
 
        ouput_file = open(_TXT_PATH  + fileList[xText])
  
        lines = ouput_file.readlines()
 
        xml_file_name = os.path.join(_ANNOTATION_SAVE_PATH, (saveName + '.xml'))
 
        img=cv2.imread(os.path.join(_IMAGE_PATH,(textName+'.bmp')))
   
        try:
            height,width,channel=img.shape
        except AttributeError:
            continue
        
        
        print(os.path.join(_IMAGE_COPY_PATH,(textName+'.jpg')))
        cv2.imwrite(os.path.join(_IMAGE_COPY_PATH,(saveName+'.jpg')),img)
        my_dom = xml.dom.getDOMImplementation()
 
        doc = my_dom.createDocument(None,_ROOT_NODE,None)
 

        root_node = doc.documentElement
 

 
        createChildNode(doc, 'folder',_FOLDER_NODE, root_node)
 

 
        createChildNode(doc, 'filename', saveName+'.jpg',root_node)
 

 
        source_node = doc.createElement('source')
 

 
        createChildNode(doc, 'database',_DATABASE_NAME, source_node)
 
        createChildNode(doc, 'annotation',_ANNOTATION, source_node)
 
        createChildNode(doc, 'image','flickr', source_node)
 
        createChildNode(doc, 'flickrid','NULL', source_node)
 
        root_node.appendChild(source_node)
 

        owner_node = doc.createElement('owner')
 

 
        createChildNode(doc, 'flickrid','NULL', owner_node)
 
        createChildNode(doc, 'name',_AUTHOR, owner_node)
 
        root_node.appendChild(owner_node)
 

 
        size_node = doc.createElement('size')
 
        createChildNode(doc, 'width',str(width), size_node)
 
        createChildNode(doc, 'height',str(height), size_node)
 
        createChildNode(doc, 'depth',str(channel), size_node)
 
        root_node.appendChild(size_node)
 
        createChildNode(doc, 'segmented',_SEGMENTED, root_node)
 
 
        for line in lines:
 
            s = line.rstrip('\n')
 
            array = s.split(' ')
            print array
 
            print(array)
 
            attrs = dict()
 
            attrs['x1']= array[1]
 
            attrs['y1']= array[2]
 
            attrs['x2']= array[3]
 
            attrs['y2']= array[4]

 
            attrs['classification'] = str('traffic_light')
 
            print(xml_file_name)
 

 
            object_node = createObjectNode(doc, attrs)
 
            root_node.appendChild(object_node)
 
     
 
            writeXMLFile(doc, xml_file_name)