dtorre38 / mujoco_opencv Goto Github PK

def get_frame(model, data, opt, scene, context, camera_name, loc_x, loc_y, width=640, height=480):

    # 1. Create a rectangular viewport in the upper right corner for example.
   offscreen_viewport = mj.MjrRect(int(loc_x), int(loc_y), width, height)    
    
    # 2. Specify a different camera view by updating the scene with mjv_updateScene.
    # Set the camera to the specified view
    camera_id = mj.mj_name2id(model, mj.mjtObj.mjOBJ_CAMERA, camera_name)
    offscreen_cam = mj.MjvCamera()
    offscreen_cam.type = mj.mjtCamera.mjCAMERA_FIXED
    offscreen_cam.fixedcamid = camera_id

    mj.mjv_updateScene(model, data, opt, None, offscreen_cam, mj.mjtCatBit.mjCAT_ALL.value, scene)
    
    # 3.Render the scene in the offscreen buffer with mjr_render.
    depth = np.zeros((height*width*3, 1), dtype=np.uint8)  # Placeholder for depth data
    mj.mjr_render(offscreen_viewport, scene, context)

    # 4. Read the depth with mjr_readPixels.
    mj.mjr_readPixels(None, depth, offscreen_viewport, context)
        
    # 5. Depth processing 
    # Shift nearest values to the origin.
    depth-= depth.min()
    # Scale by 2 mean distances of near rays.
    depth= depth/ 2*depth[depth<= 1].mean()
    # Scale to [0, 255]
    depth= 255*np.clip(depth, 0, 1)
          
    # 6. Call mjr_drawPixels using the rectangular viewport you created in step 1.
    # mj.mjr_drawPixels(depth, None, offscreen_viewport, context)
    
    return frame, offscreen_viewport