ref(dataset.ya_go_maps): add dname

datasets/ya_go_maps/.gitignore

@@ -1 +1,2 @@
-images
+images
+dataset_*

datasets/ya_go_maps/extract.py

@@ -1,15 +1,16 @@
+import argparse
 import os
 import zipfile
 from pathlib import Path
 
 
-def extract_zip_to_images():
+def extract_zip_to_images(dname):
     """
-    Разархивирует ya_go_maps.zip в папку images.
+    Разархивирует {dname}.zip в папку images.
     """
     # Определяем пути
     current_dir = Path(__file__).parent
-    zip_path = current_dir / "ya_go_maps.zip"
+    zip_path = current_dir / f"{dname}.zip"
     extract_dir = current_dir
 
     # Проверяем существование архива
@@ -49,12 +50,16 @@ def main():
     """
     Основная функция для запуска разархивирования.
     """
-    print("Начинаю разархивирование ya_go_maps.zip...")
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dname', type=str, default='ya_go_maps', help='Dataset name')
+    args = parser.parse_args()
+    
+    print(f"Начинаю разархивирование {args.dname}.zip...")
 
-    if extract_zip_to_images():
+    if extract_zip_to_images(args.dname):
         print("Разархивирование успешно завершено!")
     else:
-        print("Разархивирование не удалось.")
+        print("Разархивирование не удалось.")
 
 
 if __name__ == "__main__":

datasets/ya_go_maps/generate_dataset.py

@@ -1,15 +1,19 @@
+import argparse
 import numpy as np
 from pathlib import Path
-from ...google_map import GoogleMap
-from ...simulator import Simulator
-from ...yandex_map import YandexMap
+from google_map import GoogleMap
+from simulator import Simulator
+from yandex_map import YandexMap
+from constants import CHUNK_WIDTH
 
-LAT_MIN, LAT_MAX = 44.960236, 54.967830
-LON_MIN, LON_MAX = 53.084167, 58.677977
+LAT_MIN, LAT_MAX = 49.134520, 49.235065
+LON_MIN, LON_MAX = 55.767660, 55.825204
+
+
+
+def create_new_asset(yandex_map, google_map, dname, tile_y=None, tile_x=None):
+    folder = Path(f'dataset_{dname}')
 
-def create_new_asset(yandex_map, google_map):
-    folder = Path('dataset_ya_go_maps')
-    
     id = 0
     print(id)
     while (folder / f"{id:0{4}}_google.png").exists():
@@ -20,24 +24,34 @@ def create_new_asset(yandex_map, google_map):
 
     lat = np.random.rand() * (LAT_MAX - LAT_MIN) + LAT_MIN
     lon = np.random.rand() * (LON_MAX - LON_MIN) + LON_MIN
+    if tile_y is not None and tile_x is not None:
+        lat = tile_y * (LAT_MAX - LAT_MIN) + LAT_MIN
+        lon = tile_x * (LON_MAX - LON_MIN) + LON_MIN
 
     yandex_map.open(lat, lon, 18)
     google_map.open(lat, lon, 18)
 
     simulator = Simulator()
-    simulator._apply_perspective_transform(yandex_map.make_screenshot()).save(yandex_file)
-    simulator._apply_perspective_transform(google_map.make_screenshot()).save(google_file)
+    im_ya = simulator._apply_perspective_transform(yandex_map.make_screenshot())
+    im_go = simulator._apply_perspective_transform(google_map.make_screenshot())
+    im_ya.resize((CHUNK_WIDTH // 2, CHUNK_WIDTH // 2)).save(yandex_file)
+    im_go.resize((CHUNK_WIDTH // 2, CHUNK_WIDTH // 2)).save(google_file)
 
 def main():
-    folder = Path('dataset_ya_go_maps')
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dname', type=str, default='ya_go_maps', help='Dataset name')
+    args = parser.parse_args()
+
+    folder = Path(f'dataset_{args.dname}')
     if not folder.exists():
         folder.mkdir()
 
-    yandex_map = YandexMap(initial_zoom=15)
-    google_map = GoogleMap(initial_zoom=15)
+    yandex_map = YandexMap(initial_zoom=18)
+    google_map = GoogleMap(initial_zoom=18)
 
-    for i in range(4):
-        create_new_asset(yandex_map, google_map)
+    for tile_y in np.linspace(0, 1, 20):
+        for tile_x in np.linspace(0, 1, 20):
+            create_new_asset(yandex_map, google_map, args.dname, tile_y, tile_x)
 
 if __name__ == '__main__':
-    main()
+    main()