detect.py

import os
import sys
sys.path.insert(0, './yolov5')
import cv2
import time
import torch
import shutil
import argparse
import platform
import numpy as np
from pathlib import Path
import matplotlib.pyplot as plt
from _collections import deque
import torch.backends.cudnn as cudnn

from yolov5.models.experimental import attempt_load
from yolov5.utils.downloads import attempt_download
from yolov5.utils.datasets import LoadImages, LoadStreams
from yolov5.utils.general import check_img_size, non_max_suppression, scale_coords, check_imshow, xyxy2xywh
from yolov5.utils.torch_utils import select_device, time_sync
from yolov5.utils.plots import colors

from deep_sort_pytorch.utils.parser import get_config
from deep_sort_pytorch.deep_sort import DeepSort

pts = [deque() for _ in range(1000)]
speed_dict = {}

# Global Counters
car, motorcycle, bus, truck = 0, 0, 0, 0
car_up, car_down = 0, 0
motorcycle_up, motorcycle_down = 0, 0
bus_up, bus_down = 0, 0
truck_up, truck_down = 0, 0

counter = []
cmap = plt.get_cmap('tab20b')
colors = [cmap(i)[:3] for i in np.linspace(0, 1, 20)]

def speed(points, id, upper_distance, lower_distance, video_fps, detection_fps, vehicle):
    global car, motorcycle, bus, truck
    global car_up, car_down, motorcycle_up, motorcycle_down, bus_up, bus_down, truck_up, truck_down
    
    if len(points) < 2:
        return
        
    time_diff = points[-1][1] - points[0][1]
    if time_diff == 0:
        return
    
    if id in speed_dict.keys():
        pass
    else:
        # Determine direction based on start and end y-coordinate
        direction = 'DOWN' if points[0][0][1] < points[-1][0][1] else 'UP'
        
        if vehicle == 'car':
            car += 1
            if direction == 'UP':
                car_up += 1
            else:
                car_down += 1
        elif vehicle == 'motorcycle':
            motorcycle += 1
            if direction == 'UP':
                motorcycle_up += 1
            else:
                motorcycle_down += 1
        elif vehicle == 'bus':
            bus += 1
            if direction == 'UP':
                bus_up += 1
            else:
                bus_down += 1
        elif vehicle == 'truck':
            truck += 1
            if direction == 'UP':
                truck_up += 1
            else:
                truck_down += 1

    if points[0][0][1] < points[-1][0][1]: #  UPPER SIDE
        try:
            velocity = upper_distance / time_diff
            velocity = velocity * 3.6
            speed_dict[id] = [velocity, 'DOWN']
        except Exception as e:
            print(e)
    else: # Lower side
        try:
            velocity = lower_distance / time_diff
            velocity = velocity * 3.6
            speed_dict[id] = [velocity, 'UP']
        except Exception as e:
            print(e)

    

