#!/usr/bin/python3 import sys import gzip import time import datetime import math from PIL import Image from lxml import etree # Knihovna lxml je Python binding na C knihovnu lxml, takže většina zpracování XMLka # je implementovaná rychle v C. # Lze ji získat pomocí `pip3 install lxml` def die(exit_code, message): print(message, file=sys.stderr) sys.exit(exit_code) def hasTag(el, name): """Kontroluje jestli má element daný tag (bez kontroly hodnoty). Arguments: el {etree.Element} -- Naparsovaný XML element name {str} -- Název tagu Returns: bool """ for tag in el.findall("tag"): if tag.get('k') == name: return True return False # Vyrobíme si množinu na pamatování si bodů budov. # Pythoní set() je trochu více "nenažraný", než prosté ukládání do pole, ale # pokud máme alespoň 3GB RAMky, tak by nám to mělo na celou Evropu stačit. buildingsToDisplay = set() placedBuildings = 0 # počítadlo kolik budov jsme již umístili do heatmapy # Proměnné na spočítání obdélníku pro kreslení heatmapy lonMin, lonMax = math.inf, -math.inf latMin, latMax = math.inf, -math.inf start = time.time() phaseStart = start # Funkce na vypisování statistiky def printStat(parsed): megabytes = parsed / 1024 / 1024 elapsed = time.time() - phaseStart elapsedString = str(datetime.timedelta(seconds=elapsed)) speed = megabytes/elapsed if elapsed > 0 else 0 print( f"Elapsed: {elapsedString:20s} Read: {megabytes:10.2f}MB ({speed:5.2f}MB/s)" + f"\tFound buildings: {len(buildingsToDisplay):10d} Placed buildings: {placedBuildings:10d}", file=sys.stderr, end="\r", flush=True ) ################################################################################ if len(sys.argv) != 4: die(1, f"Usage: {sys.argv[0]} ") heatmapWidth = int(sys.argv[2]) outputImageName = sys.argv[3] ################################################################################ # PRVNÍ PRŮCHOD - získání prvního bodu pro každou budovu ####################### print("[Phase 1] Getting buildings", file=sys.stderr) with gzip.open(sys.argv[1], 'rb') as gzippedFile: # Budeme streamově parsovat pomocí iterparse(). Ten se může "chytat" na více # typů eventů - nám bude stačit defaultní "end", který nastává vždy po načtení # celého tagu (tedy třeba po načtení ). Je potřeba vyjmenovat elementy, # které nás zajímají, abychom nebyli voláni například pro vnořené elementy # a nemuseli to speciálně ošetřovat. i = 0 for event, elem in etree.iterparse(gzippedFile, tag=('node', 'way')): if i % 10000 == 0: printStat(gzippedFile.tell()) i += 1 # 2. Najdeme všechny budovy a uložíme si z nich vždy první bod + zjistíme # minimální a maximální lon/lat if elem.tag == 'node': lat = float(elem.get("lat")) lon = float(elem.get("lon")) lonMin = min(lonMin, lon) lonMax = max(lonMax, lon) latMin = min(latMin, lat) latMax = max(latMax, lat) elif elem.tag == 'way': if hasTag(elem, "building"): nodes = elem.findall("nd") if len(nodes) > 0: buildingsToDisplay.add(int(nodes[0].get("ref"))) # Odstraníme zpracovaný element z paměti (musíme udělat explicitně, jinak by tam zůstal). # Musíme smazat i všechny rodiče současného elementu, kteří byli vytvořeni při parsování # (jinak by nám elementy zůstaly a postupně užíraly paměť, souvisí to s interním fungováním # knihovny lxml) elem.clear() while elem.getprevious() is not None: del elem.getparent()[0] del elem printStat(gzippedFile.tell()) # Spočítáme velikost heatmapy lonSpan = lonMax - lonMin latSpan = latMax - latMin cellSize = lonSpan / float(heatmapWidth) heatmapHeight = int(math.ceil(latSpan / cellSize)) print(f"\nFound {len(buildingsToDisplay)} buildings", file=sys.stderr) ################################################################################ # DRUHÝ PRŮCHOD ################################################################ print("[Phase 2] Getting locations of all buildings and placing on the heatmap", file=sys.stderr) # 1. Vytvoříme si 2D pole do kterého budeme počítat počty budov heatmapMax = 0 heatmapData = [] for i in range(heatmapWidth): heatmapData.append([0] * heatmapHeight) # 2. Projdeme přes všechny body a když potkáme nějaký, který chceme, přidáme ho do heatmapy phaseStart = time.time() with gzip.open(sys.argv[1], 'rb') as gzippedFile: i = 0 # Budeme streamově parsovat pomocí iterparse(), viz první průchod. for event, elem in etree.iterparse(gzippedFile, tag=('node')): if i % 10000 == 0: printStat(gzippedFile.tell()) i += 1 ID = int(elem.get("id")) if ID in buildingsToDisplay: lat = float(elem.get("lat")) lon = float(elem.get("lon")) x = math.floor((lon - lonMin) / cellSize) y = math.floor((lat - latMin) / cellSize) heatmapData[x][y] += 1 if heatmapData[x][y] > heatmapMax: heatmapMax = heatmapData[x][y] placedBuildings += 1 # Odstraníme zpracovaný element z paměti (viz první průchod) elem.clear() while elem.getprevious() is not None: del elem.getparent()[0] del elem printStat(gzippedFile.tell()) print("\n[Phase 3] Creating heatmap", file=sys.stderr) # 3. Vygenerujeme heatmapu a uložíme img = Image.new("RGB", (heatmapWidth, heatmapHeight)) pixels = img.load() # Barvy jako trojice RGB backgroundColor = (0, 0, 0) halfHeatColor = (255, 40, 0) heatColor = (80, 255, 0) for x in range(heatmapWidth): for y in range(heatmapHeight): # Spočítáme si číslo mezi 0 a 1 udávající hustotu, pro lepší znázornění umocníme hodnotu na 1/3 value = math.pow(heatmapData[x][y]/heatmapMax, 1.0/3) # A) Černobílá varianta # pixels[x, heatmapHeight-y-1] = (int(255*value), int(255*value), int(255*value)) # B) Plnobarevná varianta fromColor, toColor = backgroundColor, halfHeatColor if value > 0.5: fromColor, toColor = halfHeatColor, heatColor pixels[x, heatmapHeight-y-1] = ( fromColor[0] + round((toColor[0]-fromColor[0])*value), fromColor[1] + round((toColor[1]-fromColor[1])*value), fromColor[2] + round((toColor[2]-fromColor[2])*value) ) img.save(outputImageName) elapsedString = str(datetime.timedelta(seconds=time.time() - start)) print(f"\nTotal time: {elapsedString}\nPlaced {placedBuildings} buildings, max heatmap value: {heatmapMax}", file=sys.stderr)