# -*- coding: utf-8 -*-

# %%capture
# #Use capture to not show the output of installing the libraries!
import gradio as gr
import requests
import torch
import torch.nn as nn
from PIL import Image
from torchvision.models import resnet50
from torchvision.transforms import functional as F
import numpy as np
import tensorflow as tf
from transformers import pipeline

# load the model from the Hugging Face Model Hub
model = pipeline('image-classification', model='image_classification/densenet')


#model = tf.keras.models.load_model('/content/drive/MyDrive/project_image_2023_NO/saved_models/saved_model/densenet')
#labels = ['Healthy', 'Patient']
labels = {0: 'healthy', 1: 'patient'}

def classify_image(inp):
  inp = inp.reshape((-1, 224, 224, 3))
  inp = tf.keras.applications.densenet.preprocess_input(inp)
  prediction = model.predict(inp)
  confidences = {labels[i]: float(prediction[0][i]) for i in range(2)}
  return confidences


gr.Interface(fn=classify_image, 
             inputs=gr.Image(shape=(224, 224)),
             outputs=gr.Label(num_top_classes = 2),
             title="Demo",
             description="Here's a sample image classification. Enjoy!",
             examples=[['path/to/example/image.jpg']]
             ).launch(share = True)