annot_to_densitymap.py 文件源码

def get_density_map(data_name, save_data, show_image = False):
    assert(data_name.endswith(".xml"))
    #xml_data = data_name + '.xml'

    with open(data_name) as xml_d:
        doc = xmltodict.parse(xml_d.read())

    img = np.zeros((240,352), np.float32)

    def add_to_image(image, bbox):
        xmin = int(bbox['xmin'])
        ymin = int(bbox['ymin'])
        xmax = int(bbox['xmax'])
        ymax = int(bbox['ymax'])
        density = 1/ float((ymax - ymin) * (xmax - xmin))   
        image[ymin:ymax, xmin:xmax] += density
        return image

    for vehicle in doc['annotation']['vehicle']:
        add_to_image(img, vehicle['bndbox'])

    if show_image: 
        show_mask(img, data_name.replace('xml','jpg'))
    if save_data:
        img.tofile(data_name.replace('xml','desmap'))