| """Shared pieces for the RL policy on the real SO101: robot/camera setup, blue-candy detection, |
| and the phone-pixel -> table homography.""" |
|
|
| import json |
| import os |
| import time |
|
|
| import cv2 |
| import numpy as np |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| CALIB = os.path.join(HERE, "table_calib.json") |
| FOLLOWER_PORT = "/dev/serial/by-id/usb-1a86_USB_Single_Serial_5AE6057813-if00" |
| FOLLOWER_ID = "my_follower_arm" |
| PHONE_URL = os.environ.get("PHONE_URL", "http://127.0.0.1:4747/video") |
| IMG_W, IMG_H = 640, 360 |
| JOINTS = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"] |
|
|
|
|
| def connect_robot(): |
| from lerobot.robots.so_follower.config_so_follower import SOFollowerRobotConfig |
| from lerobot.robots.so_follower.so_follower import SOFollower |
|
|
| robot = SOFollower(SOFollowerRobotConfig(port=FOLLOWER_PORT, id=FOLLOWER_ID)) |
| robot.connect(calibrate=False) |
| return robot |
|
|
|
|
| def read_joints(robot): |
| obs = robot.get_observation() |
| return np.array([obs[f"{n}.pos"] for n in JOINTS], dtype=np.float64) |
|
|
|
|
| class Phone: |
| """DroidCam stream, resized to 640x360 RGB. |
| |
| A background thread reads the stream continuously and keeps only the newest frame, so a photo |
| is always current (OpenCV otherwise hands back frames buffered seconds earlier, which made the |
| "before" and "after" photos identical). Reconnects if the HTTP stream drops.""" |
|
|
| def __init__(self, url=PHONE_URL): |
| import threading |
|
|
| self.url = url |
| self.cap = None |
| self.latest, self.stamp, self.count = None, 0.0, 0 |
| self.lock = threading.Lock() |
| self.running = True |
| self._open() |
| self.thread = threading.Thread(target=self._reader, daemon=True) |
| self.thread.start() |
| t0 = time.time() |
| while self.latest is None and time.time() - t0 < 10: |
| time.sleep(0.05) |
| if self.latest is None: |
| raise RuntimeError(f"no frames from {self.url} within 10 s") |
|
|
| def _open(self): |
| if self.cap is not None: |
| self.cap.release() |
| self.cap = cv2.VideoCapture(self.url) |
| if not self.cap.isOpened(): |
| raise RuntimeError(f"cannot open phone stream {self.url} (is DroidCam running? close any browser tab showing it)") |
|
|
| def _reader(self): |
| fails = 0 |
| while self.running: |
| ok, img = self.cap.read() |
| if ok and img is not None: |
| fails = 0 |
| with self.lock: |
| self.latest, self.stamp = img, time.time() |
| self.count += 1 |
| continue |
| fails += 1 |
| print(f" (phone stream dropped, reconnecting: attempt {fails})") |
| time.sleep(1.0) |
| try: |
| self._open() |
| except RuntimeError as e: |
| print(" ", e) |
|
|
| def frame(self, max_age=0.5, timeout=10.0): |
| """Newest frame (taken after this call started), 640x360 RGB.""" |
| t_call = time.time() |
| while time.time() - t_call < timeout: |
| with self.lock: |
| img, stamp = self.latest, self.stamp |
| if img is not None and stamp > t_call and time.time() - stamp < max_age: |
| return cv2.cvtColor(cv2.resize(img, (IMG_W, IMG_H), interpolation=cv2.INTER_AREA), cv2.COLOR_BGR2RGB) |
| time.sleep(0.02) |
| raise RuntimeError("no fresh phone frame for %.0f s; check the DroidCam app and connection" % timeout) |
|
|
| def close(self): |
| self.running = False |
| self.thread.join(timeout=2) |
| self.cap.release() |
|
|
|
|
| def blue_mask(rgb): |
| hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV) |
| m = cv2.inRange(hsv, (95, 90, 70), (125, 255, 255)) |
| return cv2.morphologyEx(m, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8)) |
|
|
|
|
| def blue_report(rgb): |
| """Why a photo has (no) blue blobs: blue pixel count and the largest blue component.""" |
| m = blue_mask(rgb) |
| n, _, stats, cent = cv2.connectedComponentsWithStats(m) |
| big = max(range(1, n), key=lambda i: stats[i, cv2.CC_STAT_AREA], default=None) |
| hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV) |
| return {"blue_px": int((m > 0).sum()), "largest_area": int(stats[big, cv2.CC_STAT_AREA]) if big else 0, |
| "largest_at": cent[big].round(0).tolist() if big else None, "mean_brightness": float(hsv[..., 2].mean())} |
|
|
|
|
| def blue_blobs(rgb, min_px=25, ignore=None): |
| """Blue wrapper blobs: list of dicts with pixel centroid, area and the blob's long-axis angle. |
| ``ignore``: a reference image; blue that was already there (dilated a little) is left out.""" |
| m = blue_mask(rgb) |
| if ignore is not None: |
| m = cv2.bitwise_and(m, cv2.bitwise_not(cv2.dilate(blue_mask(ignore), np.ones((9, 9), np.uint8)))) |
| n, lab, stats, cent = cv2.connectedComponentsWithStats(m) |
| out = [] |
| for i in range(1, n): |
| if stats[i, cv2.CC_STAT_AREA] < min_px: |
| continue |
| ys, xs = np.nonzero(lab == i) |
| pts = np.c_[xs, ys].astype(np.float64) |
| c = pts.mean(0) |
| evals, evecs = np.linalg.eigh(np.cov((pts - c).T)) |
| major = evecs[:, np.argmax(evals)] |
| out.append({"uv": c, "area": int(stats[i, cv2.CC_STAT_AREA]), "axis_uv": major}) |
| return out |
|
|
|
|
| def fit_homography(uv, xy): |
| """Least-squares DLT homography (u, v) -> (x, y).""" |
| A = [] |
| for (u, v), (x, y) in zip(uv, xy): |
| A.append([-u, -v, -1, 0, 0, 0, x * u, x * v, x]) |
| A.append([0, 0, 0, -u, -v, -1, y * u, y * v, y]) |
| _, _, Vt = np.linalg.svd(np.array(A)) |
| return Vt[-1].reshape(3, 3) / Vt[-1, -1] |
|
|
|
|
| def apply_h(H, uv): |
| uv = np.atleast_2d(uv) |
| q = np.c_[uv, np.ones(len(uv))] @ H.T |
| return q[:, :2] / q[:, 2:3] |
|
|
|
|
| def load_calib(): |
| if not os.path.exists(CALIB): |
| raise SystemExit(f"No {CALIB}. Run calibrate_table.py first.") |
| return np.array(json.load(open(CALIB))["H"]) |
|
|
|
|
| def candy_on_table(H, blob): |
| """Table x, y [m] and yaw [rad] of a blob (yaw from the blob's long axis mapped onto the table).""" |
| c = blob["uv"] |
| p0, p1 = apply_h(H, [c - 8 * blob["axis_uv"], c + 8 * blob["axis_uv"]]) |
| d = p1 - p0 |
| xy = apply_h(H, c)[0] |
| return xy, float(np.arctan2(d[1], d[0])) |
|
|
