import streamlit as st import streamlit.components.v1 as components import torch import torch.nn as nn import torch.nn.functional as F from torchvision import models from torchvision.transforms import ToTensor, Resize import numpy as np from PIL import Image import math from obj2html import obj2html import os # DEPTH IMAGE TO OBJ minDepth=10 maxDepth=1000 def my_DepthNorm(x, maxDepth): return maxDepth / x def vete(v, vt): if v == vt: return str(v) return str(v)+"/"+str(vt) def create_obj(img, objPath='model.obj', mtlPath='model.mtl', matName='colored', useMaterial=False): w = img.shape[1] h = img.shape[0] FOV = math.pi/4 D = (img.shape[0]/2)/math.tan(FOV/2) if max(objPath.find('\\'), objPath.find('/')) > -1: os.makedirs(os.path.dirname(mtlPath), exist_ok=True) with open(objPath, "w") as f: if useMaterial: f.write("mtllib " + mtlPath + "\n") f.write("usemtl " + matName + "\n") ids = np.zeros((img.shape[1], img.shape[0]), int) vid = 1 all_x = [] all_y = [] all_z = [] for u in range(0, w): for v in range(h-1, -1, -1): d = img[v, u] ids[u, v] = vid if d == 0.0: ids[u, v] = 0 vid += 1 x = u - w/2 y = v - h/2 z = -D norm = 1 / math.sqrt(x*x + y*y + z*z) t = d/(z*norm) x = -t*x*norm y = t*y*norm z = -t*z*norm f.write("v " + str(x) + " " + str(y) + " " + str(z) + "\n") for u in range(0, img.shape[1]): for v in range(0, img.shape[0]): f.write("vt " + str(u/img.shape[1]) + " " + str(v/img.shape[0]) + "\n") for u in range(0, img.shape[1]-1): for v in range(0, img.shape[0]-1): v1 = ids[u, v] v3 = ids[u+1, v] v2 = ids[u, v+1] v4 = ids[u+1, v+1] if v1 == 0 or v2 == 0 or v3 == 0 or v4 == 0: continue f.write("f " + vete(v1, v1) + " " + vete(v2, v2) + " " + vete(v3, v3) + "\n") f.write("f " + vete(v3, v3) + " " + vete(v2, v2) + " " + vete(v4, v4) + "\n") # MODEL class UpSample(nn.Sequential): def __init__(self, skip_input, output_features): super(UpSample, self).__init__() self.convA = nn.Conv2d(skip_input, output_features, kernel_size=3, stride=1, padding=1) self.leakyreluA = nn.LeakyReLU(0.2) self.convB = nn.Conv2d(output_features, output_features, kernel_size=3, stride=1, padding=1) self.leakyreluB = nn.LeakyReLU(0.2) def forward(self, x, concat_with): up_x = F.interpolate(x, size=[concat_with.size(2), concat_with.size(3)], mode='bilinear', align_corners=True) return self.leakyreluB( self.convB( self.convA( torch.cat([up_x, concat_with], dim=1) ) ) ) class Decoder(nn.Module): def __init__(self, num_features=1664, decoder_width = 1.0): super(Decoder, self).__init__() features = int(num_features * decoder_width) self.conv2 = nn.Conv2d(num_features, features, kernel_size=1, stride=1, padding=0) self.up1 = UpSample(skip_input=features//1 + 256, output_features=features//2) self.up2 = UpSample(skip_input=features//2 + 128, output_features=features//4) self.up3 = UpSample(skip_input=features//4 + 64, output_features=features//8) self.up4 = UpSample(skip_input=features//8 + 64, output_features=features//16) self.conv3 = nn.Conv2d(features//16, 1, kernel_size=3, stride=1, padding=1) def forward(self, features): x_block0, x_block1, x_block2, x_block3, x_block4 = features[3], features[4], features[6], features[8], features[12] x_d0 = self.conv2(F.relu(x_block4)) x_d1 = self.up1(x_d0, x_block3) x_d2 = self.up2(x_d1, x_block2) x_d3 = self.up3(x_d2, x_block1) x_d4 = self.up4(x_d3, x_block0) return self.conv3(x_d4) class Encoder(nn.Module): def __init__(self): super(Encoder, self).__init__() self.original_model = models.densenet169( pretrained=False ) def forward(self, x): features = [x] for k, v in self.original_model.features._modules.items(): features.append( v(features[-1]) ) return features class PTModel(nn.Module): def __init__(self): super(PTModel, self).__init__() self.encoder = Encoder() self.decoder = Decoder() def forward(self, x): return self.decoder( self.encoder(x) ) model = PTModel().float() path = "https://github.com/nicolalandro/DenseDepth/releases/download/0.1/nyu.pth" model.load_state_dict(torch.hub.load_state_dict_from_url(path, progress=True)) model.eval() def predict(inp): width, height = inp.size if width > height: scale_fn = Resize((640, int((640*width)/height))) else: scale_fn = Resize((int((640*height)/width), 640)) res_img = scale_fn(inp) torch_image = ToTensor()(res_img) images = torch_image.unsqueeze(0) with torch.no_grad(): predictions = model(images) output = np.clip(my_DepthNorm(predictions.numpy(), maxDepth=maxDepth), minDepth, maxDepth) / maxDepth depth = output[0,0,:,:] img = Image.fromarray(np.uint8(depth*255)) create_obj(depth, 'model.obj') html_string = obj2html('model.obj', html_elements_only=True) return res_img, img, html_string # STREAMLIT uploader = st.file_uploader('Wait the demo file to be rendered and upload your favourite image here.',type=['jpg','jpeg','png']) if uploader is not None: pil_image = Image.open(uploader) else: pil_image = Image.open('119_image.png') with st.spinner("Waiting for the predictions..."): pil_scaled, pil_depth, html_string = predict(pil_image) components.html(html_string) #st.markdown(html_string, unsafe_allow_html=True) col1, col2, col3 = st.columns(3) with col1: st.image(pil_scaled) with col2: st.image(pil_depth) with col3: with open('model.obj') as f: st.download_button('Download model.obj', f, file_name="model.obj") os.remove('model.obj') pil_depth.save('tmp.png') with open('tmp.png', "rb") as f: st.download_button('Download depth.png', f,file_name="depth.png", mime="image/png") os.remove('tmp.png')