# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: torchNew
#     language: python
#     name: torchnew
# ---

# +
import numpy as np
import pandas as pd
import os 
import torch
import random
import pickle
import collections, os, io
from PIL import Image
import torch
from torchvision.transforms import ToTensor, Resize
from torch.utils.data import Dataset
import pickle

import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
import cv2
from matplotlib import cm
import matplotlib
from itertools import combinations

# +
root_dir = '/projects/katefgroup/datasets/clevr_vqa/raw/npys/multi_obj_480_at.txt'

if root_dir.endswith("txt"):
    data  = []

    with open(root_dir) as f:
        lines = f.readlines()

        for line in lines:
            data.append(line.split()[0])

    all_files = [os.path.join(os.path.dirname(root_dir),f) for f in data if f.endswith('.p')]
else:
    all_files = [os.path.join(root_dir,f) for f in os.listdir(root_dir) if f.endswith('.p')]

all_files.sort(); 

print('Initialised.....',len(all_files),' files...')

# -

all_files[0]

data = pickle.load(open(all_files[0], "rb"))

data.keys()

data['rgb_camXs_raw'].shape

data['color_list']

hyp_B = 1
hyp_N = 7
hyp_S = 1

import utils_disco

__p = lambda x: utils_disco.pack_seqdim(x, hyp_B)
__u = lambda x: utils_disco.unpack_seqdim(x, hyp_B)
__pb = lambda x: utils_disco.pack_boxdim(x, hyp_N)
__ub = lambda x: utils_disco.unpack_boxdim(x, hyp_N)

# +
tids = torch.from_numpy(np.reshape(np.arange(hyp_B*hyp_N),[hyp_B,hyp_N]))
pix_T_cams = torch.from_numpy(data["pix_T_cams_raw"][0]).reshape(hyp_B, hyp_S, 4, 4).cuda()

camRs_T_origin = data['camR_T_origin_raw'][0]
camRs_T_origin = torch.from_numpy(camRs_T_origin).reshape(hyp_B, hyp_S, 4, 4).cuda()

origin_T_camRs = __u(utils_disco.safe_inverse(__p(camRs_T_origin)))

origin_T_camXs = torch.from_numpy(data['origin_T_camXs_raw'][0]).reshape(hyp_B, hyp_S, 4, 4).cuda()
camX0_T_camXs = utils_disco.get_camM_T_camXs(origin_T_camXs, ind=0)
camRs_T_camXs = __u(torch.matmul(utils_disco.safe_inverse(__p(origin_T_camRs)), __p(origin_T_camXs))) 

camXs_T_camRs = __u(utils_disco.safe_inverse(__p(camRs_T_camXs)))
camX0_T_camRs = camXs_T_camRs[:,0]

camR_T_camX0  = utils_disco.safe_inverse(camX0_T_camRs)

rgb_camXs = data["rgb_camXs_raw"][:,:,:,:3]
rgb_camX0 = torch.from_numpy(rgb_camXs[0]).permute(2,0,1).reshape(hyp_B, 3, 320, 480) #torch.from_numpy(np.fliplr(rgb_camXs[0,0])).reshape(1, 256, 256, 3).permute(0,3,1,2)
# -

data['bbox_camR'].shape

# +
Z, Y, X = 144, 144, 144
Z2, Y2, X2 = int(Z/2), int(Y/2), int(X/2)
Z4, Y4, X4 = int(Z/4), int(Y/4), int(X/4)

import utils_vox

gt_boxesR_corners = torch.from_numpy(data['bbox_camR']).reshape(hyp_B, hyp_N, 8, 3).cuda()

gt_boxesRMem_corners = __ub(utils_vox.Ref2Mem(__pb(gt_boxesR_corners),Z2,Y2,X2))
gt_boxesRMem_end = utils_disco.get_ends_of_corner(gt_boxesRMem_corners) 

