process_sample.py

import io
import tempfile
import zipfile

import pycolmap

from example_solutions_copy import empty_solution, predict_wireframe


def read_colmap_rec(colmap_data):
    with tempfile.TemporaryDirectory() as tmpdir:
        with zipfile.ZipFile(io.BytesIO(colmap_data), "r") as zf:
            zf.extractall(tmpdir)  # unpacks cameras.txt, images.txt, etc. to tmpdir
        # Now parse with pycolmap
        rec = pycolmap.Reconstruction(tmpdir)
        return rec


def process_sample(sample, handle_error=True):
    try:
        pred_wireframe = predict_wireframe(sample)
    except Exception:
        if handle_error:
            pred_wireframe = empty_solution()
        else:
            raise
    return {
        "order_id": sample["order_id"],
        "wf_vertices": pred_wireframe.vertices_np,
        "wf_edges": pred_wireframe.edges_np,
    }