forked from Cuda-Chen/fish-yolo-grabcut
-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
29 lines (25 loc) · 959 Bytes
/
main.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
#!/usr/bin/python
import argparse
import cv2 as cv
from utils import yolo, GrabCut
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", required=True,
help="path to input image")
ap.add_argument("-y", "--yolo", required=True,
help="base path to YOLO directory")
ap.add_argument("-c", "--confidence", type=float, default=0.25,
help="minimum probability to filter weak detections")
ap.add_argument("-t", "--threshold", type=float, default=0.45,
help="threshold when applying non-maxima suppression")
args = vars(ap.parse_args())
img, boxes, idxs = yolo.runYOLOBoundingBoxes(args)
print("boxes' length: ", len(boxes))
print("idxs' shape: ", idxs.shape)
images = GrabCut.runGrabCut(img, boxes, idxs)
# show the output image
#cv.namedWindow("Image", cv.WINDOW_NORMAL)
#cv.resizeWindow("image", 1920, 1080)
for i in range(len(images)):
#cv.imshow("Image", image)
cv.imwrite("grabcut{}.jpg".format(i), images[i])
cv.waitKey(0)