Views
No views yet
1import argparse
2import os
3from pathlib import Path
4import torch
5
6from models.common import DetectMultiBackend
7from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
8from utils.general import LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2, increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh
9from utils.plots import Annotator, colors, save_one_box, save_block_box
10from utils.torch_utils import select_device, smart_inference_mode
11
12def load_model(weights, device, dnn, data, fp16):
13 device = select_device(device)
14 model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=fp16)
15 return model
16
17def run_single_image_inference(model, img_path, stride, names, pt, conf_thres=0.35, iou_thres=0.7, max_det=100, augment=True, visualize=False, line_thickness=1, hide_labels=False, hide_conf=False, save_conf=False, save_crop=False, save_block=True, imgsz=(640, 640), vid_stride=1, bs=1, classes=None, agnostic_nms=False, save_txt=True, save_img=True):
18 dataset = LoadImages(img_path, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride) # Load image from file
19 imgsz = check_img_size(imgsz, s=stride)
20
21 # Run inference
22 model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz)) # warmup
23 seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
24 for path, im, im0s, vid_cap, s in dataset:
25 with dt[0]:
26 im = torch.from_numpy(im).to(model.device)
27 im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
28 im /= 255 # 0 - 255 to 0.0 - 1.0
29 if len(im.shape) == 3:
30 im = im[None] # expand for batch dim
31
32 # Inference
33 with dt[1]:
34 visualize = False
35 pred = model(im, augment=augment, visualize=visualize)
36
37 # NMS
38 with dt[2]:
39 pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
40
41 # Second-stage classifier (optional)
42 # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
43
44 # Process predictions
45 sorted_data_list = []
46
47 # Process predictions
48 for i, det in enumerate(pred): # per image
49 seen += 1
50 p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
51
52 p = Path(p) # to Path
53 s += '%gx%g ' % im.shape[2:] # print string
54 gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
55 imc = im0.copy() if save_crop or save_block else im0 # for save_crop
56 annotator = Annotator(im0, line_width=line_thickness, example=str(names))
57 if len(det):
58 # Rescale boxes from img_size to im0 size
59 det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
60
61 # Print results
62 for c in det[:, 5].unique():
63 n = (det[:, 5] == c).sum() # detections per class
64 s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
65
66 data_for_image=[]
67 # Write results
68 for *xyxy, conf, cls in reversed(det):
69 xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
70 line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
71 data_for_image.append((int(cls), xywh))
72 c = int(cls) # integer class
73 label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
74 annotator.box_label(xyxy, label, color=colors(c, True))
75
76 # Sort the data based on the top-left coordinates (Y first, then X)
77 sorted_data_for_image = sorted(data_for_image, key=lambda x: (x[1][1], x[1][0]))
78 sorted_data_list.extend(sorted_data_for_image)
79
80 # Return the combined sorted data as a tuple
81 return tuple(sorted_data_list)
82
83
84# Weight path
85object_detection_output_path = 'symbol_detection/runs/detect/exp/labels'
86yolo_weights_path = 'symbol_detection/runs/train/best_all/weights/best.pt'
87yolo_yaml_file = 'symbol_detection/data/mydata.yaml'
88
89yolo_model = load_model(yolo_weights_path, device='cuda:0', dnn=False, data=yolo_yaml_file, fp16=False)
90stride, names, pt = yolo_model.stride, yolo_model.names, yolo_model.pt
91
92# Example usage
93image_path = "image.png"
94labels = run_single_image_inference(yolo_model, image_path, stride, names, pt)