-
Notifications
You must be signed in to change notification settings - Fork 3
/
detect.py
109 lines (78 loc) · 2.91 KB
/
detect.py
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
102
103
104
105
106
107
108
109
import torch
import os
import numpy as np
from PIL import Image
from itertools import product
import shutil
import time
import argparse
import sys
print(sys.version)
def detect_vessels(filename, dir_in, dir_out, model):
#If dir_out doesnt exist, it is created
if os.path.exists(dir_out) is False:
os.mkdir(dir_out)
d = 800
name, ext = os.path.splitext(filename)
img = Image.open(os.path.join(dir_in, filename))
w, h = img.size
divided_filenames = []
grid = list(product(np.arange(0, h, d), np.arange(0, w, d)))
for i, j in grid:
jj, ii = w if j+d>w else j+d, h if i+d>h else i+d
box = (j, i, jj, ii)
divided_filename = f"{name}_{i}_{j}_{ext}"
out_file = os.path.join(dir_out, divided_filename)
divided_filenames.append(divided_filename)
img.crop(box).save(out_file)
preprocess_files = [os.path.join(dir_out, filename) for filename in divided_filenames]
for folder in os.listdir("runs/detect"):
shutil.rmtree(folder)
# Inference
results = model(preprocess_files)
results.save()
vessels_in_picture = sum([len(vessel_image) for vessel_image in results.xyxy])
try:
for file in preprocess_files:
os.remove(file)
except:
pass
detect_dir = r"runs/detect/exp"
detect_filenames = [os.path.splitext(filename)[0] for filename in divided_filenames]
detect_files = [os.path.join(detect_dir, filename+".jpg") for filename in detect_filenames]
imgs = [Image.open(filename) for filename in detect_files]
final_image = Image.new('RGB', size=(w, h))
for picture in imgs:
name, ext = os.path.splitext(picture.filename)
i = int(name.split("_")[-3])
j = int(name.split("_")[-2])
final_image.paste(picture, box=(j, i))
shutil.rmtree(detect_dir)
(os.remove(file) for file in detect_files)
out_filename = os.path.join(dir_out, f'{filename}')
final_image.save(out_filename)
return vessels_in_picture
def main():
# Import the library
# Create the parser
parser = argparse.ArgumentParser()
parser.add_argument('--dir_in', type=str, required=True)
parser.add_argument('--dir_out', type=str, required=True)
# Parse the argument
args = parser.parse_args()
start = time.time()
dir_in = args.dir_in
dir_out = args.dir_out
model = torch.hub.load('ultralytics/yolov5', 'custom', path=r"Models\yolov5s\weights\best.pt", force_reload=True)
filenames = os.listdir(dir_in)
n_vessels = 0
for filename in filenames:
vessels_in_picture = detect_vessels(filename, dir_in, dir_out, model)
n_vessels += vessels_in_picture
end = time.time()
minutes = (end - start)/60
format_minutes = "{:.2f}".format(minutes)
print(f"{n_vessels} vessels have been detected")
print(f"The program runs in {format_minutes} minutes")
if __name__ == "__main__":
main()