Update app.py
Browse files
app.py
CHANGED
@@ -12,8 +12,8 @@ BOX_COLORS = {
|
|
12 |
BOX_PADDING = 2
|
13 |
|
14 |
# Load models
|
15 |
-
DETECTION_MODEL = YOLO("
|
16 |
-
CLASSIFICATION_MODEL = YOLO("
|
17 |
|
18 |
def detect(image_path):
|
19 |
"""
|
@@ -29,7 +29,7 @@ def detect(image_path):
|
|
29 |
return image
|
30 |
|
31 |
# Predict on image
|
32 |
-
results = DETECTION_MODEL.predict(source=image, conf=0.
|
33 |
boxes = results[0].boxes # Get bounding boxes
|
34 |
|
35 |
if len(boxes) == 0:
|
|
|
12 |
BOX_PADDING = 2
|
13 |
|
14 |
# Load models
|
15 |
+
DETECTION_MODEL = YOLO("detector-model.pt")
|
16 |
+
CLASSIFICATION_MODEL = YOLO("classifier-model.pt") # 0: block, 1: checked, 2: unchecked
|
17 |
|
18 |
def detect(image_path):
|
19 |
"""
|
|
|
29 |
return image
|
30 |
|
31 |
# Predict on image
|
32 |
+
results = DETECTION_MODEL.predict(source=image, conf=0.5, iou=0.8) # Predict on image
|
33 |
boxes = results[0].boxes # Get bounding boxes
|
34 |
|
35 |
if len(boxes) == 0:
|