"""
Created on Sat Apr  9 04:08:02 2022
@author: Admin_with ODD Team

Edited by our team : Sat Oct 5 10:00 2024

references: https://github.com/vinvino02/GLPDepth
"""

import io
import torch
import base64
from config import CONFIG
from torchvision import transforms
from matplotlib import pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from transformers import DetrForObjectDetection, DetrImageProcessor

class DETR:
    def __init__(self):

        self.CLASSES = [
            'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
            'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A',
            'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
            'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack',
            'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
            'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
            'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
            'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
            'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
            'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A',
            'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
            'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A',
            'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
            'toothbrush'
        ]
        
        self.COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], 
                       [0.929, 0.694, 0.125], [0, 0, 1], [0.466, 0.674, 0.188], 
                       [0.301, 0.745, 0.933]]
        
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

        self.model = DetrForObjectDetection.from_pretrained(CONFIG['detr_model_path'], revision="no_timm")
        self.model.to(CONFIG['device'])
        self.model.eval()


    def box_cxcywh_to_xyxy(self, x):
        x_c, y_c, w, h = x.unbind(1)
        b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
             (x_c + 0.5 * w), (y_c + 0.5 * h)]
        return torch.stack(b, dim=1).to(CONFIG['device'])

    def rescale_bboxes(self, out_bbox, size):
        img_w, img_h = size
        b = self.box_cxcywh_to_xyxy(out_bbox)
        b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32).to(CONFIG['device'])
        return b

    
    def detect(self, im):
        img = self.transform(im).unsqueeze(0).to(CONFIG['device'])
        assert img.shape[-2] <= 1600 and img.shape[-1] <= 1600, 'Image too large'
        outputs = self.model(img)
        probas = outputs['logits'].softmax(-1)[0, :, :-1]
        keep = probas.max(-1).values > 0.7
        bboxes_scaled = self.rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)
        return probas[keep], bboxes_scaled
    
    def visualize(self, im, probas, bboxes):
        """
        Visualizes the detected bounding boxes and class probabilities on the image.

        Parameters:
            im (PIL.Image): The original input image.
            probas (Tensor): Class probabilities for detected objects.
            bboxes (Tensor): Bounding boxes for detected objects.
        """
        # Convert image to RGB format for matplotlib
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.imshow(im)
        
        # Iterate over detections and draw bounding boxes and labels
        for p, (xmin, ymin, xmax, ymax), color in zip(probas, bboxes, self.COLORS * 100):
            # Detach tensors and convert to float
            xmin, ymin, xmax, ymax = map(lambda x: x.detach().cpu().numpy().item(), (xmin, ymin, xmax, ymax))
            
            ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
                                        fill=False, color=color, linewidth=3))
            cl = p.argmax()
            text = f'{self.CLASSES[cl]}: {p[cl].detach().cpu().numpy():0.2f}'  # Detach probability as well
            ax.text(xmin, ymin, text, fontsize=15, bbox=dict(facecolor='yellow', alpha=0.5))
        
        ax.axis('off')

        # Convert the plot to a PIL Image and then to bytes
        canvas = FigureCanvas(fig)
        buf = io.BytesIO()
        canvas.print_png(buf)
        buf.seek(0)

        # Base64 encode the image
        img_bytes = buf.getvalue()
        img_base64 = base64.b64encode(img_bytes).decode('utf-8')

        # Close the figure to release memory
        plt.close(fig)

        return img_base64