text:
import math
import numpy as np
from PIL import Image
import cv2
import os
from glob import glob
import torch
import torchvision
import neural_renderer as nr

EPS = 1e-7

def get_renderer(world_ori=[0,0,1], image_size=128, fov=30, renderer_min_depth=0.1, renderer_max_depth=10, fill_back=True, device='cuda:0'):
    #### camera intrinsics
    #             (u)   (x)
    #    d * K^-1 (v) = (y)
    #             (1)   (z)

    ## renderer for visualization
    R = [[[1.,0.,0.],
          [0.,1.,0.],
          [0.,0.,1.]]]
    R = torch.FloatTensor(R).to(device)
    t = torch.FloatTensor(world_ori).to(device)
    fx = (image_size)/2/(math.tan(fov/2 *math.pi/180))
    fy = (image_size)/2/(math.tan(fov/2 *math.pi/180))
    cx = (image_size)/2
    cy = (image_size)/2
    K = [[fx, 0., cx],
         [0., fy, cy],
         [0., 0., 1.]]
    K = torch.FloatTensor(K).to(device)
    inv_K = torch.inverse(K).unsqueeze(0)
    K = K.unsqueeze(0)
    renderer = nr.Renderer(camera_mode='projection',
                            light_intensity_ambient=1.0,
                            light_intensity_directional=0.,
                            K=K, R=R, t=t,
                            near=renderer_min_depth, far=renderer_max_depth,
                            image_size=image_size, orig_size=image_size,
                            fill_back=fill_back,
                            background_color=[1.,1.,1.])
    return renderer

def get_grid(H, W, normalize=True):
    if normalize:
        h_range = torch.linspace(-1,1,H)
        w_range = torch.linspace(-1,1,W)
    else:
        h_range = torch.arange(0,H)
        w_range = torch.arange(0,W)
    grid = torch.stack(torch.meshgrid([h_range, w_range]), -1).flip(2).float() # flip h,w to x,y
    return grid

def get_sphere_vtx(n_elev, n_azim):
    elevs = ((torch.arange(n_elev).view(n_elev, 1) +0.5) /n_elev *2 -1) *np.pi/2  # -pi/2~pi/2
    azims = ((torch.arange(n_azim).view(1, n_azim) +0.5) /n_azim *2 -1) *np.pi  # -pi~pi
    xs = elevs.cos() * azims.cos()
    ys = elevs.repeat(1, n_azim).sin()
    zs = elevs.cos() * azims.sin()
    vtx = torch.stack([xs, ys, zs], 2) # ExAx3
    return vtx

def get_sor_vtx(sor_curve, T):
    b, h, _ = sor_curve.shape
    rs, hs = sor_curve.unbind(2)  # BxH
    y = hs.view(b,h,1).repeat(1,1,T)  # BxHxT
    thetas = torch.linspace(-math.pi, math.pi, T+1)[:-1].to(sor_curve.device)  # T
    x = rs.unsqueeze(2) * thetas.cos().view(1,1,T)  # BxHxT
    z = rs.unsqueeze(2) * thetas.sin().view(1,1,T)  # BxHxT
    sor_vtx = torch.stack([x, y, z], 3)  # BxHxTx3
    return sor_vtx

def get_sor_full_face_idx(h, w):
    idx_map = torch.arange(h*w).reshape(h,w)  # HxW
    idx_map = torch.cat([idx_map, idx_map[:,:1]], 1)  # Hx(W+1), connect last column to first
    faces1 = torch.stack([idx_map[:h-1,:w], idx_map[1:,:w], idx_map[:h-1,1:w+1]], -1)  # (H-1)xWx3
    faces2 = torch.stack([idx_map[1:,1:w+1], idx_map[:h-1,1:w+1], idx_map[1:,:w]], -1)  # (H-1)xWx3
    return torch.stack([faces1, faces2], 0).int()  # 2x(H-1)xWx3

