import gradio as gr
import torch
from transformers import AutoFeatureExtractor, AutoTokenizer, VisionEncoderDecoderModel, GPT2Tokenizer, pipeline
import os

HF_DATASETS_OFFLINE=1
TRANSFORMERS_OFFLINE=1

device = 'cpu'
auth_token = os.getenv("auth_token")
#auth_token = os.environ.get("auth_token")

max_length = 100
num_beams = 4
gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
def predict_step(image_paths, model):
  images = []
  for image_path in image_paths:
    i_image = Image.open(image_path)
    if i_image.mode != "RGB":
      i_image = i_image.convert(mode="RGB")

    images.append(i_image)

  pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values
  pixel_values = pixel_values.to(device)

  output_ids = model.generate(pixel_values, **gen_kwargs)

  preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
  preds = [pred.strip() for pred in preds]
  return preds

def predict_step_image(dataset_images, feature_extractor, model):
  results = []
  for i in dataset_images:
    pixel_values = feature_extractor(images=i, return_tensors="pt").pixel_values
    pixel_values = pixel_values.to(device)

    output_ids = model.generate(pixel_values, **gen_kwargs)

    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
    preds = [pred.strip() for pred in preds]
    results.append(preds)
  return results
 
def predict_step_single_image(image, tokenizer, feature_extractor, model):
    results=[]
    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
    pixel_values = pixel_values.to(device)

    output_ids = model.generate(pixel_values, **gen_kwargs)

    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
    preds = [pred.strip() for pred in preds]
    results.append(preds)
    return results

def predict_step_pixel(dataset_pixel_values, model):
  results=[]
  for pv in dataset_pixel_values:
    pixel_values = pv.reshape([1,3,224,224])
    pixel_values = pixel_values.to(device)
    output_ids = model.generate(pixel_values, **gen_kwargs)
    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
    results.append([pred.strip() for pred in preds][0])
  return results

"""
    image methods
"""
def load_image2txt_model(image_model_name):
    model = VisionEncoderDecoderModel.from_pretrained(image_model_name)
    feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/swin-large-patch4-window7-224", use_auth_token=auth_token)
    
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2", use_auth_token=auth_token)
    tokenizer.pad_token = tokenizer.eos_token

    model = model.to(device)
    return tokenizer, feature_extractor, model

def inference_image_pipe(image_input):
    image_model_name = "./checkpoint-21000"

    tokenizer, feature_extractor, image_model = load_image2txt_model(image_model_name)
    text = predict_step_single_image(image_input, tokenizer, feature_extractor, image_model)[0]
    return text

with gr.Interface(fn=inference_image_pipe, 
             inputs=gr.Image(shape=(256, 256)),
             outputs="text",
             examples=["3212210S4492629-1.png", "3216497S4499373-1.png"]) as demo:
    gr.Markdown("POC XRaySwinGen - Automatic Medical Report")
    
    
if __name__ == "__main__":
    demo.launch()