Move logic inside main(), add argparsing, improve some variable names

Author: mrhheffernan

personal_heatmap.py

@@ -1,3 +1,4 @@
+import argparse
 import glob
 import os
 
@@ -8,119 +9,143 @@
 from geopy.geocoders import Nominatim
 
 geolocator = Nominatim(user_agent="heatmap_app")
-location = geolocator.geocode("Montreal Quebec")
-lat_check = float(location.raw["lat"])
-lon_check = float(location.raw["lon"])
 
-data = glob.glob("*.gpx")
-fitdata = glob.glob("*.fit")
 
-if len(fitdata) != 0:
-    print("Converting Garmin FIT files")
-    os.system("python fit_to_csv.py")
-    os.system("mkdir fit_files")
-    os.system("mv *.fit ./fit_files")
+def parse_args() -> argparse.Namespace:
+    args = argparse.ArgumentParser()
+    args.add_argument(
+        "--dir",
+        help="Path to direcotry with .fit, .gpx files to process for the heatmap",
+        default=os.getcwd(),
+    )
+    args.add_argument(
+        "--output_path",
+        help="Path to write the heatmap .html to",
+        default="heatmap.html",
+    )
 
-csvdata = glob.glob("*.csv")
+    return args.parse_args()
 
-lat = []
-lon = []
 
-all_lat = []
-all_long = []
+def main():
+    args = parse_args()
+    location = geolocator.geocode("Montreal Quebec")
+    lat_check = float(location.raw["lat"])
+    lon_check = float(location.raw["lon"])
 
-print("Loading data")
+    gpx_files = glob.glob(args.dir + "/*.gpx")
+    fit_files = glob.glob(args.dir + "/*.fit")
 
-for activity in data:
-    gpx_filename = activity
-    with open(gpx_filename, "r") as gpx_file:
-        gpx = gpxpy.parse(gpx_file)
+    if len(fit_files) != 0:
+        print("Converting Garmin FIT files")
+        os.system("python fit_to_csv.py")
+        os.system("mkdir fit_files")
+        os.system("mv *.fit ./fit_files")
 
-        for track in gpx.tracks:
-            for segment in track.segments:
-                for point in segment.points:
-                    lat.append(point.latitude)
-                    lon.append(point.longitude)
+    csvdata = glob.glob("*.csv")
 
-    check1 = np.any(np.isclose(lat, lat_check, atol=0.5))
-    check2 = np.any(np.isclose(lon, lon_check, atol=0.5))
+    lat = []
+    lon = []
 
-    if check1 and check2:
-        all_lat.append(lat)
-        all_long.append(lon)
+    all_lat = []
+    all_long = []
 
-    lon = []
-    lat = []
+    print("Loading data")
 
-for activity in csvdata:
-    csv_filename = activity
-    csv_file = pd.read_csv(csv_filename)
+    for activity in gpx_files:
+        gpx_filename = activity
+        with open(gpx_filename, "r") as gpx_file:
+            gpx = gpxpy.parse(gpx_file)
 
-    for i in range(len(csv_file)):
-        lat.append(csv_file["position_lat"][i])
-        lon.append(csv_file["position_long"][i])
+            for track in gpx.tracks:
+                for segment in track.segments:
+                    for point in segment.points:
+                        lat.append(point.latitude)
+                        lon.append(point.longitude)
 
-    check1 = np.any(np.isclose(lat, lat_check, atol=0.5))
-    check2 = np.any(np.isclose(lon, lon_check, atol=0.5))
+        check1 = np.any(np.isclose(lat, lat_check, atol=0.5))
+        check2 = np.any(np.isclose(lon, lon_check, atol=0.5))
 
-    if check1 and check2:
-        all_lat.append(lat)
-        all_long.append(lon)
+        if check1 and check2:
+            all_lat.append(lat)
+            all_long.append(lon)
 
-    lon = []
-    lat = []
+        lon = []
+        lat = []
 
-if not all_lat or not all_long:
-    raise ValueError(
-        "No activities found within the specified location. Check lat_check/lon_check."
-    )
+    for activity in csvdata:
+        csv_filename = activity
+        csv_file = pd.read_csv(csv_filename)
 
