|
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)
|
|
|
