detr_tracker.py

# -*- coding: utf-8 -*-
"""detr_demo.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/github/facebookresearch/detr/blob/colab/notebooks/detr_demo.ipynb

# Object Detection with DETR - a minimal implementation

In this notebook we show a demo of DETR (Detection Transformer), with slight differences with the baseline model in the paper.

We show how to define the model, load pretrained weights and visualize bounding box and class predictions.

Let's start with some common imports.
"""

# Commented out IPython magic to ensure Python compatibility.
from PIL import Image, ImageDraw, ImageFont
import av
import cv2
import numpy as np
import requests
import matplotlib.pyplot as plt
import time
from pathlib import Path
import os
import sys
import shutil
# %config InlineBackend.figure_format = 'retina'
import random
from numba import cuda
import torch
from torch import nn
from torchvision.models import resnet50
import torchvision.transforms as T
torch.set_grad_enabled(False)
from deep_sort import DeepSort

deepsort = DeepSort("deep_sort/deep/checkpoint/ckpt.t7")

def bbox_rel(image_width, image_height, bbox_left, bbox_top, bbox_w, bbox_h):
    """" Calculates the relative bounding box from absolute pixel values. """
    x_c = (bbox_left + bbox_w / 2)
    y_c = (bbox_top + bbox_h / 2)
    w = bbox_w
    h = bbox_h
    return x_c, y_c, w, h



"""## DETR
Here is a minimal implementation of DETR:
"""

class DETRdemo(nn.Module):
    """
    Demo DETR implementation.

    Demo implementation of DETR in minimal number of lines, with the
    following differences wrt DETR in the paper:
    * learned positional encoding (instead of sine)
    * positional encoding is passed at input (instead of attention)
    * fc bbox predictor (instead of MLP)
    The model achieves ~40 AP on COCO val5k and runs at ~28 FPS on Tesla V100.
    Only batch size 1 supported.
    """
    def __init__(self, num_classes, hidden_dim=256, nheads=8,
                 num_encoder_layers=6, num_decoder_layers=6):
        super().__init__()

        # create ResNet-50 backbone
        self.backbone = resnet50()
        del self.backbone.fc

        # create conversion layer
        self.conv = nn.Conv2d(2048, hidden_dim, 1)

        # create a default PyTorch transformer
        self.transformer = nn.Transformer(
            hidden_dim, nheads, num_encoder_layers, num_decoder_layers)

        # prediction heads, one extra class for predicting non-empty slots
        # note that in baseline DETR linear_bbox layer is 3-layer MLP
        self.linear_class = nn.Linear(hidden_dim, num_classes + 1)
        self.linear_bbox = nn.Linear(hidden_dim, 4)

        # output positional encodings (object queries)
        self.query_pos = nn.Parameter(torch.rand(100, hidden_dim))

        # spatial positional encodings
        # note that in baseline DETR we use sine positional encodings
        self.row_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))
        self.col_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))

    def forward(self, inputs):
        # propagate inputs through ResNet-50 up to avg-pool layer
        x = self.backbone.conv1(inputs)
        x = self.backbone.bn1(x)
        x = self.backbone.relu(x)
        x = self.backbone.maxpool(x)

        x = self.backbone.layer1(x)
        x = self.backbone.layer2(x)
        x = self.backbone.layer3(x)
        x = self.backbone.layer4(x)

        # convert from 2048 to 256 feature planes for the transformer
        h = self.conv(x)

        # construct positional encodings
        H, W = h.shape[-2:]
        pos = torch.cat([
            self.col_embed[:W].unsqueeze(0).repeat(H, 1, 1),
            self.row_embed[:H].unsqueeze(1).repeat(1, W, 1),
        ], dim=-1).flatten(0, 1).unsqueeze(1)

        # propagate through the transformer
        h = self.transformer(pos + 0.1 * h.flatten(2).permute(2, 0, 1),
                             self.query_pos.unsqueeze(1)).transpose(0, 1)
        
        # finally project transformer outputs to class labels and bounding boxes yeah show me the erro
        return {'pred_logits': self.linear_class(h), 
                'pred_boxes': self.linear_bbox(h).sigmoid()}

"""As you can see, DETR architecture is very simple, thanks to the representational power of the Transformer. There are two main components:
* a convolutional backbone - we use ResNet-50 in this demo
* a Transformer - we use the default PyTorch nn.Transformer

Let's construct the model with 80 COCO output classes + 1 ⦰ "no object" class and load the pretrained weights.
The weights are saved in half precision to save bandwidth without hurting model accuracy.
"""
    
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

