import argparse
import logging
import multiprocessing as mp
import logging
import os
from detectron2.evaluation import inference_context
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from detectron2.utils.collect_env import collect_env_info
from detectron2.utils.logger import setup_logger
from fvcore.common.file_io import PathManager
from pathlib import Path
from pytorch3d.io import save_obj

from shapenet.config.config import get_shapenet_cfg
from shapenet.data.utils import imagenet_preprocess
from shapenet.modeling.heads import voxel_head
from shapenet.modeling.mesh_arch import build_model
from shapenet.utils.checkpoint import clean_state_dict

import torchvision.transforms as T

import glob
from PIL import Image

import trimesh
import pyvista as pv
import pyacvd
import numpy as np

logger = logging.getLogger('demo')

def setup_cfgs(args):
    cfg = get_shapenet_cfg()
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze()
    return cfg

def get_parser():
    parser = argparse.ArgumentParser(description="MeshRCNN Demo")
    parser.add_argument(
        "--config-file",
        default="configs/shapenet/voxmesh_R50.yaml",
        metavar="FILE",
        help="path to config file",
    )
    parser.add_argument("--input", help="A path to an input main folder")
    # parser.add_argument("--output", help="A directory to save output visualizations")
    parser.add_argument(
        "--focal-length", type=float, default=20.0, help="Focal length for the image"
    )
    parser.add_argument(
        "--onlyhighest", action="store_true", help="will return only the highest scoring detection"
    )

    parser.add_argument(
        "opts",
        help="Modify model config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
    return parser

def resample_mesh(mesh, count=2466):
    pv_mesh = pv.wrap(mesh)
    # logger.info('Original mesh:')
    # print(pv_mesh)
    
    clus = pyacvd.Clustering(pv_mesh)
    clus.subdivide(3)
    clus.cluster(count)

    # remesh
    remesh = clus.create_mesh()

    # verts = remesh.points
    # faces = remesh.faces.reshape((-1, 4))[:, 1:]
    
    return remesh

if __name__ == "__main__":
    mp.set_start_method("spawn", force=True)
    args = get_parser().parse_args()

    device = torch.device("cuda:%d" % 0)

    logger = setup_logger(name="demo shapenet")
    logger.info("Arguments: " + str(args))

    cfg = setup_cfgs(args)

    # load checkpoing and build model
    if cfg.MODEL.CHECKPOINT == "":
        raise ValueError("Invalid checkpoing provided")
    logger.info("Loading model from checkpoint: %s" % (cfg.MODEL.CHECKPOINT))
    cp = torch.load(PathManager.get_local_path(cfg.MODEL.CHECKPOINT))
    state_dict = clean_state_dict(cp["best_states"]["model"])
    model = build_model(cfg)
    model.load_state_dict(state_dict)
    logger.info("Model loaded")
    model.to(device)

    sub_dir = sorted(os.listdir(args.input))

    for sd in sub_dir:
        curr_path = os.path.join(args.input, sd)
        images = glob.glob(curr_path + "/*.png")
        
        for img_dir in images:
            # load image
            transform = [T.ToTensor()]
            transform.append(imagenet_preprocess())
            transform = T.Compose(transform)
            
            im_name = img_dir.split("/")[-1].split(".")[0]

            with PathManager.open(img_dir, "rb") as f:
                img = Image.open(f).convert("RGB")

            img = transform(img)
            img = img[None, :, :, :]
            img = img.to(device)

            with inference_context(model):
                img_feats, voxel_scores, meshes_pred, P, cubified_meshes = model(img)

            # Save voxel_score
            voxel_odir = os.path.join(curr_path, "voxel_score")
            if not Path(voxel_odir).is_dir():
                os.mkdir(voxel_odir)

            voxel_file = os.path.join(voxel_odir, "%s.pt" % (im_name))
            torch.save(voxel_scores, voxel_file)

            # Save image features
            imgfeat_odir = os.path.join(curr_path, "img_feat")
            if not Path(imgfeat_odir).is_dir():
                os.mkdir(imgfeat_odir)

            img_feat_file = os.path.join(imgfeat_odir, "%s.pt" % (im_name))
            torch.save(img_feats, img_feat_file)

            # Save P
            p_odir = os.path.join(curr_path, "P")
            if not Path(p_odir).is_dir():
                os.mkdir(p_odir)

            p_file = os.path.join(p_odir, "%s.pt" % (im_name))
            torch.save(P, p_file)

            # Save cubified mesh
            cmesh_odir = os.path.join(curr_path, "cube_mesh")
            if not Path(cmesh_odir).is_dir():
                os.mkdir(cmesh_odir)

            cube_mesh_file = os.path.join(cmesh_odir, "%s_cube.obj" % (im_name))
            c_verts, c_faces = cubified_meshes[-1].get_mesh_verts_faces(0)
            save_obj(cube_mesh_file, c_verts, c_faces)

            # Save predicted mesh
            mesh_odir = os.path.join(curr_path, "final_mesh")
            if not Path(mesh_odir).is_dir():
                os.mkdir(mesh_odir)

            save_file = os.path.join(mesh_odir, "%s.obj" % (im_name))
            verts, faces = meshes_pred[-1].get_mesh_verts_faces(0)
            save_obj(save_file, verts, faces)
            logger.info("Predictions saved for %s/%s" % (curr_path.split('/')[-1], im_name))

[((840, 865), 'logging.getLogger', 'logging.getLogger', (['"""demo"""'], {}), "('demo')\n", (857, 865), False, 'import logging\n'), ((899, 917), 'shapenet.config.config.get_shapenet_cfg', 'get_shapenet_cfg', ([], {}), '()\n', (915, 917), False, 'from shapenet.config.config import get_shapenet_cfg\n'), ((1059, 1111), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""MeshRCNN Demo"""'}), "(description='MeshRCNN Demo')\n", (1082, 1111), False, 'import argparse\n'), ((1932, 1945), 'pyvista.wrap', 'pv.wrap', (['mesh'], {}), '(mesh)\n', (1939, 1945), True, 'import pyvista as pv\n'), ((2019, 2045), 'pyacvd.Clustering', 'pyacvd.Clustering', (['pv_mesh'], {}), '(pv_mesh)\n', (2036, 2045), False, 'import pyacvd\n'), ((2273, 2313), 'torch.multiprocessing.set_start_method', 'mp.set_start_method', (['"""spawn"""'], {'force': '(True)'}), "('spawn', force=True)\n", (2292, 2313), True, 'import torch.multiprocessing as mp\n'), ((2365, 2392), 'torch.device', 'torch.device', (["('cuda:%d' % 0)"], {}), "('cuda:%d' % 0)\n", (2377, 2392), False, 'import torch\n'), ((2407, 2441), 'detectron2.utils.logger.setup_logger', 'setup_logger', ([], {'name': '"""demo shapenet"""'}), "(name='demo shapenet')\n", (2419, 2441), False, 'from detectron2.utils.logger import setup_logger\n'), ((2808, 2852), 'shapenet.utils.checkpoint.clean_state_dict', 'clean_state_dict', (["cp['best_states']['model']"], {}), "(cp['best_states']['model'])\n", (2824, 2852), False, 'from shapenet.utils.checkpoint import clean_state_dict\n'), ((2865, 2881), 'shapenet.modeling.mesh_arch.build_model', 'build_model', (['cfg'], {}), '(cfg)\n', (2876, 2881), False, 'from shapenet.modeling.mesh_arch import build_model\n'), ((2741, 2789), 'fvcore.common.file_io.PathManager.get_local_path', 'PathManager.get_local_path', (['cfg.MODEL.CHECKPOINT'], {}), '(cfg.MODEL.CHECKPOINT)\n', (2767, 2789), False, 'from fvcore.common.file_io import PathManager\n'), ((2995, 3017), 'os.listdir', 'os.listdir', (['args.input'], {}), '(args.input)\n', (3005, 3017), False, 'import os\n'), ((3063, 3091), 'os.path.join', 'os.path.join', (['args.input', 'sd'], {}), '(args.input, sd)\n', (3075, 3091), False, 'import os\n'), ((3109, 3140), 'glob.glob', 'glob.glob', (["(curr_path + '/*.png')"], {}), "(curr_path + '/*.png')\n", (3118, 3140), False, 'import glob\n'), ((3321, 3341), 'torchvision.transforms.Compose', 'T.Compose', (['transform'], {}), '(transform)\n', (3330, 3341), True, 'import torchvision.transforms as T\n'), ((3812, 3850), 'os.path.join', 'os.path.join', (['curr_path', '"""voxel_score"""'], {}), "(curr_path, 'voxel_score')\n", (3824, 3850), False, 'import os\n'), ((3960, 4003), 'os.path.join', 'os.path.join', (['voxel_odir', "('%s.pt' % im_name)"], {}), "(voxel_odir, '%s.pt' % im_name)\n", (3972, 4003), False, 'import os\n'), ((4018, 4054), 'torch.save', 'torch.save', (['voxel_scores', 'voxel_file'], {}), '(voxel_scores, voxel_file)\n', (4028, 4054), False, 'import torch\n'), ((4117, 4152), 'os.path.join', 'os.path.join', (['curr_path', '"""img_feat"""'], {}), "(curr_path, 'img_feat')\n", (4129, 4152), False, 'import os\n'), ((4269, 4314), 'os.path.join', 'os.path.join', (['imgfeat_odir', "('%s.pt' % im_name)"], {}), "(imgfeat_odir, '%s.pt' % im_name)\n", (4281, 4314), False, 'import os\n'), ((4329, 4365), 'torch.save', 'torch.save', (['img_feats', 'img_feat_file'], {}), '(img_feats, img_feat_file)\n', (4339, 4365), False, 'import torch\n'), ((4409, 4437), 'os.path.join', 'os.path.join', (['curr_path', '"""P"""'], {}), "(curr_path, 'P')\n", (4421, 4437), False, 'import os\n'), ((4535, 4574), 'os.path.join', 'os.path.join', (['p_odir', "('%s.pt' % im_name)"], {}), "(p_odir, '%s.pt' % im_name)\n", (4547, 4574), False, 'import os\n'), ((4589, 4610), 'torch.save', 'torch.save', (['P', 'p_file'], {}), '(P, p_file)\n', (4599, 4610), False, 'import torch\n'), ((4670, 4706), 'os.path.join', 'os.path.join', (['curr_path', '"""cube_mesh"""'], {}), "(curr_path, 'cube_mesh')\n", (4682, 4706), False, 'import os\n'), ((4820, 4869), 'os.path.join', 'os.path.join', (['cmesh_odir', "('%s_cube.obj' % im_name)"], {}), "(cmesh_odir, '%s_cube.obj' % im_name)\n", (4832, 4869), False, 'import os\n'), ((4959, 5001), 'pytorch3d.io.save_obj', 'save_obj', (['cube_mesh_file', 'c_verts', 'c_faces'], {}), '(cube_mesh_file, c_verts, c_faces)\n', (4967, 5001), False, 'from pytorch3d.io import save_obj\n'), ((5061, 5098), 'os.path.join', 'os.path.join', (['curr_path', '"""final_mesh"""'], {}), "(curr_path, 'final_mesh')\n", (5073, 5098), False, 'import os\n'), ((5205, 5248), 'os.path.join', 'os.path.join', (['mesh_odir', "('%s.obj' % im_name)"], {}), "(mesh_odir, '%s.obj' % im_name)\n", (5217, 5248), False, 'import os\n'), ((5330, 5363), 'pytorch3d.io.save_obj', 'save_obj', (['save_file', 'verts', 'faces'], {}), '(save_file, verts, faces)\n', (5338, 5363), False, 'from pytorch3d.io import save_obj\n'), ((3231, 3243), 'torchvision.transforms.ToTensor', 'T.ToTensor', ([], {}), '()\n', (3241, 3243), True, 'import torchvision.transforms as T\n'), ((3274, 3295), 'shapenet.data.utils.imagenet_preprocess', 'imagenet_preprocess', ([], {}), '()\n', (3293, 3295), False, 'from shapenet.data.utils import imagenet_preprocess\n'), ((3432, 3463), 'fvcore.common.file_io.PathManager.open', 'PathManager.open', (['img_dir', '"""rb"""'], {}), "(img_dir, 'rb')\n", (3448, 3463), False, 'from fvcore.common.file_io import PathManager\n'), ((3643, 3667), 'detectron2.evaluation.inference_context', 'inference_context', (['model'], {}), '(model)\n', (3660, 3667), False, 'from detectron2.evaluation import inference_context\n'), ((3913, 3933), 'os.mkdir', 'os.mkdir', (['voxel_odir'], {}), '(voxel_odir)\n', (3921, 3933), False, 'import os\n'), ((4217, 4239), 'os.mkdir', 'os.mkdir', (['imgfeat_odir'], {}), '(imgfeat_odir)\n', (4225, 4239), False, 'import os\n'), ((4496, 4512), 'os.mkdir', 'os.mkdir', (['p_odir'], {}), '(p_odir)\n', (4504, 4512), False, 'import os\n'), ((4769, 4789), 'os.mkdir', 'os.mkdir', (['cmesh_odir'], {}), '(cmesh_odir)\n', (4777, 4789), False, 'import os\n'), ((5160, 5179), 'os.mkdir', 'os.mkdir', (['mesh_odir'], {}), '(mesh_odir)\n', (5168, 5179), False, 'import os\n'), ((3492, 3505), 'PIL.Image.open', 'Image.open', (['f'], {}), '(f)\n', (3502, 3505), False, 'from PIL import Image\n'), ((3870, 3886), 'pathlib.Path', 'Path', (['voxel_odir'], {}), '(voxel_odir)\n', (3874, 3886), False, 'from pathlib import Path\n'), ((4172, 4190), 'pathlib.Path', 'Path', (['imgfeat_odir'], {}), '(imgfeat_odir)\n', (4176, 4190), False, 'from pathlib import Path\n'), ((4457, 4469), 'pathlib.Path', 'Path', (['p_odir'], {}), '(p_odir)\n', (4461, 4469), False, 'from pathlib import Path\n'), ((4726, 4742), 'pathlib.Path', 'Path', (['cmesh_odir'], {}), '(cmesh_odir)\n', (4730, 4742), False, 'from pathlib import Path\n'), ((5118, 5133), 'pathlib.Path', 'Path', (['mesh_odir'], {}), '(mesh_odir)\n', (5122, 5133), False, 'from pathlib import Path\n')]