-
Notifications
You must be signed in to change notification settings - Fork 7
/
tasks.py
209 lines (186 loc) · 6.77 KB
/
tasks.py
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
import logging
import os
from io import BytesIO
from uuid import UUID
from zipfile import ZipFile
from celery.result import AsyncResult
from celery.signals import setup_logging, worker_process_init, worker_process_shutdown
from geojson import FeatureCollection
from numpy.typing import NDArray
from segment_anything import SamPredictor, sam_model_registry
from ultralytics import YOLO
from ultralytics_4bands import YOLO as YOLO_4
from sketch_map_tool import celery_app as celery
from sketch_map_tool import get_config_value, map_generation
from sketch_map_tool.database import client_celery as db_client_celery
from sketch_map_tool.definitions import get_attribution
from sketch_map_tool.helpers import to_array
from sketch_map_tool.models import Bbox, Layer, PaperFormat, Size
from sketch_map_tool.upload_processing import (
clip,
georeference,
merge,
polygonize,
post_process,
)
from sketch_map_tool.upload_processing.detect_markings import detect_markings
from sketch_map_tool.upload_processing.ml_models import init_model
from sketch_map_tool.wms import client as wms_client
@worker_process_init.connect
def init_worker(**kwargs):
"""Initializing database connection for worker"""
db_client_celery.open_connection()
@worker_process_shutdown.connect
def shutdown_worker(**kwargs):
"""Closing database connection for worker"""
db_client_celery.close_connection()
@setup_logging.connect
def on_setup_logging(**kwargs):
level = getattr(logging, get_config_value("log-level").upper())
format = "%(asctime)s - %(levelname)s - %(filename)s - %(funcName)s - %(message)s"
logging.basicConfig(
level=level,
format=format,
)
# 1. GENERATE SKETCH MAP & QUALITY REPORT
#
@celery.task()
def generate_sketch_map(
uuid: UUID,
bbox: Bbox,
format_: PaperFormat,
orientation: str, # TODO: is not accessed
size: Size,
scale: float,
layer: Layer,
) -> BytesIO | AsyncResult:
"""Generate and returns a sketch map as PDF and stores the map frame in DB."""
map_image = wms_client.get_map_image(bbox, size, layer)
qr_code_ = map_generation.qr_code(
str(uuid),
bbox,
layer,
format_,
)
map_pdf, map_img = map_generation.generate_pdf(
map_image,
qr_code_,
format_,
scale,
layer,
)
db_client_celery.insert_map_frame(map_img, uuid)
return map_pdf
@celery.task()
def generate_quality_report(bbox: Bbox) -> BytesIO | AsyncResult:
"""Generate a quality report as PDF.
Fetch quality indicators from the OQT API
"""
# TODO: Issue #469
# report = get_report(bbox)
# return generate_report_pdf(report)
return BytesIO(
b"We can not create a Quality Report at the moment. "
b"Sorry for the inconvenience."
)
# 2. DIGITIZE RESULTS
#
@celery.task()
def georeference_sketch_maps(
file_ids: list[int],
file_names: list[str],
uuids: list[str],
map_frames: dict[str, NDArray],
bboxes: list[Bbox],
layers: list[Layer],
) -> AsyncResult | BytesIO:
def process(
sketch_map_id: int,
uuid: str,
bbox: Bbox,
) -> BytesIO:
"""Process a Sketch Map and its attribution."""
# r = interim result
r = db_client_celery.select_file(sketch_map_id)
r = to_array(r)
r = clip(r, map_frames[uuid])
r = georeference(r, bbox)
return r
def get_attribution_file(layers: list[Layer]) -> BytesIO:
attributions = []
for index, layer in enumerate(layers):
attribution = get_attribution(layer)
attribution = attribution.replace("<br />", "\n")
attributions.append(attribution)
return BytesIO("\n".join(attributions).encode())
def zip_(file: BytesIO, file_name: str):
with ZipFile(buffer, "a") as zip_file:
name = ".".join(file_name.split(".")[:-1])
zip_file.writestr(f"{name}.geotiff", file.read())
buffer = BytesIO()
for file_id, uuid, bbox, file_name in zip(file_ids, uuids, bboxes, file_names):
zip_(process(file_id, uuid, bbox), file_name)
with ZipFile(buffer, "a") as zip_file:
zip_file.writestr("attributions.txt", get_attribution_file(layers).read())
buffer.seek(0)
return buffer
@celery.task()
def digitize_sketches(
file_ids: list[int],
file_names: list[str],
uuids: list[str],
map_frames: dict[str, NDArray],
layers: list[str],
bboxes: list[Bbox],
) -> AsyncResult | FeatureCollection:
# Initialize ml-models. This has to happen inside of celery context.
#
# Prevent usage of CUDA while transforming Tensor objects to numpy arrays
# during marking detection
os.environ["CUDA_VISIBLE_DEVICES"] = ""
# Zero shot segment anything model for automatic mask generation
sam_path = init_model(get_config_value("neptune_model_id_sam"))
sam_model = sam_model_registry[get_config_value("model_type_sam")](sam_path)
sam_predictor: SamPredictor = SamPredictor(sam_model) # mask predictor
# Custom trained model for object detection (obj) and classification (cls)
# of markings and colors.
if "osm" in layers:
path = init_model(get_config_value("neptune_model_id_yolo_osm_obj"))
yolo_obj_osm: YOLO_4 = YOLO_4(path) # yolo object detection
path = init_model(get_config_value("neptune_model_id_yolo_osm_cls"))
yolo_cls_osm: YOLO = YOLO(path) # yolo classification
if "esri-world-imagery" in layers:
path = init_model(get_config_value("neptune_model_id_yolo_esri_obj"))
yolo_obj_esri: YOLO_4 = YOLO_4(path)
path = init_model(get_config_value("neptune_model_id_yolo_esri_cls"))
yolo_cls_esri: YOLO = YOLO(path)
l = [] # noqa: E741
for file_id, file_name, uuid, bbox, layer in zip(
file_ids, file_names, uuids, bboxes, layers
):
# r = interim result
r: BytesIO = db_client_celery.select_file(file_id) # type: ignore
r: NDArray = to_array(r) # type: ignore
r: NDArray = clip(r, map_frames[uuid]) # type: ignore
if layer == "osm":
yolo_obj = yolo_obj_osm
yolo_cls = yolo_cls_osm
elif layer == "esri-world-imagery":
yolo_obj = yolo_obj_esri
yolo_cls = yolo_cls_esri
else:
raise ValueError("Unexpected layer: " + layer)
r: NDArray = detect_markings(
r,
map_frames[uuid],
yolo_obj,
yolo_cls,
sam_predictor,
) # type: ignore
# m = marking
for m in r:
m: BytesIO = georeference(m, bbox, bgr=False) # type: ignore
m: FeatureCollection = polygonize(m, layer_name=file_name) # type: ignore
m: FeatureCollection = post_process(m, file_name)
l.append(m)
return merge(l)