Skip to content

Commit 3066bf6

Browse files
mtsousamtsousa
committed
Adicionei o arquivo para detecção com o modelo ONNX
1 parent 62a8985 commit 3066bf6

File tree

5 files changed

+359
-39
lines changed

5 files changed

+359
-39
lines changed

‎.gitignore

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,5 @@ __pycache__/
99
yolov7-fork-o/
1010
yolov7-fork-p/
1111
*.pt
12-
*.onnx
12+
*.onnx
13+
plot_mosaic.py

‎detect.py

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
"""
2+
@Author: Matheus Teixeira de Sousa (mtsousa14@gmail.com)
3+
4+
Detect forks from images with trained YOLOv7 ONNX model
5+
"""
6+
7+
import cv2 as cv
8+
import numpy as np
9+
import onnxruntime as ort
10+
from torch.utils.data import DataLoader
11+
from utils.dataset import TestDataset
12+
from utils.utils import plot_one_box, show_predicted_image, adjust_image
13+
from os.path import exists, isdir
14+
from os import makedirs
15+
from random import randint
16+
import argparse
17+
18+
def predict_bbox(session, images):
19+
"""
20+
Predict bounding boxes from images
21+
"""
22+
outname = [i.name for i in session.get_outputs()]
23+
24+
dict_output = {}
25+
for i, samples in enumerate(images):
26+
im, ratio, dwdh, name = samples['image'], samples['ratio'], samples['dwdh'], samples['name']
27+
im = np.ascontiguousarray(im/255)
28+
out = session.run(outname, {'images':im})[0]
29+
dict_output[f"batch {i}"] = {"preds": out, "ratio": ratio, "dwdh": dwdh, "name": name}
30+
31+
return dict_output
32+
33+
if __name__ == '__main__':
34+
# Parse command line arguments
35+
parser = argparse.ArgumentParser(
36+
description='Predict with YOLOv7-fork ONNX model')
37+
38+
parser.add_argument('--model', required=True,
39+
metavar='/path/to/model.onnx',
40+
help="Path to ONNX model")
41+
parser.add_argument('--input', required=True,
42+
help="Path to images (path/to/images) or path to image (path/to/image.jpg)")
43+
parser.add_argument('--batch', default=1,
44+
help="Batch size")
45+
parser.add_argument('--save', default=False, action='store_true',
46+
help="Save predicted image")
47+
parser.add_argument('--dontshow', default=False, action='store_true',
48+
help="Don't show predicted image")
49+
parser.add_argument('--cuda', default=False, action='store_true',
50+
help="Set execution on GPU")
51+
52+
args = parser.parse_args()
53+
for key, value in args._get_kwargs():
54+
if value is not None:
55+
print(f'{key.capitalize()}: {value}')
56+
print()
57+
58+
# Check if the input is a dir
59+
input_isdir = isdir(args.input)
60+
61+
# Load the model
62+
print('Loading model...', flush=True)
63+
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if args.cuda else ['CPUExecutionProvider']
64+
session = ort.InferenceSession(args.model, providers=providers)
65+
66+
# Get output name and input shape
67+
outname = [i.name for i in session.get_outputs()]
68+
input_shape = session.get_inputs()[0].shape
69+
h, w = input_shape[2], input_shape[3]
70+
71+
# Load the images
72+
print('Loading images...', flush=True)
73+
if input_isdir:
74+
dataset = TestDataset(args.input, shape=(h, w))
75+
images = DataLoader(dataset, batch_size=args.batch, shuffle=False, num_workers=0)
76+
else:
77+
images = [adjust_image(args.input, shape=(h, w))]
78+
79+
# Predict from images
80+
print('Making predictions...', flush=True)
81+
dict_output = predict_bbox(session, images)
82+
83+
names = ['fork']
84+
colors = {name: [randint(0, 255) for _ in range(3)] for name in names}
85+
# colors = {name: [104, 184, 82] for name in names} # green
86+
87+
if args.save and not exists(f'data/responses'):
88+
makedirs(f'data/responses')
89+
90+
# For each image, plot the results
91+
print('Plotting results...', flush=True)
92+
for i, key in enumerate(dict_output.keys()):
93+
pred, ratio, dwdh, name = dict_output[key]['preds'], dict_output[key]['ratio'][0], dict_output[key]['dwdh'], dict_output[key]['name'][0]
94+
ratio = float(ratio)
95+
dwdh = float(dwdh[0]), float(dwdh[1])
96+
97+
# Load original image
98+
if input_isdir:
99+
image = dataset.__getsrc__(i)
100+
else:
101+
image = cv.imread(args.input)
102+
103+
# Adjust bounding box to original image
104+
for prediction in pred:
105+
batch_id, x0, y0, x1, y1, cls_id, score = prediction
106+
box = np.array([x0,y0,x1,y1])
107+
box -= np.array(dwdh*2)
108+
box /= ratio
109+
box = box.round().astype(np.int32).tolist()
110+
cls_id = int(cls_id)
111+
score = round(float(score),3)
112+
label = names[cls_id]
113+
color = colors[label]
114+
label += ' ' + str(score)
115+
plot_one_box(box, image, label=label, color=color, line_thickness=1)
116+
117+
if args.save:
118+
path = 'data/responses/' + name
119+
cv.imwrite(path, image)
120+
121+
if not args.dontshow:
122+
show_predicted_image(image)
123+