gt_boxesRMem_theta = utils_disco.transform_corners_to_boxes(gt_boxesRMem_corners)
gt_boxesRUnp_corners = __ub(utils_vox.Ref2Mem(__pb(gt_boxesR_corners),Z,Y,X))

gt_boxesRUnp_end = utils_disco.get_ends_of_corner(gt_boxesRUnp_corners)

gt_boxesX0_corners = __ub(utils_disco.apply_4x4(camX0_T_camRs, __pb(gt_boxesR_corners)))

gt_boxesXs_corners = __u(__ub(utils_disco.apply_4x4(__p(camXs_T_camRs), __p(__pb(gt_boxesR_corners).unsqueeze(1).repeat(1,hyp_S,1,1)) )))
gt_boxesXs_end = __u(utils_disco.get_ends_of_corner(__p(gt_boxesXs_corners)))

gt_boxesX0Mem_corners = __ub(utils_vox.Ref2Mem(__pb(gt_boxesX0_corners),Z2,Y2,X2))
gt_boxesX0Mem_theta = utils_disco.transform_corners_to_boxes(gt_boxesX0Mem_corners)
gt_boxesX0Mem_end = utils_disco.get_ends_of_corner(gt_boxesX0Mem_corners)
gt_boxesX0_end = utils_disco.get_ends_of_corner(gt_boxesX0_corners)  

gt_cornersX0_pix = __ub(utils_disco.apply_pix_T_cam(__p(pix_T_cams), __pb(gt_boxesX0_corners)))


# +
def draw_corners_on_image(rgb, corners_cam, scores, tids, pix_T_cam,info_text=None):
    # first we need to get rid of invalid gt boxes
    # gt_boxes = trim_gt_boxes(gt_boxes)
    B, C, H, W = list(rgb.shape)
    assert(C==3)
    B2, N, D, E = list(corners_cam.shape)
    assert(B2==B)
    assert(D==8) # 8 corners
    assert(E==3) # 3D

    #rgb = back2color(rgb)
    corners_cam_ = torch.reshape(corners_cam, [B, N*8, 3])
    corners_pix_ = utils_disco.apply_pix_T_cam(pix_T_cam, corners_cam_)
    corners_pix = torch.reshape(corners_pix_, [B, N, 8, 2])
    out = draw_boxes_on_image_py(rgb[0].cpu().numpy(),
                                      corners_pix[0].cpu().numpy(),
                                      scores[0].cpu().numpy(),
                                      tids[0].cpu().numpy(),info_text)
    out = torch.from_numpy(out).type(torch.ByteTensor).permute(2, 0, 1)
    out = torch.unsqueeze(out, dim=0)
    #out = preprocess_color(out)
    out = torch.reshape(out, [1, C, H, W])
    return out, corners_pix

