| 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)
|
|
|
|
|
| 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:
|
|
|
| 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()
|
|
|
|
|
| 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()
|
|
|
|
|
| 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.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 |
