inference.py
3.37 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
import copy
import os
import sys
from pathlib import Path
import numpy as np
import torch
from utils.augmentations import letterbox
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 (check_img_size, cv2, non_max_suppression, scale_boxes)
from utils.torch_utils import select_device, smart_inference_mode
from models.yolov5_config import config
classes = ['tampered']
def gen_result_dict(boxes, label_list=[], std=False):
result = {
"error_code": 1,
"result": []
}
rs_box = {
"class": '',
"score": 0,
"left": 0,
"top": 0,
"width": 0,
"height": 0
}
if not label_list:
label_list = classes
for box in boxes:
result['error_code'] = 0
box_dict = copy.deepcopy(rs_box)
if std:
box_dict['class'] = str(int(box[-1]))
else:
box_dict['class'] = label_list[int(box[-1])]
box_dict['left'] = int(round(box[0], 0))
box_dict['top'] = int(round(box[1], 0))
box_dict['width'] = int(round(box[2], 0) - round(box[0], 0))
box_dict['height'] = int(round(box[3], 0) - (round(box[1], 0)))
box_dict['score'] = box[-2]
result['result'].append(box_dict)
return result
class Yolov5:
def __init__(self, cfg=None):
self.cfg = cfg
self.device = select_device(self.cfg.device)
self.model = DetectMultiBackend(self.cfg.weights, device=self.device, dnn=False, data=self.cfg.data, fp16=False)
def detect(self, image):
image0 = image.copy()
stride, names, pt = self.model.stride, self.model.names, self.model.pt
imgsz = check_img_size(self.cfg.imgsz, s=stride) # check image size
# Dataloader
bs = 1 # batch_size
im = letterbox(image, imgsz, stride=stride, auto=True)[0] # padded resize
im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
im = np.ascontiguousarray(im) # contiguous
# Run inference
self.model.warmup(imgsz=(1 if pt or self.model.triton else bs, 3, *imgsz)) # warmup
im = torch.from_numpy(im).to(self.model.device)
im = im.half() if self.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
pred = self.model(im, augment=False, visualize=False)
# NMS
pred = non_max_suppression(pred, self.cfg.conf_thres, self.cfg.iou_thres, None, False, max_det=self.cfg.max_det)
det = pred[0]
# if len(det):
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], image0.shape).round()
result = gen_result_dict(det.cpu().numpy().tolist())
return result
def plot(self, image, boxes):
for box in boxes:
cv2.rectangle(image, (box[0], box[1], box[2], box[3]), (0, 0, 255), 2)
return image
if __name__ == "__main__":
img = cv2.imread(
'/home/situ/qfs/invoice_tamper/09_project/project/yolov5_inference/data/images/crop_img/_1594890230.8032346page_10_img_0_hname.jpg')
detector = Yolov5(config)
result = detector.detect(img)
print(result)