def get_sor_front_face_idx(h, w):
    sor_full_face_idx = get_sor_full_face_idx(h, w)  # 2x(H-1)x(W//2)x3
    return sor_full_face_idx[:,:,:w//2,:]

def get_sor_back_face_idx(h, w):
    sor_full_face_idx = get_sor_full_face_idx(h, w)  # 2x(H-1)x(W//2)x3
    return sor_full_face_idx[:,:,w//2:,:]

def get_tex_uv_grid(ts, h, w):
    uv_grid = get_grid(h, w, normalize=True)  # -1~1, HxWx(x,y)
    ab_grid = get_grid(ts, ts, normalize=False) / (ts-1)  # 0~1, txtx(x,y)
    ab_grid_uv_offsets = ab_grid * torch.FloatTensor([2/(w-1), 2/(h-1)]).view(1,1,2)

    tex_uv_grid1 = uv_grid[:-1,:-1,:].view(h-1, w-1, 1, 1, 2) + ab_grid_uv_offsets.view(1, 1, ts, ts, 2)  # (H-1)x(W-1)xtxtx2
    tex_uv_grid2 = uv_grid[1:,1:,:].view(h-1, w-1, 1, 1, 2) - ab_grid_uv_offsets.view(1, 1, ts, ts, 2)  # (H-1)x(W-1)xtxtx2
    tex_uv_grid = torch.stack([tex_uv_grid1, tex_uv_grid2], 0)  # 2x(H-1)x(W-1)xtxtx2
    return tex_uv_grid

def get_sor_vtx_normal(sor_vtx):
    sor_vtx = torch.nn.functional.pad(sor_vtx.permute(0,3,1,2), (1,1,0,0), mode='circular').permute(0,2,3,1)  # BxHx(T+2)x3

    tu = sor_vtx[:,1:-1,2:] - sor_vtx[:,1:-1,:-2]
    tv = sor_vtx[:,2:,1:-1] - sor_vtx[:,:-2,1:-1]
    normal = tu.cross(tv, dim=3)  # Bx(H-2)xTx3
    normal = torch.nn.functional.pad(normal.permute(0,3,1,2), (0,0,1,1), mode='replicate').permute(0,2,3,1)  # BxHxTx3
    normal = normal / (((normal**2).sum(3, keepdim=True))**0.5 + EPS)
    return normal  # BxHxTx3

def get_sor_quad_center_vtx(sor_vtx):
    ## shift to quad center for shading
    sor_vtx = torch.cat([sor_vtx, sor_vtx[:,:,:1]], 2)  # Hx(T+1), connect last column to first
    sor_quad_center_vtx = torch.nn.functional.avg_pool2d(sor_vtx.permute(0,3,1,2), kernel_size=2, stride=1, padding=0).permute(0,2,3,1)
    return sor_quad_center_vtx  # Bx(H-1)xTx3

def get_sor_quad_center_normal(sor_vtx):
    ## shift to quad center for shading
    sor_vtx = torch.cat([sor_vtx, sor_vtx[:,:,:1]], 2)  # Hx(T+1), connect last column to first

    tu = sor_vtx[:,:-1,1:] - sor_vtx[:,1:,:-1]
    tv = sor_vtx[:,1:,1:] - sor_vtx[:,:-1,:-1]
    normal = tu.cross(tv, dim=3)  # Bx(H-1)xTx3
    normal = normal / (((normal**2).sum(3, keepdim=True))**0.5 + EPS)
    return normal  # Bx(H-1)xTx3

def get_rotation_matrix(tx, ty, tz):
    m_x = torch.zeros((len(tx), 3, 3)).to(tx.device)
    m_y = torch.zeros((len(tx), 3, 3)).to(tx.device)
    m_z = torch.zeros((len(tx), 3, 3)).to(tx.device)

    m_x[:, 1, 1], m_x[:, 1, 2] = tx.cos(), -tx.sin()
    m_x[:, 2, 1], m_x[:, 2, 2] = tx.sin(), tx.cos()
    m_x[:, 0, 0] = 1

    m_y[:, 0, 0], m_y[:, 0, 2] = ty.cos(), ty.sin()
    m_y[:, 2, 0], m_y[:, 2, 2] = -ty.sin(), ty.cos()
    m_y[:, 1, 1] = 1

    m_z[:, 0, 0], m_z[:, 0, 1] = tz.cos(), -tz.sin()
    m_z[:, 1, 0], m_z[:, 1, 1] = tz.sin(), tz.cos()
    m_z[:, 2, 2] = 1
    return torch.matmul(m_z, torch.matmul(m_y, m_x))

def rotate_pts(pts, rotmat):
    return pts.matmul(rotmat.transpose(2,1))

def transform_sor(sor_vtx, rxyz=None, txy=None):
    if rxyz is not None:
        rx, ry, rz = rxyz.unbind(1)
        rotmat = get_rotation_matrix(rx, ry, rz).to(sor_vtx.device)
        sor_vtx = rotate_pts(sor_vtx, rotmat)  # BxNx3
    if txy is not None:
        tz = torch.zeros(len(txy), 1).to(txy.device)
        txyz = torch.cat([txy, tz], 1)
        sor_vtx = sor_vtx + txyz.unsqueeze(1)  # BxNx3
    return sor_vtx

def render_sor(renderer, sor_vtx, sor_faces, tex_im, tx_size=4, dim_inside=False):
    b, H, T, _ = sor_vtx.shape
    tex_uv_grid = get_tex_uv_grid(tx_size, H, T+1).to(sor_vtx.device)  # Bx2x(H-1)x(W-1)xtxtx2

    tx_cube = torch.nn.functional.grid_sample(tex_im, tex_uv_grid.view(1,-1,tx_size*tx_size,2).repeat(b,1,1,1), mode='bilinear', padding_mode="reflection", align_corners=False)  # Bx3xFxT^2
    tx_cube = tx_cube.permute(0,2,3,1).view(b,-1,1,tx_size,tx_size,3).repeat(1,1,tx_size,1,1,1)  # BxFxtxtxtx3

    sor_vtx = sor_vtx.view(b,-1,3)
    sor_faces = sor_faces.view(b,-1,3)
    if dim_inside:
        fill_back = renderer.fill_back
        renderer.fill_back = False
        sor_faces = torch.cat([sor_faces, sor_faces.flip(2)], 1)
        tx_cube = torch.cat([tx_cube, tx_cube*0.5], 1)
        im_rendered = renderer.render_rgb(sor_vtx, sor_faces, tx_cube)
        renderer.fill_back = fill_back
    else:
        im_rendered = renderer.render_rgb(sor_vtx, sor_faces, tx_cube)
    return im_rendered

def save_results(root, res, name=''):
    b = res['batch_size']
    keys = res.keys()
    os.makedirs(root, exist_ok=True)
    for i in range(b):
        idx = len(os.listdir(root)) +1
        fname = '%04d' %idx + name
        out_folder = os.path.join(root, fname)
        os.makedirs(out_folder)
        for key in keys:
            fpath = os.path.join(out_folder, fname+'_'+key)
            if key == 'batch_size':
                pass
            elif res[key].dim() == 4:
                im = np.uint8(res[key][i].permute(1,2,0).flip(2).numpy() *255)
                cv2.imwrite(fpath+'.png', im)
            else:
                np.savetxt(fpath+'.txt', res[key][i].numpy(), fmt='%.6f', delimiter=', ')

GAMMA = 2.2
def HDR2LDR(img):
    return img.clamp(min=EPS) **(1/GAMMA)

def LDR2HDR(img):
    return img.clamp(min=EPS) **GAMMA

def envmap_phong_shading(point3d, albedo, spec_albedo, normal, cam_loc, ambient, env_map, spec_alpha):
    b, c, tex_h, tex_w = albedo.shape
    _, h, w, _ = point3d.shape

    view_dirs = cam_loc.view(1,1,1,3) - point3d
    view_dirs = torch.nn.functional.normalize(view_dirs, p=2, dim=3, eps=EPS)

    _, n_elev, n_azim = env_map.shape
    l_dirs = get_sphere_vtx(n_elev, n_azim).unsqueeze(0).to(point3d.device)  # BxExAx3
    l_elevs = ((torch.arange(n_elev) +0.5) /n_elev *2 -1) *np.pi/2
    l_norm = l_elevs.cos().view(1, n_elev, 1).to(point3d.device) / n_elev / n_azim  # 1xEx1
    l_ints = env_map * l_norm *50

    cos_theta = ((-normal.unsqueeze(1)) * l_dirs.view(1,n_elev*n_azim,1,1,3)).sum(4, keepdim=True)
    diffuse = l_ints.view(b,n_elev*n_azim,1,1,-1) *cos_theta.clamp(0,1)  # BxLxHxWx1
    diffuse = torch.nn.functional.interpolate(diffuse.view(b*n_elev*n_azim,h,w,1).permute(0,3,1,2), (tex_h, tex_w), mode='bilinear', align_corners=False).view(b,n_elev*n_azim,1,tex_h,tex_w)

    reflect_dirs = -l_dirs.view(1,n_elev*n_azim,1,1,3) + 2*cos_theta*(-normal.unsqueeze(1))
    specular = (view_dirs.unsqueeze(1) * reflect_dirs).sum(4,keepdim=True).clamp(min=0) * (cos_theta>0)
    specular = (spec_alpha.view(b,1,1,1,-1)+1)/2/math.pi *l_ints.view(b,n_elev*n_azim,1,1,-1) * specular.clamp(min=EPS).pow(spec_alpha.view(b,1,1,1,-1))  # BxLxHxWx1
    specular = torch.nn.functional.interpolate(specular.view(b*n_elev*n_azim,h,w,1).permute(0,3,1,2), (tex_h, tex_w), mode='bilinear', align_corners=False).view(b,n_elev*n_azim,1,tex_h,tex_w)

    colors = (ambient.view(b,-1,1,1) + diffuse.sum(1)) *albedo + specular.sum(1) *spec_albedo
    return colors, diffuse, specular

def sg_to_env_map(sg_lights, n_elev=8, n_azim=16):
    b, n_sgls, _ = sg_lights.shape

    sgl_dirs = sg_lights[:,:,:3]
    sgl_lams = sg_lights[:,:,3:4]
    sgl_Fs = sg_lights[:,:,4:5]

    l_dirs = get_sphere_vtx(n_elev, n_azim).unsqueeze(0).to(sg_lights.device)  # BxExAx3
    exps = sgl_lams.view(b,n_sgls,1,1)* ((sgl_dirs.view(b,n_sgls,1,1,3)*l_dirs.view(1,1,n_elev,n_azim,3)).sum(4)-1)  # BxLxExA
    l_ints = sgl_Fs.view(b,n_sgls,1,1) * exps.exp()  # BxLxExA
    env_map = l_ints.sum(1) / n_sgls  # BxExA
    return env_map

def get_random_sor_curve(b, H):
    t1_bottom = torch.rand(b) *math.pi -math.pi  # -pi~0
    t1_top = torch.rand(b) *1.5*math.pi + 0.5*math.pi  # pi/2~pi*2
    amp1 = torch.rand(b) *0.3  # 0~0.3
    t2_bottom = torch.rand(b) *2*math.pi  # 0~pi*2
    t2_top = t2_bottom + torch.rand(b) *1.5*math.pi + 0.5*math.pi  # bottom + pi/2~pi*2
    amp2 = torch.rand(b) *0.1  # 0~0.1
    r0 = torch.rand(b) *0.2 + 0.1  # 0.02~0.3

    ts = torch.linspace(0,1,H)
    t1s = (1-ts.view(1,H)) * t1_bottom.view(b,1) +  ts.view(1,H) * t1_top.view(b,1)
    sin1 = amp1.view(b,1) * (t1s.sin() +1)
    t2s = (1-ts.view(1,H)) * t2_bottom.view(b,1) +  ts.view(1,H) * t2_top.view(b,1)
    sin2 = amp2.view(b,1) * (t2s.sin() +1)
    # r_col = sin1 + sin2 + r0.view(b,1)  # 0.1~1
    # r_col = r_col / r_col.max(1)[0].view(b,1) *0.75  # normalize to 0.75

    ## TODO:straight axis testing
    r_col = (torch.rand(b).clamp(0.4,1)*0.75).repeat(1,32)

    h_scale = torch.zeros(b) + 0.75
    h_col = torch.linspace(-1,1,H).view(1,H) *h_scale.view(b,1)

    sor_curve = torch.stack([r_col, h_col], 2)  # BxHx(r,h)
    return sor_curve

def get_random_pitch(b):
    return torch.rand(b) *20 /180*math.pi  # 0~20

def get_random_ambient(b):
    return torch.rand(b) *0.4 +0.1  # 0.1~0.5

def get_random_spec_albedo(b):
    return torch.rand(b) *0.9 +0.1  # 0.3~1

def get_random_spec_alpha(b):
    return (torch.rand(b) *13+1)**2  # 1~196

def get_random_sg_lights(b):
    n_sgls = 3
    sgl_dirs = torch.rand(b, n_sgls, 3) *2-1

    sgl_dirs[:,:,1] = -sgl_dirs[:,:,1].abs()  # upper only
    sgl_dirs[:,:,2] = -sgl_dirs[:,:,2].abs()  # front only

    sgl_dirs = torch.nn.functional.normalize(sgl_dirs, p=2, dim=2, eps=EPS)
    sgl_lams = torch.rand(b, n_sgls) *30+10
    sgl_Fs = (torch.rand(b, n_sgls) *0.3+0.1) *sgl_lams**0.5
    sg_lights = torch.cat([sgl_dirs, sgl_lams.unsqueeze(2), sgl_Fs.unsqueeze(2)], 2)
    return sg_lights

def get_random_env_map_ambient(b):
    return torch.rand(b) *0.03  # 0~0.1

def random_crop(im_pil, size, scale=(0.08, 1.0), ratio=(0.75, 1.3333333333333333)):
    tfs_crop = torchvision.transforms.RandomResizedCrop(size, scale=scale, ratio=ratio)
    return tfs_crop(im_pil)

def random_color_jitter(im_pil, brightness=0, contrast=0, saturation=0, hue=0):
    tfs_jitter = torchvision.transforms.ColorJitter(brightness, contrast, saturation, hue)
    return tfs_jitter(im_pil)

def generate(cc0_tex_dir, out_dir):
    H = 32
    T = 96
    image_size = 256
    fov = 10
    envmap_n_elev = 32
    envmap_n_azim = 96
    tex_im_h = 256
    tex_im_w = 768
    num_im_per_tex = 5

    device = 'cuda:0'
    b = 1
    oriz = 5
    cam_loc = torch.FloatTensor([0,0,-oriz]).to(device)
    lim = math.tan(fov/2/180*math.pi) *oriz
    max_depth = oriz + lim
    min_depth = oriz - lim
    renderer = get_renderer(world_ori=[0,0,oriz], image_size=image_size, fov=fov, fill_back=True)
    sor_faces = get_sor_full_face_idx(H, T).repeat(b,1,1,1,1).to(device)  # Bx2x(H-1)xWx3
    tx_size = 8
    tex_uv_grid = get_tex_uv_grid(tx_size, H, T+1).repeat(b,1,1,1,1,1,1).to(device)  # Bx2x(H-1)x(W-1)xtxtx2

    tex_im_list = sorted(glob(os.path.join(cc0_tex_dir, '**/*Color.jpg'), recursive=True))
    num_tex_im = len(tex_im_list)
    for i, tex_fpath in enumerate(tex_im_list):
        print(f'\n\nTex {i}/{num_tex_im}:')

        tex_im_pil = Image.open(tex_fpath).convert('RGB')

        for j in range(num_im_per_tex):
            tex_im_crop_pil = random_crop(tex_im_pil, size=(tex_im_h, tex_im_w//2), scale=(0.04, 1), ratio=(0.75, 1.3333333333333333))
            tex_im_crop_pil = random_color_jitter(tex_im_crop_pil, brightness=(1.0,1.5), contrast=(1.0,2.0), saturation=0., hue=0.5)
            tex_im_crop = torch.FloatTensor(np.array(tex_im_crop_pil) / 255).to(device)
            tex_im_crop = tex_im_crop.permute(2,0,1).unsqueeze(0)
            tex_im_crop = LDR2HDR(tex_im_crop)

            sor_curve = get_random_sor_curve(b,H).to(device) *lim  # BxHx2
            sor_vtx = get_sor_vtx(sor_curve, T)  # BxHxTx3

            ## render rotation pitch rotation
            pitch = get_random_pitch(b).to(device)
            rxyz = torch.stack([pitch, torch.zeros_like(pitch), torch.zeros_like(pitch)], 1)
            posed_sor_vtx = transform_sor(sor_vtx.view(b,-1,3), rxyz, txy=None).view(b,H,T,3)

            depth_rendered = renderer.render_depth(posed_sor_vtx.view(b,-1,3), sor_faces.view(b,-1,3)).clamp(min_depth, max_depth)
            mask_rendered = renderer.render_silhouettes(posed_sor_vtx.view(b,-1,3), sor_faces.view(b,-1,3))

            normal_map = get_sor_quad_center_normal(posed_sor_vtx)  # Bx(H-1)xTx3
            normal_tx_cube = torch.nn.functional.grid_sample(normal_map.permute(0,3,1,2), tex_uv_grid.view(b,-1,tx_size*tx_size,2), mode='bilinear', padding_mode="border", align_corners=False)  # Bx3xFxT^2
            normal_tx_cube = normal_tx_cube / (normal_tx_cube**2).sum(1,keepdim=True)**0.5 /2+0.5
            normal_tx_cube = normal_tx_cube.permute(0,2,3,1).view(b,-1,1,tx_size,tx_size,3).repeat(1,1,tx_size,1,1,1)  # BxFxtxtxtx3
            normal_rendered = renderer.render_rgb(posed_sor_vtx.view(b,-1,3), sor_faces.view(b,-1,3), normal_tx_cube).clamp(0, 1)

            ## sample lighting
            albedo = torch.cat([tex_im_crop, tex_im_crop.flip(3)], 3)
            ambient = get_random_ambient(b).to(device) *0
            spec_albedo = get_random_spec_albedo(b).to(device)
            spec_alpha = get_random_spec_alpha(b).to(device)
            sg_lights = get_random_sg_lights(b).to(device)

            posed_sor_vtx_map = get_sor_quad_center_vtx(posed_sor_vtx)  # Bx(H-1)xTx3
            env_map = sg_to_env_map(sg_lights, n_elev=envmap_n_elev, n_azim=envmap_n_azim)
            env_map_ambient = get_random_env_map_ambient(b).to(device)
            env_map = env_map + env_map_ambient.view(b,1,1)
            colors, diffuse, specular = envmap_phong_shading(posed_sor_vtx_map, albedo, spec_albedo.view(1,1,1,1), normal_map, cam_loc, ambient, env_map, spec_alpha)
            colors = HDR2LDR(colors)
            colors = colors.clamp(0,1)
            albedo = HDR2LDR(albedo)
            specular = specular * spec_albedo.view(1,1,1,1,1)

            im_rendered = render_sor(renderer, posed_sor_vtx, sor_faces, colors, tx_size=tx_size, dim_inside=True).clamp(0, 1)
            albedo_rendered = render_sor(renderer, posed_sor_vtx, sor_faces, albedo, tx_size=tx_size, dim_inside=True).clamp(0, 1)
            diffuse_map = diffuse.sum(1).clamp(0, 1).repeat(1,3,1,1)
            diffuse_rendered = render_sor(renderer, posed_sor_vtx, sor_faces, diffuse_map, tx_size=tx_size).clamp(0, 1)
            specular_map = specular.sum(1).clamp(0, 1).repeat(1,3,1,1)
            specular_rendered = render_sor(renderer, posed_sor_vtx, sor_faces, specular_map, tx_size=tx_size).clamp(0, 1)

            results = {}
            results['batch_size'] = b
            results['sor_curve'] = sor_curve.detach().cpu()
            results['pitch'] = pitch.detach().cpu().unsqueeze(1) /math.pi*180
            results['depth_rendered'] = (depth_rendered.detach().cpu().unsqueeze(1).repeat(1,3,1,1) - min_depth) / (max_depth-min_depth)
            results['mask_rendered'] = mask_rendered.detach().cpu().unsqueeze(1).repeat(1,3,1,1)

            results['im_rendered'] = im_rendered.detach().cpu()
            results['normal_map'] = normal_map.detach().cpu().permute(0,3,1,2) /2+0.5
            results['normal_rendered'] = normal_rendered.detach().cpu()
            results['albedo_map'] = albedo.detach().cpu()
            results['albedo_rendered'] = albedo_rendered.detach().cpu()
            results['texture_map'] = colors.detach().cpu()
            results['diffuse_map'] = diffuse_map.detach().cpu()
            results['diffuse_rendered'] = diffuse_rendered.detach().cpu()
            results['specular_map'] = specular_map.detach().cpu()
            results['specular_rendered'] = specular_rendered.detach().cpu()
            results['sg_lights'] = sg_lights.detach().cpu()
            results['ambient'] = ambient.detach().cpu().unsqueeze(1)
            results['spec_albedo'] = spec_albedo.detach().cpu().unsqueeze(1)
            results['spec_alpha'] = spec_alpha.detach().cpu().unsqueeze(1)
            results['env_map'] = env_map.clamp(0, 1).detach().cpu().repeat(1,3,1,1)
            results['env_map_ambient'] = env_map_ambient.detach().cpu().unsqueeze(1)

            tex_id = '_' + os.path.basename(os.path.dirname(tex_fpath)) + '_%02d' %j
            save_results(out_dir, results, name=tex_id)

if __name__ == '__main__':
    cc0_tex_dir = '../cc0_textures/PhotoTexturePBR'
    out_dir = '../syn_curv_sgl5_tex_straight_20220220/rendering'

    for split in ['train', 'test']:
        print(f'Generating {split} set...')
        generate(os.path.join(cc0_tex_dir, split), os.path.join(out_dir, split))

    os.symlink('test', os.path.join(out_dir, 'val'))


text:
# ----------------------------------------------------------------------------
# Copyright (c) 2017 Massachusetts Institute of Technology (MIT)
# All rights reserved.
#
# Distributed under the terms of the BSD 3-clause license.
#
# The full license is in the LICENSE file, distributed with this software.
# ----------------------------------------------------------------------------
"""Tests for the digital_rf.digital_rf_hdf5 module."""
from __future__ import absolute_import, division, print_function

import datetime
import itertools
import os

import digital_rf
import h5py
import numpy as np
import packaging.version
import pytest

###############################################################################
#  constant fixtures  #########################################################
###############################################################################

@pytest.fixture(scope="session")
def start_datetime():
    return datetime.datetime(2014, 3, 9, 12, 30, 30, 0, None)

@pytest.fixture(scope="session")
def start_global_index(samples_per_second, start_datetime):
    return digital_rf.util.time_to_sample(start_datetime, samples_per_second)

###############################################################################
#  parametrized fixtures  #####################################################
###############################################################################

_dtype_strs = [
    "s8",
    "sc8",
    "u8",
    "uc8",
    "s16",
    "sc16",
    "u16",
    "uc16",
    "s32",
    "sc32",
    "u32",
    "uc32",
    "s64",
    "sc64",
    "u64",
    "uc64",
    "f32",
    "fc32",
    "f64",
    "fc64",
]
_num_subchannels = [(1, "x1"), (8, "x8")]
_df_params = list(itertools.product(_dtype_strs, (sc[0] for sc in _num_subchannels)))
_df_ids = list(
    "".join(p)
    for p in itertools.product(_dtype_strs, (sc[1] for sc in _num_subchannels))
)

@pytest.fixture(scope="session", params=_df_params, ids=_df_ids)
def data_params(request):
    dtypes = dict(
        s8=np.dtype("i1"),
        sc8=np.dtype([("r", "i1"), ("i", "i1")]),
        u8=np.dtype("u1"),
        uc8=np.dtype([("r", "u1"), ("i", "u1")]),
        s16=np.dtype("i2"),
        sc16=np.dtype([("r", "i2"), ("i", "i2")]),
        u16=np.dtype("u2"),
        uc16=np.dtype([("r", "u2"), ("i", "u2")]),
        s32=np.dtype("i4"),
        sc32=np.dtype([("r", "i4"), ("i", "i4")]),
        u32=np.dtype("u4"),
        uc32=np.dtype([("r", "u4"), ("i", "u4")]),
        s64=np.dtype("i8"),
        sc64=np.dtype([("r", "i8"), ("i", "i8")]),
        u64=np.dtype("u8"),
        uc64=np.dtype([("r", "u8"), ("i", "u8")]),
        f32=np.dtype("f4"),
        fc32=np.dtype("c8"),
        f64=np.dtype("f8"),
        fc64=np.dtype("c16"),
    )
    dtype_str = request.param[0]
    dtype = dtypes[dtype_str]
    is_complex = dtype.names is not None or np.issubdtype(dtype, np.complexfloating)
    num_subchannels = request.param[1]
    return dict(
        dtype_str=dtype_str,
        dtype=dtype,
        is_complex=is_complex,
        num_subchannels=num_subchannels,
    )

@pytest.fixture(scope="session", params=[True, False], ids=["continuous", "gapped"])
def file_params(request):
    return dict(is_continuous=request.param)

@pytest.fixture(
    scope="session",
    params=[
        # compression_level, checksum
        (0, False),
        (9, True),
    ],
    ids=["nozip,nochecksum", "zip,checksum"],
)
def hdf_filter_params(request):
    return dict(compression_level=request.param[0], checksum=request.param[1])

@pytest.fixture(
    scope="session",
    params=[
        # srnum, srden, sdcsec, fcms
        (200, 3, 10, 1000)
    ],
    ids=["200/3hz,10s,1000ms"],
)
def sample_params(request):
    return dict(
        sample_rate_numerator=request.param[0],
        sample_rate_denominator=request.param[1],
        subdir_cadence_secs=request.param[2],
        file_cadence_millisecs=request.param[3],
    )

###############################################################################
#  remaining fixtures  ########################################################
###############################################################################

@pytest.fixture(scope="session")
def param_dict(
    data_params, file_params, hdf_filter_params, sample_params, start_global_index
):
    params = dict(start_global_index=start_global_index)
    params.update(data_params)
    params.update(file_params)
    params.update(hdf_filter_params)
    params.update(sample_params)
    return params

@pytest.fixture(scope="session")
def param_str(param_dict):
    s = (
        "{dtype_str}_x{num_subchannels}_n{sample_rate_numerator}"
        "_d{sample_rate_denominator}_{subdir_cadence_secs}s"
        "_{file_cadence_millisecs}ms_c{is_continuous:b}"
    ).format(**param_dict)
    return s

@pytest.fixture(scope="session")
def dtype_str(data_params):
    return data_params["dtype_str"]

@pytest.fixture(scope="session")
def dtype(data_params):
    return data_params["dtype"]

@pytest.fixture(scope="session")
def is_complex(data_params):
    return data_params["is_complex"]

@pytest.fixture(scope="session")
def num_subchannels(data_params):
    return data_params["num_subchannels"]

@pytest.fixture(scope="session")
def is_continuous(file_params):
    return file_params["is_continuous"]

@pytest.fixture(scope="session")
def compression_level(hdf_filter_params):
    return hdf_filter_params["compression_level"]

@pytest.fixture(scope="session")
def checksum(hdf_filter_params):
    return hdf_filter_params["checksum"]

@pytest.fixture(scope="session")
def sample_rate_numerator(sample_params):
    return sample_params["sample_rate_numerator"]

@pytest.fixture(scope="session")
def sample_rate_denominator(sample_params):
    return sample_params["sample_rate_denominator"]

@pytest.fixture(scope="session")
def subdir_cadence_secs(sample_params):
    return sample_params["subdir_cadence_secs"]

@pytest.fixture(scope="session")
def file_cadence_millisecs(sample_params):
    return sample_params["file_cadence_millisecs"]

@pytest.fixture(scope="session")
def samples_per_second(sample_rate_numerator, sample_rate_denominator):
    return np.longdouble(sample_rate_numerator) / sample_rate_denominator

@pytest.fixture(scope="session")
def end_global_index(
    file_cadence_millisecs,
    sample_rate_numerator,
    sample_rate_denominator,
    samples_per_second,
    start_global_index,
    subdir_cadence_secs,
):
    # want data to span at least two subdirs to test creation + naming
    # also needs to span at least 8 files to accommodate write blocks (below)
    nsamples_subdirs = int(
        (1.5 * subdir_cadence_secs * sample_rate_numerator) // sample_rate_denominator
    )
    nsamples_files = int(
        np.ceil(8 * file_cadence_millisecs * (samples_per_second / 1000))
    )
    nsamples = max(nsamples_subdirs, nsamples_files)
    return start_global_index + nsamples - 1

@pytest.fixture(scope="session")
def bounds(
    end_global_index,
    file_cadence_millisecs,
    is_continuous,
    sample_rate_numerator,
    sample_rate_denominator,
    start_global_index,
):
    if is_continuous:
        srn = sample_rate_numerator
        srd = sample_rate_denominator
        fcms = file_cadence_millisecs

        # bounds are at file boundaries
        # milliseconds of file with start_global_index
        sms = (((start_global_index * srd * 1000) // srn) // fcms) * fcms
        # sample index at start of file (ceil b/c start at first whole sample)
        ss = ((sms * srn) + (srd * 1000) - 1) // (srd * 1000)
        # milliseconds of file after end_global_index
        ems = ((((end_global_index * srd * 1000) // srn) // fcms) + 1) * fcms
        # sample index at start of file (ceil b/c start at first whole sample)
        es = ((ems * srn) + (srd * 1000) - 1) // (srd * 1000)
        return (ss, es - 1)
    else:
        return (start_global_index, end_global_index)

@pytest.fixture(scope="session")
def data_block_slices(
    bounds,
    end_global_index,
    file_cadence_millisecs,
    samples_per_second,
    start_global_index,
):
    # blocks = [(start_sample, stop_sample)]
    blocks = []
    samples_per_file = file_cadence_millisecs * (samples_per_second / 1000)

    # first block stops in middle of second file
    sstart = start_global_index
    sstop = bounds[0] + int(1.5 * samples_per_file)
    blocks.append((sstart, sstop))

    # second block continues where first ended, stops in middle of third file
    sstart = sstop
    sstop = bounds[0] + int(2.5 * samples_per_file)
    blocks.append((sstart, sstop))

    # first gap, no skipped file
    # third block starts in middle of 4th file, ends in middle of 5th
    sstart = bounds[0] + int(3.5 * samples_per_file)
    sstop = bounds[0] + int(4.5 * samples_per_file)
    blocks.append((sstart, sstop))

    # second gap, skipping one file completely
    # fourth block starts in middle of 7th file, ends later in 7th file
    sstart = bounds[0] + int(6.5 * samples_per_file)
    sstop = bounds[0] + int(6.9 * samples_per_file)
    blocks.append((sstart, sstop))

    # fifth and final block continues just after previous write and goes to end
    sstart = sstop + 2
    sstop = end_global_index + 1
    blocks.append((sstart, sstop))

    return blocks

def generate_rf_data(shape, dtype, seed):
    np.random.seed(seed % 2 ** 32)
    nitems = np.product(shape)
    byts = np.random.randint(0, 256, nitems * dtype.itemsize, "u1")
    return byts.view(dtype).reshape(shape)

@pytest.fixture(scope="session")
def data(bounds, data_block_slices, dtype, num_subchannels):
    data_len = bounds[1] - bounds[0] + 1
    if num_subchannels == 1:
        shape = (data_len,)
    else:
        shape = (data_len, num_subchannels)

    # start off with data array containing its appropriate fill value
    if np.issubdtype(dtype, np.inexact):
        fill_value = np.nan * np.ones(1, dtype=dtype)[0]
    elif dtype.names is not None:
        fill_value = np.empty(1, dtype=dtype)[0]
        fill_value["r"] = np.iinfo(dtype["r"]).min
        fill_value["i"] = np.iinfo(dtype["i"]).min
    else:
        fill_value = np.iinfo(dtype).min * np.ones(1, dtype=dtype)[0]
    data = np.full(shape, fill_value, dtype=dtype)

    rdata = generate_rf_data(shape, dtype, 0)
    # fill array data for what we will be writing
    for sstart, sstop in data_block_slices:
        bstart = sstart - bounds[0]
        bend = sstop - bounds[0]
        data[bstart:bend] = rdata[bstart:bend]

    return data

@pytest.fixture(scope="session")
def data_file_list(sample_params):
    # manually generated to match output of writing data given
    # _sample_rate_params
    # get params as a tuple of values sorted in key order
    params = tuple(v for k, v in sorted(sample_params.items()))
    file_lists = {
        (1000, 3, 200, 10): [
            os.path.join("2014-03-09T12-30-30", "rf@1394368230.000.h5"),
            os.path.join("2014-03-09T12-30-30", "rf@1394368231.000.h5"),
            os.path.join("2014-03-09T12-30-30", "rf@1394368232.000.h5"),
            os.path.join("2014-03-09T12-30-30", "rf@1394368233.000.h5"),
            os.path.join("2014-03-09T12-30-30", "rf@1394368234.000.h5"),
            os.path.join("2014-03-09T12-30-30", "rf@1394368236.000.h5"),
            os.path.join("2014-03-09T12-30-30", "rf@1394368237.000.h5"),
            os.path.join("2014-03-09T12-30-30", "rf@1394368238.000.h5"),
            os.path.join("2014-03-09T12-30-30", "rf@1394368239.000.h5"),
            os.path.join("2014-03-09T12-30-40", "rf@1394368240.000.h5"),
            os.path.join("2014-03-09T12-30-40", "rf@1394368241.000.h5"),
            os.path.join("2014-03-09T12-30-40", "rf@1394368242.000.h5"),
            os.path.join("2014-03-09T12-30-40", "rf@1394368243.000.h5"),
            os.path.join("2014-03-09T12-30-40", "rf@1394368244.000.h5"),
        ]
    }
    return file_lists[params]

@pytest.fixture(scope="class")
def datadir(param_str, tmpdir_factory):
    return tmpdir_factory.mktemp("{0}_".format(param_str), numbered=True)

@pytest.fixture(scope="class")
def chdir(param_str, datadir):
    chdir = datadir.mkdir(param_str)
    return chdir

@pytest.fixture(scope="class")
def channel(chdir):
    return str(chdir.basename)

@pytest.fixture(scope="class")
def drf_writer_param_dict(chdir, param_dict):
    return dict(
        directory=str(chdir),
        dtype=param_dict["dtype"],
        subdir_cadence_secs=param_dict["subdir_cadence_secs"],
        file_cadence_millisecs=param_dict["file_cadence_millisecs"],
        start_global_index=param_dict["start_global_index"],
        sample_rate_numerator=param_dict["sample_rate_numerator"],
        sample_rate_denominator=param_dict["sample_rate_denominator"],
        uuid_str=chdir.basename,
        compression_level=param_dict["compression_level"],
        checksum=param_dict["checksum"],
        is_complex=param_dict["is_complex"],
        num_subchannels=param_dict["num_subchannels"],
        is_continuous=param_dict["is_continuous"],
        marching_periods=False,
    )

@pytest.fixture(scope="class")
def drf_writer_factory(drf_writer_param_dict):
    def writer_factory(**kwargs_updates):
        kwargs = drf_writer_param_dict.copy()
        kwargs.update(**kwargs_updates)
        return digital_rf.DigitalRFWriter(**kwargs)

    return writer_factory

@pytest.fixture(scope="class")
def drf_writer(drf_writer_factory):
    with drf_writer_factory() as dwo:
        yield dwo

@pytest.fixture(scope="class")
def drf_reader(chdir):
    with digital_rf.DigitalRFReader(chdir.dirname) as dro:
        yield dro

###############################################################################
#  tests  #####################################################################
###############################################################################

def test_get_unix_time(
    sample_rate_numerator, sample_rate_denominator, start_datetime, start_global_index
):
    global_index = start_global_index + 1
    index_dt = start_datetime + datetime.timedelta(microseconds=15000)
    dt, picoseconds = digital_rf.get_unix_time(
        global_index, sample_rate_numerator, sample_rate_denominator
    )
    assert dt == index_dt
    assert picoseconds == index_dt.microsecond * 1000000

class TestDigitalRFChannel(object):
    """Test writing and reading of a Digital RF channel."""

    @pytest.mark.firstonly(
        "data_params", "file_params", "hdf_filter_params", "sample_params"
    )
    def test_writer_init(
        self, channel, chdir, drf_writer_param_dict, drf_writer_factory
    ):
        """Test writer object's init."""
        channel = "{0}_init_test".format(channel)

        # error when directory doesn't exist
        with pytest.raises(IOError):
            drf_writer_factory(directory=str(chdir.dirpath(channel)))

        chdir = chdir.dirpath().mkdir(channel)

        # invalid dtype
        with pytest.raises(TypeError):
            drf_writer_factory(directory=str(chdir), dtype="c7")
        with pytest.raises(RuntimeError):
            drf_writer_factory(directory=str(chdir), dtype="S8")
        dtype = np.dtype([("notr", "i2"), ("noti", "i2")])
        with pytest.raises(ValueError):
            drf_writer_factory(directory=str(chdir), dtype=dtype)

        # invalid subdir_cadence_secs
        for scs in (0, -10, 1.1):
            with pytest.raises(ValueError):
                drf_writer_factory(directory=str(chdir), subdir_cadence_secs=scs)

        # invalid file_cadence_millisecs
        for fcms in (0, -10, 1.1):
            with pytest.raises(ValueError):
                drf_writer_factory(directory=str(chdir), file_cadence_millisecs=fcms)

        # invalid subdir_cadence_secs and file_cadence_secs together
        # (subdir_cadence_secs*1000 % file_cadence_millisecs) != 0
        with pytest.raises(ValueError):
            drf_writer_factory(
                directory=str(chdir),
                file_cadence_millisecs=1001,
                subdir_cadence_secs=3600,
            )

        # invalid start_global_index
        with pytest.raises(ValueError):
            drf_writer_factory(directory=str(chdir), start_global_index=-1)

        # invalid sample_rate_numerator
        for srn in (0, -10, 1.1):
            with pytest.raises(ValueError):
                drf_writer_factory(directory=str(chdir), sample_rate_numerator=srn)

        # invalid sample_rate_denominator
        for srd in (0, -10, 1.1):
            with pytest.raises(ValueError):
                drf_writer_factory(directory=str(chdir), sample_rate_denominator=srd)

        # invalid UUID
        with pytest.raises(ValueError):
            drf_writer_factory(directory=str(chdir), uuid_str=0)

        # invalid compression_level
        for cl in (-1, 0.5, 10, 100):
            with pytest.raises(ValueError):
                drf_writer_factory(directory=str(chdir), compression_level=cl)

        # invalid num_subchannels
        with pytest.raises(ValueError):
            drf_writer_factory(directory=str(chdir), num_subchannels=0)

        # now create valid writer and check attributes
        p = drf_writer_param_dict
        with drf_writer_factory(directory=str(chdir)) as dwo:
            assert dwo.directory == str(chdir)
            assert dwo.dtype == p["dtype"]
            assert dwo.is_complex == p["is_complex"]
            assert dwo.subdir_cadence_secs == p["subdir_cadence_secs"]
            assert dwo.file_cadence_millisecs == p["file_cadence_millisecs"]
            assert dwo.start_global_index == p["start_global_index"]
            assert dwo.uuid == p["uuid_str"]
            assert dwo.compression_level == p["compression_level"]
            assert dwo.checksum == p["checksum"]
            assert dwo.num_subchannels == p["num_subchannels"]
            assert dwo.is_continuous == p["is_continuous"]

        # test failure when trying to use different writer settings
        different_settings = [
            dict(dtype="c8" if p["dtype"] is not np.dtype("c8") else "f8"),
            dict(
                subdir_cadence_secs=(
                    (
                        (p["subdir_cadence_secs"] * 1000) // p["file_cadence_millisecs"]
                        + 1
                    )
                    * p["file_cadence_millisecs"]
                )
            ),
            dict(file_cadence_millisecs=p["subdir_cadence_secs"] * 1000),
            dict(sample_rate_numerator=p["sample_rate_numerator"] + 1),
            dict(sample_rate_denominator=p["sample_rate_denominator"] + 1),
            dict(num_subchannels=p["num_subchannels"] + 1),
            dict(is_continuous=not p["is_continuous"]),
        ]
        for kwargs in different_settings:
            with pytest.raises(RuntimeError):
                drf_writer_factory(directory=str(chdir), **kwargs)

        # test success when trying to use the same settings
        drf_writer_factory(directory=str(chdir))

    @pytest.mark.firstonly("file_params", "hdf_filter_params", "sample_params")
    def test_writer_write_input_checking(
        self, channel, chdir, dtype, drf_writer_factory, is_complex, num_subchannels
    ):
        """Test input checking of writer object's write method."""
        # by always using two dimensions even for num_subchannels==1,
        # we test both the 2d (here) and 1d (other test fcn) input cases
        shape = (100, num_subchannels)
        chdir = chdir.dirpath().mkdir("{0}_input_tests".format(channel))
        with drf_writer_factory(directory=str(chdir)) as dwo:
            # test interleaved input format (complex writer only)
            if is_complex:
                next_rel_index = dwo.rf_write(
                    np.ones((shape[0], 2 * shape[1]), dtype=dwo.realdtype)
                )

            # test writing data of other types
            # base_type should upcast sucessfully, other_type should fail
            # (one byte integer values will upcast to almost anything else)
            if np.issubdtype(dwo.realdtype, np.unsignedinteger):
                base_type = "u1"
                other_type = "i{0}".format(dwo.realdtype.itemsize)
            else:
                base_type = "i1"
                if np.issubdtype(dwo.realdtype, np.integer):
                    other_type = "u{0}".format(dwo.realdtype.itemsize)
                else:
                    other_type = "f{0}".format(dwo.realdtype.itemsize * 2)
            if is_complex:
                # test upcast from complex input
                # also tests struct dtype input when floating complex
                next_rel_index = dwo.rf_write(
                    np.ones(shape, dtype=np.dtype([("r", base_type), ("i", base_type)]))
                )
                # test upcast from interleaved format
                next_rel_index = dwo.rf_write(
                    np.ones((shape[0], 2 * shape[1]), dtype=base_type)
                )
                # fail writing data of complex but incompatible type
                with pytest.raises(TypeError):
                    dwo.rf_write(np.ones(shape, dtype=other_type))
                # fail writing real data when complex is expected
                with pytest.raises(ValueError):
                    dwo.rf_write(np.ones(shape, dtype=dwo.realdtype))
            else:
                next_rel_index = dwo.rf_write(np.ones(shape, dtype=base_type))
                # fail writing data of real but incompatible type
                with pytest.raises(TypeError):
                    dwo.rf_write(np.ones(shape, dtype=other_type))
                # fail writing complex data when real is expected
                with pytest.raises(TypeError):
                    dwo.rf_write(
                        np.ones(shape, dtype=np.dtype([("r", dtype), ("i", dtype)]))
                    )

            # fail at overwriting
            with pytest.raises(ValueError):
                dwo.rf_write(np.ones(shape, dtype=dtype), next_rel_index - 1)

            # fail at writing wrong shape
            with pytest.raises(ValueError):
                dwo.rf_write(np.ones(shape + (1,), dtype=dtype))

            # fail at writing wrong number of subchannels
            with pytest.raises(ValueError):
                dwo.rf_write(np.ones((shape[0], shape[1] + 1), dtype=dtype))

        # test initializing with interleaved format
        if is_complex:
            chdir = chdir.dirpath().mkdir("{0}_input_tests_interleaved".format(channel))
            if dtype.names is not None:
                realdtype = dtype["r"]
            else:
                realdtype = np.dtype("f{0}".format(dtype.itemsize // 2))
            with drf_writer_factory(
                directory=str(chdir), dtype=realdtype, is_complex=True
            ) as dwo:
                assert dwo.is_complex
                assert dwo.realdtype == realdtype
                assert dwo.dtype == dtype
                # test writing interleaved data
                dwo.rf_write(np.ones((shape[0], 2 * shape[1]), dtype=realdtype))
                # test writing complex/structured data
                dwo.rf_write(np.ones(shape, dtype=dtype))

    @pytest.mark.firstonly("hdf_filter_params")
    def test_writer_write_blocks(
        self,
        bounds,
        channel,
        chdir,
        data,
        data_block_slices,
        drf_writer_factory,
        start_global_index,
    ):
        """Test writer object's write_blocks method."""
        chdir = chdir.dirpath().mkdir("{0}_write_blocks".format(channel))
        with drf_writer_factory(directory=str(chdir)) as dwo:
            # write the specified data blocks from data_block_slices
            block_starts, block_stops = zip(*data_block_slices)
            global_sample_arr = np.asarray(block_starts) - start_global_index
            block_lengths = np.asarray(block_stops) - np.asarray(block_starts)
            block_sample_arr = np.concatenate(([0], np.cumsum(block_lengths)[:-1]))
            data_arr = np.concatenate(
                tuple(
                    data[(bs - bounds[0]) : (be - bounds[0])]
                    for bs, be in zip(block_starts, block_stops)
                )
            )

            next_rel_index = dwo.rf_write_blocks(
                data_arr,
                global_sample_arr=global_sample_arr,
                block_sample_arr=block_sample_arr,
            )
            assert next_rel_index == (block_stops[-1] - bounds[0])

            # fail with non-matching length of sample lists
            with pytest.raises(ValueError):
                dwo.rf_write_blocks(
                    data_arr,
                    global_sample_arr=[next_rel_index, next_rel_index + 10],
                    block_sample_arr=[0],
                )
            # fail with bad type for sample lists
            with pytest.raises(TypeError):
                dwo.rf_write_blocks(
                    data_arr,
                    global_sample_arr=[next_rel_index + 1.1],
                    block_sample_arr=[0],
                )
            # fail with bad shape for sample lists
            with pytest.raises(ValueError):
                dwo.rf_write_blocks(
                    data_arr,
                    global_sample_arr=[next_rel_index],
                    block_sample_arr=np.zeros((1, 1), dtype=np.uint64),
                )
            # fail with nonzero block_sample_arr[0]
            with pytest.raises(ValueError):
                dwo.rf_write_blocks(
                    data_arr, global_sample_arr=[next_rel_index], block_sample_arr=[10]
                )
            # fail with invalid increment in block_sample_arr
            with pytest.raises(ValueError):
                dwo.rf_write_blocks(
                    data_arr,
                    global_sample_arr=[next_rel_index, next_rel_index + 10],
                    block_sample_arr=[0, 11],
                )
            # fail with non-increasing block_sample_arr
            with pytest.raises(ValueError):
                dwo.rf_write_blocks(
                    data_arr,
                    global_sample_arr=[
                        next_rel_index,
                        next_rel_index + 1,
                        next_rel_index + 2,
                    ],
                    block_sample_arr=[0, 1, 0],
                )
            # fail with block_sample_arr value > len(arr)
            with pytest.raises(ValueError):
                dwo.rf_write_blocks(
                    data_arr[:5],
                    global_sample_arr=[next_rel_index, next_rel_index + 10],
                    block_sample_arr=[0, 10],
                )
            # fail at overwriting
            with pytest.raises(ValueError):
                dwo.rf_write_blocks(
                    data_arr,
                    global_sample_arr=[next_rel_index - 1],
                    block_sample_arr=[0],
                )
            # fail with non-increasing global_sample_arr
            with pytest.raises(ValueError):
                dwo.rf_write_blocks(
                    data_arr,
                    global_sample_arr=[
                        bounds[1] - start_global_index + 1,
                        bounds[1] - start_global_index,
                    ],
                    block_sample_arr=[0, 1],
                )

    @pytest.mark.firstonly("hdf_filter_params")
    def test_writer_output_files_write_blocks(self, channel, chdir, data_file_list):
        """Test output files created by writer object in write_blocks test."""
        # check that the correct subdirs and filenames are produced
        # (we don't check the contents of the written files here since testing
        #  the reader will accomplish that as long as we keep in mind that
        #  errors there could actually indicate a problem in writing)
        chdir = chdir.dirpath("{0}_write_blocks".format(channel))
        drf_files = digital_rf.lsdrf(
            str(chdir),
            include_drf=True,
            include_dmd=False,
            include_drf_properties=False,
        )
        drf_files = [os.path.relpath(p, str(chdir)) for p in sorted(drf_files)]
        assert drf_files == data_file_list

    def test_writer_data_write(
        self,
        bounds,
        chdir,
        data,
        data_block_slices,
        data_file_list,
        drf_writer,
        start_global_index,
    ):
        """Test writing specific data with writer object's write method."""
        # write the specified data blocks from data_block_slices
        next_rel_index = drf_writer.get_next_available_sample()
        for k, (sstart, sstop) in enumerate(data_block_slices):
            rel_index = sstart - start_global_index
            bstart = sstart - bounds[0]
            bstop = sstop - bounds[0]
            wdata = data[bstart:bstop]
            if k == (len(data_block_slices) - 1):
                # test writing data from a row-order (Fortran-style) array
                wdata = np.asfortranarray(wdata)
            if rel_index == next_rel_index:
                # test writing automatically to next sample
                next_rel_index = drf_writer.rf_write(wdata)
            else:
                # test writing to specified sample
                next_rel_index = drf_writer.rf_write(wdata, next_sample=rel_index)
            assert next_rel_index == rel_index + (bstop - bstart)

        # close writer to close out last file
        # (IMPORTANT: can't write using drf_writer after this test!)
        drf_writer.close()
        with pytest.raises(IOError):
            drf_writer.rf_write(data)

    def test_writer_output_files_write(
        self, bounds, chdir, data, data_file_list, num_subchannels
    ):
        """Test output files created by writer object in write test."""
        # check that the correct subdirs and filenames are produced
        # and that data that exists in the files is correct
        # Note: this does not fully check that everything we intended to write
        # was written correctly
        drf_files = digital_rf.lsdrf(
            str(chdir),
            include_drf=True,
            include_dmd=False,
            include_drf_properties=False,
        )
        drf_files = [os.path.relpath(p, str(chdir)) for p in sorted(drf_files)]
        assert drf_files == data_file_list

        for fname in data_file_list:
            fpath = str(chdir.join(fname))
            with h5py.File(fpath, "r") as f:
                rdata = f["rf_data"]
                data_index = f["rf_data_index"]
                for blk in range(data_index.shape[0]):
                    sstart = int(data_index[blk, 0])
                    bstart = int(data_index[blk, 1])
                    if blk + 1 == data_index.shape[0]:
                        bstop = rdata.shape[0]
                    else:
                        bstop = int(data_index[blk + 1, 1])
                    blk_data = rdata[bstart:bstop]
                    assert len(blk_data.shape) == 2
                    dstart = sstart - bounds[0]
                    dstop = sstart + (bstop - bstart) - bounds[0]
                    test_data = data[dstart:dstop]
                    np.testing.assert_equal(
                        blk_data, test_data.reshape((-1, num_subchannels))
                    )

    def test_writer_get_total_samples_written(self, data_block_slices, drf_writer):
        """Test writer object's get_total_samples_written method."""
        num_written = sum([se - ss for ss, se in data_block_slices])
        assert drf_writer.get_total_samples_written() == num_written

    def test_writer_get_next_available_sample(
        self, data_block_slices, drf_writer, start_global_index
    ):
        """Test writer object's get_next_available_sample method."""
        next_sample = data_block_slices[-1][1] - start_global_index
        assert drf_writer.get_next_available_sample() == next_sample

    def test_writer_get_total_gap_samples(self, drf_writer):
        """Test writer object's get_total_gap_samples method."""
        gap_samples = (
            drf_writer.get_next_available_sample()
            - drf_writer.get_total_samples_written()
        )
        assert drf_writer.get_total_gap_samples() == gap_samples

    def test_writer_get_last_file_written(self, chdir, data_file_list, drf_writer):
        """Test writer object's get_last_file_written method."""
        last_file = os.path.normpath(drf_writer.get_last_file_written())
        last_file_full_path = str(chdir.join(data_file_list[-1]))
        assert last_file == last_file_full_path

    def test_writer_get_last_dir_written(self, chdir, data_file_list, drf_writer):
        """Test writer object's get_last_dir_written method."""
        last_dir = os.path.normpath(drf_writer.get_last_dir_written())
        last_dir_full_path = chdir.join(data_file_list[-1]).dirname
        assert last_dir == last_dir_full_path

    def test_writer_get_last_utc_timestamp(self, drf_writer):
        """Test writer object's get_last_utc_timestamp method."""
        assert drf_writer.get_last_utc_timestamp()

    @pytest.mark.firstonly(
        "data_params", "file_params", "hdf_filter_params", "sample_params"
    )
    def test_writer_get_version(self, drf_writer):
        """Test writer object's get_version method."""
        assert packaging.version.Version(drf_writer.get_version())

    def test_reader_get_channels(self, channel, drf_reader):
        """Test reader object's get_channels method."""
        channels = drf_reader.get_channels()
        assert channel in channels

    @pytest.mark.firstonly("data_params", "hdf_filter_params")
    def test_reader_get_bounds(self, bounds, channel, drf_reader):
        """Test reader object's get_bounds method."""
        bounds = drf_reader.get_bounds(channel)
        assert bounds == bounds

        # fail when channel doesn't exist
        with pytest.raises(KeyError):
            drf_reader.get_bounds("not_a_channel")

    @pytest.mark.firstonly("hdf_filter_params")
    def test_reader_get_properties(
        self,
        bounds,
        channel,
        drf_reader,
        drf_writer_param_dict,
        samples_per_second,
        start_datetime,
        start_global_index,
    ):
        """Test reader object's get_properties method."""
        p = drf_writer_param_dict
        expected_properties = dict(
            H5Tget_class=None,
            H5Tget_offset=None,
            H5Tget_order=None,
            H5Tget_precision=None,
            H5Tget_size=None,
            digital_rf_time_description=None,
            digital_rf_version=None,
            epoch="1970-01-01T00:00:00Z",
            file_cadence_millisecs=p["file_cadence_millisecs"],
            is_complex=p["is_complex"],
            is_continuous=p["is_continuous"],
            num_subchannels=p["num_subchannels"],
            samples_per_second=samples_per_second,
            sample_rate_numerator=p["sample_rate_numerator"],
            sample_rate_denominator=p["sample_rate_denominator"],
            subdir_cadence_secs=p["subdir_cadence_secs"],
        )
        props = drf_reader.get_properties(channel)
        for k, v in expected_properties.items():
            if v is not None:
                assert props[k] == v
            else:
                props[k]

        expected_sample_properties = dict(
            computer_time=None,
            init_utc_timestamp=int(
                digital_rf.util.datetime_to_timestamp(start_datetime)
            ),
            sequence_num=0,
            uuid_str=p["uuid_str"],
        )
        sample_props = drf_reader.get_properties(channel, sample=start_global_index)
        for k, v in expected_sample_properties.items():
            if v is not None:
                assert sample_props[k] == v
            else:
                sample_props[k]

        # fail when channel doesn't exist
        with pytest.raises(KeyError):
            drf_reader.get_properties("not_a_channel")
        # fail to get properties when sample doesn't exist
        with pytest.raises(IOError):
            drf_reader.get_properties(channel, sample=bounds[0] - 1)
        with pytest.raises(IOError):
            drf_reader.get_properties(channel, sample=bounds[1] + 1)

    @pytest.mark.firstonly("data_params", "file_params", "hdf_filter_params")
    def test_reader_get_digital_metadata(
        self, channel, chdir, drf_reader, sample_params, start_global_index
    ):
        """Test reader object's get_digital_metadata method."""
        # first we need to write to the accompanying metadata channel
        metadata_dir = chdir.mkdir("metadata")
        dmd_writer = digital_rf.DigitalMetadataWriter(
            metadata_dir=str(metadata_dir),
            subdir_cadence_secs=sample_params["subdir_cadence_secs"],
            file_cadence_secs=1,
            sample_rate_numerator=(sample_params["sample_rate_numerator"]),
            sample_rate_denominator=(sample_params["sample_rate_denominator"]),
            file_name="metadata",
        )
        dmd_writer.write(start_global_index, dict(test=True))

        dmd_reader = drf_reader.get_digital_metadata(channel)
        assert type(dmd_reader) is digital_rf.DigitalMetadataReader

        # fail when channel doesn't exist
        with pytest.raises(IOError):
            drf_reader.get_digital_metadata("not_a_channel")

    @pytest.mark.firstonly("data_params", "file_params", "hdf_filter_params")
    def test_reader_read_metadata(
        self,
        bounds,
        channel,
        drf_reader,
        samples_per_second,
        sample_rate_numerator,
        sample_rate_denominator,
        start_global_index,
    ):
        """Test reader object's read_metadata method."""
        # must be run after test_reader_get_digital_metadata which creates
        # the metadata that we'll be reading
        expected_metadata = dict(
            samples_per_second=samples_per_second,
            sample_rate_numerator=sample_rate_numerator,
            sample_rate_denominator=sample_rate_denominator,
        )
        # read the blocks channel first, which has no Digital Metadata channel
        # all metadata
        block_channel = "{0}_write_blocks".format(channel)
        md = drf_reader.read_metadata(bounds[0], bounds[1], block_channel, method=None)
        assert list(md.keys()) == [bounds[0]]
        for k, v in expected_metadata.items():
            assert md[start_global_index][k] == v
        md2 = drf_reader.read_metadata(
            bounds[0] + 1, bounds[1], block_channel, method="ffill"
        )
        assert list(md2.keys()) == [bounds[0] + 1]
        for k, v in expected_metadata.items():
            assert md2[start_global_index + 1][k] == v

        # read the regular channel, which has a Digital Metadata channel
        expected_metadata["test"] = True
        md = drf_reader.read_metadata(bounds[0], bounds[1], channel, method=None)
        assert list(md.keys()) == [start_global_index]
        for k, v in expected_metadata.items():
            assert md[start_global_index][k] == v
        md2 = drf_reader.read_metadata(
            start_global_index + 1, bounds[1], channel, method="ffill"
        )
        assert md == md2

    @pytest.mark.firstonly("data_params", "file_params", "hdf_filter_params")
    def test_reader_get_last_write(self, channel, chdir, data_file_list, drf_reader):
        """Test reader object's get_last_write method."""
        timestamp, fpath = drf_reader.get_last_write(channel)
        assert timestamp
        assert fpath == str(chdir.join(data_file_list[-1]))

        # fail when channel doesn't exist
        with pytest.raises(KeyError):
            drf_reader.get_last_write("not_a_channel")

    @pytest.mark.firstonly("data_params", "file_params", "hdf_filter_params")
    def test_reader_get_continuous_blocks(
        self, bounds, channel, data_block_slices, drf_reader
    ):
        """Test reader object's get_continuous_blocks method."""
        # read all of the data blocks that were written in writer test
        for sstart, sstop in data_block_slices:
            data_dict = drf_reader.get_continuous_blocks(sstart, sstop - 1, channel)
            assert len(data_dict) == 1
            sample, nsamples = list(data_dict.items())[0]
            assert sample == sstart
            assert nsamples == sstop - sstart

        # check that we get all of the data if we read entire bounds
        data_dict = drf_reader.get_continuous_blocks(bounds[0], bounds[1], channel)
        assert bounds[0] in data_dict
        sample, nsamples = sorted(data_dict.items())[-1]
        assert sample + nsamples - 1 == bounds[1]

        # test read where data doesn't exist
        data_dict = drf_reader.get_continuous_blocks(
            bounds[0] - 100, bounds[0] - 1, channel
        )
        assert data_dict == {}

        # fail when channel doesn't exist
        with pytest.raises(KeyError):
            drf_reader.get_continuous_blocks(bounds[0], bounds[1], "not_a_channel")

    def test_reader_read(
        self, bounds, channel, data, data_block_slices, drf_reader, num_subchannels
    ):
        """Test reader object's read method."""
        # read all of the data blocks that were written in writer test
        for sstart, sstop in data_block_slices:
            data_dict = drf_reader.read(sstart, sstop - 1, channel)
            assert len(data_dict) == 1
            sample, rdata = list(data_dict.items())[0]
            assert sample == sstart
            assert len(rdata.shape) == 2
            bstart = sstart - bounds[0]
            bstop = sstop - bounds[0]
            np.testing.assert_equal(
                rdata, data.reshape((-1, num_subchannels))[bstart:bstop, :]
            )

        # test reading a single subchannel (the first)
        for sstart, sstop in data_block_slices:
            data_dict = drf_reader.read(sstart, sstop - 1, channel, sub_channel=0)
            assert len(data_dict) == 1
            sample, rdata = list(data_dict.items())[0]
            assert sample == sstart
            bstart = sstart - bounds[0]
            bstop = sstop - bounds[0]
            np.testing.assert_equal(
                rdata, data.reshape((-1, num_subchannels))[bstart:bstop, 0]
            )

        # check that we get all of the data if we read entire bounds
        data_dict = drf_reader.read(bounds[0], bounds[1], channel)
        for (sstart, rdata) in data_dict.items():
            bstart = sstart - bounds[0]
            bstop = sstart + len(rdata) - bounds[0]
            np.testing.assert_equal(
                rdata, data.reshape((-1, num_subchannels))[bstart:bstop, :]
            )

        # test read where data doesn't exist
        data_dict = drf_reader.read(bounds[0] - 100, bounds[0] - 1, channel)
        assert data_dict == {}

        # fail when channel doesn't exist
        with pytest.raises(KeyError):
            drf_reader.read(bounds[0], bounds[1], "not_a_channel")
        # fail when end_sample < start_sample (e.g. nsamples passed instead)
        with pytest.raises(ValueError):
            drf_reader.read(bounds[0], bounds[1] - bounds[0] + 1, channel)
        # fail when subchannel doesn't exist
        with pytest.raises(ValueError):
            drf_reader.read(bounds[0], bounds[1], channel, sub_channel=num_subchannels)

    def test_reader_read_vector_raw(
        self, bounds, channel, data, data_block_slices, drf_reader, num_subchannels
    ):
        """Test reader object's read_vector_raw method."""
        # read all of the data blocks that were written in writer test
        for sstart, sstop in data_block_slices:
            rdata = drf_reader.read_vector_raw(sstart, sstop - sstart, channel)
            bstart = sstart - bounds[0]
            bstop = sstop - bounds[0]
            np.testing.assert_equal(rdata, data[bstart:bstop])

        # test reading a single subchannel (the first)
        for sstart, sstop in data_block_slices:
            rdata = drf_reader.read_vector_raw(
                sstart, sstop - sstart, channel, sub_channel=0
            )
            bstart = sstart - bounds[0]
            bstop = sstop - bounds[0]
            np.testing.assert_equal(
                rdata, data.reshape((-1, num_subchannels))[bstart:bstop, 0]
            )

        # fail to read all data at once because of gap
        with pytest.raises(IOError):
            rdata = drf_reader.read_vector_raw(
                bounds[0], bounds[1] - bounds[0] + 1, channel
            )

        # fail when data doesn't exist
        with pytest.raises(IOError):
            drf_reader.read_vector_raw(bounds[0] - 100, 99, channel)
        # fail when channel doesn't exist
        with pytest.raises(KeyError):
            drf_reader.read_vector_raw(
                bounds[0], bounds[1] - bounds[0] + 1, "not_a_channel"
            )
        # fail when subchannel doesn't exist
        with pytest.raises(ValueError):
            drf_reader.read_vector_raw(
                bounds[0],
                bounds[1] - bounds[0] + 1,
                channel,
                sub_channel=num_subchannels,
            )

    def test_reader_read_vector(
        self, bounds, channel, data, data_block_slices, drf_reader, num_subchannels
    ):
        """Test reader object's read_vector method."""
        if data.dtype.names is not None:
            c8data = np.empty(data.shape, dtype=np.complex64)
            c8data.real = data["r"]
            c8data.imag = data["i"]
        else:
            c8data = np.array(data, dtype=np.complex64, copy=False)

        # read all of the data blocks that were written in writer test
        for sstart, sstop in data_block_slices:
            rdata = drf_reader.read_vector(sstart, sstop - sstart, channel)
            bstart = sstart - bounds[0]
            bstop = sstop - bounds[0]
            np.testing.assert_equal(rdata, c8data[bstart:bstop])

        # test reading a single subchannel (the first)
        for sstart, sstop in data_block_slices:
            rdata = drf_reader.read_vector(
                sstart, sstop - sstart, channel, sub_channel=0
            )
            bstart = sstart - bounds[0]
            bstop = sstop - bounds[0]
            np.testing.assert_equal(
                rdata, c8data.reshape((-1, num_subchannels))[bstart:bstop, 0]
            )

        # fail to read all data at once because of gap
        with pytest.raises(IOError):
            rdata = drf_reader.read_vector(
                bounds[0], bounds[1] - bounds[0] + 1, channel
            )

        # fail when data doesn't exist
        with pytest.raises(IOError):
            drf_reader.read_vector(bounds[0] - 100, 99, channel)
        # fail when channel doesn't exist
        with pytest.raises(KeyError):
            drf_reader.read_vector(
                bounds[0], bounds[1] - bounds[0] + 1, "not_a_channel"
            )
        # fail when subchannel doesn't exist
        with pytest.raises(ValueError):
            drf_reader.read_vector(
                bounds[0],
                bounds[1] - bounds[0] + 1,
                channel,
                sub_channel=num_subchannels,
            )

    def test_reader_read_vector_c81d(
        self, bounds, channel, data, data_block_slices, drf_reader, num_subchannels
    ):
        """Test reader object's read_vector_c81d method."""
        if data.dtype.names is not None:
            c8data = np.empty(data.shape, dtype=np.complex64)
            c8data.real = data["r"]
            c8data.imag = data["i"]
        else:
            c8data = np.array(data, dtype=np.complex64, copy=False)
        # make c8data have two dimensions regardless of shape
        c8data = c8data.reshape((-1, num_subchannels))

        # read all of the data blocks that were written in writer test
        for sstart, sstop in data_block_slices:
            rdata = drf_reader.read_vector_c81d(sstart, sstop - sstart, channel)
            bstart = sstart - bounds[0]
            bstop = sstop - bounds[0]
            np.testing.assert_equal(rdata, c8data[bstart:bstop, 0])

        # test reading a different subchannel (the last)
        for sstart, sstop in data_block_slices:
            rdata = drf_reader.read_vector_c81d(
                sstart, sstop - sstart, channel, sub_channel=num_subchannels - 1
            )
            bstart = sstart - bounds[0]
            bstop = sstop - bounds[0]
            np.testing.assert_equal(rdata, c8data[bstart:bstop, num_subchannels - 1])

        # fail to read all data at once because of gap
        with pytest.raises(IOError):
            rdata = drf_reader.read_vector_c81d(
                bounds[0], bounds[1] - bounds[0] + 1, channel
            )

        # fail when data doesn't exist
        with pytest.raises(IOError):
            drf_reader.read_vector_c81d(bounds[0] - 100, 99, channel)
        # fail when channel doesn't exist
        with pytest.raises(KeyError):
            drf_reader.read_vector_c81d(
                bounds[0], bounds[1] - bounds[0] + 1, "not_a_channel"
            )
        # fail when subchannel doesn't exist
        with pytest.raises(ValueError):
            drf_reader.read_vector_c81d(
                bounds[0],
                bounds[1] - bounds[0] + 1,
                channel,
                sub_channel=num_subchannels,
            )

    @pytest.mark.firstonly("data_params", "hdf_filter_params", "sample_params")
    def test_reader_multiple_topleveldirs(
        self,
        bounds,
        data,
        data_block_slices,
        drf_writer_factory,
        is_continuous,
        num_subchannels,
        start_global_index,
        tmpdir_factory,
    ):
        """Test reader object with multiple top-level directories."""
        # write data blocks alternately in multiple top-level directories
        channel = "ch0"
        tlds = [tmpdir_factory.mktemp("tld_test_") for k in range(3)]
        writers = [
            drf_writer_factory(directory=str(tld.mkdir(channel))) for tld in tlds
        ]
        for dwo, (sstart, sstop) in zip(itertools.cycle(writers), data_block_slices):
            rel_index = sstart - start_global_index
            bstart = sstart - bounds[0]
            bstop = sstop - bounds[0]
            dwo.rf_write(data[bstart:bstop], next_sample=rel_index)
        for dwo in writers:
            dwo.close()

        # get reader for all of the top-level directories
        with digital_rf.DigitalRFReader([str(d) for d in tlds]) as dro:
            # check bounds
            assert dro.get_bounds(channel) == bounds

            if not is_continuous:
                # read all of the data blocks
                for sstart, sstop in data_block_slices:
                    data_dict = dro.read(sstart, sstop - 1, channel)
                    assert len(data_dict) == 1
                    sample, rdata = list(data_dict.items())[0]
                    assert sample == sstart
                    bstart = sstart - bounds[0]
                    bstop = sstop - bounds[0]
                    np.testing.assert_equal(
                        rdata, data.reshape((-1, num_subchannels))[bstart:bstop]
                    )

ext: py

reason: import numpy


text:
import os
import shutil
from glob import glob
from math import ceil, floor

import cv2
import numpy as np
import onnxruntime as rt
import skimage

def floor_even(x):
    if isinstance(x, tuple):
        return tuple(floor_even(y) for y in x)
    if isinstance(x, list):
        return [floor_even(y) for y in x]
    return x // 2 * 2

# Does not copy img and alpha
def trim_img(img, alpha, eps):
    original_shape = alpha.shape

    trim_t = 0
    while np.all(alpha[trim_t, :] < eps):
        trim_t += 1
    trim_b = alpha.shape[0] - 1
    while np.all(alpha[trim_b, :] < eps):
        trim_b -= 1
    trim_l = 0
    while np.all(alpha[:, trim_l] < eps):
        trim_l += 1
    trim_r = alpha.shape[1] - 1
    while np.all(alpha[:, trim_r] < eps):
        trim_r -= 1

    trim_b += 1
    trim_r += 1

    trims = (trim_t, trim_b, trim_l, trim_r)
    return original_shape, trims

def untrim_img(img, alpha, original_shape, trims):
    trim_t, trim_b, trim_l, trim_r = trims
    new_img = np.zeros((original_shape[0], original_shape[1], 3))
    new_alpha = np.zeros(original_shape)
    new_img[trim_t:trim_b, trim_l:trim_r, :] = img
    new_alpha[trim_t:trim_b, trim_l:trim_r] = alpha
    return new_img, new_alpha

def get_pieces(img, piece_inner_size, pad_size):
    piece_outer_size = piece_inner_size + pad_size * 2

    max_row = ceil(img.shape[0] / piece_inner_size)
    max_col = ceil(img.shape[1] / piece_inner_size)
    img_padded_h = max_row * piece_inner_size
    img_padded_w = max_col * piece_inner_size
    pad_t = floor((img_padded_h - img.shape[0]) / 2)
    pad_b = img_padded_h - img.shape[0] - pad_t
    pad_l = floor((img_padded_w - img.shape[1]) / 2)
    pad_r = img_padded_w - img.shape[1] - pad_l
    img_full = np.pad(img, [
        (pad_t + pad_size, pad_b + pad_size),
        (pad_l + pad_size, pad_r + pad_size),
        (0, 0),
    ], 'reflect')

    pieces = []
    for i in range(max_row):
        for j in range(max_col):
            idx_t = i * piece_inner_size
            idx_b = idx_t + piece_outer_size
            idx_l = j * piece_inner_size
            idx_r = idx_l + piece_outer_size
            pieces.append(img_full[idx_t:idx_b, idx_l:idx_r, :])
    pieces = np.stack(pieces)

    max_row_col = (max_row, max_col)
    pads = (pad_t, pad_b, pad_l, pad_r)
    return pieces, max_row_col, pads

def get_batch(pieces, batch_size):
    idx = 0
    while idx < pieces.shape[0]:
        batch = pieces[idx:idx + batch_size]
        batch = batch.transpose(0, 3, 1, 2)
        idx += batch.shape[0]
        print('Piece {}/{}'.format(idx, pieces.shape[0]))
        yield batch

def merge_img(pieces,
              piece_inner_size,
              pad_size,
              max_row_col,
              pads,
              scale_shift=(1, 0)):
    max_row, max_col = max_row_col
    pad_t, pad_b, pad_l, pad_r = pads
    scale, shift = scale_shift

    piece_outer_size = piece_inner_size + pad_size * 2
    scaled_inner_size = piece_inner_size * scale
    scaled_outer_size = piece_outer_size * scale

    img = np.empty(
        (max_row * scaled_outer_size, max_col * scaled_outer_size, 3))
    for idx, piece in enumerate(pieces):
        i = idx // max_col
        j = idx % max_col

        idx_t = i * scaled_inner_size
        idx_b = idx_t + scaled_inner_size
        idx_l = j * scaled_inner_size
        idx_r = idx_l + scaled_inner_size
        piece_l = pad_size * scale - shift
        piece_r = piece_l + scaled_inner_size
        img[idx_t:idx_b, idx_l:idx_r, :] = piece[piece_l:piece_r,
                                                 piece_l:piece_r, :]

    img = img[pad_t * scale:(max_row * piece_inner_size - pad_b) * scale,
              pad_l * scale:(max_col * piece_inner_size - pad_r) * scale, :]
    return img

def read_img(filename,
             swap_rb=False,
             gamma=1,
             signed=True,
             scale=None,
             noise=0,
             return_alpha=False):
    # Use cv2 to support 16 bit image
    img = cv2.imread(filename, cv2.IMREAD_UNCHANGED)
    img = skimage.img_as_float32(img)

    alpha = None

    if img.ndim == 3:
        if img.shape[2] == 4:
            alpha = img[:, :, 3]
        else:
            assert img.shape[2] == 3
        # Remove alpha channel
        img = img[:, :, :3]
    else:
        assert img.ndim == 2
        # Convert grayscale to RGB
        img = np.repeat(img[:, :, None], 3, axis=2)

    if swap_rb:
        assert img.ndim == 3
        # BGR -> RGB
        img = img[:, :, ::-1]

    img **= gamma

    if signed:
        # [0, 1] -> [-1, 1]
        img = img * 2 - 1

    if scale is not None:
        img *= scale
        if alpha is not None:
            alpha *= scale

    if noise:
        rng = np.random.default_rng()
        img += rng.normal(scale=noise, size=img.shape)

    if return_alpha:
        return img, alpha
    else:
        return img

def write_img(filename,
              img,
              alpha=None,
              swap_rb=False,
              signed=True,
              scale=None,
              output_gray=False,
              output_8_bit=True):
    if scale is not None:
        img /= scale
        if alpha is not None:
            alpha /= scale

    if signed:
        # [-1, 1] -> [0, 1]
        img = (img + 1) / 2

    img = np.clip(img, 0, 1)

    if swap_rb:
        # RGB -> BGR
        img = img[:, :, ::-1]

    if output_gray:
        img = img.mean(axis=2, keepdims=True)

    if alpha is not None:
        if alpha.ndim == 2:
            alpha = alpha[:, :, None]
        img = np.concatenate([img, alpha], axis=2)

    if output_8_bit:
        img = skimage.img_as_ubyte(img)
    else:
        img = skimage.img_as_uint(img)

    cv2.imwrite(filename, img)

def do_imgs(fun,
            model_filename,
            in_filenames,
            out_suffix,
            out_extname=None,
            tmp_filename=None):
    if model_filename:
        sess = rt.InferenceSession(
            model_filename,
            providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
    else:
        sess = None

    if type(in_filenames) == str:
        in_filenames = [in_filenames]

    in_filenames = sum([glob(x) for x in in_filenames], [])
    if not in_filenames:
        print('No in_filenames')

    for in_filename in in_filenames:
        print(in_filename)

        basename, extname = os.path.splitext(in_filename)
        if out_extname:
            out_filename = basename + out_suffix + out_extname
        else:
            out_filename = basename + out_suffix + extname

        if tmp_filename:
            shutil.copy2(in_filename, tmp_filename)
            fun(sess, tmp_filename, tmp_filename)
            shutil.move(tmp_filename, out_filename)
        else:
            fun(sess, in_filename, out_filename)

max_stars_repo_name: Lunatic-Works/neural-tools

max_issues_repo_name: Lunatic-Works/neural-tools

max_forks_repo_name: Lunatic-Works/neural-tools


text:
#!/usr/bin/env python3

import functools
import sys

import numpy as np
from PyQt5.QtCore import *
from PyQt5.QtGui import *

from .util import file_io
from .util import project_functions as pfs
from .util.plugin import PluginDefault
from .util.plugin import WidgetDefault

class Widget(QWidget, WidgetDefault):
  class Labels(WidgetDefault.Labels):
    crop_percentage_sb_label = "Crop Percentage"

  class Defaults(WidgetDefault.Defaults):
      crop_percentage_sb_default = 25

  def __init__(self, project, plugin_position, parent=None):
    super(Widget, self).__init__(parent=parent)
    if not project or not isinstance(plugin_position, int):
        return
    self.main_button = QPushButton('Crop')
    self.crop_percentage_sb = QSpinBox()
    self.left_frame_range = QSpinBox()
    self.right_frame_range = QSpinBox()
    WidgetDefault.__init__(self, project=project, plugin_position=plugin_position)

  def setup_ui(self):
    super().setup_ui()
    self.crop_percentage_sb.setMaximum(99)
    self.vbox.addWidget(QLabel(self.Labels.crop_percentage_sb_label))
    self.vbox.addWidget(self.crop_percentage_sb)

    def min_handler(max_of_min):
        self.left_frame_range.setMaximum(max_of_min)
    def max_handler(min_of_max):
        self.right_frame_range.setMinimum(min_of_max)
    self.left_frame_range.valueChanged[int].connect(max_handler)
    self.right_frame_range.valueChanged[int].connect(min_handler)
    hbox = QHBoxLayout()
    self.vbox.addWidget(QLabel('Set range of frames cropped'))
    self.left_frame_range.setMinimum(0)
    self.left_frame_range.setMaximum(1400)
    self.left_frame_range.setValue(400)
    hbox.addWidget(self.left_frame_range)
    to = QLabel('to')
    to.setAlignment(Qt.AlignCenter)
    hbox.addWidget(to)
    self.right_frame_range.setMaximum(1000000)
    self.right_frame_range.setValue(1400)
    hbox.addWidget(self.right_frame_range)
    self.vbox.addLayout(hbox)

    self.vbox.addStretch()
    self.vbox.addWidget(self.main_button)

  def setup_signals(self):
    super().setup_signals()
    self.main_button.clicked.connect(self.execute_primary_function)

  def setup_params(self, reset=False):
      super().setup_params(reset)
      if len(self.params) == 1 or reset:
        self.update_plugin_params(self.Labels.crop_percentage_sb_label, self.Defaults.crop_percentage_sb_default)
      self.crop_percentage_sb.setValue(self.params[self.Labels.crop_percentage_sb_label])

  def setup_param_signals(self):
      super().setup_param_signals()
      self.crop_percentage_sb.valueChanged[int].connect(functools.partial(self.update_plugin_params,
                                                                    self.Labels.crop_percentage_sb_label))

  def execute_primary_function(self, input_paths=None):
      '''Primary function of plugin'''
      if not input_paths:
          if not self.selected_videos:
              return
          else:
              selected_videos = self.selected_videos
      else:
          selected_videos = input_paths
      # def crop_center(img, cropx, cropy):
      #     """Function to crop center of an image file. Output is of size (cropx, cropy)"""
      #     y, x = img.shape
      #     startx = x // 2 - (cropx // 2)
      #     starty = y // 2 - (cropy // 2)
      #     return img[starty:starty + cropy, startx:startx + cropx]
      output_paths = []
      for path in selected_videos:
          frames = file_io.load_file(path, segment=[self.left_frame_range.value(), self.right_frame_range.value()])
          percentage = self.crop_percentage_sb.value() / 100
          # frames_cropped = np.empty((frames.shape[0], frames.shape[1], frames.shape[2]),
          #               dtype=frames.dtype)
          cropped_y = round(percentage * frames.shape[1])
          cropped_x = round(percentage * frames.shape[2])
          for frame_no in range(len(frames)):
              for y in range(len(frames[frame_no])):
                  for x in range(len(frames[frame_no][y])):
                      if (x < cropped_x or x > frames.shape[2] - cropped_x) or\
                              (y < cropped_y or y > frames.shape[1] - cropped_y):
                          frames[frame_no][y][x] = 0
          output_paths = output_paths + [path]
          pfs.save_project(path, self.project, frames, 'crop-border', 'video')
          pfs.refresh_list(self.project, self.video_list,
                           self.params[self.Labels.video_list_indices_label],
                           self.Defaults.list_display_type,
                           self.params[self.Labels.last_manips_to_display_label])
      return output_paths

  def setup_whats_this(self):
      '''Setup custom help messages'''
      # todo: setup custom help messages to aid the user, each tied to one of your UI components.
      super().setup_whats_this()

class MyPlugin(PluginDefault):
  def __init__(self, project, plugin_position):
    self.name = 'Crop Border' # Define plugin name here
    self.widget = Widget(project, plugin_position)
    super().__init__(self.widget, self.widget.Labels, self.name)

    # todo: over-ride PluginDefault functions here to define custom behaviour
    # (required for automation)

if __name__=='__main__':
  app = QApplication(sys.argv)
  app.aboutToQuit.connect(app.deleteLater)
  w = QMainWindow()
  w.setCentralWidget(Widget(None, None))
  w.show()
  app.exec_()
  sys.exit()

max_stars_repo_name: oraby/Mesoscale-Brain-Explorer

max_issues_repo_name: oraby/Mesoscale-Brain-Explorer

max_forks_repo_name: oraby/Mesoscale-Brain-Explorer


text:
# -*- coding: utf-8 -*-
"""
License: MIT
@author: gaj
E-mail: anjing_guo@hnu.edu.cn
"""

import numpy as np
import cv2
import os
from scipy import signal

from methods.Bicubic import Bicubic
from methods.Brovey import Brovey
from methods.PCA import PCA
from methods.IHS import IHS
from methods.SFIM import SFIM
from methods.GS import GS
from methods.Wavelet import Wavelet
from methods.MTF_GLP import MTF_GLP
from methods.MTF_GLP_HPM import MTF_GLP_HPM
from methods.GSA import GSA
from methods.CNMF import CNMF
from methods.GFPCA import GFPCA
from methods.PNN import PNN
from methods.PanNet import PanNet

from metrics import ref_evaluate, no_ref_evaluate

'''loading data'''
original_msi = np.load('./images/GF2_BJ_mss.npy')
original_pan = np.load('./images/GF2_BJ_pan.npy')

'''normalization'''
max_patch, min_patch = np.max(original_msi, axis=(0,1)), np.min(original_msi, axis=(0,1))
original_msi = np.float32(original_msi-min_patch) / (max_patch - min_patch)

max_patch, min_patch = np.max(original_pan, axis=(0,1)), np.min(original_pan, axis=(0,1))
original_pan = np.float32(original_pan-min_patch) / (max_patch - min_patch)

'''generating ms image with gaussian kernel'''
sig = (1/(2*(2.772587)/4**2))**0.5
kernel = np.multiply(cv2.getGaussianKernel(9, sig), cv2.getGaussianKernel(9,sig).T)
new_lrhs = []
for i in range(original_msi.shape[-1]):
    temp = signal.convolve2d(original_msi[:,:, i], kernel, boundary='wrap',mode='same')
    temp = np.expand_dims(temp, -1)
    new_lrhs.append(temp)
new_lrhs = np.concatenate(new_lrhs, axis=-1)
used_ms = new_lrhs[0::4, 0::4, :]

#'''generating ms image with bicubic interpolation'''
#used_ms = cv2.resize(original_msi, (original_msi.shape[1]//4, original_msi.shape[0]//4), cv2.INTER_CUBIC)

'''generating pan image with gaussian kernel'''
used_pan = signal.convolve2d(original_pan, kernel, boundary='wrap',mode='same')
used_pan = np.expand_dims(used_pan, -1)
used_pan = used_pan[0::4, 0::4, :]

#'''generating pan image with vitual spectral kernel'''
#spectral_kernel = np.array([[0.1], [0.1], [0.4], [0.4]])
#used_pan = np.dot(original_msi, spectral_kernel)

#'''generating ms image with bicubic interpolation'''
#used_pan = cv2.resize(original_pan, (original_pan.shape[1]//4, original_pan.shape[0]//4), cv2.INTER_CUBIC)
#used_pan = np.expand_dims(used_pan, -1)

gt = np.uint8(255*original_msi)

print('ms shape: ', used_ms.shape, 'pan shape: ', used_pan.shape)

'''setting save parameters'''
save_images = True
save_channels = [0, 1, 2]#BGR-NIR for GF2
save_dir='./results/'
if save_images and (not os.path.isdir(save_dir)):
    os.makedirs(save_dir)

'''evaluating all methods'''
ref_results={}
ref_results.update({'metrics: ':'  PSNR,     SSIM,   SAM,    ERGAS,  SCC,    Q'})
no_ref_results={}
no_ref_results.update({'metrics: ':'  D_lamda, D_s,    QNR'})

'''Bicubic method'''
print('evaluating Bicubic method')
fused_image = Bicubic(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'Bicubic    ':temp_ref_results})
no_ref_results.update({'Bicubic    ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'Bicubic.tiff', fused_image[:, :, save_channels])

'''Brovey method'''
print('evaluating Brovey method')
fused_image = Brovey(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'Brovey     ':temp_ref_results})
no_ref_results.update({'Brovey     ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'Brovey.tiff', fused_image[:, :, save_channels])

'''PCA method'''
print('evaluating PCA method')
fused_image = PCA(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'PCA        ':temp_ref_results})
no_ref_results.update({'PCA        ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'PCA.tiff', fused_image[:, :, save_channels])

'''IHS method'''
print('evaluating IHS method')
fused_image = IHS(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'IHS        ':temp_ref_results})
no_ref_results.update({'IHS        ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'IHS.tiff', fused_image[:, :, save_channels])

'''SFIM method'''
print('evaluating SFIM method')
fused_image = SFIM(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'SFIM       ':temp_ref_results})
no_ref_results.update({'SFIM       ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'SFIM.tiff', fused_image[:, :, save_channels])

'''GS method'''
print('evaluating GS method')
fused_image = GS(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'GS         ':temp_ref_results})
no_ref_results.update({'GS         ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'GS.tiff', fused_image[:, :, save_channels])

'''Wavelet method'''
print('evaluating Wavelet method')
fused_image = Wavelet(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'Wavelet    ':temp_ref_results})
no_ref_results.update({'Wavelet    ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'Wavelet.tiff', fused_image[:, :, save_channels])

'''MTF_GLP method'''
print('evaluating MTF_GLP method')
fused_image = MTF_GLP(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'MTF_GLP    ':temp_ref_results})
no_ref_results.update({'MTF_GLP    ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'MTF_GLP.tiff', fused_image[:, :, save_channels])

'''MTF_GLP_HPM method'''
print('evaluating MTF_GLP_HPM method')
fused_image = MTF_GLP_HPM(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'MTF_GLP_HPM':temp_ref_results})
no_ref_results.update({'MTF_GLP_HPM':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'MTF_GLP_HPM.tiff', fused_image[:, :, save_channels])

'''GSA method'''
print('evaluating GSA method')
fused_image = GSA(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'GSA        ':temp_ref_results})
no_ref_results.update({'GSA        ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'GSA.tiff', fused_image[:, :, save_channels])

'''CNMF method'''
print('evaluating CNMF method')
fused_image = CNMF(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'CNMF       ':temp_ref_results})
no_ref_results.update({'CNMF       ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'CNMF.tiff', fused_image[:, :, save_channels])

'''GFPCA method'''
print('evaluating GFPCA method')
fused_image = GFPCA(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'GFPCA      ':temp_ref_results})
no_ref_results.update({'GFPCA      ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'GFPCA.tiff', fused_image[:, :, save_channels])

'''PNN method'''
print('evaluating PNN method')
fused_image = PNN(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'PNN        ':temp_ref_results})
no_ref_results.update({'PNN        ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'PNN.tiff', fused_image[:, :, save_channels])

'''PanNet method'''
print('evaluating PanNet method')
fused_image = PanNet(used_pan[:, :, :], used_ms[:, :, :])
temp_ref_results = ref_evaluate(fused_image, gt)
temp_no_ref_results = no_ref_evaluate(fused_image, np.uint8(used_pan*255), np.uint8(used_ms*255))
ref_results.update({'PanNet     ':temp_ref_results})
no_ref_results.update({'PanNet     ':temp_no_ref_results})
#save
if save_images:
    cv2.imwrite(save_dir+'PanNet.tiff', fused_image[:, :, save_channels])

''''print result'''
print('################## reference comparision #######################')
for index, i in enumerate(ref_results):
    if index == 0:
        print(i, ref_results[i])
    else:
        print(i, [round(j, 4) for j in ref_results[i]])
print('################## reference comparision #######################')

print('################## no reference comparision ####################')
for index, i in enumerate(no_ref_results):
    if index == 0:
        print(i, no_ref_results[i])
    else:
        print(i, [round(j, 4) for j in no_ref_results[i]])
print('################## no reference comparision ####################')

reason: import numpy,from scipy


text:
#!/usr/bin/env python3
# =============================================================================
# Developer : Shashank Sharma(shashankrnr32@gmail.com)
# License : MIT License
# Year : 2019
# =============================================================================
# =============================================================================
# Description:
#     Utilities File. Helper Script for other scripts
# =============================================================================
import scipy.io.wavfile as wav
import colorama
colorama.init(autoreset = True)

def error_log(msg):
    print(colorama.Fore.RED + msg)

def warning_log(msg):
    print(colorama.Fore.YELLOW + msg)

def WavRead(file_path = ''):
    try:
        fs,samples = wav.read(file_path)
        return [fs,samples]
    except:
        error_log('Cannot Read File : '+file_path)
        exit()

reason: import scipy


text:
#!/usr/bin/env python2.7
#
# This file is part of peakAnalysis, http://github.com/alexjgriffith/peaks/,
# and is Copyright (C) University of Ottawa, 2014. It is Licensed under
# the three-clause BSD License; see doc/LICENSE.txt.
# Contact: griffitaj@gmail.com
#
# Created : AUG262014
# File    : peakClass
# Author  : Alexander Griffith
# Lab     : Dr. Brand and Dr. Perkins

import sys
import os
import numpy
from peak_functions import *
from defineClass import define
from peakProcessingClass import peakProcessing
from UserList import UserList

class peakCapHandler(UserList):
    def __init__(self):
        self.data=[]
        self.blocked=[]
        self.blockedType=[]
        self.overlapValue=[]
    def add(self,peaks):
        for i in peaks:
            self.data.append(peakCap(i))
    def checkSort(self):
        for i in range(len(self.data)-1):
            if self.data[i].chro==self.data[i+1].chro:
                if self.data[i].max()>self.data[i+1].min():
                    print "value",i
                    return 1
            elif self.data[i].chro>self.data[i+1].chro:
                print "value chro",i,self.data[i].chro,self.data[i+1].chro
                return 1
        return 0

    def merge(self,dataA,dataB):
        temp=dataA.data+dataB.data
        return peakCap(temp)

    def split(self,data,distance):
        buffer=[data.data[0]]
        splitPeaks=[]
        for i in range(len(data.data)-1):
            if data.data[i+1].summit-data.data[i].summit<distance:
                buffer.append(data.data[i+1])
            else:
                splitPeaks.append(peakCap(buffer))
                buffer=[data.data[i+1]]
        splitPeaks.append(peakCap(buffer))
        return splitPeaks

    def overlap(self,limit):
        self.overlapValue=limit
        ran=float(limit)
        temp=[self.data[0]]
        k=0
        peakA=temp[0]
        for i in range(len(self.data)-1):
            peakB=self.data[i+1]
            if peakA.maxDiff()>=ran:
                t=self.split(peakA,ran)
                k+=len(t)-1
                temp.pop()
                temp.extend(t)
                peakA=temp[k]
            if peakA.chro==peakB.chro and peakB.min()-peakA.max()<ran:
                temp[k]=(self.merge(peakA,peakB))
                peakA=temp[k]
            else:
                temp.append(peakB)
                k=k+1
                peakA=temp[k]
        if peakA.maxDiff()>=ran:
            t=self.split(peakA,ran)
            temp.pop()
            temp.extend(t)
        self.data=temp

class peakCap():
    def __init__(self,data=None):
        self.data=[]
        self.summit=None
        self.define=define()
        self.maxum=0
        self.callSummit=lambda : [i.summit for i in self.data]
        if isinstance(data,list):
                for i in data:
                    self.define.add(i.define.data)
                    self.chro=i.chroOrder
                    self.data.append(i)
                self.maxum=self.maxDiff()
        else:
            self.define.add(data.define.data)
            self.chro=data.chroOrder
            self.data.append(data)
            self.maxum=0
        self.summit=np.mean(self.callSummit())

    def maxDiff(self):
        if self.maxum==0:
            temp=0
            for i in range(len(self.data)-1):
                if self.data[i+1].summit-self.data[i].summit>temp:
                    temp=self.data[i+1].summit-self.data[i].summit
            self.maxum=temp
            return temp
        else:
            return self.maxum
    def max(self):
        return max(self.callSummit())
    def min(self):
        return min(self.callSummit())

class peak():
    def __init__(self,chro,summit,defines):
        self.chro=chro
        self.summit=summit
        self.score=0
        self.define=define()
        self.define.add(defines)
    def addScore(self,score):
        self.score=float(score)
    def buildChroOrder(self,chromosomeList):
        self.chroOrder= np.where([self.chro==i for i in chromosomeList])[0].astype('int')[0]
    def __call__(self):
        return self.chro,self.summit

max_stars_repo_name: alexjgriffith/alpha-score

max_issues_repo_name: alexjgriffith/alpha-score

max_forks_repo_name: alexjgriffith/alpha-score


text:
"""End-effector control for bimanual Baxter robot.

This script shows how to use inverse kinematics solver from Bullet
to command the end-effectors of two arms of the Baxter robot.
"""

import os
import numpy as np

import robosuite
from robosuite.wrappers import IKWrapper

if __name__ == "__main__":

    # initialize a Baxter environment
    env = robosuite.make(
        "BaxterLift",
        ignore_done=True,
        has_renderer=True,
        gripper_visualization=True,
        use_camera_obs=False,
    )
    env = IKWrapper(env)

    obs = env.reset()

    # rotate the gripper so we can see it easily
    env.set_robot_joint_positions([
        0.00, -0.55, 0.00, 1.28, 0.00, 0.26, 0.00,
        0.00, -0.55, 0.00, 1.28, 0.00, 0.26, 0.00,
    ])

    bullet_data_path = os.path.join(robosuite.models.assets_root, "bullet_data")

    def robot_jpos_getter():
        return np.array(env._joint_positions)

    for t in range(100000):
        omega = 2 * np.pi / 1000.
        A = 5e-4
        dpos_right = np.array([A * np.cos(omega * t), 0, A * np.sin(omega * t)])
        dpos_left = np.array([A * np.sin(omega * t), A * np.cos(omega * t), 0])
        dquat = np.array([0, 0, 0, 1])
        grasp = 0.
        action = np.concatenate([dpos_right, dquat, dpos_left, dquat, [grasp, grasp]])

        obs, reward, done, info = env.step(action)
        env.render()

        if done:
            break

max_stars_repo_name: youngwoon/robosuite

max_issues_repo_name: youngwoon/robosuite

max_forks_repo_name: youngwoon/robosuite