43 lines
1.5 KiB
Python
43 lines
1.5 KiB
Python
|
import argparse
|
||
|
|
import cv2
|
||
|
|
from models import ClassificationModel
|
||
|
|
|
||
|
|
cls_model = ClassificationModel()
|
||
|
|
|
||
|
|
def main(args):
|
||
|
|
cap = cv2.VideoCapture(args.video_path)
|
||
|
|
object_detector = cv2.createBackgroundSubtractorMOG2(history=100, varThreshold=50)
|
||
|
|
while True:
|
||
|
|
ret, frame = cap.read()
|
||
|
|
if(ret):
|
||
|
|
roi = frame[100: 900,330:1900]
|
||
|
|
mask = object_detector.apply(roi)
|
||
|
|
_, mask = cv2.threshold(mask,254,255, cv2.THRESH_BINARY)
|
||
|
|
conturs, _ =cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
|
|
||
|
|
for cnt in conturs:
|
||
|
|
area = cv2.contourArea(cnt)
|
||
|
|
if area > 200:
|
||
|
|
cv2.drawContours(roi,[cnt],-1,(0,255,0),2)
|
||
|
|
x,y,w,h = cv2.boundingRect(cnt)
|
||
|
|
rectangle = cv2.rectangle(roi,(x,y),(x+w,y+h),(0,255,0),3)
|
||
|
|
image_to_predict = roi[y:y+h+10,x:x+w+10]
|
||
|
|
label = cls_model.predict(image_to_predict)
|
||
|
|
cv2.putText(rectangle, label, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 1)
|
||
|
|
|
||
|
|
roi = cv2.resize(roi, (960, 540))
|
||
|
|
cv2.imshow("roi", roi)
|
||
|
|
|
||
|
|
key = cv2.waitKey(30)
|
||
|
|
if key == 27:
|
||
|
|
break
|
||
|
|
else:
|
||
|
|
break
|
||
|
|
cap.release()
|
||
|
|
cv2.destroyAllWindows()
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--video_path", type=str, default='./test_videos/rybki2.mp4')
|
||
|
|
args = parser.parse_args()
|
||
|
|
main(args)
|