i!ddlmZddlmZmZddlmZmZddlZddl m Z ddl m cm Z Gdde jZy))cached_property) combinations permutations)DictTupleNceZdZdZdedeffd ZedeeefdZ edefdZ de jfdZ de jfd Zdd e jd ede jfd Zdd e jd ede jfd Zde jde jfdZdeedfdeedffdZedeeedfeedfffdZxZS)PowersetzPowerset to multilabel conversion, and back. Parameters ---------- num_classes : int Number of regular classes. max_set_size : int Maximum number of classes in each set. num_classes max_set_sizect|||_||_|j d|j d|j d|j dy)NmappingF) persistent cardinality)super__init__r r register_buffer build_mappingbuild_cardinality)selfr r __class__s r/Volumes/fast/ai/experiments/voice-extract-mac/.venv/lib/python3.12/site-packages/pyannote/audio/utils/powerset.pyrzPowerset.__init__0s[ &( Y(:(:(<O ]D,B,B,DQVWreturncg}td|jdzD]@}tt|j|D]}|j t |B|S)a3List of powerset classes e.g. with num_classes = 3 and max_set_size = 2: {}, {0}, {1}, {2}, {0, 1}, {0, 2}, {1, 2} Returns ------- powerset_classes : list of set[int] List of powerset classes, each represented as a set of regular class indices. r)ranger rr appendset)rpowerset_classesset_size current_sets rrzPowerset.powerset_classes9s^a!2!2Q!67H+E$2B2B,CXN  ''K(89 O8 rc,t|jS)zNumber of powerset classes)lenrrs rnum_powerset_classeszPowerset.num_powerset_classesKs4(())rctj|j|j}d}t d|j dzD]2}t t |j|D]}d|||f<|dz }4|S)atCompute powerset to regular mapping Returns ------- mapping : (num_powerset_classes, num_classes) torch.Tensor mapping[i, j] == 1 if jth regular class is a member of ith powerset class mapping[i, j] == 0 otherwise Example ------- With num_classes == 3 and max_set_size == 2, returns [0, 0, 0] # none [1, 0, 0] # class #1 [0, 1, 0] # class #2 [0, 0, 1] # class #3 [1, 1, 0] # classes #1 and #2 [1, 0, 1] # classes #1 and #3 [0, 1, 1] # classes #2 and #3 rr)torchzerosr%r rr r)rr powerset_kr r!s rrzPowerset.build_mappingPs},++d779I9IJ a!2!2Q!67H+E$2B2B,CXN 34 K/0a  O8 rcDtj|jdS)z#Compute size of each powerset classrdim)r'sumr r$s rrzPowerset.build_cardinalityosyy1--rpowersetsoftc |rtj|}nWtjjj tj |d|j j}tj||jS)a/Convert predictions from powerset to multi-label Parameter --------- powerset : (batch_size, num_frames, num_powerset_classes) torch.Tensor Soft predictions in "powerset" space. soft : bool, optional Return soft multi-label predictions. Defaults to False (i.e. hard predictions) Assumes that `powerset` are "log probabilities". Returns ------- multi_label : (batch_size, num_frames, num_classes) torch.Tensor Predictions in "multi-label" space. r+) r'expnn functionalone_hotargmaxr%floatmatmulr )rr.r/powerset_probss r to_multilabelzPowerset.to_multilabelssi" "YYx0N"XX0088 X2.))eg  ||NDLL99rc(|j||S)zAlias for `to_multilabel`)r/)r:)rr.r/s rforwardzPowerset.forwards!!(!66r multilabelctjtjtj||j j d|jS)aConvert (hard) predictions from multi-label to powerset Parameter --------- multi_label : (batch_size, num_frames, num_classes) torch.Tensor Prediction in "multi-label" space. Returns ------- powerset : (batch_size, num_frames, num_powerset_classes) torch.Tensor Hard, one-hot prediction in "powerset" space. Note ---- This method will not complain if `multilabel` is provided a soft predictions (e.g. the output of a sigmoid-ed classifier). However, in that particular case, the resulting powerset output will most likely not make much sense. r1r+)r )Fr5r'r6r8r Tr%)rr=s r to_powersetzPowerset.to_powersets?&yy LLj$,,..Ar J11  rmultilabel_permutation.c|jdd|f}tj|j|jjtj }d|zj |jdf}tj|j|zd}tj||zd}|d|dddfk(j jd}t|jS)aHelper function for `permutation_mapping` property Takes a (num_classes,)-shaped permutation in multilabel space and returns the corresponding (num_powerset_classes,)-shaped permutation in powerset space. This does not cache anything and only works on one single permutation at a time. Parameters ---------- multilabel_permutation : tuple of int Permutation in multilabel space. Returns ------- powerset_permutation : tuple of int Permutation in powerset space. Example ------- >>> powerset = Powerset(3, 2) >>> powerset._permutation_powerset((1, 0, 2)) # (0, 2, 1, 3, 4, 6, 5) N)devicedtyperr1r+r) r r'aranger rDinttiler%r-r6tupletolist)rrBpermutated_mappingrG powers_of_twobeforeafterpowerset_permutations r_permutation_powersetzPowerset._permutation_powersets6,0<<;Q8Q+R   T\\%8%8  F(($*C*CQ)GH 4<<-7R@ ,}<"E!'t ag >CCELLQRLS)00233rci}tt|j|jD]}|j||t |<!|S)aMapping between multilabel and powerset permutations Example ------- With num_classes == 3 and max_set_size == 2, returns { (0, 1, 2): (0, 1, 2, 3, 4, 5, 6), (0, 2, 1): (0, 1, 3, 2, 5, 4, 6), (1, 0, 2): (0, 2, 1, 3, 4, 6, 5), (1, 2, 0): (0, 2, 3, 1, 6, 4, 5), (2, 0, 1): (0, 3, 1, 2, 5, 6, 4), (2, 1, 0): (0, 3, 2, 1, 6, 5, 4) } )rrr rQrJ)rpermutation_mappingrBs rrSzPowerset.permutation_mappings\"!&2 $"" #T%5%5' " **+AB ,- ' #"r)F)__name__ __module__ __qualname____doc__rHrrlistrrr%r'Tensorrrboolr:r<rArrQrrS __classcell__)rs@rr r %s6XCXsX $s3x.  "*c**u||>.5<<.:ell:$:5<<:67 7D7U\\7 ell u|| 0+4&+CHo+4 sCx+4Z#T%S/5c?*J%K##rr ) functoolsr itertoolsrrtypingrrr'torch.nnr3torch.nn.functionalr4r?Moduler rrrcs.8&0 L#ryyL#r