From 811b3b693f4827b59b4afde3cdf6d57d8b48887b Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 12 Oct 2019 13:59:07 +0200 Subject: [PATCH] updates --- detect.py | 2 +- utils/utils.py | 13 ++++--------- 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/detect.py b/detect.py index 05933a55..a79b8021 100644 --- a/detect.py +++ b/detect.py @@ -31,7 +31,7 @@ def detect(save_txt=False, save_img=False): classify = False if classify: modelc = torch_utils.load_classifier(name='resnet101', n=2) # initialize - modelc.load_state_dict(torch.load('resnet101.pt', map_location=device)['model']) # load weights + modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights modelc.to(device).eval() # Fuse Conv2d + BatchNorm2d layers diff --git a/utils/utils.py b/utils/utils.py index d7182103..a4c331a2 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -739,23 +739,18 @@ def apply_classifier(x, model, img, im0): # Classes pred_cls1 = d[:, 6].long() ims = [] - j = 0 - for a in d: # per item - j += 1 + for j, a in enumerate(d): # per item cutout = im0[int(a[1]):int(a[3]), int(a[0]):int(a[2])] im = cv2.resize(cutout, (224, 224)) # BGR - cv2.imwrite('test%i.jpg' % j, cutout) + # cv2.imwrite('test%i.jpg' % j, cutout) im = im[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 - im = np.expand_dims(im, axis=0) # add batch dim im = np.ascontiguousarray(im, dtype=np.float32) # uint8 to float32 im /= 255.0 # 0 - 255 to 0.0 - 1.0 ims.append(im) - ims = torch.Tensor(np.concatenate(ims, 0)) # to torch - pred_cls2 = model(ims).argmax(1) # classifier prediction - - # x[i] = x[i][pred_cls1 == pred_cls2] # retain matching class detections + pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1) # classifier prediction + x[i] = x[i][pred_cls1 == pred_cls2] # retain matching class detections return x