|
1 | 1 | using Microsoft.ML.OnnxRuntime; |
2 | 2 | using OnnxStack.Core.Config; |
| 3 | +using OnnxStack.Core.Image; |
3 | 4 | using OnnxStack.Core.Model; |
4 | 5 |
|
5 | 6 | namespace OnnxStack.FeatureExtractor.Common |
6 | 7 | { |
7 | 8 | public class FeatureExtractorModel : OnnxModelSession |
8 | 9 | { |
9 | | - private readonly int _sampleSize; |
10 | | - private readonly bool _normalize; |
11 | | - private readonly int _channels; |
| 10 | + private readonly FeatureExtractorModelConfig _configuration; |
12 | 11 |
|
13 | 12 | public FeatureExtractorModel(FeatureExtractorModelConfig configuration) |
14 | 13 | : base(configuration) |
15 | 14 | { |
16 | | - _sampleSize = configuration.SampleSize; |
17 | | - _normalize = configuration.Normalize; |
18 | | - _channels = configuration.Channels; |
| 15 | + _configuration = configuration; |
19 | 16 | } |
20 | 17 |
|
21 | | - public int SampleSize => _sampleSize; |
22 | | - |
23 | | - public bool Normalize => _normalize; |
24 | | - |
25 | | - public int Channels => _channels; |
| 18 | + public int OutputChannels => _configuration.OutputChannels; |
| 19 | + public int SampleSize => _configuration.SampleSize; |
| 20 | + public bool NormalizeOutputTensor => _configuration.NormalizeOutputTensor; |
| 21 | + public bool SetOutputToInputAlpha => _configuration.SetOutputToInputAlpha; |
| 22 | + public ImageResizeMode InputResizeMode => _configuration.InputResizeMode; |
| 23 | + public ImageNormalizeType InputNormalization => _configuration.NormalizeInputTensor; |
26 | 24 |
|
27 | 25 | public static FeatureExtractorModel Create(FeatureExtractorModelConfig configuration) |
28 | 26 | { |
29 | 27 | return new FeatureExtractorModel(configuration); |
30 | 28 | } |
31 | 29 |
|
32 | | - public static FeatureExtractorModel Create(string modelFile, bool normalize = false, int sampleSize = 512, int channels = 3, int deviceId = 0, ExecutionProvider executionProvider = ExecutionProvider.DirectML) |
| 30 | + public static FeatureExtractorModel Create(string modelFile, int sampleSize = 0, int outputChannels = 1, bool normalizeOutputTensor = false, ImageNormalizeType normalizeInputTensor = ImageNormalizeType.ZeroToOne, ImageResizeMode inputResizeMode = ImageResizeMode.Crop, bool setOutputToInputAlpha = false, int deviceId = 0, ExecutionProvider executionProvider = ExecutionProvider.DirectML) |
33 | 31 | { |
34 | 32 | var configuration = new FeatureExtractorModelConfig |
35 | 33 | { |
36 | | - SampleSize = sampleSize, |
37 | | - Normalize = normalize, |
38 | | - Channels = channels, |
39 | 34 | DeviceId = deviceId, |
40 | 35 | ExecutionProvider = executionProvider, |
41 | 36 | ExecutionMode = ExecutionMode.ORT_SEQUENTIAL, |
42 | 37 | InterOpNumThreads = 0, |
43 | 38 | IntraOpNumThreads = 0, |
44 | | - OnnxModelPath = modelFile |
| 39 | + OnnxModelPath = modelFile, |
| 40 | + |
| 41 | + |
| 42 | + SampleSize = sampleSize, |
| 43 | + OutputChannels = outputChannels, |
| 44 | + NormalizeOutputTensor = normalizeOutputTensor, |
| 45 | + SetOutputToInputAlpha = setOutputToInputAlpha, |
| 46 | + NormalizeInputTensor = normalizeInputTensor, |
| 47 | + InputResizeMode = inputResizeMode |
45 | 48 | }; |
46 | 49 | return new FeatureExtractorModel(configuration); |
47 | 50 | } |
48 | 51 | } |
49 | | - |
50 | | - public record FeatureExtractorModelConfig : OnnxModelConfig |
51 | | - { |
52 | | - public int SampleSize { get; set; } |
53 | | - public bool Normalize { get; set; } |
54 | | - public int Channels { get; set; } |
55 | | - } |
56 | 52 | } |
0 commit comments