RЦi SrSSKJrJrJrJrJrJrJr SSK r SSK J r SSK J r J r JrJrJrJrJrJrJr SSKJrJr /SQr"S S \ R25r"S S \ R25rg) zlPyTorch FX Based Feature Extraction Helpers Using https://pytorch.org/vision/stable/feature_extraction.html )CallableDictListOptionalUnionTupleTypeN)nn) create_feature_extractorget_graph_node_namesregister_notrace_moduleregister_notrace_functionis_notrace_moduleis_notrace_functionget_notrace_functionsget_notrace_modulesFormat)_get_feature_info_get_return_layers) r rrrrrr r FeatureGraphNetGraphExtractNetc ^\rSrSr%Sr\R R\\ S'S S\ RS\ \ S4S\\S\S\4 U4S jjjrS rS rU=r$) r&zWA FX Graph based feature extractor that works with the model feature_info metadata return_dictmodel out_indices.out_map output_fmtc>[TU]5 [X5UlUb[ U5[ U5:Xde[ U5Ul[URU5n[X5Ul XPl gN) super__init__r feature_infolenrrrr graph_moduler)selfrrrrr return_nodes __class__s W/var/www/html/ai-image-ml/venv/lib/python3.13/site-packages/timm/models/_features_fx.pyr#FeatureGraphNet.__init__+sk -eA  w<3{#33 33 ,)$*;*;WE 4UI&c|URU5nUR(aU$[UR55$r!)r&rlistvaluesr'xouts r*forwardFeatureGraphNet.forward<s1"   JCJJL!!r,)r$r&rr)NNCHWF)__name__ __module__ __qualname____firstlineno____doc__torchjitFinalbool__annotations__r Modulerintrrstrr#r3__static_attributes__ __classcell__r)s@r*rr&sz&& '+$ % '99'sCx'd^ '  '  ''"""r,rc ^\rSrSr%Sr\R R\\ S'S S\ RS\ \ \\4\\4S\S\4U4SjjjrS\ \\R"\R"44S jrS rU=r$) rCaA standalone feature extraction wrapper that maps dict -> list or single tensor NOTE: * one can use feature_extractor directly if dictionary output is desired * unlike FeatureGraphNet, this is intended to be used standalone and not with model feature_info metadata for builtin feature extraction mode * create_feature_extractor can be used directly if dictionary output is acceptable Args: model: model to extract features from return_nodes: node names to return features from (dict or list) squeeze_out: if only one output, and output in list format, flatten to single tensor return_dict: return as dictionary from extractor with node names as keys, ignores squeeze_out arg rrr( squeeze_outcZ>[TU]5 X0l[X5UlX@lgr!)r"r#rHr r&r)r'rr(rHrr)s r*r#GraphExtractNet.__init__Ss) &4UI&r,returncURU5nUR(aU$[UR55nUR(a[ U5S:XaUS$U$)Nrr)r&rr.r/rHr%r0s r*r3GraphExtractNet.forward_sQ"   J3::< ))c#h!ms1vDDr,)r&rrH)TF)r6r7r8r9r:r;r<r=r>r?r r@rrrBrr#Tensorr3rCrDrEs@r*rrCs && !% % '99 ' S#XS 9: ' '  ' 'EE$u||"4ell"BCEEr,r)r:typingrrrrrrr r;r timm.layersr r r rrrrrr _featuresrr__all__r@rrr,r*rTsYFEE    = "bii":!Ebii!Er,