
    "i
E                     6   d Z ddlZddlZddlmZ ddlmZ ddlm	Z
 ddlZddlZdZdZdZd Z ej        g d	g d
g dg dg dg          dz  ZdZee         Zd Z G d de          Zedk    r ed           	  edd          Z ed ee           d           ed         Z ed            eded         j                     eded                                         dded                                         dd             ed!ed"         j                     ed#ed"                             d$%                       ed&ed'         j                     ed(ed'         dd                      ed)ed*         j                     ed+ed,                      ed-ed.                     n3# e $ r+Z! ed/e!            ddl"Z" e"j#                     Y dZ![!ndZ![!ww xY w ed0d1             ed2           dS dS )3z-Dataset for real dense trajectory prediction.    N)Dataset)Path   皙ɿ皙?c                 v    | }|dk    rd}|dk    rd}t          t          t          t          |                    S )a|  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     U/Users/cameronsmith/Projects/robotics_testing/3dkeygrip/volume_tracks_smolvla/data.pyprocess_gripper_valuer      sA     	A3ww3ww {CQ//000    )     *@{GVg/@)g
ףp=%@gfffffXr   )r   r   g/)g(\1@g=
ףpTr   )g\(6@g=
ףpQr   g     @@   c                     t          j        | d          }||z  }|d         dk    rdS ||dd         z  }|dd         |d         z  }|S )a0  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   sd     8S))Jj(I |qt 2A2&J"1"~
1-HOr   c                   <    e Zd ZdZd
dedz  dedz  fdZd Zd	 ZdS )RealTrajectoryDataseta  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.
    scratch/  Nepisodemax_episodesc                 `   t          |          | _        || _        t          | _        | j                                        st          d| j                   t          d | j                                        D                       };fd|D             }t          |          dk    rt          d d| j                   |
|d|         }t          |          dk    rt          d| j                   g | _
        |D ]c}t          d	 |                    d
          D                       }|D ]2}t          |j                  }	| j
                            ||	f           3dt          dt          |           d           t          dt          | j
                   d           dS )an  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: c                 L    g | ]!}|                                 rd |j        v |"S )r&   )is_dirname).0ds     r   
<listcomp>z2RealTrajectoryDataset.__init__.<locals>.<listcomp>{   sB     D D DQ !

D/8AF/B/B  !/B/B/Br   Nc                 *    g | ]}|j         k    |S  )r+   )r,   r-   r&   s     r   r.   z2RealTrajectoryDataset.__init__.<locals>.<listcomp>   s%    III!qv7H7HA7H7H7Hr   r   z	Episode 'z' not found in zNo episodes found in c                 D    g | ]}|j                                         |S r0   stemisdigitr,   fs     r   r.   z2RealTrajectoryDataset.__init__.<locals>.<listcomp>   s)    ![![![!&..JZJZ![!![![![r   *.pngzLoaded z	 episodeszCreated z samples (one per frame))r   dataset_root
image_sizeN_WINDOWn_windowexists
ValueErrorsortediterdirlensamplesglobintr3   r   print)
selfr8   r9   r&   r'   episode_dirsepisode_dirframe_files
frame_file	frame_idxs
      `      r   __init__zRealTrajectoryDataset.__init__j   s    !..$   '')) 	RPT=NPPQQQ D D$*;*C*C*E*E D D D E E IIII|IIIL<  A%% !XW!X!XTEV!X!XYYY#'6L|!!HT5FHHIII ' 	> 	>K ![![[-=-=g-F-F![![![\\K * > >

00	##[)$<====> 	4L))444555DT\**DDDEEEEEr   c                 *    t          | j                  S )N)r@   rA   )rE   s    r   __len__zRealTrajectoryDataset.__len__   s    4<   r   c                 .   | j         |         \  }}|d}|| dz  }t          j        |          dddf         }|j        dd         \  }}t	          d |                    d          D                       }	d	 |	D             }
|
                    |          }g }g }g }t          | j                  D ]}||z   t          |
          k    r n|
