def __init__(self,
model,
class_path,
thickness=2,
font_path=None,
font_size=10,
conf_thres=0.5,
nms_thres=0.4):
self.model = model
self.model.eval()
self.classes = load_classes(class_path)
self.num_classes = len(self.classes)
self.thickness = thickness
self.conf_thres = conf_thres
self.nms_thres = nms_thres
if font_path is not None:
self.font = ImageFont.truetype(font_path, font_size)
else:
self.font = ImageFont.load_default()
# Prepare colors for each class
hsv_color = [(1.0 * i / self.num_classes, 1., 1.)
for i in range(self.num_classes)]
colors = [colorsys.hsv_to_rgb(*x) for x in hsv_color]
random.seed(0)
random.shuffle(colors)
random.seed(None)
self.colors = np.floor(np.asarray(colors) * 255).astype(int)
def __init__(self, model, class_path,
thickness=2,
font_path=None,
font_size=10,
thres=0.7,
nms_thres=0.4,
skip_frames=-1,
fourcc=cv2.VideoWriter_fourcc('m', 'p', '4', 'v'),
class_mask=None,
win_size=None,
overlap=0.15,
tracker=None,
action_id=None,
half=False):
self.thickness = thickness
self.skip_frames = skip_frames
self.class_mask = class_mask
self.fourcc = fourcc
self.label_drawer = LabelDrawer(load_classes(class_path),
font_path,
font_size,
thickness,
img_size=model.img_size)
self.image_detector = ImageDetector(model, class_path,
thickness=thickness,
thres=thres,
nms_thres=nms_thres,
win_size=win_size,
overlap=overlap,
half=half)
self.tracker = tracker
self.action_id = action_id
def __init__(self,
model,
class_path,
thickness=2,
thres=0.5,
nms_thres=0.4,
win_size=None,
overlap=0.15,
half=False):
self.model = model
self.model.eval()
self.device = next(self.model.parameters()).device
if half:
self.model.half()
self.classes = load_classes(class_path)
self.num_classes = len(self.classes)
self.thickness = thickness
self.thres = thres
self.nms_thres = nms_thres
self.half = half
self.win_size = win_size
self.overlap = overlap
def __init__(self, model, class_path):
self.model = model.eval()
self.classes = load_classes(class_path)
help="How many layers do you prefer to freeze?")
opt = parser.parse_args()
print(opt)
logger = Logger("logs")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.makedirs("output", exist_ok=True)
os.makedirs("checkpoints", exist_ok=True)
# Get data configuration
data_config = parse_data_config(opt.data_config)
train_path = data_config["train"]
valid_path = data_config["valid"]
class_names = load_classes(data_config["names"])
# Initiate model
model = Darknet(opt.model_def).to(device)
model.apply(weights_init_normal)
# If specified we start from checkpoint
if opt.pretrained_weights:
if opt.pretrained_weights.endswith(".pth"):
model.load_state_dict(
torch.load(opt.pretrained_weights, map_location=device))
else:
model.load_darknet_weights(opt.pretrained_weights)
# Get dataloader
dataset = ListDataset(train_path,
