lcnn/misc/draw-wireframe.py

#!/usr/bin/env python3
import os
import glob
import os.path as osp

import cv2
import numpy as np
import scipy.io
import matplotlib as mpl
import matplotlib.pyplot as plt

from lcnn.utils import parmap

GT = "data/wireframe/valid/*.npz"
WF = "/data/lcnn/wirebase/result/wireframe/wireframe_1_rerun-baseline_0.5_0.5/2/*.mat"
AFM = "/data/lcnn/wirebase/result/wireframe/afm/*.npz"
IMGS = "/data/lcnn/wirebase/Wireframe/v1.1/test/*.jpg"


def imshow(im):
    sizes = im.shape
    height = float(sizes[0])
    width = float(sizes[1])

    fig = plt.figure()
    fig.set_size_inches(width / height, 1, forward=False)
    ax = plt.Axes(fig, [0.0, 0.0, 1.0, 1.0])
    ax.set_axis_off()
    fig.add_axes(ax)
    plt.xlim([-0.5, sizes[1] - 0.5])
    plt.ylim([sizes[0] - 0.5, -0.5])
    plt.imshow(im)


def main():
    gts = sorted(glob.glob(GT))
    afm = sorted(glob.glob(AFM))
    wf = sorted(glob.glob(WF))
    img = sorted(glob.glob(IMGS))

    prefix = "/data/lcnn/wirebase/myplot/"
    os.makedirs(osp.join(prefix, "GT"), exist_ok=True)
    os.makedirs(osp.join(prefix, "LSD"), exist_ok=True)
    os.makedirs(osp.join(prefix, "AFM"), exist_ok=True)
    os.makedirs(osp.join(prefix, "WF"), exist_ok=True)
    os.makedirs(osp.join(prefix, "LL"), exist_ok=True)

    def draw(args):
        i, (wf_name, gt_name, afm_name, img_name) = args
        img = cv2.imread(img_name, 0)
        lsd = cv2.createLineSegmentDetector(cv2.LSD_REFINE_ADV)
        lsd_line, _, _, lsd_score = lsd.detect(img)
        lsd_line = lsd_line.reshape(-1, 2, 2)
        lsd_score = lsd_score.flatten()
        img = cv2.imread(img_name)[:, :, ::-1]

        with np.load(gt_name) as fgt:
            gt_line = fgt["lpos"][:, :, :2]
            gt_line[:, :, 0] *= img.shape[0] / 128
            gt_line[:, :, 1] *= img.shape[1] / 128

        with np.load(afm_name) as fafm:
            afm_line = fafm["lines"].reshape(-1, 2, 2)[:, :, ::-1]

        wf_line = scipy.io.loadmat(wf_name)["lines"].reshape(-1, 2, 2)
        wf_line = wf_line[:, :, ::-1]

        plt.figure("GT")
        imshow(img)
        for a, b in gt_line - 0.5:
            plt.plot([a[1], b[1]], [a[0], b[0]], color="orange", linewidth=0.5)
            plt.scatter(a[1], a[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
            plt.scatter(b[1], b[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
        plt.savefig(osp.join(prefix, "GT", f"{i:05}"), dpi=3000, bbox_inches=0)
        plt.close()

        plt.figure("LSD")
        imshow(img)
        for a, b in lsd_line[:, :, ::-1] - 0.5:
            plt.plot([a[1], b[1]], [a[0], b[0]], color="orange", linewidth=0.5)
            plt.scatter(a[1], a[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
            plt.scatter(b[1], b[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
        plt.savefig(osp.join(prefix, "LSD", f"{i:05}"), dpi=3000, bbox_inches=0)
        plt.close()

        plt.figure("AFM")
        imshow(img)
        for a, b in afm_line - 0.5:
            plt.plot([a[1], b[1]], [a[0], b[0]], color="orange", linewidth=0.5)
            plt.scatter(a[1], a[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
            plt.scatter(b[1], b[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
        plt.savefig(osp.join(prefix, "AFM", f"{i:05}"), dpi=3000, bbox_inches=0)
        plt.close()

        plt.figure("WF")
        imshow(img)
        for a, b in wf_line - 0.5:
            plt.plot([a[1], b[1]], [a[0], b[0]], color="orange", linewidth=0.5)
            plt.scatter(a[1], a[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
            plt.scatter(b[1], b[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
        plt.savefig(osp.join(prefix, "WF", f"{i:05}"), dpi=3000, bbox_inches=0)
        plt.close()

    parmap(draw, enumerate(zip(wf, gts, afm, img)))


if __name__ == "__main__":
    main()