Utils reorganization (#1392)
* Utils reorganization * Add new utils files * cleanup * simplify * reduce datasets.py * remove evolve.sh * loadWebcam cleanup
This commit is contained in:
1 parent
379396e896
commit
fe341fa44d
14 files changed
+890
-988
No files matched your search
@@ -10,14 +10,15 @@ from numpy import random
|
||||
from models.experimental import attempt_load
|
||||
from utils.datasets import LoadStreams, LoadImages
|
||||
from utils.general import check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, \
|
||||
plot_one_box, strip_optimizer, set_logging, increment_path
|
||||
strip_optimizer, set_logging, increment_path
|
||||
from utils.plots import plot_one_box
|
||||
from utils.torch_utils import select_device, load_classifier, time_synchronized
|
||||
|
||||
|
||||
def detect(save_img=False):
|
||||
source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
|
||||
webcam = source.isnumeric() or source.endswith('.txt') or \
|
||||
source.lower().startswith(('rtsp://', 'rtmp://', 'http://'))
|
||||
webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
|
||||
('rtsp://', 'rtmp://', 'http://'))
|
||||
|
||||
# Directories
|
||||
save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
|
||||
@@ -38,8 +39,7 @@ def detect(save_img=False):
|
||||
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']) # load weights
|
||||
modelc.to(device).eval()
|
||||
modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()
|
||||
|
||||
# Set Dataloader
|
||||
vid_path, vid_writer = None, None
|
||||
@@ -53,7 +53,7 @@ def detect(save_img=False):
|
||||
|
||||
# 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 range(len(names))]
|
||||
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
|
||||
|
||||
# Run inference
|
||||
t0 = time.time()
|
||||
|
||||
Reference in new issue
Block a user