-all_lat = all_lat[0]
-all_long = all_long[0]
+        for i in range(len(csv_file)):
+            lat.append(csv_file["position_lat"][i])
+            lon.append(csv_file["position_long"][i])
 
-central_long = sum(all_long) / len(all_long)
-central_lat = sum(all_lat) / len(all_lat)
+        check1 = np.any(np.isclose(lat, lat_check, atol=0.5))
+        check2 = np.any(np.isclose(lon, lon_check, atol=0.5))
 
-print("Initializing map")
-m = folium.Map(
-    location=[central_lat, central_long], tiles="Stamen Toner", zoom_start=14.2
-)
+        if check1 and check2:
+            all_lat.append(lat)
+            all_long.append(lon)
 
-print("Plotting gpx data")
+        lon = []
+        lat = []
 
-for activity in data:
-    gpx_filename = activity
-    with open(gpx_filename, "r") as gpx_file:
-        gpx = gpxpy.parse(gpx_file)
+    if not all_lat or not all_long:
+        raise ValueError(
+            "No activities found within the specified location. Check lat_check/lon_check."
+        )
 
-        for track in gpx.tracks:
-            for segment in track.segments:
-                for point in segment.points:
-                    lat.append(point.latitude)
-                    lon.append(point.longitude)
+    all_lat = all_lat[0]
+    all_long = all_long[0]
 
-    points = list(zip(lat, lon))
+    central_long = sum(all_long) / len(all_long)
+    central_lat = sum(all_lat) / len(all_lat)
 
-    folium.PolyLine(points, color="red", weight=2.5, opacity=0.5).add_to(m)
-    lat = []
-    lon = []
+    print("Initializing map")
+    m = folium.Map(
+        location=[central_lat, central_long], tiles="Stamen Toner", zoom_start=14.2
+    )
 
-print("Plotting csv data")
-color = "red"
-hr = []
-for activity in csvdata:
-    csv_filename = activity
-    csv_file = pd.read_csv(csv_filename)
-    for i in range(len(csv_file)):
-        lat.append(csv_file["position_lat"][i])
-        lon.append(csv_file["position_long"][i])
-        if "heart_rate" in csv_file.columns:
-            hr.append(csv_file["heart_rate"][i])
-    points = list(zip(lat, lon))
-
-    folium.PolyLine(points, color=color, weight=2.5, opacity=0.5).add_to(m)
-    lat = []
-    lon = []
-    hr = []
+    print("Plotting gpx data")
 
-m.save("heatmap.html")
+    for activity in gpx_files:
+        gpx_filename = activity
+        with open(gpx_filename, "r") as gpx_file:
+            gpx = gpxpy.parse(gpx_file)
+
+            for track in gpx.tracks:
+                for segment in track.segments:
+                    for point in segment.points:
+                        lat.append(point.latitude)
+                        lon.append(point.longitude)
+
+        points = list(zip(lat, lon))
+
+        folium.PolyLine(points, color="red", weight=2.5, opacity=0.5).add_to(m)
+        lat = []
+        lon = []
+
+    print("Plotting csv data")
+    color = "red"
+    hr = []
+    for activity in csvdata:
+        csv_filename = activity
+        csv_file = pd.read_csv(csv_filename)
+        for i in range(len(csv_file)):
+            lat.append(csv_file["position_lat"][i])
+            lon.append(csv_file["position_long"][i])
+            if "heart_rate" in csv_file.columns:
+                hr.append(csv_file["heart_rate"][i])
+        points = list(zip(lat, lon))
+
+        folium.PolyLine(points, color=color, weight=2.5, opacity=0.5).add_to(m)
+        lat = []
+        lon = []
+        hr = []
+
+    m.save("heatmap.html")
+
+
+if __name__ == "__main__":
+    main()

simple_heatmap.py

@@ -25,7 +25,7 @@ def parse_args() -> argparse.Namespace:
 
 def main():
     args = parse_args()
-    
+
     gpx_list = glob.glob(args.dir + "/*.gpx")
 
     fig = plt.figure(facecolor="0.05")