app.py

import torch
from torchvision import transforms
from efficientnet_pytorch import EfficientNet
from huggingface_hub import HfFileSystem
from PIL import Image

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

# Save the EfficientNet model file to disk
efficientnet_model_file = 'efficientnet_b3.pt'
with open(efficientnet_model_file, 'wb') as f:
    f.write(model_content)

# Load the EfficientNet model onto the CPU
model = EfficientNet.from_pretrained('efficientnet-b3')
model.load_state_dict(torch.load(efficientnet_model_file, map_location=torch.device('cpu')))

# Load your custom model onto the CPU
custom_model_file = 'best_model.pth'
model.load_state_dict(torch.load(custom_model_file, map_location=torch.device('cpu')))

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

# Create a Streamlit app with an image upload input
uploaded_file = st.file_uploader('Upload an image')
if uploaded_file is not None:
    # Convert the UploadedFile object to a NumPy array
    image = Image.open(uploaded_file)
    image = np.array(image)
    
    # Use the image for inference
    predictions, predicted_label = image_classifier(image)
    
    # Display the result
    st.write(f'Predictions: {predictions}')
    st.write(f'Predicted label: {predicted_label}')