Spaces:
Running
Running
File size: 1,694 Bytes
c4c7cee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
from os.path import join, dirname
import numpy as np
import pandas as pd
if __name__ == "__main__":
# Define the list of cities
cities = [
"Walvis Bay",
"Keetmanshoop",
"Warmbad",
"Rundu",
"Outapi",
"Karibib",
"Otjimbingwe",
"Ondangwa",
"Oranjemund",
"Maltahohe",
"Otavi",
"Outjo",
"Swakopmund",
"Gobabis",
"Karasburg",
"Opuwo",
"Hentiesbaai",
"Katima Mulilo",
"Oshikango",
"Bethanie",
"Ongandjera",
"Mariental",
"Bagani",
"Nkurenkuru",
"Usakos",
"Rehoboth",
"Aranos",
"Omaruru",
"Arandis",
"Windhoek",
"Khorixas",
"Okahandja",
"Grootfontein",
"Tsumeb",
]
csv_dtype = {"category": str, "country": str, "city": str}
for split in ["train", "test"]:
fp = join(
dirname(dirname(__file__)), "datasets", "osv5m", f"{split}.csv"
)
# Read the CSV file into a pandas DataFrame
df = pd.read_csv(fp, dtype=csv_dtype)
# Check if the "country" column contains any of the cities in the list
mask = df["city"].isin(cities)
# If a city is found, set the corresponding rows in the "country" column to 'NMB'
df.loc[mask, "country"] = "NMB"
assert all(map(lambda x: isinstance(x, str), df["country"].unique().tolist()))
# Drop the columns that are all NaN
df.dropna(subset=["id", "latitude", "longitude"], inplace=True)
# Save the modified DataFrame back to the CSV file
df.to_csv(fp, index=False)
|