detr = DETRdemo(num_classes=91)
state_dict = torch.hub.load_state_dict_from_url(
    url='https://dl.fbaipublicfiles.com/detr/detr_demo-da2a99e9.pth',
    map_location='cuda', check_hash=True)
detr.load_state_dict(state_dict)
detr.eval().to('cuda')

# Load model
# detr = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=True)
# detr.eval().to(DEVICE)

"""## Computing predictions with DETR

The pre-trained DETR model that we have just loaded has been trained on the 80 COCO classes, with class indices ranging from 1 to 90 (that's why we considered 91 classes in the model construction).
In the following cells, we define the mapping from class indices to names.
"""

# COCO classes
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'
]

# colors for visualization
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
COLORS = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(CLASSES))]
"""DETR uses standard ImageNet normalization, and output boxes in relative image coordinates in $[x_{\text{center}}, y_{\text{center}}, w, h]$ format, where $[x_{\text{center}}, y_{\text{center}}]$ is the predicted center of the bounding box, and $w, h$ its width and height. Because the coordinates are relative to the image dimension and lies between $[0, 1]$, we convert predictions to absolute image coordinates and $[x_0, y_0, x_1, y_1]$ format for visualization purposes."""

# standard PyTorch mean-std input image normalization
transform = T.Compose([
    T.Resize(800),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# for output bounding box post-processing
def box_cxcywh_to_xyxy(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)

def xyxy2xywh(x):
    # Convert bounding box format from [x1, y1, x2, y2] to [x, y, w, h]
    y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)
    y[0] = (x[0] + x[2]) / 2
    y[1] = (x[1] + x[3]) / 2
    y[2] = x[2] - x[0]
    y[3] = x[3] - x[1]
    return y

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

"""Let's put everything together in a `detect` function:"""

def detect(im, model, transform):
    # mean-std normalize the input image (batch-size: 1)
    img = transform(im).unsqueeze(0).to(DEVICE)
    model.to(DEVICE)

    # demo model only support by default images with aspect ratio between 0.5 and 2
    # if you want to use images with an aspect ratio outside this range
    # rescale your image so that the maximum size is at most 1333 for best results
    assert img.shape[-2] <= 1600 and img.shape[-1] <= 1600, 'demo model only supports images up to 1600 pixels on each side'

    # propagate through the model
    outputs = model(img)

    # keep only predictions with 0.7+ confidence
    probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]
    keep = probas.max(-1).values > 0.7

    # convert boxes from [0; 1] to image scales
    bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size).to('cuda')
    # print(f"probas: {type(probas)}")
    # print(f"\r{probas}" )
    return probas[keep], bboxes_scaled

def plot_results(pil_img, prob, boxes):
    plt.figure(figsize=(16,10))
    plt.imshow(pil_img)
    ax = plt.gca()
    for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), COLORS * 100):
        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill=False, color=c, linewidth=3))
        cl = p.argmax()
        text = f'{CLASSES[cl]}: {p[cl]:0.2f}'
        ax.text(xmin, ymin, text, fontsize=15,
                bbox=dict(facecolor='yellow', alpha=0.5))
    plt.axis('off')
    plt.show()
    # plt.pause(1)
    plt.close('all')

def tracking_plot(PIL_img,prob,boxes):
    img_w,img_h = PIL_img.size
    bbox_xywh = []
    confs = []
    for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), COLORS * 100):
        # print(f"xmin {xmin} ymin {ymin} xmax {xmax} ymax {ymax}");exit()
        cl = p.argmax()
        if CLASSES[cl] == 'person':#filtering Person Class

            obj = list(xyxy2xywh([xmin,ymin,xmax,ymax]))
        # print(obj)
            conf = p[cl].tolist()
            bbox_xywh.append(obj)
            confs.append(conf)
            cl = p.argmax()
        # text = f'{CLASSES[cl]}: {p[cl]:0.2f}'
        # print(text)
        # if CLASSES[cl] == 'person':
            # print("detected person")
        # print([torch.Tensor(boxes)])
        else:
            pass
    try:
        outputs = deepsort.update(torch.Tensor(bbox_xywh),torch.Tensor(confs),np.asarray(PIL_img))
    
        if len(outputs):
            bbox_xyxy = outputs[:, :4]
            identities = outputs[:, -1]
            # print("bbox_xyxy: ",bbox_xyxy)
            # print("identities: ",identities)
            draw_boxes(PIL_img, scores,bbox_xyxy, identities)
            # draw_boxes(im,scores,boxes,index)
    except IndexError:
        pass
        # print("indexError")




