-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
114 lines (103 loc) · 4.61 KB
/
model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
import cv2 as cv
import numpy as np
import os
import sys
import xml.dom.minidom as xmldom
import argparse
from nms import py_cpu_nms
import time as T
import re
import linecache
import matplotlib.pyplot as plt
from utils.torch_utils import select_device, time_sync
from models.common import DetectMultiBackend
from util import plot_one_box, cal_iou, xyxy_to_xywh, xywh_to_xyxy, updata_trace_list, draw_trace
from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from utils.plots import Annotator, colors, save_one_box
from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox, mixup, random_perspective
from utils.general import (DATASETS_DIR, LOGGER, NUM_THREADS, check_dataset, check_requirements, check_yaml, clean_str,
cv2, is_colab, is_kaggle, segments2boxes, xyn2xy, xywh2xyxy, xywhn2xyxy, xyxy2xywhn)
from utils.torch_utils import torch_distributed_zero_first
from pathlib import Path
import ipdb
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
class Model:
def __init__(self, method="MOG2"):
self.iou_thres = 0.3
weights = "runs/train/exp10/weights/best.pt"
data = ROOT / 'data/paowu.yaml'
imgsz = (640, 640)
self.conf_thres = 0.15
self.box_color = (255, 0, 255)
# 加载yolo模型
self.device = ''
self.device = select_device(self.device)
self.model = DetectMultiBackend(weights, device=self.device, dnn=False, data=data, fp16=False)
self.stride, self.names, self.pt = self.model.stride, self.model.names, self.model.pt
self.imgsz = check_img_size(imgsz, s=self.stride) # check image size
bs = 1
self.model.warmup(imgsz=(1 if self.pt else bs, 6, *imgsz)) # warmup
self.fp16 = self.model.fp16
statistics_TP = np.zeros((5, 17))
statistics_FP = np.zeros((5, 17))
statistics_FN = np.zeros((5, 17))
statistics_TRO = np.zeros((5, 17))
statistics_num = np.zeros((5, 17))
#filename = args.dataset + line
tp = 0.00
fp = 0.00
fn = 0.00
self.framenum = 0
if method == 'MOG2':
self.backSub = cv.createBackgroundSubtractorMOG2(300, 100, False)
elif method == 'GSOC':
self.backSub = cv.bgsegm.createBackgroundSubtractorGSOC(300, 100)
elif method == 'LSBP':
self.backSub = cv.bgsegm.createBackgroundSubtractorLSBP()
elif method == 'GMG':
self.backSub = cv.bgsegm.createBackgroundSubtractorGMG()
self.backSub.setNumFrames(5)
self.backSub.setUpdateBackgroundModel(True)
elif method == 'CNT':
self.backSub = cv.bgsegm.createBackgroundSubtractorCNT()
self.backSub.setIsParallel(True)
self.backSub.setUseHistory(True)
self.backSub.setMinPixelStability(1)
self.backSub.setMaxPixelStability(4)
elif method == 'MOG':
self.backSub = cv.bgsegm.createBackgroundSubtractorMOG()
elif method == 'KNN':
self.backSub = cv.createBackgroundSubtractorKNN(300, 100, False)
else:
print("Wrong algo")
sys.exit()
def alt(self, method):
if method == 'MOG2':
self.backSub = cv.createBackgroundSubtractorMOG2(300, 100, False)
elif method == 'GSOC':
self.backSub = cv.bgsegm.createBackgroundSubtractorGSOC(300, 100)
elif method == 'LSBP':
self.backSub = cv.bgsegm.createBackgroundSubtractorLSBP()
elif method == 'GMG':
self.backSub = cv.bgsegm.createBackgroundSubtractorGMG()
self.backSub.setNumFrames(5)
self.backSub.setUpdateBackgroundModel(True)
elif method == 'CNT':
self.backSub = cv.bgsegm.createBackgroundSubtractorCNT()
self.backSub.setIsParallel(True)
self.backSub.setUseHistory(True)
self.backSub.setMinPixelStability(1)
self.backSub.setMaxPixelStability(4)
elif method == 'MOG':
self.backSub = cv.bgsegm.createBackgroundSubtractorMOG()
elif method == 'KNN':
self.backSub = cv.createBackgroundSubtractorKNN(300, 100, False)
else:
print("Wrong algo")
sys.exit()