def create_body_tarball_from_sv_tarball(instance_info, body_id):
    """
    Download a supervoxel mesh tarball from the given key-value instance,
    concatenate together the component meshes into a single body tarball,
    and upload it.
    """
    keyEncodeLevel0 = 10000000000000
    keyEncodeLevel1 = 10100000000000
    
    encoded_sv = str(body_id + keyEncodeLevel0)
    sv_tarball_path = f'/tmp/{encoded_sv}.tar'
    
    logger.info(f'Fetching {encoded_sv}.tar')
    tarball_contents = fetch_key(*instance_info, f'{encoded_sv}.tar')
    with open(sv_tarball_path, 'wb') as f:
        f.write(tarball_contents)
    
    logger.info(f'Unpacking {encoded_sv}.tar')
    sv_dir = f'/tmp/{encoded_sv}'
    os.makedirs(sv_dir, exist_ok=True)
    os.chdir(sv_dir)
    subprocess.check_call(f'tar -xf {sv_tarball_path}', shell=True)

    encoded_body = str(body_id + keyEncodeLevel1)
    body_tarball_path = f'/tmp/{encoded_body}.tar'
    
    logger.info(f"Constructing {encoded_body}.drc")
    mesh = Mesh.from_directory(sv_dir)
    mesh.serialize(f'/tmp/{encoded_body}.drc')
    subprocess.check_call(f'tar -cf {body_tarball_path} /tmp/{encoded_body}.drc', shell=True)
    
    with open(body_tarball_path, 'rb') as f:
        logger.info(f'Posting {encoded_body}.tar')
        post_key(*instance_info, f'{encoded_body}.tar', f)
Exemple #2
0
def test_createmeshes_to_keyvalue(setup_createmeshes_config,
                                  disable_auto_retry):
    template_dir, config, _dvid_address, _repo_uuid, object_boxes, _object_sizes = setup_createmeshes_config

    kv_instance = 'test_createmeshes_to_keyvalue'
    config['output'] = {
        'keyvalue': {
            'instance': kv_instance,
            'create-if-necessary': True
        }
    }
    YAML().dump(config, open(f"{template_dir}/workflow.yaml", 'w'))

    _execution_dir, _workflow = launch_flow(template_dir, 1)

    server = config['input']['dvid']['server']
    uuid = config['input']['dvid']['uuid']
    for sv in [100, 200, 300]:
        mesh_bytes = fetch_key(server, uuid, kv_instance, f'{sv}.obj')
        mesh = Mesh.from_buffer(mesh_bytes, fmt='obj')
        assert np.allclose(mesh.vertices_zyx.min(axis=0), object_boxes[sv][0],
                           1)
        assert np.allclose(mesh.vertices_zyx.max(axis=0), object_boxes[sv][1],
                           1)