#not_working
def write_Video(input_stream,frames):
    print("writing Video")
    video_stream = _input.streams.video[0]
    # print(video_stream)
    if not os.path.exists("output.mp4"):
        Path("output.mp4").touch()
    output_container = av.open('output.mp4', mode='w')
    stream = output_container.add_stream('h264', rate=video_stream.average_rate)
    stream.width = video_stream.width
    stream.height = video_stream.height
    stream.pix_fmt = 'yuv420p'
    # perform sorting
    # frames = sorted(frames, key=lambda f: f.pts)
    assert os.path.exists("output.mp4")
    for frame in frames:
        # let libav decide the correct pts and time base (in case of changing fps)
        # frame.pts = None
        # frame.time_base = None
        frame = av.VideoFrame.from_image(frame)
        output_container.mux(stream.encode(frame))
    output_container.close()



import imageio
from random import randint
import cv2
out_file = f"output/out{randint(99,1000)}.mp4"
video = imageio.get_writer(out_file,fps = 30)
# fourcc = cv2.cv.CV_FOURCC(*'XVID')
# out = cv2.VideoWriter('output.avi',fourcc, 20.0, (640,480))
def draw_boxes(im, prob, boxes,ind):
    # print(f"\ncurrent indentities: {ind}",end = '',flush=True)
    for p, (xmin, ymin, xmax, ymax), c,index in zip(prob, boxes.tolist(), COLORS * 100,ind):
        xmin,ymin,xmax,ymax = int(xmin),int(ymin),int(xmax),int(ymax)
        cl = p.argmax()
        text = f'{index} {CLASSES[cl]}'#: {p[cl]:0.2f}'
        im2_update = ImageDraw.Draw(im)
        im2_update.rectangle(((xmin,ymin),(xmax,ymax)),outline=tuple(c))
        font = ImageFont.truetype('Roboto-Black.ttf',15)
        im2_update.text((xmin,ymin),text,font=font,fill=(0,0,0,0))
        # cv2_img = np.asarray(im)[:,:,::-1].copy()
        # out.write(cv2_img)
        # im.save(f"output/Image{time.time()}.jpg")
        # exit()
    video.append_data(np.asarray(im))

# import cv2
# from PIL import Image
# import numpy 
# from random import randint
# out_file = f"output/out{randint(99,1000)}.mp4"
# result = cv2.VideoWriter('output/filename.avi',  
#                          cv2.VideoWriter_fourcc(*'mp4v'), 
#                          30, (640,480)) 
# def draw_boxes(im, prob, boxes,ind):
#     print(f"\rcurrent index: {ind}",end = '',flush=True)
#     fourcc = cv2.VideoWriter_fourcc(*'MP4V')
#     out = cv2.VideoWriter(out_file, fourcc, 20.0, (640,480))

#     for p, (xmin, ymin, xmax, ymax), c,index in zip(prob, boxes.tolist(), COLORS * 100,ind):
#         xmin,ymin,xmax,ymax = int(xmin),int(ymin),int(xmax),int(ymax)
#         cl = p.argmax()
#         text = f'{index}:{CLASSES[cl]}'#: {p[cl]:0.2f}'
#         im2_update = ImageDraw.Draw(im)
#         im2_update.rectangle(((xmin,ymin),(xmax,ymax)),outline=tuple(c))
#         font = ImageFont.truetype('Roboto-Black.ttf',15)
#         im2_update.text((xmin,ymin),text,font=font,fill=(0,0,0,0))

#         # pil_image = Image.open(im).convert('RGB') 
#         open_cv_image = numpy.array(im) 
#         open_cv_image = open_cv_image[:, :, ::-1].copy() 
#         #out.write(open_cv_image)
#         result.write(open_cv_image)
    #out.release()
"""## Using DETR
To try DETRdemo model on your own image just change the URL below.
"""

"""Let's now visualize the model predictions"""
# file_path = "bom.mp4"
file_path = "track2.mp4"
# file_path = "HD_Store.mp4"
file_path = "output.mkv"
_input = av.open(file_path)

det_fps =  0.0
print("started")
t0 = time.time()
for index,frame in enumerate(_input.decode(video=0)):
    t1 = time.time()
    im = frame.to_image()

    # print(im.pts)
    scores, boxes = detect(im, detr, transform)
    # print(boxes)
    tracking_plot(im,scores,boxes)
    # draw_boxes(im,scores,boxes,index)
    det_fps  = ( det_fps + (1./(time.time()-t1)) ) / 2
    print(f"\rFPS:{det_fps}",end = '',flush=True)

print("Successfully written to output")
print(f"took {(time.time()-t0)/60} minutes to process {file_path} to {out_file}")
# shutil.rmtree("/tmp/images")