diff --git a/.github/workflows/update_pr_description.yml b/.github/workflows/update_pr_description.yml index 732e858..fd462ba 100644 --- a/.github/workflows/update_pr_description.yml +++ b/.github/workflows/update_pr_description.yml @@ -17,16 +17,13 @@ jobs: uses: actions/checkout@v4 # We need to fetch the history up to the base branch with: - fetch-depth: 0 - ref: ${{ github.event.pull_request.head.sha }} + fetch-depth: 0 - name: ⚙️ Get Commit List and Format id: commit_list run: | - BASE_BRANCH="origin/${{ github.event.pull_request.base.ref }}" + COMMIT_LIST=$(git log --pretty=format:"* %s (%an, %ad)" ${{ github.event.pull_request.base.sha }}..HEAD) - COMMIT_LIST=$(git log --pretty=format:"* %s (%an, %ad)" $BASE_BRANCH..HEAD) - FORMATTED_LIST="## 📦 Commits Included in this PR:\n\n${COMMIT_LIST}" echo "FORMATTED_LIST<> $GITHUB_ENV diff --git a/README.md b/README.md index e9daf47..2f08462 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,10 @@ # SEN2SR Tools + +![Python Version from PEP 621 TOML](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2FMiguelPadillaR%2Fsen2sr-tools%2Fmain%2Fpyproject.toml) + [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) + ![Status: Active](https://img.shields.io/badge/Status-Active-00aa00.svg) +![Code style: Ruff](https://img.shields.io/badge/code%20style-autopep8-aa0000.svg) + This repository provides a list of comprehensive tools and utilities for the usage of the [SEN2SR](https://github.com/ESAOpenSR/SEN2SR.git) neural network super-resolution model, developed by the [ESAOpenSR](https://opensr.eu/) team. This model specializes in upscaling Sentinel-2's 10m/px images up to a x4 times improvement of 2.5m/px. This package implements a feature to crop a polygon directly from the SR image, useful for close-up satellite observation of agricultural states. @@ -15,23 +21,34 @@ This package implements a feature to crop a polygon directly from the SR image, source venv/bin/activate # On Windows use `venv\Scripts\activate` pip install -e . ``` - +## Quickstart example +The overall workflow can be run directly, since it implements a simple example in it: +```shell +cd sen2sr-tools +python -m sen2sr_tools.get_sr_image +``` +The example can also be run using commands: +```shell +python -m sen2sr_tools get-sr-image --latitude 42.465226 --longitude -2.292699 --start-date 2025-11-01 --end-date 2025-11-15 --geojson-path ./example.geojson +``` +Output will be at `sen2sr_out`, where the SR and original images can be seen. ## Usage in Python +All relevant methods are stored in the `get_sr_image` module. ### Get a SR image -It runs the full workflow, from date and location data to super resolved image. Default `size` is 128 px, and it must be a multiple of 32. Default `bands` are Near-Infrared, Red, Blue, Green and the Scene Classification Layer, that contains cloud density information among others. Cropping a polygon from the image is optional, provided the GeoJSON filepath. +Using `get_sr_image` runs the full workflow, from date and location data to super resolved image. Default `size` is the minimal requested 128px. Default `bands` are Near-Infrared, Red, Blue, Green and the Scene Classification Layer, that contains cloud density information among others. Cropping a polygon from the image is optional, provided the GeoJSON filepath. The final image outputs are in TIF and PNG formats, for usage and visualization respectively. ```python from sen2sr_tools.get_sr_image import get_sr_image lat = 37.265840 lng = -4.593406 -start_date = "2025-11-1" +start_date = "2025-11-01" end_date = "2025-11-15" # bands = ["B08", "B02", "B03", "B04", "SCL"] # Default bands # size = 128 # Default size -# geojson_path = None # Default value +# geometry = None # Default value sr_filepath = get_sr_image(lat, lng, bands, start_date, end_date) @@ -39,14 +56,14 @@ print(f"SR image successfully downloaded and saved at: {sr_filepath}") ``` ### Download CUBO -It serves the requested bands' data directly. It uses the date range to get cloudless (<=1%) images. The `crs` can be provided (i.e, "EPSG:32630"), but it can also be automatically calculated from the `lat` and `lon` arguments. If no requested data was found within the date range, it retries with by expanding it backwards in time. +Running `download_sentinel_cubo` serves the requested bands' data directly. It uses the date range to get cloudless (<=1%) images. The `crs` can be provided (i.e, "EPSG:32630"), but it can also be automatically calculated from the `lat` and `lon` arguments. If no requested data is found within the date range, it retries with by expanding it backwards in time. The CUBO data is provided as a `Dask.array`. ```python from sen2sr_tools.get_sr_image import download_sentinel_cubo lat = 37.265840 lng = -4.593406 -start_date = "2025-11-1" +start_date = "2025-11-01" end_date = "2025-11-15" # crs = None # Default CRS # bands = ["B08", "B02", "B03", "B04", "SCL"] # Default bands @@ -55,22 +72,38 @@ end_date = "2025-11-15" # max_retries = 3 # Default retries # retry_days_shift = 15 # Default shift -cloudless_image_data_array, sample_date = download_sentinel_cubo(lat, lon, start_date, end_date): +cloudless_cubo_data_array, sample_date = download_sentinel_cubo(lat, lon, start_date, end_date) + +print(f"CUBO data summary:\n\n{cloudless_cubo_data_array.coords}") + +# If you need the array as NumPy +# cubo_np_array = (cloudless_cubo_data_array.astype("float32") / 10_000).compute() ``` + ### Crop parcel from image -It crops the polygon from the given TIF file and returns the cropped image as PNG. +The `crop_parcel_from_tif` method crops the polygon from the given TIF file and returns the cropped image output and metadata. ```python -from sen2sr_tools.get_sr_image import crop_png_from_tif +import json +import geopandas as gpd +from sen2sr_tools.get_sr_image import crop_parcel_from_tif -raster_path = "path/to/file.tif" -geojson_path = "path/to/polygon.geojson" -date = "2025-11-1" +gdf = gpd.read_file(geometry) -out_png_path = crop_png_from_tif(raster_path, geojson_path, date) - -print(f"Polygon successfully cropped and saved at: {out_png_path}") +raster_path = "path/to/file.tif" +geometry = { + "type": "Polygon", + "coordinates": [[ + [-2.29544172799985, 42.46242962302221],[-2.292531924994703, 42.46538599420777], + [-2.292180580738519, 42.463584097818824], [-2.286733836169323, 42.464620125420105], + [-2.29544172799985, 42.46242962302221] + ]], + "CRS": "epsg:4258"} + +out_image, out_meta = crop_parcel_from_tif(raster_path, geojson_path) + +print(f"Polygon successfully cropped!") ``` ## Attribution and License diff --git a/example.geojson b/example.geojson new file mode 100644 index 0000000..59783f7 --- /dev/null +++ b/example.geojson @@ -0,0 +1 @@ +{"type": "Polygon", "coordinates": [[[-2.29544172799985, 42.46242962302221], [-2.295610039140137, 42.46269315594096], [-2.295614996197969, 42.46270091255462], [-2.295615258551553, 42.46270132242971], [-2.29564566306893, 42.462748897276754], [-2.295878516543866, 42.46305769415171], [-2.295991106464929, 42.46324323010404], [-2.29599140737527, 42.46324372547455], [-2.295992673880916, 42.463245811730374], [-2.295993345271399, 42.46324691814166], [-2.295993859082094, 42.46324776555212], [-2.296201055549893, 42.46358920323635], [-2.29641295175108, 42.46398073117693], [-2.296480756322637, 42.464091026970955], [-2.296481638937094, 42.46409246279106], [-2.296482275961773, 42.46409349963254], [-2.29648281827094, 42.464094381408806], [-2.296483270055548, 42.46409511566357], [-2.296678995888173, 42.46441349892175], [-2.296876671351592, 42.464743812981624], [-2.296877248864755, 42.46474477773874], [-2.296877699811172, 42.46474553127187], [-2.296899610107369, 42.464782143427115], [-2.296912614630915, 42.464802834994515], [-2.296913005227626, 42.46480345609358], [-2.29691383587428, 42.46480477791979], [-2.296914277600603, 42.464805480323314], [-2.296942342728665, 42.4648501365915], [-2.29695618879534, 42.46487216758711], [-2.297120475783675, 42.46513360502995], [-2.297162661904852, 42.465207646576026], [-2.297163078485464, 42.46520837747803], [-2.29721083354344, 42.465292194838355], [-2.297283671448152, 42.46551314684827], [-2.297310372840486, 42.46559415124173], [-2.297333542936804, 42.4656644435623], [-2.297421466591744, 42.46604649995172], [-2.297470165452073, 42.46630795248198], [-2.297470670880864, 42.46631066654239], [-2.297471100872523, 42.46631297575686], [-2.297471265996025, 42.466313861724075], [-2.297471371608011, 42.46631442498799], [-2.297471565229986, 42.46631546518224], [-2.297486759106771, 42.46639703870696], [-2.297500486150413, 42.466454595563086], [-2.297500972300826, 42.4664566306893], [-2.297501479405998, 42.46645875885464], [-2.297502269819514, 42.46646207222116], [-2.297502550613287, 42.46646324987863], [-2.297503002397895, 42.466465143350675], [-2.297613238680407, 42.46692737432509], [-2.2976548925507, 42.46733524781601], [-2.297703813531477, 42.46747465899063], [-2.297704043195638, 42.46747531194092], [-2.297704197422665, 42.46747575199087], [-2.29770560222972, 42.46747975518806], [-2.297706718699289, 42.46748294198784], [-2.297706980214684, 42.467483687977264], [-2.297707180542182, 42.46748425878491], [-2.297708436151352, 42.4674878437251], [-2.297713744410947, 42.46750300323608], [-2.297717755990044, 42.46751445710745], [-2.297794450408634, 42.46773322982241], [-2.297907858403495, 42.46790472776247], [-2.298084048533234, 42.468127859067344], [-2.29809625509955, 42.468136607260206], [-2.298096801599669, 42.46813699953329], [-2.298098146057018, 42.46813796261403], [-2.298204164565603, 42.46821394289442], [-2.298372276216583, 42.468263061688134], [-2.298373153801898, 42.46826331817439], [-2.298428218718065, 42.46827940723848], [-2.298640912038291, 42.46828855776271], [-2.298652274546908, 42.4682890464277], [-2.298295191201281, 42.468433771729245], [-2.29828402482921, 42.46843825772409], [-2.29817888474413, 42.46848085539674], [-2.29812783224525, 42.46850153942043], [-2.297961757396862, 42.46856882682853], [-2.297727035595598, 42.46866392958396], [-2.297717624394156, 42.46866774251194], [-2.296907657573084, 42.46899592086276], [-2.296604819393445, 42.46893646634364], [-2.296466300891566, 42.46890927209537], [-2.296458301202703, 42.468907702164806], [-2.29645318237413, 42.46890670388009], [-2.296432571273003, 42.46890265039147], [-2.296416943213609, 42.46889958177653], [-2.296415307904202, 42.46889926074962], [-2.296221152840325, 42.46886114320449], [-2.296216371802472, 42.46886020443128], [-2.296215757408932, 42.46886008373187], [-2.29621483959048, 42.468859903520936], [-2.296190242055967, 42.46885507470623], [-2.296186486124932, 42.46885433709871], [-2.296174779958265, 42.468852038780724], [-2.295696427230324, 42.46875812625561], [-2.295695648551473, 42.46875797370496], [-2.295695119653352, 42.468757869769355], [-2.295693719037248, 42.46875759484292], [-2.295692705665094, 42.46875739535361], [-2.295690878410094, 42.46875703660813], [-2.295680204055769, 42.46875494113221], [-2.295611219312174, 42.46874139765226], [-2.295609590708291, 42.46874107830173], [-2.295605253073141, 42.46874022670032], [-2.295590105296827, 42.4687372528009], [-2.295565594934112, 42.468732440749996], [-2.295539986542015, 42.46872741328418], [-2.295047517850102, 42.468630784184086], [-2.293047614074356, 42.46823810373176], [-2.292864166892419, 42.46820208250073], [-2.29294994813278, 42.46814604193098], [-2.292977442453202, 42.46810076791649], [-2.292978197662723, 42.468035146830836], [-2.2929726656063, 42.46802542382258], [-2.292966697690886, 42.468014943090225], [-2.292955482703774, 42.46799524813117], [-2.292929917897577, 42.46788116706944], [-2.292840907947922, 42.467251252759326], [-2.292723397849405, 42.46655604090503], [-2.292707640708693, 42.4664627478026], [-2.292682567082051, 42.4663090471586], [-2.29268169704045, 42.466303712915106], [-2.292680944345501, 42.4662991011917], [-2.292680863041035, 42.466298602468434], [-2.292680193326932, 42.46629449617383], [-2.292651783702629, 42.46612035458135], [-2.292636107028192, 42.46602426024671], [-2.292628146734276, 42.46597546583268], [-2.292626881905012, 42.46596770921902], [-2.292624460373042, 42.465952871573144], [-2.292614024064781, 42.46588890004613], [-2.292611755921648, 42.465874994467946], [-2.292532196568382, 42.46538765801564], [-2.292531924994703, 42.46538599420777], [-2.29253150841409, 42.46538344107991], [-2.292530899887884, 42.46537971364735], [-2.29253077751209, 42.46537896681973], [-2.292499108165449, 42.465184979393925], [-2.292514171284537, 42.46518367852248], [-2.293331128613271, 42.46511313223023], [-2.293332779010104, 42.46511298973787], [-2.293334302002001, 42.46511285730379], [-2.293337112454302, 42.46511261506677], [-2.293338397400135, 42.46511250442565], [-2.293343883356088, 42.46511203084808], [-2.293348163994292, 42.46511166120614], [-2.294365665073881, 42.46502379622452], [-2.294388117679245, 42.46502231346577], [-2.29439173447068, 42.46502207458151], [-2.294837438007441, 42.46499263649763], [-2.294839009614379, 42.46499253256202], [-2.29483926274787, 42.46497224751695], [-2.294839597185826, 42.46494550840604], [-2.294813193351072, 42.46494777152003], [-2.292522403152132, 42.46514520307006], [-2.292491732928409, 42.46514784672247], [-2.29247031884091, 42.46502077706282], [-2.292469913994963, 42.46501837229466], [-2.292469126096018, 42.46501369770698], [-2.292469033895078, 42.46501315120686], [-2.292468888888144, 42.465012291223545], [-2.292468796687203, 42.46501174137066], [-2.29246865084208, 42.46501087719639], [-2.292468130325862, 42.46500778595032], [-2.29237186751496, 42.46443656256789], [-2.29237045181143, 42.46442815970946], [-2.292361399355465, 42.46437443003055], [-2.292360235947236, 42.46436752669469], [-2.292332343486389, 42.46420206540164], [-2.292286468489407, 42.4639299267824], [-2.292274818481494, 42.46385490706815], [-2.292267107130117, 42.4638052015412], [-2.292255602129138, 42.46373118597902], [-2.292255502384484, 42.463730544763386], [-2.292255315468032, 42.46372934112202], [-2.292232659182409, 42.46358357897899], [-2.292146327250977, 42.46306756971062], [-2.292060333948454, 42.46257554693598], [-2.292045163540999, 42.462488707899375], [-2.292044736902102, 42.46248614806599], [-2.292044660626778, 42.462485690414056], [-2.292042724407031, 42.46247461289016], [-2.292041643141457, 42.46246842369249], [-2.291995457175848, 42.46248091272896], [-2.292000768788205, 42.46251172628324], [-2.292000855121813, 42.462512225006506], [-2.292000972468464, 42.46251290477889], [-2.29200107808045, 42.462513520848816], [-2.292001742765411, 42.46251737484812], [-2.292035018084795, 42.46271033046158], [-2.292125204015536, 42.46323295472833], [-2.292165267000512, 42.46348710075316], [-2.29218035442712, 42.463582664513304], [-2.292180452495393, 42.46358328393598], [-2.292180580738519, 42.463584097818824], [-2.292180710658026, 42.46358492092177], [-2.292182113788701, 42.46359380490147], [-2.292182350996575, 42.46359532873156], [-2.292184460721729, 42.463608892328075], [-2.292189657502005, 42.463638397467186], [-2.292304945557859, 42.464292308329334], [-2.292305218807919, 42.46429385730513], [-2.292305350403807, 42.4642945991036], [-2.292306532252224, 42.46430130630292], [-2.292314207561416, 42.4643448410724], [-2.292319823436876, 42.464376693144544], [-2.29003566332374, 42.464578025632534], [-2.290029879810207, 42.46457852854675], [-2.290018316974093, 42.464579558682715], [-2.289984699673039, 42.464582551860516], [-2.28996248679012, 42.46458447969836], [-2.289964079351817, 42.464596550477836], [-2.289964498447001, 42.46459973560123], [-2.289964852163336, 42.46460245636716], [-2.289966174827735, 42.464612643732885], [-2.28996961057005, 42.46463862847246], [-2.290049583151167, 42.464631613657275], [-2.290087498692423, 42.46462828771789], [-2.290106176088373, 42.464626649893916], [-2.290262504459173, 42.4646129379377], [-2.290263372824393, 42.46461286166238], [-2.290264154017815, 42.464612792930765], [-2.290264652741084, 42.46461274934487], [-2.290266635899492, 42.464612575001276], [-2.290271364969544, 42.46461216009704], [-2.290493659760425, 42.46459266295091], [-2.290666259082631, 42.46457752439469], [-2.291470781857269, 42.46450204451394], [-2.291890216507924, 42.46446269315259], [-2.291907628236421, 42.46446196560335], [-2.291908770689891, 42.4644619178265], [-2.291909696052056, 42.46446187926974], [-2.291911333037843, 42.464461811376324], [-2.29191235227733, 42.46446176862862], [-2.291915171949725, 42.46446165044377], [-2.291931804161179, 42.46446095558395], [-2.292068908635909, 42.46445522487642], [-2.292069683123808, 42.46445519302518], [-2.292071669634978, 42.46445511004434], [-2.292074912593509, 42.4644549742575], [-2.292087022767933, 42.464454467990514], [-2.292088953120348, 42.46445438752424], [-2.292090569151376, 42.464454319630825], [-2.29209530576514, 42.4644541209797], [-2.292101282062458, 42.46445387119898], [-2.292110812286932, 42.46445347305855], [-2.292263383050316, 42.46444709610624], [-2.292275444609698, 42.46444659235383], [-2.292288538819614, 42.464446045015514], [-2.292313368532858, 42.46444500733584], [-2.292432135078736, 42.46515891083532], [-2.291153520097793, 42.465265641805715], [-2.291150103633857, 42.465265926790444], [-2.291149074336086, 42.46526601312405], [-2.290783348409584, 42.4652965416936], [-2.290642061364922, 42.46530936684441], [-2.290464607244111, 42.46532547518688], [-2.290462147155383, 42.46532569814552], [-2.290458019906015, 42.465326072816616], [-2.290456065246079, 42.46532625051297], [-2.290453240544542, 42.465326506999226], [-2.290449475393413, 42.465326848980894], [-2.290219032553973, 42.46534776685969], [-2.29018667253847, 42.46535070471693], [-2.29010580393185, 42.46535804558817], [-2.290078821745743, 42.46536049478042], [-2.290081902933532, 42.46537073076119], [-2.290086114001937, 42.46538471932022], [-2.290467991018623, 42.46666014079587], [-2.290479558045689, 42.46669752492443], [-2.290561843194028, 42.466963446685725], [-2.290562150809893, 42.46696444161769], [-2.290562302522349, 42.46696493195906], [-2.290562799569237, 42.4669665379318], [-2.290563087906723, 42.4669674691613], [-2.290563271470414, 42.46696806260008], [-2.290563529633047, 42.46696889827587], [-2.290690529722901, 42.46737932573285], [-2.290786418700899, 42.46764922973658], [-2.290791224884464, 42.46766265586988], [-2.29079281493159, 42.467667131806444], [-2.290794099039233, 42.467670748597875], [-2.290841317655374, 42.467803717441335], [-2.290853215767635, 42.467837221586684], [-2.289249910145993, 42.46751588789841], [-2.289231675314556, 42.46751223255022], [-2.289126006331355, 42.4674910548324], [-2.288164135108241, 42.46729827607711], [-2.288103839884182, 42.46728908196697], [-2.288022996423274, 42.467276754701246], [-2.287718459231761, 42.46723031476395], [-2.287701709673656, 42.467227760797904], [-2.287700925127472, 42.467227640936684], [-2.287699724000676, 42.46722745737299], [-2.287649407432938, 42.46721978457837], [-2.287633170847337, 42.467217308564024], [-2.287624591968929, 42.46721600098705], [-2.287623728632851, 42.46721586939117], [-2.28747575366725, 42.46719331536477], [-2.287416450860587, 42.46718426123242], [-2.287343994332492, 42.467173212207], [-2.287316925812776, 42.46705995005727], [-2.287302783026712, 42.46700081488868], [-2.286733836169323, 42.464620125420105], [-2.286714781587707, 42.46455660903046], [-2.286577164140473, 42.464097540548295], [-2.286572393999093, 42.46407856308019], [-2.28657037563669, 42.464070505556194], [-2.286561383530431, 42.46403469471095], [-2.286558887399518, 42.46402475209682], [-2.286536803597915, 42.46393680581074], [-2.286412866255663, 42.46344322986935], [-2.286406092001116, 42.46341065527711], [-2.286383993112087, 42.463304405427984], [-2.286379051979873, 42.46328064608384], [-2.286372405968452, 42.4632486892379], [-2.286295813809085, 42.462880359053926], [-2.286817927617904, 42.463099304436106], [-2.286823291198062, 42.46310091376161], [-2.286863281260477, 42.46311291078034], [-2.286919281597095, 42.46312971062986], [-2.286973410254622, 42.463145949730034], [-2.287123117759712, 42.46319086248448], [-2.28724777510749, 42.46322826002408], [-2.287559799853551, 42.463321869124286], [-2.288187726814313, 42.463424033633565], [-2.288557922810043, 42.463440115992135], [-2.289043063203577, 42.46341314134974], [-2.289197773866968, 42.463382714201224], [-2.289198446095642, 42.46338258176714], [-2.289200223059221, 42.46338223224176], [-2.28920115764148, 42.46338204867808], [-2.289202950530675, 42.463381696638116], [-2.289297083499854, 42.463363182689285], [-2.289341220090015, 42.46335450238985], [-2.289372759517148, 42.46334829978113], [-2.289632969011096, 42.46329712406827], [-2.289784279136168, 42.46325800069469], [-2.289784866707615, 42.46325784898223], [-2.289884634830632, 42.463232052835494], [-2.289939771831171, 42.46321779689374], [-2.290011517573834, 42.46319627719425], [-2.290177157401422, 42.46314659429843], [-2.290443779051678, 42.46304990065768], [-2.290896614750273, 42.4628688758459], [-2.290959771556245, 42.46284362619929], [-2.291145615124441, 42.462754714317896], [-2.291154603039748, 42.46275041440132], [-2.291155250960901, 42.46275010427088], [-2.291156050594512, 42.462749722056074], [-2.291157022895337, 42.46274925686042], [-2.291157630583354, 42.46274896600836], [-2.291158519065143, 42.46274854104585], [-2.291159400841409, 42.46274811943609], [-2.291162300141889, 42.46274673223104], [-2.291166277355181, 42.4627448295389], [-2.291167637738146, 42.462744178264984], [-2.291169515284569, 42.46274327972491], [-2.291265021209569, 42.462697587453434], [-2.291346964376239, 42.46265838445177], [-2.291353346357694, 42.462655330924264], [-2.291379531424763, 42.46264280416923], [-2.291383279812084, 42.46264101044184], [-2.291783990128129, 42.462421231060254], [-2.292496464513032, 42.46206079495315], [-2.292497218884362, 42.46206041357654], [-2.29254831245457, 42.462034565462], [-2.292695367925338, 42.461932203139796], [-2.292713117444552, 42.46191984821378], [-2.29271411656747, 42.461919152515776], [-2.292943929925944, 42.46175918388421], [-2.292963197407913, 42.46174382404573], [-2.293069386907332, 42.46165917268599], [-2.29310812722791, 42.46162829040011], [-2.293146171012292, 42.461597962996244], [-2.293170351128003, 42.46157868713237], [-2.293316556673739, 42.46146213592363], [-2.293746358901045, 42.461216857952884], [-2.294156783005254, 42.46099703498539], [-2.294346522482277, 42.46090843323445], [-2.29438068460707, 42.46094337906723], [-2.294381047543498, 42.46094375122376], [-2.294381433949258, 42.46094414684961], [-2.294381741565123, 42.460944461171], [-2.294393969086198, 42.46095696864766], [-2.294485877498142, 42.46105111670426], [-2.294486353590271, 42.461051603692866], [-2.294486941999909, 42.46105220635174], [-2.294547611895061, 42.461114353976505], [-2.29459354724175, 42.46116140830733], [-2.294648450387177, 42.46121765255735], [-2.294674518945784, 42.46124807970587], [-2.29471087880572, 42.46129052147511], [-2.294711683468472, 42.46129146108651], [-2.29471213525308, 42.461291989146446], [-2.294713511561663, 42.46129359511918], [-2.294790778464274, 42.46138378607906], [-2.294917605886912, 42.46155638791584], [-2.29491848933956, 42.4615575923954], [-2.294919071043674, 42.461558381970725], [-2.294919599103606, 42.46155910029987], [-2.294989349953196, 42.461654025358946], [-2.295038668236206, 42.46173849566957], [-2.295039011056066, 42.46173908240283], [-2.295039301069934, 42.46173958028791], [-2.295191859260461, 42.46200087775295], [-2.295226980275034, 42.46206106149769], [-2.295319015253723, 42.46221897991562], [-2.295384296872276, 42.46233103842399], [-2.295384722672983, 42.46233176848779], [-2.295385052919987, 42.46233233678087], [-2.29538537897604, 42.46233289669203], [-2.29544172799985, 42.46242962302221]]], "CRS": "epsg:4258"} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 4e39172..506f10f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,16 @@ authors = [ license = "MIT" readme = "README.md" requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Development Status :: 1 - Planning", + "License :: OSI Approved :: MIT License", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: GIS", + "Topic :: Scientific/Engineering :: SR", + "Topic :: Utilities" +] dependencies = [ "cubo", "mlstac", @@ -19,14 +29,17 @@ dependencies = [ "structlog", "xarray", "pillow", + "matplotlib", "opencv-python (>=4.11, <4.12)", "numpy (>=1.9, <2.0.0)" ] [project.urls] -"Homepage" = "https://github.com/KhaosResearch/sigpac-tools" -"Bug Tracker" = "https://github.com/KhaosResearch/sigpac-tools/issues" +"Homepage" = "https://github.com/MiguelPadillaR/sen2sr-tools" +"Bug Tracker" = "https://github.com/MiguelPadillaR/sen2sr-tools/issues" +[project.optional-dependencies] +dev = ["autopep8", "pytest"] [build-system] requires = ["poetry-core>=2.0.0,<3.0.0"] diff --git a/sen2sr_tools/__main__.py b/sen2sr_tools/__main__.py index 396f72a..599d0c1 100644 --- a/sen2sr_tools/__main__.py +++ b/sen2sr_tools/__main__.py @@ -6,6 +6,11 @@ logger = structlog.get_logger() +class VerboseArgumentParser(argparse.ArgumentParser): + def error(self, message): + raise RuntimeError(f"Argument error: {message}") + + # --- Helper to load dynamic logic and run command --- def run_command(command_name: str, args: argparse.Namespace): @@ -18,27 +23,29 @@ def run_command(command_name: str, args: argparse.Namespace): crop_png_from_tif, ) except ImportError: - logger.error(f"Could not import logic from sen2sr_tools.get_sr_image. Ensure the package is installed in editable mode.") + logger.error(f"Could not import logic from sen2sr_tools.get-sr-image. Ensure the package is installed in editable mode.") sys.exit(1) # Dispatch logic based on command - if command_name == "get_sr_image": + if command_name == "get-sr-image": logger.info("Starting Super Resolution task", latitude=args.latitude, longitude=args.longitude, - start_date=args.start_date, end_date=args.end_date) - + start_date=args.start_date, end_date=args.end_date, + geometry=args.geometry + ) result = get_sr_image( lat=args.latitude, lon=args.longitude, bands=args.bands, start_date=args.start_date, end_date=args.end_date, - size=args.size + size=args.size, + geometry=args.geometry ) logger.info(f"✅ Get SR Image task completed. Result (filepath):\n{result}") return result - elif command_name == "download_cubo": + elif command_name == "download-cubo": logger.info("Starting Sentinel-2 Cubo download", start_date=args.start_date, end_date=args.end_date) # Note: The download function returns two values, we'll only print the path/date for CLI feedback @@ -53,9 +60,9 @@ def run_command(command_name: str, args: argparse.Namespace): cloud_threshold=args.cloud_threshold ) logger.info(f"✅ Download successful for acquisition date: {date}") - return date + return data, date - elif command_name == "crop_parcel": + elif command_name == "crop-polygon": logger.info("Starting parcel cropping.", tif_path=args.raster_path) # Note: The crop function needs a `date` string as the second argument, @@ -65,26 +72,26 @@ def run_command(command_name: str, args: argparse.Namespace): out_path = crop_png_from_tif( raster_path=args.raster_path, + geometry=args.geometry, date=args.acquisition_date ) logger.info(f"✅ Cropping successful. Output path: {out_path}") return out_path - - + # --- Argument Parser Setup --- def get_parser(): - parser = argparse.ArgumentParser(description="SEN2SR Tools: Command Line Interface for Sentinel-2 Super Resolution.") + parser = VerboseArgumentParser(description="SEN2SR Tools: Command Line Interface for Sentinel-2 Super Resolution.") subparsers = parser.add_subparsers( dest="command", help="Available commands") subparsers.required = True # ========================================================================= - # 1. get_sr_image (Main Function) + # 1. get-sr-image (Main Function) # ========================================================================= get_sr_image_parser = subparsers.add_parser( - "get_sr_image", help="Get SR image from downloaded Sentinel's imagery data and load up SEN2SR model to Super-Resolve it." + "get-sr-image", help="Get SR image from downloaded Sentinel's imagery data and load up SEN2SR model to Super-Resolve it." ) get_sr_image_parser.add_argument( "--latitude", type=float, help="Latitude coordinate component (e.g., 37.265840).", required=True, metavar="FLOAT", @@ -92,25 +99,29 @@ def get_parser(): get_sr_image_parser.add_argument( "--longitude", type=float, help="Longitude coordinate component (e.g., -4.593406).", required=True, metavar="FLOAT", ) - get_sr_image_parser.add_argument( - "--bands", nargs='+', default=["B08", "B02", "B03", "B04", "SCL"], - help="Image bands to retrieve (space-separated, e.g., B08 B02 B03 B04 SCL).", required=False, metavar="STR [...]", - ) get_sr_image_parser.add_argument( "--start-date", type=str, help="Initial date in search range in ISO format (yyyy-mm-dd).", required=True, metavar="STR", ) get_sr_image_parser.add_argument( "--end-date", type=str, help="Final date in search range in ISO format (yyyy-mm-dd).", required=True, metavar="STR", ) + get_sr_image_parser.add_argument( + "--bands", nargs='+', default=["B08", "B02", "B03", "B04", "SCL"], + help="Image bands to retrieve (space-separated, e.g., B08 B02 B03 B04 SCL).", required=False, + ) get_sr_image_parser.add_argument( "--size", type=int, default=128, help="Size of the image tile in pixels (e.g., 128 for 128x128).", required=False, metavar="INT", ) + get_sr_image_parser.add_argument( + "--geojson-path", type=str, default="", help="GeoJSON to crop SR image filepath. If none is provided, it downloads the image in sizexsize", metavar="STR", + ) + # ========================================================================= # 2. download_cubo (Helper Function) # ========================================================================= download_cubo_parser = subparsers.add_parser( - "download_cubo", help="Download raw Sentinel-2 imagery data using Cubo and filter for cloudless data." + "download-cubo", help="Download raw Sentinel-2 imagery data using Cubo and filter for cloudless data." ) # Re-use common arguments download_cubo_parser.add_argument( @@ -119,15 +130,15 @@ def get_parser(): download_cubo_parser.add_argument( "--longitude", type=float, help="Longitude coordinate component.", required=True, metavar="FLOAT", ) - download_cubo_parser.add_argument( - "--bands", nargs='+', default=["B08", "B02", "B03", "B04", "SCL"], help="Image bands to retrieve.", metavar="STR [...]", - ) download_cubo_parser.add_argument( "--start-date", type=str, help="Initial date in search range (yyyy-mm-dd).", required=True, metavar="STR", ) download_cubo_parser.add_argument( "--end-date", type=str, help="Final date in search range (yyyy-mm-dd).", required=True, metavar="STR", ) + download_cubo_parser.add_argument( + "--bands", nargs='+', default=["B08", "B02", "B03", "B04", "SCL"], help="Image bands to retrieve.", + ) download_cubo_parser.add_argument( "--size", type=int, default=128, help="Size of the image tile in pixels.", metavar="INT", ) @@ -139,7 +150,7 @@ def get_parser(): "--cloud-threshold", type=float, default=0.01, help="Max percentage of cloud density tolerated (0.0 to 1.0).", metavar="FLOAT", ) download_cubo_parser.add_argument( - "--polygon-path", type=str, default="", help="GeoJSON to crop SR image filepath. If none is provided, it downloads the image in sizexsize", metavar="STR", + "--geojson-path", type=str, default="", help="GeoJSON to crop SR image filepath. If none is provided, it downloads the image in sizexsize", metavar="STR", ) @@ -147,7 +158,7 @@ def get_parser(): # 3. crop_parcel (Helper Function) # ========================================================================= crop_parcel_parser = subparsers.add_parser( - "crop_parcel", help="Crops a parcel from a Super-Resolved TIF using a local GeoJSON file." + "crop-polygon", help="Crops a parcel from a Super-Resolved TIF using a local GeoJSON file." ) crop_parcel_parser.add_argument( "--raster-path", type=Path, help="Path to the uncropped SR TIF image.", required=True, metavar="PATH", @@ -161,13 +172,16 @@ def get_parser(): "--acquisition-date", type=str, help="Acquisition date (yyyy-mm-dd) used to name the output file.", required=True, metavar="STR", ) + return parser def main(): args = get_parser().parse_args() # Run the command - run_command(args.command, args) + result = run_command(args.command, args) + logger.info(f"RESULT:") + logger.info(f"{result}") if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/sen2sr_tools/constants.py b/sen2sr_tools/constants.py index c4743c1..0b337b3 100644 --- a/sen2sr_tools/constants.py +++ b/sen2sr_tools/constants.py @@ -5,7 +5,9 @@ CURR_SCRIPT_DIR = Path.Path(os.path.dirname(os.path.abspath(__file__))) -MODEL_DIR = str(CURR_SCRIPT_DIR / "model") +MODEL_DIR = CURR_SCRIPT_DIR / "model" +RGBN_MODEL_DIR = str(MODEL_DIR / "SEN2SRLite_RGBN") +MAIN_MODEL_DIR = str(MODEL_DIR / "SEN2SRLite") SEN2SR_SR_DIR = CURR_SCRIPT_DIR / "sen2sr_out" PNG_DIR = SEN2SR_SR_DIR / "png" TIF_DIR = SEN2SR_SR_DIR / "tif" diff --git a/sen2sr_tools/get_sr_image.py b/sen2sr_tools/get_sr_image.py index 3008f5d..05b45fe 100644 --- a/sen2sr_tools/get_sr_image.py +++ b/sen2sr_tools/get_sr_image.py @@ -13,14 +13,15 @@ from datetime import datetime, timedelta from rasterio.mask import mask +from xarray import DataArray from .constants import * -from .utils import lonlat_to_utm_epsg, save_to_png, save_to_tif, get_cloudless_time_indices, make_pixel_faithful_comparison, reorder_bands +from .utils import lonlat_to_utm_epsg, reproject_geometry, save_to_png, save_to_tif, get_cloudless_time_indices, make_pixel_faithful_comparison, reorder_bands logger = structlog.get_logger() -def get_sr_image(lat: float, lon: float, start_date: str, end_date: str, bands: list=["B08", "B02", "B03", "B04", "SCL"], size: int=128, geojson_path: str=None): +def get_sr_image(lat: float, lon: float, start_date: str, end_date: str, bands: list=["B08", "B02", "B03", "B04", "SCL"], size: int=128, geometry: dict=None): """ Get SR image from downloaded Sentinel's imagery data and load up SEN2SR model from HuggingFace to Super-Resolve it Arguments: @@ -30,43 +31,57 @@ def get_sr_image(lat: float, lon: float, start_date: str, end_date: str, bands: end_date (str): Final date in search range. ISO format (`yyyy-mm-dd`). bands (list): List of bands Defaults are NIR + RGB + Clouds (`B02`, `B03`, `B04`, `B08` and `SCL`) size (int): Image size in px. Default is `128` and must be a multiple of 32. - geojson_path (str): GeoJSON to crop SR image filepath. If `None`, it downloads the full image in `size`x`size` px. + geometry (dict): GeoJSON to crop SR image from. If `None`, it downloads the full image in `size`x`size` px. Returns: sr_image_filepath (str): Local filepath to SR image. """ try: - # Ensure sizeis right (minimum for SEN2SR) + # Ensure size is right (minimum for SEN2SR) + if size < 128: + logger.warning(f"Warning: size cannot be less thatn 128px for SR process ({size}px indicated).") + logger.warning(f"Resetting size to 128px...") + size = 128 logger.debug(f"Image size {size}x{size}px") - # Download model - if not os.path.exists(MODEL_DIR) or len(os.listdir(MODEL_DIR)) == 0: - mlstac.download( - file="https://huggingface.co/tacofoundation/sen2sr/resolve/main/SEN2SRLite/NonReference_RGBN_x4/mlm.json", - output_dir=MODEL_DIR, - ) # Prepare data crs = lonlat_to_utm_epsg(lon, lat) - cloudless_image_data, sample_date = download_sentinel_cubo( + cubo_image_data, sample_date = download_sentinel_cubo( lat, lon, start_date, end_date, crs, bands, size) - original_s2_reordered, superX_reordered = apply_sen2sr(size, cloudless_image_data) + original_s2_reordered, superX_reordered = apply_sen2sr(size, cubo_image_data) # Save original and super-res images in TIF & PNG save_to_tif(original_s2_reordered, OG_TIF_FILEPATH, - cloudless_image_data, crs) + cubo_image_data, crs) save_to_tif(superX_reordered, SR_TIF_FILEPATH, - cloudless_image_data, crs) + cubo_image_data, crs) - save_to_png(original_s2_reordered, OG_PNG_FILEPATH, lat) - save_to_png(superX_reordered, SR_PNG_FILEPATH, lat) + save_to_png(original_s2_reordered, OG_PNG_FILEPATH) + save_to_png(superX_reordered, SR_PNG_FILEPATH) # Make comparison grid make_pixel_faithful_comparison(original_s2_reordered, superX_reordered) # Get and save cropped sr parcel image - if geojson_path: - sr_image_filepath = str( - crop_png_from_tif(SR_TIF_FILEPATH, geojson_path, sample_date)) + if geometry: + out_image, out_meta = crop_parcel_from_tif(SR_TIF_FILEPATH, geometry) + + # Save cropped SR cropped TIF & PNG + year, month, day = sample_date.split("-") + filename = f"SR_{year}-{month}-{day}" + + out_tif_path = TIF_DIR / f"{filename}.tif" + with rasterio.open(out_tif_path, "w", **out_meta) as dest: + dest.write(out_image) + + out_png_path = SEN2SR_SR_DIR / f"{filename}.png" + save_to_png(out_image, out_png_path, apply_gamma_correction=True) + + logger.info( + f"✅ Clipped raster saved to {out_tif_path} and PNG saved to {out_png_path}") + + sr_image_filepath = out_png_path + else: sr_image_filepath = SR_PNG_FILEPATH return sr_image_filepath @@ -80,7 +95,7 @@ def get_sr_image(lat: float, lon: float, start_date: str, end_date: str, bands: # -------------------- -def download_sentinel_cubo(lat: float, lon: float, start_date: str, end_date: str, crs: str=None, bands: list=["B08", "B02", "B03", "B04", "SCL"], size: int=128, cloud_threshold: float = 0.01, max_retries: int = 3, retry_days_shift: int = 15): +def download_sentinel_cubo(lat: float, lon: float, start_date: str, end_date: str, crs: str=None, bands: list=["B08", "B02", "B03", "B04", "SCL"], size: int=128, resolution: int = RESOLUTION, cloud_threshold: float = 0.01, max_retries: int = 3, retry_days_shift: int = 15): """ Download Sentinel's imagery data cubo and uses SCL band to filter the least cloudy data within date range. Arguments: @@ -89,8 +104,9 @@ def download_sentinel_cubo(lat: float, lon: float, start_date: str, end_date: st start_date (str): Intial date in search range end_date (str): Final date in search range crs (str): Coordinate Reference System for the image. If `None` it is calculated from `land` and `lon`. - bands (list): List of bands Defaults are NIR + RGB + Clouds (`B02`, `B03`, `B04`, `B08` and `SCL`) - size (int): Image size in px. Default is `128` and must be a multiple of 32. + bands (list): List of bands Defaults are NIR + RGB + Clouds (`B02`, `B03`, `B04`, `B08` and `SCL`). + resolution (int): Resolution in m/px. Default is `10`. + size (int): Image size in px. Default is `128` and must be at least `32`. cloud_threshold (float) : Max percentage of cloud density tolerated (0.0 to 1.0). Default is `0.01`. max_retries (int): Max number of retries if requested image is not found within date range. Shifts `retry_days_shift` back from `start_date` for new date range. Default is `3`. retry_days_shift (int): Number of days to shift back from `start_date`. Default is `15`. @@ -102,7 +118,7 @@ def download_sentinel_cubo(lat: float, lon: float, start_date: str, end_date: st logger.info( f"🌍 Attempt {attempt+1}/{max_retries}: {start_date} → {end_date}") - da = cubo.create( + cube_data = cubo.create( lat=lat, lon=lon, collection="sentinel-2-l2a", @@ -110,25 +126,31 @@ def download_sentinel_cubo(lat: float, lon: float, start_date: str, end_date: st start_date=start_date, end_date=end_date, edge_size=size, - resolution=RESOLUTION, + resolution=resolution, ) - # Find cloudless time index - scl = da.sel(band="SCL") - cloudless_date = get_cloudless_time_indices( - scl, cloud_threshold)[-1] - cloudless_image_data = da.isel(time=cloudless_date).sel( - band=bands[:-1]) # drop SCL band + # Find cloudless time index for RBG upscale + if "SCL" in bands and len(bands) == 5: + scl = cube_data.sel(band="SCL") + cloudless_date = get_cloudless_time_indices( + scl, cloud_threshold)[-1] + cloudless_image_data = cube_data.isel(time=cloudless_date).sel( + band=bands[:-1]) # drop SCL band + + # Get acquisition date and reproject + crs = lonlat_to_utm_epsg(lat,lon) if crs is None else crs # Calculate CRS if not provided + cloudless_image_data = cloudless_image_data.rio.write_crs(crs).rio.reproject(crs) + cube_data = cloudless_image_data + msg = "☁️ Cloudless image found on " + else: + cube_data = cube_data[-1] + msg = "Image data found on " - # Get acquisition date and reproject - acq_date = cloudless_image_data["time"].values + acq_date = cube_data["time"].values acq_date_str = np.datetime_as_string(acq_date, unit='D') - crs = lonlat_to_utm_epsg(lat,lon) if crs is None else crs # Calculate CRS if not provided - cloudless_image_data = cloudless_image_data.rio.write_crs( - crs).rio.reproject(crs) + logger.info(f"{msg}{acq_date_str}!") - logger.info(f"☁️ Cloudless image found on {acq_date_str}!") - return cloudless_image_data, str(acq_date_str) + return cube_data, acq_date_str except ValueError as e: logger.warning(f"⚠️ {e}") @@ -153,17 +175,48 @@ def download_sentinel_cubo(lat: float, lon: float, start_date: str, end_date: st # -------------------- -def apply_sen2sr(size, cloudless_image_data): +def apply_sen2sr(size: int, cubo_sample: DataArray): + """ + Super resolves the given `cubo_sample` data for both RGB upscaling (`B02`, `B03`, `B04`, and `B08`, optional `SCL`) and +10 bands options. + Anything in between will lead to error. + + Arguments: + size (int): Edge size of the images + cubo_sample (DataAray): `Cubo.cube` sample with the image data. + Returns: + (original_s2_reordered, superX_reordered) (tuple): Reordered original and SR version of the image data + """ + # Get CUBO sample bands + bands = list(cubo_sample.band) + is_rgbn_upscale = all(b in BANDS for b in bands) + # Select model type + if is_rgbn_upscale: # True color upscale + file="https://huggingface.co/tacofoundation/sen2sr/resolve/main/SEN2SRLite/NonReference_RGBN_x4/mlm.json" + model_dir = RGBN_MODEL_DIR + elif len(bands) >= 10: # Everything else + file="https://huggingface.co/tacofoundation/sen2sr/resolve/main/SEN2SRLite/main/mlm.json" + model_dir = MAIN_MODEL_DIR + else: + raise ValueError( + f'Error: Bands must be {BANDS} for RGBN upscaling or have at least have 10 bands for other upscaling.' + ) + + # Download model + if not os.path.exists(model_dir) or len(os.listdir(model_dir)) == 0: + mlstac.download( + file=file, + output_dir=model_dir, + ) + # Prepare the data to be used in the model device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - original_s2_numpy = (cloudless_image_data.compute( - ).to_numpy() / 10_000).astype("float32") + original_s2_numpy = (cubo_sample.compute().to_numpy() / 10_000).astype("float32") X = torch.from_numpy(original_s2_numpy).float().to(device) X = torch.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0) - # Load - model = mlstac.load((MODEL_DIR)).compiled_model(device=device) + # Load + model = mlstac.load((model_dir)).compiled_model(device=device) - # Apply model for normal or large size images + # Apply model for normal or large size images if size <= 128: superX = model(X[None]).squeeze(0) else: @@ -172,11 +225,12 @@ def apply_sen2sr(size, cloudless_image_data): X=X, # The input tensor overlap=32, # The overlap between the patches ) - # Reorder bands ( [NIR, B, G, R] -> [R, G, B, NIR]) - original_s2_reordered, superX_reordered = reorder_bands( - original_s2_numpy, superX) + if is_rgbn_upscale: # True color upscale + # Reorder bands ( [NIR, B, G, R] -> [R, G, B, NIR]) + original_s2_numpy = reorder_bands(original_s2_numpy, False) + superX = reorder_bands(superX, True) - return original_s2_reordered,superX_reordered + return original_s2_numpy, superX # -------------------- @@ -184,30 +238,30 @@ def apply_sen2sr(size, cloudless_image_data): # -------------------- -def crop_png_from_tif(raster_path: str, geojson_path: str, date: str): +def crop_parcel_from_tif(raster_path: str, geometry: dict): """ Crops the parcel from the SR image, using the stored parcel's geometry and`rasterio` Arguments: raster_path (str): Path to uncropped SR image. - geojson_filepath (str): Path to uncropped SR image. - date (str): Image date. For naming the file. + geometry (dict): Parcel geometry to crop with. Returns: - out_png_path (str): Path to cropped SR parcel image + Tuple: Cropped image output and metadata info """ with rasterio.open(raster_path) as src: - + # Reproject geometry + geometry_clean = { + "type": geometry["type"], + "coordinates": geometry["coordinates"] + } raster_crs = src.crs - gdf = gpd.read_file(geojson_path) - if raster_crs: - gdf = gdf.to_crs(raster_crs) - gdf["geometry"] = gdf["geometry"].buffer(1) - logger.info( - f"Reprojected polygon to match raster CRS: {raster_crs}") + from_crs = geometry.get("CRS", "epsg:4326").upper() + + geom_reproj = reproject_geometry(geometry_clean, from_crs, raster_crs) + logger.info(f"Reprojected polygon to match raster CRS: {raster_crs}") # Get parcel's geom and apply mask on SR image logger.info(f"Cropping parcel's geometry from raster...") - geom = [json.loads(gdf.to_json())["features"][0]["geometry"]] - out_image, out_transform = mask(src, geom, crop=True) + out_image, out_transform = mask(src, [geom_reproj ], crop=True) out_meta = src.meta.copy() logger.info(f"Cropping successful!") @@ -219,32 +273,24 @@ def crop_png_from_tif(raster_path: str, geojson_path: str, date: str): "transform": out_transform }) - # Save cropped TIF & PNG - year, month, day = date.split("-") - filename = f"SR_{year}-{month}-{day}" - - out_tif_path = TIF_DIR / f"{filename}.tif" - with rasterio.open(out_tif_path, "w", **out_meta) as dest: - dest.write(out_image) - - out_png_path = SEN2SR_SR_DIR / f"{filename}.png" - save_to_png(out_image, out_png_path, apply_gamma_correction=True) + return out_image, out_meta - logger.info( - f"✅ Clipped raster saved to {out_tif_path} and PNG saved to {out_png_path}") - - return out_png_path if __name__ == "__main__": lat, lon = 42.465774, -2.292634 - # lat, lon = 46.256440, 2.315916 delta = 20 now = datetime.today().strftime("%Y-%m-%d") look_from = (datetime.today() - timedelta(days=delta)).strftime("%Y-%m-%d") start_time = time.time() + + # geojson_path = CURR_SCRIPT_DIR.parent / "example.geojson" + # with open(geojson_path, "r") as f: + # geom = json.load(f) + # get_sr_image(lat, lon, look_from, now, geometry=geom) + get_sr_image(lat, lon, look_from, now) finish_time = time.time() logger.info(f"Total time:\t{(finish_time - start_time)/60:.1f} minutes") diff --git a/sen2sr_tools/model/SEN2SRLite/example_data.safetensor b/sen2sr_tools/model/SEN2SRLite/example_data.safetensor new file mode 100644 index 0000000..f7a95ad Binary files /dev/null and b/sen2sr_tools/model/SEN2SRLite/example_data.safetensor differ diff --git a/sen2sr_tools/model/SEN2SRLite/f2_hard_constraint.safetensor b/sen2sr_tools/model/SEN2SRLite/f2_hard_constraint.safetensor new file mode 100644 index 0000000..d4efe55 Binary files /dev/null and b/sen2sr_tools/model/SEN2SRLite/f2_hard_constraint.safetensor differ diff --git a/sen2sr_tools/model/SEN2SRLite/f2_model.safetensor b/sen2sr_tools/model/SEN2SRLite/f2_model.safetensor new file mode 100644 index 0000000..b1fb3af Binary files /dev/null and b/sen2sr_tools/model/SEN2SRLite/f2_model.safetensor differ diff --git a/sen2sr_tools/model/SEN2SRLite/hard_constraint.safetensor b/sen2sr_tools/model/SEN2SRLite/hard_constraint.safetensor new file mode 100644 index 0000000..564dbea Binary files /dev/null and b/sen2sr_tools/model/SEN2SRLite/hard_constraint.safetensor differ diff --git a/sen2sr_tools/model/SEN2SRLite/load.py b/sen2sr_tools/model/SEN2SRLite/load.py new file mode 100644 index 0000000..94f3756 --- /dev/null +++ b/sen2sr_tools/model/SEN2SRLite/load.py @@ -0,0 +1,211 @@ +import pathlib +import safetensors.torch +import matplotlib.pyplot as plt + +import sen2sr + +sen2sr.__version__ + +from sen2sr.models.opensr_baseline.cnn import CNNSR +from sen2sr.models.tricks import HardConstraint +from sen2sr.nonreference import srmodel as rgbn_model +from sen2sr.referencex2 import srmodel as rswir_modelx2 +from sen2sr.referencex4 import srmodel as rswir_modelx4 + +# MLSTAC API ----------------------------------------------------------------------- +def example_data(path: pathlib.Path, *args, **kwargs): + data_f = path / "example_data.safetensor" + sample = safetensors.torch.load_file(data_f) + return sample["lr"] + +def trainable_model(path, device: str = "cpu", *args, **kwargs): + ## RGBN 10m to 2.5m model ---------------------------------- + sr_rgbn_model_weights = safetensors.torch.load_file(path / "sr_model.safetensor") + sr_rgbn_model = CNNSR(4, 4, 24, 4, True, False, 6) + sr_rgbn_model.load_state_dict(sr_rgbn_model_weights) + sr_rgbn_model.to(device) + sr_rgbn_model = sr_rgbn_model.eval() + for param in sr_rgbn_model.parameters(): + param.requires_grad = False + sr_rgbn_hard_constraint_weights = safetensors.torch.load_file(path / "sr_hard_constraint.safetensor") + sr_rgbn_hard_constraint = HardConstraint(low_pass_mask=sr_rgbn_hard_constraint_weights["weights"].to(device), device=device) + sr_rgbn_model = rgbn_model( + sr_model=sr_rgbn_model, + hard_constraint=sr_rgbn_hard_constraint, + device=device + ) + + ## RSWIRs 20m to 10m model ---------------------------------- + sr_rswir_model_weights = safetensors.torch.load_file(path / "f2_model.safetensor") + params = { + "in_channels": 10, + "out_channels": 6, + "feature_channels": 24, + "upscale": 1, + "bias": True, + "train_mode": False, + "num_blocks": 6 + } + sr_rswir_model = CNNSR(**params) + sr_rswir_model.load_state_dict(sr_rswir_model_weights) + sr_rswir_model = sr_rswir_model.to(device) + sr_rswir_model = sr_rswir_model.eval() + for param in sr_rswir_model.parameters(): + param.requires_grad = False + sr_rswir_hard_constraint_weights = safetensors.torch.load_file(path / "f2_hard_constraint.safetensor") + sr_rswir_hard_constraint = HardConstraint( + low_pass_mask=sr_rswir_hard_constraint_weights["weights"].to(device), + bands= [0, 1, 2, 3, 4, 5], + device=device + ) + reference_srx2 = rswir_modelx2( + sr_model=sr_rswir_model, + hard_constraint=sr_rswir_hard_constraint, + device=device + ) + + + ## RSWIRs 10m to 2.5m model ---------------------------------- + # Load model parameters + sr_model_weights = safetensors.torch.load_file(path / "model.safetensor") + params = { + "in_channels": 10, + "out_channels": 6, + "feature_channels": 24, + "upscale": 1, + "bias": True, + "train_mode": False, + "num_blocks": 6 + } + sr_model = CNNSR(**params) + sr_model.load_state_dict(sr_model_weights) + sr_model = sr_model.to(device) + + # Load HardConstraint + hard_constraint_weights = safetensors.torch.load_file(path / "hard_constraint.safetensor") + hard_constraint = HardConstraint( + low_pass_mask=hard_constraint_weights["weights"].to(device), + bands= [0, 1, 2, 3, 4, 5], + device=device + ) + return rswir_modelx4(sr_rgbn_model, reference_srx2, sr_model, hard_constraint, device=device) + + +def compiled_model(path, device: str = "cpu", *args, **kwargs): + ## RGBN 10m to 2.5m model ---------------------------------- + sr_rgbn_model_weights = safetensors.torch.load_file(path / "sr_model.safetensor") + sr_rgbn_model = CNNSR(4, 4, 24, 4, True, False, 6) + sr_rgbn_model.load_state_dict(sr_rgbn_model_weights) + sr_rgbn_model.to(device) + sr_rgbn_model = sr_rgbn_model.eval() + for param in sr_rgbn_model.parameters(): + param.requires_grad = False + sr_rgbn_hard_constraint_weights = safetensors.torch.load_file(path / "sr_hard_constraint.safetensor") + sr_rgbn_hard_constraint = HardConstraint(low_pass_mask=sr_rgbn_hard_constraint_weights["weights"].to(device), device=device) + sr_rgbn_model = rgbn_model( + sr_model=sr_rgbn_model, + hard_constraint=sr_rgbn_hard_constraint, + device=device + ) + + ## RSWIRs 20m to 10m model ---------------------------------- + sr_rswir_model_weights = safetensors.torch.load_file(path / "f2_model.safetensor") + params = { + "in_channels": 10, + "out_channels": 6, + "feature_channels": 24, + "upscale": 1, + "bias": True, + "train_mode": False, + "num_blocks": 6 + } + sr_rswir_model = CNNSR(**params) + sr_rswir_model.load_state_dict(sr_rswir_model_weights) + sr_rswir_model = sr_rswir_model.to(device) + sr_rswir_model = sr_rswir_model.eval() + for param in sr_rswir_model.parameters(): + param.requires_grad = False + sr_rswir_hard_constraint_weights = safetensors.torch.load_file(path / "f2_hard_constraint.safetensor") + sr_rswir_hard_constraint = HardConstraint( + low_pass_mask=sr_rswir_hard_constraint_weights["weights"].to(device), + bands= [0, 1, 2, 3, 4, 5], + device=device + ) + reference_srx2 = rswir_modelx2( + sr_model=sr_rswir_model, + hard_constraint=sr_rswir_hard_constraint, + device=device + ) + + + ## RSWIRs 10m to 2.5m model ---------------------------------- + # Load model parameters + sr_model_weights = safetensors.torch.load_file(path / "model.safetensor") + params = { + "in_channels": 10, + "out_channels": 6, + "feature_channels": 24, + "upscale": 1, + "bias": True, + "train_mode": False, + "num_blocks": 6 + } + sr_model = CNNSR(**params) + sr_model.load_state_dict(sr_model_weights) + sr_model = sr_model.to(device) + sr_model = sr_model.eval() + for param in sr_model.parameters(): + param.requires_grad = False + + # Load HardConstraint + hard_constraint_weights = safetensors.torch.load_file(path / "hard_constraint.safetensor") + hard_constraint = HardConstraint( + low_pass_mask=hard_constraint_weights["weights"].to(device), + bands= [0, 1, 2, 3, 4, 5], + device=device + ) + return rswir_modelx4(sr_rgbn_model, reference_srx2, sr_model, hard_constraint, device=device) + + +def display_results(path: pathlib.Path, device: str = "cpu", *args, **kwargs): + # Load model + model = compiled_model(path, device) + + # Load data + lr = example_data(path) + + # Run model + sr = model(lr.to(device)) + + # Create the viz + lr_rgb = lr[0, [2, 1, 0]].cpu().numpy().transpose(1, 2, 0) + sr_rgb = sr[0, [2, 1, 0]].cpu().numpy().transpose(1, 2, 0) + + lr_swirs = lr[0, [9, 8, 7]].cpu().numpy().transpose(1, 2, 0) + sr_swirs = sr[0, [9, 8, 7]].cpu().numpy().transpose(1, 2, 0) + + lr_reds = lr[0, [6, 5, 4]].cpu().numpy().transpose(1, 2, 0) + sr_reds = sr[0, [6, 5, 4]].cpu().numpy().transpose(1, 2, 0) + + + #Display results + lr_slice = slice(16, 32+80) + hr_slice = slice(lr_slice.start*4, lr_slice.stop*4) + fig, ax = plt.subplots(3, 2, figsize=(8, 12)) + ax = ax.flatten() + ax[0].imshow(lr_rgb[lr_slice]*2) + ax[0].set_title("LR RGB") + ax[1].imshow(sr_rgb[hr_slice]*2) + ax[1].set_title("SR RGB") + ax[2].imshow(lr_swirs[lr_slice]*2) + ax[2].set_title("LR SWIR") + ax[3].imshow(sr_swirs[hr_slice]*2) + ax[3].set_title("SR SWIR") + ax[4].imshow(lr_reds[lr_slice]*2) + ax[4].set_title("LR RED") + ax[5].imshow(sr_reds[hr_slice]*2) + ax[5].set_title("SR RED") + for a in ax: + a.axis("off") + fig.tight_layout() + return fig diff --git a/sen2sr_tools/model/SEN2SRLite/mlm.json b/sen2sr_tools/model/SEN2SRLite/mlm.json new file mode 100644 index 0000000..031941d --- /dev/null +++ b/sen2sr_tools/model/SEN2SRLite/mlm.json @@ -0,0 +1,220 @@ +{ + "type": "Feature", + "stac_version": "1.1.0", + "stac_extensions": [ + "https://stac-extensions.github.io/mlm/v1.4.0/schema.json" + ], + "id": "SEN2SRLite_manin", + "geometry": { + "type": "Polygon", + "coordinates": [ + [ + [ + -180.0, + -90.0 + ], + [ + -180.0, + 90.0 + ], + [ + 180.0, + 90.0 + ], + [ + 180.0, + -90.0 + ], + [ + -180.0, + -90.0 + ] + ] + ] + }, + "bbox": [ + -180, + -90, + 180, + 90 + ], + "properties": { + "start_datetime": "1900-01-01T00:00:00Z", + "end_datetime": "9999-01-01T00:00:00Z", + "description": "A Swift Parameter-free Attention Network (SPAN) trained on the SEN2NAIPv2 dataset to enhance RSWIRs Sentinel-2 bands, improving spatial resolution from 10 meters to 2.5 meters.", + "dependencies": [ + "torch", + "safetensors.torch", + "supers2" + ], + "mlm:framework": "pytorch", + "mlm:framework_version": "2.1.2+cu121", + "file:size": 1889984, + "mlm:memory_size": 1, + "mlm:accelerator": "cuda", + "mlm:accelerator_constrained": false, + "mlm:accelerator_summary": "Unknown", + "mlm:name": "CNN_Light_F4", + "mlm:architecture": "SPAN", + "mlm:tasks": [ + "super-resolution" + ], + "mlm:input": [ + { + "name": "10 Band Sentinel-2", + "bands": [ + "B04", + "B03", + "B02", + "B08", + "B05", + "B06", + "B07", + "B8A", + "B11", + "B12" + ], + "input": { + "shape": [ + -1, + 10, + 128, + 128 + ], + "dim_order": [ + "batch", + "channel", + "height", + "width" + ], + "data_type": "float16" + }, + "pre_processing_function": null + } + ], + "mlm:output": [ + { + "name": "super-resolution", + "bands": [ + "B05", + "B06", + "B07", + "B8A", + "B11", + "B12" + ], + "tasks": [ + "super-resolution" + ], + "result": { + "shape": [ + -1, + 6, + 128, + 128 + ], + "dim_order": [ + "batch", + "channel", + "height", + "width" + ], + "data_type": "float16" + }, + "classification:classes": [], + "post_processing_function": null + } + ], + "mlm:total_parameters": 472496, + "mlm:pretrained": true, + "datetime": null + }, + "links": [], + "assets": { + "trainable": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/main/model.safetensor", + "type": "application/octet-stream; application=safetensor", + "title": "Pytorch model weights checkpoint", + "description": "The weights of the model in safetensor format.", + "mlm:artifact_type": "safetensor.torch.save_file", + "roles": [ + "mlm:model", + "mlm:weights", + "data" + ] + }, + "hardconstraint": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/main/hard_constraint.safetensor", + "type": "application/octet-stream; application=safetensor", + "title": "Pytorch hard constraint weights checkpoint. It is used to load faster the hard constraint module.", + "description": "The weights of the model in safetensor format.", + "mlm:artifact_type": "safetensor.torch.save_file", + "roles": [ + "mlm:model", + "mlm:weights", + "data" + ] + }, + "source_code": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/main/load.py", + "type": "text/x-python", + "title": "Model load script", + "description": "Python script to load the model.", + "roles": [ + "mlm:source_code", + "code" + ] + }, + "example_data": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/main/example_data.safetensor", + "type": "application/octet-stream; application=safetensors", + "title": "Example Sentinel-2 image", + "description": "Example Sentinel-2 image for model inference.", + "roles": [ + "mlm:example_data", + "data" + ] + }, + "sr_model": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/main/sr_model.safetensor", + "type": "application/octet-stream; application=safetensors", + "title": "Auxiliar RGBN super-resolution model", + "description": "Auxiliar RGBN super-resolution model used to convert RGBN band to 2.5m resolution.", + "roles": [ + "mlm:weights", + "data" + ] + }, + "sr_hard_constraint": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/main/sr_hard_constraint.safetensor", + "type": "application/octet-stream; application=safetensors", + "title": "Auxiliar RGBN super-resolution hard constraint", + "description": "Auxiliar RGBN super-resolution hard constraint used to convert RGBN band to 2.5m resolution.", + "roles": [ + "mlm:weights", + "data" + ] + }, + "f2_model": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/main/f2_model.safetensor", + "type": "application/octet-stream; application=safetensors", + "title": "Auxiliar RSWIRs super-resolution model", + "description": "Auxiliar RSWIRs super-resolution model used to convert RSWIRs band to 10m resolution.", + "roles": [ + "mlm:weights", + "data" + ] + }, + "f2_hard_constraint": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/main/f2_hard_constraint.safetensor", + "type": "application/octet-stream; application=safetensors", + "title": "Auxiliar RSWIRs super-resolution hard constraint", + "description": "Auxiliar RSWIRs super-resolution hard constraint used to convert RSWIRs band to 10m resolution.", + "roles": [ + "mlm:weights", + "data" + ] + } + }, + "collection": "SEN2SRLite" +} \ No newline at end of file diff --git a/sen2sr_tools/model/SEN2SRLite/model.safetensor b/sen2sr_tools/model/SEN2SRLite/model.safetensor new file mode 100644 index 0000000..99cdf10 Binary files /dev/null and b/sen2sr_tools/model/SEN2SRLite/model.safetensor differ diff --git a/sen2sr_tools/model/SEN2SRLite/sr_hard_constraint.safetensor b/sen2sr_tools/model/SEN2SRLite/sr_hard_constraint.safetensor new file mode 100644 index 0000000..7320aea Binary files /dev/null and b/sen2sr_tools/model/SEN2SRLite/sr_hard_constraint.safetensor differ diff --git a/sen2sr_tools/model/SEN2SRLite/sr_model.safetensor b/sen2sr_tools/model/SEN2SRLite/sr_model.safetensor new file mode 100644 index 0000000..53a4539 Binary files /dev/null and b/sen2sr_tools/model/SEN2SRLite/sr_model.safetensor differ diff --git a/sen2sr_tools/model/SEN2SRLite_RGBN/example_data.safetensor b/sen2sr_tools/model/SEN2SRLite_RGBN/example_data.safetensor new file mode 100644 index 0000000..f998dc2 Binary files /dev/null and b/sen2sr_tools/model/SEN2SRLite_RGBN/example_data.safetensor differ diff --git a/sen2sr_tools/model/SEN2SRLite_RGBN/hard_constraint.safetensor b/sen2sr_tools/model/SEN2SRLite_RGBN/hard_constraint.safetensor new file mode 100644 index 0000000..6d00b94 Binary files /dev/null and b/sen2sr_tools/model/SEN2SRLite_RGBN/hard_constraint.safetensor differ diff --git a/sen2sr_tools/model/SEN2SRLite_RGBN/load.py b/sen2sr_tools/model/SEN2SRLite_RGBN/load.py new file mode 100644 index 0000000..5f81358 --- /dev/null +++ b/sen2sr_tools/model/SEN2SRLite_RGBN/load.py @@ -0,0 +1,75 @@ +import pathlib +import safetensors.torch +import matplotlib.pyplot as plt +from sen2sr.models.opensr_baseline.cnn import CNNSR +from sen2sr.models.tricks import HardConstraint +from sen2sr.nonreference import srmodel + + +# MLSTAC API ----------------------------------------------------------------------- +def example_data(path: pathlib.Path, *args, **kwargs): + data_f = path / "example_data.safetensor" + sample = safetensors.torch.load_file(data_f) + return sample["lr"], sample["hr"] + +def trainable_model(path, device: str = "cpu", *args, **kwargs): + trainable_f = path / "model.safetensor" + + # Load model parameters + sr_model_weights = safetensors.torch.load_file(trainable_f) + sr_model = CNNSR(4, 4, 24, 4, True, False, 6) + sr_model.load_state_dict(sr_model_weights) + sr_model.to(device) + + # Load HardConstraint + hard_constraint_weights = safetensors.torch.load_file(path / "hard_constraint.safetensor") + hard_constraint = HardConstraint(low_pass_mask=hard_constraint_weights["weights"].to(device), device=device) + + return srmodel(sr_model, hard_constraint, device) + + +def compiled_model(path, device: str = "cpu", *args, **kwargs): + trainable_f = path / "model.safetensor" + + # Load model parameters + sr_model_weights = safetensors.torch.load_file(trainable_f) + sr_model = CNNSR(4, 4, 24, 4, True, False, 6) + sr_model.load_state_dict(sr_model_weights) + sr_model = sr_model.eval() + sr_model.to(device) + for param in sr_model.parameters(): + param.requires_grad = False + + # Load HardConstraint + hard_constraint_weights = safetensors.torch.load_file(path / "hard_constraint.safetensor") + hard_constraint = HardConstraint(low_pass_mask=hard_constraint_weights["weights"].to(device), device=device) + + return srmodel(sr_model, hard_constraint, device) + + +def display_results(path: pathlib.Path, device: str = "cpu", *args, **kwargs): + # Load model + model = compiled_model(path, device) + + # Load data + lr, hr = example_data(path) + + # Run model + SuperX = model(lr.to(device)) + + #Display results + Xrgb = lr[0, 0:3].cpu().numpy().transpose(1, 2, 0) + SuperXrgb = SuperX[0, 0:3].cpu().numpy().transpose(1, 2, 0) + lr_slice = slice(16, 32+80) + hr_slice = slice(lr_slice.start*4, lr_slice.stop*4) + fig, ax = plt.subplots(1, 3, figsize=(12, 4)) + ax[0].imshow(Xrgb[lr_slice, lr_slice]*3) + ax[0].set_title("Sentinel-2") + ax[1].imshow(SuperXrgb[hr_slice, hr_slice]*3) + ax[1].set_title("Super-Resolved") + ax[2].imshow(hr[0, 0:3].cpu().numpy().transpose(1, 2, 0)[hr_slice, hr_slice]*3) + ax[2].set_title("True HR") + for a in ax: + a.axis("off") + fig.tight_layout() + return fig \ No newline at end of file diff --git a/sen2sr_tools/model/SEN2SRLite_RGBN/mlm.json b/sen2sr_tools/model/SEN2SRLite_RGBN/mlm.json new file mode 100644 index 0000000..98e3cab --- /dev/null +++ b/sen2sr_tools/model/SEN2SRLite_RGBN/mlm.json @@ -0,0 +1,166 @@ +{ + "type": "Feature", + "stac_version": "1.1.0", + "stac_extensions": [ + "https://stac-extensions.github.io/mlm/v1.4.0/schema.json" + ], + "id": "SEN2SRLite_NonReference_RGBN_x4", + "geometry": { + "type": "Polygon", + "coordinates": [ + [ + [ + -180.0, + -90.0 + ], + [ + -180.0, + 90.0 + ], + [ + 180.0, + 90.0 + ], + [ + 180.0, + -90.0 + ], + [ + -180.0, + -90.0 + ] + ] + ] + }, + "bbox": [ + -180, + -90, + 180, + 90 + ], + "properties": { + "start_datetime": "1900-01-01T00:00:00Z", + "end_datetime": "9999-01-01T00:00:00Z", + "description": "A Swift Parameter-free Attention Network (SPAN) trained on the SEN2NAIPv2 dataset to enhance RGBN Sentinel-2 bands, improving spatial resolution from 10 meters to 2.5 meters.", + "dependencies": [ + "torch", + "safetensors.torch", + "sen2sr" + ], + "mlm:framework": "pytorch", + "mlm:framework_version": "2.1.2+cu121", + "file:size": 1889984, + "mlm:memory_size": 1, + "mlm:accelerator": "cuda", + "mlm:accelerator_constrained": false, + "mlm:accelerator_summary": "Unknown", + "mlm:name": "CNN_Light_SR", + "mlm:architecture": "SPAN", + "mlm:tasks": [ + "super-resolution" + ], + "mlm:input": [ + { + "name": "4 Band Sentinel-2 10m bands", + "bands": [ + "B04", + "B03", + "B02", + "B08" + ], + "input": { + "shape": [ + -1, + 4, + 128, + 128 + ], + "dim_order": [ + "batch", + "channel", + "height", + "width" + ], + "data_type": "float16" + }, + "pre_processing_function": null + } + ], + "mlm:output": [ + { + "name": "super-resolution", + "tasks": [ + "super-resolution" + ], + "result": { + "shape": [ + -1, + 4, + 512, + 512 + ], + "dim_order": [ + "batch", + "channel", + "height", + "width" + ], + "data_type": "float16" + }, + "classification:classes": [], + "post_processing_function": null + } + ], + "mlm:total_parameters": 472496, + "mlm:pretrained": true, + "datetime": null + }, + "links": [], + "assets": { + "trainable": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/NonReference_RGBN_x4/model.safetensor", + "type": "application/octet-stream; application=safetensor", + "title": "Pytorch model weights checkpoint", + "description": "The weights of the model in safetensor format.", + "mlm:artifact_type": "safetensor.torch.save_file", + "roles": [ + "mlm:model", + "mlm:weights", + "data" + ] + }, + "hardconstraint": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/NonReference_RGBN_x4/hard_constraint.safetensor", + "type": "application/octet-stream; application=safetensor", + "title": "Pytorch hard constraint weights checkpoint. It is used to load faster the hard constraint module.", + "description": "The weights of the model in safetensor format.", + "mlm:artifact_type": "safetensor.torch.save_file", + "roles": [ + "mlm:model", + "mlm:weights", + "data" + ] + }, + "source_code": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/NonReference_RGBN_x4/load.py", + "type": "text/x-python", + "title": "Model load script", + "description": "Python script to load the model.", + "roles": [ + "mlm:source_code", + "code" + ] + }, + "example_data": { + "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SRLite/NonReference_RGBN_x4/example_data.safetensor", + "type": "application/octet-stream; application=safetensors", + "title": "Example Sentinel-2 image", + "description": "Example Sentinel-2 image for model inference.", + "roles": [ + "mlm:example_data", + "data" + ] + } + }, + "collection": "SEN2SRLite" +} \ No newline at end of file diff --git a/sen2sr_tools/model/SEN2SRLite_RGBN/model.safetensor b/sen2sr_tools/model/SEN2SRLite_RGBN/model.safetensor new file mode 100644 index 0000000..53a4539 Binary files /dev/null and b/sen2sr_tools/model/SEN2SRLite_RGBN/model.safetensor differ diff --git a/sen2sr_tools/utils.py b/sen2sr_tools/utils.py index ef12828..c6a26cd 100644 --- a/sen2sr_tools/utils.py +++ b/sen2sr_tools/utils.py @@ -6,9 +6,12 @@ import numpy as np +from shapely.geometry import shape, mapping +from shapely.ops import transform +from PIL import Image, ImageEnhance +from pyproj import Transformer from rasterio.transform import from_bounds from xarray import DataArray -from PIL import Image, ImageEnhance from .constants import BRIGHTNESS_FACTOR, COMPARISON_PNG_FILEPATH, GAMMA, PNG_DIR, TIF_DIR @@ -18,16 +21,29 @@ # -------------------- -def reorder_bands(original_s2_numpy, superX): +def reorder_bands(image_data, is_sr_data): """ - Reorders image bands from NIR,B,G,R to R,G,B,NIR + Rearrange Sentinel-2 bands from the original order [NIR, B, G, R] into standard display order [R, G, B, NIR]. + + Args + --- + image_data: np.ndarray or torch.Tensor + Input image array + is_sr_data: bool + Set `True` if `image_data` is a torch tensor. If so, detach and convert it to a NumPy array before reordering. + + Returns + ------- + np.ndarray + Array with reordered bands. """ # Original: [NIR, B, G, R] -> reorder to [R, G, B, NIR] - band_order_tif = [3, 2, 1, 0] # indices in original array - original_s2_reordered = original_s2_numpy[band_order_tif] - superX_np = superX.detach().cpu().numpy() - superX_reordered = superX_np[band_order_tif] - return original_s2_reordered, superX_reordered + band_order = [3, 2, 1, 0] # [R, G, B, NIR] + + if is_sr_data: + image_data = image_data.detach().cpu().numpy() + + return image_data[band_order] def save_to_tif(array, filepath, sample, crs: str): @@ -272,3 +288,9 @@ def lonlat_to_utm_epsg(lon, lat): return f"EPSG:{32600 + zone}" # Northern hemisphere else: return f"EPSG:{32700 + zone}" # Southern hemisphere + +def reproject_geometry(geometry_dict, from_crs, to_crs): + geom = shape(geometry_dict) + transformer = Transformer.from_crs(from_crs, to_crs, always_xy=True) + geom_reproj = transform(lambda x, y: transformer.transform(x, y), geom) + return mapping(geom_reproj)