From 88885a6a25c2e120f99957c58a1a1647bcd428e2 Mon Sep 17 00:00:00 2001 From: tangger Date: Tue, 8 Apr 2025 12:11:54 +0800 Subject: [PATCH] add gui & modify camera view --- .../__pycache__/main.cpython-310.pyc | Bin 0 -> 1621 bytes .../__pycache__/agilex_robot.cpython-310.pyc | Bin 8474 -> 8573 bytes .../robot_components.cpython-310.pyc | Bin 13668 -> 13714 bytes .../__pycache__/rosrobot.cpython-310.pyc | Bin 4618 -> 4618 bytes .../rosrobot_factory.cpython-310.pyc | Bin 1948 -> 1948 bytes lerobot_aloha/common/agilex_robot.py | 5 + lerobot_aloha/common/robot_components.py | 3 +- lerobot_aloha/common/rosrobot.py | 2 +- .../__pycache__/control_utils.cpython-310.pyc | Bin 6820 -> 7078 bytes lerobot_aloha/common/utils/control_utils.py | 77 ++++--- lerobot_aloha/gui_app.py | 208 ++++++++++++++++++ 11 files changed, 259 insertions(+), 36 deletions(-) create mode 100644 lerobot_aloha/__pycache__/main.cpython-310.pyc create mode 100644 lerobot_aloha/gui_app.py diff --git a/lerobot_aloha/__pycache__/main.cpython-310.pyc b/lerobot_aloha/__pycache__/main.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..02d7524837cc6ad3a43bfc0f465ba631f8f87b31 GIT binary patch literal 1621 zcmb7EOOM<{5O#Y!9?x?L;i2$I97P<4!%TwUvO*{jDJQf-NiO)(^tgLwys_O@cY9$* z8wqK-@*l#HU%-FqYfgLS#sR1*@3ISApe=u0Z+BH!Sx={<0M;+Be%;)OgWwf1*H;M4 zLwM~9Kn9YPL56!+vMS6%hIFs&RZ$iZ9hLnm&SIijImiaGFXJDAERh45z&=zXIeZ?P zrAp3%Y^;WAe8#|;d*=1Q^RcAe<5~kul&<4gg0o0UXOA;Do%U zB%jFGzkFy7rK2h9l6^F99Qt`FX74fFn zY#n!+7@G^wk!U|yETwSb6T=|Fc?9t6TmlyN+aTCez!UMW#)5K7~ltf3S z2j25HeR6qtKl}V5S1NW-X zt!`^+{edhj@m&Pa+M?b-)ts=;{XuILm#T5w?pQuS!Ae-=$4zUuSm$=$9|^AO~CtHwuIgP#NQ9+d32kC$6rsb#xrYu7Hn{QB9lR0jK>i&AgJKe$R1 z^`hDN=|+LUWit&Rko^!K86+(J4^PAlwoZ<>f4Ojmun}mEz5v1n$2dDM@jr$yU_TCJ z59ar#gBdyv&hQV=e+p9u@(AR298A@f9Ch<}{|zdS|ePc z>uVDMopWs*^7ZC0g$P%ssa_qc+)?QL!~a#l?vr*YHbtqvTkJ$t!p)33F<*nwybHNg zI~y#=(Y$Z4G(OIu!JrFo{Y?$dJvgB6Q)+aj9lVb|Q;16)K-v#LWG4W#WYn9n5u35- z4|5L<6VCB!IiF8Vf`Yg4w*f^l+IHQr?tX`Kmh4j7x}vlToMcWU9Xq>q;oE6py0F7L zJ!S+4KKY35|F3V?Ffg2(yP&<60`{k&Jd#f&1)M+6rUNbR?=!$t4~{Y6KSI+X`O&mXj5o$DON=fAryiS2_dPPcobEeN^8xo zp&zr6!8w%D5W(cskg%eM^pHz1G>4k}0fCz4RM4rHo?B?4r$FDVBwKdU?lN!Yz4v?X z{pRgszpFm4rX0i25Ik$&{N22N_EGAIhQAd*YZrPoj_hcl4p0qU&Rils_K6qcWD`^M zI$_IRoa4KQ#_oXyqwX6tvVv%wCT=1x0g*T$1GI_5b98i`g-n~+kAOGAnmiI{RHNi3 z0$q^^=^a2W(38yXK`~_OVL65Q05_yglXJ*U*=CR)APM#|JOW9T|0{ag^a&@^kzBcs z*nNy7>F#aF8)1(@j%kGW-irZ0jsqS5G|6`M5!^XQE-bMl3$&0Qg{1-lkNbHF`WXvJ zse9K6BmUwf$PN%jJz#vm>HZEWazTCsq56RMajy`}@c01pC^dNOF1kmSQ1@5BMf#Yn zz^;r-R{yIc;1(xEr)u_~!IDuQyX6G(nU-1qF zuW}7~a8&jn>_UDC@g(?v+x3HI%3lD#4*owU{vOx&0cHST^a(=~05_)mDxC#Cj3pqd z(}O8>u|xWx&?H@$*g;ii2I-sARb=le)5lBS2 zdX220g}JRo^Lu=fVfeXxdDObX+E&N$tW~GeX{>q;YB|j6t~b^jR?D>-?KSTU@hrWZ z{2Jh_EMrazGhGqyXU?d%rL7|FWVZ1MaV}e)e@7ZjDNaf8rW9{UaawF-E8;=+r}VN^ zYf@CD_&~IC7sS)t+w%`WEvgDG5d{}XmgqQz$%cYeMV!`7h^2f5m$vHpOHusC*7n@H z3Lc7o4)=|rwzhutvzGg5gSkU-r&K$3F@B9X?RtxTeo^(Da8F)Y^;%A+?lw9elMII9 zdFdnkp*UOqVg4)0qb-FiMR?JWeoj0}=fn{!U*I$6QBO0N{39aIuYZMW$zf?%pNnNZ4g*q${etzz5y_+QbAkmdjY delta 1361 zcmaJ>O=uHA6rS1LB>R(OlQjLWO-DTdqdEYnhee-tMKfMon z-!AJ8r@IxDZE$V}ag8>!I zOB8z%A`@W&!CK-Bj`|q$wQS~W1!i>~6JMCj3^`ymE?_ryx2}10eQJSq%&Ev(#53+WszBM-aQare(QBg$ zbYMU4K+uD#E<{zpKbCPow%}>NHNd}?ajgQc20RM*k1`(Jf`>?r3k0~8MZr9>^FQ0AIQmL`XE06Z%m{ z-SUY31hbVyl65DK<9Rlnys6!TKZn`I4QA`fH@Jg6Pi5Gbl#X@wJ=KO!vA?NzS}!M@ z<))9VRt<*C9`IRv`ef=UICWtOHwzN3gAo?lZ2AIjTK#f`tA&cFYJ3hS?Dg=C(UenOj7sq^QbxTFq(^D=p?xfV_>WT4r2yCY7WfQ`CeUG1Fxj zeUjZSEn#l=?cE(X+yMZgZSW9XJ=wU{Pm^((M(jrR)8NlL(%N85AJ52nN57Ox#ijlc zx=<`KtLIK54huBSu53J&)P{h?w#Y)CffQaQ=6>~@A?tC>q*cTym7v?sh&qJHhZb4{8XhU!#_z~I>o&v~Q zl8T%ccODQ!gCCEv}zNW_$Hbcz+tiyui5Vo+WkddC*71g?o%x5doX(+f*Al7 z_Y&Is0Ypo`lz#lsdRdJ>u%aGw7Am6=w{cKi=-!Ob$oBhd%K1w#)5ewlSVGR|L$T4U zoRu|N&(8Z-kuKKlD`8*yzYvfvU&cNSmgR1h?jYm<`@O8R2&*Zl+{t`_FxkSI1MTE6 z8w~tbWVnr)`hwL0-{UwE5{X?5b_fF)OU!*6B&3IZSH42%L&rg83q9fG2PMEk$&}J| zwkcE{GSV@HYlV{gr9SipXCgLj{ z!vwy(K~Ao-kE-vIDVC`Tlb5+C%Wl+cl)0wJu1Lv$%L>I}XsaKs!w3=rms%7p8DSW~ zh9DpuLC_Jhs5!ESiQziH<32x|3I71c{K#IFo38C8KHQ>3tf5aN?wJOypUFPm1>l4%oKC1Jy6?-rKr?GN8>tT z=)H*W7C)K`?9$r%^~gcrL@)@ziqf|M3!RMDUg8xq>}JEZo+b?X{*Th$K-#kiTx93a zdQ4gu&^Na%bjfegjBwRRaXEC7o`{- zhL?**ZE7b^u^UYhqOf~S4di3?S5w23Zn|nq2)bj^H0&I9ogi0Cv*v4rn2d3!sc*&6 zLbPa7C^wxT3HykN_VGH$q;SQQaL!lH$HC*8wIN5sdelZt3K`wX+1&Ire|na%PC17? zOK4AVS|D0NJ0JfS+5t}IrPt>dPz;N6!{V-U7SEy=YQ3|T!jdL!^cK%3zDAh0xq|rE zn&v;qE>_vn?Vng4&8D@mIuM3$flQ@vQVX+m%ghuiShVg-rscdvQO86k%Qk@>n8$x)wrs6O1kvC zq)5>8*BSXy_};z?ywDAraoQGyR)jVL{;A!A7Rrt{#=LDHrAj*Gz3@6pGH@AUhjC&D s7%0gybI?OCvomc~WE*?GtqE58b=$ho2ISzByv;81UIv-~+t}Xt50JMU(EtDd delta 2343 zcma)7U2IcT9KWZ%z5VVA>(;Jp-NtaN?)wWwAajU9AgRppIWoF#z3aNQz2)3n=dfcS zkq{pQcLWgvY)-)W@qm4hm|%P}@rC%p!wq6gOpNhGV&qAp|NrT_0%Bs5{_g+$KYr)j zbAIQ4XD5ChcPdV&O@QCYP4{Do8&0>G=qzBqM;h5fvleKCh;EL`qmfKPqbp&Qhf=g! z*lEkG)SNeaYkavL>nTHrZjGdr)KpT*s6mmI@F;+8qKcZH(yj4STnnd^s2nuV6&T&j zZj_ZWhvUY!tvGxM075-jW9!cht0Oc%EYqMlpZy~GElD*N)~3?3ZW)Or62p;^3A$XY zGtqgo8HWWL>#odjBCd`Bi>b&$w?PUo({sN%-zIH*N^?Sq#Z#k-Zi&gT&PY^jh%pJxj^b_WUFiMeg)bJ7m9I}y4Nx)B@*9)xuWy#Tr~p3-!2 zSWyym2RcIt{S0=rY6F;N?|OqaMwr30Q;e0wIrh>%1r)eT>PLUI4GG?x{f{P5BT%x zn8G#1mC8lBou&P;>SB01=J4>?N zG_a$UlVm3MpwdF%&iq}qp+c83aFp%q=t(Sm6X7j{6Ra=LRLi$x`98oze5Es(z?Y|R z@+|uzaEqK`>FOGClzXPx)#~*!*A#gg^tE`G>6^~5f=E- zTxRDR@2o`*I**Vi04qd40xWbgUVD{STw#k%oBO%g-2X>u?<4Iw1TL}*XgwjV%jmnp zKw7nD!2ARLDqcWrXDbc1Qk&@Vm)rfefS;a0(dAAye?=;+F4Cf{40O2)8Tso5?nSL zu^B7|m%%}7L?oioVPs-g(0Es@XKh_on-(DnZg~ehp85&hivg4qZ+N(M2G+rimH~#}f CQx&2B diff --git a/lerobot_aloha/common/__pycache__/rosrobot.cpython-310.pyc b/lerobot_aloha/common/__pycache__/rosrobot.cpython-310.pyc index 9f5fa9360352140c89676bdb7ff593a9c675f0a5..af49d76ad30e83554bc5d8ef512bf7e7cc6137b1 100644 GIT binary patch delta 26 gcmeBD=~Cg!=jG*M00O%g8@VEQ8SOSF^Dbur08jV^O#lD@ delta 26 gcmeBD=~Cg!=jG*M00PEc8@VEQ8C5nX^Dbur07@GL$N&HU diff --git a/lerobot_aloha/common/__pycache__/rosrobot_factory.cpython-310.pyc b/lerobot_aloha/common/__pycache__/rosrobot_factory.cpython-310.pyc index d058a8e54a09126695051a28deee8c9407c5fb9b..786d65ce2804dcaee3e429259a4a81dbeaa3f18c 100644 GIT binary patch delta 47 zcmbQkKZk$AR#wKS&D&Uo85tEOpJMZ3RG%!$t|Ab{QIua?l%JGeQpCZ)zz{XrlRX#! DQ@{;f delta 47 zcmbQkKZk$AR#wJao42tFGcqbpKE>w6s4-cTT}9v)M^S!pQGQZ>Nf8GF1H-M!p6tN@ DbNUY7 diff --git a/lerobot_aloha/common/agilex_robot.py b/lerobot_aloha/common/agilex_robot.py index ce63d76..e7d11c5 100644 --- a/lerobot_aloha/common/agilex_robot.py +++ b/lerobot_aloha/common/agilex_robot.py @@ -125,7 +125,12 @@ class AgilexRobot(Robot): return None, None if any(len(q) == 0 for q in self.sync_arm_queues.values()): + # 遍历字典，检查并报告每个空队列 + for arm_name, queue in self.sync_arm_queues.items(): + if len(queue) == 0: + print(f"{arm_name} arm not connected or queue is empty") return None, None + # 计算最小时间戳 timestamps = [ diff --git a/lerobot_aloha/common/robot_components.py b/lerobot_aloha/common/robot_components.py index 7669f71..1d3957e 100644 --- a/lerobot_aloha/common/robot_components.py +++ b/lerobot_aloha/common/robot_components.py @@ -365,7 +365,7 @@ class RobotDataManager: self.sensors = sensors self.actuators = actuators - def warmup(self, timeout: float = 10.0) -> bool: + def warmup(self, timeout: float = 30.0) -> bool: """ Wait until all data queues have sufficient messages @@ -391,6 +391,7 @@ class RobotDataManager: all_ready = True # Check camera image queues + rospy.loginfo(f"Nums of camera is {len(self.sensors.cameras)}") for cam_name in self.sensors.cameras: if len(self.sensors.sync_img_queues[cam_name]) < 50: rospy.loginfo(f"Waiting for camera {cam_name} (current: {len(self.sensors.sync_img_queues[cam_name])}/50)") diff --git a/lerobot_aloha/common/rosrobot.py b/lerobot_aloha/common/rosrobot.py index 30f80dd..fe3bc26 100644 --- a/lerobot_aloha/common/rosrobot.py +++ b/lerobot_aloha/common/rosrobot.py @@ -112,7 +112,7 @@ class Robot: - def warmup(self, timeout: float = 10.0) -> bool: + def warmup(self, timeout: float = 30.0) -> bool: """Wait until all data queues have at least 20 messages. Args: diff --git a/lerobot_aloha/common/utils/__pycache__/control_utils.cpython-310.pyc b/lerobot_aloha/common/utils/__pycache__/control_utils.cpython-310.pyc index 7fa908c5bf58ef09e66d7c05d349e58b3163db11..b2af2e15a85ad1a1da25170ff7b513e5c1ddf816 100644 GIT binary patch delta 2059 zcmZ8hO>7%Q6rS1j&)Tt_#Q%2E=C@7VmQo6V0#S&XCjAQyO47fjm2o{2XOp$pnVtM< ztOzx|AR)Eg8v+zd2&sUCP>#KUgt%7RRv^I%6$B>^NWJi8(q{oOxLZDkiI248-e*Pm_(n4DZK7&Yvws|^A+ohfjkQiy*^y_ib zM%p(~UQduXNdQTbBuNT%ilj*fbo#qL-^YGenOSh1i%J1U9j@F2`B>?|t=oSqqiR9|=1LVf)j1r^Fi={F*tXh{d8F&Q`I3EU%#}0vCo_-XD{g3k~@StUp`?rEkbPwCM zb3OR=#t={3k+UJvO4>G&h@C}ua3ju>A`Uts5+uCfyNB5jX>TM+ypiH5p5Ds9$cjW0 zH!(z|cWfdd(jCaNfPX$GG9r~kA}x|hL=v+q*zzR#vC_!$4#|EN-J0CEsWftpyvT`k z9yL0oPNPPrh>3g-ZFM$L;{ZRfbx?qu?`kUEtXmuqNs*dWV17_^$Rv#}(IrwD)aVwS zGF_ub^oS(Pd&NPK1z-1F1uVUHm04v7i5>#6qKl+66*Q|f`XEIQgyjXG`CgKNi_A5H z7%lz|o<5Qn*`|UK;C`@Uj6{&&B@9)WT zvqTks0{ID^Mj|Be0u-JtImbge+_60#>Q!*TFUE67ry-o=anTE~FR<%C-k?aTs zL^oRRE?YIuEVF1@R{g|~X_ZT+J?y+3iO{bwLdFP)Co04iFeKt@b^62)vnQ<}Dlp*t z?VWq~rwu5vdyK6!bj zcz$eZ`uy09;`D`gd*7QFyY2>zsy%0|xWSc@!7tG@H*|Z+UMtF3;gG|luV;^V+o#L4 z#%*ep;4uiFgSX|`bdDO8Ij)B) z%pk?GW!ldmvsP>04v z(Sxhc4Xa=vdW;*~q{X^r8}bd&eS{b77uzSYzfZ;?tG(&`5xX#=;jrRUd^oAb@j{c-Xm;=sx0|71M_byCKU-rP%9l&8ijQF|_SpozD`*ff{%VLFP42 z(6lNQqgIi($7%rdKz*$SH_oWjm+VQ&*d3;-MBc-ib{;3v_y?z+EMUI1Xtx+k3hE2 z&e`_*=bE$~1yR7PB$dXzRlBinjMfd?BDBbjDus*lsP+it+7~Y~#%S4MBvgF1G~|qA zyYQ;>LH3uy(=z@|nMrYjHpJD4?S{P$@-n=@#bzM%C1)|0Iw?^>yF5OAS$>5Kvm}+r zmyzf#DJG=2B*mEXP0qL`i^{zJ|5-q;YIEkDlX2;gHDs$&Nb-HDcgyT*Lsg3`nRFN) XDr66X(2$0m`TVQFhBm0RNBaK%i2Jijo|slQ60A7w){Mgi=-#!D9gFmJ#8!)g$nH0vV`hJ#J=D2K92=Uet%O z>!8o1ZvU(%#pOYt&r1awY? zMi*PiZj)xKO>D0-z%C?ShI`rV#0$nSmnZl)$v^)C_D$l$)N3TZp=^RX(0`p>Q9fBy zN%odFsi+5K)&Vt+gIh4@L%l&3^#wVSqka6GCwZDh*{YBgXr2~|0Od-67Ahbpl8AFl z3QD9Xw2yFq3|H;_(WUe@hWEw}(I6&|`ph*g^@s8l5?4x8E7Pmf4 z4oPqydz>1DBTUzaFP3;W7XSs9)c_(uA+J!P@UC=SIu9BUXSg^TPeTf&122X0HS5Oq zxSq~T8+UAwaa+tr;-JLt=_7}4Dgf9=VjSG$Oq}I{FpjJ)S+swK05s}heBSX| zB-CbG&XRQkV-Mqftd_o7+A9Lbg+Ze2b>{5_JSaFr?0&lcowVS~j@1ceyF)_RwL4+b z>MXeyZV@Y#+smV&+FmA;p5x&#&IR2YVS4i1nR7GEnUlvxkDZ(e6SnIwdMlykddt=Y zyMx$-QG~;6);L^xE{ba{TOH!#x`?|^fc@-&@kXL1$PxC7F|wxnZERVc=8BETqM3ej zTD0tiMPh0^Y}0mI3zl!D+}3Kd?Kw+s#}8xJTk%6>=P27~6~8E)ie{Vo3eUw`+41@0 zwCyvLxhXHtb4+cu*&X+MTWpD`-Dq}q=olhR|Jf1gyoUsPPSf$c?g<{of1mx>bL!QY z9D}N?!wl45MpmH&3$iXpT&PBGM%H9a5*oUsOB(0J;LvAycI!%J48q3Nz25Jj(%>Ga z@CDY*PQ#1r%j_XI$9~Vw!6R%sSKlvc3!(K;!f$a3n*uz&J*wAf@}>9v?-|QA)i#$l z6r5&%=LVq3%K5>JU~L;*6_n0S=d18@HlJ@md+Sa<5B2Ci#C+m)`9*(JC=Z_!ii-lA z*)`q`iRU=)H}-3x20{WXC}eYKar%No}HRHE1pIkBjGyopBL020giGI zYT}w$tHiOdvM(xTv}ZWl#0epmeS)Lnu`NtQ8?h^Zi2a{@o))N%yB=C{tO=eN@W(k& Q0duMcl_hmd?bS#A1<^jHJpcdz diff --git a/lerobot_aloha/common/utils/control_utils.py b/lerobot_aloha/common/utils/control_utils.py index 11ae4d9..990d60b 100644 --- a/lerobot_aloha/common/utils/control_utils.py +++ b/lerobot_aloha/common/utils/control_utils.py @@ -3,6 +3,7 @@ import time import torch import rospy import cv2 +import numpy as np from contextlib import nullcontext from copy import copy from lerobot.common.datasets.lerobot_dataset import LeRobotDataset @@ -134,42 +135,50 @@ def control_loop( dataset.add_frame(frame) if display_cameras and not is_headless(): - image_keys = [key for key in observation if "image" in key] - - # 获取屏幕分辨率（假设屏幕分辨率为 1920x1080，可以根据实际情况调整） - screen_width = 1920 - screen_height = 1080 - - # 计算窗口的排列方式 + image_keys = [key for key in observation if "image" in key and "depth" not in key] num_images = len(image_keys) - max_columns = int(screen_width / 640) # 假设每个窗口宽度为 640 - rows = (num_images + max_columns - 1) // max_columns # 计算需要的行数 - columns = min(num_images, max_columns) # 实际使用的列数 + + if num_images > 0: + # 设置每个图像的显示尺寸 + display_width = 426 # 更小的宽度 + display_height = 320 # 更小的高度 + + # 确定网格布局的行列数 (尽量接近正方形布局) + grid_cols = int(np.ceil(np.sqrt(num_images))) + grid_rows = int(np.ceil(num_images / grid_cols)) + + # 创建一个大的画布来容纳所有图像 + canvas = np.zeros((grid_rows * display_height, grid_cols * display_width, 3), dtype=np.uint8) + + # 在画布上放置每个图像 + for idx, key in enumerate(image_keys): + row = idx // grid_cols + col = idx % grid_cols + + # 获取图像并转换为BGR + image = observation[key].numpy() + image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + + # 调整图像大小 + resized_image = cv2.resize(image, (display_width, display_height)) + + # 计算在画布上的位置 + y_start = row * display_height + y_end = y_start + display_height + x_start = col * display_width + x_end = x_start + display_width + + # 将图像放置到画布上 + canvas[y_start:y_end, x_start:x_end] = resized_image + + # 添加图像标题 + title_position = (x_start + 5, y_start + 15) + cv2.putText(canvas, key, title_position, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) + + # 显示合并后的画布 + cv2.imshow("Camera Views", canvas) + cv2.waitKey(1) - # 遍历所有图像键并显示 - for idx, key in enumerate(image_keys): - if "depth" in key: - continue # 跳过深度图像 - - # 将图像从 RGB 转换为 BGR 格式 - image = cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR) - - # 创建窗口 - cv2.imshow(key, image) - - # 计算窗口位置 - window_width = 640 - window_height = 480 - row = idx // max_columns - col = idx % max_columns - x_position = col * window_width - y_position = row * window_height - - # 移动窗口到指定位置 - cv2.moveWindow(key, x_position, y_position) - - # 等待 1 毫秒以处理事件 - cv2.waitKey(1) if fps is not None: dt_s = time.perf_counter() - start_loop_t diff --git a/lerobot_aloha/gui_app.py b/lerobot_aloha/gui_app.py new file mode 100644 index 0000000..a4aa87f --- /dev/null +++ b/lerobot_aloha/gui_app.py @@ -0,0 +1,208 @@ +import sys +import numpy as np +import cv2 +from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, + QHBoxLayout, QLabel, QLineEdit, QSpinBox, QCheckBox, + QPushButton, QGroupBox, QTabWidget, QScrollArea, QGridLayout) +from PyQt5.QtCore import Qt, QTimer +import cv2 +from PyQt5.QtGui import QImage, QPixmap +from main import get_arguments, control_robot + +class ConfigGroup(QGroupBox): + """Group of configuration widgets""" + def __init__(self, title, parent=None): + super().__init__(title, parent) + self.layout = QVBoxLayout() + self.setLayout(self.layout) + + def add_config(self, name, value, widget_type="lineedit"): + """Add a configuration widget""" + row = QHBoxLayout() + label = QLabel(name) + + if isinstance(value, bool): + widget = QCheckBox() + widget.setChecked(value) + elif isinstance(value, int): + widget = QSpinBox() + widget.setRange(0, 999999) + widget.setValue(value) + else: # string or other + widget = QLineEdit(str(value)) + + row.addWidget(label) + row.addWidget(widget) + self.layout.addLayout(row) + return widget + +class MainWindow(QMainWindow): + def __init__(self): + super().__init__() + self.setWindowTitle("MindRobot-V1 Control GUI") + self.setGeometry(100, 100, 400, 600) # Adjusted window size + self.robot = None + + # Get default arguments + self.cfg = get_arguments() + + # Main layout + main_widget = QWidget() + main_layout = QHBoxLayout() + main_widget.setLayout(main_layout) + + # Left panel - configuration + config_scroll = QScrollArea() + config_widget = QWidget() + config_layout = QVBoxLayout() + config_widget.setLayout(config_layout) + + # Add configuration groups + general_group = ConfigGroup("General Settings") + self.fps_widget = general_group.add_config("FPS", self.cfg.fps, "spinbox") + self.resume_widget = general_group.add_config("Resume", self.cfg.resume, "checkbox") + self.repo_id_widget = general_group.add_config("Repo ID", self.cfg.repo_id) + + # Config file with browse button + config_row = QHBoxLayout() + config_label = QLabel("Config File") + self.config_widget = QLineEdit("/home/ubuntu/LYT/lerobot_aloha/lerobot_aloha/configs/agilex.yaml") + config_browse_button = QPushButton("Browse...") + config_browse_button.clicked.connect(self.browse_config_file) + config_row.addWidget(config_label) + config_row.addWidget(self.config_widget) + config_row.addWidget(config_browse_button) + general_group.layout.addLayout(config_row) + # Root directory with browse button + root_row = QHBoxLayout() + root_label = QLabel("Root Directory") + self.root_widget = QLineEdit(str(self.cfg.root)) + browse_button = QPushButton("Browse...") + browse_button.clicked.connect(self.browse_root_directory) + root_row.addWidget(root_label) + root_row.addWidget(self.root_widget) + root_row.addWidget(browse_button) + general_group.layout.addLayout(root_row) + config_layout.addWidget(general_group) + + recording_group = ConfigGroup("Recording Settings") + self.num_episodes_widget = recording_group.add_config("Number of Episodes", self.cfg.num_episodes, "spinbox") + self.episode_time_widget = recording_group.add_config("Episode Time (ms)", self.cfg.episode_time_s, "spinbox") + self.video_widget = recording_group.add_config("Save Video", self.cfg.video, "checkbox") + self.display_cameras_widget = recording_group.add_config("Display Cameras", self.cfg.display_cameras, "checkbox") + config_layout.addWidget(recording_group) + + advanced_group = ConfigGroup("Advanced Settings") + self.num_writer_processes_widget = advanced_group.add_config("Writer Processes", self.cfg.num_image_writer_processes, "spinbox") + self.num_writer_threads_widget = advanced_group.add_config("Threads per Camera", self.cfg.num_image_writer_threads_per_camera, "spinbox") + self.use_depth_widget = advanced_group.add_config("Use Depth Image", self.cfg.use_depth_image, "checkbox") + self.use_base_widget = advanced_group.add_config("Use Base", self.cfg.use_base, "checkbox") + self.push_to_hub_widget = advanced_group.add_config("Push to Hub", self.cfg.push_to_hub, "checkbox") + config_layout.addWidget(advanced_group) + + # Control buttons + control_buttons = QHBoxLayout() + self.record_button = QPushButton("Record") + self.record_button.clicked.connect(self.start_recording) + control_buttons.addWidget(self.record_button) + + self.stop_button = QPushButton("Stop") + self.stop_button.clicked.connect(self.stop_recording) + control_buttons.addWidget(self.stop_button) + config_layout.addLayout(control_buttons) + + config_scroll.setWidget(config_widget) + config_scroll.setWidgetResizable(True) + main_layout.addWidget(config_scroll, stretch=1) # Left panel stretch factor + + # Remove camera view panel completely + + self.setCentralWidget(main_widget) + + # Robot control flag + self.is_recording = False + + + def browse_config_file(self): + """Open file dialog to select config file""" + from PyQt5.QtWidgets import QFileDialog + file_path, _ = QFileDialog.getOpenFileName( + self, + "Select Config File", + self.config_widget.text(), + "YAML Files (*.yaml *.yml)" + ) + if file_path: + self.config_widget.setText(file_path) + + def browse_root_directory(self): + """Open file dialog to select root directory""" + from PyQt5.QtWidgets import QFileDialog + dir_path = QFileDialog.getExistingDirectory( + self, + "Select Root Directory", + self.root_widget.text() + ) + if dir_path: + self.root_widget.setText(dir_path) + + def get_config_values(self): + """Get current configuration values from UI""" + self.cfg.fps = self.fps_widget.value() + self.cfg.resume = self.resume_widget.isChecked() + self.cfg.repo_id = self.repo_id_widget.text() + self.cfg.root = self.root_widget.text() + self.cfg.num_episodes = self.num_episodes_widget.value() + self.cfg.episode_time_s = self.episode_time_widget.value() + self.cfg.video = self.video_widget.isChecked() + self.cfg.display_cameras = self.display_cameras_widget.isChecked() + self.cfg.num_image_writer_processes = self.num_writer_processes_widget.value() + self.cfg.num_image_writer_threads_per_camera = self.num_writer_threads_widget.value() + self.cfg.use_depth_image = self.use_depth_widget.isChecked() + self.cfg.use_base = self.use_base_widget.isChecked() + self.cfg.push_to_hub = self.push_to_hub_widget.isChecked() + self.cfg.control_type = "record" + + def start_recording(self): + """Start recording with current configuration""" + self.get_config_values() + + try: + # Create robot instance with current config + from common.rosrobot_factory import RobotFactory + self.cfg.config_file = self.config_widget.text() + self.robot = RobotFactory.create( + config_file=self.cfg.config_file, + args=self.cfg + ) + from common.utils.data_utils import record + record(self.robot, self.cfg) + + self.is_recording = True + self.record_button.setEnabled(False) + self.stop_button.setEnabled(True) + print("Recording started with configuration:", vars(self.cfg)) + except Exception as e: + print(f"Failed to start recording: {e}") + self.is_recording = False + + def stop_recording(self): + """Stop recording""" + self.is_recording = False + self.record_button.setEnabled(True) + self.stop_button.setEnabled(False) + + # 模拟ESC键按下事件 + from pynput.keyboard import Key, Controller + keyboard = Controller() + keyboard.press(Key.esc) + keyboard.release(Key.esc) + + print("Recording stopped") + + +if __name__ == "__main__": + app = QApplication(sys.argv) + window = MainWindow() + window.show() + sys.exit(app.exec_())