<jupyter_start><jupyter_text>Digging loop<jupyter_code>n = 0  # number of parts

curr_box_grid = box_grid.copy()
box_grid.sum()
%%time

%matplotlib inline

def objective_f(xs, box_grid, part, voxel_size, grid_pos, grid_shape):
    if xs.size > 0:
        return np.apply_along_axis(eval_position, -1, xs, box_grid=box_grid, part=part, voxel_size=voxel_size, grid_pos=grid_pos, grid_shape=grid_shape)
    else:
        return np.apply_along_axis(eval_position, -1, np.array([0.,0.,0.,0.,0.,0.]), box_grid=box_grid, part=part, voxel_size=voxel_size, grid_pos=grid_pos, grid_shape=grid_shape)

ax = plt.subplot()
while True:
    optimizer = ps.single.LocalBestPSO(
        n_particles=18,
        dimensions=6,
        bounds=bounds,
        options={'c1': 0.5, 'c2': 0.2, 'w':0.9, 'k': 2, 'p': 2},
    )

    cost, pos = optimizer.optimize(objective_f, iters=500, n_processes=6, box_grid=curr_box_grid, part=_PointCloudTransmissionFormat(part), voxel_size=voxel_size, grid_pos=grid_pos, grid_shape=grid_shape)

    ax = plot_cost_history(cost_history=optimizer.cost_history, ax=ax, label=n+1)
    
    # get part at optimal position
    tx, ty, tz, rx, ry, rz = pos

    part_ = o3d.geometry.PointCloud(part)

    part_.translate(np.array([tx, ty, tz]))

    part_.rotate(
        o3d.geometry.get_rotation_matrix_from_xyz(np.array([rx, ry, rz]))
    )

    part_voxel = o3d.geometry.VoxelGrid.create_from_point_cloud(part_, voxel_size)
    part_grid = voxel2np(part_voxel, grid_pos, grid_shape, mode='floating')

    # find out how much of the part is immersed at the box content
    immersion_ratio = (part_grid & curr_box_grid).sum() / part_grid.sum()

    if immersion_ratio > 0.60:
        curr_box_grid = dig_box(curr_box_grid, part, grid_pos, grid_shape, pos, voxel_size)
        n += 1
    else:
        print('no more parts!')
        print('parts digged: ',n)
        break

plt.show()
%matplotlib widget

print(curr_box_grid.sum())

ax = plt.figure().add_subplot(projection='3d')
ax.voxels(curr_box_grid, facecolors='gray')

plt.show()<jupyter_output>75881<jupyter_text>Refactored<jupyter_code>dig_and_predict(box, part, 0.01)
import json

preds_fpath = project_dir/'models/pso_train_preds.json'

if not preds_fpath.exists():
    preds = dict()
else:
    with open(preds_fpath, 'r') as f:
        preds = json.load(f)

preds
from tqdm.notebook import tqdm

for img_fpath in tqdm(list(raw_data_dir.glob('*/*.exr'))):
    ## LOADING
    rgbd = load_rgbd(img_fpath)

    box_mask = box_mask_from_rgbd(rgbd)
    
    vol = mask_selection_volume(rgbd, box_mask)

    pcd = load_pcd(rgbd)

    box = vol.crop_point_cloud(pcd)

    part_fpath = project_dir/'data/raw/part.stl'

    part_mesh = o3d.io.read_triangle_mesh(str(part_fpath), enable_post_processing=True)

    part_mesh.paint_uniform_color([1., 0., 0.,])

    part = part_mesh.sample_points_uniformly(number_of_points=10000)

    part_points = np.array(part.points) / 1000  # mm to meter conversion
    part_points = part_points + np.array([0,0,0.3])
    part_points = o3d.utility.Vector3dVector(part_points)
    part.points = part_points
    
    # PREDICTING
    n = dig_and_predict(box, part, 0.02)
    
    # STORING
    preds[img_fpath.name] = n
    
    with open(preds_fpath, 'w') as f:
        json.dump(preds, f)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12673506336753168, 'fraction_numerical': 0.020217260108630054, 'mean_word_length': 3.9330357142857144, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '73389', 'license': 'CNRI-Python', 'n_tokens_mistral': 1363, 'n_tokens_neox': 1292, 'n_words': 227, 'path': 'notebooks/03.02-bmp-digging-through-pso.ipynb', 'repo_name': 'brunompacheco/part-counting', 'size': '3314'}