demo.py

from __future__ import annotations

import cv2
import matplotlib.pyplot as plt
import numpy as np
import open3d as o3d
import supervision as sv
import torch
from PIL import Image

import api.utils.o3d_extension as o3de
import api.utils.sv_extension as sve
from api import THREAD_POOL
from api.utils import Path, Timer
from api.utils.camera import Frame, Camera, Pinhole, ObservationSequence

D435I = Camera.from_yaml("config/camera/RS-D435i.yaml")
PCD_PROCESSOR = o3de.PointCloudProcessor(voxel_size=2e-3, nn_radius=1e-2, nn_k_thresh=10)


def test_ai_client():
    from api.utils import chat_tool
    from langchain_openai import ChatOpenAI
    from langchain_core.messages import HumanMessage

    content = chat_tool.make_content([
        "这副图描绘了什么颜色的内容？",
        ("image_url", cv2.imread("assets/fruits-c.jpg"))
    ])

    llm = ChatOpenAI(
        model="qwen3.6-plus-2026-04-02",
        **chat_tool.CLIENT_KWARGS["bailian"]
    )
    print(llm.invoke([HumanMessage(content)]))


def test_clip():
    from api.encoders import OpenCLIP
    # print(load_dinov2())

    clip = OpenCLIP()

    # Image-text similarity
    img = Image.open("assets/cat.jpg").convert("RGB")
    feat_img = clip.encode_images([img])
    feat_text = clip.encode_texts(["a photo of 5 cats", "a dog", "desktop"])
    print(feat_img.shape)
    score = (feat_img @ feat_text.T).softmax(dim=-1)
    print(score)

    # Multi-view similarity
    for obj in ("green-pot", "mustard"):
        root = Path(f"assets/{obj}/color")
        img = [Image.open(file).convert("RGB") for file in root.glob("*.png")]
        feat_img = clip.encode_images(img)
        print(feat_img @ feat_img.T)


def test_codet():
    from api.codet import CoDet
    from api.segment_anything import SegmentAnythingV2

    codet = CoDet()
    sam2 = SegmentAnythingV2("tiny")

    img = cv2.imread("/media/tongzj/Data/Workbench/Lab/assets/takeout.jpg")
    img = sv.resize_image(img, [640] * 2, keep_aspect_ratio=True)
    with Timer("Main"):
        dets = sam2.inference_with_box(img, THREAD_POOL.submit(codet, img))
    cv2.imwrite("runs/codet-sam2.jpg", sve.annotate(img, dets))

    # Infer a video stream
    from api.utils.realsense import RGBDCamera

    for color, depth in RGBDCamera():
        with Timer("Main"):
            dets = sam2.inference_with_box(color, THREAD_POOL.submit(codet, color))
        plt.imshow(sve.annotate(color, dets)[..., ::-1])
        plt.pause(1e-3)


def test_cgn():
    from api.contact_graspnet import ContactGraspNet

    # Load RGB-D data
    struct = np.load("assets/rgb-depth-K.npy", allow_pickle=True).item()
    camera = Pinhole(img_size=struct["depth"].shape[::-1],
                     intrinsics=struct["K"][[0, 1, 0, 1], [0, 1, 2, 2]])
    pointmap = camera.pointmap(struct["depth"]).astype(np.float32)
    mask = (pointmap[..., 2] > 0.2) & (pointmap[..., 2] < 1.8)

    # Infer grasp poses
    cgn = ContactGraspNet()
    grasps = cgn(pointmap[mask])

    # contact point
    cpt = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(grasps.left_finger_tip))

    # Visualize point cloud
    if grasps:
        vis = o3d.visualization.Visualizer()
        vis.create_window()

        vis.add_geometry(o3de.rgbd2pcd(pointmap, struct["rgb"], mask))
        vis.add_geometry(cpt)
        for mesh in grasps.to_meshes(): vis.add_geometry(mesh)

        vis.run()
        vis.destroy_window()


def test_curobo():
    from api.utils.curobo_planning import CuRoboPlanner
    from api.utils import pose_tf

    Tce = pose_tf.SE3.from_config([0.048098, -0.985089, -0.165184, 0.06595,
                                   0.998825, 0.048422, 0.002063, -0.000357,
                                   0.005966, -0.165089, 0.986261, 0.048961])

    CuRoboPlanner.generate_ws_for_lookat(
        "runs/franka_ws.npy", focus=np.array([.5, 0, .1]), Tec=np.linalg.inv(Tce))


def test_da():
    from api.depth_anything import DepthAnythingV3

    camera = D435I
    camera.depth_range[1] = .5
    obs_seq = ObservationSequence(root="assets/green-pot", img_ext="png")
    fabs = [Frame(camera, **data) for data in obs_seq.load()]

    if 0:
        model = DepthAnythingV3()
        with Timer("Reconstruct"):
            scene = model.mapping(fabs)
        pcds = [PCD_PROCESSOR(scene.get_aligned_pcd("first"))]

    else:
        model = DepthAnythingV3("mono")
        with Timer("Monocular"):
            frel = model.mapping(fabs[:1]).frames_rel[0]
        pcds = [frel.to_pcd()]

    pcds += [fabs[0].to_pcd()]
    pcds[-1].paint_uniform_color([0., 0., 1.])
    o3d.visualization.draw_geometries(pcds)


