infer.py

def infer(input_data, model=model):
    inference = []
    
    classes = np.array(['desert', 'mountain', 'sea', 'sunset', 'trees'])
    
    y_pred = model.predict(input_data)
    
    # Performing masking
    y_pred = (y_pred > 0.5) * 1.0
    
    for i in range(y_pred.shape[0]):
        # select the indices
        indices = np.where(y_pred[i] == 1.0)[0]
        # Adding the results 
        inference.append(classes[indices].tolist())
        
    return inference