Update app.py

app.py

@@ -10,13 +10,102 @@ from typing import Dict
 import functools
 import inspect
 from types import SimpleNamespace
-import torch
 from torch.utils.data import Dataset
 from torchvision import transforms
 import rasterio
 from pathlib import Path
 from torchvision.transforms import ToPILImage
-import numpy as np
+from base64 import b64encode
+import gc
+from datasets import load_dataset
+import torchvision
+import torch.nn.functional as F
+from IPython.display import HTML
+from matplotlib import pyplot as plt
+from pathlib import Path
+from torch import autocast
+from torchvision import transforms as tfms
+from tqdm.auto import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer, logging
+import os
+import csv
+from torchvision.utils import save_image
+import torch
+import cv2
+from PIL import Image
+import os
+from django.conf import settings
+import torch.nn.functional as F
+import os
+import torch
+from transformers import AutoImageProcessor, SwinModel
+from diffusers import UNet2DConditionModel
+
+def load_models():
+    torch_device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    image_processor_model_path = os.path.join(settings.BASE_DIR, 'depthAPI', 'models', 'image_processor')
+    swin_transformer_model_path = os.path.join(settings.BASE_DIR, 'depthAPI', 'models', 'swin_transformer')
+    vae_model_path = os.path.join(settings.BASE_DIR, 'depthAPI', 'models', 'vae', 'MonoChannelVAE.pth')
+    unet_model_path = os.path.join(settings.BASE_DIR, 'depthAPI', 'models', 'unet')
+
+    image_processor = AutoImageProcessor.from_pretrained(image_processor_model_path)
+    swin_transformer = SwinModel.from_pretrained(swin_transformer_model_path)
+
+    vae = Autoencoder()
+    vae.load_state_dict(torch.load(vae_model_path, map_location=torch.device('cpu')))
+    unet = UNet2DConditionModel.from_pretrained(unet_model_path)
+    scheduler = DDIMScheduler(beta_start=0.0001, beta_end=0.02, beta_schedule='linear',
+                                         num_train_timesteps=1000)
+
+    vae = vae.to(torch_device)
+    swin_transformer = swin_transformer.to(torch_device)
+    unet = unet.to(torch_device)
+
+    return image_processor, swin_transformer, vae, unet, scheduler
+
+def tensor_to_latent(input_im,vae):
+  with torch.no_grad():
+    latent = vae.encoder(input_im)
+  return latent
+
+def latent_to_tensor(input_im,vae):
+  with torch.no_grad():
+    images = vae.decoder(input_im)
+  return images
+
+def upscale_resolution(image):
+  sr = cv2.dnn_superres.DnnSuperResImpl_create()
+  path = os.path.join(settings.BASE_DIR, 'depthAPI', 'models', 'FSRCNN','FSRCNN_x2.pb')
+  sr.readModel(path)
+  sr.setModel("fsrcnn",2)
+  result = sr.upsample(image)
+  resized = cv2.resize(image,dsize=None,fx=2,fy=2)
+  img = Image.fromarray(resized.astype('uint8'))
+  return img
+
+def extract_features(image,torch_device,swin_transformer):
+  image.to(torch_device)
+  with torch.no_grad():
+    swin_output = swin_transformer(**image)
+  del image
+  image_fea = swin_output.last_hidden_state.squeeze(0)
+  return image_fea
+
+def rescale(image):
+  max_val = torch.max(image)
+  min_val = torch.min(image)
+
+  image = (((image - min_val) / (max_val - min_val)) * 2) - 1
+  return image
+
+def normalize(x):
+  return 2 * (x - x.min()) / (x.max() - x.min()) - 1
+
+def upscale_tensor(image):
+  output = F.interpolate(image.unsqueeze(0), size=(512, 512), mode='bilinear', align_corners=False)
+  return output.squeeze(0)
+
 
 class UAHiRISEDataset(Dataset):
     def __init__(self, root, stage, transform=None):
@@ -457,13 +546,49 @@ class DDIMScheduler():
     def __len__(self):
         return self.config.num_train_timesteps
 
+
+image_processor, swin_transformer, vae, unet, scheduler = load_models()
+
+def MonoGeoDepthModelRun(image):
+    batch_size=1
+    torch_device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+
+
+    image = image.convert("RGB")
+    extracted_image = image_processor(image, return_tensors="pt")
+    image_embeddings = extract_features(extracted_image, torch_device, swin_transformer)
+    image_embeddings = image_embeddings.unsqueeze(0)
+
+    torch.manual_seed(0)
+    random_noise = normalize(torch.randn(1, 1, 512, 512).to(torch_device))
+
+    image_embeddings = image_embeddings.to(torch_device)
+
+    with torch.no_grad():
+        noisy_latents = tensor_to_latent(random_noise, vae)
+        del random_noise
+        t = torch.tensor(1000)
+        model_input = scheduler.scale_model_input(noisy_latents, t)
+        noise_pred = unet(model_input, t, encoder_hidden_states=image_embeddings, return_dict=False)
+        noisy_latents = model_input - noise_pred[0]
+        predicted_dtm = latent_to_tensor(noisy_latents, vae)
+        predicted_dtm = predicted_dtm.detach().cpu()
+
+        image_ = predicted_dtm.squeeze(0)
+        image_ = (image_ - image_.min()) / (image_.max() - image_.min())
+
+        to_pil = ToPILImage()
+        predicted_dtm = to_pil(image_)
+
+    return predicted_dtm
         
-def dummy_model(img):
+def model(img):
     img_array = np.array(img)
     return img_array
 
 iface = gr.Interface(
-    fn=dummy_model,
+    fn=model,
     inputs="image",
     outputs="image"
 )