1
- using Microsoft . Azure . CognitiveServices . Vision . CustomVision . Training ;
1
+ using Fritz . StreamLib . Core ;
2
+ using Fritz . StreamTools . Hubs ;
3
+ using Microsoft . AspNetCore . SignalR ;
4
+ using Microsoft . Azure . CognitiveServices . Vision . CustomVision . Training ;
2
5
using Microsoft . Azure . CognitiveServices . Vision . CustomVision . Training . Models ;
3
6
using Microsoft . Extensions . Configuration ;
7
+ using Microsoft . Extensions . DependencyInjection ;
4
8
using Microsoft . Extensions . Hosting ;
5
9
using Microsoft . Extensions . Logging ;
6
10
using System ;
7
11
using System . Collections . Generic ;
8
12
using System . ComponentModel . DataAnnotations ;
13
+ using System . IO ;
14
+ using System . Linq ;
9
15
using System . Text ;
10
16
using System . Threading ;
11
17
using System . Threading . Tasks ;
@@ -18,29 +24,27 @@ public class ScreenshotTrainingService : IHostedService, ITrainHat
18
24
// TODO: Track how many images are loaded -- 5k is the maximum for the FREE service
19
25
private string _CustomVisionKey = "" ;
20
26
private string _AzureEndpoint = "" ;
21
- private string _TwitchChannel = "" ;
22
27
private Guid _AzureProjectId ;
23
- private readonly ILogger _Logger ;
28
+ private ILogger _Logger ;
29
+ private IServiceProvider _Services ;
24
30
private CancellationTokenSource _TokenSource ;
25
31
26
32
private bool _CurrentlyTraining = false ;
27
33
private byte _TrainingCount = 0 ;
28
34
private Task _TrainingTask ;
29
35
private byte _RetryCount = 0 ;
30
36
31
- public string TwitchScreenshotUrl => $ "https://static-cdn.jtvnw.net/previews-ttv/live_user_{ _TwitchChannel } -1280x720.jpg?_=";
32
-
33
- public ScreenshotTrainingService ( IConfiguration configuration , ILoggerFactory loggerFactory )
37
+ public ScreenshotTrainingService ( IConfiguration configuration , ILoggerFactory loggerFactory , IServiceProvider services )
34
38
{
35
39
36
40
_CustomVisionKey = configuration [ "AzureServices:HatDetection:Key" ] ;
37
41
_AzureEndpoint = configuration [ "AzureServices:HatDetection:CustomVisionEndpoint" ] ;
38
- _TwitchChannel = configuration [ "StreamServices:Twitch:Channel" ] ;
39
42
_AzureProjectId = Guid . Parse ( configuration [ "AzureServices:HatDetection:ProjectId" ] ) ;
40
43
_Logger = loggerFactory . CreateLogger ( "ScreenshotTraining" ) ;
41
-
44
+ _Services = services ;
42
45
}
43
46
47
+
44
48
public Task StartAsync ( CancellationToken cancellationToken )
45
49
{
46
50
@@ -113,11 +117,12 @@ private async Task AddScreenshot(bool @internal)
113
117
Endpoint = _AzureEndpoint
114
118
} ;
115
119
116
- var result = await trainingClient . CreateImagesFromUrlsAsync ( _AzureProjectId , new ImageUrlCreateBatch (
117
- new List < ImageUrlCreateEntry > {
118
- new ImageUrlCreateEntry ( TwitchScreenshotUrl + Guid . NewGuid ( ) . ToString ( ) )
119
- }
120
- ) ) ;
120
+ var imageStream = await GetScreenshotFromObs ( ) ;
121
+ // TODO: If imageStream is null, handle gracefully
122
+
123
+ var result = await trainingClient . CreateImagesFromDataAsync ( _AzureProjectId ,
124
+ imageStream
125
+ ) ;
121
126
122
127
if ( ! result . IsBatchSuccessful && _RetryCount < 3 ) {
123
128
_Logger . LogWarning ( $ "Error while adding screenshot #{ _TrainingCount } - trying again in 10 seconds") ;
@@ -144,11 +149,37 @@ private async Task AddScreenshot(bool @internal)
144
149
}
145
150
catch ( Exception ex )
146
151
{
147
- _Logger . LogError ( $ "Error while adding screenshot: { ex . Message } ") ;
152
+
153
+ _Logger . LogError ( $ "Error while adding screenshot: { ex . Message } ") ;
148
154
}
149
155
150
156
}
151
157
158
+ internal async Task < Stream > GetScreenshotFromObs ( )
159
+ {
160
+
161
+ Stream result = null ;
162
+
163
+ ScreenshotSink . Instance . ScreenshotReceived += ( obj , args ) =>
164
+ {
165
+ result = args . Screenshot ;
166
+ } ;
167
+
168
+ using ( var scope = _Services . CreateScope ( ) )
169
+ {
170
+ var obsContext = scope . ServiceProvider . GetRequiredService < IHubContext < ObsHub , ITakeScreenshots > > ( ) ;
171
+ await obsContext . Clients . All . TakeScreenshot ( ) ;
172
+ }
173
+ var i = 0 ;
174
+ while ( result == null ) {
175
+ await Task . Delay ( 100 ) ;
176
+ i ++ ;
177
+ if ( i >= 100 ) break ;
178
+ }
179
+
180
+ return result ;
181
+
182
+ }
152
183
153
184
public Task AddScreenshot ( )
154
185
{
0 commit comments