Commit d47c9894 authored by flfuchs's avatar flfuchs
Browse files

Improve performance

parent a8c5bd73
# encoding: utf8
import copy
import sys
......@@ -69,19 +70,25 @@ class App(QMainWindow, UiMainWindow):
self.yolo_model.load_state_dict(torch.load(weights_path))
# counter Thread
self.counterThread = CounterThread(self.yolo_model, self.yolo_class_names, self.device)
self.counterThread = CounterThread(
self.yolo_model, self.yolo_class_names, self.device
)
self.counterThread.sin_counterResult.connect(self.show_image_label)
self.counterThread.sin_done.connect(self.done)
self.counterThread.sin_counter_results.connect(self.update_counter_results)
def open_video(self):
openfile_name = QFileDialog.getOpenFileName(self, "Open video", "", "Video files(*.avi , *.mp4)")
openfile_name = QFileDialog.getOpenFileName(
self, "Open video", "", "Video files(*.avi , *.mp4)"
)
self._video_name = openfile_name[0]
vid = cv2.VideoCapture(openfile_name[0])
while vid.isOpened():
vid.set(cv2.CAP_PROP_POS_FRAMES, int(vid.get(cv2.CAP_PROP_FRAME_COUNT) * 0.5))
vid.set(
cv2.CAP_PROP_POS_FRAMES, int(vid.get(cv2.CAP_PROP_FRAME_COUNT) * 0.5)
)
ret, frame = vid.read()
if ret:
self.exampleImage = frame
......@@ -108,11 +115,17 @@ class App(QMainWindow, UiMainWindow):
if self.get_points_flag:
x = event.x()
y = event.y()
self.countArea.append([int(x * self.imgScale[1]), int(y * self.imgScale[0])])
self.countArea.append(
[int(x * self.imgScale[1]), int(y * self.imgScale[0])]
)
exampleImageWithArea = copy.deepcopy(self.exampleImage)
for point in self.countArea:
exampleImageWithArea[point[1] - 10 : point[1] + 10, point[0] - 10 : point[0] + 10] = (0, 255, 255)
cv2.fillConvexPoly(exampleImageWithArea, np.array(self.countArea), (0, 0, 255))
exampleImageWithArea[
point[1] - 10 : point[1] + 10, point[0] - 10 : point[0] + 10
] = (0, 255, 255)
cv2.fillConvexPoly(
exampleImageWithArea, np.array(self.countArea), (0, 0, 255)
)
self.show_image_label(exampleImageWithArea)
print(self.countArea)
......
# coding:utf-8
import os
from collections import defaultdict
from math import sqrt
import cv2
......@@ -9,7 +10,13 @@ from tqdm import tqdm
import predict
from config import color_dict, names
from magic_constants import SECONDS_PER_FRAME, FRAME_COUNT_STEP, UPDATE_DISPLAY_AT_EVERY_FRAME, DO_WRITE_OUTPUT_TO_MP4
from magic_constants import (
SECONDS_PER_FRAME,
FRAME_COUNT_STEP,
UPDATE_DISPLAY_AT_EVERY_FRAME,
DO_WRITE_OUTPUT_TO_MP4,
MINIMAL_NUMBER_OF_OBSERVATIONS_PER_DETECTION,
)
from utils.sort import KalmanBoxTracker, Sort
......@@ -34,14 +41,13 @@ class CounterThread(QThread):
self._color_dict = color_dict.copy()
# create instance of SORT
self.mot_tracker = Sort(max_age=32, min_hits=4)
self.mot_tracker = Sort(max_age=256, min_hits=4)
self.count_area = None
self._running_flag = 0
self._pause_flag = 0
self._video_name = []
self.last_max_id = 0
self.history = {} # save history
# history = {id:{"no_update_count": int, "his": list}}
self._tracker_history = {} # save history
self.sin_runningFlag.connect(self.update_flag)
self.sin_videoList.connect(self.update_video_name)
......@@ -53,7 +59,7 @@ class CounterThread(QThread):
os.makedirs(self.save_dir)
def run(self):
# with cProfile.Profile() as profiler:
self.last_max_id = 0
cap = cv2.VideoCapture(self._video_name)
if DO_WRITE_OUTPUT_TO_MP4:
......@@ -95,7 +101,7 @@ class CounterThread(QThread):
print("Update countArea!")
self.count_area = np.array(area)
def counter(self, current_frame, frame_count):
def counter(self, current_frame: np.array, frame_count: int):
# painting area
AreaBound = [
min(self.count_area[:, 0]),
......@@ -103,16 +109,22 @@ class CounterThread(QThread):
max(self.count_area[:, 0]),
max(self.count_area[:, 1]),
]
painting = np.zeros((AreaBound[3] - AreaBound[1], AreaBound[2] - AreaBound[0]), dtype=np.uint8)
painting = np.zeros(
(AreaBound[3] - AreaBound[1], AreaBound[2] - AreaBound[0]), dtype=np.uint8
)
CountArea_mini = self.count_area - AreaBound[0:2]
cv2.fillConvexPoly(painting, CountArea_mini, (1,))
objects = predict.yolo_prediction(self.model, self.device, current_frame, self.class_names)
objects = predict.yolo_prediction(
self.model, self.device, current_frame, self.class_names
)
objects = filter(lambda x: x[0] in self.permission, objects)
objects = filter(lambda x: x[1] > 0.5, objects)
objects = list(
filter(
lambda x: point_in_count_area(painting, AreaBound, [int(x[2][0]), int(x[2][1] + x[2][3] / 2)]),
lambda x: point_in_count_area(
painting, AreaBound, [int(x[2][0]), int(x[2][1] + x[2][3] / 2)]
),
objects,
)
)
......@@ -146,25 +158,25 @@ class CounterThread(QThread):
for bb in track_bbs_ids: # add all bbox to history
id = int(bb[-1])
object_found = get_obj(bb, objects)
if id not in self.history.keys(): # add new id
self.history[id] = {}
self.history[id]["no_update_count"] = 0
self.history[id]["his"] = [object_found[0]]
self.history[id]["first_frame"] = frame_count
self.history[id]["last_frame"] = frame_count
self.history[id]["last_position_x_y"] = (
if id not in self._tracker_history.keys(): # add new id
self._tracker_history[id] = {}
self._tracker_history[id]["no_update_count"] = 0
self._tracker_history[id]["his"] = [object_found[0]]
self._tracker_history[id]["first_frame"] = frame_count
self._tracker_history[id]["last_frame"] = frame_count
self._tracker_history[id]["last_position_x_y"] = (
object_found[2][0] + object_found[2][2] * 0.5,
object_found[2][1] + object_found[2][3] * 0.5,
)
self.history[id]["first_position_x_y"] = (
self._tracker_history[id]["first_position_x_y"] = (
object_found[2][0] + object_found[2][2] * 0.5,
object_found[2][1] + object_found[2][3] * 0.5,
)
else:
self.history[id]["no_update_count"] = 0
self.history[id]["his"].append(object_found[0])
self.history[id]["last_frame"] = frame_count
self.history[id]["last_position_x_y"] = (
self._tracker_history[id]["no_update_count"] = 0
self._tracker_history[id]["his"].append(object_found[0])
self._tracker_history[id]["last_frame"] = frame_count
self._tracker_history[id]["last_position_x_y"] = (
object_found[2][0] + object_found[2][2] * 0.5,
object_found[2][1] + object_found[2][3] * 0.5,
)
......@@ -175,7 +187,7 @@ class CounterThread(QThread):
id = bb[-1]
x1, y1, x2, y2 = bb[:4]
his = self.history[id]["his"]
his = self._tracker_history[id]["his"]
result = {i: his.count(i) for i in set(his)}
res = sorted(result.items(), key=lambda d: d[1], reverse=True)
objectName = res[0][0]
......@@ -194,37 +206,53 @@ class CounterThread(QThread):
counter_results = []
removed_id_list = []
for id in self.history.keys(): # extract id after tracking
if self.history[id]["no_update_count"] > 20:
his = self.history[id]["his"]
if len(his) > 4:
result = {i: his.count(i) for i in set(his)}
res = sorted(result.items(), key=lambda d: d[1], reverse=True)
first_time_stamp = self.history[id]["first_frame"] * SECONDS_PER_FRAME
last_time_stamp = self.history[id]["last_frame"] * SECONDS_PER_FRAME
for id in self._tracker_history.keys(): # extract id after tracking
if self._tracker_history[id]["no_update_count"] > 80:
his = self._tracker_history[id]["his"]
if len(his) >= MINIMAL_NUMBER_OF_OBSERVATIONS_PER_DETECTION:
results = defaultdict(int)
for detected_class in his:
results[detected_class] += 1
most_frequent_class = sorted(
results.items(),
key=lambda key_value_pair: key_value_pair[1],
reverse=True,
)[0][0]
first_time_stamp = (
self._tracker_history[id]["first_frame"] * SECONDS_PER_FRAME
)
last_time_stamp = (
self._tracker_history[id]["last_frame"] * SECONDS_PER_FRAME
)
distance_in_pixels = sqrt(
(self.history[id]["last_position_x_y"][0] - self.history[id]["first_position_x_y"][0]) ** 2
+ (self.history[id]["last_position_x_y"][1] - self.history[id]["first_position_x_y"][1]) ** 2
(
self._tracker_history[id]["last_position_x_y"][0]
- self._tracker_history[id]["first_position_x_y"][0]
)
** 2
+ (
self._tracker_history[id]["last_position_x_y"][1]
- self._tracker_history[id]["first_position_x_y"][1]
)
** 2
)
density = len([self.history[i] for i in self.history.keys() if len(self.history[i]["his"]) > 6])
objectName = res[0][0]
counter_results.append(
[
id,
objectName,
most_frequent_class,
first_time_stamp,
last_time_stamp,
density,
distance_in_pixels,
]
)
# del id
removed_id_list.append(id)
else:
self.history[id]["no_update_count"] += 1
self._tracker_history[id]["no_update_count"] += 1
for id in removed_id_list:
self.history.pop(id)
for tracked_id in removed_id_list:
self._tracker_history.pop(tracked_id)
if len(counter_results) > 0:
self.sin_counter_results.emit(counter_results)
......@@ -277,7 +305,11 @@ def cal_iou(box1, box2):
x2 = min(box1[2], box2[2])
y2 = min(box1[3], box2[3])
i = max(0, (x2 - x1)) * max(0, (y2 - y1))
u = (box1[2] - box1[0]) * (box1[3] - box1[1]) + (box2[2] - box2[0]) * (box2[3] - box2[1]) - i
u = (
(box1[2] - box1[0]) * (box1[3] - box1[1])
+ (box2[2] - box2[0]) * (box2[3] - box2[1])
- i
)
return float(i) / float(u)
......
......@@ -2,3 +2,4 @@ SECONDS_PER_FRAME = 1 / 25
FRAME_COUNT_STEP = 1
UPDATE_DISPLAY_AT_EVERY_FRAME = True
DO_WRITE_OUTPUT_TO_MP4 = False
MINIMAL_NUMBER_OF_OBSERVATIONS_PER_DETECTION = 10
from __future__ import division
import numpy as np
import torch
import torch.nn as nn
......@@ -47,7 +45,9 @@ def create_modules(module_defs):
kernel_size = int(module_def["size"])
stride = int(module_def["stride"])
if kernel_size == 2 and stride == 1:
modules.add_module(f"_debug_padding_{module_i}", nn.ZeroPad2d((0, 1, 0, 1)))
modules.add_module(
f"_debug_padding_{module_i}", nn.ZeroPad2d((0, 1, 0, 1))
)
maxpool = nn.MaxPool2d(
kernel_size=kernel_size,
stride=stride,
......@@ -130,8 +130,12 @@ class YOLOLayer(nn.Module):
self.stride = self.img_dim / self.grid_size
# Calculate offsets for each grid
self.grid_x = torch.arange(g).repeat(g, 1).view([1, 1, g, g]).type(FloatTensor)
self.grid_y = torch.arange(g).repeat(g, 1).t().view([1, 1, g, g]).type(FloatTensor)
self.scaled_anchors = FloatTensor([(a_w / self.stride, a_h / self.stride) for a_w, a_h in self.anchors])
self.grid_y = (
torch.arange(g).repeat(g, 1).t().view([1, 1, g, g]).type(FloatTensor)
)
self.scaled_anchors = FloatTensor(
[(a_w / self.stride, a_h / self.stride) for a_w, a_h in self.anchors]
)
self.anchor_w = self.scaled_anchors[:, 0:1].view((1, self.num_anchors, 1, 1))
self.anchor_h = self.scaled_anchors[:, 1:2].view((1, self.num_anchors, 1, 1))
......@@ -187,7 +191,18 @@ class YOLOLayer(nn.Module):
if targets is None:
return output, 0
else:
(iou_scores, class_mask, obj_mask, noobj_mask, tx, ty, tw, th, tcls, tconf,) = build_targets(
(
iou_scores,
class_mask,
obj_mask,
noobj_mask,
tx,
ty,
tw,
th,
tcls,
tconf,
) = build_targets(
pred_boxes=pred_boxes,
pred_cls=pred_cls,
target=targets,
......@@ -202,7 +217,9 @@ class YOLOLayer(nn.Module):
loss_h = self.mse_loss(h[obj_mask], th[obj_mask])
loss_conf_obj = self.bce_loss(pred_conf[obj_mask], tconf[obj_mask])
loss_conf_noobj = self.bce_loss(pred_conf[noobj_mask], tconf[noobj_mask])
loss_conf = self.obj_scale * loss_conf_obj + self.noobj_scale * loss_conf_noobj
loss_conf = (
self.obj_scale * loss_conf_obj + self.noobj_scale * loss_conf_noobj
)
loss_cls = self.bce_loss(pred_cls[obj_mask], tcls[obj_mask])
total_loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls
......@@ -245,7 +262,9 @@ class Darknet(nn.Module):
super(Darknet, self).__init__()
self.module_defs = parse_model_config(config_path)
self.hyperparams, self.module_list = create_modules(self.module_defs)
self.yolo_layers = [layer[0] for layer in self.module_list if hasattr(layer[0], "metrics")]
self.yolo_layers = [
layer[0] for layer in self.module_list if hasattr(layer[0], "metrics")
]
self.img_size = img_size
self.seen = 0
self.header_info = np.array([0, 0, 0, self.seen, 0], dtype=np.int32)
......@@ -254,12 +273,17 @@ class Darknet(nn.Module):
img_dim = x.shape[2]
loss = 0
layer_outputs, yolo_outputs = [], []
for i, (module_def, module) in enumerate(zip(self.module_defs, self.module_list)):
for i, (module_def, module) in enumerate(
zip(self.module_defs, self.module_list)
):
if module_def["type"] in ["convolutional", "upsample", "maxpool"]:
x = module(x)
elif module_def["type"] == "route":
x = torch.cat(
[layer_outputs[int(layer_i)] for layer_i in module_def["layers"].split(",")],
[
layer_outputs[int(layer_i)]
for layer_i in module_def["layers"].split(",")
],
1,
)
elif module_def["type"] == "shortcut":
......@@ -278,7 +302,9 @@ class Darknet(nn.Module):
# Open the weights file
with open(weights_path, "rb") as f:
header = np.fromfile(f, dtype=np.int32, count=5) # First five are header values
header = np.fromfile(
f, dtype=np.int32, count=5
) # First five are header values
self.header_info = header # Needed to write header when saving weights
self.seen = header[3] # number of images seen during training
weights = np.fromfile(f, dtype=np.float32) # The rest are weights
......@@ -289,7 +315,9 @@ class Darknet(nn.Module):
cutoff = 75
ptr = 0
for i, (module_def, module) in enumerate(zip(self.module_defs, self.module_list)):
for i, (module_def, module) in enumerate(
zip(self.module_defs, self.module_list)
):
if i == cutoff:
break
if module_def["type"] == "convolutional":
......@@ -299,30 +327,42 @@ class Darknet(nn.Module):
bn_layer = module[1]
num_b = bn_layer.bias.numel() # Number of biases
# Bias
bn_b = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(bn_layer.bias)
bn_b = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(
bn_layer.bias
)
bn_layer.bias.data.copy_(bn_b)
ptr += num_b
# Weight
bn_w = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(bn_layer.weight)
bn_w = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(
bn_layer.weight
)
bn_layer.weight.data.copy_(bn_w)
ptr += num_b
# Running Mean
bn_rm = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(bn_layer.running_mean)
bn_rm = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(
bn_layer.running_mean
)
bn_layer.running_mean.data.copy_(bn_rm)
ptr += num_b
# Running Var
bn_rv = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(bn_layer.running_var)
bn_rv = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(
bn_layer.running_var
)
bn_layer.running_var.data.copy_(bn_rv)
ptr += num_b
else:
# Load conv. bias
num_b = conv_layer.bias.numel()
conv_b = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(conv_layer.bias)
conv_b = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(
conv_layer.bias
)
conv_layer.bias.data.copy_(conv_b)
ptr += num_b
# Load conv. weights
num_w = conv_layer.weight.numel()
conv_w = torch.from_numpy(weights[ptr : ptr + num_w]).view_as(conv_layer.weight)
conv_w = torch.from_numpy(weights[ptr : ptr + num_w]).view_as(
conv_layer.weight
)
conv_layer.weight.data.copy_(conv_w)
ptr += num_w
......@@ -336,7 +376,9 @@ class Darknet(nn.Module):
self.header_info.tofile(fp)
# Iterate through layers
for i, (module_def, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])):
for i, (module_def, module) in enumerate(
zip(self.module_defs[:cutoff], self.module_list[:cutoff])
):
if module_def["type"] == "convolutional":
conv_layer = module[0]
# If batch norm, load bn first
......
from __future__ import division
import cv2
from PIL import Image
from torch.nn.functional import interpolate
......@@ -37,6 +35,10 @@ def yolo_prediction(model, device, image, class_names):
model.eval()
with torch.no_grad():
outputs = model(_resize(images, 416).unsqueeze(0).to(device))
if outputs[0] is None:
return []
outputs = non_max_suppression(outputs, conf_thres=0.5, nms_thres=0.40)
if outputs[0] is None:
return []
return _extract_class_names(class_names, img_scale, outputs[0].cpu().data)
......@@ -22,12 +22,10 @@ import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np
from filterpy.kalman import KalmanFilter
from numba import jit
from scipy.optimize import linear_sum_assignment
from skimage import io
@jit
def iou(bb_test, bb_gt):
"""
Computes IUO between two bboxes in the form [x1,y1,x2,y2]
......@@ -45,7 +43,6 @@ def iou(bb_test, bb_gt):
return o
@jit
def convert_bbox_to_z(bbox):
"""
Takes a bounding box in the form [x1,y1,x2,y2] and returns z in the form
......@@ -199,7 +196,7 @@ def associate_detections_to_trackers(detections, candidate_trackers, iou_thresho
return matches, np.array(unmatched_detections), np.array(unmatched_trackers)
class Sort(object):
class Sort:
def __init__(self, max_age=10, min_hits=3):
"""
Sets key parameters for SORT
......
from __future__ import division
import numpy as np
import torch
import tqdm
def to_cpu(tensor):
return tensor.detach().cpu()
return tensor #.detach().cpu()
def load_classes(path):
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment