1+ using Microsoft . Extensions . Logging ;
2+ using OnnxStack . Core . Config ;
3+ using OnnxStack . Core . Image ;
4+ using OnnxStack . Core . Video ;
5+ using OnnxStack . FeatureExtractor . Common ;
6+ using OnnxStack . FeatureExtractor . Pipelines ;
7+ using System ;
8+ using System . Collections . Generic ;
9+ using System . Threading ;
10+ using System . Threading . Tasks ;
11+
12+ namespace OnnxStack . UI . Services
13+ {
14+ public class FeatureExtractorService : IFeatureExtractorService
15+ {
16+ private readonly ILogger < FeatureExtractorService > _logger ;
17+ private readonly Dictionary < IOnnxModel , FeatureExtractorPipeline > _pipelines ;
18+
19+ /// <summary>
20+ /// Initializes a new instance of the <see cref="FeatureExtractorService"/> class.
21+ /// </summary>
22+ /// <param name="configuration">The configuration.</param>
23+ /// <param name="modelService">The model service.</param>
24+ /// <param name="imageService">The image service.</param>
25+ public FeatureExtractorService ( ILogger < FeatureExtractorService > logger )
26+ {
27+ _logger = logger ;
28+ _pipelines = new Dictionary < IOnnxModel , FeatureExtractorPipeline > ( ) ;
29+ }
30+
31+
32+ /// <summary>
33+ /// Loads the model.
34+ /// </summary>
35+ /// <param name="model">The model.</param>
36+ /// <returns></returns>
37+ public async Task < bool > LoadModelAsync ( FeatureExtractorModelSet model )
38+ {
39+ if ( _pipelines . ContainsKey ( model ) )
40+ return true ;
41+
42+ var pipeline = CreatePipeline ( model ) ;
43+ await pipeline . LoadAsync ( ) ;
44+ return _pipelines . TryAdd ( model , pipeline ) ;
45+ }
46+
47+
48+ /// <summary>
49+ /// Unloads the model.
50+ /// </summary>
51+ /// <param name="model">The model.</param>
52+ /// <returns></returns>
53+ public async Task < bool > UnloadModelAsync ( FeatureExtractorModelSet model )
54+ {
55+ if ( _pipelines . Remove ( model , out var pipeline ) )
56+ {
57+ await pipeline ? . UnloadAsync ( ) ;
58+ }
59+ return true ;
60+ }
61+
62+
63+ /// <summary>
64+ /// Determines whether [is model loaded] [the specified model options].
65+ /// </summary>
66+ /// <param name="model">The model.</param>
67+ /// <returns>
68+ /// <c>true</c> if [is model loaded] [the specified model options]; otherwise, <c>false</c>.
69+ /// </returns>
70+ /// <exception cref="System.NotImplementedException"></exception>
71+ public bool IsModelLoaded ( FeatureExtractorModelSet model )
72+ {
73+ return _pipelines . ContainsKey ( model ) ;
74+ }
75+
76+
77+ /// <summary>
78+ /// Generates the feature image.
79+ /// </summary>
80+ /// <param name="model">The model.</param>
81+ /// <param name="inputImage">The input image.</param>
82+ /// <returns></returns>
83+ public async Task < OnnxImage > GenerateAsync ( FeatureExtractorModelSet model , OnnxImage inputImage , CancellationToken cancellationToken = default )
84+ {
85+ if ( ! _pipelines . TryGetValue ( model , out var pipeline ) )
86+ throw new Exception ( "Pipeline not found or is unsupported" ) ;
87+
88+ return await pipeline . RunAsync ( inputImage , cancellationToken ) ;
89+ }
90+
91+
92+ /// <summary>
93+ /// Generates the feature video.
94+ /// </summary>
95+ /// <param name="model">The model.</param>
96+ /// <param name="inputVideo">The input video.</param>
97+ /// <param name="cancellationToken">The cancellation token.</param>
98+ /// <returns></returns>
99+ /// <exception cref="System.Exception">Pipeline not found or is unsupported</exception>
100+ public async Task < OnnxVideo > GenerateAsync ( FeatureExtractorModelSet model , OnnxVideo inputVideo , CancellationToken cancellationToken = default )
101+ {
102+ if ( ! _pipelines . TryGetValue ( model , out var pipeline ) )
103+ throw new Exception ( "Pipeline not found or is unsupported" ) ;
104+
105+ return await pipeline . RunAsync ( inputVideo , cancellationToken ) ;
106+ }
107+
108+
109+ /// <summary>
110+ /// Creates the pipeline.
111+ /// </summary>
112+ /// <param name="modelSet">The model.</param>
113+ /// <returns></returns>
114+ private FeatureExtractorPipeline CreatePipeline ( FeatureExtractorModelSet model )
115+ {
116+ return FeatureExtractorPipeline . CreatePipeline ( model , _logger ) ;
117+ }
118+ }
119+ }
0 commit comments