||z            }|d}|| d
z  }|
                                s nt          j        |          }|ddddf         }|dddf         }|t          z  |z   }|                    |           || dz  }|
                                rNt          j        |          }t          |d                   }t!          |          }|                    |           n|                    d           || dz  }|| dz  }|
                                r|
                                s nt          j        |          }t          j        |          }|                                }|dxx         |z  cc<   |dxx         |z  cc<   t%          |||          }| n|                    |           t          |          dk    rt          j        || d
z            }|ddddf         }|dddf         }|t          z  |z   }|g}|| dz  }|
                                r<t          j        |          }t          |d                   }t!          |          }|g}ndg}t          j        || dz            }t          j        || dz            }|                                }|dxx         |z  cc<   |dxx         |z  cc<   t%          |||          }||g}t          |          dk    rt'          d|j         d|           t          j        |          }t          j        |          }t          j        |          }t          |          | j        k     r|dd         } |dd         }!|dd         }"| j        t          |          z
  }#t          j        |t          j        | |#df          gd          }t          j        |t          j        |!|#df          gd          }t          j        |t          j        |"|#f          gd          }nEt          |          | j        k    r-|d| j                 }|d| j                 }|d| j                 }|d         }$t          j        || dz            }t          j        || dz            }|| j        k    s|| j        k    rZt3          j        || j        | j        ft2          j                  }| j        |z  }%| j        |z  }&|t          j        |%|&g          z  }g }'t          | j                  D ]}(||(         })t          j        t          j        |)d         d| j        dz
            t          j        |)d         d| j        dz
            g          })t          j        | j        | j        ft          j                  }*t?          tA          |)d                             }+t?          tA          |)d                             },d|+cxk    r| j        k     rn nd|,cxk    r| j        k     r
n nd|*|,|+f<   |'                    |*           t          j        |'          }'tC          j"        |          #                    ddd                                          }-|-$                                dk    r|-dz  }-tC          j%        g d          &                    ddd          }.tC          j%        g d          &                    ddd          }/|-|.z
  |/z  }-tC          j"        |'                                          }0tC          j"        |                                          }1tC          j"        |                                          }2tC          j"        |                                          }3tC          j"        |$                                          }4tC          j"        |                                          }5tC          j"        |                                          }6|| dz  }7d}8|7
                                r(	 t          j        |7          }8n# tN          $ r d}8Y nw xY w|dz  }9d}:|9
                                r(	 t          j        |9          }:n# tN          $ r d}:Y nw xY w|82tC          j"        t          j(        |8t          j                            nd};|:2tC          j"        t          j(        |:t          j                            nd}<|-|0|1|2|3|4|5|6|j         d| |j        |tS          |          |;|<dS )a  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   c                 D    g | ]}|j                                         |S r0   r2   r5   s     r   r.   z5RealTrajectoryDataset.__getitem__.<locals>.<listcomp>   s)    WWWAafnnFVFVWaWWWr   r7   c                 6    g | ]}t          |j                  S r0   )rC   r3   r5   s     r   r.   z5RealTrajectoryDataset.__getitem__.<locals>.<listcomp>   s     :::QV:::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_namerJ   rG   joint_stateepisode_start_joint_state)*rA   pltimreadshaper>   rB   indexranger;   r@   r<   r   loadkp_localr   floatr   copyr!   r=   r+   arrayconcatenatetiler9   cv2resizeINTER_LINEARclipzerosfloat32rC   roundtorch
