from transformers import pipeline
from transformers import DetrFeatureExtractor, DetrForObjectDetection
from PIL import Image, ImageDraw, ImageFont
import gradio as gr

# Initialize another model and feature extractor
feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-50')
model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')

# Initialize the object detection pipeline
object_detector = pipeline("object-detection", model = model, feature_extractor = feature_extractor)

# Draw bounding box definition
def draw_bounding_box(im, score, label, xmin, ymin, xmax, ymax, index, num_boxes):
    """ Draw a bounding box. """
    # Draw the actual bounding box
    outline = ' '
    if label in ['truck', 'car', 'motorcycle', 'bus']:
      outline = 'red'
    elif label in ['person', 'bicycle']:
      outline = 'green'
    else:
      outline = 'blue'
    im_with_rectangle = ImageDraw.Draw(im)  
    im_with_rectangle.rounded_rectangle((xmin, ymin, xmax, ymax), outline = outline, width = 3, radius = 10)

    # Return the result
    return im

def detect_image(im):
    # Perform object detection
    bounding_boxes = object_detector(im)

    # Iteration elements
    num_boxes = len(bounding_boxes)
    index = 0

    # Draw bounding box for each result
    for bounding_box in bounding_boxes:
        if bounding_box['label'] in ['person','motorcycle','bicycle', 'truck', 'car','bus']:
            box = bounding_box['box']

            #Draw the bounding box
            output_image = draw_bounding_box(im, bounding_box['score'],
                                  bounding_box['label'],
                                  box['xmin'], box['ymin'],
                                  box['xmax'], box['ymax'],
                                  index, num_boxes)
            index += 1

    return output_image
    
iface = gr.Interface(detect_image, gr.inputs.Image(type = 'pil'), gr.outputs.Image()).launch(share = True)