import gradio as gr
from PIL import Image
import cv2
import numpy as np

demo_inf_ofa = gr.Interface.load(name="spaces/ICML2022/OFA") #"spaces/freddyaboulton/blocks_inputs")
#demo_inf_lama = gr.Interface.load(name="spaces/CVPR/lama-example")  #gr.Blocks.load
#img = gr.Interface.load("spaces/multimodalart/latentdiffusion")(ent[0],'50','256','256','1',10)[0] 

#assert demo_inf("Foo", "bar") == "Foo bar"

def get_ofa(img, txt, ques):
  print("********Inside OFA ********")
  Img = cv2.imread(img) #no use
  #txt = "Visual Grounding"
  print(f"text is :{txt}") #= cv2.imread(img2) #Image.open(img2)
  img_pil = Image.fromarray(Img)  #no use
  img_ofa_out, dummy_txt = gr.Interface.load(name="spaces/ICML2022/OFA")(img, txt, ques)
  img_ofa_out = cv2.imread(img_ofa_out)
  img = Image.open(img) #cv2.imread(img)
  print("Calling Workflow function!")
  img_final = workflow(img, img_ofa_out)
  
  return img_final #img_ofa_out


def get_lama(img1, img2):
  print("********Inside LAMA ********")
  #Img1 = cv2.imread(img1) #Image.open(img1)
  #Img2 = cv2.imread(img2) #Image.open(img2)
  #img1_pil = Image.fromarray(Img1)
  #img2_pil = Image.fromarray(Img2)
  img, mask = gr.Interface.load(name="spaces/CVPR/lama-example")(img1, img2, "automatic (U2net)")
  #TypeError: expected str, bytes or os.PathLike object, not Image
  #TypeError: expected str, bytes or os.PathLike object, not JpegImageFile
  #TypeError: expected str, bytes or os.PathLike object, not Image
  #img = cv2.imread(img)
  print("Going back to Workflow")
  return img
  

def workflow(img, img_ofa_out):
  print("********Inside Workflow ********")
  #Image with visual grounding
  im_vg = img_ofa_out #cv2.imread(image_path_vg)
  green_color = (0,255,0)
  #Getting cordinates for bounding box
  Y, X = np.where(np.all(im_vg==green_color,axis=2))
  zipped = np.column_stack((X,Y))
  img_crop_cords = tuple(list(zipped[0]) + list(zipped[-1]))  
  print(f">>>img_crop_cordinates are : {img_crop_cords}")
  width = img_crop_cords[2] - img_crop_cords[0] #x2-x1
  height = img_crop_cords[3] - img_crop_cords[1] #y2-y1
  print(f">>>Width is:{width}, Height is:{height}")  
  
  #creating mask
  print(">>>creating mask")
  blank_image_mask = 255 * np.ones(shape=(height, width), dtype=np.uint8)
  #cv2_imshow(blank_image_mask)
  cv2.imwrite("image_mask.png", blank_image_mask)
  print(">>>mask created")
  
  #Crop using these derived cordinates
  print(">>>Cropping original image based on cordinates for preparing input to LAMA")
  #img = Image.open(image_path)
  imgcrop = img.crop(img_crop_cords)  
  imgcrop.save("image_crop.jpg")
  #imgcrop
  
  #Call LAMA Space for inference
  print(">>>Calling LAMA with cropped image ")
  img_lama = get_lama("image_crop.jpg", "image_mask.png") #(imgcrop, blank_image_mask) 
  
  #Resize the LAMA image to have cols (width) and rows (height) as original cropped image
  print(">>>Resizing LAMA image to meet original crop size")
  img_lama = cv2.imread(img_lama)
  img_lama_resize = cv2.resize(img_lama, (width, height)) #- src is not a numpy array, neither a scalar
  print(">>>LAMA image resized")
  
  #Paste resized LAMA image on top of original image at correct cordinates
  print(">>>Pasting LAMA image on original image at desired cordinates")
  #img = cv2.imread(img) #getting TypeError: Can't convert object to 'str' for 'filename'
  
  #converting to cv2 image
  #img = np.array(img) 
  img.save('temp.jpg')
  img = cv2.imread("temp.jpg")
  x, y = zipped[0][0], zipped[0][1]
  img[y: y + height, x: x + width] = img_lama_resize
  #cv2_imshow(im_orig)  
  print(">>>LAMA image Pasted on original image")
  cv2.imwrite("image_ofa_lama_out.png", img)
  print(">>>Workflow function complete")
  return img
  
  
demo = gr.Blocks()

with demo:
  gr.Markdown("<h1><center>Testing</center></h1>")
  gr.Markdown(
        """Testing Inference for Gradio. Work in Progress."""
        )
  with gr.Row():
    in_image = gr.Image(type='filepath') #(visible=False) type='numpy'
    in_image_mask = gr.Image(type='filepath') #, source="canvas") #, visible=False) #type='numpy',
    out_image = gr.outputs.Image(type='file') #(type='file')
    
    #in_text1 = gr.Textbox()
    #in_text2 = gr.Textbox()
    #out_text = gr.Textbox()
    
    b1 = gr.Button("Image Button") 
    #b2 = gr.Button("Text Button") 

    b1.click(get_lama,inputs=[in_image, in_image_mask], outputs=out_image) 
    #b2.click(txt_fun, inputs=[in_text1, in_text2], outputs=out_text)
    #examples=examples 

  with gr.Row():
    in_image_ofa = gr.Image(type='filepath') #(visible=False) type='numpy'
    in_txt_ofa = gr.Textbox(value="Visual Grounding", visible=False)
    in_txt_ofa_ques = gr.Textbox(label="Who do you want to remove from your picture?")
    out_image_ofa = gr.outputs.Image(type='file') #(type='file')
    
    #in_text1 = gr.Textbox()
    #in_text2 = gr.Textbox()
    #out_text = gr.Textbox()
    
    b2 = gr.Button("Image Button OFA")
    #b2 = gr.Button("Text Button") 
  
    b2.click(get_ofa,inputs=[in_image_ofa, in_txt_ofa, in_txt_ofa_ques], outputs=out_image_ofa) 
  
demo.launch(enable_queue=True, debug=True)