Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-10-12 13:59:07 +02:00
parent 8397fa7a2a
commit 811b3b693f
2 changed files with 5 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
detect.py
View File

@ -31,7 +31,7 @@ def detect(save_txt=False, save_img=False):
classify = False classify = False
if classify: if classify:
modelc = torch_utils.load_classifier(name='resnet101', n=2) # initialize 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() modelc.to(device).eval()
# Fuse Conv2d + BatchNorm2d layers # Fuse Conv2d + BatchNorm2d layers

13
utils/utils.py
View File

@ -739,23 +739,18 @@ def apply_classifier(x, model, img, im0):
# Classes # Classes
pred_cls1 = d[:, 6].long() pred_cls1 = d[:, 6].long()
ims = [] ims = []
j = 0 for j, a in enumerate(d): # per item
for a in d: # per item
j += 1
cutout = im0[int(a[1]):int(a[3]), int(a[0]):int(a[2])] cutout = im0[int(a[1]):int(a[3]), int(a[0]):int(a[2])]
im = cv2.resize(cutout, (224, 224)) # BGR 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 = 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 = np.ascontiguousarray(im, dtype=np.float32) # uint8 to float32
im /= 255.0 # 0 - 255 to 0.0 - 1.0 im /= 255.0 # 0 - 255 to 0.0 - 1.0
ims.append(im) ims.append(im)
ims = torch.Tensor(np.concatenate(ims, 0)) # to torch pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1) # classifier prediction
pred_cls2 = model(ims).argmax(1) # classifier prediction x[i] = x[i][pred_cls1 == pred_cls2] # retain matching class detections
# x[i] = x[i][pred_cls1 == pred_cls2] # retain matching class detections
return x return x