def RenderFitting_allSamples_panoptic()

in eft/db_processing/AMT_sample_gen.py [0:0]


def RenderFitting_allSamples_panoptic(inputDir_list, bVisSpinOriginal=False, bSampling= True, bRenderInit = False):

    #1K random selection file
    g_rootDirName = '/run/media/hjoo/disk/data/cocoPose3D_amt/0_SPIN/0_allSampleVis'
    # selectionFileDir ='/run/media/hjoo/disk/data/cocoPose3D_amt/0_SPIN/0_quant/AMT_1Kset_fromChallenging'
    # selectionOutoutTxt = os.path.join(selectionFileDir, 'selected1K.json')


    metaDir_list ={}
    metaDir_list['coco'] ='/run/media/hjoo/disk/data/cocoPose3D_amt/0_SPIN/0_exemplarOutput/meta_coco_2dskeletons'
    metaDir_list['mpii'] = '/run/media/hjoo/disk/data/cocoPose3D_amt/0_SPIN/0_exemplarOutput/meta_mpii_2dskeletons'
    metaDir_list['lspet'] = '/run/media/hjoo/disk/data/cocoPose3D_amt/0_SPIN/0_exemplarOutput/meta_lspet_2dskeletons'
    # metaDir_list['pennaction'] = '/run/media/hjoo/disk/data/cocoPose3D_amt/0_SPIN/0_exemplarOutput/meta_lspet_2dskeletons'

    imgDir ={}
    imgDir['mpii'] = '/run/media/hjoo/disk/data/mpii_human_pose_v1/images'
    imgDir['lspet'] = '/run/media/hjoo/disk/data/lspet_dataset/images_highres'
    imgDir['coco'] =  '/run/media/hjoo/disk/data/coco/train2014'
    imgDir['pennaction'] =  '/run/media/hjoo/disk/data/Penn_Action/frames'
    imgDir['panoptic'] =  '/run/media/hjoo/disk/data/panoptic_mtc/a4_release/hdImgs'

    
    totalFileList = []
    for dbName in inputDir_list:
        
        rootDirName_db = os.path.join(g_rootDirName, os.path.basename(inputDir_list[dbName]))
        if bVisSpinOriginal:
            glViewer.SetMeshColor('red')
        else:
            glViewer.SetMeshColor('blue')
        
        if False== os.path.exists(rootDirName_db):
            os.mkdir(rootDirName_db)
            
        #Folder Name
        imageFolder = os.path.join(rootDirName_db, 'images')
        if not os.path.exists(imageFolder):
            os.mkdir(imageFolder)

        croppedimageFolder = os.path.join(rootDirName_db, 'croppedImage')
        if not os.path.exists(croppedimageFolder):
            os.mkdir(croppedimageFolder)

        skelCroppedimageFolder_ = os.path.join(rootDirName_db, 'skelCroppedImage')
        if not os.path.exists(skelCroppedimageFolder_):
            os.mkdir(skelCroppedimageFolder_)

        overlaidImageFolder = os.path.join(rootDirName_db, 'overlaid')
        if not os.path.exists(overlaidImageFolder):
            os.mkdir(overlaidImageFolder)

        sideImageFolder = os.path.join(rootDirName_db, 'side')
        if not os.path.exists(sideImageFolder):
            os.mkdir(sideImageFolder)

        mergedImageFolder = os.path.join(rootDirName_db, 'merged')
        if not os.path.exists(mergedImageFolder):
            os.mkdir(mergedImageFolder)




        inputDir = inputDir_list[dbName]
        fileList  = listdir(inputDir)
        for fileName in fileList:
            
            if bRenderInit:
                if ("_init" in fileName):
                    totalFileList.append( (dbName,fileName))
            else:
                if ("_init" in fileName) == False:
                    totalFileList.append( (dbName,fileName))
    
        selectedSampleList = totalFileList

        # #Generate Output
        rejectedNum = 0
        
        loss_keypoints_2d_list =[]
        maxShape_list =[]
        for idx, sample in enumerate(sorted(selectedSampleList)):

            dbName = sample[0]
            fileName = sample[1]
            inputDir = inputDir_list[dbName]
            fileFullPath = join(inputDir, fileName)


            # if '246' in fileName:
            #     print(data['loss_keypoints_2d'])
            #     print(data['pred_shape'])
            #     here=0

            print(fileFullPath)
            with open(fileFullPath,'rb') as f:
                dataList = pickle.load(f)

            # if not isinstance(dataList, dict):
            #     dataListDict ={}
            #     dataListDict['0'] = dataList
            #     dataList = dataListDict


            # dataList = dataList[5:6]
            for dataKey in dataList:
                data = dataList[dataKey]

                # loss_keypoints_2d_list.append(data['loss_keypoints_2d'])
                # maxShape_list.append( np.max(abs(data['pred_shape'])))
                # continue
                

                if bRenderInit:
                    fileFullPath_newName = join(inputDir, fileName.replace("_init",""))
                    with open(fileFullPath_newName,'rb') as f:
                        datameta = pickle.load(f)
                        data['pred_camera'] = datameta['pred_camera']


                if False:#IsValid(data) == False:
                    rejectedNum+=1
                    print("Rejected: so far{}".format(rejectedNum))
                    continue

                # #Load meta data
                # metaDir = metaDir_list[dbName]
                # meta_fileFullPath = join(metaDir, fileName)
                # with open(meta_fileFullPath,'rb') as f:
                #     metadata = pickle.load(f)
                #     data['keypoint2d'] = metadata['keypoint2d']

                imagePath = data['imageName'][0]
                if dbName=='pennaction':
                    imgFullPath =os.path.join(imgDir[dbName], os.path.basename(os.path.dirname(imagePath)), os.path.basename(imagePath) )
                elif dbName=='panoptic':
                    imgFullPath =os.path.join(imgDir[dbName], os.path.basename( os.path.dirname(os.path.dirname(imagePath))),  os.path.basename(os.path.dirname(imagePath)), os.path.basename(imagePath) )
                else:
                    # imgFullPath =os.path.join(imgDir, os.path.basename(imgFullPath) )
                    imgFullPath =os.path.join(imgDir[dbName], os.path.basename(imagePath) )

                print(imgFullPath)

                # outputFileName = '{0}-{1}-{2}.jpg'.format(dbName, os.path.basename(imgFullPath), idx )
                outputFileName = imgFullPath.replace('/','-')
                mergedImgFileName = '{0}/{1}'.format(mergedImageFolder,outputFileName )
                if os.path.exists(mergedImgFileName):
                    continue

                assert os.path.exists(imgFullPath)
                rawImg = cv2.imread(imgFullPath)

                scale = data['scale'][0]
                center = data['center'][0]
                croppedImg_highRes = crop(rawImg, center, scale, 
                            [constants.IMG_RES*3, constants.IMG_RES*3])

                #Crop image
                croppedImg = crop(rawImg, center, scale, 
                            [constants.IMG_RES, constants.IMG_RES])

                validity = data['keypoint2d'][0,:,2] ==1
                valid_keypoints = data['keypoint2d'][0,validity,:2]
                min_pt = np.min(valid_keypoints, axis=0)
                max_pt = np.max(valid_keypoints, axis=0)
                bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]                    

                bbox[0] = max(bbox[0]-20,0)
                bbox[1] = max(bbox[1]-20,0)
                bbox[2] = bbox[2]+20
                bbox[3] = bbox[3]+20

                rawImage_wSkel = viewer2D.Vis_Skeleton_2D_SPIN49(data['keypoint2d'][0,:,:2], pt2d_visibility=data['keypoint2d'][0,:,2], image= rawImg.copy())
                rawImage_wSkel = viewer2D.Vis_Bbox(rawImage_wSkel,bbox)
                croppedImg_highRes_withSkel = crop(rawImage_wSkel, center, scale, 
                            [constants.IMG_RES*3, constants.IMG_RES*3])

                if False:
                    viewer2D.ImShow(croppedImg_highRes,name='croppedImg_highRes', waitTime=1)
                    viewer2D.ImShow(croppedImg_highRes_withSkel,waitTime=1,name='croppedImg_highRes_withSkel')

                    viewer2D.ImShow(croppedImg,name='croppedImg', waitTime=1)
                    viewer2D.ImShow(rawImg,waitTime=0)


                #Export Image
                if True:
                    imgFilePath = os.path.join(imageFolder,outputFileName)
                    cv2.imwrite(imgFilePath,rawImg)

                    imgFilePath = os.path.join(croppedimageFolder,outputFileName) 
                    cv2.imwrite(imgFilePath,croppedImg_highRes)

                    imgFilePath = os.path.join(skelCroppedimageFolder_,outputFileName)  
                    cv2.imwrite(imgFilePath,croppedImg_highRes_withSkel)

                #Vis overlaid Img
                ours_betas = torch.from_numpy(data['pred_shape'])
                ours_pose_rotmat = torch.from_numpy(data['pred_pose_rotmat'])
                spin_betas = torch.from_numpy(data['opt_beta'])
                spin_pose = torch.from_numpy(data['opt_pose'])
                pred_camera_vis = data['pred_camera']


                #Visualize SMPL output
                # Note that gt_model_joints is different from gt_joints as it comes from SMPL
                ours_output = smpl(betas=ours_betas, body_pose=ours_pose_rotmat[:,1:], global_orient=ours_pose_rotmat[:,0].unsqueeze(1), pose2rot=False )
                ours_vertices = ours_output.vertices.detach().cpu().numpy() 
                ours_joints_3d = ours_output.joints.detach().cpu().numpy() 

                spin_output = smpl(betas=spin_betas, body_pose=spin_pose[:,3:], global_orient=spin_pose[:,:3])
                spin_vertices = spin_output.vertices.detach().cpu().numpy() 
                spin_joints_3d = spin_output.joints.detach().cpu().numpy() 

                #Centering
                if False:
                    spin_pelvis = (ours_joints_3d[:, 27:28,:3] + ours_joints_3d[:,28:29,:3]) / 2
                    ours_joints_3d[:,:,:3] = ours_joints_3d[:,:,:3] - spin_pelvis        #centering
                    ours_vertices = ours_vertices - spin_pelvis

                    spin_model_pelvis = (spin_joints_3d[:, 27:28,:3] + spin_joints_3d[:,28:29,:3]) / 2
                    spin_joints_3d[:,:,:3] = spin_joints_3d[:,:,:3] - spin_model_pelvis        #centering
                    spin_vertices = spin_vertices - spin_model_pelvis

                b =0
                ############### Visualize Mesh ############### 
                pred_vert_vis = ours_vertices[b]
                # meshVertVis = gt_vertices[b].detach().cpu().numpy() 
                # meshVertVis = meshVertVis-pelvis        #centering
                pred_vert_vis *=pred_camera_vis[b,0]
                pred_vert_vis[:,0] += pred_camera_vis[b,1]        #no need +1 (or  112). Rendernig has this offset already
                pred_vert_vis[:,1] += pred_camera_vis[b,2]        #no need +1 (or  112). Rendernig has this offset already
                pred_vert_vis*=112
                pred_meshes = {'ver': pred_vert_vis, 'f': smpl.faces}


                opt_vert_vis = spin_vertices[b]
                opt_vert_vis *=pred_camera_vis[b,0]
                opt_vert_vis[:,0] += pred_camera_vis[b,1]        #no need +1 (or  112). Rendernig has this offset already
                opt_vert_vis[:,1] += pred_camera_vis[b,2]        #no need +1 (or  112). Rendernig has this offset already
                opt_vert_vis*=112
                opt_meshes = {'ver': opt_vert_vis, 'f': smpl.faces}

                # glViewer.setMeshData([pred_meshes, opt_meshes], bComputeNormal= True)
                # glViewer.setMeshData([pred_meshes], bComputeNormal= True)
                # glViewer.setMeshData([pred_meshes, opt_meshes], bComputeNormal= True)

                if bVisSpinOriginal:
                    glViewer.setMeshData([opt_meshes], bComputeNormal= True)
                else:
                    glViewer.setMeshData([pred_meshes], bComputeNormal= True)

                glViewer.setBackgroundTexture(croppedImg)
                glViewer.setWindowSize(croppedImg.shape[1]*3, croppedImg.shape[0]*3)
                glViewer.SetOrthoCamera(True)

                # glViewer.show(0)
                # continue


                if True:#g_bRenderFile:   #Save to File
                    glViewer.SetNearPlane(500)
                    glViewer.setWindowSize(croppedImg.shape[1]*3, croppedImg.shape[0]*3)
                    glViewer.show_SMPL_cameraView(True)


                    overlaidImageFile = os.path.join(overlaidImageFolder,outputFileName)  
                    rendered_img_source = g_renderDir + '/scene_00000000.jpg'
                    cmd = "cp -v {0} {1}".format(rendered_img_source, overlaidImageFile)
                    subprocess.call(cmd,shell=True)


                    glViewer.SetNearPlane(50)
                    glViewer.setWindowSize(croppedImg.shape[1]*3, croppedImg.shape[0]*3)
                    glViewer.show_SMPL_sideView(True)

                    sideImageFileName = os.path.join(sideImageFolder,outputFileName)  
                    rendered_img_source = g_renderDir + '/scene_00000000.jpg'
                    cmd = "cp -v {0} {1}".format(rendered_img_source, sideImageFileName)
                    subprocess.call(cmd,shell=True)

                else:
                    glViewer.SetNearPlane(50)
                    glViewer.setWindowSize(croppedImg.shape[1]*3, croppedImg.shape[0]*3)
                    glViewer.show_SMPL_sideView(True)

                    # glViewer.show(0)

                
                #Merge images
                if True:
                    # croppedImg_highRes  + overlaidImageFile + sideImageFileName
                    overlaidImg = cv2.imread(overlaidImageFile)          
                    sideImg = cv2.imread(sideImageFileName)

                    overlaidImg_vis = cv2.resize(overlaidImg, (croppedImg_highRes.shape[1], croppedImg_highRes.shape[0]))
                    sideImg_vis = cv2.resize(sideImg, (croppedImg_highRes.shape[1], croppedImg_highRes.shape[0]))
                    mergedImg = np.concatenate( (croppedImg_highRes, overlaidImg_vis, sideImg_vis), axis=1)

                    # viewer2D.ImShow(overlaidImg)
                    mergedImgFileName = '{0}/{1}'.format(mergedImageFolder,outputFileName )
                    cv2.imwrite(mergedImgFileName, mergedImg)
                break