Update helpers/functional.py

helpers/functional.py

@@ -46,6 +46,22 @@ df_dict = {
     'COP-DEM' : dem_df
 }
 
+'''
+    --- SigLIP Search ---
+'''
+from TextSearch import *
+
+BASE_DIR='./'
+INDEX_OUTPUT = f'{BASE_DIR}/siglip_ivfpq.index'
+METADATA_OUTPUT = f'{BASE_DIR}/siglip_ivfpq_metadata.parquet'
+
+search = SearchSigLIP(index_path=INDEX_OUTPUT,
+                      metadata_path=METADATA_OUTPUT)
+
+'''
+    --- Helper Functions ---
+'''
+
 def pretty_date(input):
     template = '%Y%m%dT%H%M%S' if 'T' in input else '%Y%m%d%H%M%S'
     return datetime.datetime.strptime(input, template).strftime('%H:%M:%S - %d %b %Y')
@@ -91,7 +107,7 @@ def cell2row(grid_string, meta_df, return_row = False):
             return R.parquet_url.item(), R.parquet_row.item()
     else:
         return None
-
+        
 def map_to_image(map, return_centre=False, return_gridcell=False, return_timestamp=False, source='Sentinel-2 L2A'):
 
     try:
@@ -104,9 +120,28 @@ def map_to_image(map, return_centre=False, return_gridcell=False, return_timesta
     # 2. translate coordinate to major-tom tile
     rows, cols = grid.latlon2rowcol([center[0]], [center[1]])
 
+    return cell_to_image(cell="{}_{}".format(rows[0],cols[0]),
+                         return_centre=return_centre,
+                         return_gridcell=return_gridcell,
+                         return_timestamp=return_timestamp,
+                         source=source)
+
+def text_to_image(text, return_centre=False, return_gridcell=False, return_timestamp=False, source='Sentinel-2 L2A'):
+
+    cell = search.faiss(text)[0]["grid_cell"]
+
+    return cell_to_image(cell=cell,
+                         return_centre=return_centre,
+                         return_gridcell=return_gridcell,
+                         return_timestamp=return_timestamp,
+                         source=source)
+    
+
+def cell_to_image(cell, return_centre=False, return_gridcell=False, return_timestamp=False, source='Sentinel-2 L2A'):
+
     # 3. translate major-tom cell to row in parquet
     df = df_dict[source]
-    row = cell2row("{}_{}".format(rows[0],cols[0]), df, return_row = True)
+    row = cell2row(cell, df, return_row = True)
 
     if row is not None:
         parquet_url, parquet_row, meta_row = row
@@ -115,6 +150,7 @@ def map_to_image(map, return_centre=False, return_gridcell=False, return_timesta
             img = row2image(parquet_url, parquet_row)
         except:
             return None
+            
         # 4. acquire image # X. update map
         lat, lon = meta_row.centre_lat.item(), meta_row.centre_lon.item()