def detect(opt):
    out, source, yolo_weights, deep_sort_weights, show_vid, save_vid, save_txt, imgsz, evaluate, detection_track, detection_bbox = \
        opt.output, opt.source, opt.yolo_weights, opt.deep_sort_weights, opt.show_vid, opt.save_vid, \
            opt.save_txt, opt.img_size, opt.evaluate, opt.track, opt.bbox
    UPPER_DISTANCE, LOWER_DISTANCE, LINE = opt.upper, opt.lower, opt.line
    UP_DETECT_LIMIT, DOWN_DETECT_LIMIT, SHOW_DETECT_LIMIT = opt.uplimit, opt.downlimit, opt.showlimit
    WIDTH, FIRST_FRAME = 0, True
    webcam = source == '0' or source.startswith(
        'rtsp') or source.startswith('http') or source.endswith('.txt')

    # Initialize device first to set CUDA status
    device_str = str(opt.device) if hasattr(opt, 'device') and opt.device is not None else ''
    if device_str != 'cpu' and not torch.cuda.is_available():
        print("CUDA is not available. Falling back to CPU.")
        device_str = 'cpu'
    device = select_device(device_str)
    use_cuda = device.type != 'cpu'

    # initialize deepsort
    cfg = get_config()
    cfg.merge_from_file(opt.config_deepsort)
    attempt_download(deep_sort_weights, repo='mikel-brostrom/Yolov5_DeepSort_Pytorch')
    deepsort = DeepSort(cfg.DEPOSORT.REID_CKPT if hasattr(cfg, 'DEPOSORT') else cfg.DEEPSORT.REID_CKPT,
                        max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
                        max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
                        max_age=cfg.DEPOSORT.MAX_AGE if hasattr(cfg, 'DEPOSORT') else cfg.DEEPSORT.MAX_AGE, 
                        n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
                        use_cuda=use_cuda)

    if not evaluate:
        if os.path.exists(out):
            pass
            shutil.rmtree(out)  # delete output folder
        os.makedirs(out)  # make new output folder

    half = device.type != 'cpu'  # half precision only supported on CUDA
    
    # Load model
    model = attempt_load(yolo_weights, map_location=device)  # load FP32 model
    stride = int(model.stride.max())  # model stride
    imgsz = check_img_size(imgsz, s=stride)  # check img_size

    names = model.module.names if hasattr(model, 'module') else model.names  # get class names
    if half:
        model.half()  # to FP16

    # Set Dataloader
    vid_path, vid_writer = None, None
    # Check if environment supports image displays
    if show_vid:
        show_vid = check_imshow()

    if webcam:
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=imgsz, stride=stride)
    else:
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadImages(source, img_size=imgsz, stride=stride)

    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names

    # Run inference
    if device.type != 'cpu':
        model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
    t0 = time.time()

    save_path = str(Path(out))


    for frame_idx, (path, img, im0s, vid_cap, s) in enumerate(dataset):
        if FIRST_FRAME == True:
            FPS = vid_cap.get(cv2.CAP_PROP_FPS) if vid_cap is not None else 30.0
            if FPS == 0 or FPS is None:
                FPS = 30.0
            HEIGHT, WIDTH, _ = im0s.shape
            FIRST_FRAME = False

        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        
        # Inference
        
        t1 = time_sync()
        pred = model(img, augment=opt.augment)[0]

        # Apply NMS
        pred = non_max_suppression(
            pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
        
        # Process detections
        for i, det in enumerate(pred):  # detections per image
            if webcam:  # batch_size >= 1
                p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
            else:
                p, s, im0 = path, '', im0s
            
            s += '%gx%g ' % img.shape[2:]  # print string
            save_path = str(Path(out) / Path(p).name)

            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(
                    img.shape[2:], det[:, :4], im0.shape).round()

                xywhs = xyxy2xywh(det[:, 0:4])
                confs = det[:, 4]
                clss = det[:, 5]

                # pass detections to deepsort
                outputs = deepsort.update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)

                # draw boxes for visualization
                if len(outputs) > 0:
                    for j , (output, conf) in enumerate(zip(outputs, confs)):
                        
                        bboxes = output[0:4]
                        id = output[4]
                        cls = output[5]
                        c = int(cls)
                    
                        color = colors[int(id) % len(colors)]
                        color = [i*255 for i in color]

                        center = (int(((output[0]) + (output[2]))/2), int(((output[1]) + (output[3]))/2))
                        if center[1] <= UP_DETECT_LIMIT or center[1] >= DOWN_DETECT_LIMIT:
                            pass
                        else:
                            pts[id].append([center, time.time()])

                        # if id in speed_dict.keys():
                        #     if speed_dict[id][1] == 'DOWN' and center[1] > LINE and (len(pts[id]) == 0 or pts[id][0][0][1] < LINE):
                        #         pts[id] = deque()
                        #     elif speed_dict[id][1] == 'UP' and center[1] < LINE and (len(pts[id]) == 0 or pts[id][0][0][1] > LINE):
                        #         pts[id] = deque()



                        if center[1] - 7 < LINE and center[1] + 7 > LINE:
                            speed(pts[id], id, UPPER_DISTANCE, LOWER_DISTANCE, FPS, fps, names[c])

                        if id in speed_dict.keys():
                            try:
                                if (center[1] < UP_DETECT_LIMIT) and (pts[id][0][0][1] > LINE):
                                    pts[id] = deque()
                                elif (center[1] > DOWN_DETECT_LIMIT) and (pts[id][0][0][1] < LINE):
                                    pts[id] = deque()
                            except:
                                pass

                        if detection_track:
                            for j in range(1, len(pts[id])):
                                if j > 20:
                                    break
                                if pts[id][j-1][0] is None or pts[id][j][0] is None:
                                    continue
                                cv2.circle(im0, (pts[id][-(j-1)][0]), 1, color)

                        if id in speed_dict.keys():
                            try:
                                speed_string = '{:.1f} KM/H'.format(speed_dict[id][0])
                                cv2.rectangle(im0, (int(bboxes[0]), int(bboxes[1]) - 20), (int(bboxes[0] + len(speed_string)*10)  , int(bboxes[1])), color, -1)
                                cv2.putText(im0, speed_string , (int(bboxes[0]), int(bboxes[1]) - 7), cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.7, (255, 255, 255), 1, cv2.FILLED)
                            except:
                                pass



                        if detection_bbox:
                            cv2.rectangle(im0, (int(bboxes[0]), int(bboxes[1])), (int(bboxes[2]), int(bboxes[3])), color, 1)
            
            else:
                deepsort.increment_ages()


            t2 = time_sync()
            fps = 1./(t2-t1)

            # Draw premium dashboard
            cv2.rectangle(im0, (10, 10), (320, 190), (0, 0, 0), -1)  # Premium black background card
            cv2.rectangle(im0, (10, 10), (320, 190), (0, 255, 255), 2)  # Cyan border
            cv2.putText(im0, f"DASHBOARD | FPS: {int(fps)}", (20, 30), cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.9, (0, 255, 255), 1)
            cv2.line(im0, (15, 40), (315, 40), (255, 255, 255), 1)
            cv2.putText(im0, f"CAR:        {car:<3} (UP: {car_up:<2} DN: {car_down:<2})", (20, 70), cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.8, (255, 255, 255), 1)
            cv2.putText(im0, f"MOTORCYCLE: {motorcycle:<3} (UP: {motorcycle_up:<2} DN: {motorcycle_down:<2})", (20, 100), cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.8, (255, 255, 255), 1)
            cv2.putText(im0, f"TRUCK:      {truck:<3} (UP: {truck_up:<2} DN: {truck_down:<2})", (20, 130), cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.8, (255, 255, 255), 1)
            cv2.putText(im0, f"BUS:        {bus:<3} (UP: {bus_up:<2} DN: {bus_down:<2})", (20, 160), cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.8, (255, 255, 255), 1)
            
            
        if SHOW_DETECT_LIMIT == True:
            im0 = cv2.line(im0, (0, UP_DETECT_LIMIT), (WIDTH, UP_DETECT_LIMIT), (0, 0, 255), 1)    
            im0 = cv2.line(im0, (0, DOWN_DETECT_LIMIT), (WIDTH, DOWN_DETECT_LIMIT), (0, 0, 255), 1)    

        im0 = cv2.line(im0, (0, LINE), (WIDTH, LINE), (0, 0, 255), 3)    
        if show_vid:
            cv2.imshow(p, im0)
            if cv2.waitKey(1) == ord('q'):  # q to quit
                cv2.destroyAllWindows()
                raise StopIteration


        # Save results (image with detections)
        if save_vid:
            if vid_path != save_path:  # new video
                vid_path = save_path
                if isinstance(vid_writer, cv2.VideoWriter):
                    vid_writer.release()  # release previous video writer
                if vid_cap:  # video
                    fps = vid_cap.get(cv2.CAP_PROP_FPS)
                    w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                    h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                else:  # stream
                    fps, w, h = 30, im0.shape[1], im0.shape[0]
                    save_path += '.mp4'

                vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
            vid_writer.write(im0)
        
    if save_txt or save_vid:
        print('Results saved to %s' % os.getcwd() + os.sep + out)
        if platform == 'darwin':  # MacOS
            os.system('open ' + save_path)

    print('Done. (%.3fs)' % (time.time() - t0))


