Add GeoTIFF support

app.py

@@ -7,6 +7,9 @@ from pathlib import Path
 import gradio as gr
 from PIL import Image
 import spaces
+import numpy as np
+import rasterio
+from inference.geotiff import normalize_rgb, save_sr_geotiff
 
 # ============================================================
 # PATHS
@@ -223,12 +226,12 @@ body, .gradio-container {
 # INFERENCE FUNCTION
 # ============================================================
 
-def run_super_resolution(image, model_choice, progress=gr.Progress()):
+def run_super_resolution(image_path, model_choice, progress=gr.Progress()):
 
     # --------------------------------------------------------
     # Validate input
     # --------------------------------------------------------
-    if image is None:
+    if image_path is None:
         raise gr.Error("⚠️ Please upload a satellite image before running.")
 
     # --------------------------------------------------------
@@ -241,13 +244,50 @@ def run_super_resolution(image, model_choice, progress=gr.Progress()):
         raise gr.Error("⚠️ Please select a valid model.")
 
     # --------------------------------------------------------
-    # Ensure PIL image + RGB
+    # Load image (PNG/JPG or GeoTIFF)
     # --------------------------------------------------------
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
+    if isinstance(image_path, list):
+        image_path = image_path[0]
+    
+    file_path = str(image_path)
+    is_tiff = file_path.lower().endswith(('.tif', '.tiff'))
+    
+    if is_tiff:
+        try:
+            with rasterio.open(file_path) as src:
+                profile = src.profile.copy()
+                count = src.count
+                if count not in [1, 3, 4]:
+                    raise gr.Error(f"⚠️ Unsupported band count: {count}. Expected 1, 3, or 4 bands.")
+                
+                if count == 3:
+                    data = src.read([1, 2, 3])
+                    data = np.transpose(data, (1, 2, 0))
+                elif count == 4:
+                    data = src.read([1, 2, 3])
+                    data = np.transpose(data, (1, 2, 0))
+                else:
+                    data = src.read(1)
+                    data = np.stack([data, data, data], axis=-1)
+                
+                input_w, input_h = src.width, src.height
+        except rasterio.errors.RasterioIOError:
+            raise gr.Error("⚠️ Corrupted TIFF or unsupported format.")
+        except Exception as e:
+            raise gr.Error(f"⚠️ Error reading TIFF: {str(e)}")
+
+        if data.dtype == np.uint8:
+            pass
+        else:
+            data = normalize_rgb(data)
 
-    image = image.convert("RGB")
-    input_w, input_h = image.size
+        image = Image.fromarray(data, mode="RGB")
+    else:
+        try:
+            image = Image.open(file_path).convert("RGB")
+            input_w, input_h = image.size
+        except Exception:
+            raise gr.Error("⚠️ Could not read image.")
 
     # --------------------------------------------------------
     # Basic sanity check
@@ -286,14 +326,24 @@ def run_super_resolution(image, model_choice, progress=gr.Progress()):
     output_w, output_h = result.size
 
     # --------------------------------------------------------
-    # Save to a temp PNG for download
+    # Save to a temp file for download
     # --------------------------------------------------------
-    tmp = tempfile.NamedTemporaryFile(
-        suffix=f"_sr_{selected_model}_4x.png",
-        delete=False
-    )
-    result.save(tmp.name, format="PNG")
-    tmp.close()
+    if is_tiff:
+        tmp = tempfile.NamedTemporaryFile(
+            suffix=f"_sr_{selected_model}_4x_georef.tif",
+            delete=False
+        )
+        tmp.close()
+        save_sr_geotiff(result, profile, tmp.name, scale=4)
+        out_file_path = tmp.name
+    else:
+        tmp = tempfile.NamedTemporaryFile(
+            suffix=f"_sr_{selected_model}_4x.png",
+            delete=False
+        )
+        result.save(tmp.name, format="PNG")
+        tmp.close()
+        out_file_path = tmp.name
 
     # --------------------------------------------------------
     # Info markdown
@@ -316,7 +366,7 @@ def run_super_resolution(image, model_choice, progress=gr.Progress()):
 """
 
     progress(1.0, desc="Done.")
-    return result, info, tmp.name
+    return result, info, out_file_path
 
 
 # ============================================================
@@ -360,10 +410,9 @@ with gr.Blocks(
         with gr.Column(scale=5):
 
             gr.HTML('<div class="section-label">📂 Input Image</div>')
-            input_image = gr.Image(
-                label="Upload Satellite Image",
-                type="pil",
-                sources=["upload", "clipboard"],
+            input_image = gr.File(
+                label="Upload Satellite Image (PNG, JPG, TIFF)",
+                file_types=["image", ".tif", ".tiff"],
                 elem_id="input-img",
             )
 
@@ -410,7 +459,7 @@ with gr.Blocks(
             )
 
             download_file = gr.File(
-                label="⬇️ Download Result (PNG)",
+                label="⬇️ Download Result",
                 interactive=False,
                 visible=True,
                 elem_id="download-btn",