‎utils/dataset.py

Lines changed: 46 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@
88
import json
99
import cv2 as cv
1010
from torch import as_tensor, float64
11+
from .utils import letterbox
12+
import numpy as np
13+
import os
1114

1215
class CocoForkDataset(Dataset):
1316
def __init__(self, img_dir, ann_dir, transforms=None, mode='train', device='cpu'):
@@ -99,25 +102,46 @@ def read_ids(self, ann_dir, mode):
99102
img_ids = list(dict.fromkeys(img_ids))
100103
return img_ids
101104

102-
# def show_image(img, bbox):
103-
# """
104-
# Draw the bbox and show the image
105-
# """
106-
# import numpy as np
107-
# img = np.float32(img)
108-
# img = cv.cvtColor(img, cv.COLOR_RGB2BGR)
109-
# for box in bbox:
110-
# x0, y0, x1, y1 = int(box[0]), int(box[1]), int(box[2]), int(box[3])
111-
# cv.rectangle(img, (x0, y0), (x1, y1), (255,0,0), 2)
112-
# # cv.imwrite('img_name.jpg', img*255)
113-
# cv.imshow(f'IMG', img)
114-
# cv.waitKey(0)
115-
# cv.destroyAllWindows()
116-
117-
# from utils import get_transforms
118-
119-
# dataset = CocoForkDataset('../data/train', '../data/annotations', get_transforms(train=True, size=(640, 640)), 'train')
120-
# for k in range(20, 30):
121-
# img, ann = dataset.__getitem__(k)
122-
# # print(img.shape, ann)
123-
# show_image(img.numpy().transpose(1, 2, 0), ann['boxes'])
105+
class TestDataset(Dataset):
106+
def __init__(self, path, shape):
107+
"""
108+
Initiliaze the dataset
109+
110+
Args
111+
- path: Path to images
112+
- shape: Input shape
113+
"""
114+
self.path = path
115+
self.imgs_list = os.listdir(path)
116+
self.shape = shape
117+
118+
def __len__(self):
119+
"""
120+
Return the length of the dataset
121+
"""
122+
return len(self.imgs_list)
123+
124+
def __getitem__(self, idx):
125+
"""
126+
Get an item from dataset by index with transformations
127+
"""
128+
img = cv.imread(self.path + "/" + self.imgs_list[idx])
129+
img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
130+
image = img.copy()
131+
132+
# Adjust image shape
133+
image, ratio, dwdh = letterbox(image, new_shape=self.shape, auto=False)
134+
135+
# Transpose and set as a contiguous array
136+
image = image.transpose((2, 0, 1))
137+
image = np.ascontiguousarray(image)
138+
im = image.astype(np.float32)
139+
140+
return {"image": im, "ratio": ratio,"dwdh": dwdh, "name": self.imgs_list[idx]}
141+
142+
def __getsrc__(self, idx):
143+
"""
144+
Get an item from dataset by index without transformations
145+
"""
146+
img = cv.imread(self.path + "/" + self.imgs_list[idx])
147+
return img

‎utils/images_exploration/show_images.py

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,16 @@
99
import argparse
1010
from random import shuffle
1111

12+
import sys
13+
from pathlib import Path
14+
15+
file = Path(__file__).resolve()
16+
package_root_directory = file.parents[2]
17+
print(package_root_directory)
18+
sys.path.append(str(package_root_directory))
19+
20+
from utils.utils import plot_one_box, show_predicted_image
21+
1222
def read_json(input_name):
1323
"""
1424
Read json annotations file
@@ -48,19 +58,6 @@ def find_image(coco, img_ids):
4858
images[image_id]['bbox'].append(bbox)
4959
return images
5060

51-
def show_image(id, dataset, file_name, bbox):
52-
"""
53-
Draw the bbox and show the image
54-
"""
55-
img = cv.imread(f'../../data/{dataset}/{file_name}')
56-
# "bbox": [x,y,width,height]
57-
for box in bbox:
58-
x, y, w, h = int(box[0]), int(box[1]), int(box[2]), int(box[3])
59-
cv.rectangle(img, (x, y), (x+w, y+h), (255,0,0), 2)
60-
cv.imshow(f'IMG {id}', img)
61-
cv.waitKey(0)
62-
cv.destroyAllWindows()
63-
6461
if __name__ == '__main__':
6562
# Parse command line arguments
6663
parser = argparse.ArgumentParser(description='Explore COCO dataset.')
@@ -86,4 +83,12 @@ def show_image(id, dataset, file_name, bbox):
8683
# For each image, show the bbox
8784
print('Loading images...')
8885
for k in images.keys():
89-
show_image(k, dataset, images[k]['name'], images[k]['bbox'])
86+
image = images[k]['name']
87+
img = cv.imread(f'../../data/{dataset}/{image}')
88+
89+
for j in images[k]['bbox']:
90+
x, y, w, h = int(j[0]), int(j[1]), int(j[2]), int(j[3])
91+
box = [x, y, x+w, y+h]
92+
plot_one_box(box, img, label='fork', color=[104, 184, 82], line_thickness=1)
93+
94+
show_predicted_image(img)

0 commit comments

Comments
 (0)