
    i
E                        S r SSKrSSKrSSKJr  SSKJr  SSKJ	r
  SSKrSSKrSrSrSrS r\R"                  " / S	Q/ S
Q/ SQ/ SQ/ SQ/5      S-  rSr\\   rS r " S S\5      r\S:X  Ga  \" S5         \" SSS9r\" S\" \5       S35        \S   r\" S5        \" S\S   R8                   35        \" S\S   R;                  5       S S\S   R=                  5       S S 35        \" S!\S"   R8                   35        \" S#\S"   R?                  S$S%9 35        \" S&\S'   R8                   35        \" S(\S'   SS  35        \" S)\S*   R8                   35        \" S+\S,    35        \" S-\S.    35        \" S0S1 35        \" S25        gg! \  a*  r!\" S/\! 35        SSK"r"\"RF                  " 5          Sr!C!N@Sr!C!ff = f)3z-Dataset for real dense trajectory prediction.    N)Dataset)Path   皙ɿ皙?c                 b    U nUS:  a  SnUS:  a  Sn[        [        [        [        U5      5      $ )aP  Process gripper value: fix wrapping and clamp to range.

Args:
    gripper_value: scalar gripper joint value

Returns:
    Processed gripper value in [MIN_GRIPPER, MAX_GRIPPER] = [-0.2, 0.8]
    Processing steps:
    1. Map values > 4.0 to -0.2 (wraparound fix)
    2. Clamp values > 0.8 to 0.8
    3. Ensure final range is [-0.2, 0.8]
g      @r   r   )maxMIN_GRIPPERminMAX_GRIPPER)gripper_valuexs     _/Users/cameronsmith/Projects/robotics_testing/3dkeygrip/volume_dino_tracks_act_baseline/data.pyprocess_gripper_valuer      s7     	A3w3w {CQ/00    )     *@{GVg/@)g
ףp=%@gfffffXr   )r   r   g/)g(\1@g=
ףpTr   )g\(6@g=
ףpQr   g     @@   c                 t    [         R                  " U S5      nX-  nUS   S::  a  gX$SS -  nUSS US   -  nU$ )a  Project 3D point to 2D pixel coordinates.

Args:
    point_3d: (3,) 3D point in world coordinates
    camera_pose: (4, 4) camera pose matrix (world-to-camera)
    cam_K: (3, 3) camera intrinsics

Returns:
    (2,) 2D pixel coordinates [x, y], or None if behind camera
      ?   r   Nr   )npappend)point_3dcamera_posecam_K
point_3d_h	point_cam
point_2d_hpoint_2ds          r   project_3d_to_2dr!   D   sV     8S)J(I |q 2A&J"1~
1-HOr   c                   F    \ rS rSrSrS
S\S-  S\S-  4S jjrS rS r	S	r
g)RealTrajectoryDataset]   a  Dataset for real dense trajectory prediction.

Each sample contains:
    - RGB image (from any frame)
    - Dense 2D trajectory (N_WINDOW waypoints starting from that frame)
    - Camera parameters (pose and intrinsics)
    - Ground truth heatmaps for each timestep

