Spaces:
Running
Running
| """Model registry and lazy fetch from the Hugging Face model repo. | |
| The Space and the FraXtex QGIS plugin consume the same published artifacts, so what | |
| you see here is what QGIS users get. Weights are not committed to this Space: they | |
| are downloaded on first use and cached by `huggingface_hub`. | |
| `channels` below drives the UI only — whether the DEM upload is shown before any | |
| model has been opened. At inference time the channel count is read from the ONNX | |
| graph itself (`inference.load_session`), which stays authoritative. | |
| """ | |
| from huggingface_hub import hf_hub_download | |
| REPO_ID = "ayoubft/fraXteX" | |
| MODELS = { | |
| "unet-rgb": { | |
| "filename": "rgb/unet-rgb.onnx", | |
| "label": "U-Net (RGB)", | |
| "channels": 3, | |
| "note": ( | |
| "U-Net trained from scratch, 31 M parameters. Normalizes inside the graph " | |
| "with dataset statistics, and accepts any patch size that is a multiple of 16." | |
| ), | |
| }, | |
| "sam2-rgb": { | |
| "filename": "rgb/sam2-rgb.onnx", | |
| "label": "SAM 2 (RGB) — slow", | |
| "channels": 3, | |
| "note": ( | |
| "SAM 2 (hiera-tiny) with a frozen encoder, prompt-free. Runs at 512x512 " | |
| "patches, so it is noticeably slower than the U-Net on this CPU Space." | |
| ), | |
| }, | |
| "unet-rgbdem": { | |
| "filename": "rgbdem/unet-rgbdem.onnx", | |
| "label": "U-Net (RGB+DEM)", | |
| "channels": 4, | |
| "note": ( | |
| "Same U-Net architecture, taking a DEM as a 4th channel alongside the RGB " | |
| "bands. Fixed 256x256 patches, inputs scaled to 0-1." | |
| ), | |
| }, | |
| "segformer-rgbdem": { | |
| "filename": "rgbdem/segformer-rgbdem.onnx", | |
| "label": "SegFormer (RGB+DEM)", | |
| "channels": 4, | |
| "note": ( | |
| "SegFormer with a transformer encoder, taking a DEM as a 4th channel " | |
| "alongside the RGB bands. Fixed 256x256 patches, inputs scaled to 0-1." | |
| ), | |
| }, | |
| } | |
| # Dropdown order: the RGB models first, since they need no DEM upload. | |
| CHOICES = [(spec["label"], key) for key, spec in MODELS.items()] | |
| DEFAULT = "unet-rgb" | |
| _cache = {} | |
| def needs_dem(key): | |
| return MODELS[key]["channels"] == 4 | |
| def note_for(key): | |
| return MODELS[key]["note"] | |
| def get_model_path(key, feedback=None): | |
| """Download the model on first use and return its local path.""" | |
| if key in _cache: | |
| return _cache[key] | |
| spec = MODELS[key] | |
| if feedback is not None: | |
| feedback.pushInfo(f"Fetching {spec['filename']} from {REPO_ID} (first use only)...") | |
| path = hf_hub_download(repo_id=REPO_ID, filename=spec["filename"]) | |
| _cache[key] = path | |
| return path | |