-
Notifications
You must be signed in to change notification settings - Fork 0
/
labelDetector.py
79 lines (69 loc) · 2.35 KB
/
labelDetector.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
from ultralytics import YOLO
# import matplotlib.pyplot as plt
from PIL import Image
import cv2
class LabelDetector():
def __init__(self):
self.BookModel = YOLO('./weights/best2.pt')
self.LabelModel = YOLO('./weights/label.pt')
def getBookCrops(self,image_path):
crop_paths = []
result = self.BookModel(image_path)
if(len(result) == 0):
return (0,[])
result = result[0]
for i in range(len(result.boxes)):
print(result.boxes.xywh[i])
print(result.boxes.xyxy[i])
x1,y1,x2,y2=result.boxes.xyxy[i].numpy()
x1=round(x1)
y1=round(y1)
x2=round(x2)
y2=round(y2)
print(x1,x2)
img=result.orig_img
# plt.imshow(img[y1:y2,x1:x2,:])
crop=img[y1:y2,x1:x2,:]
print(img.shape)
im = Image.fromarray(crop)
im.save(f"./crops/Book/{i}.jpg")
crop_paths.append(f"./crops/Book/{i}.jpg")
return (1, crop_paths)
def getLabelCrops(self, image, index):
label_paths = []
result = self.LabelModel(image)
if(len(result) == 0):
return (0,[])
result = result[0]
for i in range(len(result)):
# plt.imshow(result.plot())
x1,y1,x2,y2=result.boxes.xyxy[i].numpy()
x1=round(x1)
y1=round(y1)
x2=round(x2)
y2=round(y2)
print(x1,x2)
img=result.orig_img
# plt.imshow(img[y1:y2,x1:x2,:])
crop=img[y1:y2,x1:x2,:]
print(img.shape)
im = Image.fromarray(crop)
im.save(f"./crops/Label/{index}_{i}.jpg")
label_paths.append(f"./crops/Label/{index}_{i}.jpg")
return (1, label_paths)
def run(self,image_path):
Book_crop=self.getBookCrops(image_path)
labels = []
if(Book_crop[0] == 1):
got_label = False
for j, path in enumerate(Book_crop[1]):
Label_crop = self.getLabelCrops(path, j)
if(Label_crop[0] == 1):
got_label = True
labels += Label_crop[1]
if(got_label):
return (1, labels)
else:
return (0,[])
else:
return (0,[])