def autogen_points(input_seg,
                   count,
                   roi,
                   body,
                   tbars,
                   use_skeleton,
                   random_seed=None,
                   minimum_distance=0):
    """
    Generate a list of points within the input segmentation, based on the given criteria.
    See the main help text below for details.
    """
    if tbars and not body:
        sys.exit(
            "If you want to auto-generate tbar points, please specify a body.")

    if not tbars and not count:
        sys.exit(
            "You must supply a --count unless you are generating all tbars of a body."
        )

    if use_skeleton:
        if not body:
            sys.exit(
                "You must supply a body ID if you want to use a skeleton.")
        if tbars:
            sys.exit(
                "You can't select both tbar points and skeleton points.  Pick one or the other."
            )
        if not count and minimum_distance > 0:
            sys.exit(
                "You must supply a --count if you want skeleton point samples to respect the minimum distance."
            )
        if not count and not roi and minimum_distance == 0:
            logger.warning(
                "You are using all nodes of a skeleton without any ROI filter! Is that what you meant?"
            )

    rng = default_rng(random_seed)

    if tbars:
        logger.info(f"Fetching synapses for body {body}")
        syn_df = fetch_annotation_label(*input_seg[:2],
                                        'synapses',
                                        body,
                                        format='pandas')
        tbars = syn_df.query('kind == "PreSyn"')[[*'zyx']]

        if roi:
            logger.info(f"Filtering tbars for roi {roi}")
            determine_point_rois(*input_seg[:2], [roi], tbars)
            tbars = tbars.query('roi == @roi')[[*'zyx']]

        if minimum_distance:
            logger.info(
                f"Pruning close points from {len(tbars)} total tbar points")
            tbars = prune_close_pairs(tbars, minimum_distance, rng)
            logger.info(f"After pruning, {len(tbars)} tbars remain.")

        if count:
            count = min(count, len(tbars))
            logger.info(f"Sampling {count} tbars")
            choices = rng.choice(tbars.index, size=count, replace=False)
            tbars = tbars.loc[choices]

        logger.info(f"Returning {len(tbars)} tbar points")
        return tbars

    elif use_skeleton:
        assert body
        logger.info(f"Fetching skeleton for body {body}")
        skeleton_instance = f'{input_seg[2]}_skeletons'
        swc = fetch_key(*input_seg[:2], skeleton_instance, f'{body}_swc')
        skeleton_df = swc_to_dataframe(swc)
        skeleton_df['x'] = skeleton_df['x'].astype(int)
        skeleton_df['y'] = skeleton_df['y'].astype(int)
        skeleton_df['z'] = skeleton_df['z'].astype(int)

        if roi:
            logger.info(f"Filtering skeleton for roi {roi}")
            determine_point_rois(*input_seg[:2], [roi], skeleton_df)
            skeleton_df = skeleton_df.query('roi == @roi')[[*'zyx']]

        if minimum_distance:
            assert count
            # Distance-pruning is very expensive on a huge number of close points.
            # If skeleton is large, first reduce the workload by pre-selecting a
            # random sample of skeleton points, and prune more from there.
            if len(skeleton_df) > 10_000:
                # FIXME: random_state can't use rng until I upgrade to pandas 1.0
                skeleton_df = skeleton_df.sample(min(4 * count,
                                                     len(skeleton_df)),
                                                 random_state=None)
            logger.info(
                f"Pruning close points from {len(skeleton_df)} skeleton points"
            )
            prune_close_pairs(skeleton_df, minimum_distance, rng)
            logger.info(
                f"After pruning, {len(skeleton_df)} skeleton points remain.")

        if count:
            count = min(count, len(skeleton_df))
            logger.info(f"Sampling {count} skeleton points")
            choices = rng.choice(skeleton_df.index, size=count, replace=False)
            skeleton_df = skeleton_df.loc[choices]

        logger.info(f"Returning {len(skeleton_df)} skeleton points")
        return skeleton_df

    elif body:
        assert count
        if roi:
            # TODO: intersect the ranges with the ROI.
            raise NotImplementedError(
                "Sorry, I haven't yet implemented support for "
                "body+roi filtering.  Pick one or the other, "
                "or ask Stuart to fix this.")

        logger.info(f"Fetching sparsevol for body {body}")
        ranges = fetch_sparsevol(*input_seg, body, format='ranges')
        logger.info("Sampling from sparsevol")

        if minimum_distance > 0:
            # Sample 4x extra so we still have enough after pruning.
            points = sample_points_from_ranges(ranges, 4 * count, rng)
        else:
            points = sample_points_from_ranges(ranges, count, rng)

        points = pd.DataFrame(points, columns=[*'zyx'])

        if minimum_distance > 0:
            logger.info(f"Pruning close points from {len(points)} body points")
            prune_close_pairs(points, minimum_distance, rng)
            logger.info(f"After pruning, {len(points)} body points remain")

        points = points.iloc[:count]
        logger.info(f"Returning {len(points)} body points")
        return points

    elif roi:
        assert count
        logger.info(f"Fetching roi {roi}")
        roi_ranges = fetch_roi_roi(*input_seg[:2], roi, format='ranges')
        logger.info("Sampling from ranges")

        if minimum_distance > 0:
            # Sample 4x extra so we can prune some out if necessary.
            points_s5 = sample_points_from_ranges(roi_ranges, 4 * count, rng)
        else:
            points_s5 = sample_points_from_ranges(roi_ranges, count, rng)

        corners_s0 = points_s5 * (2**5)
        points_s0 = rng.integers(corners_s0, corners_s0 + (2**5))
        points = pd.DataFrame(points_s0, columns=[*'zyx'])

        if minimum_distance > 0:
            logger.info(f"Pruning close points from {len(points)} roi points")
            prune_close_pairs(points, minimum_distance, rng)
            logger.info(
                f"After pruning, points from {len(points)} roi points remain")

        points = points.iloc[:count]
        logger.info(f"Returning {len(points)} roi points")
        return points
    else:
        # No body or roi specified, just sample from the whole non-zero segmentation area
        assert count
        logger.info("Sampling random points from entire input segmentation")
        logger.info("Fetching low-res input volume")
        box_s6 = round_box(fetch_volume_box(*input_seg), 2**6, 'out') // 2**6
        seg_s6 = fetch_labelmap_voxels(*input_seg, box_s6, scale=6)
        mask_s6 = seg_s6.astype(bool)
        logger.info("Encoding segmentation as ranges")
        seg_ranges = runlength_encode_mask_to_ranges(mask_s6, box_s6)

        logger.info("Sampling from ranges")

        if minimum_distance > 0:
            # Sample 4x extra so we can prune some out if necessary.
            points_s6 = sample_points_from_ranges(seg_ranges, 4 * count, rng)
        else:
            points_s6 = sample_points_from_ranges(seg_ranges, count, rng)

        corners_s0 = points_s6 * (2**6)
        points_s0 = rng.integers(corners_s0, corners_s0 + (2**6))

        points = pd.DataFrame(points_s0, columns=[*'zyx'])

        if minimum_distance > 0:
            logger.info(
                f"Pruning close points from {len(points)} segmentation points")
            prune_close_pairs(points, minimum_distance, rng)
            logger.info(
                f"After pruning, points from {len(points)} segmentation points remain"
            )

        points = points.iloc[:count]
        logger.info(f"Returning {len(points)} segmentation points")
        return points