ModelsAPI/assets_reader.py at master · Instinct323/ModelsAPI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
from __future__ import annotations

from functools import lru_cache

import cv2
import numpy as np

import api.utils.o3d_extension as o3de
from api.utils import Path
from api.utils.camera import Frame, Camera, ObservationSequence


@lru_cache()
def CAM_d435i() -> Camera:
    return Camera.from_yaml("config/camera/RS-D435i.yaml")


@lru_cache()
def COLOR_cat() -> np.ndarray:
    return cv2.imread("assets/cat.jpg")


@lru_cache()
def COLOR_fruit() -> np.ndarray:
    return cv2.imread("assets/fruits-c.jpg")


@lru_cache()
def COLOR_kitchen(i: int) -> np.ndarray:
    return cv2.imread(f"assets/kitchen/{str(i).zfill(2)}.png")


@lru_cache()
def FRAME_aruco() -> Frame:
    return Frame(
        camera=CAM_d435i(),
        color=Path("assets/aruco_obj/c.png"),
        depth=Path("assets/aruco_obj/d.png"),
    )


@lru_cache()
def FRAME_desktop() -> Frame:
    return Frame(
        camera=CAM_d435i(),
        color=Path("assets/desktop-c.png"),
        depth=Path("assets/desktop-d.png"),
    )


@lru_cache()
def FRAMES_green_pot() -> list[Frame]:
    obs_seq = ObservationSequence(root="assets/green-pot", img_ext="png")
    return [Frame(CAM_d435i(), **data) for data in obs_seq.load()]


@lru_cache()
def FRAMES_mustard() -> list[Frame]:
    obs_seq = ObservationSequence(root="assets/mustard", img_ext="png")
    return [Frame(CAM_d435i(), **data) for data in obs_seq.load()]


@lru_cache()
def PCD_PROCESSOR() -> o3de.PointCloudProcessor:
    return o3de.PointCloudProcessor(voxel_size=2e-3, nn_radius=1e-2, nn_k_thresh=10)