source, weights, view_img, save_txt, imgsz, trace = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, not opt.no_trace

#source = '/home/akhil_kk/WORKING_PROJECT/computer_vision/DATASET/Road_traffic_video.mp4'

#view_img=True

save_img = not opt.nosave and not source.endswith('.txt') # save inference images

webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(

('rtsp://', 'rtmp://', 'http://', 'https://'))

# Directories

save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run

(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir

# Initialize

set_logging()

device = select_device(opt.device)

half = device.type != 'cpu' # half precision only supported on CUDA

# Load model

model = attempt_load(weights, map_location=device) # load FP32 model

stride = int(model.stride.max()) # model stride

imgsz = check_img_size(imgsz, s=stride) # check img_size

if trace:

model = TracedModel(model, device, opt.img_size)

if half:

model.half() # to FP16

# Second-stage classifier

classify = False

if classify:

modelc = load_classifier(name='resnet101', n=2) # initialize

modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()

# Set Dataloader

if webcam:

view_img = check_imshow()

cudnn.benchmark = True # set True to speed up constant image size inference

dataset = LoadStreams(source, img_size=imgsz, stride=stride)

else:

dataset = LoadImages(source, img_size=imgsz, stride=stride)

# Get names and colors

names = model.module.names if hasattr(model, 'module') else model.names

colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]

# Run inference

if device.type != 'cpu':

model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once

old_img_w = old_img_h = imgsz

old_img_b = 1

t0 = time.time()

#Create a tracker object for tracking (The constructor need list of class label as input)

sel_classes=None # selective tracking : list of claases to be tracked, if None then all class objects are tracked

#sel_classes=['bus','truck'] # selective tracking : list of claases to be tracked, only provided class objects are tracked

tracker_m = tracker(names,sel_classes) #create tracker object for tracking

for path, img, im0s, vid_cap in dataset:

str_time=time.time()

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)

# Warmup

if device.type != 'cpu' and (old_img_b != img.shape[0] or old_img_h != img.shape[2] or old_img_w != img.shape[3]):

old_img_b = img.shape[0]

old_img_h = img.shape[2]

old_img_w = img.shape[3]

for i in range(3):

model(img, augment=opt.augment)[0]

# Inference

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)

# Apply Classifier

if classify:

pred = apply_classifier(pred, modelc, img, im0s)

# Process detections

for i, det in enumerate(pred): # detections per image

if webcam: # batch_size >= 1

p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count

else:

p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)

p = Path(p) # to Path

if len(det):

# Rescale boxes from img_size to im0 size

det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

# track the objects based on the history

#the tracker will return bbox and the label of the tracked objects and new label for new objects

new_det = tracker_m.track(im0,det)

#plotiing the result from the tracker

for xyxy,conf,cls,label in new_det:

plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1)

if view_img:

t_diff=time.time()-str_time

txt="Speed: "+str(int(1.00/t_diff))+" fps"

#print(txt)

cv2.putText(im0, txt, (10,40), cv2.FONT_HERSHEY_SIMPLEX,1, (255,0,0))

cv2.imshow(str(p), im0)