if __name__ == '__main__':
        parser = argparse.ArgumentParser()
        parser.add_argument('--yolo_weights', nargs='+', type=str, default='yolov5/yolov5s.pt', help='model.pt path(s)')
        parser.add_argument('--deep_sort_weights', type=str, default='deep_sort_pytorch/deep_sort/deep/checkpoint/ckpt.t7', help='ckpt.t7 path')
        parser.add_argument('--source', type=str, default='yolov5\Relaxing highway traffic.mp4', help='source')
        parser.add_argument('--output', type=str, default='inference/output', help='output folder')  # output folder
        parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
        parser.add_argument('--conf-thres', type=float, default=0.1, help='object confidence threshold')
        parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
        parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
        parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
        parser.add_argument('--show-vid', default=True, help='display tracking video results')
        parser.add_argument('--save-vid', action='store_true', help='save video tracking results')
        parser.add_argument('--save-txt', action='store_true', help='save MOT compliant results to *.txt')
        parser.add_argument('--classes', nargs='+', type=int, default = [2, 3, 5, 7], help='filter by class: --class 0, or --class 16 17')
        parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
        parser.add_argument('--augment', action='store_true', help='augmented inference')
        parser.add_argument('--evaluate', action='store_true', help='augmented inference')
        parser.add_argument("--config_deepsort", type=str, default="deep_sort_pytorch/configs/deep_sort.yaml")
        parser.add_argument('--upper', type=int, default=30, help='Upper detection (m)')
        parser.add_argument('--lower', type=int, default=20, help='Lower detection (m)')
        parser.add_argument('--line', type=int, default=450, help="Line placement")
        parser.add_argument('--uplimit', type=int, default=350, help="Line placement")
        parser.add_argument('--downlimit', type=int, default=650, help="Line placement")
        parser.add_argument('--showlimit', default=True, help="Line placement")
        parser.add_argument('--track', default=True, help="show the detection track")
        parser.add_argument('--bbox', action='store_true', help="show the detection box")
        args = parser.parse_args()
        args.img_size = check_img_size(args.img_size)

        with torch.no_grad():
            detect(args)