|
|
import os |
|
|
import json |
|
|
import joblib |
|
|
import pandas as pd |
|
|
import numpy as np |
|
|
import reverse_geocoder |
|
|
from os.path import join, dirname |
|
|
|
|
|
|
|
|
class QuadTree(object): |
|
|
def __init__( |
|
|
self, data, mins=None, maxs=None, id="", depth=3, min_split=0, do_split=1000 |
|
|
): |
|
|
self.id = id |
|
|
self.data = data |
|
|
|
|
|
if mins is None: |
|
|
mins = data[["latitude", "longitude"]].to_numpy().min(0) |
|
|
if maxs is None: |
|
|
maxs = data[["latitude", "longitude"]].to_numpy().max(0) |
|
|
|
|
|
self.mins = np.asarray(mins) |
|
|
self.maxs = np.asarray(maxs) |
|
|
self.sizes = self.maxs - self.mins |
|
|
|
|
|
self.children = [] |
|
|
|
|
|
mids = 0.5 * (self.mins + self.maxs) |
|
|
xmin, ymin = self.mins |
|
|
xmax, ymax = self.maxs |
|
|
xmid, ymid = mids |
|
|
|
|
|
if depth > 0 and len(self.data) >= do_split: |
|
|
|
|
|
data_q1 = data[(data["latitude"] < mids[0]) & (data["longitude"] < mids[1])] |
|
|
data_q2 = data[ |
|
|
(data["latitude"] < mids[0]) & (data["longitude"] >= mids[1]) |
|
|
] |
|
|
data_q3 = data[ |
|
|
(data["latitude"] >= mids[0]) & (data["longitude"] < mids[1]) |
|
|
] |
|
|
data_q4 = data[ |
|
|
(data["latitude"] >= mids[0]) & (data["longitude"] >= mids[1]) |
|
|
] |
|
|
|
|
|
|
|
|
if data_q1.shape[0] > min_split: |
|
|
self.children.append( |
|
|
QuadTree(data_q1, [xmin, ymin], [xmid, ymid], id + "0", depth - 1) |
|
|
) |
|
|
if data_q2.shape[0] > min_split: |
|
|
self.children.append( |
|
|
QuadTree(data_q2, [xmin, ymid], [xmid, ymax], id + "1", depth - 1) |
|
|
) |
|
|
if data_q3.shape[0] > min_split: |
|
|
self.children.append( |
|
|
QuadTree(data_q3, [xmid, ymin], [xmax, ymid], id + "2", depth - 1) |
|
|
) |
|
|
if data_q4.shape[0] > min_split: |
|
|
self.children.append( |
|
|
QuadTree(data_q4, [xmid, ymid], [xmax, ymax], id + "3", depth - 1) |
|
|
) |
|
|
|
|
|
def unwrap(self): |
|
|
if len(self.children) == 0: |
|
|
return {self.id: [self.mins, self.maxs, self.data.copy()]} |
|
|
else: |
|
|
d = dict() |
|
|
for child in self.children: |
|
|
d.update(child.unwrap()) |
|
|
return d |
|
|
|
|
|
|
|
|
def extract(qt): |
|
|
cluster = qt.unwrap() |
|
|
boundaries, data = {}, [] |
|
|
for id, vs in cluster.items(): |
|
|
(min_lat, min_lon), (max_lat, max_lon), points = vs |
|
|
points["category"] = id |
|
|
data.append(points) |
|
|
boundaries[id] = ( |
|
|
float(min_lat), |
|
|
float(min_lon), |
|
|
float(max_lat), |
|
|
float(max_lon), |
|
|
) |
|
|
|
|
|
data = pd.concat(data) |
|
|
return boundaries, data |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
|
|
data_path = join(dirname(dirname(__file__)), "datasets", "osv5m") |
|
|
train_fp = join(data_path, f"train.csv") |
|
|
test_fp = join(data_path, f"test.csv") |
|
|
|
|
|
df_train = pd.read_csv(train_fp) |
|
|
df_train["split"] = "train" |
|
|
|
|
|
df_test = pd.read_csv(test_fp) |
|
|
df_test["split"] = "test" |
|
|
|
|
|
df = pd.concat([df_train, df_test]) |
|
|
size_before = df.shape[0] |
|
|
qt = QuadTree(df, depth=15) |
|
|
boundaries, df = extract(qt) |
|
|
assert df.shape[0] == size_before |
|
|
|
|
|
location = reverse_geocoder.search( |
|
|
[(lat, lon) for lat, lon in zip(df["latitude"], df["longitude"])] |
|
|
) |
|
|
df["city"] = [l.get("name", "") for l in location] |
|
|
df["country"] = [l.get("cc", "") for l in location] |
|
|
del location |
|
|
|
|
|
df_train = df[df["split"] == "train"].drop(["split"], axis=1) |
|
|
df_test = df[df["split"] == "test"].drop(["split"], axis=1) |
|
|
assert (df_train.shape[0] + df_test.shape[0]) == size_before |
|
|
|
|
|
json.dump(boundaries, open(join(data_path, "borders.json"), "w")) |
|
|
df_train.to_csv(train_fp, index=False) |
|
|
df_test.to_csv(test_fp, index=False) |
|
|
|
