import streamlit as st
from PIL import Image
import torch
from torchvision import transforms, utils
from facenet_pytorch import MTCNN
from torchvision.transforms.functional import to_pil_image

# Function to load the ViT model and MTCNN
def load_model_and_mtcnn(model_path):
    model = torch.load(model_path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    mtcnn = MTCNN(keep_all=True, device=device)
    return model, device, mtcnn

def detect_and_process_skin(image_bytes):
    """Detects faces in an image, crops the skin region, and returns it as an image object."""
    # Load image from bytes
    img = Image.open(io.BytesIO(image_bytes))
    img_np = np.array(img)
    img_rgb = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)

    # Detect faces in the image
    detections = mtcnn.detect_faces(img_rgb)

    # Check if any faces were detected
    if detections:
        x, y, width, height = detections[0]['box']

        # Crop the face region
        face_img_np = img_np[y:y+height, x:x+width]

        # Convert to PIL Image for return
        pil_img = Image.fromarray(face_img_np)
        return pil_img
    else:
        # Return original image if no face was detected
        return img

# Function to preprocess the image and return both the tensor and the final PIL image for display
def preprocess_image(image, mtcnn, device):
    processed_image = image  # Initialize with the original image
    try:
        # Directly call mtcnn with the image to get cropped faces
        cropped_faces = mtcnn(image)
        if cropped_faces is not None and len(cropped_faces) > 0:
            # Convert the first detected face tensor back to PIL Image for further processing
            processed_image = to_pil_image(cropped_faces[0].cpu(),mode='BGR;16')
    except Exception as e:
        st.write(f"Exception in face detection: {e}")
        processed_image = image
    
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    image_tensor = transform(processed_image).to(device)
    image_tensor = image_tensor.unsqueeze(0)  # Add a batch dimension
    return image_tensor, processed_image

# Function for inference
def predict(image_tensor, model, device):
    model.eval()
    with torch.no_grad():
        outputs = model(image_tensor)
        # Adjust for your model's output if it does not have a 'logits' attribute
        probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
        predicted_class = torch.argmax(probabilities, dim=1)
    return predicted_class, probabilities

# Streamlit UI setup
st.title("Face Detection and Classification with ViT")
st.write("Upload an image, and the model will detect faces and classify the image.")

model_path = "model_v1.0.pt"  # Adjust this path as necessary
model, device, mtcnn = load_model_and_mtcnn(model_path)

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption='Uploaded Image', use_column_width=True)
    image1 = image.getvalue()
    image_ten = detect_and_process_skin(image1)
    image_tensor, final_image = preprocess_image(image, mtcnn, device)
    predicted_class, probabilities = predict(image_ten, model, device)
    
    st.write(f"Predicted class: {predicted_class.item()}")
    # Display the final processed image
   # st.image(final_image, caption='Processed Image', use_column_width=True)
    img_bytes = io.BytesIO()
    detect_and_process_skin(image1.getvalue()).save(img_bytes, format='JPEG')
    st.image(img_bytes.getvalue(), width=250, caption="Processed Image")