コード例 #1
0
ファイル: reencode.py プロジェクト: HomerW/ShapeAssembly
def get_train_pc(prog, ns):
    full_pc = prog_to_pc(prog, ns // 2)
    partial_prog = Program()
    pc_list = [get_line_pc(full_pc, partial_prog, ns)]
    for i in range(len(prog['prog'])):
        partial_prog.execute(prog['prog'][i])
        pc_list.append(get_line_pc(full_pc, partial_prog, ns))
    train_pc = torch.cat(pc_list, 0)
    return train_pc
コード例 #2
0
ファイル: proc_gen.py プロジェクト: HomerW/ShapeAssembly
def main(ind):
    sa = ShapeAssembly()
    root_lines, has_mid = make_skel()

    prog_lines = ['Assembly Program_0 {']
    for b in root_lines:
        prog_lines.append('\t'+b)
    prog_lines.append('}')

    RP = Program()

    for l in root_lines:
        RP.execute(l)

    base_par = 'mid' if has_mid else 'top'

    base_lines = make_base_prog(
        RP.cuboids['base'],
        RP.cuboids[base_par]
    )

    for i in range(len(prog_lines)):
        prog_lines[i] = prog_lines[i].replace('base', 'Program_1')



    prog_lines += base_lines

    #for b in base_lines:
    #    prog_lines.append('\t'+b)
    #prog_lines.append('}')

    cube_count = -1
    switches = []
    for line in prog_lines:
        if 'Cuboid' in line:
            if not ('Program_' in line or "bbox" in line):
                switches.append((
                    f'cube{cube_count}', line.split()[0]
                ))
            if "bbox" in line:
                cube_count = -1

            cube_count += 1
    for a, b in switches:
        prog_lines = [line.replace(b,a) for line in prog_lines]

    hier_prog = make_hier_prog(prog_lines)
    verts, faces = hier_execute(hier_prog)
    if check_rooted(verts, faces) and check_stability(verts, faces):
        # writeObj(verts, faces, f'out_{ind}.obj')
        writeHierProg(hier_prog, f"random_hier_data/{ind}.txt")
        return True
    return False
コード例 #3
0
def lines_to_pc(lines):
    P = Program()
    for i, line in enumerate(lines):
        P.execute(line)
    verts, faces = P.getShapeGeo()
    for i in range(3):
        verts[:, i] = verts[:, i] - verts[:, i].mean()
    pc = utils.sample_surface(faces,
                              verts.unsqueeze(0),
                              num_samps,
                              return_normals=False)[0]
    return pc
コード例 #4
0
def rand_program(prog, max_cuboids, bbox_dims, hier_index):
    size_max = max(bbox_dims)
    num_cuboids = int(samp((max_cuboids / 1.5), 1, 3, max_cuboids))
    num_source = int(samp(1.75, 0.5, 1, min(max_sources, num_cuboids)))
    cdim_mean = np.mean(bbox_dims) / 2
    cdim_std = cdim_mean / 2.5
    leaf_prob = 0.5**(hier_index + 1)

    print(f"NUM SOURCE: {num_source}")
    print(f"NUM CUBOIDS: {num_cuboids}")

    def attach_source(name, P):
        counter = 0
        while counter < 100:
            old_prog = deepcopy(P)
            aligned = random.random() < align_prob
            cuboid_line, dims = make_cuboid(name, cdim_mean, cdim_std,
                                            size_max, aligned)
            # check if sqeezing into bbox is a good fit
            good_squeeze = ((dims[1] / bbox_dims[1]) > 0.9)

            if random.random() < squeeze_prob and (not aligned
                                                   or good_squeeze):
                attach_line = make_squeeze(name, "bbox", "bot")
            else:
                # bottom of bbox and bottom of (unattached) source cuboid
                attach_line = make_attach(name, "bbox", "bot", "bot")
            lines = [cuboid_line, attach_line]
            P.execute(cuboid_line)
            P.execute(attach_line)
            if valid(P):
                break
            P = old_prog
            counter += 1
        if counter >= 100:
            print("GAVE UP ON SOURCE")
            return P, [], None
        return P, lines, dims

    def extend_cuboid(src_cuboid, cuboids, P):
        faces = ['right', 'left', 'top', 'bot', 'front', 'back']
        lines = []
        new_cuboids = []
        next_q = []

        num_new_cuboids = int(samp(1, 0.5, 1, max_out))

        for _ in range(num_new_cuboids):
            if new_cuboids == []:
                index = max([x["id"] for x in cuboids]) + 1
            else:
                index = max([x["id"] for x in new_cuboids]) + 1
            name = f"cube{index}"

            # make initial attachment to src cuboid
            if (len(cuboids) + len(new_cuboids)) < (num_cuboids - 1):
                counter = 0
                while counter < 100:
                    old_prog = deepcopy(P)
                    old_new_cuboids = deepcopy(new_cuboids)
                    old_next_q = deepcopy(next_q)

                    aligned = random.random() < align_prob
                    cuboid_line, dims = make_cuboid(name, cdim_mean, cdim_std,
                                                    size_max, aligned)
                    nc = {
                        "name": f"cube{index}",
                        "id": index,
                        "ancestor": src_cuboid["name"],
                        "dims": dims
                    }
                    new_cuboids.append(nc)

                    attach_lines = []
                    if random.random() < squeeze_prob:
                        # faces = ['top', 'bot']
                        # face = random.choice(faces)
                        attach_line = make_squeeze(name, src_cuboid["name"],
                                                   'top')
                    else:
                        next_q.append(nc)
                        attach_line = make_attach(name, src_cuboid["name"],
                                                  random.choice(faces),
                                                  random.choice(faces))

                    P.execute(cuboid_line)
                    P.execute(attach_line)
                    if valid(P):
                        lines += [cuboid_line, attach_line]
                        break
                    P = old_prog
                    new_cuboids = old_new_cuboids
                    next_q = old_next_q
                    counter += 1
                if counter >= 100:
                    print("GAVE UP ON EXTENSION")
                    return [], [], P, []

                # if cuboid is not aligned potentially add some more attachments
                counter = 0
                while counter < 100 and not aligned:
                    old_prog = deepcopy(P)
                    num_extra_attaches = int(samp(1, 0.5, 0, 3))
                    print(f"EXTRA ATTACHES {num_extra_attaches}")

                    attach_lines = []
                    for _ in range(num_extra_attaches):
                        possible_cuboids = [
                            c["name"] for c in cuboids
                            if not c["name"] == c["ancestor"]
                        ]
                        attach_cuboid = random.choice(possible_cuboids)
                        attach_lines.append(
                            make_attach(name, attach_cuboid,
                                        random.choice(faces),
                                        random.choice(faces)))
                    for l in attach_lines:
                        P.execute(l)
                    if valid(P):
                        lines += attach_lines
                        break
                    P = old_prog
                    counter += 1
                    if counter >= 100:
                        print("GAVE UP ON EXTENSION II")

        return new_cuboids, next_q, P, lines

    P = Program()
    bbox_line = f"bbox = Cuboid({bbox_dims[0]}, {bbox_dims[1]}, {bbox_dims[2]}, True)"
    print(bbox_line)
    P.execute(bbox_line)
    lines = [bbox_line]
    q = []
    src_count = 0
    for _ in range(num_source):
        P, new_lines, dims = attach_source(f"cube{src_count}", P)
        if len(new_lines) > 0:
            lines += new_lines
            for l in new_lines:
                print(l)
            q.append({
                "name": f"cube{src_count}",
                "id": src_count,
                "ancestor": "bbox",
                "dims": dims
            })
            src_count += 1
    if len(q) == 0:
        print("COULDNT FIT ANY CUBOIDS")
        return None
    cuboids = deepcopy(q)
    while len(q) > 0 and len(cuboids) < (num_cuboids - 1):
        c = q.pop(0)
        new_cuboids, next_q, P, new_lines = extend_cuboid(c, cuboids, P)
        lines += new_lines
        for l in new_lines:
            print(l)
        q += next_q
        cuboids += new_cuboids

    # add some symmetry macros
    num_sym = int(samp(1, 0.5, 0, 3))
    for _ in range(num_sym):
        counter = 0
        while counter < 100:
            old_prog = deepcopy(P)
            sym_cuboid = random.choice([c["name"] for c in cuboids])
            if random.random() < reflect_prob:
                new_line = make_reflect(sym_cuboid)
            else:
                new_line = make_translate(sym_cuboid)
            P.execute(new_line)
            if valid(P):
                lines.append(new_line)
                print(new_line)
                break
            P = old_prog
            counter += 1
            if counter >= 100:
                print("GAVE UP ON MACRO")

    # correct dimensions since cuboids might have been scaled during execution
    for c in cuboids:
        new_dims = [round(x, 3) for x in P.cuboids[c['name']].dims.tolist()]
        c['dims'] = new_dims
        new_lines = []
        for l in lines:
            if (c['name'] in l) and ("Cuboid" in l):
                aligned = P.parseCuboid(l)[-1]
                new_lines.append(
                    f"{c['name']} = Cuboid({new_dims[0]}, {new_dims[1]}, {new_dims[2]}, {aligned})"
                )
            else:
                new_lines.append(l)
        lines = new_lines

    # choose from the largest cuboids to expand
    non_bbox_cuboids = [
        x for x in cuboids
        if not x['name'] == "bbox" and np.prod(x["dims"]) > 0.02
    ]
    sorted_cuboids = sorted(non_bbox_cuboids,
                            key=lambda x: -np.prod(x["dims"]))
    num_sub = len([
        _ for _ in range(len(non_bbox_cuboids)) if random.random() < leaf_prob
    ])
    # num_sub = len(sorted_cuboids)
    sub_cuboids = sorted_cuboids[:num_sub]

    # prog['prog'] = canonical(lines)
    prog['prog'] = lines
    next_q = []
    # start of with bbox child
    children = [{}]
    for c in cuboids:
        if c in sub_cuboids:
            cprog = {"prog": None, "children": None}
            children.append(cprog)
            next_q.append((cprog, hier_index + 1, c["dims"]))
        else:
            children.append({})
    prog['children'] = children

    return next_q
コード例 #5
0
def semValidGen(pc, prog, encoder, decoder, max_lines, input_dim, device, gt_prog, rejection_sample):
    q = []

    prog_out = []
    preds = []
    children = []
    out = torch.zeros((1, 1, input_dim), dtype = torch.float).to(device)
    out[0][0][0] = 1.0
    bb_dims = prog["bb_dims"]
    gt_nc_ind = 0
    gt_children = []

    if gt_prog is not None and len(gt_prog) > 0:
        gt_children = gt_prog["children"]

    P = Program()
    P.mode = 'start'
    P.grounded = set([0])
    P.atts = {}
    c = 0

    stop = False

    loops = 0

    num_rejects = 0

    h = decoder.init_gru
    full_pc = pc[0, :, :(num_samps // 2)]

    while(not stop and loops < max_lines):
        loops += 1

        prev_out = out.clone().detach()
        reencoding = reencode_prog(full_pc, P, num_samps, encoder)

        if loops == 1:
            bb_pred = decoder.bbdimNet(reencoding.squeeze())
            prog["bb_dims"] = bb_pred
            bb_dims = bb_pred

        out, h, valid = clean_forward(
            decoder,
            out,
            reencoding,
            h,
            bb_dims,
            input_dim,
            device,
            P
        )

        if not valid:
            if rejection_sample and DO_REJECT:
                assert False, "Couldn't clean line"
            num_rejects += 1
            if MASK_BAD_OUT and num_rejects < MAX_REJECT:
                out = prev_out
            else:
                num_rejects = 0

            continue

        prog_out.append(out)
        line = out.clone().detach().squeeze()
        preds.append(line)

        command = torch.argmax(line[:7])

        pline = None

        if command == 1:
            P.mode = 'cuboid'
            pline = getCuboidLine(line, c)
            c += 1

        elif command == 2:
            P.mode = 'attach'
            cub1 = torch.argmax(line[7:18]).item()
            cub2 = torch.argmax(line[18:29]).item()

            P.grounded.add(cub1)

            if cub2 in P.atts:
                P.atts[cub2].append(cub1)
            else:
                P.atts[cub2] = [cub1]

            if cub1 in P.atts:
                P.atts[cub1].append(cub2)
            else:
                P.atts[cub1] = [cub2]

            pline = getAttachLines(line)

        elif command == 3:
            P.mode = 'sym'
            pline = getReflectLine(line)

        elif command == 4:
            P.mode = 'sym'
            pline = getTranslateLine(line)

        elif command == 5:
            P.mode = 'attach'
            cub1 = torch.argmax(line[7:18]).item()
            cub2 = torch.argmax(line[18:29]).item()
            cub3 = torch.argmax(line[29:40]).item()

            P.grounded.add(cub1)

            if cub2 in P.atts:
                P.atts[cub2].append(cub1)
            else:
                P.atts[cub2] = [cub1]

            if cub3 in P.atts:
                P.atts[cub3].append(cub1)
            else:
                P.atts[cub3] = [cub1]

            if cub1 in P.atts:
                P.atts[cub1].append(cub2)
                P.atts[cub1].append(cub3)
            else:
                P.atts[cub1] = [cub2, cub3]

            pline = getSqueezeLine(line)

        try:
            if pline is not None:
                P.execute(pline)

        except Exception:
            if VERBOSE:
                print("Unexpectedly, failed to execute line")
            pass

        # Stop at end token or when we have gone past max lines

        if command == 6:
            stop = True

        # If make a new Cuboid, use l to decide if it should have a child program or be a leaf
        if command == 1:
            children.append({})

    fc_preds, fc_prog_out, fchildren = cuboid_line_clean(preds, prog_out, children, P)
    prog["children"] = fchildren
    return fc_preds, fc_prog_out
コード例 #6
0
def rand_program(prog, max_cuboids, bbox_dims, hier_index):
    size_max = max(bbox_dims)
    num_cuboids = int(samp((max_cuboids / 3), 1, 2, max_cuboids))
    num_source = int(samp(1.5, 1, 1, min(max_sources, num_cuboids)))
    cdim_mean = np.mean(bbox_dims) / 2
    cdim_std = np.std(bbox_dims)
    leaf_prob = 0.5**(hier_index + 1)

    print(f"NUM SOURCE: {num_source}")
    print(f"NUM CUBOIDS: {num_cuboids}")

    def attach_source(name, P):
        counter = 0
        while counter < 100:
            old_prog = deepcopy(P)
            cuboid_line = make_cuboid(name, cdim_mean, cdim_std, size_max,
                                      align_prob)
            # bottom of bbox and bottom of (unattached) source cuboid
            attach_line = make_attach(name, "bbox", "bot", "bot")
            lines = [cuboid_line, attach_line]
            P.execute(cuboid_line)
            P.execute(attach_line)
            if valid(P):
                break
            P = old_prog
            counter += 1
        if counter >= 100:
            print("GAVE UP ON SOURCE")
            return P, []
        return P, lines

    def create_edge(src_cuboid, cuboids, P, new_src_cuboids, prev_attach,
                    new_cuboid):
        lines = []
        faces = ['right', 'left', 'top', 'bot', 'front', 'back']

        new_attach = None
        possible_attachments = []
        for c in cuboids:
            # don't attach to ancestors, sources, or cuboids already attached to in this cycle
            if not (c["id"] in src_cuboid["ancestors"] or c["level"] == 1
                    or c["id"] in prev_attach):
                if c['name'] == 'bbox' and src_cuboid['level'] == 1:
                    continue
                possible_attachments.append(c)
        if new_cuboid or possible_attachments == []:
            if len(P.cuboids) >= num_cuboids:
                print("TOO MANY CUBOIDS")
                return P, new_src_cuboids, [], new_attach
            if new_src_cuboids == []:
                index = max([x["id"] for x in cuboids]) + 1
            else:
                index = max([x["id"] for x in new_src_cuboids]) + 1
            new_src_cuboids.append({
                "name":
                f"cube{index}",
                "id":
                index,
                "ancestors":
                src_cuboid["ancestors"] | set([index]),
                "level":
                src_cuboid["level"] + 1
            })
            cuboid_line = make_cuboid(f"cube{index}", cdim_mean, cdim_std,
                                      size_max, align_prob)
            lines = [cuboid_line]
            P.execute(cuboid_line)
            c_new = f"cube{index}"

            face1 = faces[random.choice(range(6))]
            face2 = faces[random.choice(range(6))]
            attach_line = make_attach(c_new, src_cuboid["name"], face1, face2)
        else:
            # pos_levels = [c['level'] for c in possible_attachments]
            # p = [-1*(x - max(pos_levels))+1.0 for x in pos_levels]
            # p = p / np.sum(p)
            selected_cuboid = np.random.choice(possible_attachments)
            # selected_cuboid['ancestors'] = selected_cuboid['ancestors'] | src_cuboid['ancestors']
            src_cuboid['ancestors'] = selected_cuboid[
                'ancestors'] | src_cuboid['ancestors']
            c_new = selected_cuboid["name"]
            new_attach = selected_cuboid["id"]

            if not c_new == "bbox":
                face1 = faces[random.choice(range(6))]
                face2 = faces[random.choice(range(6))]
                attach_line = make_attach(src_cuboid["name"], c_new, face1,
                                          face2)
            # if bbox only attach to top face and reverse direction of edge
            else:
                attach_line = make_attach(src_cuboid["name"], c_new, "top",
                                          "top")

        lines.append(attach_line)
        P.execute(attach_line)

        return P, new_src_cuboids, lines, new_attach

    def add_attachment(src_cuboid, cuboids, P, num_sources):
        p = np.full((max_out + 1, ), 1.0)
        p[0] = p[0] / 5
        p[3] = p[3] / 5
        p[2] = p[2] / 2
        p = p / np.sum(p)
        out = np.random.choice(list(range(max_out + 1)), p=p)
        print(f"OUT: {out}")

        new_src_cuboids = []
        lines = []
        prev_attach = []

        for edge in range(out):
            edge_counter = 0
            new_cuboid = random.random() < new_cuboid_prob
            if len(P.cuboids) == num_cuboids:
                new_cuboid = False
            while edge_counter < 100:
                old_prog = deepcopy(P)
                old_new_src_cuboids = deepcopy(new_src_cuboids)

                P, new_src_cuboids, new_lines, new_attach = create_edge(src_cuboid, cuboids, \
                                                                                 P, new_src_cuboids, prev_attach, new_cuboid)
                if valid(P):
                    if not new_cuboid:
                        print("MADE ATTACHMENT WITH PREV CUBOID")
                    else:
                        print("MADE NEW CUBOID")
                    prev_attach.append(new_attach)
                    lines += new_lines
                    break

                P = old_prog
                new_src_cuboids = old_new_src_cuboids
                edge_counter += 1

            if not out == 0 and edge_counter >= 100:
                print("GAVE UP ON EDGE")

        return new_src_cuboids, P, lines

    P = Program()
    bbox_line = f"bbox = Cuboid({bbox_dims[0]}, {bbox_dims[1]}, {bbox_dims[2]}, True)"
    P.execute(bbox_line)
    lines = [bbox_line]
    current_cuboids = []
    src_count = 0
    for _ in range(num_source):
        P, new_lines = attach_source(f"cube{src_count}", P)
        if len(new_lines) > 0:
            lines += new_lines
            current_cuboids.append({
                "name": f"cube{src_count}",
                "id": src_count,
                "ancestors": set([src_count]),
                "level": 1
            })
            src_count += 1
    if len(current_cuboids) <= 1:
        print("COULDNT FIT MORE THAN ONE CUBOID")
        return None
    cuboids = [{
        "name": "bbox",
        "id": -1,
        "ancestors": set([-1]),
        "level": 0
    }] + current_cuboids
    while len(current_cuboids) > 0:
        c = current_cuboids.pop(0)
        new_cuboids, P, new_lines = add_attachment(c, cuboids, P, num_source)
        lines += new_lines
        current_cuboids += new_cuboids
        cuboids += new_cuboids

    # choose from the largest cuboids to expand
    non_bbox_cuboids = [
        x for x in P.cuboids
        if not x == "bbox" and torch.prod(P.cuboids[x].dims) > 0.02
    ]
    sorted_cuboids = sorted(non_bbox_cuboids,
                            key=lambda x: -torch.prod(P.cuboids[x].dims))
    # num_sub = len([_ for _ in range(len(non_bbox_cuboids)) if random.random() < leaf_prob])
    num_sub = len(sorted_cuboids)
    sub_cuboids = sorted_cuboids[:num_sub]
    print(sub_cuboids)

    prog['prog'] = canonical(lines)
    next_q = []
    children = []
    for c in P.cuboids:
        if c in sub_cuboids:
            cprog = {"prog": None, "children": None}
            children.append(cprog)
            next_q.append((cprog, hier_index + 1, P.cuboids[c].dims.tolist()))
        else:
            children.append({})
    prog['children'] = children

    return next_q
コード例 #7
0
def semValidGen(prog, rnn, h, hier_ind, max_lines, input_dim, device, gt_prog,
                rejection_sample):
    q = []

    prog_out = []
    preds = []
    children = []
    out = torch.zeros((1, 1, input_dim), dtype=torch.float).to(device)
    out[0][0][0] = 1.0
    h_start = h.clone()
    bb_dims = prog["bb_dims"]
    gt_nc_ind = 0
    gt_children = []

    if gt_prog is not None and len(gt_prog) > 0:
        gt_children = gt_prog["children"]

    P = Program()
    P.mode = 'start'
    P.grounded = set([0])
    P.atts = {}
    c = 0

    stop = False

    loops = 0

    num_rejects = 0

    while (not stop and loops < max_lines):
        loops += 1

        prev_out = out.clone().detach()

        out, pnext, pleaf, h, valid = clean_forward(rnn, out, h, h_start,
                                                    bb_dims, hier_ind,
                                                    input_dim, device, P)

        if not valid:
            if rejection_sample and DO_REJECT:
                assert False, "Couldn't clean line"
            num_rejects += 1
            if MASK_BAD_OUT and num_rejects < MAX_REJECT:
                out = prev_out
            else:
                num_rejects = 0

            continue

        prog_out.append(out)
        line = out.clone().detach().squeeze()
        preds.append(line)

        command = torch.argmax(line[:7])

        pline = None

        if command == 1:
            P.mode = 'cuboid'
            pline = getCuboidLine(line, c)
            c += 1

        elif command == 2:
            P.mode = 'attach'
            cub1 = torch.argmax(line[7:18]).item()
            cub2 = torch.argmax(line[18:29]).item()

            P.grounded.add(cub1)

            if cub2 in P.atts:
                P.atts[cub2].append(cub1)
            else:
                P.atts[cub2] = [cub1]

            if cub1 in P.atts:
                P.atts[cub1].append(cub2)
            else:
                P.atts[cub1] = [cub2]

            pline = getAttachLines(line)

        elif command == 3:
            P.mode = 'sym'
            pline = getReflectLine(line)

        elif command == 4:
            P.mode = 'sym'
            pline = getTranslateLine(line)

        elif command == 5:
            P.mode = 'attach'
            cub1 = torch.argmax(line[7:18]).item()
            cub2 = torch.argmax(line[18:29]).item()
            cub3 = torch.argmax(line[29:40]).item()

            P.grounded.add(cub1)

            if cub2 in P.atts:
                P.atts[cub2].append(cub1)
            else:
                P.atts[cub2] = [cub1]

            if cub3 in P.atts:
                P.atts[cub3].append(cub1)
            else:
                P.atts[cub3] = [cub1]

            if cub1 in P.atts:
                P.atts[cub1].append(cub2)
                P.atts[cub1].append(cub3)
            else:
                P.atts[cub1] = [cub2, cub3]

            pline = getSqueezeLine(line)

        try:
            if pline is not None:
                P.execute(pline)

        except Exception:
            if VERBOSE:
                print("Unexpectedly, failed to execute line")
            pass

        # Stop at end token or when we have gone past max lines

        if command == 6:
            stop = True

        # If make a new Cuboid, use l to decide if it should have a child program or be a leaf
        if command == 1:

            gt_child = None

            if gt_nc_ind < len(gt_children):
                gt_child = gt_children[gt_nc_ind]
                gt_nc_ind += 1

            # Skip BBox line
            if pleaf.squeeze().item() < 0 and len(preds) > 1:
                d = {"children": [], "bb_dims": line[40:43]}
                children.append(d)
                q.append((pnext, d, hier_ind + 1, gt_child))

            else:
                children.append({})

    fc_preds, fc_prog_out, fchildren = cuboid_line_clean(
        preds, prog_out, children, P)
    prog["children"] = fchildren
    return fc_preds, fc_prog_out, q