Use bitwise operations to calculate RGB values where allowed to improve efficiency

README.md

@@ -42,6 +42,7 @@ Options:
                          (.mbtiles output only)
   --min-z INTEGER        Minimum zoom to tile (.mbtiles output only)
   --format [png|webp]    Output tile format (.mbtiles output only)
+  --chinaoffset          Use GCJ02 coordinates within China as bounding box
   -j, --workers INTEGER  Workers to run [DEFAULT=4]
   -v, --verbose
   --co NAME=VALUE        Driver specific creation options.See the

rio_rgbify/coords.py

@@ -0,0 +1,79 @@
+import math
+
+mercator_a = 6378137.0
+mercator_max = 20037508.342789244
+
+rad = 0.017453292519943295
+a = 6378245.0
+ee = 0.00669342162296594323
+
+
+def mercator_to_wgs(x, y):
+    lon = x / rad / mercator_a
+    lat = (math.pi * 0.5 - 2.0 * math.atan(math.exp(-y / mercator_a))) / rad
+
+    return [lon, lat]
+
+
+def wgs_to_mercator(lon, lat):
+    adjusted = lon if abs(lon) <= 180 else lon - _sign(lon) * 360
+
+    x = mercator_a * adjusted * rad
+    y = mercator_a * math.log(math.tan(math.pi * 0.25 + 0.5 * lat * rad))
+
+    if (x > mercator_max):
+        x = mercator_max
+    if (x < -mercator_max):
+        x = -mercator_max
+    if (y > mercator_max):
+        y = mercator_max
+    if (y < -mercator_max):
+        y = -mercator_max
+
+    return [x, y];
+
+
+def gcj_to_wgs(lon, lat):
+    if _out_of_china(lon, lat):
+        return [lon, lat]
+    dlat = _transformlat(lon - 105.0, lat - 35.0)
+    dlon = _transformlon(lon - 105.0, lat - 35.0)
+    radlat = lat * rad
+    magic = math.sin(radlat)
+    magic = 1 - ee * magic * magic
+    sqrtmagic = math.sqrt(magic)
+    dlat = dlat / ((a * (1 - ee)) / (magic * sqrtmagic)* rad)
+    dlon = dlon / (a / sqrtmagic * math.cos(radlat)* rad)
+    return [lon - dlon, lat - dlat]
+
+
+def _sign(x):
+    return -1 if x < 0 else (1 if x > 0 else 0)
+
+
+def _out_of_china(lon, lat):
+    return not (lon > 73.66 and lon < 135.05 and lat > 3.86 and lat < 53.55)
+
+
+def _transformlat(lon, lat):
+    ret = -100.0 + 2.0 * lon + 3.0 * lat + 0.2 * lat * lat + \
+          0.1 * lon * lat + 0.2 * math.sqrt(abs(lon))
+    ret += (20.0 * math.sin(6.0 * lon * math.pi) + 20.0 *
+            math.sin(2.0 * lon * math.pi)) * 2.0 / 3.0
+    ret += (20.0 * math.sin(lat * math.pi) + 40.0 *
+            math.sin(lat / 3.0 * math.pi)) * 2.0 / 3.0
+    ret += (160.0 * math.sin(lat / 12.0 * math.pi) + 320 *
+            math.sin(lat * math.pi / 30.0)) * 2.0 / 3.0
+    return ret
+
+
+def _transformlon(lon, lat):
+    ret = 300.0 + lon + 2.0 * lat + 0.1 * lon * lon + \
+          0.1 * lon * lat + 0.1 * math.sqrt(abs(lon))
+    ret += (20.0 * math.sin(6.0 * lon * math.pi) + 20.0 *
+            math.sin(2.0 * lon * math.pi)) * 2.0 / 3.0
+    ret += (20.0 * math.sin(lon * math.pi) + 40.0 *
+            math.sin(lon / 3.0 * math.pi)) * 2.0 / 3.0
+    ret += (150.0 * math.sin(lon / 12.0 * math.pi) + 300.0 *
+            math.sin(lon / 30.0 * math.pi)) * 2.0 / 3.0
+    return ret

rio_rgbify/encoders.py

@@ -41,9 +41,16 @@ def data_to_rgb(data, baseval, interval, round_digits=0):
 
     rgb = np.zeros((3, rows, cols), dtype=np.uint8)
 