def test_deberta():
    from api.encoders import DeBERTa

    texts = ["Hello, I love using DeBERTa for NLP tasks!", "Grasp the cup"]
    model = DeBERTa()
    ret = model(texts)
    print(ret)

    for k, v in ret.items():
        print(f"{k}: {v.shape}")


def test_decompose():
    from api import decompose
    from tqdm import tqdm

    root = Path(r"../assets/objects")
    output_dir = Path("runs/coacd")
    output_dir.mkdir(parents=True, exist_ok=True)

    img_size = (640, 480)
    acd = decompose.CoACD(threshold=0.07, max_convex_hull=10, verbose=False)

    # Process all *.glb files
    for file in tqdm(list(root.rglob("*.glb"))):
        output = output_dir / "_".join(file.relative_to(root).parts)
        output = output.with_suffix(".bin")

        # Load original mesh and render
        org = o3d.io.read_triangle_mesh(file, enable_post_processing=True)
        rgb_org = o3de.auto_render_color(img_size, [org])
        cv2.imwrite(str(output.with_suffix(".png")), rgb_org[..., ::-1])

        # Decompose and render parts
        parts = output.lazy_obj(lambda: acd(org))
        rgb_parts = o3de.auto_render_color(img_size, [decompose.part_to_o3d(parts)])
        cv2.imwrite(str(output.with_suffix(".jpg")), rgb_parts[..., ::-1])

    # Visualize
    for f in output_dir.rglob("*.bin"):
        parts = decompose.part_to_o3d(f.binary())
        o3d.visualization.draw_geometries([parts])


def test_gdino():
    from api.grounding_dino import GroundingDINO

    gdino = GroundingDINO(box_thresh=0.5)
    image = cv2.imread("/media/tongzj/Data/Downloads/1.jpg")

    dets = gdino(image, "robot. bag")
    cv2.imwrite("runs/gdino.jpg", sve.annotate(image, dets))

    # SAM2
    from api.segment_anything import SegmentAnythingV2
    sam2 = SegmentAnythingV2("tiny")
    dets = sam2.inference_with_box(image, dets)
    cv2.imwrite("runs/gsam.png", sve.annotate(image, dets))

    # other
    mask = dets.mask.any(axis=0)
    image[~mask] = 255
    cv2.imwrite("runs/debug.png", image)


def test_handeye_calib():
    from api.utils.code_det import Charuco, hand_in_eye_calib, hand_to_eye_calib

    len_marker = .021
    detector = Charuco(len_marker=len_marker, len_square=.028, shape=(5, 7))

    # cv2.imwrite("runs/charuco.png", detector.to_image(100)), exit()

    def info_result(T):
        print(T)
        T = T[:3].flatten().tolist()
        print([round(i, 8) for i in T])

    if 0:
        root = Path("runs/handeye1")
        cam = D435I
        bgrs = [cv2.imread(str(file)) for file in (root / "color").iterdir()]
        Teb = np.stack([Path(file).binary()["transform"]["T"] for file in (root / "metadata").iterdir()])

        Toc = [detector.pose_global(bgr, cam) for bgr in bgrs]
        Tce = hand_in_eye_calib(Toc, Teb)
        info_result(Tce)

    else:
        root = Path("runs/handeye2")
        cam = Camera.from_yaml("config/camera/RS-L515.yaml")
        frames = [Path(f).binary() for f in root.iterdir()]
        bgrs = [f["cam2"][1] for f in frames]
        Teb = np.stack([f["transform"]["T"] for f in frames])

        Toc = np.stack([detector.pose_global(bgr, cam) for bgr in bgrs])
        Tcb = hand_to_eye_calib(Toc, Teb)
        info_result(Tcb)

        for f in frames:
            frame = Frame(cam, *f["cam2"][1:], Tcw=Tcb)
            o3d.visualization.draw_geometries([
                frame.to_pcd(), o3d.geometry.TriangleMesh.create_coordinate_frame(0.2)
            ])


def test_obj_tracker():
    from api.grounding_dino import GroundingDINO
    from api.segment_anything import SegmentAnythingV2
    from api.utils.object_tracker import ArucoObjectTracker

    tracker = ArucoObjectTracker("config/aruco_obj.yaml")

    gdino = GroundingDINO()
    sam2 = SegmentAnythingV2("tiny")

    frame = Frame(
        camera=D435I,
        color="assets/aruco_obj/c.png",
        depth="assets/aruco_obj/d.png",
    )

    # track objects
    tracker.track(frame)
    geos = [obj.o3dmesh() for obj in tracker.objects]
    geos = list(filter(None, geos))
    if geos:

        # segment objects
        dets = sam2.inference_with_box(frame.color, THREAD_POOL.submit(gdino, frame.color, caption="bag"))
        color_anno = sve.annotate(frame.color, dets, anno_box=False)
        cv2.imwrite("runs/obj_tracker.png", color_anno)

        # visualize
        pcds = [frame.to_pcd(mask) for mask in dets.mask]
        idx = np.arange(len(pcds))
        for i, color in zip(idx, plt.get_cmap("tab20")(idx)[:, :3]):
            pcds[i].paint_uniform_color(color)

        # save objects
        Path("runs/shelf.bin").binary((tracker.objects[0].pose, [o3de.to_pcd_array(x) for x in pcds]))

        o3d.visualization.draw_geometries(geos + pcds)


