population.py - develop - Seler1ty/DataSet_Pipeline

population.py
164 строки5 KB
Seler1ty
-m create main.py json_to_csv.py, area.py, population.py
11 дек 2025, 11:33
11 дек 2025, 11:33dcb73d9
import json
import csv
import ast
import time
import requests
import math

# Настройки
OVERPASS_URLS = [
    'https://overpass-api.de/api/interpreter',
    'https://lz4.overpass-api.de/api/interpreter',
    'https://z.overpass-api.de/api/interpreter',
]
REQUEST_TIMEOUT = 60
DELAY = 2
BUILDINGS_PER_FLOOR = 7.5
AVG_FAMILY_SIZE = 2.5

def get_bbox(coordinates):
    """Получает bounding box для координат"""
    lats = [coord[0] for coord in coordinates]
    lons = [coord[1] for coord in coordinates]

    return {
        'south': min(lats),
        'north': max(lats),
        'west': min(lons),
        'east': max(lons)
    }

def make_overpass_request(bbox):
    """Запрос зданий из Overpass API"""
    query = f"""
[out:json][timeout:{REQUEST_TIMEOUT}];
(
  way["building"]({bbox['south']:.6f},{bbox['west']:.6f},{bbox['north']:.6f},{bbox['east']:.6f});
  relation["building"]({bbox['south']:.6f},{bbox['west']:.6f},{bbox['north']:.6f},{bbox['east']:.6f});
);
out body;
>;
out skel qt;
"""

    for url in OVERPASS_URLS:
        try:
            response = requests.post(
                url,
                data={'data': query},
                timeout=REQUEST_TIMEOUT
            )

            if response.status_code == 200:
                data = response.json()
                buildings = [e for e in data.get('elements', []) if e.get('type') == 'way']
                return buildings

        except Exception:
            continue

    return []

def estimate_population(buildings):
    """Оценка населения по зданиям"""
    if not buildings:
        return 0

    population = 0

    for building in buildings:
        tags = building.get('tags', {})

        # Пытаемся получить количество этажей
        levels = None
        for key in ['building:levels', 'levels', 'level']:
            if key in tags:
                try:
                    levels = int(float(str(tags[key])))
                    break
                except:
                    continue

        # Оцениваем население
        if levels:
            population += levels * BUILDINGS_PER_FLOOR * AVG_FAMILY_SIZE
        else:
            # Базовая оценка для здания без данных
            population += 3 * BUILDINGS_PER_FLOOR * AVG_FAMILY_SIZE

    return int(population)

def add_population_to_csv(input_csv, output_csv):
    """Добавляет колонку population в CSV файл"""
    print(f"Добавление population в {input_csv}...")

    try:
        with open(input_csv, 'r', encoding='utf-8') as infile:
            reader = csv.DictReader(infile)
            rows = list(reader)

        headers = ['id', 'lat', 'lon', 'coordinates', 'area', 'population']

        with open(output_csv, 'w', newline='', encoding='utf-8') as outfile:
            writer = csv.DictWriter(outfile, fieldnames=headers)
            writer.writeheader()

            processed = 0
            removed = 0

            for i, row in enumerate(rows):
                try:
                    # Пропускаем если нет area
                    if 'area' not in row or float(row['area']) == 0:
                        removed += 1
                        continue

                    # Преобразуем координаты
                    coords_str = row['coordinates']
                    coordinates = ast.literal_eval(coords_str)

                    # Получаем bounding box
                    bbox = get_bbox(coordinates)

                    # Запрашиваем здания
                    buildings = make_overpass_request(bbox)

                    # Оцениваем население
                    population = estimate_population(buildings)

                    # Удаляем если население 0
                    if population == 0:
                        removed += 1
                        continue

                    # Формируем новую строку
                    new_row = {
                        'id': row['id'],
                        'lat': row['lat'],
                        'lon': row['lon'],
                        'coordinates': coords_str,
                        'area': row['area'],
                        'population': population
                    }

                    writer.writerow(new_row)
                    processed += 1

                    # Задержка между запросами
                    if i < len(rows) - 1:
                        time.sleep(DELAY)

                    # Прогресс каждые 10 объектов
                    if processed % 10 == 0:
                        print(f"  Обработано {processed} объектов")

                except Exception:
                    removed += 1
                    continue

        print(f"✓ {input_csv}: обработано {processed} строк, удалено {removed}")
        return True

    except Exception as e:
        print(f"✗ Ошибка при обработке {input_csv}: {e}")
        return False