def predict_fn(input_object, model):
half = device.type != 'cpu'
imgsz = opt.img_size
img0 = input_object
stride = int(model.stride.max()) # model stride
img = letterbox(img0, imgsz, stride=stride)[0]
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
# Run inference
if device.type != 'cpu':
# run once
model(
torch.zeros(1, 3, imgsz,
imgsz).to(device).type_as(next(model.parameters())))
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
# Process detections
pred_labels = []
for i, det in enumerate(pred): # detections per image
gn = torch.tensor(img0.shape)[[1, 0, 1, 0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
img0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
# Write results
for *xyxy, conf, cls in reversed(det):
# normalized xywh
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist()
line = {}
line['label'] = int(cls.item())
line['x'] = xywh[0]
line['y'] = xywh[1]
line['width'] = xywh[2]
line['height'] = xywh[3]
line['conf'] = conf.item()
pred_labels.append(line)
return pred_labels
def detect(opt):
source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
save_img = not opt.nosave and not source.endswith(
'.txt') # save inference images
webcam = source.isnumeric() or source.endswith(
'.txt') or source.lower().startswith(
('rtsp://', 'rtmp://', 'http://', 'https://'))
# Directories
save_dir = increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
stride = int(model.stride.max()) # model stride
imgsz = check_img_size(imgsz, s=stride) # check img_size
names = model.module.names if hasattr(
model, 'module') else model.names # get class names
if half:
model.half() # to FP16
# Second-stage classifier
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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride)
else:
dataset = LoadImages(source, img_size=imgsz, stride=stride)
# Run inference
if device.type != 'cpu':
model(
torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(
next(model.parameters()))) # run once
t0 = time.time()
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], f'{i}: ', im0s[i].copy(
), dataset.count
else:
p, s, im0, frame = path, '', im0s.copy(), getattr(
dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / 'labels' / p.stem) + (
'' if dataset.mode == 'image' else f'_{frame}') # img.txt
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
imc = im0.copy() if opt.save_crop else im0 # for opt.save_crop
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (
cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or opt.save_crop or view_img: # Add bbox to image
c = int(cls) # integer class
label = None if opt.hide_labels else (
names[c]
if opt.hide_conf else f'{names[c]} {conf:.2f}')
plot_one_box(xyxy,
im0,
label=label,
color=colors(c, True),
line_thickness=opt.line_thickness)
if opt.save_crop:
save_one_box(xyxy,
imc,
file=save_dir / 'crops' / names[c] /
f'{p.stem}.jpg',
BGR=True)
# Print time (inference + NMS)
print(f'{s}Done. ({t2 - t1:.3f}s)')
# Stream results
if view_img:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections)
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video' or 'stream'
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
save_path += '.mp4'
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps,
(w, h))
vid_writer.write(im0)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
print(f'Done. ({time.time() - t0:.3f}s)')
def detect(save_img=False):
data = []
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://'))
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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 names]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = Path(path[i]), '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = Path(path), '', im0s
save_path = str(save_dir / p.name)
# print(save_path)
txt_path = str(save_dir / 'labels' / p.stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
# img_name = p.stem + '_%g' % dataset.frame
# print(img_name)
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (
cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
score = '%.2f' % conf
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
# columns = ['scores','video','frame','class','x1','y1','x2','y2']
data.append(
(score, p.stem, dataset.frame, names[int(cls)],
int(xyxy[0]), int(xyxy[1]), int(xyxy[2]),
int(xyxy[3])))
# Print time (inference + NMS)
# print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow(str(p), im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
fps = 4
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps,
(w, h))
vid_writer.write(im0)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
columns = ['scores', 'video', 'frame', 'class', 'x1', 'y1', 'x2', 'y2']
df = pd.DataFrame(data=data, columns=columns)
df.to_csv(f'{opt.csv_name}', index=False)
def test(
data,
weights=None,
batch_size=32,
imgsz=640,
conf_thres=0.001,
iou_thres=0.6, # for NMS
save_json=False,
single_cls=False,
augment=False,
verbose=False,
model=None,
dataloader=None,
save_dir=Path(''), # for saving images
save_txt=False, # for auto-labelling
save_conf=False,
plots=True,
log_imgs=0): # number of logged images
# Initialize/load model and set device
training = model is not None
if training: # called by train.py
device = next(model.parameters()).device # get model device
else: # called directly
set_logging()
device = select_device(opt.device, batch_size=batch_size)
save_txt = opt.save_txt # save *.txt labels
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
# model = nn.DataParallel(model)
# Half
half = device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half()
# Configure
model.eval()
is_coco = data.endswith('coco.yaml') # is COCO dataset
with open(data) as f:
data = yaml.load(f, Loader=yaml.FullLoader) # model dict
check_dataset(data) # check
nc = 1 if single_cls else int(data['nc']) # number of classes
iouv = torch.linspace(0.5, 0.95,
10).to(device) # iou vector for [email protected]:0.95
niou = iouv.numel()
# Logging
log_imgs, wandb = min(log_imgs, 100), None # ceil
try:
import wandb # Weights & Biases
except ImportError:
log_imgs = 0
# Dataloader
if not training:
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
path = data['test'] if opt.task == 'test' else data[
'val'] # path to val/test images
dataloader = create_dataloader(path,
imgsz,
batch_size,
model.stride.max(),
opt,
pad=0.5,
rect=True)[0]
seen = 0
confusion_matrix = ConfusionMatrix(nc=nc)
names = {
k: v
for k, v in enumerate(
model.names if hasattr(model, 'names') else model.module.names)
}
coco91class = coco80_to_coco91_class()
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R',
'[email protected]', '[email protected]:.95')
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
for batch_i, (img, targets, paths,
shapes) in enumerate(tqdm(dataloader, desc=s)):
img = img.to(device, non_blocking=True)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
nb, _, height, width = img.shape # batch size, channels, height, width
with torch.no_grad():
# Run model
t = time_synchronized()
inf_out, train_out = model(
img, augment=augment) # inference and training outputs
t0 += time_synchronized() - t
# Compute loss
if training:
loss += compute_loss([x.float() for x in train_out], targets,
model)[1][:3] # box, obj, cls
# Run NMS
targets[:, 2:] *= torch.Tensor([width, height, width,
height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:]
for i in range(nb)] if save_txt else [] # for autolabelling
t = time_synchronized()
output = non_max_suppression(inf_out,
conf_thres=conf_thres,
iou_thres=iou_thres,
labels=lb)
t1 += time_synchronized() - t
# Statistics per image
for si, pred in enumerate(output):
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class
path = Path(paths[si])
seen += 1
if len(pred) == 0:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool),
torch.Tensor(), torch.Tensor(), tcls))
continue
# Predictions
predn = pred.clone()
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0],
shapes[si][1]) # native-space pred
# Append to text file
if save_txt:
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0
]] # normalization gain whwh
for *xyxy, conf, cls in predn.tolist():
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh,
conf) if save_conf else (cls,
*xywh) # label format
with open(save_dir / 'labels' / (path.stem + '.txt'),
'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
# W&B logging
if plots and len(wandb_images) < log_imgs:
box_data = [{
"position": {
"minX": xyxy[0],
"minY": xyxy[1],
"maxX": xyxy[2],
"maxY": xyxy[3]
},
"class_id": int(cls),
"box_caption": "%s %.3f" % (names[cls], conf),
"scores": {
"class_score": conf
},
"domain": "pixel"
} for *xyxy, conf, cls in pred.tolist()]
boxes = {
"predictions": {
"box_data": box_data,
"class_labels": names
}
} # inference-space
wandb_images.append(
wandb.Image(img[si], boxes=boxes, caption=path.name))
# Append to pycocotools JSON dictionary
if save_json:
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
image_id = int(
path.stem) if path.stem.isnumeric() else path.stem
box = xyxy2xywh(predn[:, :4]) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for p, b in zip(pred.tolist(), box.tolist()):
jdict.append({
'image_id':
image_id,
'category_id':
coco91class[int(p[5])] if is_coco else int(p[5]),
'bbox': [round(x, 3) for x in b],
'score':
round(p[4], 5)
})
# Assign all predictions as incorrect
correct = torch.zeros(pred.shape[0],
niou,
dtype=torch.bool,
device=device)
if nl:
detected = [] # target indices
tcls_tensor = labels[:, 0]
# target boxes
tbox = xywh2xyxy(labels[:, 1:5])
scale_coords(img[si].shape[1:], tbox, shapes[si][0],
shapes[si][1]) # native-space labels
if plots:
confusion_matrix.process_batch(
pred, torch.cat((labels[:, 0:1], tbox), 1))
# Per target class
for cls in torch.unique(tcls_tensor):
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(
-1) # prediction indices
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(
-1) # target indices
# Search for detections
if pi.shape[0]:
# Prediction to target ious
ious, i = box_iou(predn[pi, :4], tbox[ti]).max(
1) # best ious, indices
# Append detections
detected_set = set()
for j in (ious > iouv[0]).nonzero(as_tuple=False):
d = ti[i[j]] # detected target
if d.item() not in detected_set:
detected_set.add(d.item())
detected.append(d)
correct[
pi[j]] = ious[j] > iouv # iou_thres is 1xn
if len(
detected
) == nl: # all targets already located in image
break
# Append statistics (correct, conf, pcls, tcls)
stats.append(
(correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
# Plot images
if plots and batch_i < 3:
f = save_dir / f'test_batch{batch_i}_labels.jpg' # labels
Thread(target=plot_images,
args=(img, targets, paths, f, names),
daemon=True).start()
f = save_dir / f'test_batch{batch_i}_pred.jpg' # predictions
Thread(target=plot_images,
args=(img, output_to_target(output), paths, f, names),
daemon=True).start()
# Compute statistics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
if len(stats) and stats[0].any():
p, r, ap, f1, ap_class = ap_per_class(*stats,
plot=plots,
save_dir=save_dir,
names=names)
p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(
1) # [P, R, [email protected], [email protected]:0.95]
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
nt = np.bincount(stats[3].astype(np.int64),
minlength=nc) # number of targets per class
else:
nt = torch.zeros(1)
# Print results
pf = '%20s' + '%12.3g' * 6 # print format
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
# Print results per class
if verbose and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
# Print speeds
t = tuple(x / seen * 1E3
for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
if not training:
print(
'Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g'
% t)
# Plots
if plots:
confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
if wandb and wandb.run:
wandb.log({"Images": wandb_images})
wandb.log({
"Validation": [
wandb.Image(str(f), caption=f.name)
for f in sorted(save_dir.glob('test*.jpg'))
]
})
# Save JSON
if save_json and len(jdict):
w = Path(weights[0] if isinstance(weights, list) else weights
).stem if weights is not None else '' # weights
anno_json = '../coco/annotations/instances_val2017.json' # annotations json
pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)
with open(pred_json, 'w') as f:
json.dump(jdict, f)
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
anno = COCO(anno_json) # init annotations api
pred = anno.loadRes(pred_json) # init predictions api
eval = COCOeval(anno, pred, 'bbox')
if is_coco:
eval.params.imgIds = [
int(Path(x).stem) for x in dataloader.dataset.img_files
] # image IDs to evaluate
eval.evaluate()
eval.accumulate()
eval.summarize()
map, map50 = eval.stats[:
2] # update results ([email protected]:0.95, [email protected])
except Exception as e:
print(f'pycocotools unable to run: {e}')
# Return results
if not training:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
model.float() # for training
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
return (mp, mr, map50, map,
*(loss.cpu() / len(dataloader)).tolist()), maps, t
def __getitem__(self, index):
index = self.indices[index] # linear, shuffled, or image_weights
hyp = self.hyp
mosaic = self.mosaic and random.random() < hyp['mosaic']
if mosaic:
# Load mosaic
img, labels = load_mosaic(self, index)
shapes = None
# MixUp https://arxiv.org/pdf/1710.09412.pdf
if random.random() < hyp['mixup']:
img2, labels2 = load_mosaic(self, random.randint(0, self.n - 1))
r = np.random.beta(8.0, 8.0) # mixup ratio, alpha=beta=8.0
img = (img * r + img2 * (1 - r)).astype(np.uint8)
labels = np.concatenate((labels, labels2), 0)
else:
# Load image
img, (h0, w0), (h, w) = load_image(self, index)
# Letterbox
shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape
img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
# Load labels
labels = []
x = self.labels[index]
if x.size > 0:
# Normalized xywh to pixel xyxy format
labels = x.copy()
labels[:, 1] = ratio[0] * w * (x[:, 1] - x[:, 3] / 2) + pad[0] # pad width
labels[:, 2] = ratio[1] * h * (x[:, 2] - x[:, 4] / 2) + pad[1] # pad height
labels[:, 3] = ratio[0] * w * (x[:, 1] + x[:, 3] / 2) + pad[0]
labels[:, 4] = ratio[1] * h * (x[:, 2] + x[:, 4] / 2) + pad[1]
if self.augment:
# Augment imagespace
if not mosaic:
img, labels = random_perspective(img, labels,
degrees=hyp['degrees'],
translate=hyp['translate'],
scale=hyp['scale'],
shear=hyp['shear'],
perspective=hyp['perspective'])
# Augment colorspace
augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])
# Apply cutouts
# if random.random() < 0.9:
# labels = cutout(img, labels)
nL = len(labels) # number of labels
if nL:
labels[:, 1:5] = xyxy2xywh(labels[:, 1:5]) # convert xyxy to xywh
labels[:, [2, 4]] /= img.shape[0] # normalized height 0-1
labels[:, [1, 3]] /= img.shape[1] # normalized width 0-1
if self.augment:
# flip up-down
if random.random() < hyp['flipud']:
img = np.flipud(img)
if nL:
labels[:, 2] = 1 - labels[:, 2]
# flip left-right
if random.random() < hyp['fliplr']:
img = np.fliplr(img)
if nL:
labels[:, 1] = 1 - labels[:, 1]
labels_out = torch.zeros((nL, 6))
if nL:
labels_out[:, 1:] = torch.from_numpy(labels)
# Convert
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
return torch.from_numpy(img), labels_out, self.img_files[index], shapes
def detect(save_img=True):
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://'))
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
elif source == 'own_camera':
pass
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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 names]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
if source == 'own_camera':
camera_matrix = np.array(
[[3.50379164e+03, 0.00000000e+00, 7.96197449e+02],
[0.00000000e+00, 3.49926605e+03, 6.94741115e+02],
[0.00000000e+00, 0.00000000e+00, 1.00000000e+00]])
dist_coefs = np.array([
-6.93155538e-01, 2.44136043e+00, -8.62122690e-03, 6.01043129e-03,
-1.64531179e+01
])
camera = pylon.InstantCamera(
pylon.TlFactory.GetInstance().CreateFirstDevice())
camera.Open()
converter = pylon.ImageFormatConverter()
# camera = pylon.InstantCamera(pylon.TlFactory.GetInstance().CreateFirstDevice())
# ========== Grabing Continusely (video) with minimal delay ==========
camera.StartGrabbing(pylon.GrabStrategy_LatestImageOnly)
camera.ExposureTimeAbs = 80000
# ========== converting to opencv bgr format ==========
converter.OutputPixelFormat = pylon.PixelType_BGR8packed
converter.OutputBitAlignment = pylon.OutputBitAlignment_MsbAligned
pygame.init()
# clock = pygame.time.Clock()
display_width = 1920
display_height = 1080
screen = pygame.display.get_surface()
screen = pygame.display.set_mode((display_width, display_height),
pygame.FULLSCREEN)
modes = pygame.display.list_modes()
pygame.display.set_mode(max(modes))
black = (0, 0, 0)
white = (255, 255, 255)
# screen.fill(black) # show black image
while camera.IsGrabbing():
# for projector
for event in pygame.event.get():
if event.type == pygame.QUIT:
crashed = True
screen.fill(black) # show black image
grabResult = camera.RetrieveResult(
5000, pylon.TimeoutHandling_ThrowException)
if grabResult.GrabSucceeded():
# image = converter.Convert(grabResult)
image = converter.Convert(grabResult)
img = image.GetArray()
width = int(img.shape[1])
height = int(img.shape[0])
dim = (width, height)
newcameramtx, roi = cv2.getOptimalNewCameraMatrix(
camera_matrix, dist_coefs, (width, height), 1,
(width, height))
dst = cv2.undistort(img, camera_matrix, dist_coefs, None,
newcameramtx)
x, y, width, height = roi
img0 = dst[y:y + height, x:x + width]
frame = img0
# frame = cv2.resize(frame, (int(frame.shape[1]/1.5),int(frame.shape[0]/1.5)))
img0 = frame.copy()
im0s = img0
img_size = 640
img, ratio, (dw, dh) = letterbox(img0, new_shape=img_size)
img = img[:, :, ::-1].transpose(2, 0,
1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = Path('0'), '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = Path('0'), '', im0s
save_path = str(save_dir / p.name)
# txt_path = str(save_dir / 'labels' / p.stem) + (
# '_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(img0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
classes = []
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
s1 = '%g' % (n)
classes.append([names[int(c)], s1])
df_cls_info = pd.DataFrame(classes,
columns=['Class', 'number'])
df_cls_info.to_csv('class_num.csv')
print("\n number of object : ", len(det[:, -1]))
print(df_cls_info)
# Write results
box_info = []
box3D = {}
indice_3D = 0
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (
cls, *xywh) # label format
# with open(txt_path + '.txt', 'a') as f:
# f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
img0,
label=label,
color=colors[int(cls)],
line_thickness=3)
c1, c2 = (int(xyxy[0]),
int(xyxy[1])), (int(xyxy[2]),
int(xyxy[3]))
# print("name : ",names[int(cls)])
# print("opt",opt.classflash)
if names[int(cls)] == opt.classflash:
box3D[indice_3D] = (int(xyxy[0]), int(xyxy[1]),
int(xyxy[2]), int(xyxy[3]))
elif opt.classflash == "all":
box3D[indice_3D] = (int(xyxy[0]), int(xyxy[1]),
int(xyxy[2]), int(xyxy[3]))
indice_3D += 1
confident = '%g' % (conf)
box_info.append([label, c1, c2, confident])
# print(box3D)
if box3D != {}:
depth, box_min, dists = D435camera(box3D, img0)
# print(depth, box_min, dists)
# screen.fill(black) # show black image
if dists != {}:
img = drawimage(box3D, dists,
(200, 0, 0)) # set color
# _,img = cv2.threshold(img,0,255,cv2.THRESH_BINARY_INV)
imgpro = pygame.image.frombuffer(
img.tostring(), img.shape[1::-1], "RGB")
screen.blit(imgpro, (50, 50))
else:
screen.fill(black)
pygame.display.flip()
pygame.display.update()
# df_box = pd.DataFrame(box_info,columns=['Class','top left','bottom_right','confidence'])
# df_box.to_csv('box_info.csv')
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
im0 = cv2.resize(
im0, (int(im0.shape[1] / 1.5), int(im0.shape[0] / 1.5)))
cv2.imshow("asd", im0)
if cv2.waitKey(1) == ord('q'): # q to quit
break
if save_txt or save_img:
print('Results saved to %s' % save_dir)
print('Done. (%.3fs)' % (time.time() - t0))
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = Path(path[i]), '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = Path(path), '', im0s
save_path = str(save_dir / p.name)
txt_path = str(save_dir / 'labels' / p.stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (
cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow("asd", im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
# cv2.imwrite(save_path, im0)
pass
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps,
(w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % save_dir)
print('Done. (%.3fs)' % (time.time() - t0))
def detect(save_img=False):
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source.isnumeric() or source.startswith(
'rtsp') or source.startswith('http') or source.endswith('.txt')
# Initialize
set_logging()
device = select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per img
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
txt_path = str(Path(out) / Path(p).stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
handles_ymax = []
#handles_xmid = []
handles_ymid = []
handle_mids = []
tailgates_ymin = []
tailgates_ymax = []
tailgate_ythird_coord = []
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
det_sorted = sorted(det, key=lambda x: x[
-1]) # sort detected items by last index which is class
# Write results
for *xyxy, conf, cls in reversed(
det_sorted
): #coords, confidence, classes.... reversed for some reason? But actually helpful since plate is cls 2
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * 5 + '\n') %
(cls, *xywh)) # label format
if int(cls) == 2: # license plate
license_width = abs(int(xyxy[2]) - int(xyxy[0]))
px_ratio = license_width / 12 # number of pixels per inch as license plates are 12"
if save_img or view_img: # Add bbox to image
#label = '%s %.2f' % (names[int(cls)], conf) #confidence not needed
label = '%s ' % (names[int(cls)])
coord1, coord2, dim_label = plot_one_box(
xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3,
px_ratio=px_ratio)
# get important points for line drawing
if int(cls) == 1: #handle
ymax = max(coord1[1], coord2[1])
handles_ymax.append(ymax)
xmid = int((coord1[0] + coord2[0]) / 2)
ymid = int((coord1[1] + coord2[1]) / 2)
handle_mids.append([xmid, ymid])
handle_img = img[:, coord1[0]:coord2[0],
coord1[1]:coord2[1]]
elif int(cls) == 0: #tailgate
tailgate_xmin = min(coord1[0], coord2[0])
ymax = max(coord1[1], coord2[1])
tailgates_ymax.append(ymax)
ymin = min(coord1[1], coord2[1])
tailgates_ymin.append(ymin)
tailgate_ythird = int(
abs(coord1[1] - coord2[1]) / 3 + ymin)
tailgate_ythird_coord.append(
[tailgate_xmin, tailgate_ythird])
#tailgate_img = img[coord1[0]:coord2[0], coord1[1]:coord2[1]]
#
# added ability to measure between bottom of handle and bottom of tailgate if handle in top 1/3
for i, (handle_mid,
max_point) in enumerate(zip(handle_mids,
handles_ymax)):
hyps = [
hypotenuse(handle_mid, b)
for b in tailgate_ythird_coord
]
closest_index = np.argmin(hyps)
if handle_mid[1] < tailgate_ythird_coord[closest_index][1]:
min_dist_tg = min(
[int(abs(max_point - x)) for x in tailgates_ymax])
start_point = (handle_mid[0], handles_ymax[i])
end_point = (handle_mid[0],
handles_ymax[i] + min_dist_tg)
cv2.line(im0, start_point, end_point, (100, 100, 0), 4)
line_mid = int((start_point[1] + end_point[1]) / 2)
label = f'Distance: {((end_point[1] - start_point[1]) / px_ratio):.4f}"'
cv2.putText(im0,
label, (start_point[0], line_mid),
0,
1, [0, 0, 0],
thickness=2,
lineType=cv2.LINE_AA)
### Previous ability to measure between bottom of handle and tailgate --- was not robust.
### Keeping until determined not needed
# for i, (mid_point, max_point) in enumerate(zip(handles_ymid, handles_ymax)):
# print(f'\nmidpoint: {mid_point}')
# min_y_dist = min([int(abs(mid_point - x)) for x in tailgate_ythird]) # gets min distance from handle midpoint to tailgate third
# print(f'min y dist: {min_y_dist}')
# print(f'tailgate third: {tailgate_ythird}')
# min_dist_third = min([x for x in tailgate_ythird if abs(x - min_y_dist) in handles_ymid])
# print(f'min_dist_third: {min_dist_third}')
# if mid_point < min_dist_third: #handle mid point in top 1/3 of truck
# min_dist_tg = min([int(abs(max_point - x)) for x in tailgates_ymax])
# print(f'min_dist_tg {min_dist_tg}')
# start_point = (handles_xmid[i], handles_ymax[i])
# print(f'start point: {start_point}')
# end_point = (handles_xmid[i], handles_ymax[i] + min_dist_tg)
# print(f'end point: {end_point}')
# cv2.line(im0, start_point, end_point, (100,100,0), 4)
# label = f'Distance: {min_dist_tg/300:.4f}"L'
# line_mid = int((start_point[1] + end_point[1])/2)
# cv2.putText(im0, label, (start_point[0], line_mid), 0, 1, [0, 0, 0],
# thickness=2, lineType=cv2.LINE_AA)
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps,
(w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % Path(out))
if platform.system() == 'Darwin' and not opt.update: # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
return tailgate_img, handle_img
def detect(opt):
save_img = False
out, source, weights, view_img, save_txt, imgsz = \
opt.save_dir, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source.isnumeric() or source.startswith(
('rtsp://', 'rtmp://', 'http://')) or source.endswith('.txt')
# Initialize
set_logging()
device = select_device(opt.device)
flag = False
if os.path.exists(out): # output dir
shutil.rmtree(out) # delete dir
os.makedirs(out) # make new dir
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
names[0], names[1] = names[1], names[0]
colors = [[0, 0, 255], [0, 255, 0]]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
flag = False
if webcam: # batch_size >= 1
p, s, im0 = path[i], ' ', im0s[i].copy()
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
txt_path = str(Path(out) / Path(p).stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%ss : %g, ' % (names[int(c)], n) # add to string
if (names[int(c)] == 'without_mask'):
flag = True
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, conf, *xywh) if opt.save_conf else (
cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line) + '\n') % line)
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
# Print time (inference + NMS)
now = datetime.now()
dt_string = now.strftime("%d_%b_%y %H_%M_%S")
dt_folder = now.strftime("%d_%b_%y")
dt_file = now.strftime("%H_%M_%S")
# if (s != ' '):
# print(dt_string," ",(s))
try:
os.makedirs("inference/data/" + dt_folder)
except FileExistsError:
# directory already exists
pass
# Stream results
if view_img:
# im0 = cv2.resize(im0,(1120,840))
if (flag):
cv2.imwrite(
"inference\\data\\" + dt_folder + "\\" + dt_string +
".png", im0)
cv2.imshow(p, im0)
# if cv2.waitKey(1) == ord('q'): # q to quit
# cv2.VideoCapture(source).release()
# cv2.destroyAllWindows()
# return
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps,
(w, h))
vid_writer.write(im0)
def predict(opt, model, img):
out, source, view_img, save_img, save_txt, imgsz = \
opt['output'], opt['source'], opt['view_img'], opt['save_img'], opt['save_txt'], opt['imgsz']
#webcam = source.isnumeric() or source.startswith(('rtsp://', 'rtmp://', 'http://')) or source.endswith('.txt')
# Initialize
device = select_device(opt['device']) # 选择设备
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
im0_shape = img.shape # 记下原始图片的尺寸
#print('im0_shape = %s \n' % str(im0_shape))
# Load model
# model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Set Dataloader
dataset = LoadImages(opt['source'], img_size=imgsz)
# 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))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for path, img, im0s, _ in dataset:
print('input_path:', path)
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
# 前向推理
pred = model(img, augment=opt['augment'])[0]
# Apply NMS（非极大抑制）
pred = non_max_suppression(pred,
opt['conf_thres'],
opt['iou_thres'],
classes=opt['classes'],
agnostic=opt['agnostic_nms'])
t2 = time_synchronized()
# Process detections
for i, det in enumerate(pred): # detections per image
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name) # 保存路径
txt_path = str(Path(out) / Path(p).stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
boxes_detected = [] #检测结果
for *xyxy, conf, cls in reversed(det):
xyxy_list = (torch.tensor(xyxy).view(1,
4)).view(-1).tolist()
boxes_detected.append({
"name":
id2name[int(cls.item())],
"conf":
str(conf.item()),
"bbox": [
int(xyxy_list[0]),
int(xyxy_list[1]),
int(xyxy_list[2]),
int(xyxy_list[3])
]
})
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * 5 + '\n') %
(cls, *xywh)) # label format
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
# Print time (inference + NMS)
# print('%sDone. (%.3fs)' % (s, t2 - t1))
# Save results (image with detections)
if save_img:
if dataset.mode == 'image':
print('output_path:', save_path)
cv2.imwrite(save_path, im0)
# Print time (inference + NMS)
print('predect device: %s' % device.type)
print('predect time: %.3fs' % (t2 - t1))
results = {"results": boxes_detected}
# print(results)
return results
def detect(number_person):
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# Initialize
device = select_device(opt.device)
# if os.path.exists(out):
# shutil.rmtree(out) # delete output folder
if not os.path.exists(out):
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
half = False
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
# model=torch.load(weights, map_location=device)['model'].float().eval()
# stride = [8, 16, 32]
# imgsz = check_img_size(imgsz, s=max(stride)) # check img_size
# model = Darknet('cfg/prune_0.8_yolov3-spp.cfg', (opt.img_size, opt.img_size)).to(device)
# initialize_weights(model)
# model.load_state_dict(torch.load('weights/prune_0.8_yolov3-spp-ultralytics.pt')['model'])
# model.eval()
# stride = [8, 16, 32]
# imgsz = check_img_size(imgsz, s=max(stride)) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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()
# Set Dataloader
vid_path, vid_writer = None, None
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
# names = ['1', '2']
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for dirs in os.listdir(source):
# if dirs !='WH':
# continue
src = os.path.join(source, dirs)
save_img = True
save_xml = False
dataset = LoadImages(src, img_size=imgsz)
for path, img, im0s, vid_cap in dataset:
# if os.path.basename(path)!='2_31746253093C100D_2018-12-10-21-56-37-998_0_75_636_307_6.jpg':
# continue
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
t2 = time_synchronized()
# Apply NMS
pred = non_max_suppression_test(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
# pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes,
# agnostic=opt.agnostic_nms)
# t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
# save_path = str(Path(out) / Path(p).name)
# txt_path = str(Path(out) / Path(p).stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
# results = [0, 0]
minconf = 1
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
results = [0, 0]
for *xyxy, conf, cls in det:
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
# with open(txt_path + '.txt', 'a') as f:
# f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format
if save_img or view_img: # Add bbox to image
if names[int(cls)] == "1":
results[0] += 1
elif names[int(cls)] == "2":
results[1] += 1
else:
plot_one_box(xyxy,
im0,
label=None,
color=colors[int(cls)],
line_thickness=3)
continue
# else:
# results[1] += 1
minconf = min(conf, minconf)
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
if len(results) == 2:
if (results[0] > number_person) | (results[1] >
number_person):
tmp = "err"
elif (results[0] == number_person) | (results[1]
== number_person):
tmp = "corr"
# if len(det) == 1:
# tmp = 'single'
# elif minconf < 0.5:
# tmp='05'
# elif minconf < 0.6:
# tmp='06'
# elif minconf < 0.7:
# tmp='07'
# elif minconf < 0.8:
# tmp = '08'
else:
tmp = "miss"
elif len(results) == 1:
if (results[0] == number_person):
tmp = "corr"
elif (results[0] > number_person):
tmp = "err"
else:
tmp = "miss"
else:
tmp = "miss"
save_path = os.path.join(Path(out), dirs, tmp) #, tmp
if not os.path.exists(save_path):
os.makedirs(save_path)
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(os.path.join(save_path, Path(p).name), im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc),
fps, (w, h))
vid_writer.write(im0)
if det is not None and len(det) and save_xml:
# xml
output_name = os.path.join(save_path, Path(p).name)
create_tree(output_name)
for *xyxy, conf, cls in det:
label = names[int(cls)]
left, top, right, bottom = torch.tensor(xyxy).view(
1, 4).view(-1).tolist()
create_object(annotation, label, left, top, right,
bottom)
# 将树模型写入xml文件
tree = ET.ElementTree(annotation)
# tree.write('%s.xml' % output_name.rstrip('.jpg'))
# tree = ET.ElementTree.parse('%s.xml' % output_name.rstrip('.jpg')) # 解析movies.xml这个文件
root = tree.getroot() # 得到根元素，Element类
pretty_xml(root, '\t', '\n') # 执行美化方法
tree.write('%s.xml' % output_name.rstrip('.jpg'))
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
# if platform == 'darwin' and not opt.update: # MacOS
# os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
def detect(save_img=True):
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://'))
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
elif source == 'own_camera':
pass
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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 names]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
# work space
if source == 'own_camera':
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1440)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
cap.set(cv2.CAP_PROP_FPS, 30)
img_name = 0
while True:
# print(source)
_, frame = cap.read()
print(img_name)
print(img_name % 30)
if opt.collect == 'True' and img_name % 5 == 0:
cv2.imwrite("./data/collect/" + str(img_name) + ".jpg", frame)
img_name += 1
im0 = []
img0 = frame.copy()
im0s = img0
img_size = 640
img, ratio, (dw, dh) = letterbox(img0, new_shape=img_size)
img = img[:, :, ::-1].transpose(2, 0,
1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
t1 = time_synchronized()
cmd = cv2.waitKey(1)
if opt.command == "True" and cmd == ord("c"):
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = Path('0'), '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = Path('0'), '', im0s
save_path = str(save_dir / p.name)
# txt_path = str(save_dir / 'labels' / p.stem) + (
# '_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(
img0.shape)[[1, 0, 1, 0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
classes = []
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]
) # add to string
s1 = '%g' % (n)
classes.append([names[int(c)], s1])
df_cls_info = pd.DataFrame(classes,
columns=['Class', 'number'])
now = datetime.now()
dt_string = now.strftime("%d/%m/%Y %H:%M:%S")
dt_list = list(dt_string)
print(dt_list)
for dt_i in range(len(dt_list)):
if dt_list[dt_i] == ":" or dt_list[dt_i] == "/":
dt_list[dt_i] = "_"
dt_string = "".join(dt_list)
df_cls_info.to_csv('X:/Desktop/test/class_num_' +
dt_string + '.csv')
print("\n number of object : ", len(det[:, -1]))
print(df_cls_info)
# Write results
box_info = []
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (
xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh,
conf) if opt.save_conf else (
cls, *xywh) # label format
# with open(txt_path + '.txt', 'a') as f:
# f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
img0,
label=label,
color=colors[int(cls)],
line_thickness=3)
c1, c2 = (int(xyxy[0]),
int(xyxy[1])), (int(xyxy[2]),
int(xyxy[3]))
confident = '%g' % (conf)
box_info.append([label, c1, c2, confident])
df_box = pd.DataFrame(box_info,
columns=[
'Class', 'top left',
'bottom_right', 'confidence'
])
df_box.to_csv('box_info.csv')
cv2.imwrite('X:/Desktop/test/img_' + dt_string + '.jpg', im0)
elif opt.command == "False":
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = Path('0'), '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = Path('0'), '', im0s
save_path = str(save_dir / p.name)
# txt_path = str(save_dir / 'labels' / p.stem) + (
# '_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(
img0.shape)[[1, 0, 1, 0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
classes = []
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]
) # add to string
s1 = '%g' % (n)
classes.append([names[int(c)], s1])
df_cls_info = pd.DataFrame(classes,
columns=['Class', 'number'])
df_cls_info.to_csv('X:/Desktop/test/class_num.csv')
print("\n number of object : ", len(det[:, -1]))
print(df_cls_info)
# Write results
box_info = []
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (
xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh,
conf) if opt.save_conf else (
cls, *xywh) # label format
# with open(txt_path + '.txt', 'a') as f:
# f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
img0,
label=label,
color=colors[int(cls)],
line_thickness=3)
c1, c2 = (int(xyxy[0]),
int(xyxy[1])), (int(xyxy[2]),
int(xyxy[3]))
confident = '%g' % (conf)
box_info.append([label, c1, c2, confident])
df_box = pd.DataFrame(box_info,
columns=[
'Class', 'top left',
'bottom_right', 'confidence'
])
df_box.to_csv('box_info.csv')
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
if im0 == []:
im0 = frame.copy()
cv2.imshow("asd", im0)
if cv2.waitKey(1) == ord('q'): # q to quit
break
if save_txt or save_img:
print('Results saved to %s' % save_dir)
print('Done. (%.3fs)' % (time.time() - t0))
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = Path(path[i]), '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = Path(path), '', im0s
save_path = str(save_dir / p.name)
txt_path = str(save_dir / 'labels' / p.stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (
cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow("asd", im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
# cv2.imwrite(save_path, im0)
pass
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps,
(w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % save_dir)
print('Done. (%.3fs)' % (time.time() - t0))
def predict():
file = request.files['image']
file_name = file.filename
path = os.path.join('./upload', file_name)
file.save(path)
source = f'./upload/{file_name}'
weights, view_img, save_txt, imgsz = './model/best-m.pt', opt.view_img, opt.save_txt, 640
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
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True,
exist_ok=True) # make dir
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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 names]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
# run once
_ = model(img.half() if half else img) if device.type != 'cpu' else None
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(
pred, 0.7, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(
), dataset.count
else:
p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / 'labels' / p.stem) + \
('' if dataset.mode == 'image' else f'_{frame}') # img.txt
s += '%gx%g ' % img.shape[2:] # print string
# normalization gain whwh
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(
img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f'{n} {names[int(c)]}s, ' # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)
) / gn).view(-1).tolist() # normalized xywh
# label format
line = (
cls, *xywh, conf) if opt.save_conf else (cls, *xywh)
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = f'{names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy, im0, label=label,
color=colors[int(cls)], line_thickness=3)
# Print time (inference + NMS)
print(f'{s}Done. ({t2 - t1:.3f}s)')
# Stream results
if view_img:
cv2.imshow(str(p), im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video'
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
vid_writer.write(im0)
time.sleep(2)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
print(f'Done. ({time.time() - t0:.3f}s)')
# prepare image for response
# _, img_encoded = cv2.imencode('.png', im0)
# response = img_encoded.tostring()
# return Response(response=response, status=200, mimetype='image/png')
resultPath = str(save_dir)[13:]+'/'+file_name
return jsonify({"imgsrc": resultPath})
def test(
data,
weights=None,
batch_size=16,
imgsz=640,
conf_thres=0.001,
iou_thres=0.6, # for NMS
save_json=False,
single_cls=False,
augment=False,
verbose=False,
model=None,
dataloader=None,
save_dir='',
merge=False,
save_txt=False):
# Initialize/load model and set device
training = model is not None
if training: # called by train.py
device = next(model.parameters()).device # get model device
else: # called directly
set_logging()
device = select_device(opt.device, batch_size=batch_size)
merge, save_txt = opt.merge, opt.save_txt # use Merge NMS, save *.txt labels
if save_txt:
out = Path('inference/output')
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Remove previous
for f in glob.glob(str(Path(save_dir) / 'test_batch*.jpg')):
os.remove(f)
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
# model = nn.DataParallel(model)
# Half
half = device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half()
# Configure
model.eval()
with open(data) as f:
data = yaml.load(f, Loader=yaml.FullLoader) # model dict
check_dataset(data) # check
nc = 1 if single_cls else int(data['nc']) # number of classes
iouv = torch.linspace(0.5, 0.95,
10).to(device) # iou vector for [email protected]:0.95
niou = iouv.numel()
# Dataloader
if not training:
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
path = data['test'] if opt.task == 'test' else data[
'val'] # path to val/test images
dataloader = create_dataloader(path,
imgsz,
batch_size,
model.stride.max(),
opt,
hyp=None,
augment=False,
cache=True,
pad=0.5,
rect=True)[0]
seen = 0
names = model.names if hasattr(model, 'names') else model.module.names
coco91class = coco80_to_coco91_class()
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R',
'[email protected]', '[email protected]:.95')
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class = [], [], [], []
evaluator = COCOEvaluator(root=DATA_ROOT,
model_name=opt.weights.replace('.pt', ''))
for batch_i, (img, targets, paths,
shapes) in enumerate(tqdm(dataloader, desc=s)):
img = img.to(device, non_blocking=True)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
nb, _, height, width = img.shape # batch size, channels, height, width
whwh = torch.Tensor([width, height, width, height]).to(device)
# Disable gradients
with torch.no_grad():
# Run model
t = time_synchronized()
inf_out, train_out = model(
img, augment=augment) # inference and training outputs
t0 += time_synchronized() - t
# Compute loss
if training: # if model has loss hyperparameters
loss += compute_loss([x.float() for x in train_out], targets,
model)[1][:3] # GIoU, obj, cls
# Run NMS
t = time_synchronized()
output = non_max_suppression(inf_out,
conf_thres=conf_thres,
iou_thres=iou_thres,
merge=merge)
t1 += time_synchronized() - t
# Statistics per image
for si, pred in enumerate(output):
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class
seen += 1
if pred is None:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool),
torch.Tensor(), torch.Tensor(), tcls))
continue
# Append to text file
if save_txt:
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0
]] # normalization gain whwh
x = pred.clone()
x[:, :4] = scale_coords(img[si].shape[1:], x[:, :4],
shapes[si][0],
shapes[si][1]) # to original
for *xyxy, conf, cls in x:
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(str(out / Path(paths[si]).stem) + '.txt',
'a') as f:
f.write(
('%g ' * 5 + '\n') % (cls, *xywh)) # label format
# Clip boxes to image bounds
clip_coords(pred, (height, width))
# Append to pycocotools JSON dictionary
if save_json:
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
image_id = Path(paths[si]).stem
box = pred[:, :4].clone() # xyxy
scale_coords(img[si].shape[1:], box, shapes[si][0],
shapes[si][1]) # to original shape
box = xyxy2xywh(box) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for p, b in zip(pred.tolist(), box.tolist()):
result = {
'image_id':
int(image_id) if image_id.isnumeric() else image_id,
'category_id': coco91class[int(p[5])],
'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)
}
jdict.append(result)
#evaluator.add([result])
#if evaluator.cache_exists:
# break
# # Assign all predictions as incorrect
# correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)
# if nl:
# detected = [] # target indices
# tcls_tensor = labels[:, 0]
#
# # target boxes
# tbox = xywh2xyxy(labels[:, 1:5]) * whwh
#
# # Per target class
# for cls in torch.unique(tcls_tensor):
# ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # prediction indices
# pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # target indices
#
# # Search for detections
# if pi.shape[0]:
# # Prediction to target ious
# ious, i = box_iou(pred[pi, :4], tbox[ti]).max(1) # best ious, indices
#
# # Append detections
# detected_set = set()
# for j in (ious > iouv[0]).nonzero(as_tuple=False):
# d = ti[i[j]] # detected target
# if d.item() not in detected_set:
# detected_set.add(d.item())
# detected.append(d)
# correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn
# if len(detected) == nl: # all targets already located in image
# break
#
# # Append statistics (correct, conf, pcls, tcls)
# stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
# # Plot images
# if batch_i < 1:
# f = Path(save_dir) / ('test_batch%g_gt.jpg' % batch_i) # filename
# plot_images(img, targets, paths, str(f), names) # ground truth
# f = Path(save_dir) / ('test_batch%g_pred.jpg' % batch_i)
# plot_images(img, output_to_target(output, width, height), paths, str(f), names) # predictions
evaluator.add(jdict)
evaluator.save()
def detect(save_img=False, object_csv_path, output_video_path, mode='video'):
object_data = []
if mode == 'video':
v = Video(opt.source)
vname = Path(opt.source).name
object_image_dir = Path(opt.source).stem
u.mkdir_if_not_exists(object_image_dir)
fps = v.get_fps()
det_data = []
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://'))
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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 names]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for ffid, (path, img, im0s, vid_cap) in enumerate(dataset):
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
pred = model(img, augment=opt.augment)[0]
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
frame_id = i
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(
), dataset.count
else:
p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / 'labels' / p.stem) + (
'' if dataset.mode == 'image' else f'_{frame}') # img.txt
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f'{n} {names[int(c)]}s, ' # add to string
# Write results
new_det = []
for obj_id, (xmin, ymin, xmax, ymax, conf,
cls) in enumerate(det):
xmin_val = xmin.cpu().numpy().astype(int).item()
ymin_val = ymin.cpu().numpy().astype(int).item()
xmax_val = xmax.cpu().numpy().astype(int).item()
ymax_val = ymax.cpu().numpy().astype(int).item()
w_val = xmax_val - xmin_val
h_val = ymax_val - ymin_val
conf = conf.cpu().numpy().item()
cls = cls.cpu().numpy().item()
rows, cols = im0.shape[:2]
if ((ymin_val / rows) > 0.05) and (
(w_val / h_val) < 10) and ((h_val / w_val) < 10):
if (w_val / cols > 0.02) or (h_val / rows > 0.02):
obj_im = im0[ymin_val:ymax_val, xmin_val:xmax_val]
green_percent, _ = get_green_percent_BGR(obj_im)
green_percent = round(green_percent, 2)
if green_percent > 0.3:
if mode == 'video':
ftime = fid_2_time(ffid, v)
cv2.imwrite(
f'./{object_image_dir}/ftime{ftime:.2f}s_fid{ffid}_obj{obj_id}_percent{green_percent}.png',
obj_im)
new_det.append(
[xmin, ymin, xmax, ymax, conf, cls])
det = new_det
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (
cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = f'{names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
# Save results (image with detections)
raise ValueError(save_img)
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video'
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps,
(w, h))
vid_writer.write(im0)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
print(f'Done. ({time.time() - t0:.3f}s)')
def run(
weights=ROOT / 'yolov5s.pt', # model.pt path(s)
source=ROOT / 'data/images', # file/dir/URL/glob, 0 for webcam
data=ROOT / 'data/coco128.yaml', # dataset.yaml path
imgsz=(640, 640), # inference size (height, width)
conf_thres=0.25, # confidence threshold
iou_thres=0.45, # NMS IOU threshold
max_det=1000, # maximum detections per image
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=False, # show results
save_txt=False, # save results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_crop=False, # save cropped prediction boxes
nosave=False, # do not save images/videos
classes=None, # filter by class: --class 0, or --class 0 2 3
agnostic_nms=False, # class-agnostic NMS
augment=False, # augmented inference
visualize=False, # visualize features
update=False, # update all models
project=ROOT / 'runs/detect', # save results to project/name
name='exp', # save results to project/name
exist_ok=False, # existing project/name ok, do not increment
line_thickness=3, # bounding box thickness (pixels)
hide_labels=False, # hide labels
hide_conf=False, # hide confidences
half=False, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
):
source = str(source)
save_img = not nosave and not source.endswith(
'.txt') # save inference images
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
is_url = source.lower().startswith(
('rtsp://', 'rtmp://', 'http://', 'https://'))
webcam = source.isnumeric() or source.endswith('.txt') or (is_url
and not is_file)
if is_url and is_file:
source = check_file(source) # download
# Directories
save_dir = increment_path(Path(project) / name,
exist_ok=exist_ok) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Load model
device = select_device(device)
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data)
stride, names, pt, jit, onnx, engine = model.stride, model.names, model.pt, model.jit, model.onnx, model.engine
imgsz = check_img_size(imgsz, s=stride) # check image size
# Half
half &= (
pt or jit or onnx or engine
) and device.type != 'cpu' # FP16 supported on limited backends with CUDA
if pt or jit:
model.model.half() if half else model.model.float()
# Dataloader
if webcam:
view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)
bs = len(dataset) # batch_size
else:
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
bs = 1 # batch_size
vid_path, vid_writer = [None] * bs, [None] * bs
# Run inference
model.warmup(imgsz=(1, 3, *imgsz), half=half) # warmup
dt, seen = [0.0, 0.0, 0.0], 0
for path, im, im0s, vid_cap, s in dataset:
t1 = time_sync()
im = torch.from_numpy(im).to(device)
im = im.half() if half else im.float() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
if len(im.shape) == 3:
im = im[None] # expand for batch dim
t2 = time_sync()
dt[0] += t2 - t1
# Inference
visualize = increment_path(save_dir / Path(path).stem,
mkdir=True) if visualize else False
pred = model(im, augment=augment, visualize=visualize)
t3 = time_sync()
dt[1] += t3 - t2
# NMS
pred = non_max_suppression(pred,
conf_thres,
iou_thres,
classes,
agnostic_nms,
max_det=max_det)
dt[2] += time_sync() - t3
# Second-stage classifier (optional)
# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
# Process predictions
for i, det in enumerate(pred): # per image
seen += 1
if webcam: # batch_size >= 1
p, im0, frame = path[i], im0s[i].copy(), dataset.count
s += f'{i}: '
else:
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # im.jpg
txt_path = str(save_dir / 'labels' / p.stem) + (
'' if dataset.mode == 'image' else f'_{frame}') # im.txt
s += '%gx%g ' % im.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
imc = im0.copy() if save_crop else im0 # for save_crop
annotator = Annotator(im0,
line_width=line_thickness,
example=str(names))
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(im.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh,
conf) if save_conf else (cls,
*xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or save_crop or view_img: # Add bbox to image
c = int(cls) # integer class
label = None if hide_labels else (
names[c]
if hide_conf else f'{names[c]} {conf:.2f}')
annotator.box_label(xyxy, label, color=colors(c, True))
if save_crop:
save_one_box(xyxy,
imc,
file=save_dir / 'crops' / names[c] /
f'{p.stem}.jpg',
BGR=True)
# Print time (inference-only)
LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
# Stream results
im0 = annotator.result()
if view_img:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections)
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video' or 'stream'
if vid_path[i] != save_path: # new video
vid_path[i] = save_path
if isinstance(vid_writer[i], cv2.VideoWriter):
vid_writer[i].release(
) # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
save_path += '.mp4'
vid_writer[i] = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps,
(w, h))
vid_writer[i].write(im0)
# Print results
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
LOGGER.info(
f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}'
% t)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
if update:
strip_optimizer(weights) # update model (to fix SourceChangeWarning)
def detect(self, save_img=False) -> list:
source, weights, view_img, save_txt, imgsz = self.opt.source, self.opt.weights, self.opt.view_img, self.opt.save_txt, self.opt.img_size
webcam = source.isnumeric() or source.endswith(
'.txt') or source.lower().startswith(
('rtsp://', 'rtmp://', 'http://'))
# Initialize
set_logging()
device = select_device(self.opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# print(half)
# Load model
t0 = time.time()
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
# if half:
# model.half() # to FP16
# Second-stage classifier
# 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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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 names]
# Run inference
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
print(time.time() - t0)
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
return_imgs = []
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=False)[0]
# Apply NMS
pred = non_max_suppression(pred,
self.opt.conf_thres,
self.opt.iou_thres,
classes=self.opt.classes,
agnostic=self.opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
# if classify:
# pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0, frame = Path(
path[i]), '%g: ' % i, im0s[i].copy(), dataset.count
else:
p, s, im0, frame = Path(path), '', im0s, getattr(
dataset, 'frame', 0)
save_path = str(self.opt.save_dir / p.name)
txt_path = str(self.opt.save_dir / 'labels' / p.stem) + (
'' if dataset.mode == 'image' else f'_{frame}')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f'{n} {names[int(c)]}s, ' # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh,
conf) if self.opt.save_conf else (
cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line +
'\n')
if save_img or view_img: # Add bbox to image
label = f'{names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=1)
# Print time (inference + NMS)
print(f'{s}Done. ({t2 - t1:.3f}s)')
# Stream results
if view_img:
cv2.imshow(str(p), im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
return_imgs.append(im0.copy())
if save_txt or save_img:
s = f"\n{len(list(self.opt.save_dir.glob('labels/*.txt')))} labels saved to {self.opt.save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {self.opt.save_dir}{s}")
print(f'Done. ({time.time() - t0:.3f}s)')
return return_imgs
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://'))
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
stride = int(model.stride.max()) # model stride
imgsz = check_img_size(imgsz, s=stride) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz, stride=stride)
# Get names and colors
if opt.labels:
names = load_labels(opt.labels)
else:
names = model.module.names if hasattr(model, 'module') else model.names
if opt.color == "same":
# same color for all
colors = [[255, 0, 255] for _ in range(len(names))]
else:
if opt.color == "det":
random.seed(2)
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
# Run inference
if device.type != 'cpu':
model(
torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(
next(model.parameters()))) # run once
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
if opt.bboxfilt:
# Detections filter
bbox_filter = BboxFilter(30, 5) # 30x30 pixel grids, 5 deep in time
# Object sizes
# se, comml, jet, heli, drone
OBj_SIZES = (10.0, 30.0, 10.0, 5.0, 0.3)
CAM_FOVH = 60 # degree
quit = False
for path, img, im0s, vid_cap in dataset:
if quit: break
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
#print("pred shape", pred.shape)
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(
pred
): # detections per image, this is always 1 except for maybe (batched images?)
#print("i det",i, det)
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(
), dataset.count
else:
p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
imh, imw, _ = im0.shape
p = Path(p) # to Path
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / 'labels' / p.stem) + (
'' if dataset.mode == 'image' else f'_{frame}') # img.txt
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
x1, y1, x2, y2 = int(xyxy[0]), int(xyxy[1]), int(
xyxy[2]), int(xyxy[3])
if opt.bboxfilt:
bbox_filter.add((x1, y1, x2, y2), conf.item(),
int(cls))
#print("imw", im0.shape, int(cls))
# if int(cls) > len(OBj_SIZES):
# sz = 1.0
# else:
# sz = OBj_SIZES[int(cls)]
# eul, quat, dist = estimate_yaw_pitch_dist(CAM_FOVH, sz, imw, imh, x1, y1, x2, y2)
else:
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (
cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line +
'\n')
if save_img or view_img: # Add bbox to image
label = f'{names[int(cls)]} {conf:.2f}'
# label = '%s %.2f %.1f %.2f %.2f %.2f %.2f' % (names[int(cls)], conf, dist, quat[0], quat[1], quat[2], quat[3])
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=2)
# Print time (inference + NMS)
print(f'{s}Done. ({t2 - t1:.3f}s)')
# Stream results
# if view_img:
# cv2.imshow(p, im0)
# if cv2.waitKey(1) == ord('q'): # q to quit
# raise StopIteration
# Save results (image with detections)
# if save_img:
# if dataset.mode == 'images':
# cv2.imwrite(save_path, im0)
# else:
# if vid_path != save_path: # new video
# vid_path = save_path
# if isinstance(vid_writer, cv2.VideoWriter):
# vid_writer.release() # release previous video writer
# fourcc = 'mp4v' # output video codec
# fps = vid_cap.get(cv2.CAP_PROP_FPS)
# w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
# h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
# vid_writer.write(im0)
if opt.bboxfilt:
print("-----------------BBOX")
boxes = bbox_filter.get_boxes()
#print("filt boxes",boxes)
for box in boxes:
print("b", box)
xyxy, conf, cls = box
# label = '%s %.2f' % (names[cls], conf)
# print(label, conf)
# plot_one_box(b, im0, label=label, color=colors[cls], line_thickness=3)
if int(cls) > len(OBj_SIZES):
sz = 1.0
else:
sz = OBj_SIZES[int(cls)]
eul, quat, dist = estimate_yaw_pitch_dist(
CAM_FOVH, sz, imw, imh, x1, y1, x2, y2)
if save_img or view_img: # Add bbox to image
# with quaternion
# label = '%s %.2f %.1f %.2f %.2f %.2f %.2f' % (names[int(cls)], conf, dist, quat[0], quat[1], quat[2], quat[3])
label = '%s %.2f %.1f' % (names[int(cls)], conf, dist)
print(label)
print("xyxy", xyxy)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=2)
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * 5 + '\n') %
(cls, *xywh)) # label format
if view_img:
cv2.imshow(str(p), im0)
k = cv2.waitKey(1) # 1 millisecond
if k == 27: # ESC
quit = True
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video'
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps,
(w, h))
vid_writer.write(im0)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
print(f'Done. ({time.time() - t0:.3f}s)')
def detect(save_img=False):
logging.basicConfig(filename='detect.log', level=logging.INFO)
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.lower().startswith(
('rtsp://', 'rtmp://', 'http://'))
imglist = opt.imlist
print("imglist : ", imglist)
source_list = source.split('\n')
#source_list = source
# Directories
save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# save_dir = Path(increment_path(Path(opt.project) / opt.name)) # increment run
# (save_dir / 'labels' if save_txt else save_dir).mkdir # make dir
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
stride = int(model.stride.max()) # model stride
imgsz = check_img_size(imgsz, s=stride) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride)
else:
save_img = True #이미지 저장을 하지 않기위해 주석처
if imglist:
dataset_list = []
for i in range(len(source_list)):
#print("source_list : ", source_list[i])
try:
dataset_list.append(LoadImages(source_list[i], img_size=imgsz, stride=stride))
except:
print("error!!!!!: ", source_list[i])
else:
dataset = LoadImages(source, img_size=imgsz, stride=stride)
# 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 names]
# Run inference
if device.type != 'cpu':
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
t0 = time.time()
#sujin
if imglist:
j = len(dataset_list)
#print("imglist True, j = ", j)
else:
print(dataset.nf)
detect_count =0
time_sum=0.0 #시간의 평균을 구하기위한 변
if imglist:
for k in range(j):
dataset = dataset_list[k] ## 필요없는 부분
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
prevTime = 0
# Process detections
for i, det in enumerate(pred): # detections per image
curTime = time.time() * 1000
sec = curTime - prevTime
prevTime = curTime #이전 시간을 현재시간으로 다시 저장시킴
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
fps_ = 1/(sec)
else:
p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
p = Path(p) # to Path
print("p.name=", p.name)
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = f'{names[int(cls)]} {conf:.2f}'
pt_start = time.time()*1000
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
pt_end = time.time()*1000
if os.path.exists(result_dir) == False:
os.makedirs(result_dir)
with open(result_dir + p.name[:-4] + '.txt', 'w') as f:
xmin=(int(xyxy[0]))
ymin=(int(xyxy[1]))
xmax=(int(xyxy[2]))
ymax=(int(xyxy[3]))
h, w, bs = im0.shape
print("bs h w = ",bs, h, w)
absolute_x = xmin + 0.5 * (xmax - xmin)
absolute_y = ymin + 0.5 * (ymax - ymin)
absolute_width = xmax - xmin
absolute_height = ymax - ymin
x = str(absolute_x / w)
y = str(absolute_y / h)
width = str(absolute_width / w)
height = str(absolute_height / h)
f.write(str(int(cls))+" "+x + " " + y + " " + width + " " + height)
detect_count += 1
print("detect_count = ",detect_count)
else:
tl = 3 or round(0.002 * (im0.shape[0] + im0.shape[1]) / 2) + 1 # line/font thickness
tf = max(tl - 1, 1) # font thickness
cv2.putText(im0, "0", (0, 100), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
# Print time (inference + NMS)
time_sum += t2-t1
print(f'{s}Done. ({t2 - t1:.3f}s)', ' avg fps=', time_sum/detect_count)
# Stream results
if view_img:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections)
if save_img:
if dataset.mode == 'image':
print("Save path=", save_path)
cv2.imwrite(save_path, im0)
else: # 'video'
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
fourcc = 'mp4v' # output video codec
fps = int(vid_cap.get(cv2.CAP_PROP_FPS))
print("fps is : ", fps)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
vid_writer.write(im0)
else:
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float()
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
else:
p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = f'{names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
# Print time (inference + NMS)
print(f'{s}Done. ({t2 - t1:.3f}s)')
# Stream results
if view_img:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections)
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video'
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
vid_writer.write(im0)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
if imglist:
print('detect rate = ',(detect_count),' / ',j-1,' =', detect_count/(j-1))
print('average fps = ', 1/(time_sum/(j-1)))
def do_detect(source = ' ', weights =' ', imlist= True, result = ' '):
print('<<Detection Start>>')
parser = argparse.ArgumentParser()
parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--view-img', action='store_true', help='display results')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
parser.add_argument('--augment', action='store_true', help='augmented inference')
parser.add_argument('--update', action='store_true', help='update all models')
parser.add_argument('--project', default='runs/detect', help='save results to project/name')
parser.add_argument('--name', default='exp', help='save results to project/name')
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
parser.add_argument('--imlist', action='store_true', help='existing project/name ok, do not increment')
global opt
global result_dir
result_dir = result
opt = parser.parse_args()
opt.source = source
opt.imlist = imlist
opt.weights= [weights]
check_requirements()
detect_start = time.time()*1000
print("DETECT START: ", detect_start)
with torch.no_grad():
detect()
detect_end = time.time()*1000-detect_start
logging.info("hfgf"+str(detect_end))
print("DETECT TIME : ", time.time()*1000-detect_start)
print("<<DETECT DONE>>")
def detect(opt, dp, save_img=False):
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# Initialize
device = select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max(
)) # check img_size 如果不是32的倍数，就向上取整调整至32的倍数并答应warning
if half:
model.half() # to FP16
# Second-stage classifier
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()
# Set Dataloader
vid_path, vid_writer = None, None
if opt.use_roi:
# print(dp.cl)
# print(dp.cl[0], dp.cl[1])
# cl = opt.control_line
cl = dp.cl
roi_in_pixels = np.array([0, cl[0], 1280,
cl[1]]) # two points coor, x1, y1, x2, y2
else:
roi_in_pixels = None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz, roi=roi_in_pixels)
# Get names and colors
names = model.module.names if hasattr(
model, 'module') else model.names # 解决GPU保存的模型多了module属性的问题
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))] # 随机颜色，对应names，names是class
# fix issue: when single cls, names = ['item'] rather than names = ['crosswalk']
if 'item' in names:
names = ['crosswalk']
# prune
# torch_utils.prune(model, 0.7)
model.eval()
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once 空跑一次，释放！！牛逼
detected_img_id = 0
time_list = [None] * len(dataset)
bar = tqdm(dataset)
for iii, (path, img, im0s, vid_cap, recover) in enumerate(bar):
# print(img.shape, im0s.shape, vid_cap)
# exit()
# img.shape [3, 384, 640] im0s.shape [720, 1280, 3] None
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0) # 从[3, h, w]转换为[batch_size, 3, h, w]的形式
# Inference
t1 = time_synchronized()
# print('aug', opt.augment) # False
pred = model(img, augment=opt.augment)[0]
# print(pred.shape) [1, 15120, 25]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
infer_time = t2 - t1
time_list[iii] = t2 - t1
# print('\n', len(pred), pred, recover) # list 长度是bs，代表每张图, 元素tensor，代表检测到的目标，每个tensor.shape [n, 6] xy4, conf, cls
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if opt.use_roi and det is not None:
small_img_shape = torch.from_numpy(
np.array([recover[1], recover[0]]).astype(np.float))
det[:,
0], det[:,
2] = det[:, 0] + recover[2], det[:, 2] + recover[2]
det[:,
1], det[:,
3] = det[:, 1] + recover[3], det[:, 3] + recover[3]
else:
small_img_shape = img.shape[2::]
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s # im0s是原图
save_path = str(Path(out) /
Path(p).name) # output/filenamexxxx.jpg
txt_path = str(Path(out) / Path(p).stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
# output/filenamexxxx.txt
s += '%gx%g ' % img.shape[2:] # print string, 640x640
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
# 本来是[720, 1280, 3]，重复取，变成[1280, 720, 1280, 720]
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(
small_img_shape, det[:, :4],
im0.shape).round() # 转换成原图的x1 y1 x2 y1，像素值
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]
) # add to string # i.e. 1 crosswalk
# s += f'{det[:, 4].item():.4f} '
# print(n)
# Write results
for *xyxy, conf, cls in det:
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(txt_path + '.txt', 'a') as f:
x, y, w, h = xywh
string = f"{int(cls)} {conf.item():.4f} {x:.6f} {y:.6f} {w:.6f} {h:.6f}\n"
f.write(string) # label format
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
# print(type(im0), im0.shape) array, 720, 1280, 3
if names[int(cls)] in opt.plot_classes:
# color = colors[int(cls)]
color = (255, 85, 33)
plot_one_box(xyxy,
im0,
label=label,
color=color,
line_thickness=5)
# Print time (inference + NMS)
prt_str = '%sDone. (%.5fs)' % (s, t2 - t1)
# print(prt_str)
os.system(f'echo "{prt_str}" >> {opt.output}/detect.log')
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
im0 = dp.dmpost(im0,
det,
det_id=detected_img_id,
filename=Path(p).name,
names=names)
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps,
(w, h))
# print(detected_img_id, p, txt_path)
tmp_filename = Path(txt_path).stem
im0 = dp.dmpost(im0,
det,
det_id=detected_img_id,
filename=tmp_filename,
names=names)
vid_writer.write(im0)
detected_img_id += 1
bar.set_description(f'inf_time: {infer_time*1000:.2f}ms {prt_str:<40}')
if save_txt or save_img:
print('Results saved to %s' % out)
if platform == 'darwin' and not opt.update: # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
time_arr = np.array(time_list)
prnt = f'Done. Network mean inference time: {np.mean(time_arr)*1000:.2f}ms, Mean FPS: {1/np.mean(time_arr):.2f}.'
print(f'\nModel size {opt.img_size} inference {prnt}')
os.system(f'echo "{prnt}" >> {opt.output}/detect.log')
os.system(f'echo "useroi {opt.img_size} {prnt}" >> detect2.log')
def detect(model, device, frame, imgsz, iou_thresh, conf_thresh):
weights = 'yolov5l.pt'
# Initialize
set_logging()
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
print("IMGSZ:", imgsz)
if half:
model.half() # to FP16
# Set Dataloader
cudnn.benchmark = True
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
img = [letterbox(frame, new_shape=imgsz, auto=True)[0]]
img = np.stack(img, 0)
# Convert
img = img[:, :, :, ::-1].transpose(0, 3, 1,
2) # BGR to RGB, to bsx3x416x416
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=False)[0]
# Apply NMS
pred = non_max_suppression(pred, conf_thresh, iou_thresh, agnostic=False)
t2 = time_synchronized()
det = pred[0]
s = '%gx%g ' % img.shape[2:]
gn = torch.tensor(img.shape)[[1, 0, 1, 0]]
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
bboxes = []
for *xyxy, conf, cls in reversed(det):
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
label = names[int(cls)]
bboxes.append((xyxy, label))
return bboxes
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://'))
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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 names]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset:
print('THIS IS THE PATH YOU WANT : ' + path)
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
#creating/opening the text doc corresponding to the image
txt_name = 'hbr.txt'
txt_file = open('/static/texts/' + txt_name, 'a')
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(
), dataset.count
else:
p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / 'labels' / p.stem) + (
'' if dataset.mode == 'image' else f'_{frame}') # img.txt
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f'{n} {names[int(c)]}s, ' # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (
cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = f'{names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
if opt.plate:
coords = xyxy
txt_word = recognize_plate(path, coords)
print(txt_word)
txt_file.write(txt_word + ' ')
#ADDING THE RECO ALGO
#"""Here I am tempting to implement the licence plate recognition function from the AI Guy,
#and to apply it to our problematic. This is the V.1"""
# Print time (inference + NMS)
print(f'{s}Done. ({t2 - t1:.3f}s)')
# Stream results
if view_img:
cv2.imshow(str(p), im0)
# Save results (image with detections)
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video'
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps,
(w, h))
vid_writer.write(im0)
txt_file.close()
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
print(f'Done. ({time.time() - t0:.3f}s)')
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://"))
)
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)
) # increment run
(save_dir / "labels" if save_txt else save_dir).mkdir(
parents=True, exist_ok=True
) # make dir
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != "cpu" # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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"]
).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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 names]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img) if device.type != "cpu" else None # run once
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(
pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms,
)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], "%g: " % i, im0s[i].copy(), dataset.count
else:
p, s, im0, frame = path, "", im0s, getattr(dataset, "frame", 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / "labels" / p.stem) + (
"" if dataset.mode == "image" else f"_{frame}"
) # img.txt
s += "%gx%g " % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}s, " # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (
(xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn)
.view(-1)
.tolist()
) # normalized xywh
line = (
(cls, *xywh, conf) if opt.save_conf else (cls, *xywh)
) # label format
with open(txt_path + ".txt", "a") as f:
f.write(("%g " * len(line)).rstrip() % line + "\n")
if save_img or view_img: # Add bbox to image
label = f"{names[int(cls)]} {conf:.2f}"
plot_one_box(
xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3,
)
# Print time (inference + NMS)
print(f"{s}Done. ({t2 - t1:.3f}s)")
# Stream results
if view_img:
cv2.imshow(str(p), im0)
if cv2.waitKey(1) == ord("q"): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == "image":
cv2.imwrite(save_path, im0)
else: # 'video'
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
fourcc = "mp4v" # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)
)
vid_writer.write(im0)
if save_txt or save_img:
s = (
f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}"
if save_txt
else ""
)
print(f"Results saved to {save_dir}{s}")
print(f"Done. ({time.time() - t0:.3f}s)")
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://'))
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
cnt = 0
mat = [0, 0, 0, 0, 0] # can pls gls trsh none
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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 names]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
# run once
_ = model(img.half() if half else img) if device.type != 'cpu' else None
for path, img, im0s, vid_cap in dataset:
if ARD.readable(): # readable을 통해 값을 받을 수 있으면
cur = ARD.readline().decode().strip() # Serial로 받은 한 줄을 읽고
# 받은 값이 1이거나, cnt가 0이 아니면 카메라에 촬영된 한 프레임을 판별함.
# cnt는 0으로 관리하다가 Serial로 1을 받으면 1씩 증가시키고 들어온 물건을 판별했다고 판단되면 다시 cnt를 0으로 만듬
if cur == "1" or cnt != 0:
cnt += 1
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = Path(path[i]), '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = Path(path), '', im0s
save_path = str(save_dir / p.name)
txt_path = str(save_dir / 'labels' / p.stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
# normalization gain whwh
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
# 한 객체가 여러개 판별된 경우 문자로는 한번만 출력하기 위해 unique사용
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# 여기서 c에 현재 프레임에 판별된 클래스가 들어있음
# 판별된 객체마다 mat리스트에 카운트해줌 0인덱스는 can 1인덱스는 pls 2인덱스는 gls 3인덱스는 trsh
mat[int(c)] += 1
else: # det이 None일 때 -> 4(none)인덱스에 1증가
mat[4] += 1
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
# label format
line = (cls, *xywh, conf) if opt.save_conf else (cls,
*xywh)
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
# 프레임에 라벨 씌우기(can 0.7)이런식으로
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
if cnt == 10: # 10프레임의 결과를 가지고
cnt = 0
# can, pls, gls, (trsh + none) 중 젤 많이 나온 것을 serial로 보냄
max_mat = -1
max_value = 0
for i in range(3):
if mat[i] > max_value:
max_value = mat[i]
max_mat = i
if sum(mat[3:5]) > max_value:
max_value = sum(mat[3:5])
max_mat = 3 # trsh, none은 합쳐서 3으로 취급
print(mat)
for i in range(5):
mat[i] = 0
# 여기서 serial로 max_mat 전송. 0을 넘기면 아두이노에서 인식을 못함. 그래서 1을 더해서 넘겨주기로 함
ARD.write(str(max_mat + 1).encode())
print("print", max_mat + 1)
# 물건이 한번 들어왔을 때 10프레임이나 기다리는 건 비효율적이라 생각되어 0~3인덱스 중 하나라도 4번 이상 판별되면
# 해당 번호의 물건으로 분류된다고 확신하고 값을 출력시킴
elif mat[0] >= 4 or mat[1] >= 4 or mat[2] >= 4 or mat[
3] >= 4:
cnt = 0
max_mat = -1
max_value = 0
for i in range(4):
if mat[i] > max_value:
max_value = mat[i]
max_mat = i
for i in range(5):
mat[i] = 0
ARD.write(str(max_mat + 1).encode())
print("print", max_mat + 1)
# Stream results
if view_img:
cv2.imshow("video", im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc),
fps, (w, h))
vid_writer.write(im0)
else:
if view_img:
cv2.imshow("video", im0s[0].copy())
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
def trt_detect(save_img=False):
# yolov3-ssp with evolve
# anchor_nums = 4
# nc = 1
# anchors = np.array([
# [[11, 10], [17, 9], [18, 16], [29, 16]],
# [[34, 28], [48, 24], [59, 33], [46, 64]],
# [[69, 45], [86, 59], [96, 80], [150, 106]]
# ])
# output_shapes = [
# (1, anchor_nums, 80, 80, nc + 5),
# (1, anchor_nums, 40, 40, nc + 5),
# (1, anchor_nums, 20, 20, nc + 5)
# ]
# yolov5s
anchor_nums = 3
nc = 1
anchors = np.array([
[[10, 13], [16, 30], [33, 23]], # P3/8
[[30, 61], [62, 45], [59, 119]], # P4/16
[[116, 90], [156, 198], [373, 326]]
])
strides = np.array([8., 16., 32.])
output_shapes = [
(1, anchor_nums, 60, 80, nc + 5),
(1, anchor_nums, 30, 40, nc + 5),
(1, anchor_nums, 15, 20, nc + 5)
# (1, anchor_nums*60*80, nc + 5),
# (1, anchor_nums*30*40, nc + 5),
# (1, anchor_nums*15*20, nc + 5)
]
source, weights, view_img, save_txt, imgsz = \
opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
save_img = not opt.nosave and not source.endswith('.txt') # save inference images
webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
('rtsp://', 'rtmp://', 'http://', 'https://'))
# Directories
save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# Initialize
set_logging()
stride = int(strides.max()) # model stride
print(f"Loading trt engine!")
# imgsz = check_img_size(imgsz, s=stride) # check img_size
# Set Dataloader
vid_path, vid_writer = None, None
bird_transform, pts = False, None
if opt.plot_move_routes:
bird_transform = True
# Coordinates of chessboard region needs to apply brid transform
pts = np.array([(567, 28), (1458, 60), (1890, 639), (638, 1068)])
center_point = queue.Queue()
if webcam:
view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreamsTrt(source, img_size=imgsz, bird_transform=bird_transform, pts=pts)
else:
dataset = LoadImagesTrt(source, img_size=imgsz, bird_transform=bird_transform, pts=pts)
# Get names and colors
names = ['Robot']
# colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
colors = [[96, 171, 132]]
# Run inference
t0 = time.time()
img_index = 0
# Create an empty picture and draw the route on it, which is independent from the real picture.
route_mask = None
cover_heatmap_accum = None
processor = Processor(weights[0], anchor_nums, nc, anchors, output_shapes, imgsz)
for path, img, im0s, vid_cap in dataset:
if opt.plot_move_routes and route_mask is None:
# initialize something for ploting move routes
route_img = im0s.copy()
route_mask = np.zeros((im0s.shape[0], im0s.shape[1], 3), np.uint8)
route_mask.fill(255)
route_mask_bg_color = 'white'
if opt.cover_heatmap:
if cover_heatmap_accum is None:
# initialize something for ploting cover heatmap
cover_heatmap_img = im0s.copy()
# The template with a black base color is used to
# continuously accumulate the covered elliptical areas going in.
cover_heatmap_accum = np.zeros((im0s.shape[0], im0s.shape[1], 3), np.uint8)
cover_heatmap_accum.fill(0)
cover_heatmap_accum_bg_color = 'black'
# The same img as img0s used to host cumulative heatmaps.
cover_heatmap_accum_img0s = im0s.copy()
# restore cover_heatmap_tmp every loop
cover_heatmap_tmp = np.zeros((im0s.shape[0], im0s.shape[1], 3), np.uint8)
cover_heatmap_tmp.fill(0)
# Inference
# t1 = time_synchronized()
pred = processor.detect(img)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
else:
p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
det[:, 6:8] = scale_coords(img.shape[2:], det[:, 6:8], im0.shape).round()
# plot move routes
if opt.plot_move_routes and img_index % opt.move_routes_interval == 0: # 设置绘制间隔
assert pred[0].shape[0] == 1, "Only one robot can exist in the scene when drawing movement routes!"
center_point.put(list(map(int, det[:, 6:8].tolist()[0])))
# elements of center_point always Less than or equal to 2
if center_point.qsize() <= 1:
pass
else:
pts1 = center_point.get()
pts2 = center_point.get()
plot_move_routes([pts1, pts2], route_mask, colors[int(cls)], 3)
route_img = one_cover_two_with_mask(route_mask, route_img, bg_color=route_mask_bg_color)
cv2.imwrite('/home/yousixia/project/yolov3/runs/detect/tmp/im0s.jpg', im0s)
center_point.put(pts2)
# plot cover heatmap
if opt.cover_heatmap and img_index % opt.cover_heatmap_interval == 0:
x1, y1, x2, y2 = list(map(int, det[:, :4].tolist()[0]))
# Draws a white ellipse with a center at center_point of bbox,
# two axes is (x2-x1, y2-y1), and a line width of 3.
cv2.ellipse(cover_heatmap_tmp, list(map(int, det[:, 6:8].tolist()[0])),
(int((x2 - x1) * 0.9 / 2), int((y2 - y1) * 0.9 / 2)), 0, 0, 360, (32, 16, 16),
-1) # 画椭圆
cover_heatmap_accum = cv2.addWeighted(cover_heatmap_accum, 1, cover_heatmap_tmp, 1, 0) # 累加覆盖面积
cover_heatmap_accum_colormap = cv2.applyColorMap(cover_heatmap_accum, cv2.COLORMAP_JET)
one_cover_two_with_mask(cover_heatmap_accum, cover_heatmap_img, cover_heatmap_accum_colormap,
bg_color='black')
# cv2.imwrite('/home/yousixia/project/yolov3/runs/detect/tmp/original_img.jpg',
# original_img)
# Print results
for c in np.unique(det[:, 5]):
n = (det[:, 5] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
for *xyxy, conf, cls, center_x, center_y in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = f'{names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
# if opt.plot_move_routes:
# plot_one_box(xyxy, route_img, label=label, color=colors[int(cls)], line_thickness=3)
# if opt.cover_heatmap:
# plot_one_box(xyxy, cover_heatmap_img, label=label, color=colors[int(cls)], line_thickness=3)
# cv2.imwrite('/home/yousixia/project/yolov3/runs/detect/tmp/cover_heatmap_img.jpg',
# cover_heatmap_img)
# plot_center_point((center_x, center_y), im0, color=colors[int(cls)], line_thickness=3)
# Print time (inference + NMS)
# print(f'{s}Done. ({t2 - t1:.3f}s)')
# drawing cover rate plot
cover_rate_plot_data = []
if opt.cover_rate:
heatmap_non_zero_pixels = cv2.countNonZero(cover_heatmap_accum)
all_pixels = cv2.countNonZero(im0s)
cover_heatmap_img
cover_rate = 1.0 * 100 * heatmap_non_zero_pixels / all_pixels
secend = img_index / vid_cap.get(cv2.CAP_PROP_FPS)
cover_rate_plot_data.append([cover_rate, secend])
l1 = plt.plot(secend, cover_rate, 'b--', label='覆盖率')
plt.xlabel('时间/s')
plt.ylabel('覆盖率/%')
plt.legend()
# plt.show()
# stack bbox, movement routes, cover heatmap, cover rate to one matrix.
# TODO
# Stream results
if view_img:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections)
if save_img:
# cv2.imwrite('/home/yousixia/project/yolov3/runs/detect/tmp/123.jpg', im0)
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video' or 'stream'
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)
if opt.plot_move_routes:
w, h = im0.shape[1], im0.shape[0]
else:
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
save_path += '.mp4'
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
if opt.plot_move_routes:
filename = save_path.split('/')[-1]
movement_routes_vid_writer = cv2.VideoWriter(
os.path.join('/'.join(save_path.split('/')[:-1]), 'movement_routes_' + filename),
cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
if opt.cover_heatmap:
cover_heatmap_vid_writer = cv2.VideoWriter(
os.path.join('/'.join(save_path.split('/')[:-1]), 'cover_heatmap_' + filename),
cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
vid_writer.write(im0)
movement_routes_vid_writer.write(route_img)
cover_heatmap_vid_writer.write(cover_heatmap_img)
img_index += 1
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
print(f'Done. ({time.time() - t0:.3f}s)')
def run(weights='yolov5s.pt', # model.pt path(s)
source='data/images', # file/dir/URL/glob, 0 for webcam
imgsz=640, # inference size (pixels)
conf_thres=0.25, # confidence threshold
iou_thres=0.45, # NMS IOU threshold
max_det=1000, # maximum detections per image
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=False, # show results
save_txt=False, # save results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_crop=False, # save cropped prediction boxes
nosave=False, # do not save images/videos
classes=None, # filter by class: --class 0, or --class 0 2 3
agnostic_nms=False, # class-agnostic NMS
augment=False, # augmented inference
visualize=False, # visualize features
update=False, # update all models
project='runs/detect', # save results to project/name
name='exp', # save results to project/name
exist_ok=False, # existing project/name ok, do not increment
line_thickness=3, # bounding box thickness (pixels)
hide_labels=False, # hide labels
hide_conf=False, # hide confidences
half=False, # use FP16 half-precision inference
):
save_img = not nosave and not source.endswith('.txt') # save inference images
webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
('rtsp://', 'rtmp://', 'http://', 'https://'))
# Directories
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# Initialize
set_logging()
device = select_device(device)
half &= device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
stride = int(model.stride.max()) # model stride
imgsz = check_img_size(imgsz, s=stride) # check image size
names = model.module.names if hasattr(model, 'module') else model.names # get class names
if half:
model.half() # to FP16
# Second-stage classifier
classify = False
if classify:
modelc = load_classifier(name='resnet50', n=2) # initialize
modelc.load_state_dict(torch.load('resnet50.pt', map_location=device)['model']).to(device).eval()
# Dataloader
if webcam:
view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride)
bs = len(dataset) # batch_size
else:
dataset = LoadImages(source, img_size=imgsz, stride=stride)
bs = 1 # batch_size
vid_path, vid_writer = [None] * bs, [None] * bs
# Run inference
if device.type != 'cpu':
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
t0 = time.time()
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img,
augment=augment,
visualize=increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False)[0]
# Apply NMS
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], f'{i}: ', im0s[i].copy(), dataset.count
else:
p, s, im0, frame = path, '', im0s.copy(), getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
imc = im0.copy() if save_crop else im0 # for save_crop
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or save_crop or view_img: # Add bbox to image
c = int(cls) # integer class
label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_thickness=line_thickness)
if save_crop:
save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
# Print time (inference + NMS)
print(f'{s}Done. ({t2 - t1:.3f}s)')
# Stream results
if view_img:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections)
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video' or 'stream'
if vid_path[i] != save_path: # new video
vid_path[i] = save_path
if isinstance(vid_writer[i], cv2.VideoWriter):
vid_writer[i].release() # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
save_path += '.mp4'
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
vid_writer[i].write(im0)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
if update:
strip_optimizer(weights) # update model (to fix SourceChangeWarning)
print(f'Done. ({time.time() - t0:.3f}s)')
def detect_one(model, image_path, device):
# Load model
img_size = 640
conf_thres = 0.3
iou_thres = 0.5
orgimg = cv2.imread(image_path) # BGR
img0 = copy.deepcopy(orgimg)
assert orgimg is not None, 'Image Not Found ' + image_path
h0, w0 = orgimg.shape[:2] # orig hw
r = img_size / max(h0, w0) # resize image to img_size
if r != 1: # always resize down, only resize up if training with augmentation
interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR
img0 = cv2.resize(img0, (int(w0 * r), int(h0 * r)),
interpolation=interp)
imgsz = check_img_size(img_size, s=model.stride.max()) # check img_size
img = letterbox(img0, new_shape=imgsz)[0]
# Convert
img = img[:, :, ::-1].transpose(2, 0, 1).copy() # BGR to RGB, to 3x416x416
# Run inference
t0 = time.time()
img = torch.from_numpy(img).to(device)
img = img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img)[0]
# Apply NMS
pred = non_max_suppression_face(pred, conf_thres, iou_thres)
print('pred: ', pred)
t2 = time_synchronized()
print('img.shape: ', img.shape)
print('orgimg.shape: ', orgimg.shape)
# Process detections
for i, det in enumerate(pred): # detections per image
gn = torch.tensor(orgimg.shape)[[1, 0, 1, 0]].to(
device) # normalization gain whwh
gn_lks = torch.tensor(orgimg.shape)[[1, 0, 1, 0, 1, 0, 1, 0, 1, 0]].to(
device) # normalization gain landmarks
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
orgimg.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
det[:, 5:15] = scale_coords_landmarks(img.shape[2:], det[:, 5:15],
orgimg.shape).round()
for j in range(det.size()[0]):
xywh = (xyxy2xywh(torch.tensor(det[j, :4]).view(1, 4)) /
gn).view(-1).tolist()
conf = det[j, 4].cpu().numpy()
landmarks = (det[j, 5:15].view(1, 10) /
gn_lks).view(-1).tolist()
class_num = det[j, 15].cpu().numpy()
orgimg = show_results(orgimg, xywh, conf, landmarks, class_num)
# Stream results
print(f'Done. ({time.time() - t0:.3f}s)')
cv2.imshow('orgimg', orgimg)
if cv2.waitKey(0) == ord('q'): # q to quit
raise StopIteration
def detect(save_img=False):
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# Initialize
device = select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
txt_path = str(Path(out) / Path(p).stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
for *xyxy, conf, cls in det:
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * 5 + '\n') %
(cls, *xywh)) # label format
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps,
(w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin' and not opt.update: # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
def detect(save_img=False):
atama = 0
flag1 = False
flag2 = False
flag3 = 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://'))
# Directories
save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
stride = int(model.stride.max()) # model stride
imgsz = check_img_size(imgsz, s=stride) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz, stride=stride)
# 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 names]
# Run inference
if device.type != 'cpu':
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
t0 = time.time()
for path, img, im0s, vid_cap in dataset: # her frame burda dönüyor
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
else:
p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
sonuc = Tespit()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
for *xyxy, conf, cls in reversed(det): # xyxy koordinatlar imiş
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh)
# label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = f'{names[int(cls)]} {conf:.2f}'
(plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3))
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()
ab = torch.tensor(xyxy).view(1,4)[0]
ab = ab.numpy()
isim = f'{names[int(cls)]}'
sonuc.label_list.append(isim)
sonuc.koordinat_list.append(ab)
#sol cıkıs ve kalıp label bulma
i = 0
cıkıs_list = []
while i < len(sonuc.label_list):
if('CIKIS' == sonuc.label_list[i]):
cıkıs_list.append(sonuc.koordinat_list[i])
if('KALIP' == sonuc.label_list[i]):
kalıp = sonuc.koordinat_list[i]
i += 1
if(len(cıkıs_list) == 2):
if(cıkıs_list[0][0] < cıkıs_list[1][0]):
sol_cıkıs = cıkıs_list[0]
else:
sol_cıkıs = cıkıs_list[1]
elif(len(cıkıs_list) == 1):
sol_cıkıs = cıkıs_list[0]
else:
sol_cıkıs = [0,0,0,0]
sol_alan = (sol_cıkıs[2] - sol_cıkıs[0]) * (sol_cıkıs[3] - sol_cıkıs[1])
cv2.putText(im0, "SOL ALAN" + str(sol_alan), (800, 500), cv2.FONT_HERSHEY_SIMPLEX, 1,
(209, 80, 0, 255), 3)
try:
x1 = kalıp[0]
x2 = kalıp[2]
y1 = kalıp[1]
y2 = kalıp[3]
print("x1 degeri : " , x1 , " x2 degeri : " , x2 , " y1 degeri : " , y1 , " y2 degeri : " , y2)
if ((640 < x1 < 675) and (855 < x2 < 874) and (290 < y1 < 305) and (835 < y2 < 855)) or flag1:
flag1 = True
print("butun kosullar saglandı")
cv2.putText(im0, "UYGUN KONUM" ,(200,200), cv2.FONT_HERSHEY_SIMPLEX, 1, (209,80, 0 ,255), 3)
print("sol kapak alanı : " , str(sol_alan))
if(2000 < sol_alan < 2400) or flag2:
# print("cıkmıs")
flag2 = True
if (900 < sol_alan < 1200) or flag3:
# print("en aşşa indi")
flag3 = True
if (2000 < sol_alan):
if 'BOS' in sonuc.label_list:
cv2.putText(im0, "SIKINTI YOK", (400, 400), cv2.FONT_HERSHEY_SIMPLEX, 1,
(209, 80, 0, 255), 3)
print("******************sıkıntı yok***********************")
else:
cv2.putText(im0, "KALIP DUSMEDI", (400, 400), cv2.FONT_HERSHEY_SIMPLEX, 1,
(209, 80, 0, 255), 3)
print("bunu bi şekilde halletmemiz gerek")
time.sleep(1)
#print("flag yazdırdık: "+ str(flag1)+ " "+ str(flag2)+ " "+ str(flag3))
else:
print("kosullar saglanmadı")
if ((sol_cıkıs[0] > 800) and flag3):
flag1 = False
flag2 = False
flag3 = False
except IndexError:
pass
# Print time (inference + NMS)
print(f'{s}Done. ({t2 - t1:.3f}s)')
atama += (t2 - t1)
print("toplam zaman", str(atama))
# Stream results
if True: # önceden view_img idi, şimdi True oldu yani resimleri video gibi oynatma sağlandı
cv2.imshow("Result", im0) # farklı isim olursa ayrı pencerelerde açılır, aynı isimle aynı pencerede açar
cv2.waitKey(1) # 1 millisecond - 0 girilir ise oynaması için imagein input bekler -- 1 kalması yeterli bizim için
# Save results (image with detections)
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video'
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
vid_writer.write(im0)
#print("gecen zaman :" , str(t2-t1))
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
print(f'Done. ({time.time() - t0:.3f}s)')
def test(
data,
weights=None,
batch_size=16,
imgsz=640,
conf_thres=0.001,
iou_thres=0.6, # for NMS
save_json=False,
single_cls=False,
augment=False,
verbose=False,
model=None,
dataloader=None,
save_dir='',
merge=False,
save_txt=False):
# Initialize/load model and set device
training = model is not None
if training: # called by train.py
device = next(model.parameters()).device # get model device
else: # called directly
set_logging()
device = select_device(opt.device, batch_size=batch_size)
merge, save_txt = opt.merge, opt.save_txt # use Merge NMS, save *.txt labels
if save_txt:
out = Path('inference/output')
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Remove previous
for f in glob.glob(str(Path(save_dir) / 'test_batch*.jpg')):
os.remove(f)
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
# model = nn.DataParallel(model)
# Half
half = device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half()
# Configure
model.eval()
with open(data) as f:
data = yaml.load(f, Loader=yaml.FullLoader) # model dict
check_dataset(data) # check
nc = 1 if single_cls else int(data['nc']) # number of classes
iouv = torch.linspace(0.5, 0.95,
10).to(device) # iou vector for [email protected]:0.95
niou = iouv.numel()
# Dataloader
if not training:
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
path = data['test'] if opt.task == 'test' else data[
'val'] # path to val/test images
dataloader = create_dataloader(path,
imgsz,
batch_size,
model.stride.max(),
opt,
hyp=None,
augment=False,
cache=False,
pad=0.5,
rect=True)[0]
seen = 0
names = model.names if hasattr(model, 'names') else model.module.names
coco91class = coco80_to_coco91_class()
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R',
'[email protected]', '[email protected]:.95')
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class = [], [], [], []
for batch_i, (img, targets, paths,
shapes) in enumerate(tqdm(dataloader, desc=s)):
img = img.to(device, non_blocking=True)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
nb, _, height, width = img.shape # batch size, channels, height, width
whwh = torch.Tensor([width, height, width, height]).to(device)
# Disable gradients
with torch.no_grad():
# Run model
t = time_synchronized()
inf_out, train_out = model(
img, augment=augment) # inference and training outputs
t0 += time_synchronized() - t
# Compute loss
if training: # if model has loss hyperparameters
loss += compute_loss([x.float() for x in train_out], targets,
model)[1][:3] # GIoU, obj, cls
# Run NMS
t = time_synchronized()
output = non_max_suppression(inf_out,
conf_thres=conf_thres,
iou_thres=iou_thres,
merge=merge)
t1 += time_synchronized() - t
# Statistics per image
for si, pred in enumerate(output):
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class
seen += 1
if pred is None:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool),
torch.Tensor(), torch.Tensor(), tcls))
continue
# Append to text file
if save_txt:
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0
]] # normalization gain whwh
x = pred.clone()
x[:, :4] = scale_coords(img[si].shape[1:], x[:, :4],
shapes[si][0],
shapes[si][1]) # to original
for *xyxy, conf, cls in x:
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(str(out / Path(paths[si]).stem) + '.txt',
'a') as f:
f.write(
('%g ' * 5 + '\n') % (cls, *xywh)) # label format
# Clip boxes to image bounds
clip_coords(pred, (height, width))
# Append to pycocotools JSON dictionary
if save_json:
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
image_id = Path(paths[si]).stem
box = pred[:, :4].clone() # xyxy
scale_coords(img[si].shape[1:], box, shapes[si][0],
shapes[si][1]) # to original shape
box = xyxy2xywh(box) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for p, b in zip(pred.tolist(), box.tolist()):
jdict.append({
'image_id':
int(image_id) if image_id.isnumeric() else image_id,
'category_id':
coco91class[int(p[5])],
'bbox': [round(x, 3) for x in b],
'score':
round(p[4], 5)
})
# Assign all predictions as incorrect
correct = torch.zeros(pred.shape[0],
niou,
dtype=torch.bool,
device=device)
if nl:
detected = [] # target indices
tcls_tensor = labels[:, 0]
# target boxes
tbox = xywh2xyxy(labels[:, 1:5]) * whwh
# Per target class
for cls in torch.unique(tcls_tensor):
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(
-1) # prediction indices
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(
-1) # target indices
# Search for detections
if pi.shape[0]:
# Prediction to target ious
ious, i = box_iou(pred[pi, :4], tbox[ti]).max(
1) # best ious, indices
# Append detections
detected_set = set()
for j in (ious > iouv[0]).nonzero(as_tuple=False):
d = ti[i[j]] # detected target
if d.item() not in detected_set:
detected_set.add(d.item())
detected.append(d)
correct[
pi[j]] = ious[j] > iouv # iou_thres is 1xn
if len(
detected
) == nl: # all targets already located in image
break
# Append statistics (correct, conf, pcls, tcls)
stats.append(
(correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
# Plot images
if batch_i < 10:
f = Path(save_dir) / ('test_batch%g_gt.jpg' % batch_i) # filename
plot_images(img, targets, paths, str(f), names) # ground truth
f = Path(save_dir) / ('test_batch%g_pred.jpg' % batch_i)
plot_images(img, output_to_target(output, width, height), paths,
str(f), names) # predictions
# Compute statistics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
if len(stats) and stats[0].any():
p, r, ap, f1, ap_class = ap_per_class(*stats)
p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(
1) # [P, R, [email protected], [email protected]:0.95]
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
nt = np.bincount(stats[3].astype(np.int64),
minlength=nc) # number of targets per class
else:
nt = torch.zeros(1)
# Print results
pf = '%20s' + '%12.3g' * 6 # print format
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
# Print results per class
if verbose and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
# Print speeds
t = tuple(x / seen * 1E3
for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
if not training:
print(
'Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g'
% t)
# Save JSON
if save_json and len(jdict):
f = '%s_results.json' % \
(weights.split(os.sep)[-1].replace('.pt', '') if isinstance(weights, str) else '') # filename
print('\nCOCO mAP with pycocotools... saving %s...' % f)
with open(f, 'w') as file:
json.dump(jdict, file)
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files]
cocoGt = COCO(
glob.glob('../coco/annotations/instances_val*.json')
[0]) # initialize COCO ground truth api
cocoDt = cocoGt.loadRes(f) # initialize COCO pred api
cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
cocoEval.params.imgIds = imgIds # image IDs to evaluate
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
map, map50 = cocoEval.stats[:
2] # update results ([email protected]:0.95, [email protected])
except Exception as e:
print('ERROR: pycocotools unable to run: %s' % e)
# Return results
model.float() # for training
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
return (mp, mr, map50, map,
*(loss.cpu() / len(dataloader)).tolist()), maps, t
def detect(save_img=False):
out, source, weights, view_img, save_txt, imgsz, conn = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, opt.conn
webcam = source.isnumeric() or source.startswith(
('rtsp://', 'rtmp://', 'http://')) or source.endswith('.txt')
# Initialize
set_logging()
device = select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
#Parameter definition
print(conn)
angle = 0
rho = 0
saw_ball = 0
theta = 0
theta_increment = 3 * math.pi / 4
increment = math.pi / 48
if conn:
serial_port_1 = serial.Serial(
port="/dev/ttyACM0",
baudrate=115200,
bytesize=serial.EIGHTBITS,
parity=serial.PARITY_NONE,
stopbits=serial.STOPBITS_ONE,
)
# Wait a second to let the port initialize
serial_port = utils.openContinuous(timeout=0.01)
time.sleep(1)
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
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()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# 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))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
hasBall = 0
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
#Add deadzone to the image
cv2.rectangle(im0, (round(im0.shape[1] * 0.45), 0),
(round(im0.shape[1] * 0.55), im0.shape[0]),
(0, 0, 255),
thickness=3,
lineType=cv2.LINE_AA)
save_path = str(Path(out) / Path(p).name)
txt_path = str(Path(out) / Path(p).stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * 5 + '\n') %
(cls, *xywh)) # label format
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
#Get Object Coordinates
mid_x = xywh[0]
mid_y = xywh[1]
box_w = xywh[2]
box_h = xywh[3]
#print mid coordinates of the box
print('Mid coordinates of box,w, h: %.3f,%.3f,%.3f,%.3f' %
(mid_x, mid_y, box_w, box_h))
if (mid_x < 0.45):
#ROBOT TURN LEFT
angle = angle + increment #rad/s
print('LEFT %.4f', angle)
rho = 0
angularControl = 0 #angular velocity control
elif (mid_x > 0.55):
#ROBOT TURN RIGHT
angle = angle - increment #rad/s
print('RIGHT %.4f', angle)
rho = 0
angularControl = 0 #angular velocity control
else:
#DRIVE STRAIGHT TO THE BALL
angle = angle + 0 #rad/s
print('STRAIGHT %.4f', angle)
rho = rho + 0.1 #m/s
#if saw_ball:
#theta = theta #----> INSERT VALUE FOR THE ROBOT TO MOVE IN THE DIRECTION OF THE CAMERA
angularControl = 0 #angular velocity control
saw_ball = 1
else:
if not (hasBall):
angle = angle + increment #rad/s
print('NOBALL %.4f', angle)
rho = 0
angularControl = 0 #angular velocity control
else:
#STOP THE ROBOT
angle = angle + 0 #rad/s
rho = 0
angularControl = 0 #angular velocity control
if conn:
#Constructing the packet
cmd = rem.ffi.new("RobotCommand*")
cmd.header = rem.lib.PACKET_TYPE_ROBOT_COMMAND
cmd.id = 3
#check value of angle
if angle > math.pi * 2:
angle = angle - math.pi * 2
if angle < -math.pi * 2:
angle = angle + math.pi * 2
cmd.angle = angle
#cmd.rho = rho
#cmd.theta = theta
#cmd. angularControl = angularControl
packet = rem.ffi.new("RobotCommandPayload*")
rem.lib.encodeRobotCommand(packet, cmd)
#Sending the packet
try:
serial_port.write(packet.payload)
print(" angle : %.4f" % cmd.angle)
except KeyboardInterrupt:
print("Exiting Program")
except Exception as exception_error:
print("Error occurred. Exiting Program")
print("Error: " + str(exception_error))
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
cv2.destroyAllWindows()
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*fourcc), fps,
(w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % Path(out))
if platform.system() == 'Darwin' and not opt.update: # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
serial_port.close()
