import io
import torch
from torchvision import transforms
from huggingface_hub import HfFileSystem
from PIL import Image


# Authenticate and download the EfficientNet model from Hugging Face
fs = HfFileSystem()
efficientnet_model_path = 'dhhd255/efficientnet_b3/efficientnet_b3.pt'
with fs.open(efficientnet_model_path, 'rb') as f:
    efficientnet_model_content = f.read()

# Load the EfficientNet model onto the CPU
efficientnet_model_file = io.BytesIO(efficientnet_model_content)
efficientnet_model = torch.load(efficientnet_model_file, map_location=torch.device('cpu'))

# Authenticate and download your custom model from Hugging Face
custom_model_path = 'dhhd255/efficient_net_parkinsons/best_model.pth'
with fs.open(custom_model_path, 'rb') as f:
    custom_model_content = f.read()

# Load your custom model onto the CPU
custom_model_file = io.BytesIO(custom_model_content)
custom_model_state_dict = torch.load(custom_model_file, map_location=torch.device('cpu'))

# Create a new instance of your custom model
model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model=efficientnet_model)

# Load your custom model into the new instance
model.load_state_dict(custom_model_state_dict)
model.eval()

# Define a function that takes an image as input and uses the model for inference
def image_classifier(image):
    # Preprocess the input image
    data_transform = transforms.Compose([
        transforms.Lambda(lambda x: x.convert('RGB')),
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ])
    image = Image.fromarray(image)
    image = data_transform(image)
    image = image.unsqueeze(0)
    
    # Use your custom model for inference
    with torch.no_grad():
        outputs = model(image)
        _, predicted = torch.max(outputs.data, 1)
    
    # Map the index to a class label
    labels = ['Healthy', 'Parkinson']
    predicted_label = labels[predicted.item()]
    
    # Return the result
    return outputs[0].numpy(), predicted_label

# Load and preprocess the image
img_path = '/content/test_image_healthy.png'
img = Image.open(img_path)
img = data_transform(img)

# Add a batch dimension and move the image to the device
img = img.unsqueeze(0)

# Perform inference
with torch.no_grad():
 outputs = model(img)
 _, predicted = torch.max(outputs.data, 1)
 print(f'Predicted class: {predicted.item()}')