合并几张图

def awesome_plot(rgb_and_grey):
    """
    注意!输入是一个tuple!!函数的输入输出装进tuple里面比较安全。
    画图函数:输入几张图片(可以彩色或者灰色, 也可以是不同size), 拼成一张图片, HWC

    :param rgb_and_grey: tuple[0]:img, tuple[1]:gd, tuple[2]:pred ...
    :return: stacked ndarray
    """
    if not isinstance(rgb_and_grey, tuple) or len(rgb_and_grey) <=1 :
        raise  ValueError("input must be a tuple with len > 2")
    pics_list = []
    for pic in rgb_and_grey:
        if len(pic.shape) == 2:
            pic = pic.repeat(3).reshape(pic.shape[0], pic.shape[1], 3)
            pic = cv2.resize(pic, (320, 256))
            pics_list.append(pic)
        elif len(pic.shape) == 3:
            if pic.shape[2] == 3:
                pic = cv2.resize(pic, (320, 256))
                pics_list.append(pic)
            else:
                raise ValueError("input format must be HxWxC and C must be 3!")
        else:
            raise ValueError("wrong input: must be pics")

    # print("stacking ", len(pics_list), " pics~")

    if len(pics_list) == 2:
        res = np.hstack((pics_list[0], pics_list[1]))
        return res
    elif len(pics_list) == 3:
        res = np.hstack((pics_list[0], pics_list[1], pics_list[2]))
        return res
    elif len(pics_list) == 4:
        row1 = np.hstack((pics_list[0], pics_list[1]))
        row2 = np.hstack((pics_list[2], pics_list[3]))
        res = np.vstack((row1, row2))
        return res
    elif len(pics_list) == 6:
        row1 = np.hstack((pics_list[0], pics_list[1], pics_list[2]))
        row2 = np.hstack((pics_list[3], pics_list[4], pics_list[5]))
        # row3 = np.hstack((pics_list[4], pics_list[5]))
        res = np.vstack((row1, row2))
        return res
    else:
        raise ValueError("too man pics( more than 6)")
posted @ 2018-06-01 13:36  一条图图犬  阅读(214)  评论(0编辑  收藏  举报