韩国PR2数据集和UM数据集

annotation_to_voc.py

# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import os
import cv2
from skimage import io, draw
import os
import re
import random
import lxml.etree as ElementTree
import dicttoxml
from xml.dom.minidom import parseString
from collections import OrderedDict

MIN_X = 1
MIN_Y = 1
MAX_X = 640
MAX_Y = 480
BASE_DIR = "/home/haoyu/VOC2007"
annotation_dir = "/home/haoyu/pr2/umannotations/_2011-06-20-15-36-20_0"


class Target:
    def __init__(self, data):
        self.data = data.values
        self.time_stamp = (self.data[:, 0]).ravel()
        self.time_stamp_length = len(self.time_stamp)
        self.time_stamp_min = np.min(self.time_stamp)
        self.time_stamp_max = np.max(self.time_stamp)

    def __linear_box(self, pos, timestamp):
        # 两帧之间做一个线性插值
        t_0 = self.data[pos - 1, 0]
        t_1 = self.data[pos, 0]
        scale = (timestamp - t_0) / (t_1 - t_0)
        return self.data[pos - 1, 1:] + (self.data[pos, 1:] - self.data[pos - 1, 1:]) * scale

    def get_box(self, timestamp):
        pos = np.searchsorted(self.time_stamp, timestamp)
        if pos <= 0 or pos >= self.time_stamp_length:
            return None
        if np.abs(self.time_stamp[pos] - timestamp) > 0.3 and \
                        np.abs(self.time_stamp[pos - 1] - timestamp) > 0.3:
            return None
        return self.__linear_box(pos, timestamp)


def get_filenames(file_dir):
    filenames = os.listdir(file_dir)
    filenames.sort()
    return filenames


def read_annotation(anno_dir):
    annotations = list()
    filenames = get_filenames(anno_dir)

    for filename in filenames:
        file_path = os.path.join(anno_dir, filename)
        t = pd.read_csv(
            file_path,
            sep=" ",
            header=None,
            names=['time', 'x_3d', 'y_3d', 'z_3d', 'x', 'y', 'width', 'height'],
            # dtype={
            #     0: np.string_
            # }
        )
        tt = t.loc[:, ['time', 'x', 'y', 'width', 'height']]
        # 新生成一个Target对象扔进去
        annotations.append(Target(tt))

    return annotations


def mark(image_dir, targets):
    filename_jar = list()
    sets = {
        'train': 0.8,
        'val': 0.1,
        'test': 0.1,
    }
    files = get_filenames(image_dir)

    for one_file in files:
        # img = io.imread(os.path.join(image_dir, one_file))
        # img = cv2.imread(os.path.join(image_dir, one_file))
        tm, _ = os.path.splitext(one_file)
        tsm = np.float(tm)
        filename_jar.append(tm)
        AnnotationFile = os.path.join(
            BASE_DIR, "Annotations",
            "{}{}".format(tm, ".xml")
        )
        data = {
            'folder': "VOC2007",
            'filename': "{}{}".format(tm, ".png"),
            'size': {
                'width': 640,
                'height': 480,
                'depth': 3,
            },
            'segmented': 0
        }
        xml = dicttoxml.dicttoxml(OrderedDict(data), attr_type=False, custom_root='annotation')
        dom = parseString(xml)

        for tar in targets:
            box = tar.get_box(tsm)
            if box is not None:
                x, y, w, h = box
                x1 = int(max(x, MIN_X))
                y1 = int(max(y, MIN_Y))
                x2 = int(min(x + w, MAX_X))
                y2 = int(min(y + h, MAX_Y))
                obj = {
                    'name': 'person',
                    'pose': 'Left',
                    'truncated': 1,
                    'difficult': 0,
                    'bndbox': {
                        'xmin': x1,
                        'ymin': y1,
                        'xmax': x2,
                        'ymax': y2,
                    }
                }
                xml_obj = parseString(dicttoxml.dicttoxml(
                    OrderedDict(obj),
                    attr_type=False,
                    custom_root='object')
                )
                x = dom.importNode(xml_obj.childNodes[0], True)
                dom.childNodes[0].appendChild(x)

        with open(AnnotationFile, "w") as anno:
            print(dom.toprettyxml(), file=anno)
            # cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 0), 2)
            # cv2.imshow("img", img)
            # cv2.waitkey(10)
            #         print(x, y, w, h)
            #         rr, cc = draw.polygon_perimeter([y, y, y + h, y + h],
            #                                         [x, x + w, x + w, x],
            #                                         shape=img.shape, clip=True)
            #         img[rr, cc] = [255, 0, 0]
            # io.imshow(img)
            # io.show()
            # 分数据集
        total = len(filename_jar)
        random.shuffle(filename_jar)
        sets_counter = 0
        for (set_name, set_scale) in sets.items():
            with open(os.path.join(BASE_DIR, "ImageSets/Main/", set_name + ".txt"), 'w') as st:
                tot = int(total * set_scale)
                for ith in range(sets_counter, sets_counter + tot):
                    print("{}".format(filename_jar[ith]), file=st)
                sets_counter += tot

        with open(os.path.join(BASE_DIR, "ImageSets/Main/trainval.txt"), 'w') as train_val:
            for set_name in ["train", "val"]:
                for line in open(os.path.join(BASE_DIR, "ImageSets/Main/", set_name + ".txt"), 'r'):
                    print(line, end="", file=train_val)


