| import os |
| os.system("python -m pip install --upgrade pip") |
| import gradio as gr |
| import numpy as np |
| import os |
| import torch |
| os.system('pip install --upgrade --no-cache-dir gdown') |
| import openai |
|
|
| from PIL import Image |
|
|
| from utils.logger import get_logger |
| from config.defaults import get_config |
| from inference import preprocess, run_one_inference |
| from models.build import build_model |
| from argparse import Namespace |
| import gdown |
|
|
| |
| def summarize_and_classify(text): |
| |
| response = openai.ChatCompletion.create( |
| model="gpt-3.5-turbo", |
| messages=[ |
| {"role": "system", "content": "고객 취향을 듣고, 어떤 인테리어가 어울릴지 추천해줘."}, |
| {"role": "user", "content": text}, |
| ], |
| max_tokens=300, |
| temperature=0.4 |
| ) |
|
|
| return response.choices[0].message["content"] |
|
|
| def down_ckpt(model_cfg, ckpt_dir): |
| model_ids = [ |
| ['src/config/mp3d.yaml', '1CCEEwY12EkMrokF6TEeGjMao5SEl26FQ'], |
| ['src/config/zind.yaml', '1ie4KQsB3uhd4ghZ7OTYt-XYbVRJ8w4ce'], |
| ['src/config/pano.yaml', '1Ria7Bl4THDqQ3z8pL6bIGUmkTgtEzsuK'], |
| ['src/config/s2d3d.yaml', '1P_EK8SreFCmA_LLScWCcvvpn-0l3sU1m'], |
| ['src/config/ablation_study/full.yaml', '16wQFUC76-tfVlGes_avh7QwE9rcBrDXX'] |
| ] |
|
|
| for model_id in model_ids: |
| if model_id[0] != model_cfg: |
| continue |
| path = os.path.join(ckpt_dir, 'best.pkl') |
| if not os.path.exists(path): |
| logger.info(f"Downloading {model_id}") |
| os.makedirs(ckpt_dir, exist_ok=True) |
| gdown.download(f"https://drive.google.com/uc?id={model_id[1]}", path, False) |
| gdown.download(f"https://drive.google.com/uc?id={model_id[1]}", path, False) |
|
|
|
|
| def greet(img_path, pre_processing, weight_name, post_processing, visualization, mesh_format, mesh_resolution): |
| args.pre_processing = pre_processing |
| args.post_processing = post_processing |
| if weight_name == 'mp3d': |
| model = mp3d_model |
| elif weight_name == 'zind': |
| model = zind_model |
| else: |
| logger.error("unknown pre-trained weight name") |
| raise NotImplementedError |
|
|
| img_name = os.path.basename(img_path).split('.')[0] |
| img = np.array(Image.open(img_path).resize((1024, 512), Image.Resampling.BICUBIC))[..., :3] |
|
|
| vp_cache_path = 'src/demo/default_vp.txt' |
| if args.pre_processing: |
| vp_cache_path = os.path.join('src/output', f'{img_name}_vp.txt') |
| logger.info("pre-processing ...") |
| img, vp = preprocess(img, vp_cache_path=vp_cache_path) |
|
|
| img = (img / 255.0).astype(np.float32) |
| run_one_inference(img, model, args, img_name, |
| logger=logger, show=False, |
| show_depth='depth-normal-gradient' in visualization, |
| show_floorplan='2d-floorplan' in visualization, |
| mesh_format=mesh_format, mesh_resolution=int(mesh_resolution)) |
|
|
| return [os.path.join(args.output_dir, f"{img_name}_pred.png"), |
| os.path.join(args.output_dir, f"{img_name}_3d{mesh_format}"), |
| os.path.join(args.output_dir, f"{img_name}_3d{mesh_format}"), |
| vp_cache_path, |
| os.path.join(args.output_dir, f"{img_name}_pred.json")] |
|
|
|
|
| def get_model(args): |
| config = get_config(args) |
| down_ckpt(args.cfg, config.CKPT.DIR) |
| if ('cuda' in args.device or 'cuda' in config.TRAIN.DEVICE) and not torch.cuda.is_available(): |
| logger.info(f'The {args.device} is not available, will use cpu ...') |
| config.defrost() |
| args.device = "cpu" |
| config.TRAIN.DEVICE = "cpu" |
| config.freeze() |
| model, _, _, _ = build_model(config, logger) |
| return model |
|
|
|
|
| if __name__ == '__main__': |
| logger = get_logger() |
| args = Namespace(device='cuda', output_dir='src/output', visualize_3d=False, output_3d=True) |
| os.makedirs(args.output_dir, exist_ok=True) |
|
|
| args.cfg = 'src/config/mp3d.yaml' |
| mp3d_model = get_model(args) |
|
|
| args.cfg = 'src/config/zind.yaml' |
| zind_model = get_model(args) |
|
|
| description = "(주)플랜포스_인트로우 3D 스캐닝 고도화" |
|
|
| demo = gr.Interface(fn=greet, |
| inputs=[gr.Image(type='filepath', label='input rgb panorama', value='src/demo/demo.jpg'), |
| gr.Checkbox(label='pre-processing', value=True), |
| gr.Radio(['mp3d', 'zind'], |
|
|
| value='mp3d', |
| visible=False), |
| gr.Radio(['manhattan', 'atalanta', 'original'], |
|
|
| value='manhattan', |
| visible=False), |
| gr.CheckboxGroup(['depth-normal-gradient', '2d-floorplan'], |
|
|
| value=['2d-floorplan'], |
| visible=False), |
| gr.Radio(['.gltf', '.obj', '.glb'], |
| |
| value='.obj', |
| visible=False), |
| gr.Radio(['128', '256', '512', '1024'], |
|
|
| value='128', |
| visible=False), |
| ], |
| outputs=[gr.Image(label='predicted result 2d-visualization', type='filepath'), |
| gr.Model3D(label='3d mesh reconstruction', clear_color=[1.0, 1.0, 1.0, 1.0]), |
| gr.File(label='3d mesh') |
| |
| |
| |
| ], |
| |
| title='room3d', allow_flagging="never", cache_examples=False, description=description) |
|
|
| demo.launch(debug=True, enable_queue=False) |
