Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import numpy as np | |
| import PIL.Image as pil | |
| import torch | |
| import os | |
| import io | |
| from torchvision import transforms | |
| import matplotlib as mpl | |
| import matplotlib.cm as cm | |
| import networks | |
| from layers import disp_to_depth_no_scaling | |
| # Function to load the model | |
| def load_model(device, model_path): | |
| encoder_path = os.path.join(model_path, "encoder.pth") | |
| depth_decoder_path = os.path.join(model_path, "depth.pth") | |
| encoder = networks.ResnetEncoder(18, False) | |
| loaded_dict_enc = torch.load(encoder_path, map_location=device) | |
| feed_height = loaded_dict_enc['height'] | |
| feed_width = loaded_dict_enc['width'] | |
| filtered_dict_enc = {k: v for k, v in loaded_dict_enc.items() if k in encoder.state_dict()} | |
| encoder.load_state_dict(filtered_dict_enc) | |
| encoder.to(device) | |
| encoder.eval() | |
| depth_decoder = networks.DepthDecoder(num_ch_enc=encoder.num_ch_enc, scales=range(4)) | |
| loaded_dict = torch.load(depth_decoder_path, map_location=device) | |
| depth_decoder.load_state_dict(loaded_dict, strict=False) | |
| depth_decoder.to(device) | |
| depth_decoder.eval() | |
| return encoder, depth_decoder, feed_height, feed_width | |
| # Function to apply the model to an image | |
| def predict_depth(image, encoder, depth_decoder, device, feed_height, feed_width): | |
| input_image = image.resize((feed_width, feed_height), pil.LANCZOS) | |
| input_image = transforms.ToTensor()(input_image).unsqueeze(0).to(device) | |
| features = encoder(input_image) | |
| outputs = depth_decoder(features) | |
| disp = outputs[("disp", 0)] | |
| disp_resized = disp | |
| disp_resized_np = disp_resized.squeeze().cpu().numpy() | |
| vmax = np.percentile(disp_resized_np, 95) | |
| normalizer = mpl.colors.Normalize(vmin=disp_resized_np.min(), vmax=vmax) | |
| mapper = cm.ScalarMappable(norm=normalizer, cmap='magma') | |
| colormapped_im = (mapper.to_rgba(disp_resized_np)[:, :, :3] * 255).astype(np.uint8) | |
| return colormapped_im | |
| # Streamlit app | |
| def main(): | |
| st.title("Monodepth2 Depth Estimation") | |
| st.write("Upload a PNG image to get depth estimation") | |
| if not os.path.exits("depth.pth"): | |
| subprocess.run(["wget", "https://github.com/meghakalia/depthEstimationColonoscopy/releases/download/0.0.1/depth.pth"]) | |
| if not os.path.exits("encoder.pth"): | |
| subprocess.run(["wget", "https://github.com/meghakalia/depthEstimationColonoscopy/releases/download/0.0.1/encoder.pth"]) | |
| uploaded_file = st.file_uploader("Choose a PNG file", type="png") | |
| if uploaded_file is not None: | |
| image = pil.open(uploaded_file).convert('RGB') | |
| st.image(image, caption='Uploaded Image', use_column_width=True) | |
| if st.button('Predict Depth'): | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model_path = '.' # Update this to the path of your model | |
| encoder, depth_decoder, feed_height, feed_width = load_model(device, model_path) | |
| colormapped_im = predict_depth(image, encoder, depth_decoder, device, feed_height, feed_width) | |
| depth_image = pil.fromarray(colormapped_im) | |
| st.image(depth_image, caption='Predicted Depth Image', use_column_width=True) | |
| # Convert depth image to bytes | |
| img_bytes = io.BytesIO() | |
| depth_image.save(img_bytes, format='PNG') | |
| img_bytes = img_bytes.getvalue() | |
| # Provide download link | |
| st.download_button( | |
| label="Download Depth Image", | |
| data=img_bytes, | |
| file_name="depth_image.png", | |
| mime="image/png" | |
| ) | |
| if __name__ == "__main__": | |
| main() | |