pytthon学习脚本

convert_fence_label.py
import os
import cv2
import resource_pb2
import hadmap_pb2
import google
import google.protobuf
import google.protobuf.text_format
import threading
import numpy as np
import argparse

FLAGS = None
def convert(label_dir, image_dir, out_dir, lane_num, crop_h, resize_w, resize_h, debug):
    data = None
    line_num = 0
    filenames = os.listdir(label_dir)
    for filename in filenames:
        if os.path.splitext(filename)[-1] != '.road':
            continue
        filepath = os.path.join(label_dir, filename)
        proto_file = open(filepath, 'r')
        proto_str = proto_file.read()
        pbout = resource_pb2.UplaodLaneLine2dRequest()
        google.protobuf.text_format.Merge(proto_str, pbout)
        image_name = filename[0:-5] + '.jpg'
        image = cv2.imread(os.path.join(image_dir, image_name))
        if image is None:
            print 'can not open image'
            continue
        height = image.shape[0]
        width = image.shape[1]
        #print("h=%d,w=%d" % (height, width))
        step = float(crop_h) / float(lane_num)
        resize_ratio_w = float(resize_w) / width
        resize_ratio_h = float(resize_h) / height
        image = cv2.resize(image, (resize_w, resize_h))
        h = image.shape[0]
        w = image.shape[1]

        #print len(pbout.lane_line)
        print(image_name)
        has_fence = False
        with open(os.path.join(out_dir , image_name.replace('jpg','txt')), 'w') as outfile:
            lane_group = {}
            for lane in pbout.lane_line:
                if lane.type != hadmap_pb2.LINE_TYPE_FENCE:
                    continue
                has_fence = True
                tab = int(lane.lane_id)
                print tab
                if tab not in lane_group:
                    lane_group[tab] = {"x":[], "y":[]}
                lane_group[tab]['x'] += (np.array([(p.x + width / 2) for p in lane.center_line]) * resize_ratio_w).astype(np.int32).tolist()
                lane_group[tab]['y'] += (np.array([(-p.y + height / 2) for p in lane.center_line]) * resize_ratio_h).astype(np.int32).tolist()
            #print(lane_group)
            #print([x for x in lane_group])
            #print([lane_group[x]['type'] for x in lane_group])

            if debug:
              for tbs in lane_group:
                for p in range(len(lane_group[tbs]['x'])):
                    cv2.circle(image, (lane_group[tbs]['x'][p], lane_group[tbs]['y'][p]), 2, (255, 0, 0), 2)

        if debug:
            print(image_name)
            if has_fence:
                cv2.imwrite('debug/'+ image_name, image)
            #cv2.imshow("press ESC to quit", image)
            #k = cv2.waitKey(0)
            #if k == 27:
            #   break

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-d','--imagedir',
        type=str,
        default='checked_images/',
        help='The images dir.'
    )
    parser.add_argument(
        '-i','--labeldir',
        type=str,
        default='checked_labels/',
        help='The images dir.'
    )
    parser.add_argument(
        '-o','--outdir',
        type=str,
        required=False,
        default='./labels',
        help='The output labels dir.'
    )
    parser.add_argument(
        '-l','--lanenum',
        type=int,
        default=11, # 64 = 352 * 2 / 11, downsample=2, input_h = 352, feature_map_h = 11
        help='The step size of lane gap, pixels in raw image, default=57.1=314 * 2 / 11'
    )
    parser.add_argument(
        '-c','--crop',
        type=int,
        default=352, # 704 = 352 * 2
        help='croped height of the image, default=628=314*2'
    )
    parser.add_argument(
        '-y','--resize_h',
        type=int,
        default=604,
        help='height of the image after resize'
    )
    parser.add_argument(
        '-x','--resize_w',
        type=int,
        default=960,
        help='width of the image after resize'
    )
    parser.add_argument(
        '--debug',
        action='store_true',
        default=False
    )

    FLAGS, unparsed = parser.parse_known_args()
    print(FLAGS)
    if not os.path.exists(FLAGS.outdir):
        os.mkdir(FLAGS.outdir)
    convert(FLAGS.labeldir, FLAGS.imagedir, FLAGS.outdir, FLAGS.lanenum, FLAGS.crop, FLAGS.resize_w, FLAGS.resize_h, FLAGS.debug)
    os.system('find ' + FLAGS.outdir + ' -size 0 | xargs -r -i rm "{}"') #delete the empty label
    os.system('ls ' + FLAGS.outdir + ' | sed s/txt/jpg/g >list.txt') #generate the data list


import os
label_dir = 'labels_fence'
output_dir = 'output'
files = os.listdir(label_dir)
for i, filename in enumerate(files):
    print("%d/%d:%s" % (i, len(files),filename))
    path = os.path.join(label_dir, filename)
    ok = True
    with open(path) as f:
        points = []
        while True:
            line = f.readline().strip()
            if line == '':
                break
            if len(line.split(' ')) != 4:
                os.system('echo %s >>error.txt' % filename)
                ok = False
                break
            x1,y1,x2,y2 = line.split(' ')
            points.append([float(x1),float(y1),float(x2),float(y2)])
        if not ok:
            continue
        out_path = os.path.join(output_dir, filename)
        with open(out_path, 'w') as o:
            for p in points:
                t = 4
                if (p[0] == 0 and p[1] == 0) or (p[2] == 0 and p[3] ==0):
                    t = -1
                o.write('3 %d %f %f %f\n' % (t, (p[0] + p[2]) / 2, p[1], p[2] - p[0]))
            #for p in points:
            #    t = 4
            #    if p[0] == 0 and p[1] == 0:
            #        t = -1
            #    o.write('4 %f %f %f %f\n' % (t, p[2], p[3], 0))

