create app.py

app.py

@@ -0,0 +1,68 @@
+import gradio as gr
+from pathlib import Path
+import torch
+import torchvision.transforms as T
+from PIL import Image
+from utils import SSLModule
+from io import BytesIO
+import matplotlib.pyplot as plt
+import os
+
+# Load the model
+checkpoints_dir = Path("saved_checkpoints")
+if not checkpoints_dir.exists():
+	os.system("aws s3 --no-sign-request cp --recursive s3://dataforgood-fb-data/forests/v1/models/ .")
+checkpoint = "SSLhuge_satellite.pth"
+device = "cpu"
+ckpt_path = checkpoints_dir / checkpoint
+model = SSLModule(ssl_path=str(ckpt_path))
+model.to(device)
+model = model.eval()
+
+# Define the normalization transform
+norm = T.Normalize((0.420, 0.411, 0.296), (0.213, 0.156, 0.143))
+norm = norm.to(device)
+
+
+# Define a function to make predictions
+def predict(image):
+    # Convert PIL Image to tensor
+    image_t = torch.tensor(image).permute(2, 0, 1)[:3].float().to(device) / 255
+    # Normalize the image
+    with torch.no_grad():
+        pred = model(norm(image_t.unsqueeze(0)))
+        pred = pred.cpu().detach().relu()
+    # Convert tensor to numpy array
+    pred_np = pred[0, 0].numpy()
+    # Save the image to an in-memory buffer
+    buffer = BytesIO()
+    plt.imsave(buffer, pred_np, cmap="Greens")
+    buffer.seek(0)  # Rewind the buffer to the beginning
+    # Read the image back from the buffer
+    image_from_buffer = Image.open(buffer)
+    return image_from_buffer
+
+
+# create a Gradio interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(label="Upload a Satellite Image"),
+    outputs=gr.Image(label="Estimated Canopy Height"),
+	title="Estimate 🌳 Canopy Height from Satellite Image 🛰️",
+	description="""
+    <div style='display: flex; justify-content: center; align-items: center;'>
+        <img src='https://sustainability.fb.com/wp-content/uploads/2024/04/worldmap-2500.jpg?w=1536' style='max-width: 500px'/>
+    </div>
+    <p>This application uses a pre-trained model to estimate canopy height from satellite images. Upload an image and see the result!</p>
+    """,
+    examples=[
+        ["examples/image.png"],
+        ["examples/image2.png"],
+        ["examples/image3.png"],
+    ],
+	article="<p style='text-align: center'>Find more information <a href='https://sustainability.fb.com/blog/2024/04/22/using-artificial-intelligence-to-map-the-earths-forests/'>here</a>.</p>",
+	allow_flagging=False,
+)
+
+# launch the interface
+demo.launch()