-    rgb[2] = ((data / 256) - (data // 256)) * 256
-    rgb[1] = (((data // 256) / 256) - ((data // 256) // 256)) * 256
-    rgb[0] = ((((data // 256) // 256) / 256) - (((data // 256) // 256) // 256)) * 256
+    if data.min() >= 0:
+        udata = np.array(data, dtype=np.uint32)
+
+        rgb[0] = np.right_shift(udata, 16)
+        rgb[1] = np.right_shift(np.bitwise_and(udata, 0x00FF00), 8)
+        rgb[2] = np.bitwise_and(udata, 0x0000FF)
+    else:
+        rgb[2] = ((data / 256) - (data // 256)) * 256
+        rgb[1] = (((data // 256) / 256) - ((data // 256) // 256)) * 256
+        rgb[0] = ((((data // 256) // 256) / 256) - (((data // 256) // 256) // 256)) * 256
 
     return rgb
 

rio_rgbify/mbtiler.py

@@ -24,6 +24,7 @@
 from rasterio.enums import Resampling
 
 from rio_rgbify.encoders import data_to_rgb
+from rio_rgbify.coords import mercator_to_wgs, wgs_to_mercator, gcj_to_wgs
 
 buffer = bytes if sys.version_info > (3,) else buffer
 
@@ -130,6 +131,19 @@ def _tile_worker(tile):
         for c in i
     ]
 
+    if global_args["chinaoffset"]:
+        wgs_min = mercator_to_wgs(bounds[0], bounds[1])
+        wgs_max = mercator_to_wgs(bounds[2], bounds[3])
+        gcj_min = gcj_to_wgs(wgs_min[0], wgs_min[1])
+        gcj_max = gcj_to_wgs(wgs_max[0], wgs_max[1])
+        mercator_min = wgs_to_mercator(gcj_min[0], gcj_min[1])
+        mercator_max = wgs_to_mercator(gcj_max[0], gcj_max[1])
+        bounds = [
+            c
+            for i in (mercator_min, mercator_max)
+            for c in i
+        ]
+
     toaffine = transform.from_bounds(*bounds + [512, 512])
 
     out = np.empty((512, 512), dtype=src.meta["dtype"])
@@ -170,7 +184,7 @@ def _tile_range(min_tile, max_tile):
     return itertools.product(range(min_x, max_x + 1), range(min_y, max_y + 1))
 
 
-def _make_tiles(bbox, src_crs, minz, maxz):
+def _make_tiles(bbox, src_crs, minz, maxz, chinaoffset):
     """
     Given a bounding box, zoom range, and source crs,
     find all tiles that would intersect
@@ -260,6 +274,7 @@ def __init__(
         base_val=0,
         round_digits=0,
         bounding_tile=None,
+        chinaoffset=False,
         **kwargs
     ):
         self.run_function = _tile_worker
@@ -268,6 +283,7 @@ def __init__(
         self.min_z = min_z
         self.max_z = max_z
         self.bounding_tile = bounding_tile
+        self.chinaoffset = chinaoffset
 
         if not "format" in kwargs:
             writer_func = _encode_as_png
@@ -297,6 +313,7 @@ def __init__(
             "interval": interval,
             "round_digits": round_digits,
             "writer_func": writer_func,
+            "chinaoffset": chinaoffset,
         }
 
     def __enter__(self):
@@ -366,11 +383,12 @@ def run(self, processes=4):
 
         # generator of tiles to make
         if self.bounding_tile is None:
-            tiles = _make_tiles(bbox, src_crs, self.min_z, self.max_z)
+            tiles = _make_tiles(bbox, src_crs, self.min_z, self.max_z, self.chinaoffset)
         else:
             constrained_bbox = list(mercantile.bounds(self.bounding_tile))
-            tiles = _make_tiles(constrained_bbox, "EPSG:4326", self.min_z, self.max_z)
+            tiles = _make_tiles(constrained_bbox, "EPSG:4326", self.min_z, self.max_z, self.chinaoffset)
 
+        tiles_count = 0
         for tile, contents in self.pool.imap_unordered(self.run_function, tiles):
             x, y, z = tile
 
@@ -385,8 +403,11 @@ def run(self, processes=4):
                 (z, x, tiley, buffer(contents)),
             )
 
-            conn.commit()
+            tiles_count += 1
+            if (tiles_count % 100 == 0):
+                conn.commit()
 
+        conn.commit()
         conn.close()
 
         self.pool.close()

rio_rgbify/scripts/cli.py

@@ -67,6 +67,7 @@ def _rgb_worker(data, window, ij, g_args):
     default="png",
     help="Output tile format (.mbtiles output only)",
 )
+@click.option("--chinaoffset", is_flag=True, default=False, help="Use GCJ02 coordinates within China as bounding box")
 @click.option("--workers", "-j", type=int, default=4, help="Workers to run [DEFAULT=4]")
 @click.option("--verbose", "-v", is_flag=True, default=False)
 @click.pass_context
@@ -83,6 +84,7 @@ def rgbify(
     min_z,
     bounding_tile,
     format,
+    chinaoffset,
     workers,
     verbose,
     creation_options,
@@ -132,6 +134,7 @@ def rgbify(
             bounding_tile=bounding_tile,
             max_z=max_z,
             min_z=min_z,
+            chinaoffset=chinaoffset,
         ) as tiler:
             tiler.run(workers)
 

test/test_coords.py

@@ -0,0 +1,21 @@
+from rio_rgbify.coords import mercator_to_wgs, wgs_to_mercator, gcj_to_wgs
+import pytest
+
+
+def test_mercator_to_wgs_and_wgs_to_mercator():
+    x = 116.39
+    y = 39.9
+
+    p = wgs_to_mercator(x, y)
+    p = mercator_to_wgs(p[0], p[1])
+
+    assert abs(p[0] - x) < 1e-5
+    assert abs(p[1] - y) < 1e-5
+
+
+def test_gcj_to_wgs():
+    p = gcj_to_wgs(116.39623950248456, 39.901400287519685)
+
+    assert abs(p[0] - 116.39) < 1e-5
+    assert abs(p[1] - 39.9) < 1e-5
+