include/fsdk/Estimators/IMedicalMaskEstimator.h

#pragma once

#pragma once

#include <fsdk/IObject.h>
#include <fsdk/FSDKError.h>
#include <fsdk/Optional.h>
#include <fsdk/Types.h>

namespace fsdk {

#ifndef DOXYGEN_SHOULD_SKIP_THIS
    DECLARE_SMARTPTR(IMedicalMaskEstimator);
#endif

    enum class MedicalMask {
        Mask = 0,    
        NoMask,      
        OccludedFace 
    };

    enum class MedicalMaskExtended {
        Mask = 0,       
        NoMask,         
        MaskNotInPlace, 
        OccludedFace    
    };

    enum class DetailedMaskType {
        CorrectMask = 0,            
        MouthCoveredWithMask,       
        ClearFace,                  
        ClearFaceWithMaskUnderChin, 
        PartlyCoveredFace,          
        FullMask,                   
        Count
    };

    struct MedicalMaskEstimation {
        MedicalMask result;        
        DetailedMaskType maskType; 

        // scores
        float maskScore;         
        float noMaskScore;       
        float occludedFaceScore; 

        float scores[static_cast<int>(DetailedMaskType::Count)]; 

        inline float getScore(DetailedMaskType type) const;
    };

    struct MedicalMaskEstimationExtended {
        MedicalMaskExtended result; 
        // scores
        float maskScore;         
        float noMaskScore;       
        float maskNotInPlace;    
        float occludedFaceScore; 
    };

    struct IMedicalMaskEstimator : IRefCounted {
        virtual Result<FSDKError>
        estimate(const Image& warp, MedicalMaskEstimation& estimation) const noexcept = 0;

        virtual Result<FSDKError>
        estimate(const Image& warp, MedicalMaskEstimationExtended& estimation) const noexcept = 0;

        virtual Result<FSDKError> estimate(
            const Image& image,
            const Detection& detection,
            MedicalMaskEstimation& estimation) const noexcept = 0;

        virtual Result<FSDKError> estimate(
            const Image& image,
            const Detection& detection,
            MedicalMaskEstimationExtended& estimation) const noexcept = 0;

        virtual Result<FSDKError>
        estimate(Span<const Image> warps, Span<MedicalMaskEstimation> estimations) const noexcept = 0;

        virtual Result<FSDKError>
        estimate(Span<const Image> warps, Span<MedicalMaskEstimationExtended> estimations) const noexcept = 0;

        virtual Result<FSDKError> estimate(
            Span<const Image> images,
            Span<const Detection> detections,
            Span<MedicalMaskEstimation> estimations) const noexcept = 0;

        virtual Result<FSDKError> estimate(
            Span<const Image> images,
            Span<const Detection> detections,
            Span<MedicalMaskEstimationExtended> estimations) const noexcept = 0;
        virtual Result<FSDKError>
        validate(Span<const Image> warps, Span<Result<FSDKError>> errors) const noexcept = 0;

        virtual Result<FSDKError> validate(
            Span<const Image> images,
            Span<const Detection> detections,
            Span<Result<FSDKError>> errors) const noexcept = 0;
    };

    /*
        Implementation details.
    */
    float MedicalMaskEstimation::getScore(DetailedMaskType type) const {
        switch(type) {
        case DetailedMaskType::CorrectMask:
        case DetailedMaskType::MouthCoveredWithMask:
        case DetailedMaskType::ClearFace:
        case DetailedMaskType::ClearFaceWithMaskUnderChin:
        case DetailedMaskType::PartlyCoveredFace:
        case DetailedMaskType::FullMask:
            return scores[static_cast<int>(type)];
        default:
            return 0.f;
        }
    }

} // namespace fsdk