Spaces:
Build error
Build error
File size: 5,361 Bytes
4d85df4 |
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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
from pytorch_lightning import Trainer
from monoscene.models.monoscene import MonoScene
from monoscene.data.NYU.nyu_dm import NYUDataModule
from monoscene.data.semantic_kitti.kitti_dm import KittiDataModule
from monoscene.data.kitti_360.kitti_360_dm import Kitti360DataModule
# import hydra
from omegaconf import DictConfig
import torch
import numpy as np
import os
from hydra.utils import get_original_cwd
import gradio as gr
import numpy as np
import plotly.express as px
import pandas as pd
# @hydra.main(config_name="../config/monoscene.yaml")
def plot(input_img):
torch.set_grad_enabled(False)
# Setup dataloader
# if config.dataset == "kitti" or config.dataset == "kitti_360":
feature = 64
project_scale = 2
full_scene_size = (256, 256, 32)
# if config.dataset == "kitti":
# data_module = KittiDataModule(
# root=config.kitti_root,
# preprocess_root=config.kitti_preprocess_root,
# frustum_size=config.frustum_size,
# batch_size=int(config.batch_size / config.n_gpus),
# num_workers=int(config.num_workers_per_gpu * config.n_gpus),
# )
# data_module.setup()
# data_loader = data_module.val_dataloader()
# # data_loader = data_module.test_dataloader() # use this if you want to infer on test set
# else:
# data_module = Kitti360DataModule(
# root=config.kitti_360_root,
# sequences=[config.kitti_360_sequence],
# n_scans=2000,
# batch_size=1,
# num_workers=3,
# )
# data_module.setup()
# data_loader = data_module.dataloader()
# elif config.dataset == "NYU":
# project_scale = 1
# feature = 200
# full_scene_size = (60, 36, 60)
# data_module = NYUDataModule(
# root=config.NYU_root,
# preprocess_root=config.NYU_preprocess_root,
# n_relations=config.n_relations,
# frustum_size=config.frustum_size,
# batch_size=int(config.batch_size / config.n_gpus),
# num_workers=int(config.num_workers_per_gpu * config.n_gpus),
# )
# data_module.setup()
# data_loader = data_module.val_dataloader()
# # data_loader = data_module.test_dataloader() # use this if you want to infer on test set
# else:
# print("dataset not support")
# Load pretrained models
# if config.dataset == "NYU":
# model_path = os.path.join(
# get_original_cwd(), "trained_models", "monoscene_nyu.ckpt"
# )
# else:
# model_path = os.path.join(
# get_original_cwd(), "trained_models", "monoscene_kitti.ckpt"
# )
model_path = "trained_models/monoscene_kitti.ckpt"
model = MonoScene.load_from_checkpoint(
model_path,
feature=feature,
project_scale=project_scale,
fp_loss=False,
full_scene_size=full_scene_size,
)
model.cuda()
model.eval()
print(input_img.shape)
x = np.arange(12).reshape(4, 3) / 12
data = pd.DataFrame(data=x, columns=['x', 'y', 'z'])
fig = px.scatter_3d(data, x="x", y="y", z="z")
return fig
demo = gr.Interface(plot, gr.Image(shape=(200, 200)), gr.Plot())
demo.launch()
# Save prediction and additional data
# to draw the viewing frustum and remove scene outside the room for NYUv2
# output_path = os.path.join(config.output_path, config.dataset)
# with torch.no_grad():
# for batch in tqdm(data_loader):
# batch["img"] = batch["img"].cuda()
# pred = model(batch)
# y_pred = torch.softmax(pred["ssc_logit"], dim=1).detach().cpu().numpy()
# y_pred = np.argmax(y_pred, axis=1)
# for i in range(config.batch_size):
# out_dict = {"y_pred": y_pred[i].astype(np.uint16)}
# if "target" in batch:
# out_dict["target"] = (
# batch["target"][i].detach().cpu().numpy().astype(np.uint16)
# )
# if config.dataset == "NYU":
# write_path = output_path
# filepath = os.path.join(write_path, batch["name"][i] + ".pkl")
# out_dict["cam_pose"] = batch["cam_pose"][i].detach().cpu().numpy()
# out_dict["vox_origin"] = (
# batch["vox_origin"][i].detach().cpu().numpy()
# )
# else:
# write_path = os.path.join(output_path, batch["sequence"][i])
# filepath = os.path.join(write_path, batch["frame_id"][i] + ".pkl")
# out_dict["fov_mask_1"] = (
# batch["fov_mask_1"][i].detach().cpu().numpy()
# )
# out_dict["cam_k"] = batch["cam_k"][i].detach().cpu().numpy()
# out_dict["T_velo_2_cam"] = (
# batch["T_velo_2_cam"][i].detach().cpu().numpy()
# )
# os.makedirs(write_path, exist_ok=True)
# with open(filepath, "wb") as handle:
# pickle.dump(out_dict, handle)
# print("wrote to", filepath) |