For each episode, creates samples from every frame, with trajectories
padded with the last observed keypoint if there aren't enough subsequent frames.
Nepisodemax_episodesc                 (   [        U5      U l        X l        [        U l        U R                  R                  5       (       d  [        SU R                   35      e[        U R                  R                  5        Vs/ s H.  nUR                  5       (       d  M  SUR                  ;   d  M,  UPM0     sn5      nUbM  U Vs/ s H  oUR                  U:X  d  M  UPM     nn[        U5      S:X  a  [        SU SU R                   35      eUb  USU n[        U5      S:X  a  [        SU R                   35      e/ U l        U H  n[        UR                  S5       Vs/ s H%  oR                  R                  5       (       d  M#  UPM'     sn5      n	U	 H4  n
[!        U
R                  5      nU R                  R#                  X{45        M6     M     [%        S	[        U5       S
35        [%        S[        U R                  5       S35        gs  snf s  snf s  snf )a6  Initialize dataset.

Args:
    dataset_root: Root directory containing episodes
    image_size: Size to resize images to (will be square, default 448)
    episode: Optional episode directory name (e.g. "episode_001") to load only.
    max_episodes: Optional limit on how many episodes to load (after sorting).
zDataset directory not found: r%   Nr   z	Episode 'z' not found in zNo episodes found in *.pngzLoaded z	 episodeszCreated z samples (one per frame))r   dataset_root
image_sizeN_WINDOWn_windowexists
ValueErrorsortediterdiris_dirnamelensamplesglobstemisdigitintr   print)selfr)   r*   r%   r&   depisode_dirsepisode_dirfframe_files
frame_file	frame_idxs               r   __init__RealTrajectoryDataset.__init__j   s    !.$    ''))<T=N=N<OPQQ$*;*;*C*C*E D*EQ !
  !/8AFF/B  !*E D E '3I|!vv7HA|LI< A% 9WI_TEVEVDW!XYY#'6L|!4T5F5F4GHII 'K [-=-=g-F![-F&&..JZ!-F![\K *

0	##[$<= * ( 	L)*)45T\\*++CDE5D J "\s*   :HH)H=H
H
"H
?H
c                 ,    [        U R                  5      $ )N)r3   r4   )r:   s    r   __len__RealTrajectoryDataset.__len__   s    4<<  r   c                 T   U R                   U   u  p#US nX$ S3-  n[        R                  " U5      SSS24   nUR                  SS u  px[	        UR                  S5       V	s/ s H%  oR                  R                  5       (       d  M#  U	PM'     sn	5      n
U
 V	s/ s H  n	[        U	R                  5      PM     nn	UR                  U5      n/ n/ n/ n[        U R                  5       GH  nUU-   [        U5      :  a    GOXU-      nUS nUU S3-  nUR                  5       (       d    GOn[        R                  " U5      nUSS2SS24   nUSS2S4   nU[         -  U-   nUR#                  U5        UU S	3-  nUR                  5       (       aA  [        R                  " U5      n[%        US
   5      n['        U5      nUR#                  U5        OUR#                  S5        X$ S3-  nX$ S3-  nUR                  5       (       a  UR                  5       (       d    O~[        R                  " U5      n[        R                  " U5      nUR)                  5       nUS==   U-  ss'   US==   U-  ss'   [+        UUU5      n U c    OUR#                  U 5        GM     [        U5      S:X  Ga  [        R                  " X$ S3-  5      nUSS2SS24   nUSS2S4   nU[         -  U-   nU/nX$ S	3-  nUR                  5       (       a3  [        R                  " U5      n[%        US
   5      n['        U5      nU/nOS/n[        R                  " X$ S3-  5      n[        R                  " X$ S3-  5      nUR)                  5       nUS==   U-  ss'   US==   U-  ss'   [+        UUU5      n U b  U /n[        U5      S:X  a  [-        SUR.                   SU 35      e[        R0                  " U5      n[        R0                  " U5      n[        R0                  " U5      n[        U5      U R                  :  a  US
S n!US
S n"US
S n#U R                  [        U5      -
  n$[        R2                  " U[        R4                  " U!U$S45      /SS9n[        R2                  " U[        R4                  " U"U$S45      /SS9n[        R2                  " U[        R4                  " U#U$45      /SS9nOF[        U5      U R                  :  a-  USU R                   nUSU R                   nUSU R                   nUS
   n%[        R                  " X$ S3-  5      n[        R                  " X$ S3-  5      nXpR6                  :w  d  XR6                  :w  ar  [8        R:                  " X`R6                  U R6                  4[8        R<                  S9nU R6                  U-  n&U R6                  U-  n'U[        R0                  " U&U'/5      -  n/ n([        U R                  5       GH   n)UU)   n*[        R0                  " [        R>                  " U*S   SU R6                  S-
  5      [        R>                  " U*S   SU R6                  S-
  5      /5      n*[        R@                  " U R6                  U R6                  4[        RB                  S9n+[        [E        U*S   5      5      n,[        [E        U*S   5      5      n-SU,s=::  a  U R6                  :  a$  O  O!SU-s=::  a  U R6                  :  a
  O  OSU+U-U,4'   U(R#                  U+5        GM#     [        R0                  " U(5      n([F        RH                  " U5      RK                  SSS5      R%                  5       n.U.RM                  5       S:  a  U.S-  n.[F        RN                  " / SQ5      RQ                  SSS5      n/[F        RN                  " / SQ5      RQ                  SSS5      n0U.U/-
  U0-  n.[F        RH                  " U(5      R%                  5       n1[F        RH                  " U5      R%                  5       n2[F        RH                  " U5      R%                  5       n3[F        RH                  " U5      R%                  5       n4[F        RH                  " U%5      R%                  5       n5[F        RH                  " U5      R%                  5       n6[F        RH                  " U5      R%                  5       n7X$ S	3-  n8Sn9U8R                  5       (       a   [        R                  " U85      n9US-  n:Sn;U:R                  5       (       a   [        R                  " U:5      n;U9b7  [F        RH                  " [        RT                  " U9[        RB                  S95      OSn<U;b7  [F        RH                  " [        RT                  " U;[        RB                  S95      OSn=U.U1U2U3U4U5U6U7UR.                   SU 3UR.                  U[W        U5      U<U=S.$ s  sn	f s  sn	f ! [R         a    Sn9 Nf = f! [R         a    Sn; Nf = f)ae  Get a single sample.

Returns:
    dict with keys:
        - rgb: (3, H, W) normalized RGB image
        - heatmap_target: (N_WINDOW, H, W) one-hot heatmaps for each timestep
        - trajectory_2d: (N_WINDOW, 2) 2D pixel locations for each timestep
        - trajectory_3d: (N_WINDOW, 3) 3D world positions for each timestep
        - trajectory_gripper: (N_WINDOW,) gripper values for each timestep
        - target_3d: (3,) final target 3D world position (last waypoint)
        - camera_pose: (4, 4) camera pose matrix (from first frame)
        - cam_K_norm: (3, 3) normalized intrinsics (from first frame)
06dz.png.Nr   r   r(   z_gripper_pose.npyz.npyr   z
_cam_K.npyz_camera_pose.npyr      z!Could not compute trajectory for z frame )axis)interpolation)dtypeg     o@)g
ףp=
?gv/?gCl?)gZd;O?gy&1?g?z
000000.npy_frame_)rgbheatmap_targettrajectory_2dtrajectory_3dtrajectory_gripper	target_3dr   
cam_K_norm
episode_idepisode_namerA   r=   joint_stateepisode_start_joint_state),r4   pltimreadshaper/   r5   r6   r7   r8   indexranger,   r3   r-   r   loadkp_localr   floatr   copyr!   r.   r2   arrayconcatenatetiler*   cv2resizeINTER_LINEARclipzerosfloat32roundtorch
from_numpypermuter	   tensorview	Exceptionasarraystr)>r:   idxr=   rA   	frame_strrgb_pathrO   H_origW_origr>   r?   frame_indicesstart_frame_idxrQ   rR   rS   inext_frame_idxnext_frame_strgripper_pose_pathgripper_posegripper_rotgripper_poskp_3djoint_state_pathrX   r   cam_K_norm_pathcamera_pose_pathrU   r   r   kp_2dlast_point_2dlast_point_3dlast_grippern_padrT   scale_xscale_yheatmap_targetst	target_2drP   target_xtarget_y
rgb_tensormeanstdheatmap_tensortrajectory_2d_tensortrajectory_3d_tensortrajectory_gripper_tensortarget_3d_tensorcamera_pose_tensorcam_K_norm_tensorjoint_state_path_curjoint_state_curjoint_state_path_ep0joint_state_ep0joint_state_cur_tensorjoint_state_ep0_tensors>                                                                 r   __getitem__!RealTrajectoryDataset.__getitem__   s	    "&c!2 o	 ;d!33jj"37+2A )9)9')BW)BAffnnFVa)BWX.9:kQVVk:'--i8  t}}%A"c-&88*Q+>?N .s3N !,0@@Q.R R$++--77#45L&rr2A2v.K&rr1u-K  (*[8E  '  +/?t-DD&&(( gg&67 %k"o 6 5m D"))-8 #))#. *k,DDO*{:J-KK#**,,1A1H1H1J1J1J''"23K OO%E!HH!HH %UK?E}  'g &j }"77;;>O1P#PQL&rr2A2v.K&rr1u-K(*[8E"GM  +{$-??&&(( gg&67 %k"o 6 5m D&3_"&)U"J/G!GHJ''++=M0N"NOKOO%E!HH!HH$UK?E !&}"@AQAQ@RRYZcYdeff//XX&89 }-)"#.M)"#.M-bc2LMMC$66ENNM277=SXZ[R\;]+^efgMNNM277=SXZ[R\;]+^efgM!#1CRWW\\a[cEd0elm!n$--/).4==9M).4==9M!3NT]]!C ""%	 ggkk9I,JJKWW[[
+CCD
 __$//(A**S??DOO"D),)9)9;C oo.Goo.G)BHHgw5G,HHM t}}%A%a(I 	!a1)<=	!a1)<=" I  XXt&HPRPZPZ[N51./H51./HH.t.13R4??3R58x12"">2! &$ ((?3 %%c*221a;AAC
>>c!#e+J ||12771a@ll0166q!Q? 4'3.
))/:@@B$//>DDF$//>DDF$)$4$45G$H$N$N$P! ++I6<<>"--k:@@B!,,Z8>>@  +{$-??&&(('"$''*>"?  +\9&&(('"$''*>"? * RZZrzzJK 	 * RZZrzzJK 	 ,11";)-+(--.gi[A',,"{+1)?
 	
O X:`  '"&'  '"&'s6   #"g<	g<hh ;h hhh'&h')r)   r*   r,   r4   )zscratch/  NN)__name__
__module____qualname____firstlineno____doc__rt   r8   rB   rE   r   __static_attributes__ r   r   r#   r#   ]   s4    
+Ft +Fjmptjt +FZ!m
r   r#   __main__z Testing RealTrajectoryDataset...r   z&scratch/parsed_moredata_pickplace_home)r*   r)   u   ✓ Loaded z samplesz

Sample 0:z  RGB shape: rO   z  RGB range: [z.3fz, ]z  Heatmap shape: rP   z  Heatmap sum per timestep: )rJ   r   )dimz  Trajectory 2D shape: rQ   z  Trajectory 2D (first 3): z  Trajectory 3D shape: rR   z  Target 3D: rT   z  Episode ID: rV   u   ✗ Error: 
z<============================================================u   ✓ Dataset test complete!)$r   rm   numpyr   torch.utils.datar   pathlibr   matplotlib.pyplotpyplotrZ   rf   mathr+   r
   r   r   rc   KEYPOINTS_LOCAL_M_ALLKP_INDEXr`   r!   r#   r   r9   datasetr3   sampler\   r   r	   sumrr   e	traceback	print_excr   r   r   <module>r      s:   3   $   
   1\ "79JLb#57I"K LNTU  *2j
G j
Z	 z	
,-'3EmnCL>23 fUm11234ve}00237r&-:K:K:Mc9RRSTU!&)9":"@"@!ABC,V4D-E-I-Ie-I-T,UVW'(?(E(E'FGH+F?,CBQ,G+HIJ'(?(E(E'FGHf[1234vl3456 
Bvh-	
&'7 *  A3 s   7DF F= F88F=