added robot position, toggle how many to display, just show data pair

main.py

@@ -9,6 +9,9 @@ import matplotlib.pyplot as plt
 import pandas as pd
 import geopandas as gpd
 from pyproj.transformer import Transformer
+import cv2
+from matplotlib import patches as mpatches
+from matplotlib import gridspec
 
 sys.path.append(os.path.dirname(os.path.realpath(__file__)))
 from MapItAnywhere.mia.bev import get_bev
@@ -72,8 +75,7 @@ def split_dataframe(df, chunk_size = 100):
         chunks.append(df[i*chunk_size:(i+1)*chunk_size])
     return chunks
 
-async def fetch(location, filter_undistort, disable_cam_filter, map_length, mpp):
-    N=1
+async def fetch(location, num_images, filter_undistort, disable_cam_filter, map_length, mpp):
     TOTAL_LOOKED_INTO_LIMIT = 10000
 
     ################ FPV
@@ -109,36 +111,37 @@ async def fetch(location, filter_undistort, disable_cam_filter, map_length, mpp)
                 
                 dfs_meta.append(df_meta)
                 total_rows = sum([len(x) for x in dfs_meta])
-                if total_rows > N:
+                if total_rows > num_images:
                     break
                 elif total_looked_into > TOTAL_LOOKED_INTO_LIMIT:
                     yield (f"Went through {total_looked_into} images and could not find images satisfying the filters."
                            "\nPlease rerun or run the data engine locally for bulk time consuming operations.", None, None)
                     return
-            if total_rows > N:
+            if total_rows > num_images:
                 break
         except:
             pass
 
     df_meta = pd.concat(dfs_meta)
-    df_meta = df_meta.sample(N)
+    df_meta = df_meta.sample(num_images)
 
     # Calc derrivative attributes
