import cv2
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt
from transformers import pipeline


model = pipeline("object-detection", "facebook/detr-resnet-50")  #loading model

#render function

def render_results(raw_image, model_output):

  raw_image = np.array(raw_image)

  for detection in model_output:
      label = detection['label']
      score = detection['score']
      box = detection['box']
      xmin, ymin, xmax, ymax = box['xmin'], box['ymin'], box['xmax'], box['ymax']

      #Drawing the bounding box
      cv2.rectangle(raw_image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)

      #Puting label and score near the bounding box
      cv2.putText(raw_image, f"{label}: {score:.2f}", (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
  return raw_image



def get_object_counts(detections):  ##to get count of object detected in the image
    object_counts = {}
    for detection in detections:
        label = detection['label']
        if label in object_counts:
            object_counts[label] += 1
        else:
            object_counts[label] = 1

    return object_counts

def generate_output_text(object_counts):  ##to get the output string
    output_text = "In this image there are"
    for label, count in object_counts.items():
        output_text += f" {count} {label},"
    output_text = output_text.rstrip(',') + "."
    return output_text



def main(pil_image):
  pipeline_output = model(pil_image)  #model output
  processed_image = render_results(pil_image, pipeline_output)  ##process image by drawing bounding boxes
  output_text = generate_output_text(get_object_counts(pipeline_output))  ##output string

  return processed_image, output_text


demo = gr.Interface(
    fn = main,
    inputs = gr.Image(label = "Input Image", type = "pil"),
    outputs = [gr.Image(label = "Modle output Predictions", type = "numpy"), gr.Text(label="Output Text")]
)

demo.launch()