Update streamlit_app.py

Pre-generate thumbnails and save in AWS S3 bucket. Read thumbnails' url in click function.

src/streamlit_app.py

@@ -142,94 +142,6 @@ chips_df = pd.read_csv("data/embeddings_df_v0.11_test.csv")
 # set anonymous S3FileSystem to read files from public bucket 
 s3 = s3fs.S3FileSystem(anon=True)
 
-## helper function
-def gen_chip_urls(row,  s3_prefix):
-    '''
-    Generate S3 urls for chips
-    :param row: dictionary with chip_id and dates
-    :param s3_prefix: S3 url prefix 
-    :return s3_urls: a list of urls
-    '''
-    s3_urls = []
-    dates = ast.literal_eval(row["dates"])
-    for date in dates:
-        filename = f"s2_{row['chip_id']:06}_{date}.tif"
-        s3_url = f"{s3_prefix}/{filename}"
-        s3_urls.append(s3_url)
-    return s3_urls
-
-def mask_nodata(band, nodata_values=(-999,)):
-    '''
-    Mask nodata to nan
-    :param band
-    :param nodata_values:nodata values in chips is -999
-    :return band
-    '''
-    band = band.astype(float)
-    for val in nodata_values:
-        band[band == val] = np.nan
-    return band
-
-def normalize(band):
-    '''
-    Normalize a band to 0-1 range(float)
-    :param band (ndarray)
-    return normalize band (ndarray); when max equals min, returns zeros.
-    '''
-    if np.nanmean(band) >= 4000:
-        band = band / 6000
-    else:
-        band = band / 4000
-    band = np.clip(band, None, 1)
-
-    return band
-
-def create_thumbnail(url):
-    '''
-    Read S3 file into memory, using rasterio to create a png thumbnail then encode as a base64 string url
-    :param url: chip url
-    :return a base64-encoded png string, returns an empty string when an error occurs
-    '''
-    try:
-        # read raw bytes from s3 file
-        with s3.open(url, "rb") as f:
-            data = f.read()
-
-        # wrap the raw bytes into an memory file
-        with MemoryFile(data) as memfile:
-            
-            # read memory file with rasterio
-            with memfile.open() as src:
-                # mask nodata to have correct calculate normalization
-                # band1->blue, band2->green, band3->red
-
-                blue = src.read(1).astype(float)
-                green = src.read(2).astype(float)
-                red = src.read(3).astype(float)
-
-                blue = normalize(mask_nodata(blue))
-                green = normalize(mask_nodata(green))
-                red = normalize(mask_nodata(red))
-
-                # stack in RGB
-                rgb = np.dstack((red, green, blue))
-
-                # convert float(0-1) to uint8 (0-255)
-                rgb_8bit = (rgb * 255).astype(np.uint8)
-
-                # convert to png in memory
-                pil_img = Image.fromarray(rgb_8bit)
-                buf = io.BytesIO()
-                pil_img.save(buf, format='PNG')
-
-                # encoded into base64 HTML format
-                encoded = base64.b64encode(buf.getvalue()).decode('utf-8')
-                return f"data:image/png;base64,{encoded}"
-
-    except Exception as e:
-        # return an empty string for Exception
-        return ""
-
 def get_lat(geometry):
     lat = wkt.loads(geometry).coords.xy[1][0]
 
@@ -240,7 +152,6 @@ def get_lon(geometry):
     
     return lon
 
-
 ## generate json
 # title: plot title
 # xaxis_title: x axis title
@@ -255,12 +166,6 @@ title_js = json.dumps(config["title"])
 xaxis_js = json.dumps(config["xaxis_title"])
 yaxis_js = json.dumps(config["yaxis_title"])
 
-# set prefix
-s3_prefix="s3://gfm-bench"
-
-# generate S3 file URLs 
-chips_df["urls"] = chips_df.apply(lambda row: gen_chip_urls(row, s3_prefix), axis=1)
-
 # set lc(str) for categorical data for plotting
 chips_df["lc"] = chips_df["lc"].astype(str)
 # add latitude and longitude
@@ -300,13 +205,21 @@ color_dict_label = {
     'Rangeland': '#f7980a'
 }
 
-# create thumbnail 
-chips_df["thumbs"] = chips_df["urls"].apply(
-    lambda urls: [create_thumbnail(p) for p in urls]
-)
 # create dates Python list
 chips_df["dates_list"] = chips_df["dates"].apply(ast.literal_eval)
 
+# set prefix
+s3_url="https://gfm-bench.s3.amazonaws.com/thumbnails"
+
+# create thumb_urls column
+chips_df["thumb_urls"] = chips_df.apply(
+    lambda r: [
+        f"{s3_url}/s2_{r.chip_id:06}_{date}.png"
+        for date in r.dates_list
+    ],
+    axis=1
+)
+
 # build a list of points dictionary
 points = (
     chips_df
@@ -314,12 +227,13 @@ points = (
         "cls_dim1": "x",
         "cls_dim2": "y",
         "Land Cover": "category"
-    })[["x","y","chip_id", "latitude", "longitude","category","thumbs","dates_list"]]
+    })[["x","y","chip_id", "latitude", "longitude","category","dates_list"]]
     .assign(
         id = chips_df["chip_id"], 
         lat = chips_df["latitude"],
         lon = chips_df["longitude"],
-        color=chips_df["Land Cover"].map(color_dict_label))
+        color=chips_df["Land Cover"].map(color_dict_label),
+        thumbs = chips_df["thumb_urls"])
     .to_dict(orient="records")
 )
 
@@ -380,8 +294,7 @@ plot_html = f"""
         customdata:pts.map(p=>[
                             p.id,        
                             p.lat,       
-                            p.lon,      
-                            p.thumbs     
+                            p.lon   
                             ]),
         mode:      'markers',
         type:      'scatter',
@@ -438,8 +351,7 @@ plot_html = f"""
     gd.on('plotly_click', evt => {{
       // grab thumbs and dates through point index
       const idx = evt.points[0].pointIndex;
-      const cds   = evt.points[0].customdata;
-      const thumbs = cds[3];
+      const thumbs = points[idx].thumbs; 
       const dates  = points[idx].dates_list;
       // grab image container and clear out old thumbs
       const container = document.getElementById('image-container');