From 70bd021e96c178a180b502f00043c1901b1a4f45 Mon Sep 17 00:00:00 2001 From: VangengLab <145450894+VangengLab@users.noreply.github.com> Date: Mon, 4 Nov 2024 16:45:48 +0800 Subject: [PATCH] Add files via upload --- src/modules/__init__.py | 0 .../__pycache__/__init__.cpython-310.pyc | Bin 0 -> 178 bytes ...pearance_feature_extractor.cpython-310.pyc | Bin 0 -> 1856 bytes .../__pycache__/convnextv2.cpython-310.pyc | Bin 0 -> 4790 bytes .../__pycache__/dense_motion.cpython-310.pyc | Bin 0 -> 3659 bytes .../motion_extractor.cpython-310.pyc | Bin 0 -> 1661 bytes .../spade_generator.cpython-310.pyc | Bin 0 -> 2051 bytes ...tching_retargeting_network.cpython-310.pyc | Bin 0 -> 1991 bytes src/modules/__pycache__/util.cpython-310.pyc | Bin 0 -> 14833 bytes .../warping_network.cpython-310.pyc | Bin 0 -> 2204 bytes src/modules/appearance_feature_extractor.py | 48 ++ src/modules/convnextv2.py | 149 ++++++ src/modules/dense_motion.py | 104 ++++ src/modules/motion_extractor.py | 35 ++ src/modules/spade_generator.py | 59 +++ src/modules/stitching_retargeting_network.py | 38 ++ src/modules/util.py | 452 ++++++++++++++++++ src/modules/warping_network.py | 77 +++ 18 files changed, 962 insertions(+) create mode 100644 src/modules/__init__.py create mode 100644 src/modules/__pycache__/__init__.cpython-310.pyc create mode 100644 src/modules/__pycache__/appearance_feature_extractor.cpython-310.pyc create mode 100644 src/modules/__pycache__/convnextv2.cpython-310.pyc create mode 100644 src/modules/__pycache__/dense_motion.cpython-310.pyc create mode 100644 src/modules/__pycache__/motion_extractor.cpython-310.pyc create mode 100644 src/modules/__pycache__/spade_generator.cpython-310.pyc create mode 100644 src/modules/__pycache__/stitching_retargeting_network.cpython-310.pyc create mode 100644 src/modules/__pycache__/util.cpython-310.pyc create mode 100644 src/modules/__pycache__/warping_network.cpython-310.pyc create mode 100644 src/modules/appearance_feature_extractor.py create mode 100644 src/modules/convnextv2.py create mode 100644 src/modules/dense_motion.py create mode 100644 src/modules/motion_extractor.py create mode 100644 src/modules/spade_generator.py create mode 100644 src/modules/stitching_retargeting_network.py create mode 100644 src/modules/util.py create mode 100644 src/modules/warping_network.py diff --git a/src/modules/__init__.py b/src/modules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/modules/__pycache__/__init__.cpython-310.pyc b/src/modules/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fef71516bc08d73f437faa291f4d358337aa4d28 GIT binary patch literal 178 zcmd1j<>g`k0{1eFbP)X*L?8o3AjbiSi&=m~3PUi1CZpdO3W;YFEi3FE=tzV i%}*)K0ji0Q&&&gg#_JVS-r}%eq)RYjHgi%dv%7ZYbe+uYx)}Y4n9ICd#5_(0-i_7u z897C5|2hr69}amWqCDmyzfvNKRUvkd+acq{_LI0DX3%inP?x+zJyQsB>vt_?G5Z#E?RUsVT?dZp z=%)`SJ)R}@!{WTLzqEM+^cr*%L;eLqA%#nHLndU3Fe*W(R!KnGr3KO{ZIEv1fb>ci zq_1jVMNBTtdT&;PGc@*9y|ixd5&Hd;aZNSAi!i&at)X?;dmKE>FlQ6ZXX~XuX{twM zP}Uc_O)*C2yQ(!=DgCL*<*{lFR?B(`*A73oinpfL2d&AIDT4J4i0f5xSI4yy4K^m5 z%u`!s?FMfmwOu;Q2S4k*H3jgTGlrmAwUaivGX)T)GHoI%-H0~P^*3DPJf$+Z=Gq%Z zj3xQm9WmCGg{|vJPUC)*gO${cVx;az!c|n zLTRU$2&uFKBshbdvLx58L2=H6V-96brh}ggHcI)+1T5>z0o}tQzX0sIlE+2Pv;+K! z`#Lz~=Odo0BuaGyg3i*T2TTW0Fq&1nX$uM=T-d+e(6;1hPpq1VYeuXa@x+J?BQ`|hm)$__h+u^y%>GKx<+Z7Mp9Dx0sWnHL{X?tG*Sgi_f$6Q92Qyc?W>yzh`y*d+{^p{*&n_Yk0H57(B&6K z!S2#;mDr?#SMdhk!VSDm8Zg?y_CMF3{x?VW7nq36eJq-RlYz`3blC?{q6tuRipNA@ z1tc8+J{+hb?!j=hphTtn)C2i*rsdsAoZ zgLTbN4m2Mqf=}9k*n)^+8-%t}o>%uFJ~l(Af}OT65`Ot`e_{vL+BNaHF=@;mP+B2e zTfh?WrI{zMrN*(weT^@*)uU28gXx*rGX`F-5SNi)rmk`YM1WVx3J&liyg>qNnRPQ zebgW;Th&~T} MO?(?bgg2~z0cSn%O8@`> literal 0 HcmV?d00001 diff --git a/src/modules/__pycache__/convnextv2.cpython-310.pyc b/src/modules/__pycache__/convnextv2.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09ecdbad2189979d7a8ef4b2a25e8f214b7e44a6 GIT binary patch literal 4790 zcmb7H&2JmW72nw}E-8tYWhJ)LG*~BT6E?OaD?z?$xM}RhP3x#Zov*SDw(FfCwGz2Y z%?@Q-rO^U1;P@Z|MSIMF3VQLS=%x3bdg@;=x1xtYhn`!sK+@lvrAW$74qbBJ&b)au z^JeC~-+PnK&DAVCZ*4!{_7^SdA5@usTvXn`TeLv9)wDR{c84{Y`PyBl=`f2s-0iqs zujzIDrr)hJE8U>g8#sinC!aF9beu#zS$lCq-fUbAb`MF?C$39bmE6+5rP?xVm~Z zz9r-unssnl=Dl}gwXWxs>}N@oQKwF%ohv_G)5j9%l|h5;)kvk;t%Qb|erPx8DZT?D zEbOAyv^i@!ymHrSG1|vF)uzim?qfH+yOwJ;ee6h}y*E2~vbl4Adg+F>RI7z}zAxLw z>!w6Co_53MvD(J#;ab0YT}ZPRS+B3cWIfI@(J7WD!dOz&!x_s`90KoaM$gfEpVqi=B^1!e78L$+@*OW7OFfDMfG zfI55Lz=~Yup^R1U*{J!-9#&BI2W-*WCr$uPb>Bj%re;e1E<g^QJSSHinN_&+9%^&<=W>r$Q&fiqFu5R z?O#nZ5lc*Bn$BzY$`7u-uN~N{^hwNV%j~dqRnk;Cc$jSoNa)+G7#0qG@Yh)}5!NLms$UdV1=33-ZIo*?ohk#7=tibx33 zaO4?$G`p$2lw3-Fi<%7iHZ@cLBsXHoNve1p#Ki{9<30S%v%2lFow;(eM$!#dT2+P2 zDT})T=2VNK@)S^>jiRl7+!?i0qKM~76q)moFVM7eqw`Nhr-O!xQqQ1Ol4ddhIfl2O zLIAz$mHC?CCSV+S#z^%eu#X^>-3H!&#z5WA824JPK`uggLN0^rO%HfJ0|fK=ET6*{ z=;m6@04Gq_HS&)VyyGu$8e>EYKUn=LEK9S9Fst}FKuf!ETTEbW0#YxI5iaX@OWzP* zF^X{-G*AFB<+PP1CJuz_BIZ-5cYbB1ggpbXgo_wX^?My*u=WVpUMILl+)|;4ll8Dr zP;zRhpTBrE#Q!zWkH>1lnaBFQC1TZ=f{uHRr zJGBGn%K&7&@PS)8BYETjVe5!{@JnL+1tk3qEG4DLV!R~rz&m`X7bP+DE&c|%Nw^tM z_^_2TbPfD_DEZ3S04RI{usLmU8$V~yRl%rkLhqilZ{PL@)`mB%4*Y$FuvO!gU3S^} z`00C&83Xd4IOso;r!h+oQ68MM&{o9{V`Hq!Xsim2wRLW+&rIniYIZos*}yG>*1mJu z+GY2U7TlwG5ymj@8Jc&)ziZuPckR1Q%Y|dj+zC`|V3Qssv&-&Vw*$;H1FZv^=@YAL zeZ{(GZKdc>tDG}h*G;HBHt>e?{1~5yj#X1W&KFRwLBr!C*~*|c3d0Li+`xp}6BEv$ zYB-CJa2|Vv^Z10rE2h`UY0h}QQxv)zCrZec3pm||Eobr3exYQ_g(NJ~^}~!NYNfJJ zkU38586K-RDs#FC)XV44mcT~F4+Q`MyK8t0if(Nr8Xx;6AX)p?Z|yClqy|kJOuMNP z-G;60r4`wL>=eox=&xPqq=icIZtvoIG#(AUXze^VfxM-|*tm272oZORizpoSc|bIG)My6)tfDzkppw^yUD%xOJ|I-jILlraN4%T{V2X9i;#BP(JgIa z?{u__f--De1sXGDWDgEQj{G4h^ac%DA17aC%e23V5Z^A4hruHB{KM-)YDgN(%QV1S zL`)RENF@`B$f^&Bs;%ElfDX~U_6R#*m#|my=%)S6q2hAk1lq5HZYHE z_>?|v1RQvX%N$(9c%-}5{?+0iul#-I_uqp{IECUSscKYpEiG^`D$WvN&b~ZQIuuEO zt^G10z>N?)bQPH;39CrE6dquvXGn^5A9ky`G60R@Ne4nEgIPe@xrj~LIW@g_lZuZ- zkL&<~$0PxzlTf(K8g!kfM0A4|V+ z6q5wI@_iy75c!Y@fsgzd5dzUh&6q{GbF?XPlL%RX>=Ma{3on~i_A|8I9@XW`M^ZdB!2be89cED>2ILcT29M9h+CBDq3j zqQNs%HENtki5g9;p&Z;DZ_0&m`zB$419KjgaPKR?D>gdi|n7x_0cS-6cX z-1!v8`3H92+@o*|xaX`wPQC;Jw;IEk(FivUngaY^;G;cUF+^udSED9Hp5Ngu2r~$Q z8ufCqKcPN>!QrJWDDmD>@s#b`z`QfF$^QVD7L4Tp literal 0 HcmV?d00001 diff --git a/src/modules/__pycache__/dense_motion.cpython-310.pyc b/src/modules/__pycache__/dense_motion.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f588c827c6a176678d56609f393380e26857a5d0 GIT binary patch literal 3659 zcmZ`+&5s*N74PcrZhvMnNl3Cp76K73B*SKdkRaMcWRH=Y(tyNGl}%pt>PW^Dl6L->bG~GP6lrey)#KudC|!e(yaOola=rd31Dp z^3N-V@gM5E{n+UI9N+RO1UE(oXWX2z%p94F+SbgrF z@{*>jw8#fzS!9E99?P>KDWG18yo$BPBs~)O;MvQ;;XEpfg-paC&N=nCOpjnPJi$?n zAya0a@1HPgTPpY^R422zEY)rnABrd`3dz$vu0%AEDOat-`MpWJD9bd?6Eg7jL$^dv z`3ZzDaDHQC^7dE87`Mf(*TKl58+mPv?B)(e4)=M0zRN@2!msz0VH+bK_t8;1pV5iG z&9gK1~Asf<{o=#n-~O{0Bs3RID~$nJaS*C@QOFwawWzqxMMd8nsoM+k zt)fU0Xev0^;x0|+Iz^Jq7IeTUi|38;dt3#l!^1$?rI?MSkC)393F4!?L3`VSqwHb2mhxnF%gQ(fm_=B^wm(J2% zdP^T}?$%!IkFDC`?2Wk$>R<)e=4La_t+h*IHK;(YI<;sHAB=y?ScY{tHMn`*Shni0 zw!W!z%~-Z!(FQe|b&zCH6CCSDk(MnI*+o`?c&~$GCWlr`gbh*AZ6I z!Jx8_Qt`Y&)kVC5qbsK>WOAUaB(8>Cd5g5X4d5y~i&iQ)jbqAuRw@=L#uYnI?78Zs zz@kd4mraQ3U5r`mP+i&ty{IZS(NJ}p@VP;-l;6dx8}U8b#Fa#1dm=io^=(X*gmjqCtsQ2wwzg>3Y{fWRZ>7l*;(#Xt8p6}& zvxZO9KJY@YYUU7v%-SJuSUZ?A&(Ae8+#EY=wlp=W(4i3c#?rp@mZhzou?t=Jl*U^F zeeyp%j{toqo~4rV>!&0>B5{wz$0Y8P_%Vd)Zezj5dwLsqxELe+2!0X8Ya0Nz^V@KG2gXWbYoFa_ zAAqcVNZUOA@ckHGG;xlxgdoK(&Dg)7K_~(p_^Z12yv~;ft;{|MmFa$FDyfnyQT}tI~LrFp$rqqgYZl7m|vA23bl1^7|0LhQO4RXXOtelmk2q zB!lYZTquX@AG0cy%f(R&BXqTr3V>no7OZ(R>Q-d|xu)g3GQc4N2$DxwYYVa1Sb6A$YZKAfSow8pGt;gD6dmol1$0qA zHPLp3(t}$~ln1j>e5~4a2k%4ivFg-aw0f0=Qly9SgVyxwULA1ffl>EXJ&Ge330Hm8 zUpC4aM)emWJ2FDu+cl<_^eVRO*L|&{=ednHVFY$U*fFL%@~?IGd`7Cfo1Gymt?lKd zDp*~{j&|!yYkKESy@S;)tfu(U-j|OERB|x+fCNQEd6fkD;0YxQ?YJkD7FDa^o-S-v z=d@~|Vq$ou@lPrO@ub*&xLC>Sw` zBOwwRr8VDl%|2Rvj3Ff6cn4!PW{2!4DmRppd7OzTQelK_N`GI_Zb#9Vi+I+|=;G)i z5PzhTbT-vlB4#rbqh~b#4jqRwv94FCq$t}EJ}B+?{JtOd+j?f$R$dhGBEf16t*&1$ z>#E{zp5GlWa{PlT@_44gv-ZYI7F9aa1+XqpbdXc-7y1a=_*?tYJJXec|DZYKC*-Pn PQ?}W*ac{W4tJePk+j_;% literal 0 HcmV?d00001 diff --git a/src/modules/__pycache__/motion_extractor.cpython-310.pyc b/src/modules/__pycache__/motion_extractor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a19f2900bcf69d1a9bb6c576214228c57d3242c3 GIT binary patch literal 1661 zcmah}PirJa6tC*8p6PTZe^gx8RkQ`cZOCRYc+w>Ziz~ur5P}{Wr08_jWIE~Y8mnq% zH^YdB1opIkfkjBbOTJQFJ?&QzLEo#HY_joSK^3oFRlPsI_j{$*))E5a@$~-TSB#K9 zadUHkj(iQ%3}6u?CxVKgrcICs6y>m{ITcJqb=a_+fsTnz9W|Z219Dfgiss#KiHOC@ z84)WoIEqihd_{!&Bui$$CVR?Os!ioHtBcYqeSa_8N}rCZa+Hcn%hJ{-=~zoqmDZ$o zB-65JRa=!sogT}Rv8q~Yw$hO-L^@VRrbR2j7`!kB?Fu^I+tJ+Gh<+N_>5Dgv*WD0YMS~ zhUCF>GNcIFS)7L=lx!FvTxTqgL`QVbNZxr)7|FW;NbEKiczkmOM~OjQ4K$B_S7Dkj zV6nu~IXMl^g9}2*oSuj1=py(fu(;NpX#JU&pZ^l4m@ ztm7gx8B47@VAaK(S8ZiEcj2M5E@_mIwNKl{LRu(_(pl)3k@e7Zj-M9#(C99_yrgc0 z^Rh0C;ata{`Rk{?R?7B^iB+P0V4HFOk!pq~KYZ6OC&sFVw_w~{DIYwkrt*8Gfehe( zy3;qh?1TGB4R;z(75;!UUI5N>augmh*aN>|AKFMHbNZ5q z)p_`m{LW4(bOE65k?#$%VBp9bzA+B0@jm*-K;|G`(t$@W3tSgK(m+)g)~fW%#qu^> zb@9WxXbwcN`yswYQ;f_$zny-q3c;Zx5Y$$R?T28gsLd`2i`T&LUQ>)Y>akN?Dy9qbeRE~ql7;EUijK3 zl+oFHw}V;Ug6@aLLyV#E63!R22f4@j&y%8FoOs5&5GtTQ1+)4ySlrsO=}KAG%f?$^ zw24=cD!#E%FbH5wdUUZ1E7`8p1JK5cablV0BXps==N^Q0uAaOuYXdyy^sv_!T1+u1N?`D!2&(!v$ zX*4;dQcpQ@fK$S)KZU>GE2sPg4v5#YyU8FCkLKCOuV268_sd)KdJ%%Lxp!-p0Q<{M z+QWs(LwMB>01QnK!;EBjg5lT6a84!!&t*TNzOvJJAb(vA3K zEK5+D@>o}r4|Oq|alvJ*fy1s?q&aY~lE*W?5uL%&(UTE2M2IsC9O&xS)2|u<6m*-flcdsxryP@e#B1RkRaVVyT&}stcP^1fm&jKx zk#{eV_b!q5m-0F|Kg$&j9@Osnd=-3o2P(`zw!W^_XooOVA@ys+$S#(rjkhlNdEx0* zhPIIQP6ms$pjiu#U6@}3Yn@|!O2A&DH;SG}8~P?k=9;pVs7jt_8E2_bIu;2(vw4AH z31E?>Vg|IUk~rgZ8YjAtXI;Ba%*rm8f@f5v2VAbq_w=i>tvb4HJQ<4_H%^`k6D(GZ zrz$9M6H=OrR8uN#jTrw?A@=SwL#Bzeb!ehZnx~Axm+!Xty%xXU;tyK<(-!~CxK#-j zEZ#HWvr@%*nQ=p^(u4_olBW60L|gnt#f47e%ygfo`#js)uBOwBLwUQ3XH&UmkJQcM zeY$8tWv)E8O+@*A84Hyb!gPyD(_}jq5Exh%RZdy)QqV?E8JwDa)6f?>z9?i)H_ciW zKZju29n1Z1vFT8nWU*3|zD0lkJeEbF$8n_#mfg{LIet{+(^t>FA19U4MNS1=C~7IZ z^FzAFpB55Q3nAXSKUOjs=LM@W&}kA+7xzaYH_BIX0PgkSRks1UcpdvV!Y(}PWF2P9 zXMk669}ht0l2y=%{yDhT?*5HsP7|Bc<*@0rSHL59)gJ&gK7^t?f!|Yn=m5jDQ`HSR*>yR9c8-A1xTAZO%Y@SluK7n{AN z>!rw`2P5D1qKI=&sfj3^mkIDLr7x;Do3At_B0mRx(_7SK!m|w8_u{yBtX+F% z=>hcMRrY%!>f_a$jSTuJpwJZw?0aApo8G_L1ZG2sjcFwk4ZQ??7+r9A4^*Lrb(+oZ k2LH#VlpB++?_b0${9^9Km-fF;**nW6q1(l45Vq^izjv$@oB#j- literal 0 HcmV?d00001 diff --git a/src/modules/__pycache__/stitching_retargeting_network.cpython-310.pyc b/src/modules/__pycache__/stitching_retargeting_network.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5d523f6e6f7359f608c4e18580c82290008cf0b1 GIT binary patch literal 1991 zcmah~L2Dy76qd&0v3J(l4Q&ZY3a#31OI(r|dg!5)QrN;$NJB|>FNPwNk!&lOktUKH zufqfi3GAu;0X-$h{H1}OdfLBGDDBgX?YO%%l^KnsC%yUJ_r53Vm6aGldv>-p{Jn|L zAAYjjLYRCDkNpffj{1mW9;7%6`T<7I5O47C3h__`$BoNIKjh6r)Qy-45Q-0-bje7n zA;~l!r()}{OPJ!sO*ApWF*6kI&T6*~x`c~?RDu(!h#Lu#vs{?nINl-8VP@a(2WIPp zt#&dcGv?rAu8lKHI&#F46Y%47W!FXyrD7?8Z9{OCvqX@IbR#`>B$qEmN-}9#Du*f) z%54)nPDX@T@DO_>ci%Q*D79LE5>r6%b~Hp<2*!dtzyuf==SJ8XtF(43h#r9NL}z&_ zoFL$aB&oJzBM2WGPfKF7wL6l7ccEC~YFxtRB9Sajr}1PY6k!A?vRYNa5SxL;JAf$m`D)GbR)|gG~HVg#nW6*N&3JtK=^HsawsFmy$%%dyRZ@xky>PMjFR{8!y z$p?3`-4~!^bCUQ1nYRcAZFv00>~GLH7wgafs3GMZlIgs24kciFV(tU#g;cMSL2hv)QpcJa&_21CC56h+waF@_Bng}@@xiUPG z%yt8_>bKy`ti!wRHvGz_KLnYB9MSrOrA|}tOV literal 0 HcmV?d00001 diff --git a/src/modules/__pycache__/util.cpython-310.pyc b/src/modules/__pycache__/util.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0efd2cbcfa8f5fbf6828438b4fe7da8443c8cd33 GIT binary patch literal 14833 zcmbVTd5|2}S?_Cljy<$mZKRc@n0Bm0nnYg3@{!mvj%~?OM4ZL3WRh%RCcV?WyR$o& zU-#@v!z`GrLjr|EfKZT7P`iH+P!S*;sUoBTDb7&b1qH>UprD{CCBl&^a})^`^ZR|> zJ;xqphuP}a-QRoty5IMH-|@cpZEIqpWZ-vs{mk;C4;#im@@Dww;N}q=(Wgwq2#f_I zFaxV@Hmn89_Ef%SKQPO3xcXDYRA(Z`ljepo|6Mpp2)KBcMzKlb}qdlo?PC z1yi6*rIc5JayXazIs5O_pYOA^4s-sD(iO1Pog;6c&;AXMjs;)(6OU;&Q`1M-X zPakHZeqY{FOL692amzA4AbIHtpExnLT^U+H_)^SFJ7Ja@N?)1y(ONWAvO=>zJ|W;%xWa zj~JT{+XQwm_d)Bp5jd#Jqt5B&_tfR&IXivM?%_%FTZnVpZs7I`LH>xb4fO7qc#Czy z*fBSYy&~qM*nOnr;ddC3&Gnq*c&!uVT3$- z3vy`OG!+GNzU;_(SFwkBi=0}xO!4tz99?{ALvo zH$uOu5AE*%WV6#)3{|DIRH=tPfmI1=4NN`ZrpEZLd@{+AF0hMA)LGPas}tWe)NxdQ z^_}Nbs}-N~J8>(h-xW97=PtAwOE;dra<1Bm;#Q;5Yz1MI5$}4UwjMs&QUF{nuB@Lw z7pdyG6!YhZ_qp~Be)#!}{fF_6XcmRzl&nc}(kz+NJnWKLuufWzrB0x-N=|0K30&lH zyd6j6V3vBu`|M47)7i|$#tE&s?=}c5uf1rYkkhJ~er%VcXr$6aus7_JHA8 z)GL6FR!Mub zcBK6nzOK+KZ5IQ4d#{+D$CC514XU*{-WvzvbMuN!Nf27#6{WOXQe>BQeF(xbU=rT| z(sF^EthlA!AY89iLp9DP1bcSXk9DcjjGpa;VK=OlCln`1+bUer?%Ia0mLqKj+Fa4* zhBmKj^M*FN+PZN=Ti35^tJ~FXBaBwUKme2!x3T*oiZTwF?~-|&{LR=i=8QRO&6>03 zDb6=rRdI~`D1`jI_<9tf0lQ}`Sb@1<2gMH=OQu+|;)0V}vfPIZ$5?Q|m6Fb1fFiwi zKKSaKNcGe0$Zxdk)UcvPMNM|15J~~6y9z9;q@l2Hz|tz2oBeoA7L#Am{U~Asde7Lh zwyhml=-7rXGkaE#`YP80W?pc6MkN>9+iuUoJQ|gJVBqU4=oLsqsG4Ctk!d%Ip%y}& zZ~FnIRZl_`+nuWz4tSDCIqdbC?NWfjFl>{-HO8VfzHv2T+@?R zno&$OF%J^}vO~*J&TA(M>q~m9)~qZ+&3B*y^*G+wzfe!UxT7&Y{Qm7Vl7;O zhrEidZ9BQ@Sy%A^)W#2>t$GUzZC=tP=(e>R6&2PyH*NJ0tIn{XUfBggorf6&hJ!1p zdy$h^g1RZ0j@eDPcRrX)Ti`7ep#cU&wr=Lh^`_qlD-~U;R1)Qe>+wqESy<<^CSR!p zttvP~zNp7h=!ukV)v#XASSS2@E1L)tIjP@_Lq4ovIOe1?m6xfmf;#de*Z6ypBOR6q zkq79P=L?VrC&YbT=G3QyMg@zCcRSFbzc-waprGARDjg#ir<8| zq)5%~5~LEO*8ou0U!oHR$OWmmn2-%Zif4kbi=8}B(?>Rl`@vXF5M21C5Q4@WeGA`y_JEKpFU zhlZ5Oin=`7%K*}3{*&yAvPMXc6Vjyafpmtgy-L0ypeC59YeC8J;g%QB*CZMW!37m zqigCa?#hJ(f6uT@AAgD{(GE4UdN*sR#csMo7{l^#s898N-tL38I_N3dN}_FNtnPH5 z5eI{fe7V49scm9o)r9at{9dp`=k5Uu~Hb1s(wAJEcsP< zuH{?;H;8K-s&)&ji%|J|ghN&DWpR5at5dex-gY*sXMAepns==#+ko7%$ z8=aF(pFxuy8UEHxl`XNC1Mx&Yeh0hC^chUFt4RF@aCJ|hR``}ah!2JeEy5^Nin|=l zgEZC4fkp9yRSiXk5k*2DBciAhKKJ#M&VYdt2)7Ty#Sf#aXcEPwX_+M_aRx`waZm8_ z=ul055$%|2u{daMf0C7L6hkwL{$JUQa?~He(YH{T(J@8p!83XewOCw?U4__M@x+Cx*ltG@iK=k-6Nq6U$0<%~G>Ht67K%=5b>SOI$Dwo~48V6pcyQQqSR}HhG<{IoOYOuqKamX)xgZ9H%d$rGU(Z z+cYEV+3a&Or^!m}O_wjZ%wphw4g7V^#pklV(j{m?k^1t0y_f`0_5e1e4*;w94I_YM z)Q5)v_6xW;2(W>HKl6k3adPzyfUcW_-T>T_fz`xDN$d!85o!()c21x=gAnt8p3pV{ z*XI9nN*MpRFzz z6V3eH?1qw{KEPs&#rLv!9)&J3`jd9@eWdQAsQ&`=Xa>ck=~zdt*Tc=5MTDn>u-#FJ zMLNCviv_h34-pIc5j}I;*x?ctD-TTJ5eNo^VvC}Z)Htz)ct0Yb%Mg`hUGi)m>O5P8 z7~AP$H+B(OO_s+LBYR$Wjn#;a7&jticDvP1M#>ndx`Tz!g2qgiDB?C^x$+PRlr;$% z=|aY2M99`^?IAh2jcvF({wEyKG>QVmt$@E7S*;|2q>Le%|M4L{b)dv?>i>k@71&)8 z^qEqmzO2L*`$nOA$1Wq`V*{ZUhgF94h9jDLBrlsFkvuYDRPtaZBwEvRwiv;onHauD zU~|<){3Z=)CNYO}-TP5=mSCFdL+8W9p2^U3K1h^>BIvRy*bl1v;PF54h;VBXYIzE7 zLwCB*i2YG$;fX#9r*I=n4n`QEhEJIqg(8dZNf-n&f!Ll6+v^|<;z+DJ2mI%B?&_0|UVOZo zJ<7)&9K@B(E`7S#*ut&CFL<8pMFG(Ea)^m98V6CI50yD33u-bSKRDGz8jfzHPNrfRcmQ+KRu7=o$ zXmS@*iA%b{W~D~GUBRxI3Iko_^|Ie+_}Ib1B1 z1YvVOJx4c6gt}x-A)?|~Q>JUWU!SpO&F-Y|c7U#7CX>N^bcv2Lwv#Bf9h<9gL!sI1 z9hj}akR4BkI&F!zBsvX(1)T;zY7-$R>=E#aW)05$F?e4+S9VDAoa1Dd2crlv_$V%0 zP*L0rbP$IQ6aAN0KZ@^3^jpZQZ5AJA@d*|`jsl^nKDxCVh06`U4Mz#w;)XJ~>L8M{ zu5$@r9rAR)Xu7DvadM4Iy^?t{Rf1FWKPD26WlpoJ-p03IVeu_2Wa=O1)e|fRGB$8V zegqXS;fQED3&yxPHLggj;u!f6=KRTSOeQVCn8*v|qzHLQvK{(+L`km0f-CYv-j3@j z`C>#jv0E>zx=X^*k%bxOn%=bt)F0zwAW&3{Yc@o}N|Yw{!}D04iPD7FVCPVF6ycUg zXfYq>K!pVxp%$P}Ep?yNp+8DF9TN%~cD(I%L(Ru@LGoL&k0jRxxlW!;YKuZC#$$M^ zvqjMz$2E2?<(lq_rNW_FaamHK6-h1I12ySQ+Fv zlv7uR^Qxz?C4+6P08)zG;A$rpmqxqTYO6$H;?Pi@waY+g<&Ek2?(`h#iL6`E12XH{ z6>&yeIiszdqZG{AM$lh7OK^+Qn#<}Z@Cx-Q7C*z6<(C=|{(88d@E;?LscaD+f@|ZN zN8#MKudDueq%qb)*q z1bv-^4z0kpDINQX&;+?2w0#=f2lq4n{u~FGqu4gpFXCGLJd0mKfjxkMuw&0}`EG4r zuSD7kt{*7WbL2(|50Asun}vS^fu7wn=2W)#Se*Vki2MZD3^j=XC4)@|O1JMSlw@*F zhGQ#+1k^9E_+=Est047jq<)UYKvRe{{#DTK%RbkIZ|d02OZeQqQ&z{=~=m`~9~9_K%^BVnfCF*6t7aENBn$L)hiG**&_4UotyoD^V+e z?SUGRV^ZwzVBPrBs6Cnuy8zd`aIBEW)(Yv`TB?o=k^}4y@(gZj;(^2ep1=fPSqi-} z_AF4>!@`Qj{~-f+U_cUlBge5-UBWIhGFO6H73*&mHd^bj(ujcR$!54w8LG6AYqN_x z`f#$!7;5$izn<|QL+& zZX>BB^yX75CaFASukPOnl{m*E9N`=laIoiyJ1zq~&i&@t+0MaU;}0Gp>@hT`djLJ# z%TKIt>iddu<2!WmNxC#4R%uBpMy?yGR-}Fkb?TES04tHs?4ks&%=Ul(?x!DyV=rs3 zxA5^JEC%Bq4slqQkvTGdeHUiLy?81c|1@?O;E0WJ{F$45nOzR_!N7C)EmU0LlbmDO zcC@E*iqtBOkspV~U$*^Ie1eD(6kS3OUjNpMSNzyR`b8TiE%K1i=LLT3Glc|+b)Z?;;JiC6?QcXp1vLL&*5~2!j8?#YbOXxoI2H9E6NFk-xME*xK zH|Rk+ueKEYKL9?!hqTg0-bO865vWT>nRjN9d)ks&a@Jc3Ys)L~S+oh_m9zcs%HGBb z^A<-x5b1~e^eLXc+z3NartlR>c|9Y5cU=(PlKHH^h&495ppE zoM9cNr54RJW!W3A5nW5$}JPEgBZk-$(qQpIVcc8ni*%rv+&)I%CB z6JrrxLy(_w(f+-J)X%yV(R3TE_wye3eu;)qV72uyLnAphIU zh6JE?V2lC`0wA&Hzl0|EedVp`h4Bask^}wAbXm2e$T@6%bDDQeurA1Qp&Bf0a5iW; zPg>nGfDv+W{h$&n>+_+ca=%99+JHqEfvxrG`;g)6eICk zwHn@g{j=WUGt}J8<{UHg=0-d9sPY4BU#Qxmq;^y7ulw~5Qm`8S4R4X^UQ};2o*XL#h-Vy=;W93$1!*XE-7J2p1!>aY0Zz~)kR z^wIsf7?aVrqd~#;9M`lg_vabxTICw9b z?ESc4%xcYwfXfI=UT6>KC#2?5lL5uYnGVWpNY5<%8CfuKF~ziCU{yjZ4tj9Z*`AHV z#bOIo;+L`jkmdX`z~Nm~t+s~*`jyB#b2VRk-m@k-f44yGp=A)$9>VO7VomLFE7`l^RTZ)` z(KF`JdmNc<)f%kJd+8EQa@J-Bs{5J!Vk~{1`Mk*fBX@0j%!f6j*yn^WxX>HW5 z%gkmx1|R%njgWAO|3rZxg}m%^)KTj-$R~RUPiIq>BA!%cie3!-UT(i4 z(-sb$E^>JkHrnwGT|zc~v(>5BYY@bswqC*BrI@55zqis>twmj6_B=%zK@dr+TD(!A zr>%=DrA^ZN1a?aUq~NClyGth19x?VGqZrbS7%y?DbsIt$X&0iY{CD{0^><~3Fx-|hqXjHgU#NagqHwhx-5=ZWXubAc3C9XZRiL9XX^^{cq?y9`n}@-caJu`CnxHwa+BiKEjZ{MpdS0 zzo0bkMtl&Cq`!hBU@#hvM{**!`e7$i6gDt^3r&(`-9Pi{%Pjtd#V#mLjF`dC;@i~6 z*zA8%jCcbaaDT@60fQZn#8rd`F!BJp66-*hFPcNaFJ{mkoVV1J+#%>0O<~^(UeHw| z=8muNVWc$hE`dxQyP#1m&5_q*)mFU@^~-;e(YC*c-GeJkU|g(+G=~cRozsDfFaIr} z{uZx9#2sfFe-zDe(?MoB>uDO{B!8YKTh!2{Yw&YLA91GD*I00t%db;!=d~=x$z***{>5kx zeoFQ)V}4zi2GX0<|Bv<8A zM3^v37vomtd@^$bbjp=XTH&E2lmB5bHVFtJ?;|&jx#RrtZyb@}!tC_S^pWXfQqE4l e3g9PFZ({F(j$UVHLCD2fK{P`g_{C}=d~kTb)XZ$@RmA0QYHW_KpHy9oVh zCKrbdlW*abkAW~WLJTvK;hc;JHoTQtIUxw6+#WfwVlyXmb8qCeyqo!XXVk&y0b(BW zUn1snvKPFtMqOt8h@#HXAHkDYRB15@$0^TP$fZj2Se4?=lPF|-oEDsgX%VU^539K1 z;%;~_O_OOj;RP45;xdfGgG(xEY1EkmXWG`Yna$X@Dja-Oppa;vaE-mcrCVs1NK|1j9;DwAM`5K5q3ODEl zIU%Qz%f>&kOrB3kgQ(rkz+yXd8jqnGi;tAcx(#le#%0cHtMPuvzk&wp6&6G0tNg}l z2y+|zI@)Nn57_7&J0R<*V6S^FJLedICpflYPp@%-_8EfP^mesZS2Bq+E-Y|XZ;rEg zLTkxsnIswHyDVrPS3gCTc8fZv`;}f>W``!zxF~q0HClfrYBU^XVVr4<Ud4rplJ4d4bGi(T)N3YKtUM@ayAKwiIHpi#t6Q`x_4P~c zRI!M2soj03x`}*)b%)X0JgY=F_fpAJP?s$pEKl4bW)UJdV;`S|eXhr^_ns?2Fo zGA>t(_kKuc{BbFih#?-c?V%LOFfUo1!JcJduR^d@9g0mbu?DZa4P?+AkOB6IjY<31 zZ;2iJ<9fHq2RI-{?_4f|tyRzylkL?oyovgKc;z|}h4wHYo0tkXHeJfv-GRF8$5}do zY?pusAtR#kj_yt%yHv)|GBZHu~QCW}=&4%7W1^}pPwPmcD zY_8*@n^)YuxG!0?1lWoUvQ#Ur{E~nJK0Y@QaqJ`kH@f{A-ZXhLq1 za9(2P7Q|;pNS;ZJ6OEb1Q;iR_C%B?4QCBnRL5=~U>~lCI?*Iv~OsyBNNe{1+ z4IIETAU-CRL8<%Jc8@l%q#8 z>n*=k6kFrENDQ9iOb6$@_!zWlyH;t|b|~?svDLOQ_fhM5tNudhOp|Cm(*D Bx64x256x256 + + for i in range(len(self.down_blocks)): + out = self.down_blocks[i](out) + out = self.second(out) + bs, c, h, w = out.shape # ->Bx512x64x64 + + f_s = out.view(bs, self.reshape_channel, self.reshape_depth, h, w) # ->Bx32x16x64x64 + f_s = self.resblocks_3d(f_s) # ->Bx32x16x64x64 + return f_s diff --git a/src/modules/convnextv2.py b/src/modules/convnextv2.py new file mode 100644 index 0000000..83ea126 --- /dev/null +++ b/src/modules/convnextv2.py @@ -0,0 +1,149 @@ +# coding: utf-8 + +""" +This moudle is adapted to the ConvNeXtV2 version for the extraction of implicit keypoints, poses, and expression deformation. +""" + +import torch +import torch.nn as nn +# from timm.models.layers import trunc_normal_, DropPath +from .util import LayerNorm, DropPath, trunc_normal_, GRN + +__all__ = ['convnextv2_tiny'] + + +class Block(nn.Module): + """ ConvNeXtV2 Block. + + Args: + dim (int): Number of input channels. + drop_path (float): Stochastic depth rate. Default: 0.0 + """ + + def __init__(self, dim, drop_path=0.): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv + self.norm = LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.grn = GRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + + def forward(self, x): + input = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) + + x = input + self.drop_path(x) + return x + + +class ConvNeXtV2(nn.Module): + """ ConvNeXt V2 + + Args: + in_chans (int): Number of input image channels. Default: 3 + num_classes (int): Number of classes for classification head. Default: 1000 + depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] + dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] + drop_path_rate (float): Stochastic depth rate. Default: 0. + head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. + """ + + def __init__( + self, + in_chans=3, + depths=[3, 3, 9, 3], + dims=[96, 192, 384, 768], + drop_path_rate=0., + **kwargs + ): + super().__init__() + self.depths = depths + self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers + stem = nn.Sequential( + nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4), + LayerNorm(dims[0], eps=1e-6, data_format="channels_first") + ) + self.downsample_layers.append(stem) + for i in range(3): + downsample_layer = nn.Sequential( + LayerNorm(dims[i], eps=1e-6, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2), + ) + self.downsample_layers.append(downsample_layer) + + self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks + dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] + cur = 0 + for i in range(4): + stage = nn.Sequential( + *[Block(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + + self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer + + # NOTE: the output semantic items + num_bins = kwargs.get('num_bins', 66) + num_kp = kwargs.get('num_kp', 24) # the number of implicit keypoints + self.fc_kp = nn.Linear(dims[-1], 3 * num_kp) # implicit keypoints + + # print('dims[-1]: ', dims[-1]) + self.fc_scale = nn.Linear(dims[-1], 1) # scale + self.fc_pitch = nn.Linear(dims[-1], num_bins) # pitch bins + self.fc_yaw = nn.Linear(dims[-1], num_bins) # yaw bins + self.fc_roll = nn.Linear(dims[-1], num_bins) # roll bins + self.fc_t = nn.Linear(dims[-1], 3) # translation + self.fc_exp = nn.Linear(dims[-1], 3 * num_kp) # expression / delta + + def _init_weights(self, m): + if isinstance(m, (nn.Conv2d, nn.Linear)): + trunc_normal_(m.weight, std=.02) + nn.init.constant_(m.bias, 0) + + def forward_features(self, x): + for i in range(4): + x = self.downsample_layers[i](x) + x = self.stages[i](x) + return self.norm(x.mean([-2, -1])) # global average pooling, (N, C, H, W) -> (N, C) + + def forward(self, x): + x = self.forward_features(x) + + # implicit keypoints + kp = self.fc_kp(x) + + # pose and expression deformation + pitch = self.fc_pitch(x) + yaw = self.fc_yaw(x) + roll = self.fc_roll(x) + t = self.fc_t(x) + exp = self.fc_exp(x) + scale = self.fc_scale(x) + + ret_dct = { + 'pitch': pitch, + 'yaw': yaw, + 'roll': roll, + 't': t, + 'exp': exp, + 'scale': scale, + + 'kp': kp, # canonical keypoint + } + + return ret_dct + + +def convnextv2_tiny(**kwargs): + model = ConvNeXtV2(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs) + return model diff --git a/src/modules/dense_motion.py b/src/modules/dense_motion.py new file mode 100644 index 0000000..c1a7f9a --- /dev/null +++ b/src/modules/dense_motion.py @@ -0,0 +1,104 @@ +# coding: utf-8 + +""" +The module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving +""" + +from torch import nn +import torch.nn.functional as F +import torch +from .util import Hourglass, make_coordinate_grid, kp2gaussian + + +class DenseMotionNetwork(nn.Module): + def __init__(self, block_expansion, num_blocks, max_features, num_kp, feature_channel, reshape_depth, compress, estimate_occlusion_map=True): + super(DenseMotionNetwork, self).__init__() + self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp+1)*(compress+1), max_features=max_features, num_blocks=num_blocks) # ~60+G + + self.mask = nn.Conv3d(self.hourglass.out_filters, num_kp + 1, kernel_size=7, padding=3) # 65G! NOTE: computation cost is large + self.compress = nn.Conv3d(feature_channel, compress, kernel_size=1) # 0.8G + self.norm = nn.BatchNorm3d(compress, affine=True) + self.num_kp = num_kp + self.flag_estimate_occlusion_map = estimate_occlusion_map + + if self.flag_estimate_occlusion_map: + self.occlusion = nn.Conv2d(self.hourglass.out_filters*reshape_depth, 1, kernel_size=7, padding=3) + else: + self.occlusion = None + + def create_sparse_motions(self, feature, kp_driving, kp_source): + bs, _, d, h, w = feature.shape # (bs, 4, 16, 64, 64) + identity_grid = make_coordinate_grid((d, h, w), ref=kp_source) # (16, 64, 64, 3) + identity_grid = identity_grid.view(1, 1, d, h, w, 3) # (1, 1, d=16, h=64, w=64, 3) + coordinate_grid = identity_grid - kp_driving.view(bs, self.num_kp, 1, 1, 1, 3) + + k = coordinate_grid.shape[1] + + # NOTE: there lacks an one-order flow + driving_to_source = coordinate_grid + kp_source.view(bs, self.num_kp, 1, 1, 1, 3) # (bs, num_kp, d, h, w, 3) + + # adding background feature + identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1, 1) + sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1) # (bs, 1+num_kp, d, h, w, 3) + return sparse_motions + + def create_deformed_feature(self, feature, sparse_motions): + bs, _, d, h, w = feature.shape + feature_repeat = feature.unsqueeze(1).unsqueeze(1).repeat(1, self.num_kp+1, 1, 1, 1, 1, 1) # (bs, num_kp+1, 1, c, d, h, w) + feature_repeat = feature_repeat.view(bs * (self.num_kp+1), -1, d, h, w) # (bs*(num_kp+1), c, d, h, w) + sparse_motions = sparse_motions.view((bs * (self.num_kp+1), d, h, w, -1)) # (bs*(num_kp+1), d, h, w, 3) + sparse_deformed = F.grid_sample(feature_repeat, sparse_motions, align_corners=False) + sparse_deformed = sparse_deformed.view((bs, self.num_kp+1, -1, d, h, w)) # (bs, num_kp+1, c, d, h, w) + + return sparse_deformed + + def create_heatmap_representations(self, feature, kp_driving, kp_source): + spatial_size = feature.shape[3:] # (d=16, h=64, w=64) + gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=0.01) # (bs, num_kp, d, h, w) + gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=0.01) # (bs, num_kp, d, h, w) + heatmap = gaussian_driving - gaussian_source # (bs, num_kp, d, h, w) + + # adding background feature + zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1], spatial_size[2]).type(heatmap.dtype).to(heatmap.device) + heatmap = torch.cat([zeros, heatmap], dim=1) + heatmap = heatmap.unsqueeze(2) # (bs, 1+num_kp, 1, d, h, w) + return heatmap + + def forward(self, feature, kp_driving, kp_source): + bs, _, d, h, w = feature.shape # (bs, 32, 16, 64, 64) + + feature = self.compress(feature) # (bs, 4, 16, 64, 64) + feature = self.norm(feature) # (bs, 4, 16, 64, 64) + feature = F.relu(feature) # (bs, 4, 16, 64, 64) + + out_dict = dict() + + # 1. deform 3d feature + sparse_motion = self.create_sparse_motions(feature, kp_driving, kp_source) # (bs, 1+num_kp, d, h, w, 3) + deformed_feature = self.create_deformed_feature(feature, sparse_motion) # (bs, 1+num_kp, c=4, d=16, h=64, w=64) + + # 2. (bs, 1+num_kp, d, h, w) + heatmap = self.create_heatmap_representations(deformed_feature, kp_driving, kp_source) # (bs, 1+num_kp, 1, d, h, w) + + input = torch.cat([heatmap, deformed_feature], dim=2) # (bs, 1+num_kp, c=5, d=16, h=64, w=64) + input = input.view(bs, -1, d, h, w) # (bs, (1+num_kp)*c=105, d=16, h=64, w=64) + + prediction = self.hourglass(input) + + mask = self.mask(prediction) + mask = F.softmax(mask, dim=1) # (bs, 1+num_kp, d=16, h=64, w=64) + out_dict['mask'] = mask + mask = mask.unsqueeze(2) # (bs, num_kp+1, 1, d, h, w) + sparse_motion = sparse_motion.permute(0, 1, 5, 2, 3, 4) # (bs, num_kp+1, 3, d, h, w) + deformation = (sparse_motion * mask).sum(dim=1) # (bs, 3, d, h, w) mask take effect in this place + deformation = deformation.permute(0, 2, 3, 4, 1) # (bs, d, h, w, 3) + + out_dict['deformation'] = deformation + + if self.flag_estimate_occlusion_map: + bs, _, d, h, w = prediction.shape + prediction_reshape = prediction.view(bs, -1, h, w) + occlusion_map = torch.sigmoid(self.occlusion(prediction_reshape)) # Bx1x64x64 + out_dict['occlusion_map'] = occlusion_map + + return out_dict diff --git a/src/modules/motion_extractor.py b/src/modules/motion_extractor.py new file mode 100644 index 0000000..b2982e5 --- /dev/null +++ b/src/modules/motion_extractor.py @@ -0,0 +1,35 @@ +# coding: utf-8 + +""" +Motion extractor(M), which directly predicts the canonical keypoints, head pose and expression deformation of the input image +""" + +from torch import nn +import torch + +from .convnextv2 import convnextv2_tiny +from .util import filter_state_dict + +model_dict = { + 'convnextv2_tiny': convnextv2_tiny, +} + + +class MotionExtractor(nn.Module): + def __init__(self, **kwargs): + super(MotionExtractor, self).__init__() + + # default is convnextv2_base + backbone = kwargs.get('backbone', 'convnextv2_tiny') + self.detector = model_dict.get(backbone)(**kwargs) + + def load_pretrained(self, init_path: str): + if init_path not in (None, ''): + state_dict = torch.load(init_path, map_location=lambda storage, loc: storage)['model'] + state_dict = filter_state_dict(state_dict, remove_name='head') + ret = self.detector.load_state_dict(state_dict, strict=False) + print(f'Load pretrained model from {init_path}, ret: {ret}') + + def forward(self, x): + out = self.detector(x) + return out diff --git a/src/modules/spade_generator.py b/src/modules/spade_generator.py new file mode 100644 index 0000000..147a9ae --- /dev/null +++ b/src/modules/spade_generator.py @@ -0,0 +1,59 @@ +# coding: utf-8 + +""" +Spade decoder(G) defined in the paper, which input the warped feature to generate the animated image. +""" + +import torch +from torch import nn +import torch.nn.functional as F +from .util import SPADEResnetBlock + + +class SPADEDecoder(nn.Module): + def __init__(self, upscale=1, max_features=256, block_expansion=64, out_channels=64, num_down_blocks=2): + for i in range(num_down_blocks): + input_channels = min(max_features, block_expansion * (2 ** (i + 1))) + self.upscale = upscale + super().__init__() + norm_G = 'spadespectralinstance' + label_num_channels = input_channels # 256 + + self.fc = nn.Conv2d(input_channels, 2 * input_channels, 3, padding=1) + self.G_middle_0 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_1 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_2 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_3 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_4 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_5 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.up_0 = SPADEResnetBlock(2 * input_channels, input_channels, norm_G, label_num_channels) + self.up_1 = SPADEResnetBlock(input_channels, out_channels, norm_G, label_num_channels) + self.up = nn.Upsample(scale_factor=2) + + if self.upscale is None or self.upscale <= 1: + self.conv_img = nn.Conv2d(out_channels, 3, 3, padding=1) + else: + self.conv_img = nn.Sequential( + nn.Conv2d(out_channels, 3 * (2 * 2), kernel_size=3, padding=1), + nn.PixelShuffle(upscale_factor=2) + ) + + def forward(self, feature): + seg = feature # Bx256x64x64 + x = self.fc(feature) # Bx512x64x64 + x = self.G_middle_0(x, seg) + x = self.G_middle_1(x, seg) + x = self.G_middle_2(x, seg) + x = self.G_middle_3(x, seg) + x = self.G_middle_4(x, seg) + x = self.G_middle_5(x, seg) + + x = self.up(x) # Bx512x64x64 -> Bx512x128x128 + x = self.up_0(x, seg) # Bx512x128x128 -> Bx256x128x128 + x = self.up(x) # Bx256x128x128 -> Bx256x256x256 + x = self.up_1(x, seg) # Bx256x256x256 -> Bx64x256x256 + + x = self.conv_img(F.leaky_relu(x, 2e-1)) # Bx64x256x256 -> Bx3xHxW + x = torch.sigmoid(x) # Bx3xHxW + + return x \ No newline at end of file diff --git a/src/modules/stitching_retargeting_network.py b/src/modules/stitching_retargeting_network.py new file mode 100644 index 0000000..5f50b7c --- /dev/null +++ b/src/modules/stitching_retargeting_network.py @@ -0,0 +1,38 @@ +# coding: utf-8 + +""" +Stitching module(S) and two retargeting modules(R) defined in the paper. + +- The stitching module pastes the animated portrait back into the original image space without pixel misalignment, such as in +the stitching region. + +- The eyes retargeting module is designed to address the issue of incomplete eye closure during cross-id reenactment, especially +when a person with small eyes drives a person with larger eyes. + +- The lip retargeting module is designed similarly to the eye retargeting module, and can also normalize the input by ensuring that +the lips are in a closed state, which facilitates better animation driving. +""" +from torch import nn + + +class StitchingRetargetingNetwork(nn.Module): + def __init__(self, input_size, hidden_sizes, output_size): + super(StitchingRetargetingNetwork, self).__init__() + layers = [] + for i in range(len(hidden_sizes)): + if i == 0: + layers.append(nn.Linear(input_size, hidden_sizes[i])) + else: + layers.append(nn.Linear(hidden_sizes[i - 1], hidden_sizes[i])) + layers.append(nn.ReLU(inplace=True)) + layers.append(nn.Linear(hidden_sizes[-1], output_size)) + self.mlp = nn.Sequential(*layers) + + def initialize_weights_to_zero(self): + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.zeros_(m.weight) + nn.init.zeros_(m.bias) + + def forward(self, x): + return self.mlp(x) diff --git a/src/modules/util.py b/src/modules/util.py new file mode 100644 index 0000000..dc6b925 --- /dev/null +++ b/src/modules/util.py @@ -0,0 +1,452 @@ +# coding: utf-8 + +""" +This file defines various neural network modules and utility functions, including convolutional and residual blocks, +normalizations, and functions for spatial transformation and tensor manipulation. +""" + +from torch import nn +import torch.nn.functional as F +import torch +import torch.nn.utils.spectral_norm as spectral_norm +import math +import warnings +import collections.abc +from itertools import repeat + +def kp2gaussian(kp, spatial_size, kp_variance): + """ + Transform a keypoint into gaussian like representation + """ + mean = kp + + coordinate_grid = make_coordinate_grid(spatial_size, mean) + number_of_leading_dimensions = len(mean.shape) - 1 + shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape + coordinate_grid = coordinate_grid.view(*shape) + repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 1) + coordinate_grid = coordinate_grid.repeat(*repeats) + + # Preprocess kp shape + shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 3) + mean = mean.view(*shape) + + mean_sub = (coordinate_grid - mean) + + out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance) + + return out + + +def make_coordinate_grid(spatial_size, ref, **kwargs): + d, h, w = spatial_size + x = torch.arange(w).type(ref.dtype).to(ref.device) + y = torch.arange(h).type(ref.dtype).to(ref.device) + z = torch.arange(d).type(ref.dtype).to(ref.device) + + # NOTE: must be right-down-in + x = (2 * (x / (w - 1)) - 1) # the x axis faces to the right + y = (2 * (y / (h - 1)) - 1) # the y axis faces to the bottom + z = (2 * (z / (d - 1)) - 1) # the z axis faces to the inner + + yy = y.view(1, -1, 1).repeat(d, 1, w) + xx = x.view(1, 1, -1).repeat(d, h, 1) + zz = z.view(-1, 1, 1).repeat(1, h, w) + + meshed = torch.cat([xx.unsqueeze_(3), yy.unsqueeze_(3), zz.unsqueeze_(3)], 3) + + return meshed + + +class ConvT2d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, stride=2, padding=1, output_padding=1): + super(ConvT2d, self).__init__() + + self.convT = nn.ConvTranspose2d(in_features, out_features, kernel_size=kernel_size, stride=stride, + padding=padding, output_padding=output_padding) + self.norm = nn.InstanceNorm2d(out_features) + + def forward(self, x): + out = self.convT(x) + out = self.norm(out) + out = F.leaky_relu(out) + return out + + +class ResBlock3d(nn.Module): + """ + Res block, preserve spatial resolution. + """ + + def __init__(self, in_features, kernel_size, padding): + super(ResBlock3d, self).__init__() + self.conv1 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding) + self.conv2 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding) + self.norm1 = nn.BatchNorm3d(in_features, affine=True) + self.norm2 = nn.BatchNorm3d(in_features, affine=True) + + def forward(self, x): + out = self.norm1(x) + out = F.relu(out) + out = self.conv1(out) + out = self.norm2(out) + out = F.relu(out) + out = self.conv2(out) + out += x + return out + + +class UpBlock3d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(UpBlock3d, self).__init__() + + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = nn.BatchNorm3d(out_features, affine=True) + + def forward(self, x): + out = F.interpolate(x, scale_factor=(1, 2, 2)) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class DownBlock2d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups) + self.norm = nn.BatchNorm2d(out_features, affine=True) + self.pool = nn.AvgPool2d(kernel_size=(2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class DownBlock3d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock3d, self).__init__() + ''' + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups, stride=(1, 2, 2)) + ''' + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = nn.BatchNorm3d(out_features, affine=True) + self.pool = nn.AvgPool3d(kernel_size=(1, 2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class SameBlock2d(nn.Module): + """ + Simple block, preserve spatial resolution. + """ + + def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1, lrelu=False): + super(SameBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups) + self.norm = nn.BatchNorm2d(out_features, affine=True) + if lrelu: + self.ac = nn.LeakyReLU() + else: + self.ac = nn.ReLU() + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = self.ac(out) + return out + + +class Encoder(nn.Module): + """ + Hourglass Encoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Encoder, self).__init__() + + down_blocks = [] + for i in range(num_blocks): + down_blocks.append(DownBlock3d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), min(max_features, block_expansion * (2 ** (i + 1))), kernel_size=3, padding=1)) + self.down_blocks = nn.ModuleList(down_blocks) + + def forward(self, x): + outs = [x] + for down_block in self.down_blocks: + outs.append(down_block(outs[-1])) + return outs + + +class Decoder(nn.Module): + """ + Hourglass Decoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Decoder, self).__init__() + + up_blocks = [] + + for i in range(num_blocks)[::-1]: + in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1))) + out_filters = min(max_features, block_expansion * (2 ** i)) + up_blocks.append(UpBlock3d(in_filters, out_filters, kernel_size=3, padding=1)) + + self.up_blocks = nn.ModuleList(up_blocks) + self.out_filters = block_expansion + in_features + + self.conv = nn.Conv3d(in_channels=self.out_filters, out_channels=self.out_filters, kernel_size=3, padding=1) + self.norm = nn.BatchNorm3d(self.out_filters, affine=True) + + def forward(self, x): + out = x.pop() + for up_block in self.up_blocks: + out = up_block(out) + skip = x.pop() + out = torch.cat([out, skip], dim=1) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class Hourglass(nn.Module): + """ + Hourglass architecture. + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Hourglass, self).__init__() + self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features) + self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features) + self.out_filters = self.decoder.out_filters + + def forward(self, x): + return self.decoder(self.encoder(x)) + + +class SPADE(nn.Module): + def __init__(self, norm_nc, label_nc): + super().__init__() + + self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False) + nhidden = 128 + + self.mlp_shared = nn.Sequential( + nn.Conv2d(label_nc, nhidden, kernel_size=3, padding=1), + nn.ReLU()) + self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1) + self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1) + + def forward(self, x, segmap): + normalized = self.param_free_norm(x) + segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest') + actv = self.mlp_shared(segmap) + gamma = self.mlp_gamma(actv) + beta = self.mlp_beta(actv) + out = normalized * (1 + gamma) + beta + return out + + +class SPADEResnetBlock(nn.Module): + def __init__(self, fin, fout, norm_G, label_nc, use_se=False, dilation=1): + super().__init__() + # Attributes + self.learned_shortcut = (fin != fout) + fmiddle = min(fin, fout) + self.use_se = use_se + # create conv layers + self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=dilation, dilation=dilation) + self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=dilation, dilation=dilation) + if self.learned_shortcut: + self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False) + # apply spectral norm if specified + if 'spectral' in norm_G: + self.conv_0 = spectral_norm(self.conv_0) + self.conv_1 = spectral_norm(self.conv_1) + if self.learned_shortcut: + self.conv_s = spectral_norm(self.conv_s) + # define normalization layers + self.norm_0 = SPADE(fin, label_nc) + self.norm_1 = SPADE(fmiddle, label_nc) + if self.learned_shortcut: + self.norm_s = SPADE(fin, label_nc) + + def forward(self, x, seg1): + x_s = self.shortcut(x, seg1) + dx = self.conv_0(self.actvn(self.norm_0(x, seg1))) + dx = self.conv_1(self.actvn(self.norm_1(dx, seg1))) + out = x_s + dx + return out + + def shortcut(self, x, seg1): + if self.learned_shortcut: + x_s = self.conv_s(self.norm_s(x, seg1)) + else: + x_s = x + return x_s + + def actvn(self, x): + return F.leaky_relu(x, 2e-1) + + +def filter_state_dict(state_dict, remove_name='fc'): + new_state_dict = {} + for key in state_dict: + if remove_name in key: + continue + new_state_dict[key] = state_dict[key] + return new_state_dict + + +class GRN(nn.Module): + """ GRN (Global Response Normalization) layer + """ + + def __init__(self, dim): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x): + Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * Nx) + self.beta + x + + +class LayerNorm(nn.Module): + r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. + The ordering of the dimensions in the inputs. channels_last corresponds to inputs with + shape (batch_size, height, width, channels) while channels_first corresponds to inputs + with shape (batch_size, channels, height, width). + """ + + def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + if self.data_format not in ["channels_last", "channels_first"]: + raise NotImplementedError + self.normalized_shape = (normalized_shape, ) + + def forward(self, x): + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + elif self.data_format == "channels_first": + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x + + +def _no_grad_trunc_normal_(tensor, mean, std, a, b): + # Cut & paste from PyTorch official master until it's in a few official releases - RW + # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + def norm_cdf(x): + # Computes standard normal cumulative distribution function + return (1. + math.erf(x / math.sqrt(2.))) / 2. + + if (mean < a - 2 * std) or (mean > b + 2 * std): + warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " + "The distribution of values may be incorrect.", + stacklevel=2) + + with torch.no_grad(): + # Values are generated by using a truncated uniform distribution and + # then using the inverse CDF for the normal distribution. + # Get upper and lower cdf values + l = norm_cdf((a - mean) / std) + u = norm_cdf((b - mean) / std) + + # Uniformly fill tensor with values from [l, u], then translate to + # [2l-1, 2u-1]. + tensor.uniform_(2 * l - 1, 2 * u - 1) + + # Use inverse cdf transform for normal distribution to get truncated + # standard normal + tensor.erfinv_() + + # Transform to proper mean, std + tensor.mul_(std * math.sqrt(2.)) + tensor.add_(mean) + + # Clamp to ensure it's in the proper range + tensor.clamp_(min=a, max=b) + return tensor + + +def drop_path(x, drop_prob=0., training=False, scale_by_keep=True): + """ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, + the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for + changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use + 'survival rate' as the argument. + + """ + if drop_prob == 0. or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0 and scale_by_keep: + random_tensor.div_(keep_prob) + return x * random_tensor + + +class DropPath(nn.Module): + """ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + """ + + def __init__(self, drop_prob=None, scale_by_keep=True): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) + + +def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.): + return _no_grad_trunc_normal_(tensor, mean, std, a, b) + +# From PyTorch internals +def _ntuple(n): + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return tuple(x) + return tuple(repeat(x, n)) + return parse + +to_2tuple = _ntuple(2) diff --git a/src/modules/warping_network.py b/src/modules/warping_network.py new file mode 100644 index 0000000..9191a19 --- /dev/null +++ b/src/modules/warping_network.py @@ -0,0 +1,77 @@ +# coding: utf-8 + +""" +Warping field estimator(W) defined in the paper, which generates a warping field using the implicit +keypoint representations x_s and x_d, and employs this flow field to warp the source feature volume f_s. +""" + +from torch import nn +import torch.nn.functional as F +from .util import SameBlock2d +from .dense_motion import DenseMotionNetwork + + +class WarpingNetwork(nn.Module): + def __init__( + self, + num_kp, + block_expansion, + max_features, + num_down_blocks, + reshape_channel, + estimate_occlusion_map=False, + dense_motion_params=None, + **kwargs + ): + super(WarpingNetwork, self).__init__() + + self.upscale = kwargs.get('upscale', 1) + self.flag_use_occlusion_map = kwargs.get('flag_use_occlusion_map', True) + + if dense_motion_params is not None: + self.dense_motion_network = DenseMotionNetwork( + num_kp=num_kp, + feature_channel=reshape_channel, + estimate_occlusion_map=estimate_occlusion_map, + **dense_motion_params + ) + else: + self.dense_motion_network = None + + self.third = SameBlock2d(max_features, block_expansion * (2 ** num_down_blocks), kernel_size=(3, 3), padding=(1, 1), lrelu=True) + self.fourth = nn.Conv2d(in_channels=block_expansion * (2 ** num_down_blocks), out_channels=block_expansion * (2 ** num_down_blocks), kernel_size=1, stride=1) + + self.estimate_occlusion_map = estimate_occlusion_map + + def deform_input(self, inp, deformation): + return F.grid_sample(inp, deformation, align_corners=False) + + def forward(self, feature_3d, kp_driving, kp_source): + if self.dense_motion_network is not None: + # Feature warper, Transforming feature representation according to deformation and occlusion + dense_motion = self.dense_motion_network( + feature=feature_3d, kp_driving=kp_driving, kp_source=kp_source + ) + if 'occlusion_map' in dense_motion: + occlusion_map = dense_motion['occlusion_map'] # Bx1x64x64 + else: + occlusion_map = None + + deformation = dense_motion['deformation'] # Bx16x64x64x3 + out = self.deform_input(feature_3d, deformation) # Bx32x16x64x64 + + bs, c, d, h, w = out.shape # Bx32x16x64x64 + out = out.view(bs, c * d, h, w) # -> Bx512x64x64 + out = self.third(out) # -> Bx256x64x64 + out = self.fourth(out) # -> Bx256x64x64 + + if self.flag_use_occlusion_map and (occlusion_map is not None): + out = out * occlusion_map + + ret_dct = { + 'occlusion_map': occlusion_map, + 'deformation': deformation, + 'out': out, + } + + return ret_dct