Spaces:
Sleeping
Sleeping
File size: 2,376 Bytes
13d5a13 |
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 |
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()
|