utils/predict.py

import torch

classes = [
    "Dark Circle",
    "Oiliness",
    "Spots",
    "Wrinkles"
]

def predict(model, image_tensor, threshold=0.5):

    with torch.no_grad():

        outputs = model(image_tensor)

        probabilities = torch.sigmoid(outputs)

    probs = probabilities.squeeze().tolist()

    results = {}

    detected_conditions = []

    for cls, prob in zip(classes, probs):

        results[cls] = prob

        if prob >= threshold:
            detected_conditions.append(cls)

    return detected_conditions, results