from_numpypermuter	   tensorview	Exceptionasarraystr)=rE   idxrG   rJ   	frame_strrgb_pathrX   H_origW_origrH   frame_indicesstart_frame_idxrZ   r[   r\   inext_frame_idxnext_frame_strgripper_pose_pathgripper_posegripper_rotgripper_poskp_3djoint_state_pathra   r   cam_K_norm_pathcamera_pose_pathr^   r   r   kp_2dlast_point_2dlast_point_3dlast_grippern_padr]   scale_xscale_yheatmap_targetst	target_2drY   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__z!RealTrajectoryDataset.__getitem__   s    "&c!2Y &&	 I!3!3!33j""37+2A2 WW)9)9')B)BWWWXX::k:::'--i88  t}%% 3	( 3	(A"c-&8&888*?Q+>?N .44N !,.R.R.R R$++-- 7#455L&rr2A2v.K&rr1u-K  (*[8E  '''  +-D-D-DD&&(( / g&677 %k"o 6 6 5m D D"))-8888 #))#... *y,D,D,DDO*	-K-K-KK#**,, 1A1H1H1J1J 11J'"233K OO%%E!HHHHHH!HHHHHH %UK??E}  ''''}""7;I1P1P1P#PQQL&rr2A2v.K&rr1u-K(*[8E"GM  +	-?-?-??&&(( + g&677 %k"o 6 6 5m D D&3_""&)U")/G/G/G!GHHJ'+90N0N0N"NOOKOO%%E!HHHHHH!HHHHHH$UK??E !&}""eAQeeZceefff////X&899 }--)"##.M)"##.M-bcc2LMC$6$66ENM27=SXZ[R\;];]+^efgggMNM27=SXZ[R\;];]+^efgggM!#1CRW\\a[cEdEd0elm!n!n!n$-//).4=.9M).4=.9M!3NT]N!C ""%	 gky,J,J,JJKKW[i+C+C+CCDD
 T_$$$/(A(A*S4?DO"D),)9; ; ;C o.Go.G)BHgw5G,H,HHM t}%% 	3 	3A%a(I 	!a1)<==	!a1)<=="  I  Xt&HPRPZ[[[N51..//H51..//HH....t.....13R3R3R3R4?3R3R3R3R3R58x12"">2222(?33 %c**221a;;AACC
>>c!!#e+J |11122771a@@l0001166q!Q?? 4'3.
)/::@@BB$/>>DDFF$/>>DDFF$)$45G$H$H$N$N$P$P! +I66<<>>"-k::@@BB!,Z88>>@@  +	-?-?-??&&(( 	''"$'*>"?"? ' ' '"&'  +\9&&(( 	''"$'*>"?"? ' ' '"&'
 * RZrzJJJKKK 	 * RZrzJJJKKK 	 ,11";)-+(-AAiAA',"{++1)?
 
 	
s$   d) )d87d8e, ,e;:e;)r$   r%   NN)	__name__
__module____qualname____doc__r}   rC   rK   rM   r   r0   r   r   r#   r#   ]   s{        
 
+F +Ft +Fjmptjt +F +F +F +FZ! ! !m
 m
 m
 m
 m
r   r#   __main__z Testing RealTrajectoryDataset...r%   z&scratch/parsed_moredata_pickplace_home)r9   r8   u   ✓ Loaded z samplesz

Sample 0:z  RGB shape: rX   z  RGB range: [z.3fz, ]z  Heatmap shape: rY   z  Heatmap sum per timestep: )rS   r   )dimz  Trajectory 2D shape: rZ   z  Trajectory 2D (first 3): z  Trajectory 3D shape: r[   z  Target 3D: r]   z  Episode ID: r_   u   ✗ Error: 
z<============================================================u   ✓ Dataset test complete!)$r   rv   numpyr   torch.utils.datar   pathlibr   matplotlib.pyplotpyplotrc   ro   mathr:   r
   r   r   rl   KEYPOINTS_LOCAL_M_ALLKP_INDEXri   r!   r#   r   rD   datasetr@   samplere   r   r	   sumr{   e	traceback	print_excr0   r   r   <module>r      s   3 3      $ $ $ $ $ $             



   1 1 1\ !"7"7"79J9J9JLbLbLb#5#5#57I7I7I"K L LNTU  *  2j
 j
 j
 j
 j
G j
 j
 j
Z	 z	E
,---''3Emnnn2CCLL222333 n3fUm133444Tve}0022TTT&-:K:K:M:MTTTTUUUB&)9":"@BBCCCVV4D-E-I-Ie-I-T-TVVWWWG(?(EGGHHHIF?,CBQB,GIIJJJG(?(EGGHHH3f[1334445vl3556666   A   	
 
E-v--	E
&'''''7 s   >EG G;!G66G;