-    df_meta["loc_descrip"] = filters.haversine_np(
+    df_meta["loc_discrepancy"] = filters.haversine_np(
         lon1=df_meta["geometry.long"], lat1=df_meta["geometry.lat"],
         lon2=df_meta["computed_geometry.long"], lat2=df_meta["computed_geometry.lat"]
     )
 
-    df_meta["angle_descrip"] = filters.angle_dist(
+    df_meta["angle_discrepancy"] = filters.angle_dist(
         df_meta["compass_angle"],
         df_meta["computed_compass_angle"]
     )
-
+    img_list_to_show = list()
     for index, row in df_meta.iterrows():
+        print("Processing image", row["id"])
         desc = list()
         # Display attributes
         keys = ["id", "geometry.long", "geometry.lat", "compass_angle",
-                "loc_descrip", "angle_descrip",
+                "loc_discrepancy", "angle_discrepancy",
                 "make", "model", "camera_type", 
                 "quality_score"]
         for k in keys:
@@ -148,96 +151,166 @@ async def fetch(location, filter_undistort, disable_cam_filter, map_length, mpp)
             bullet = f"{k}: {v}"
             desc.append(bullet)
         metadata_fmt = "\n".join(desc)
-    
-    yield metadata_fmt, None, None
-    image_urls = list(df_meta.set_index("id")["thumb_2048_url"].items())
-    num_fail = await get_fpv.fetch_images_pixels(image_urls, downloader, raw_image_dir)
-    if num_fail > 0:
-        logger.error(f"Failed to download {num_fail} images.")
-    
-    seq_to_image_ids = df_meta.groupby('sequence')['id'].agg(list).to_dict()
-    lon_center = (bbox['east'] + bbox['west']) / 2
-    lat_center = (bbox['north'] + bbox['south']) / 2
-    projection = get_fpv.Projection(lat_center, lon_center, max_extent=200e3)
+        # yield metadata_fmt, None, None
+        image_urls = list(df_meta.set_index("id")["thumb_2048_url"].items())
+        num_fail = await get_fpv.fetch_images_pixels(image_urls, downloader, raw_image_dir)
+        if num_fail > 0:
+            logger.error(f"Failed to download {num_fail} images.")
+        
+        seq_to_image_ids = df_meta.groupby('sequence')['id'].agg(list).to_dict()
+        lon_center = (bbox['east'] + bbox['west']) / 2
+        lat_center = (bbox['north'] + bbox['south']) / 2
+        projection = get_fpv.Projection(lat_center, lon_center, max_extent=200e3)
 
-    df_meta.index = df_meta["id"]
-    image_infos = df_meta.to_dict(orient="index")
-    process_sequence_args = get_fpv.default_cfg
+        df_meta.index = df_meta["id"]
+        image_infos = df_meta.to_dict(orient="index")
+        process_sequence_args = get_fpv.default_cfg
 
-    if filter_undistort:
-        for seq_id, seq_image_ids in seq_to_image_ids.items():
-            try: 
-                d, pi = get_fpv.process_sequence(
-                    seq_image_ids,
-                    image_infos,
-                    projection,
-                    process_sequence_args,
-                    raw_image_dir,
-                    out_image_dir,
-                )
-                if d is None or pi is None:
-                    raise Exception("process_sequence returned None")
-            except Exception as e:
-                logger.error(f"Failed to process sequence {seq_id} skipping it. Error: {repr(e)}.")
-        
-        fpv = plt.imread(out_image_dir/ f"{row['id']}_undistorted.jpg")
-    else:
-        fpv = plt.imread(raw_image_dir/ f"{row['id']}.jpg")
-    yield metadata_fmt, fpv, None
-    ################ BEV
-    df = df_meta
-    # convert pandas dataframe to geopandas dataframe
-    gdf = gpd.GeoDataFrame(df, 
-                        geometry=gpd.points_from_xy(
-                            df['computed_geometry.long'], 
-                            df['computed_geometry.lat']), 
-                            crs=4326)
-        
-    # convert the geopandas dataframe to UTM
-    utm_crs = gdf.estimate_utm_crs()
-    gdf_utm = gdf.to_crs(utm_crs)
-    transformer = Transformer.from_crs(utm_crs, 4326)
-        # load OSM data, if available
-    padding = 50
-    # calculate the required distance from the center to the edge of the image 
-    # so that the image will not be out of bounds when we rotate it
-    map_length = map_length
-    map_length = np.ceil(np.sqrt(map_length**2 + map_length**2))
-    distance = map_length * mpp
+        if filter_undistort:
+            for seq_id, seq_image_ids in seq_to_image_ids.items():
+                try: 
+                    d, pi = get_fpv.process_sequence(
+                        seq_image_ids,
+                        image_infos,
+                        projection,
+                        process_sequence_args,
+                        raw_image_dir,
+                        out_image_dir,
+                    )
+                    if d is None or pi is None:
+                        raise Exception("process_sequence returned None")
+                except Exception as e:
+                    logger.error(f"Failed to process sequence {seq_id} skipping it. Error: {repr(e)}.")
+            
+            fpv = plt.imread(out_image_dir/ f"{row['id']}_undistorted.jpg")
+        else:
+            print("Loading raw image")
+            fpv = plt.imread(raw_image_dir/ f"{row['id']}.jpg")
+        # yield metadata_fmt, fpv, None
+        ################ BEV
+        df = df_meta
+        # convert pandas dataframe to geopandas dataframe
+        gdf = gpd.GeoDataFrame(df, 
+                            geometry=gpd.points_from_xy(
+                                df['computed_geometry.long'], 
+                                df['computed_geometry.lat']), 
+                                crs=4326)
+            
+        # convert the geopandas dataframe to UTM
+        utm_crs = gdf.estimate_utm_crs()
+        gdf_utm = gdf.to_crs(utm_crs)
+        transformer = Transformer.from_crs(utm_crs, 4326)
+            # load OSM data, if available
+        padding = 50
+        # calculate the required distance from the center to the edge of the image 
+        # so that the image will not be out of bounds when we rotate it
+        map_length = map_length
+        map_length = np.ceil(np.sqrt(map_length**2 + map_length**2))
+        distance = map_length * mpp
+
+        # create bounding boxes for each point
+        gdf_utm['bounding_box_utm_p1'] = gdf_utm.apply(lambda row: (
+            row.geometry.x - distance - padding,
+            row.geometry.y - distance - padding,
+        ), axis=1)
 
-    # create bounding boxes for each point
-    gdf_utm['bounding_box_utm_p1'] = gdf_utm.apply(lambda row: (
-        row.geometry.x - distance - padding,
-        row.geometry.y - distance - padding,
-    ), axis=1)
+        gdf_utm['bounding_box_utm_p2'] = gdf_utm.apply(lambda row: (
+            row.geometry.x + distance + padding,
+            row.geometry.y + distance + padding,
+        ), axis=1)
 
-    gdf_utm['bounding_box_utm_p2'] = gdf_utm.apply(lambda row: (
-        row.geometry.x + distance + padding,
-        row.geometry.y + distance + padding,
-    ), axis=1)
+        # convert the bounding box back to lat, long
+        gdf_utm['bounding_box_lat_long_p1'] = gdf_utm.apply(lambda row: transformer.transform(*row['bounding_box_utm_p1']), axis=1)
+        gdf_utm['bounding_box_lat_long_p2'] = gdf_utm.apply(lambda row: transformer.transform(*row['bounding_box_utm_p2']), axis=1)
+        gdf_utm['bbox_min_lat'] = gdf_utm['bounding_box_lat_long_p1'].apply(lambda x: x[0])
+        gdf_utm['bbox_min_long'] = gdf_utm['bounding_box_lat_long_p1'].apply(lambda x: x[1])
+        gdf_utm['bbox_max_lat'] = gdf_utm['bounding_box_lat_long_p2'].apply(lambda x: x[0])
+        gdf_utm['bbox_max_long'] = gdf_utm['bounding_box_lat_long_p2'].apply(lambda x: x[1])
+        gdf_utm['bbox_formatted'] = gdf_utm.apply(lambda row: f"{row['bbox_min_long']},{row['bbox_min_lat']},{row['bbox_max_long']},{row['bbox_max_lat']}", axis=1)
+
+        # iterate over the dataframe and get BEV images
+        jobs = gdf_utm[['id', 'bbox_formatted', 'computed_compass_angle']] # only need the id and bbox_formatted columns for the jobs
+        jobs = jobs.to_dict(orient='records').copy()
+            
+        get_bev.get_bev_from_bbox_worker_init(osm_cache_dir, bev_dir, semantic_mask_dir, rendered_mask_dir, 
+                                            "MapItAnywhere/mia/bev/styles/mia.yml", map_length, mpp, 
+                                            None, True, False, True, True, 1)
+        for job_dict in jobs:
+            get_bev.get_bev_from_bbox_worker(job_dict)
+        
+        bev = cv2.imread(rendered_mask_dir / f"{row['id']}.png")
+        bev = cv2.cvtColor(bev, cv2.COLOR_BGR2RGB)
+        print("BEV shape", bev.shape)
+        
+        img_list_to_show_i = [fpv, bev, metadata_fmt]
+        
+        img_list_to_show.append(img_list_to_show_i)
+    
+    # Make plt figure 
+    plt_row = len(img_list_to_show)
+    print("plt_row", plt_row)
+    plt_col = 3
+    
+    for i in range(plt_row):
+        
+        fpv, bev, metadata_fmt = img_list_to_show[i]
+        if i == 0:
+            imgs = [fpv, bev]
+            ratios = [i.shape[1] / i.shape[0] for i in imgs]  # W / H
+            ratios.append(0.5)  # Metadata
+            figsize = [sum(ratios) * 4.5, 4.5 * plt_row]
+            dpi = 100
+            fig, ax = plt.subplots(
+                plt_row, plt_col, figsize=figsize, dpi=dpi, gridspec_kw={"width_ratios": ratios}
+            )
+        
+        # Plot FPV image
+        if plt_row == 1:
+            ax0 = ax[0]
+            ax1 = ax[1]
+            ax2 = ax[2]
+        else:
+            ax0 = ax[i, 0]
+            ax1 = ax[i, 1]
+            ax2 = ax[i, 2]
+        ax0.imshow(fpv)
+        ax0.set_title("First Person View Image")
+        ax0.axis('off')
+        
+        # Plot BEV image
+        ax1.imshow(bev)
+        # Put a white upward triangle at the center of the image
+        ax1.scatter(bev.shape[1]//2, bev.shape[0]//2, s=200, c='white', marker='^', edgecolors='black')
+        ax1.set_title("Bird's Eye View Map")
+        ax1.axis('off')
 
-    # convert the bounding box back to lat, long
-    gdf_utm['bounding_box_lat_long_p1'] = gdf_utm.apply(lambda row: transformer.transform(*row['bounding_box_utm_p1']), axis=1)
-    gdf_utm['bounding_box_lat_long_p2'] = gdf_utm.apply(lambda row: transformer.transform(*row['bounding_box_utm_p2']), axis=1)
-    gdf_utm['bbox_min_lat'] = gdf_utm['bounding_box_lat_long_p1'].apply(lambda x: x[0])
-    gdf_utm['bbox_min_long'] = gdf_utm['bounding_box_lat_long_p1'].apply(lambda x: x[1])
-    gdf_utm['bbox_max_lat'] = gdf_utm['bounding_box_lat_long_p2'].apply(lambda x: x[0])
-    gdf_utm['bbox_max_long'] = gdf_utm['bounding_box_lat_long_p2'].apply(lambda x: x[1])
-    gdf_utm['bbox_formatted'] = gdf_utm.apply(lambda row: f"{row['bbox_min_long']},{row['bbox_min_lat']},{row['bbox_max_long']},{row['bbox_max_lat']}", axis=1)
+        # Add legend to BEV image
+        class_colors = {
+            'Road': (68, 68, 68),           # 0: Black
+            'Crossing': (244, 162, 97),     # 1; Red
+            'Sidewalk': (233, 196, 106),    # 2: Yellow
+            'Building': (231, 111, 81),     # 5: Magenta
+            'Terrain': (42, 157, 143),      # 7: Cyan
+            'Parking': (204, 204, 204),     # 8: Dark Grey
+        }
+        patches = [mpatches.Patch(color=[c/255.0 for c in color], label=label) for label, color in class_colors.items()]
+        ax1.legend(handles=patches, loc='upper center', bbox_to_anchor=(0.5, -0.05), ncol=3)
+        
+        # Plot metadata text
+        ax2.axis('off')
+        ax2.text(0.1, 0.5, metadata_fmt, fontsize=12, va='center', ha='left', wrap=True)
+        ax2.set_title("Metadata")
 
-    # iterate over the dataframe and get BEV images
-    jobs = gdf_utm[['id', 'bbox_formatted', 'computed_compass_angle']] # only need the id and bbox_formatted columns for the jobs
-    jobs = jobs.to_dict(orient='records').copy()
+    plt.tight_layout(pad=2.0)
+        
         
-    get_bev.get_bev_from_bbox_worker_init(osm_cache_dir, bev_dir, semantic_mask_dir, rendered_mask_dir, 
-                                          "MapItAnywhere/mia/bev/styles/mia.yml", map_length, mpp, 
-                                          None, True, False, True, True, 1)
-    for job_dict in jobs:
-        get_bev.get_bev_from_bbox_worker(job_dict)
+    # Save figure and then read 
+    fig_img_path = 'fpv_bev.png'
+    fig.savefig(fig_img_path)
+    fig_img = plt.imread(fig_img_path)
     
-    bev = plt.imread(rendered_mask_dir / f"{row['id']}.png")
 
-    yield metadata_fmt, fpv, bev
+    yield fig_img
 
 filter_pipeline = filters.FilterPipeline.load_from_yaml("MapItAnywhere/mia/fpv/filter_pipelines/mia.yaml")
 filter_pipeline.verbose=False
@@ -261,15 +334,16 @@ logger.info(f"Current working directory: {os.getcwd()}, listdir: {os.listdir('.'
 
 demo = gr.Interface(
     fn=fetch,
-    inputs=[gr.Text("Pittsburgh, PA, United States", label="Location"),
-            gr.Checkbox(value=False, label="Filter & Undistort"),
-            gr.Checkbox(value=False, label="Disable camera model filtering"),
+    inputs=[gr.Text("Pittsburgh, PA, United States", label="Location (City, {Optional: State,} Country)"), 
+            gr.Number(value=1, label="Number of Data Pairs to Generate (Max: 3)", minimum=1, maximum=3),
+            gr.Checkbox(value=False, label="Filter & Undistort (True in paper. Results in better robot position estimate, but slower.)"),
+            gr.Checkbox(value=False, label="Disable camera model filtering (Enabled in paper. Results in better quality labels, but slower.)"),
             gr.Slider(minimum=64, maximum=512, step=1, label="BEV Dimension", value=224),
             gr.Slider(minimum=0.1, maximum=2, label="Meters Per Pixel", value=0.5)],
-    outputs=[gr.Text(label="METADATA"), gr.Image(label="FPV"), gr.Image(label="BEV")],
-    title="MapItAnywhere (Data Engine)",
-    description="A demo showcasing samples of MIA's capability to retrieve FPV-BEV pairs worldwide."
-                "For bulk download/heavy filtering please visit the github and follow the instructions to run locally"
+    outputs=[gr.Image(label="Data Pair")],
+    title="MapItAnywhere (MIA) Data Engine",
+     
+    description="Use our Data Engine to sample first-person view images and bird's-eye view semantic map pairs from locations around the globe. Simply pick a location to see the results! For faster bulk downloads and more stringent filtering, visit our repository and follow the instructions to run the data curation locally."
 )
 
 logger.info("Starting server")