152 lines
6.1 KiB
Python
152 lines
6.1 KiB
Python
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
|
"""
|
|
Run YOLOv5 detection inference on images, videos, directories, globs, YouTube, webcam, streams, etc.
|
|
|
|
Usage - sources:
|
|
$ python detect.py --weights yolov5s.pt --source 0 # webcam
|
|
img.jpg # image
|
|
vid.mp4 # video
|
|
screen # screenshot
|
|
path/ # directory
|
|
list.txt # list of images
|
|
list.streams # list of streams
|
|
'path/*.jpg' # glob
|
|
'https://youtu.be/Zgi9g1ksQHc' # YouTube
|
|
'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
|
|
|
|
Usage - formats:
|
|
$ python detect.py --weights yolov5s.pt # PyTorch
|
|
yolov5s.torchscript # TorchScript
|
|
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
|
yolov5s_openvino_model # OpenVINO
|
|
yolov5s.engine # TensorRT
|
|
yolov5s.mlmodel # CoreML (macOS-only)
|
|
yolov5s_saved_model # TensorFlow SavedModel
|
|
yolov5s.pb # TensorFlow GraphDef
|
|
yolov5s.tflite # TensorFlow Lite
|
|
yolov5s_edgetpu.tflite # TensorFlow Edge TPU
|
|
yolov5s_paddle_model # PaddlePaddle
|
|
"""
|
|
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
FILE = Path(__file__).resolve()
|
|
ROOT = FILE.parents[0] # YOLOv5 root directory
|
|
if str(ROOT) not in sys.path:
|
|
sys.path.append(str(ROOT)) # add ROOT to PATH
|
|
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
|
|
|
from models.common import DetectMultiBackend
|
|
from utils.general import (LOGGER, Profile, check_img_size, cv2,
|
|
non_max_suppression, scale_boxes)
|
|
from utils.augmentations import (letterbox)
|
|
from utils.torch_utils import select_device, smart_inference_mode
|
|
|
|
|
|
@smart_inference_mode()
|
|
def run(
|
|
model = '', # model path or triton URL
|
|
imagePath ='', # file/dir/URL/glob/screen/0(webcam)
|
|
imgsz=(512, 512), # inference size (height, width)
|
|
conf_thres=0.25, # confidence threshold
|
|
iou_thres=0.45, # NMS IOU threshold
|
|
max_det=1, # maximum detections per image
|
|
save_crop=True, # save cropped prediction boxes
|
|
save_dir='data/QR2023_roi/images/', # 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
|
|
):
|
|
|
|
if save_crop:
|
|
# Directories
|
|
Path(save_dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
stride, names, pt = model.stride, model.names, model.pt
|
|
imgsz = check_img_size(imgsz, s=stride) # check image size
|
|
im0 = cv2.imread(imagePath) # BGR
|
|
im = letterbox(im0, imgsz, stride=32, auto=True)[0] # padded resize
|
|
im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
|
|
im = np.ascontiguousarray(im) # contiguous
|
|
p = imagePath
|
|
|
|
bs = 1
|
|
# Run inference
|
|
model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz)) # warmup
|
|
seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
|
|
|
|
with dt[0]:
|
|
im = torch.from_numpy(im).to(model.device)
|
|
im = im.half() if model.fp16 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
|
|
|
|
# Inference
|
|
with dt[1]:
|
|
visualize = False
|
|
pred = model(im, augment=augment, visualize=visualize)
|
|
|
|
# NMS
|
|
with dt[2]:
|
|
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
|
|
|
|
|
|
# Process predictions
|
|
for i, det in enumerate(pred): # per image
|
|
|
|
p = Path(p)
|
|
save_path = os.path.join(save_dir, p.name)
|
|
|
|
if len(det):
|
|
# Rescale boxes from img_size to im0 size
|
|
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
|
|
|
|
# Write results
|
|
a = det[:, :4]
|
|
b = det[:,4:5]
|
|
for *xyxy, conf, cls in reversed(det):
|
|
|
|
x_min, y_min, x_max, y_max = xyxy[:4]
|
|
x_min, y_min, x_max, y_max = int(x_min), int(y_min), int(x_max), int(y_max)
|
|
quarter_width = (x_max - x_min) // 2
|
|
quarter_height = (y_max - y_min) // 2
|
|
|
|
# Save results (image with detections)
|
|
if save_crop:
|
|
# 以左上顶点坐标为原点截图4分之一原图
|
|
# Convert im0 (NumPy array) to PIL image
|
|
im0 = Image.fromarray(np.uint8(im0))
|
|
cropped_im = im0.crop((x_min, y_min, x_min + quarter_width, y_min + quarter_height))
|
|
cropped_im.save(save_path)
|
|
|
|
# Print time (inference-only)
|
|
LOGGER.info(f"{p}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# check_requirements(exclude=('tensorboard', 'thop'))
|
|
meta_info = '/project/dataset/QR2023/terminal-box/meta_info_big_box_terminal.txt'
|
|
with open(meta_info) as fin:
|
|
paths = [line.strip() for line in fin]
|
|
data_len = len(paths)
|
|
|
|
|
|
# Load model
|
|
weights = '/project/yolov5-qr/runs/train_QR/exp10/weights/qr_roi_cloud_detect_20230831.pt'
|
|
device = '4'
|
|
device = select_device(device)
|
|
qrbox_model = DetectMultiBackend(weights, device=device)
|
|
|
|
|
|
for index in range(0, data_len):
|
|
imagePath = paths[index]
|
|
run(model=qrbox_model, imagePath=imagePath)
|