import numpy as np
import gradio as gr
import math

demo = gr.Blocks()

def estimate(input_img):
  threshold = 0.9
  img  = Image.fromarray(input_img.astype('uint8'),'RGB')

  (w, h) = img.size
  image_intrinsics = np.array([[w + h, 0, w // 2], [0, w + h, h // 2], [0, 0, 1]])
      
  res = img2pose_model.predict([transform(img)])[0]

  all_bboxes = res["boxes"].cpu().numpy().astype('float')

  poses = []
  bboxes = []
  for i in range(len(all_bboxes)):
      if res["scores"][i] > threshold:
          bbox = all_bboxes[i]
          pose_pred = res["dofs"].cpu().numpy()[i].astype('float')
          pose_pred = pose_pred.squeeze()

          poses.append(pose_pred)  
          bboxes.append(bbox)

  trans_vertices = renderer.transform_vertices(img, poses)
  img = renderer.render(img, trans_vertices, alpha=1)    

  plt.figure(figsize=(8, 8))     
    
  for bbox in bboxes:
      plt.gca().add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2] - bbox[0], bbox[3] - bbox[1],linewidth=3,edgecolor='b',facecolor='none'))            
    
  plt.imshow(img)        
  plt.show()
  pose_pred[0] = math.degrees(pose_pred[0])
  pose_pred[1] = math.degrees(pose_pred[1])
  pose_pred[2] = math.degrees(pose_pred[2])
  listToStr = ' '.join([str(elem) for elem in pose_pred])
  return img,pose_pred[0],pose_pred[1],pose_pred[2],pose_pred[3],pose_pred[4],pose_pred[5]


def estimate_crowd(input_img):
  threshold = 0.5
  img  = Image.fromarray(input_img.astype('uint8'),'RGB')

  (w, h) = img.size
  image_intrinsics = np.array([[w + h, 0, w // 2], [0, w + h, h // 2], [0, 0, 1]])
  
  res = img2pose_model.predict([transform(img)])[0]

  all_bboxes = res["boxes"].cpu().numpy().astype('float')

  poses = []
  bboxes = []
  for i in range(len(all_bboxes)):
      if res["scores"][i] > threshold:
          bbox = all_bboxes[i]
          pose_pred = res["dofs"].cpu().numpy()[i].astype('float')
          pose_pred = pose_pred.squeeze()

          poses.append(pose_pred)  
          bboxes.append(bbox)


  #render_plot(img.copy(), poses, bboxes)
  
  n = str((len(bboxes)))
  (w, h) = img.size
  image_intrinsics = np.array([[w + h, 0, w // 2], [0, w + h, h // 2], [0, 0, 1]])
    
  trans_vertices = renderer.transform_vertices(img, poses)
  img = renderer.render(img, trans_vertices, alpha=1) 
  return img, n


with demo:

    gr.Markdown("Head Pose Estimation in Crowd")

    with gr.Tabs():
        with gr.TabItem("What is Head Pose Estimation?"):
           gr.Markdown(
        """
            ### Head pose estimation in crowd (HPECS) is a task of estimating the head pose of people in the crowd scenes.

            Six degrees of freedom of head pose is the freedom of movement of the head pose in three-dimensional space. 6DoF is the 
            abbreviation of six degrees of freedoms. 6DoF of head pose records the rotational motion such as
            pitch, yaw and roll , and translational motions of the head pose including up/down, left/right and front/back. 
            6DoF of 3D head pose is being regressed as it can harvest more versatile 
            spatial information which able to track the movement of object even in a still image. 

    """
          )
           gr.Image("/content/gdrive/MyDrive/project_folder/img2pose/6DoF.gif",size=20)

 
        with gr.TabItem("Head Pose Estimation"):
            gr.Markdown(
        """
            ### Please insert image with single head

    """
          )
            image = gr.inputs.Image()
            text_button = gr.Button("Estimate")
            with gr.Row():
              outputs=[gr.outputs.Image(),gr.outputs.Textbox(type="auto", label="Yaw"),
                                              gr.outputs.Textbox(type="auto", label="Pitch"),
                                              gr.outputs.Textbox(type="auto", label="Roll"),
                                              gr.outputs.Textbox(type="auto", label="tx"),
                                              gr.outputs.Textbox(type="auto", label="ty"),
                                              gr.outputs.Textbox(type="auto", label="tz")]
            
        with gr.TabItem("Estimate Head Pose in Crowd"):
            gr.Markdown(
          """
              ### Please insert crowd scene image

      """
            )
            image_input = gr.Image()
            with gr.Row():
                
                image_output = [gr.outputs.Image(),gr.outputs.Textbox(type="auto", label="Number of head pose estimated"),
                                            ]
            image_button = gr.Button("Estimate")

    text_button.click(estimate, inputs=image, outputs=outputs)
    image_button.click(estimate_crowd, inputs=image_input, outputs=image_output)

demo.launch()