def test_o3de():
    camera = D435I
    camera.depth_scale = 5000
    camera.depth_range = None

    color = cv2.imread("assets/desktop-c.png")
    depth = camera.process_depth("assets/desktop-d.png")
    pcd = o3de.rgbd2pcd(camera.pointmap(depth), color, mask=depth > 0)

    pcd = o3de.estimate_pcd_normal(pcd)  # , cam_pos=[0] * 3)
    o3d.visualization.draw_geometries([pcd])


def test_radio():
    from api.encoders import RADIO

    color = cv2.imread("assets/cat.jpg")
    color = torch.as_tensor(color / 255).permute(2, 0, 1)[None]

    model = RADIO(enable_sam=True)
    ret = model.encode_images(color)

    if model.enable_text:
        feat_text = model.encode_texts(["cats", "cat", "dog"])
        print(torch.cosine_similarity(ret[1][0], feat_text))

    if model.enable_sam:
        model.sam()


def test_sam2():
    from api.segment_anything import SegmentAnythingV2
    from api.utils.realsense import RGBDCamera

    sam2 = SegmentAnythingV2("tiny")

    image = cv2.imread("assets/desktop-c.png")
    point = np.array([[310, 160]])

    # Automatic mask generation
    pred = sam2.build_predictor(image, auto_mode=True)
    dets = pred()
    cv2.imwrite("runs/sam2_all.png", sve.annotate(image, dets, anno_box=False))

    # Predict with prompts
    dets = pred(point_coords=point, point_labels=np.ones(len(point), dtype=np.bool_))
    cv2.imwrite("runs/sam2.png", sve.annotate(image, dets, anno_box=False))

    for c, d in RGBDCamera():
        dets = sam2(c)
        print(dets.confidence)
        plt.imshow(sve.annotate(c, dets))
        plt.pause(1e-3)


def test_server():
    import logging
    import sys
    import time
    from api.utils import func_invoker

    # config
    FUNCTIONS = {
        "add": lambda x, y: x + y,
        "sleep": lambda x: time.sleep(x) or x,
    }

    # FastAPI
    if len(sys.argv) > 0 and sys.argv[0].split("/")[-1] == "uvicorn":
        import fastapi

        SERVER = func_invoker.FastapiServer(functions=FUNCTIONS, loginfo=logging.getLogger("uvicorn").info)

        # uvicorn server:app
        app = fastapi.FastAPI(
            title="Functions API",
            description="author: [TongZJ](https://github.com/Instinct323)\n\n" + SERVER.func_docs(),
            version="1.0.0"
        )

        @app.post("/invoke")
        async def invoke(request):
            return await SERVER.get_response(request)

    # Zmq
    else:
        import zmq
        from api.utils.zmq_socket import ZmqSocket

        SERVER = func_invoker.ZmqSocketServer(
            functions=FUNCTIONS, loginfo=print,
            sock=ZmqSocket(zmq.REP, addr="0.0.0.0:5555", queue_size=10)
        )
        print("Ready")
        SERVER.main()


def test_tag2text():
    from api.grounding_dino import GroundingDINO
    from api.tag2text import Tag2Text

    gdino = GroundingDINO()
    tt = Tag2Text(tag_thresh=0.5)
    image = cv2.imread("/media/tongzj/Data/Workbench/Lab/Lan-grasp/assets/green-pot/c1.png")

    text = tt(sv.cv2_to_pillow(image))
    print(text)
    dets = gdino(image, text)
    cv2.imwrite("runs/tt-gdino.jpg", sve.annotate(image, dets))


def test_zero123plus():
    from api.zero123plus import Zero123plus
    from api.grounding_dino import GroundingDINO
    from api.segment_anything import SegmentAnythingV2

    zero123p = Zero123plus()
    gdino = GroundingDINO()
    sam2 = SegmentAnythingV2("large")

    target = "cup"
    img = cv2.imread("assets/cup2-c.png")

    with Timer("Main"):
        # Ground and segment the image
        mask = sam2(img, box=gdino(img, target).xyxy[0][None]).mask[0]
        img[~mask] = 114
        # Inference
        result = zero123p(Image.fromarray(img), 75)
        bboxes = [gdino(r, target)[0] for r in result]

    # Show the results
    for i, (img, box) in enumerate(zip(result, bboxes)):
        plt.subplot(2, 3, i + 1)
        plt.imshow(sve.annotate(img, box))
    plt.show()


if __name__ == '__main__':
    test_ai_client()