scripts/sort_blocks.py

"""RGB-guided SO-101 pick and place using calibrated camera and joint targets."""
import argparse
import json
from pathlib import Path
import urllib.request

import cv2
import numpy as np
from PIL import Image

from arm_kinematics import SO101Kinematics
from block_vision import detect_blocks, detections_from_mask, annotate
from calibrate_camera import command, server_url

ROOT = Path(__file__).resolve().parent.parent
RESULTS = ROOT / "results"


class Controller:
    def __init__(self):
        state = self.state()
        if state.get("scene", "lego") != "lego":
            raise RuntimeError("The block sorter requires the SO-101 block-sorting scene.")
        self.kin = SO101Kinematics(ROOT / "scene/scene.xml")
        calibration = json.loads((RESULTS / "camera-calibration.json").read_text())
        self.K = np.array(calibration["intrinsic"])
        self.R, _ = cv2.Rodrigues(np.array(calibration["rotation"]))
        self.t = np.array(calibration["translation"])
        self.center = -self.R.T @ self.t
        self.q = state["qpos"]
        self.grip = float(self.q[5])

    def state(self):
        with urllib.request.urlopen(server_url() + "/state", timeout=90) as response:
            return json.load(response)

    def pixel_xyz(self, pixel, height):
        ray = self.R.T @ np.linalg.solve(self.K, np.r_[pixel, 1.])
        return self.center + ray * ((height - self.center[2]) / ray[2])

    def observe(self, label):
        result = command("/camera", label=label)
        self.q = result["qpos"]
        image = np.asarray(Image.open(result["image"]).convert("RGB"))
        return image, result

    def move(self, target, seconds=1.5, label="move"):
        result = command("/move", target=np.asarray(target).tolist(), seconds=seconds, label=label)
        self.q = result["qpos"]
        if any(result["warning_counts"]):
            raise RuntimeError("Physics warning: " + str(result["warning_counts"]))
        return result

    def cartesian(self, xyz, seconds=2., label="cartesian"):
        self.q = self.state()["qpos"]
        start = self.kin.forward_grasp(self.q)
        xyz = np.asarray(xyz)
        count = max(5, int(np.ceil(np.linalg.norm(xyz - start) / .012)))
        # Solve the complete path before sending motor commands.
        targets = []
        seed = self.q
        for amount in np.linspace(0, 1, count + 1)[1:]:
            seed = self.kin.solve(start * (1 - amount) + xyz * amount, seed=seed, grip=self.grip, yaw=0)
            targets.append(seed)
        for index, target in enumerate(targets):
            result = self.move(target, seconds=seconds / count, label=f"{label}-{index}")
        return result

    def gripper(self, angle, label):
        self.grip = angle
        target = np.asarray(self.state()["ctrl"])
        target[5] = angle
        return self.move(target, seconds=1.4, label=label)

    def locate_box(self, rgb):
        hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
        h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
        mask = ((h >= 10) & (h <= 23) & (s > 80) & (v > 40)).astype(np.uint8) * 255
        detections = detections_from_mask(mask, "box", min_area=600)
        if not detections:
            raise RuntimeError("Box not visible in the front camera")
        detection = detections[0]
        return self.pixel_xyz(detection.centroid, .023)[:2], detection

    def pick_place(self, color, destination):
        rgb, observation = self.observe(f"before-{color}")
        detection = detect_blocks(rgb)[color]
        if detection is None:
            raise RuntimeError(f"{color} block not visible")
        # Known block height supplies the ray/plane intersection; XY comes
        # only from the live RGB image, never from simulator body positions.
        target = self.pixel_xyz(detection.centroid, .013)
        xy = target[:2]
        print(f"{color}: camera pixel={detection.centroid.round(2).tolist()}, XY={xy.round(5).tolist()}", flush=True)
        self.gripper(.6, f"{color}-open")
        self.cartesian([*xy, .062], seconds=2., label=f"{color}-approach")
        self.cartesian([*xy, .016], seconds=1.8, label=f"{color}-descend")
        closed = self.gripper(-.15, f"{color}-close")
        self.cartesian([*xy, .062], seconds=1.8, label=f"{color}-lift")
        lifted, lift_state = self.observe(f"{color}-lift-check")
        visible = detect_blocks(lifted)[color]
        if visible is None:
            raise RuntimeError(f"Cannot verify {color} grasp from the camera")
        # A held block must project near the expected lifted grasp point.
        grasp_xyz = self.kin.forward_grasp(self.q)
        lifted_center = grasp_xyz.copy(); lifted_center[2] -= .005
        prediction, _ = cv2.projectPoints(lifted_center.reshape(1, 3), cv2.Rodrigues(self.R)[0], self.t, self.K, None)
        distance = float(np.linalg.norm(visible.centroid - prediction.reshape(2)))
        if distance > 24:
            raise RuntimeError(f"Camera does not confirm lift for {color}: {distance:.1f}px discrepancy")
        print(f"{color}: camera confirms lifted block ({distance:.1f}px from predicted gripper location)", flush=True)
        self.cartesian([*destination, .062], seconds=3.2, label=f"{color}-carry")
        self.gripper(.6, f"{color}-release")
        self.cartesian([*destination, .070], seconds=.8, label=f"{color}-clear")
        print(f"{color}: released over box", flush=True)
        return {"color": color, "detected_xy": xy.tolist(), "destination": np.asarray(destination).tolist(),
                "source_image": observation["image"], "lift_image": lift_state["image"], "lift_pixel_error": distance}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--block", choices=("blue", "red", "green", "all", "remaining"), default="all")
    args = parser.parse_args()
    RESULTS.mkdir(parents=True, exist_ok=True)
    c = Controller()
    c.move([0, -.6, .4, .6, 0, .6], seconds=1.6, label="clear-camera-view")
    c.grip = .6
    rgb, observation = c.observe("scene-localization")
    box_xy, box_detection = c.locate_box(rgb)
    print("Box center from RGB:", box_xy.round(5).tolist(), flush=True)
    detections = detect_blocks(rgb)
    Image.fromarray(annotate(rgb, {**detections, "box": box_detection})).save(RESULTS / "localized-scene.png")
    offsets = {"blue": [-.020, -.023], "red": [.021, -.021], "green": [-.001, .026]}
    report_path = RESULTS / "sorting-progress.json"
    report = {"camera_input": observation["image"], "box_center_from_rgb": box_xy.tolist(), "placements": []}
    if args.block != "all" and report_path.exists():
        report = json.loads(report_path.read_text())
        box_xy = np.array(report["box_center_from_rgb"])
    if args.block == "remaining":
        completed = {item["color"] for item in report["placements"]}
        colors = [name for name in ("blue", "red", "green") if name not in completed]
    else:
        colors = ["blue", "red", "green"] if args.block == "all" else [args.block]
    if not colors:
        print("All blocks already recorded as placed.", flush=True)
        return
    report_path.write_text(json.dumps(report, indent=2))
    # Convert from the initial raised pose into a reachable vertical approach.
    first = detect_blocks(rgb)[colors[0]]
    if first is None:
        raise RuntimeError(f"{colors[0]} block not visible in the front camera")
    first_xy = c.pixel_xyz(first.centroid, .013)[:2]
    c.move(c.kin.solve([*first_xy, .062], seed=c.q, grip=.6), seconds=2.0, label="vertical-start")
    for color in colors:
        report["placements"].append(c.pick_place(color, box_xy + offsets[color]))
        report_path.write_text(json.dumps(report, indent=2))
    c.move([0, -.6, .4, .6, 0, .6], seconds=1.6, label="finished-arm-clear")
    final_rgb, final_state = c.observe("finished")
    Image.fromarray(final_rgb).save(RESULTS / "blocks-in-box.png")
    print("Finished commanded placements; final camera:", final_state["image"], flush=True)


if __name__ == "__main__":
    main()