annotation = read_annotation(annotation_dir)
mark(os.path.join(BASE_DIR, 'JPEGImages-1'), annotation)

dat_to_png3.py

import os
import numpy as np
import skimage.color as color
import skimage.io as io

dtype = np.dtype([
    ('fcc', np.dtype('a4')),
    ('nrows', np.uint16),
    ('ncols', np.uint16),
    ('nsize', np.uint16),
    ('data', np.dtype((np.float32, (480, 640))))
])

MAX_RANGE = 10


def convert_dat_to_png3(dat_filename, png_filename):
    with open(dat_filename, 'rb') as f:
        data = np.fromfile(f, dtype=dtype)
        float_img = data['data'][0, :, :]
        float_img = np.nan_to_num(float_img) / MAX_RANGE
        float_img = color.gray2rgb(float_img)
        io.imsave(png_filename, float_img)


def generate_filename(dat_dir, png_dir):
    for parent, dirnames, filenames in os.walk(dat_dir):
        for filename in filenames:
            f, _ = os.path.splitext(filename)
            yield (os.path.join(dat_dir, filename), os.path.join(png_dir, f + ".png"))


dir_dat = '/home/haoyu/pr2/pr2dataset/2011-06-22-17-06-38_0/2011-06-22-17-06-38_0_depth'
dir_dst = '/home/haoyu/pr2/pr2dataset/2011-06-22-17-06-38_0/2011-06-22-17-06-38_0_depth_png3'

counter = 0
for (dat_file, png_file) in generate_filename(dir_dat, dir_dst):
    counter += 1
    print(counter)
    try:
        convert_dat_to_png3(dat_file, png_file)
    except:
        print(dat_file)
        print(png_file)

filename_equal.py

# -*- coding: utf-8 -*-
import numpy as np
import os
import shutil

# 参考文件名
dir_ref = "/home/haoyu/Data/RGB/"
# 需要修改的文件名
dir_ori = "/home/haoyu/Data/Depth3/"
# 目的地址
dir_dst = "/home/haoyu/Data/Depth3_rename/"

# 记录原始文件名
fn_name = list()
# 记录数字
fn_ori = list()


# 找出最近邻的数字
def find_nearest(array, value):
    idx = (np.abs(array - value)).argmin()
    return array[idx], idx


def generate_filename(dat_dir):
    for parent, dirnames, filenames in os.walk(dat_dir):
        for filename in filenames:
            f, _ = os.path.splitext(filename)
            yield f


# 读出参考文件名
for x in generate_filename(dir_ref):
    fn_name.append(x)
    fn_ori.append(np.double(x))
fn_ori = np.array(fn_ori)

# 移动文件
foo = list()
err = 0
yes = 0
for x in generate_filename(dir_ori):
    val, idx = find_nearest(fn_ori, np.double(x))
    frm = os.path.join(dir_ori, x + ".png")
    dst = os.path.join(dir_dst, fn_name[idx] + ".png")
    print(x, "=>", fn_name[idx])
    shutil.copy2(frm, dst)
    if fn_name[idx] in foo:
        print("[dup]")
        err += 1
    else:
        foo.append(fn_name[idx])
        yes += 1
print(err, "/", yes + err, ",", err / (yes + err), "%")