# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['pil_to_tensor', 'tensor_to_pil', 'iterate_modules', 'tensor_stats_df', 'get_torch_device']

# %% ../nbs/00_core.ipynb 3
# Import necessary modules from the standard library
from pathlib import Path  # For working with file paths
import logging  # For logging messages
import hashlib

# Disable logging warnings
logging.disable(logging.WARNING)

import numpy as np  # For working with arrays

from PIL import Image  # For working with images

import torch  # PyTorch module for deep learning
from torchvision import transforms  # PyTorch module for image transformations

# %% ../nbs/00_core.ipynb 5
def pil_to_tensor(img:Image, # The input PIL image.
                  mean=[0.485, 0.456, 0.406], # The mean values for normalization.
                  std=[0.229, 0.224, 0.225] # The standard deviation values for normalization.
                 ):
    """
    Converts a PIL image to a normalized and batched PyTorch tensor.
    
    Returns:
        The normalized and batched tensor.
    """
    
    return transforms.Normalize(mean, std)(transforms.ToTensor()(img))[None]

# %% ../nbs/00_core.ipynb 13
def tensor_to_pil(tensor: torch.Tensor # The tensor to be converted
                 ):
    """
    Convert a tensor to a PIL image.
    
    Returns:
    img (PIL.Image): The PIL image
    """
    
    # Remove the first dimension if the tensor has 4 dimensions
    if len(tensor.shape) == 4: tensor.squeeze_(0)
        
    # Use the ToPILImage() function from the transforms module to convert the tensor to a PIL image
    return transforms.ToPILImage()(tensor)

# %% ../nbs/00_core.ipynb 17
def iterate_modules(module: torch.nn.Module): # A PyTorch module that contains child modules to be iterated over.
    """
    A generator function that yields the children and grandchildren of a PyTorch module.    
    """
    
    for child in module.children():
        yield child
        yield from iterate_modules(child)

# %% ../nbs/00_core.ipynb 20
import pandas as pd

# %% ../nbs/00_core.ipynb 21
def tensor_stats_df(tensor, # Input tensor for which statistics are to be calculated.
                    attrs = ["mean", "std", "min", "max"], # List of statistics to be calculated.
                    shape=True): # If True, include shape of the tensor in the output.
    """
    Calculate and return statistics of a given tensor as a pandas DataFrame.
    """
    
    attr_dict = {attr: getattr(tensor, attr)().item() for attr in attrs}
    if shape: attr_dict["shape"] = tensor.shape
    return pd.DataFrame.from_dict(attr_dict, orient='index')

# %% ../nbs/00_core.ipynb 24
def get_torch_device():
    """
    This function returns the device to be used for PyTorch computations.
    
    Returns:
    str: "mps" if Metal Performance Shaders (MPS) for MacOS is available, 
         "cuda" if CUDA is available, 
         "cpu" otherwise
    """
    device = (
        "mps"
        if torch.backends.mps.is_available()
        else "cuda"
        if torch.cuda.is_available()
        else "cpu"
    )
    return device
