1- import argparse
2- import os
3- import platform
4- import sys
5- from pathlib import Path
6-
7- import torch
8-
9- # YOLOv5 root directory
10- ROOT = Path (os .path .abspath (os .path .join (os .path .dirname (__file__ ), '..' , 'yolov5' )))
11- if str (ROOT ) not in sys .path :
12- sys .path .append (str (ROOT )) # add ROOT to PATH
13- ROOT = Path (os .path .relpath (ROOT , Path .cwd ())) # relative
14-
15- from yolov5 .models .common import DetectMultiBackend
16- from yolov5 .utils .dataloaders import IMG_FORMATS , VID_FORMATS , LoadImages , LoadScreenshots , LoadStreams
17- from yolov5 .utils .general import (LOGGER , Profile , check_file , check_img_size , check_imshow , check_requirements , colorstr , cv2 ,
18- increment_path , non_max_suppression , print_args , scale_boxes , strip_optimizer , xyxy2xywh )
19- from yolov5 .utils .plots import Annotator , colors , save_one_box
20- from yolov5 .utils .torch_utils import select_device , smart_inference_mode
21-
22- @smart_inference_mode ()
23- def run (
24- weights = ROOT / 'yolov5s.pt' , # model path or triton URL
25- source = ROOT / 'data/images' , # file/dir/URL/glob/screen/0(webcam)
26- data = ROOT / 'data/coco128.yaml' , # dataset.yaml path
27- imgsz = (640 , 640 ), # inference size (height, width)
28- conf_thres = 0.25 , # confidence threshold
29- iou_thres = 0.45 , # NMS IOU threshold
30- max_det = 1000 , # maximum detections per image
31- device = '' , # cuda device, i.e. 0 or 0,1,2,3 or cpu
32- view_img = False , # show results
33- save_txt = False , # save results to *.txt
34- save_conf = False , # save confidences in --save-txt labels
35- save_crop = False , # save cropped prediction boxes
36- nosave = False , # do not save images/videos
37- # 滤除暂时不分类的零件类别
38- classes = None , # filter by class: --class 0, or --class 0 2 3
39- agnostic_nms = False , # class-agnostic NMS
40- augment = False , # augmented inference
41- visualize = False , # visualize features
42- update = False , # update all models
43- project = ROOT / 'runs/detect' , # save results to project/name
44- name = 'exp' , # save results to project/name
45- exist_ok = False , # existing project/name ok, do not increment
46- line_thickness = 3 , # bounding box thickness (pixels)
47- hide_labels = False , # hide labels
48- hide_conf = False , # hide confidences
49- half = False , # use FP16 half-precision inference
50- dnn = False , # use OpenCV DNN for ONNX inference
51- vid_stride = 1 , # video frame-rate stride
52- ):
53- source = str (source )
54- save_img = not nosave and not source .endswith ('.txt' ) # save inference images
55- is_file = Path (source ).suffix [1 :] in (IMG_FORMATS + VID_FORMATS )
56- is_url = source .lower ().startswith (('rtsp://' , 'rtmp://' , 'http://' , 'https://' ))
57- webcam = source .isnumeric () or source .endswith ('.streams' ) or (is_url and not is_file )
58- screenshot = source .lower ().startswith ('screen' )
59- if is_url and is_file :
60- source = check_file (source ) # download
61-
62- # Directories
63- save_dir = increment_path (Path (project ) / name , exist_ok = exist_ok ) # increment run
64- (save_dir / 'labels' if save_txt else save_dir ).mkdir (parents = True , exist_ok = True ) # make dir
65-
66- # Load model
67- device = select_device (device )
68- model = DetectMultiBackend (weights , device = device , dnn = dnn , data = data , fp16 = half )
69- stride , names , pt = model .stride , model .names , model .pt
70- imgsz = check_img_size (imgsz , s = stride ) # check image size
71-
72- # Dataloader
73- bs = 1 # batch_size
74- # 摄像头
75- if webcam :
76- view_img = check_imshow (warn = True )
77- dataset = LoadStreams (source , img_size = imgsz , stride = stride , auto = pt , vid_stride = vid_stride )
78- bs = len (dataset )
79- elif screenshot :
80- dataset = LoadScreenshots (source , img_size = imgsz , stride = stride , auto = pt )
81- else :
82- dataset = LoadImages (source , img_size = imgsz , stride = stride , auto = pt , vid_stride = vid_stride )
83- vid_path , vid_writer = [None ] * bs , [None ] * bs
84-
85- # Run inference
86- model .warmup (imgsz = (1 if pt or model .triton else bs , 3 , * imgsz )) # warmup
87- seen , windows , dt = 0 , [], (Profile (), Profile (), Profile ())
88- for path , im , im0s , vid_cap , s in dataset :
89- with dt [0 ]:
90- im = torch .from_numpy (im ).to (model .device )
91- im = im .half () if model .fp16 else im .float () # uint8 to fp16/32
92- im /= 255 # 0 - 255 to 0.0 - 1.0
93- if len (im .shape ) == 3 :
94- im = im [None ] # expand for batch dim
95-
96- # Inference
97- with dt [1 ]:
98- visualize = increment_path (save_dir / Path (path ).stem , mkdir = True ) if visualize else False
99- pred = model (im , augment = augment , visualize = visualize )
100-
101- # NMS
102- with dt [2 ]:
103- pred = non_max_suppression (pred , conf_thres , iou_thres , classes , agnostic_nms , max_det = max_det )
104-
105- # Second-stage classifier (optional)
106- # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
107-
108- # Process predictions
109- for i , det in enumerate (pred ): # per image
110- seen += 1
111- if webcam : # batch_size >= 1
112- p , im0 , frame = path [i ], im0s [i ].copy (), dataset .count
113- s += f'{ i }
114- else :
115- p , im0 , frame = path , im0s .copy (), getattr (dataset , 'frame' , 0 )
116-
117- p = Path (p ) # to Path
118- save_path = str (save_dir / p .name ) # im.jpg
119- txt_path = str (save_dir / 'labels' / p .stem ) + ('' if dataset .mode == 'image' else f'_{ frame } ) # im.txt
120- s += '%gx%g ' % im .shape [2 :] # print string
121- gn = torch .tensor (im0 .shape )[[1 , 0 , 1 , 0 ]] # normalization gain whwh
122- imc = im0 .copy () if save_crop else im0 # for save_crop
123- annotator = Annotator (im0 , line_width = line_thickness , example = str (names ))
124- if len (det ):
125- # Rescale boxes from img_size to im0 size
126- det [:, :4 ] = scale_boxes (im .shape [2 :], det [:, :4 ], im0 .shape ).round ()
127-
128- # Print results
129- for c in det [:, 5 ].unique ():
130- n = (det [:, 5 ] == c ).sum () # detections per class
131- s += f"{ n } { names [int (c )]} { 's' * (n > 1 )} # add to string
132-
133- # Write results
134- for * xyxy , conf , cls in reversed (det ):
135- if save_txt : # Write to file
136- xywh = (xyxy2xywh (torch .tensor (xyxy ).view (1 , 4 )) / gn ).view (- 1 ).tolist () # normalized xywh
137- line = (cls , * xywh , conf ) if save_conf else (cls , * xywh ) # label format
138- with open (f'{ txt_path } , 'a' ) as f :
139- f .write (('%g ' * len (line )).rstrip () % line + '\n ' )
140-
141- if save_img or save_crop or view_img : # Add bbox to image
142- c = int (cls ) # integer class
143- label = None if hide_labels else (names [c ] if hide_conf else f'{ names [c ]} { conf :.2f} )
144- annotator .box_label (xyxy , label , color = colors (c , True ))
145- if save_crop :
146- save_one_box (xyxy , imc , file = save_dir / 'crops' / names [c ] / f'{ p .stem } , BGR = True )
147-
148- # Stream results
149- im0 = annotator .result ()
150- if view_img :
151- if platform .system () == 'Linux' and p not in windows :
152- windows .append (p )
153- cv2 .namedWindow (str (p ), cv2 .WINDOW_NORMAL | cv2 .WINDOW_KEEPRATIO ) # allow window resize (Linux)
154- cv2 .resizeWindow (str (p ), im0 .shape [1 ], im0 .shape [0 ])
155- cv2 .imshow (str (p ), im0 )
156- cv2 .waitKey (1 ) # 1 millisecond
157-
158- # Save results (image with detections)
159- if save_img :
160- if dataset .mode == 'image' :
161- cv2 .imwrite (save_path , im0 )
162- else : # 'video' or 'stream'
163- if vid_path [i ] != save_path : # new video
164- vid_path [i ] = save_path
165- if isinstance (vid_writer [i ], cv2 .VideoWriter ):
166- vid_writer [i ].release () # release previous video writer
167- if vid_cap : # video
168- fps = vid_cap .get (cv2 .CAP_PROP_FPS )
169- w = int (vid_cap .get (cv2 .CAP_PROP_FRAME_WIDTH ))
170- h = int (vid_cap .get (cv2 .CAP_PROP_FRAME_HEIGHT ))
171- else : # stream
172- fps , w , h = 30 , im0 .shape [1 ], im0 .shape [0 ]
173- save_path = str (Path (save_path ).with_suffix ('.mp4' )) # force *.mp4 suffix on results videos
174- vid_writer [i ] = cv2 .VideoWriter (save_path , cv2 .VideoWriter_fourcc (* 'mp4v' ), fps , (w , h ))
175- vid_writer [i ].write (im0 )
176-
177- # Print time (inference-only)
178- LOGGER .info (f"{ s } { '' if len (det ) else '(no detections), ' } { dt [1 ].dt * 1E3 :.1f} )
179-
180- # Print results
181- t = tuple (x .t / seen * 1E3 for x in dt ) # speeds per image
182- LOGGER .info (f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape { (1 , 3 , * imgsz )} % t )
183- if save_txt or save_img :
184- s = f"\n { len (list (save_dir .glob ('labels/*.txt' )))} { save_dir / 'labels' } if save_txt else ''
185- LOGGER .info (f"Results saved to { colorstr ('bold' , save_dir )} { s } )
186- if update :
187- strip_optimizer (weights [0 ]) # update model (to fix SourceChangeWarning)
1+ """Perform test request"""
2+ import cv2
3+ import requests
4+
5+ class Detect :
6+ def __init__ (self ):
7+ self .detection_url = "http://localhost:5000/v1/object-detection/yolov5"
8+
9+ def detect (self , image ):
10+ # 将OpenCV的Mat转换为JPEG格式
11+ _ , img_encoded = cv2 .imencode (".jpg" , image )
12+ image_bytes = img_encoded .tobytes ()
13+
14+ # 发送POST请求并获取响应
15+ response = requests .post (self .detection_url , files = {"image" : image_bytes }).json ()
16+
17+ return response
18+
19+ # # 解析响应并返回结果
20+ # results = []
21+ # for obj in response["predictions"]:
22+ # class_name = obj["label"]
23+ # confidence = obj["confidence"]
24+ # bbox = obj["bbox"]
25+ # x1, y1, x2, y2 = bbox
26+ # results.append({
27+ # "class": class_name,
28+ # "confidence": confidence,
29+ # "bbox": [x1, y1, x2, y2]
30+ # })
31+ # return results
0 commit comments