This commit is contained in:
Glenn Jocher 2019-02-12 16:58:07 +01:00
parent e4e64a9ff6
commit cc5e9a5a85
5 changed files with 5 additions and 12 deletions

1
detect.py Executable file → Normal file
View File

@ -1,5 +1,6 @@
import argparse import argparse
import time import time
from sys import platform
from models import * from models import *
from utils.datasets import * from utils.datasets import *

View File

@ -243,7 +243,7 @@ class Darknet(nn.Module):
def __init__(self, cfg_path, img_size=416): def __init__(self, cfg_path, img_size=416):
super(Darknet, self).__init__() super(Darknet, self).__init__()
self.module_defs = parse_model_config(cfg_path) self.module_defs = parse_model_cfg(cfg_path)
self.module_defs[0]['cfg'] = cfg_path self.module_defs[0]['cfg'] = cfg_path
self.module_defs[0]['height'] = img_size self.module_defs[0]['height'] = img_size
self.hyperparams, self.module_list = create_modules(self.module_defs) self.hyperparams, self.module_list = create_modules(self.module_defs)

View File

@ -2,7 +2,6 @@ import glob
import math import math
import os import os
import random import random
from sys import platform
import cv2 import cv2
import numpy as np import numpy as np

View File

@ -1,4 +1,4 @@
def parse_model_config(path): def parse_model_cfg(path):
"""Parses the yolo-v3 layer configuration file and returns module definitions""" """Parses the yolo-v3 layer configuration file and returns module definitions"""
file = open(path, 'r') file = open(path, 'r')
lines = file.read().split('\n') lines = file.read().split('\n')

View File

@ -1,16 +1,9 @@
import torch import torch
def check_cuda():
return torch.cuda.is_available()
CUDA_AVAILABLE = check_cuda()
def init_seeds(seed=0): def init_seeds(seed=0):
torch.manual_seed(seed) torch.manual_seed(seed)
if CUDA_AVAILABLE: if torch.cuda.is_available():
torch.cuda.manual_seed(seed) torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed) torch.cuda.manual_seed_all(seed)
# torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available # torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available
@ -20,6 +13,6 @@ def select_device(force_cpu=False):
if force_cpu: if force_cpu:
device = torch.device('cpu') device = torch.device('cpu')
else: else:
device = torch.device('cuda:0' if CUDA_AVAILABLE else 'cpu') device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Using ' + str(device) + '\n') print('Using ' + str(device) + '\n')
return device return device