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))

[((32, 9, 32, 34), 'logging.getLogger', 'logging.getLogger', ({(32, 27, 32, 33): '"""demo"""'}, {}), "('demo')", False, 'import logging\n'), ((35, 10, 35, 28), 'shapenet.config.config.get_shapenet_cfg', 'get_shapenet_cfg', ({}, {}), '()', False, 'from shapenet.config.config import get_shapenet_cfg\n'), ((42, 13, 42, 65), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((67, 14, 67, 27), 'pyvista.wrap', 'pv.wrap', ({(67, 22, 67, 26): 'mesh'}, {}), '(mesh)', True, 'import pyvista as pv\n'), ((71, 11, 71, 37), 'pyacvd.Clustering', 'pyacvd.Clustering', ({(71, 29, 71, 36): 'pv_mesh'}, {}), '(pv_mesh)', False, 'import pyacvd\n'), ((84, 4, 84, 44), 'torch.multiprocessing.set_start_method', 'mp.set_start_method', (), '', True, 'import torch.multiprocessing as mp\n'), ((87, 13, 87, 40), 'torch.device', 'torch.device', ({(87, 26, 87, 39): "'cuda:%d' % 0"}, {}), "('cuda:%d' % 0)", False, 'import torch\n'), ((89, 13, 89, 47), 'detectron2.utils.logger.setup_logger', 'setup_logger', (), '', False, 'from detectron2.utils.logger import setup_logger\n'), ((99, 17, 99, 61), 'shapenet.utils.checkpoint.clean_state_dict', 'clean_state_dict', ({(99, 34, 99, 60): "cp['best_states']['model']"}, {}), "(cp['best_states']['model'])", False, 'from shapenet.utils.checkpoint import clean_state_dict\n'), ((100, 12, 100, 28), 'shapenet.modeling.mesh_arch.build_model', 'build_model', ({(100, 24, 100, 27): 'cfg'}, {}), '(cfg)', False, 'from shapenet.modeling.mesh_arch import build_model\n'), ((98, 20, 98, 68), 'fvcore.common.file_io.PathManager.get_local_path', 'PathManager.get_local_path', ({(98, 47, 98, 67): 'cfg.MODEL.CHECKPOINT'}, {}), '(cfg.MODEL.CHECKPOINT)', False, 'from fvcore.common.file_io import PathManager\n'), ((105, 21, 105, 43), 'os.listdir', 'os.listdir', ({(105, 32, 105, 42): 'args.input'}, {}), '(args.input)', False, 'import os\n'), ((108, 20, 108, 48), 'os.path.join', 'os.path.join', ({(108, 33, 108, 43): 'args.input', (108, 45, 108, 47): 'sd'}, {}), '(args.input, sd)', False, 'import os\n'), ((109, 17, 109, 48), 'glob.glob', 'glob.glob', ({(109, 27, 109, 47): "curr_path + '/*.png'"}, {}), "(curr_path + '/*.png')", False, 'import glob\n'), ((115, 24, 115, 44), 'torchvision.transforms.Compose', 'T.Compose', ({(115, 34, 115, 43): 'transform'}, {}), '(transform)', True, 'import torchvision.transforms as T\n'), ((130, 25, 130, 63), 'os.path.join', 'os.path.join', ({(130, 38, 130, 47): 'curr_path', (130, 49, 130, 62): '"""voxel_score"""'}, {}), "(curr_path, 'voxel_score')", False, 'import os\n'), ((134, 25, 134, 70), 'os.path.join', 'os.path.join', ({(134, 38, 134, 48): 'voxel_odir', (134, 50, 134, 69): "'%s.pt' % im_name"}, {}), "(voxel_odir, '%s.pt' % im_name)", False, 'import os\n'), ((135, 12, 135, 48), 'torch.save', 'torch.save', ({(135, 23, 135, 35): 'voxel_scores', (135, 37, 135, 47): 'voxel_file'}, {}), '(voxel_scores, voxel_file)', False, 'import torch\n'), ((138, 27, 138, 62), 'os.path.join', 'os.path.join', ({(138, 40, 138, 49): 'curr_path', (138, 51, 138, 61): '"""img_feat"""'}, {}), "(curr_path, 'img_feat')", False, 'import os\n'), ((142, 28, 142, 75), 'os.path.join', 'os.path.join', ({(142, 41, 142, 53): 'imgfeat_odir', (142, 55, 142, 74): "'%s.pt' % im_name"}, {}), "(imgfeat_odir, '%s.pt' % im_name)", False, 'import os\n'), ((143, 12, 143, 48), 'torch.save', 'torch.save', ({(143, 23, 143, 32): 'img_feats', (143, 34, 143, 47): 'img_feat_file'}, {}), '(img_feats, img_feat_file)', False, 'import torch\n'), ((146, 21, 146, 49), 'os.path.join', 'os.path.join', ({(146, 34, 146, 43): 'curr_path', (146, 45, 146, 48): '"""P"""'}, {}), "(curr_path, 'P')", False, 'import os\n'), ((150, 21, 150, 62), 'os.path.join', 'os.path.join', ({(150, 34, 150, 40): 'p_odir', (150, 42, 150, 61): "'%s.pt' % im_name"}, {}), "(p_odir, '%s.pt' % im_name)", False, 'import os\n'), ((151, 12, 151, 33), 'torch.save', 'torch.save', ({(151, 23, 151, 24): 'P', (151, 26, 151, 32): 'p_file'}, {}), '(P, p_file)', False, 'import torch\n'), ((154, 25, 154, 61), 'os.path.join', 'os.path.join', ({(154, 38, 154, 47): 'curr_path', (154, 49, 154, 60): '"""cube_mesh"""'}, {}), "(curr_path, 'cube_mesh')", False, 'import os\n'), ((158, 29, 158, 80), 'os.path.join', 'os.path.join', ({(158, 42, 158, 52): 'cmesh_odir', (158, 54, 158, 79): "'%s_cube.obj' % im_name"}, {}), "(cmesh_odir, '%s_cube.obj' % im_name)", False, 'import os\n'), ((160, 12, 160, 54), 'pytorch3d.io.save_obj', 'save_obj', ({(160, 21, 160, 35): 'cube_mesh_file', (160, 37, 160, 44): 'c_verts', (160, 46, 160, 53): 'c_faces'}, {}), '(cube_mesh_file, c_verts, c_faces)', False, 'from pytorch3d.io import save_obj\n'), ((163, 24, 163, 61), 'os.path.join', 'os.path.join', ({(163, 37, 163, 46): 'curr_path', (163, 48, 163, 60): '"""final_mesh"""'}, {}), "(curr_path, 'final_mesh')", False, 'import os\n'), ((167, 24, 167, 69), 'os.path.join', 'os.path.join', ({(167, 37, 167, 46): 'mesh_odir', (167, 48, 167, 68): "'%s.obj' % im_name"}, {}), "(mesh_odir, '%s.obj' % im_name)", False, 'import os\n'), ((169, 12, 169, 45), 'pytorch3d.io.save_obj', 'save_obj', ({(169, 21, 169, 30): 'save_file', (169, 32, 169, 37): 'verts', (169, 39, 169, 44): 'faces'}, {}), '(save_file, verts, faces)', False, 'from pytorch3d.io import save_obj\n'), ((113, 25, 113, 37), 'torchvision.transforms.ToTensor', 'T.ToTensor', ({}, {}), '()', True, 'import torchvision.transforms as T\n'), ((114, 29, 114, 50), 'shapenet.data.utils.imagenet_preprocess', 'imagenet_preprocess', ({}, {}), '()', False, 'from shapenet.data.utils import imagenet_preprocess\n'), ((119, 17, 119, 48), 'fvcore.common.file_io.PathManager.open', 'PathManager.open', ({(119, 34, 119, 41): 'img_dir', (119, 43, 119, 47): '"""rb"""'}, {}), "(img_dir, 'rb')", False, 'from fvcore.common.file_io import PathManager\n'), ((126, 17, 126, 41), 'detectron2.evaluation.inference_context', 'inference_context', ({(126, 35, 126, 40): 'model'}, {}), '(model)', False, 'from detectron2.evaluation import inference_context\n'), ((132, 16, 132, 36), 'os.mkdir', 'os.mkdir', ({(132, 25, 132, 35): 'voxel_odir'}, {}), '(voxel_odir)', False, 'import os\n'), ((140, 16, 140, 38), 'os.mkdir', 'os.mkdir', ({(140, 25, 140, 37): 'imgfeat_odir'}, {}), '(imgfeat_odir)', False, 'import os\n'), ((148, 16, 148, 32), 'os.mkdir', 'os.mkdir', ({(148, 25, 148, 31): 'p_odir'}, {}), '(p_odir)', False, 'import os\n'), ((156, 16, 156, 36), 'os.mkdir', 'os.mkdir', ({(156, 25, 156, 35): 'cmesh_odir'}, {}), '(cmesh_odir)', False, 'import os\n'), ((165, 16, 165, 35), 'os.mkdir', 'os.mkdir', ({(165, 25, 165, 34): 'mesh_odir'}, {}), '(mesh_odir)', False, 'import os\n'), ((120, 22, 120, 35), 'PIL.Image.open', 'Image.open', ({(120, 33, 120, 34): 'f'}, {}), '(f)', False, 'from PIL import Image\n'), ((131, 19, 131, 35), 'pathlib.Path', 'Path', ({(131, 24, 131, 34): 'voxel_odir'}, {}), '(voxel_odir)', False, 'from pathlib import Path\n'), ((139, 19, 139, 37), 'pathlib.Path', 'Path', ({(139, 24, 139, 36): 'imgfeat_odir'}, {}), '(imgfeat_odir)', False, 'from pathlib import Path\n'), ((147, 19, 147, 31), 'pathlib.Path', 'Path', ({(147, 24, 147, 30): 'p_odir'}, {}), '(p_odir)', False, 'from pathlib import Path\n'), ((155, 19, 155, 35), 'pathlib.Path', 'Path', ({(155, 24, 155, 34): 'cmesh_odir'}, {}), '(cmesh_odir)', False, 'from pathlib import Path\n'), ((164, 19, 164, 34), 'pathlib.Path', 'Path', ({(164, 24, 164, 33): 'mesh_odir'}, {}), '(mesh_odir)', False, 'from pathlib import Path\n')]