opensr-utils/flowchart.txt at master · ESAOpenSR/opensr-utils · GitHub

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
Large File SR Processing Workflow
=================================

User: large_file_processing(root, model, ...)
--------------------------------------------

   ┌─────────────────────────────┐
   │ Verify input type (file/SAFE│
   │ /S2GM)                      │
   └───────────┬─────────────────┘
               │
               ▼
   ┌─────────────────────────────┐
   │ Extract metadata:            │
   │  - width, height, dtype      │
   │  - CRS, transform, bands     │
   │  - build sliding windows     │
   └───────────┬─────────────────┘
               │
               ▼
   ┌─────────────────────────────┐
   │ Create placeholder GeoTIFF  │
   │ (HR dimensions, ZSTD, tiled)│
   └───────────┬─────────────────┘
               │
               ▼
   ┌─────────────────────────────┐
   │ Build PredictionDataModule  │
   │ (wraps Dataset of windows)  │
   └───────────┬─────────────────┘
               │
               ▼
   ┌─────────────────────────────┐
   │ Preprocess model (wrap in   │
   │ Lightning + add hooks)      │
   └───────────┬─────────────────┘
               │
               ▼
   ┌─────────────────────────────┐
   │ Trainer.predict()           │
   │ - Each rank runs predict_step│
   │ - SR patches saved → temp/  │
   │ - Per-rank index JSONs      │
   └───────────┬─────────────────┘
               │
               ▼
   ┌─────────────────────────────┐
   │ On rank 0:                  │
   │ - Merge shard indices        │
   │ - Write global index.json    │
   │ - Clean up temp shards       │
   └───────────┬─────────────────┘
               │
               ▼
   ┌─────────────────────────────┐
   │ Stitch patches into         │
   │ placeholder GeoTIFF:        │
   │ - Load patch from .npz      │
   │ - Blend with feathered edges│
   │ - Write into placeholder    │
   └───────────┬─────────────────┘
               │
               ▼
   ┌─────────────────────────────┐
   │ Rename placeholder → sr.tif │
   │ (final stitched SR image)   │
   └─────────────────────────────┘


Multi-GPU SR Inference Workflow
================================

   ┌─────────────┐
   │ Trainer.run │
   └──────┬──────┘
          │
          ▼
   ┌─────────────────────────┐
   │ Each GPU rank executes: │
   └──────────┬──────────────┘
              │
              ▼
   [on_predict_start]
        │  Initialize:
        │   - rank/world_size
        │   - temp folder
        │   - local entries[]
        ▼
   [predict_step]
        │  For each batch:
        │   - Run model (SR patch)
        │   - Save patch → temp/sr_rXXXX_cXXXX.npz
        │   - Append offsets + path to entries[]
        ▼
   [on_predict_end]
        │  Write shard index_rank{r}.json
        │
        │───► DDP barrier ───(all ranks wait)───┐
        │                                      │
        ▼                                      │
   Rank >0:                                    │ Rank 0:
   - Finished                                  │  - Collect all shard JSONs
   - Optionally delete shards                  │  - Merge entries[]
                                               │  - Deduplicate + sort
                                               │  - Write global index.json
                                               │  - Clean up shard files
                                               ▼
                                        ┌───────────────┐
                                        │ index.json    │
                                        │ (global list) │
                                        └───────────────┘