wijddlZddlmZmZmZddlZddlZddlZddl m Z m Z m Z Gdde Z dS)N)ListOptionalUnion)VaeImageProcessoris_valid_imageis_valid_image_imagelistc eZdZdZd deedeedejfdZ d dejd e de e j eje ejjffd ZdS) VideoProcessorzSimple video processor.Nheightwidthreturnct|trat|dtjrA|djdkr0t jdttj|d}t|trat|dtj rA|djdkr0t jdttj |d}t|tjtj fr|jdkrt|}nyt|trt|dst|r|g}nz3VideoProcessor.preprocess_video..Rs,___RUT__Su_MM___)dimr) isinstancelistnpndarrayndimwarningswarn FutureWarning concatenatetorchTensorcatrr ValueErrorstackpermute)rvideor r s` ``rpreprocess_videozVideoProcessor.preprocess_videos2 eT " " 2z%(BJ'G'G 2ERSHM]^L^L^ Mn    N5q111E eT " " -z%(EL'I'I -eTUhm_`N`N` Mp    Ie!,,,E ebj%,7 8 8 UZ1__KKEE t $ $ a)A)A E]^cEdEd GEE t $ $ )A%()K)K EEw  ______Y^___efggg aAq!,, rr r- output_typecv|jd}g}t|D]K}||dddd}|||}||L|dkrt j|}n3|dkrtj|}n|dkst|d|S) z Converts a video tensor to a list of frames for export. Args: video (`torch.Tensor`): The video as a tensor. output_type (`str`, defaults to `"np"`): Output type of the postprocessed `video` tensor. rrrrr ptpilz9 does not exist. Please choose one of ['np', 'pt', 'pil']) shaperanger, postprocessappendr r+r'r*)rr-r/ batch_sizeoutputs batch_idx batch_vid batch_outputs rpostprocess_videoz VideoProcessor.postprocess_videoYs[^ z** ) )Ii(00Aq!<> ?rr )r#typingrrrnumpyr rCr'image_processorrrr r rErrrIs(((((((((( XXXXXXXXXXXXXXX&XXXXXr