Spaces:
Building
Building
File size: 5,370 Bytes
4187c6f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
import argparse
import os
from pathlib import Path
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.patches import Patch
import pandas as pd
import numpy as np
import tqdm
from ..bev.get_bev import mask2rgb, PRETTY_COLORS as COLORS, VIS_ORDER
from ..fpv.filters import haversine_np, angle_dist
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--dataset_dir", '-d', type=str, required=True, help="Dataset directory")
parser.add_argument("--locations", '-l', type=str, default="all",
help="Location names in CSV format. Set to 'all' to traverse all locations.")
parser.add_argument("--rows", type=int, default=5, help="How many samples per PDF page")
parser.add_argument("--n_samples", type=int, default=30, help="How many samples to visualize?")
parser.add_argument("--store_sat", action="store_true", help="Add sattelite column")
args = parser.parse_args()
MAX_ROWS = args.rows
MAX_COLS = 4 if args.store_sat else 3
MAX_TEXT_LEN=30
locations = list()
if args.locations.lower() == "all":
locations = os.listdir(args.dataset_dir)
locations = [l for l in locations if os.path.isdir(os.path.join(args.dataset_dir, l))]
else:
locations = args.locations.split(",")
print(f"Parsing {len(locations)} locations..")
all_locs_stats = dict()
for location in tqdm.tqdm(locations):
dataset_dir = Path(args.dataset_dir)
location_dir = dataset_dir / location
semantic_mask_dir = location_dir / "semantic_masks"
sat_dir = location_dir / "sattelite"
comp_dir = location_dir / "images"
pq_name = 'image_metadata_filtered_processed.parquet'
df = pd.read_parquet(location_dir / pq_name)
# Calc derrivative attributes
df["loc_descrip"] = haversine_np(
lon1=df["geometry.long"], lat1=df["geometry.lat"],
lon2=df["computed_geometry.long"], lat2=df["computed_geometry.lat"]
)
df["angle_descrip"] = angle_dist(
df["compass_angle"],
df["computed_compass_angle"]
)
with PdfPages(location_dir / 'compare.pdf') as pdf:
# Plot legend page
plt.figure()
key2mask_i = dict(zip(COLORS.keys(), range(len(COLORS))))
patches = [Patch(color=COLORS[key], label=f"{key}") for i,key in enumerate(VIS_ORDER) if COLORS[key] is not None]
plt.legend(handles=patches, loc='center', title='Legend')
plt.axis("off")
plt.tight_layout()
pdf.savefig()
plt.close()
# Plot pairs
row_cnt = 0
fig = plt.figure(figsize=(MAX_COLS*2, MAX_ROWS*2))
for index, row in tqdm.tqdm(df.iterrows()):
id = row["id"]
mask_fp = semantic_mask_dir / f"{id}.npz"
comp_fp = comp_dir / f"{id}_undistorted.jpg"
sat_fp = sat_dir / f"{id}.png"
if not os.path.exists(mask_fp) or not os.path.exists(comp_fp) or \
(args.store_sat and not os.path.exists(sat_fp)):
continue
plt.subplot(MAX_ROWS, MAX_COLS, (row_cnt % MAX_ROWS)*MAX_COLS + 1)
plt.axis("off")
desc = list()
# Display attributes
keys = ["geometry.long", "geometry.lat", "compass_angle",
"loc_descrip", "angle_descrip",
"make", "model", "camera_type",
"quality_score"]
for k in keys:
v = row[k]
if isinstance(v, float):
v = f"{v:.4f}"
bullet = f"{k}: {v}"
if len(bullet) > MAX_TEXT_LEN:
bullet = bullet[:MAX_TEXT_LEN-2] + ".."
desc.append(bullet)
plt.text(0,0, "\n".join(desc), fontsize=7)
plt.title(id)
plt.subplot(MAX_ROWS, MAX_COLS, (row_cnt % MAX_ROWS)*MAX_COLS + 2)
mask = np.load(mask_fp)["arr_0"]
mask_rgb = mask2rgb(mask)
plt.imshow(mask_rgb); plt.axis("off")
plt.title(f"BEV")
H,W,_ = mask_rgb.shape
plt.scatter(np.array([H/2]), np.array([W/2]), marker="x")
plt.subplot(MAX_ROWS, MAX_COLS, (row_cnt % MAX_ROWS)*MAX_COLS + 3)
plt.imshow(plt.imread(comp_fp)); plt.axis("off")
plt.title(f"FPV")
if args.store_sat:
sat_fp = sat_dir / f"{id}.png"
plt.subplot(MAX_ROWS, MAX_COLS, (row_cnt % MAX_ROWS)*MAX_COLS + 4)
plt.imshow(plt.imread(sat_fp)); plt.axis("off")
plt.title(f"SAT")
row_cnt += 1
if row_cnt % MAX_ROWS == 0:
#plt.suptitle(location)
plt.tight_layout()
fig.align_titles()
pdf.savefig()
plt.close()
fig = plt.figure(figsize=(MAX_COLS*2, MAX_ROWS*2))
if row_cnt == args.n_samples:
break |