annotations/Azaion.Common/Services/InferenceService.cs

using Azaion.Common.Database;
using Azaion.Common.DTO.Config;
using Azaion.Common.Events;
using Azaion.Common.Extensions;
using Azaion.CommonSecurity.DTO.Commands;
using Azaion.CommonSecurity.Services;
using MediatR;
using MessagePack;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;

namespace Azaion.Common.Services;

public interface IInferenceService
{
    Task RunInference(List<string> mediaPaths, CancellationToken ct = default);
    void StopInference();
}

public class InferenceService : IInferenceService
{
    private readonly IInferenceClient _client;
    private readonly IAzaionApi _azaionApi;
    private readonly IOptions<AIRecognitionConfig> _aiConfigOptions;
    private readonly IAnnotationService _annotationService;
    private readonly IMediator _mediator;
    private CancellationTokenSource _inferenceCancelTokenSource = new();

    public InferenceService(
        ILogger<InferenceService> logger,
        IInferenceClient client,
        IAzaionApi azaionApi,
        IOptions<AIRecognitionConfig> aiConfigOptions,
        IAnnotationService annotationService,
        IMediator mediator)
    {
        _client = client;
        _azaionApi = azaionApi;
        _aiConfigOptions = aiConfigOptions;
        _annotationService = annotationService;
        _mediator = mediator;

        client.InferenceDataReceived += async (sender, command) =>
        {
            try
            {
                if (command.Message == "DONE")
                {
                    _inferenceCancelTokenSource?.Cancel();
                    return;
                }

                var annImage = MessagePackSerializer.Deserialize<AnnotationImage>(command.Data);
                await ProcessDetection(annImage);
            }
            catch (Exception e)
            {
                logger.LogError(e, e.Message);
            }
        };
    }

    private async Task ProcessDetection(AnnotationImage annotationImage, CancellationToken ct = default)
    {
        var annotation = await _annotationService.SaveAnnotation(annotationImage, ct);
        await _mediator.Publish(new AnnotationAddedEvent(annotation), ct);
    }

    public async Task RunInference(List<string> mediaPaths, CancellationToken ct = default)
    {
        _inferenceCancelTokenSource = new CancellationTokenSource();
        _client.Send(RemoteCommand.Create(CommandType.Login, _azaionApi.Credentials));

        var aiConfig = _aiConfigOptions.Value;
        aiConfig.Paths = mediaPaths;
        _client.Send(RemoteCommand.Create(CommandType.Inference, aiConfig));

        using var combinedTokenSource = CancellationTokenSource.CreateLinkedTokenSource(ct, _inferenceCancelTokenSource.Token);
        await combinedTokenSource.Token.AsTask();
    }

    public void StopInference()
    {
        _client.Send(RemoteCommand.Create(CommandType.StopInference));
    }
}