def draw_boxes_on_image_py(rgb, corners_pix, scores, tids,info_text=None, boxes=None, thickness=1,text=False):
    # all inputs are numpy tensors
    # rgb is H x W x 3
    # corners_pix is N x 8 x 2, in xy order
    # scores is N
    # tids is N
    # boxes is N x 9 < this is only here to print some rotation info
    # pix_T_cam is 4 x 4
    rgb = np.transpose(rgb, [1, 2, 0]) # put channels last
    

    rgb = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
   
  

    H, W, C = rgb.shape
    assert(C==3)
    N, D, E = corners_pix.shape
    assert(D==8)
    assert(E==2)

    if boxes is not None:
        rx = boxes[:,6]
        ry = boxes[:,7]
        rz = boxes[:,8]
    else:
        rx = 0
        ry = 0
        rz = 0

    color_map = matplotlib.cm.get_cmap('tab20')
    color_map = color_map.colors

    # draw
    for ind, corners in enumerate(corners_pix):
        # corners is 8 x 2
        # st()
        if not np.isclose(scores[ind], 0.0):
            # print 'score = %.2f' % scores[ind]
            color_id = tids[ind] % 20
            color = color_map[2]
            color_text = color_map[2]

            # st()

            color = np.array(color)*255.0
            # print 'tid = %d; score = %.3f' % (tids[ind], scores[ind])
            if info_text is not None:
                text_to_put = info_text[ind]
                cv2.putText(rgb,
                            text_to_put, 
                            (np.min(corners[:,0]), np.min(corners[:,1])),
                            cv2.FONT_HERSHEY_SIMPLEX,
                            0.5, # font size
                            color_text,
                            2) # font weight

            for c in corners:

                # rgb[pt1[0], pt1[1], :] = 255
                # rgb[pt2[0], pt2[1], :] = 255
                # rgb[np.clip(int(c[0]), 0, W), int(c[1]), :] = 255

                c0 = np.clip(int(c[0]), 0,  W-1)
                c1 = np.clip(int(c[1]), 0,  H-1)
                rgb[c1, c0, :] = 255

            # we want to distinguish between in-plane edges and out-of-plane ones
            # so let's recall how the corners are ordered:
            xs = np.array([-1/2., -1/2., -1/2., -1/2., 1/2., 1/2., 1/2., 1/2.])
            ys = np.array([-1/2., -1/2., 1/2., 1/2., -1/2., -1/2., 1/2., 1/2.])
            zs = np.array([-1/2., 1/2., -1/2., 1/2., -1/2., 1/2., -1/2., 1/2.])
            xs = np.reshape(xs, [8, 1])
            ys = np.reshape(ys, [8, 1])
            zs = np.reshape(zs, [8, 1])
            offsets = np.concatenate([xs, ys, zs], axis=1)

            corner_inds = list(range(8))
            combos = list(combinations(corner_inds, 2))

            for combo in combos:
                pt1 = offsets[combo[0]]
                pt2 = offsets[combo[1]]
                # draw this if it is an in-plane edge
                eqs = pt1==pt2
                if np.sum(eqs)==2:
                    i, j = combo
                    pt1 = (corners[i, 0], corners[i, 1])
                    pt2 = (corners[j, 0], corners[j, 1])
                    print(W, H, pt1, pt2)
                    retval, pt1, pt2 = cv2.clipLine((0, 0, W, H), int(pt1), int(pt2))
                    if retval:
                        cv2.line(rgb, pt1, pt2, color, thickness, cv2.LINE_AA)

                    # rgb[pt1[0], pt1[1], :] = 255
                    # rgb[pt2[0], pt2[1], :] = 255
    rgb = cv2.cvtColor(rgb.astype(np.uint8), cv2.COLOR_BGR2RGB)
    # utils_basic.print_stats_py('rgb_uint8', rgb)
    # imageio.imwrite('boxes_rgb.png', rgb)
    return rgb

def summ_box_by_corners(rgbR, corners, scores, tids, pix_T_cam, only_return=False):
    # rgb is B x H x W x C
    # corners is B x N x 8 x 3 
    # scores is B x N
    # tids is B x N
    # pix_T_cam is B x 4 x 4
    # st()
    B, C, H, W = list(rgbR.shape)
    boxes_vis = draw_corners_on_image(rgbR,corners,scores,tids,pix_T_cam,None)
    return boxes_vis
# -





# +
boxes_vis_q, corners_pix_q = summ_box_by_corners(rgb_camX0, gt_boxesX0_corners, torch.from_numpy(np.ones([1,hyp_N])), tids, __p(pix_T_cams))


boxes_q = torch.zeros([hyp_N, 4])

for n in range(hyp_N):
    boxes_q[n][0] = torch.min(corners_pix_q[0, n, :, 0]) 
    boxes_q[n][1] = torch.min(corners_pix_q[0, n, :, 1]) 
    boxes_q[n][2] = torch.max(corners_pix_q[0, n, :, 0]) 
    boxes_q[n][3] = torch.max(corners_pix_q[0, n, :, 1])
# -





StarcoderdataJupyterScriptsDedupFiltered