viewpoint.py

import torch

def get_viewpoint(model, image, device):
    model.eval()  # Set model to evaluate mode

    with torch.no_grad():
        image = image.to(device)
        # Get the prediction probabilities, the highest is the viewpoint
        pred_label = model(image)
        viewpoint = torch.argmax(pred_label, 1)

    return viewpoint