import os
label_dir = 'labels'
output_dir = 'labels_new'
files = os.listdir(label_dir)
for i, filename in enumerate(files):
    print("%d/%d:%s" % (i, len(files),filename))
    path = os.path.join(label_dir, filename)
    ok = True
    with open(path) as f:
        points = []
        lanes = {}
        has_left = False
        while True:
            line = f.readline().strip()
            if line == '':
                break
            if len(line.split(' ')) != 5:
                os.system('echo %s >>error.txt' % filename)
                ok = False
                break
            lane = line.split(' ')
            group, t, x, y, w = lane
            if has_left == False and group == '1' and t != '-1':
                has_left = True
            if group not in lanes:
                lanes[group] = [lane]
            else:
                lanes[group].append(lane)
        if not ok:
            continue
        out_path = os.path.join(output_dir, filename)
        with open(out_path, 'w') as o:
            for l in ['1','0','2']:
                for p in lanes[l]:
                    g = 0
                    if has_left:
                        if l == '0':
                            g = 1
                        elif l == '1':
                            g = 0
                        else:
                            g = 2
                    else:
                        if l == '0':
                            g = 0
                        elif l == '1':
                            g = 0
                        else:
                            g = 1
                    group, t, x, y, w = p
                    o.write('%d %s %s %s %s\n' % (g,t,x,y,w))
            #for p in points:
            #    t = 4
            #    if p[0] == 0 and p[1] == 0:
            #        t = -1
            #    o.write('4 %f %f %f %f\n' % (t, p[2], p[3], 0))

import os
import math
label_dir = 'labels_only_bollard_center'
output_dir = 'output'
files = os.listdir(label_dir)
for i, filename in enumerate(files):
    print("%d/%d:%s" % (i, len(files),filename))
    path = os.path.join(label_dir, filename)
    ok = True
    with open(path) as f:
        points = []
        while True:
            line = f.readline().strip()
            if line == '':
                break
            if len(line.split(' ')) != 4:
                os.system('echo %s >>error.txt' % filename)
                ok = False
                break
            g, t, xs,ys,ws = line.split(' ')

            x = float(xs)
            y = float(ys)
            w = float(ws)
            if t != 0 and x==0 and y ==0:
                print line
                exit()
            if x > 3 or y > 3 or x < -3 or y < -3 or w > 5 or w < 0:
                print line
                exit()

#! /usr/bin python
# -*- coding:utf-8 -*-
import cv2
import copy as cp
import os

infilename="/home/zhao/Documents/20180117_rain_chery_4.h2640658.jpg"
image = cv2.imread(infilename, 1)
img_row_start = 768#192#628#560 448#560#530    
img_height_in = 768#768#718#640#640
resize_height = 1536#960#1346#996#1200
resize_width =  1536#1920

'''
628   560
718   640
1346
1536

'''
'''
718   640
718   640
1436
1536

'''

'''
402   360
718   640
1120
1536

'''

image_cut = cv2.resize(image, (resize_width, resize_height))
image_resize = image_cut[img_row_start:img_row_start + img_height_in, :, :]
cv2.namedWindow('before', cv2.WINDOW_NORMAL)
cv2.resizeWindow("before", image_resize.shape[1]/2, image_resize.shape[0]/2);
print "before shape: " + str(image_resize.shape[0]) + " " + str(image_resize.shape[1])
cv2.imshow('before', image_resize)


infilename="/home/zhao/Documents/43465069059273.jpg"
image = cv2.imread(infilename, 1)
img_row_start = 432#168#424#168#424
img_height_in = 768#768#768#768  936
resize_height = 1200#1284
resize_width =  1536
image_cut = cv2.resize(image, (resize_width, resize_height))
image_resize = image_cut[img_row_start:img_row_start + img_height_in, :, :]
cv2.namedWindow('now', cv2.WINDOW_NORMAL)
cv2.resizeWindow("now", image_resize.shape[1]/2, image_resize.shape[0]/2);
print "now shape: " + str(image_resize.shape[0]) + " " + str(image_resize.shape[1])
cv2.imshow('now', image_resize)

'''
infilename="/home/zhao/Documents/43465069059273.jpg"
image = cv2.imread(infilename, 1)
img_row_start = 424#168#424
img_height_in = 511#768  936
resize_height = 1285
resize_width =  1536
image_cut = cv2.resize(image, (resize_width, resize_height))
image_resize = image_cut[img_row_start:img_row_start + img_height_in, :, :]
image_cut = cv2.resize(image, (resize_width, 768))
cv2.namedWindow('now_', cv2.WINDOW_NORMAL)
cv2.resizeWindow("now_", image_resize.shape[1]/2, image_resize.shape[0]/2);
print "now_ shape: " + str(image_resize.shape[0]) + " " + str(image_resize.shape[1])
cv2.imshow('now_', image_resize)
'''
cv2.waitKey(0)