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