#!/usr/bin/env python3 #import matplotlib.pyplot as plt import argparse as ap import numpy as np import cv2 as cv import os import sys class CustomFormatter(ap.HelpFormatter): def _format_action_invocation(self, action): if not action.option_strings: metavar, = self._metavar_formatter(action, action.dest)(1) return metavar else: parts = [] if action.nargs == 0: parts.extend(action.option_strings) else: default = action.dest.upper() args_string = self._format_args(action, default) for option_string in action.option_strings: #parts.append('%s %s' % (option_string, args_string)) parts.append('%s' % option_string) parts[-1] += ' %s'%args_string return ', '.join(parts) # Parser Arguments parser = ap.ArgumentParser(description='Cascade Classifier', formatter_class=CustomFormatter) parser.add_argument("-s", "--save", metavar='', help="specify output name") parser.add_argument("-c", "--cas", metavar='', help="specify specific trained cascade", default="./stage_outputs/cascade.xml") parser.add_argument("-i", "--img", metavar='', help="specify image to be classified") parser.add_argument("-d", "--dir", metavar='', help="specify directory of images to be classified") parser.add_argument("-v", "--vid", metavar='', help="specify video to be classified") parser.add_argument("-w", "--cam", metavar='', help="enable camera access for classification") parser.add_argument("-f", "--fps", help="enable frames text (TODO)", action="store_true") parser.add_argument("-o", "--circle", help="enable circle detection", action="store_true") parser.add_argument("-z", "--scale", metavar='', help="decrease video scale by scale factor", type=int, default=1) parser.add_argument("-t", "--track", metavar='', help="select tracking algorithm [KCF, CSRT, MEDIANFLOW]", choices=['KCF', 'CSRT', 'MEDIANFLOW']) args = parser.parse_args(sys.argv[1:]) # Load the trained cascade cascade = cv.CascadeClassifier() if not cascade.load(args.cas): print("Can't find cascade file. Do you have the directory ./stage_outputs/cascade.xml") exit(0) def plot(): pass def detect_circles(src): img = src img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) img_blur = cv.medianBlur(img_gray, 5) rows = img_blur.shape[0] #Images circles = cv.HoughCircles(img_blur, cv.HOUGH_GRADIENT, 1, rows / 3, param1=100, param2=40, maxRadius=40) circles = cv.HoughCircles(img_blur, cv.HOUGH_GRADIENT, 1, rows/3, param1=100, param2=15, minRadius=10, maxRadius=15) if circles is not None: circles = np.uint16(np.around(circles)) for i in circles[0, :]: center = (i[0], i[1]) # circle center cv.circle(img, center, 1, (0, 100, 100), 3) # circle outline radius = i[2] cv.circle(img, center, radius, (255, 0, 255), 3) return img def choose_tracker(): OPENCV_TRACKERS = { 'KCF': cv.TrackerKCF_create(), 'CSRT': cv.TrackerCSRT_create(), 'MEDIANFLOW': cv.TrackerMedianFlow_create() } tracker = OPENCV_TRACKERS[args.track] return tracker def tracking(vid, tracker): ok, frame = vid.read() frame = scale(frame, args.scale) ok, roi = tracker.update(frame) if ok: p1 = (int(roi[0]), int(roi[1])) p2 = (int(roi[0] + roi[2]), int(roi[1] + roi[3])) cv.rectangle(frame, p1, p2, (0,255,0), 2, 1) cpoint_circle = cv.circle(frame, (int(roi[0]+(roi[2]/2)), int(roi[1]+(roi[3]/2))), 3, (0,255,0), 3) return frame def save(frame): # Need dimensions of frame to determine proper video output fourcc = cv.VideoWriter_fourcc(*'XVID') height, width, channels = frame.shape out = cv.VideoWriter(args.save + '.avi', fourcc, 30.0, (width, height)) return out def get_roi(frame): # Get initial bounding box by running cascade detection on first frame frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) frame_gray = cv.GaussianBlur(frame_gray, (3, 3), 0) cas_object = cascade.detectMultiScale(frame_gray) if len(cas_object) == 0: return [] roi = (cas_object[0][0], cas_object[0][1], cas_object[0][2], cas_object[0][3]) return roi def get_cascade(frame): frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) frame_gray = cv.GaussianBlur(frame_gray, (3, 3), 0) cas_object = cascade.detectMultiScale(frame_gray) for (x, y, w, h) in cas_object: cv.rectangle(frame, (x,y), (x+w, y+h), (0,0,255), 2) cpoint_circle = cv.circle(frame, (int(x+(w/2)), int(y+(h/2))), 3, (0,0,255), 3) return frame def scale(frame, scale_factor): height, width, channels = frame.shape scaled_height = int(height/scale_factor) scaled_width = int(width/scale_factor) resized_frame = cv.resize(frame, (scaled_width, scaled_height)) return resized_frame def img_classifier(): # Read image, convert to gray, equalize histogram, and detect. img = cv.imread(args.img, cv.IMREAD_COLOR) img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) #img_gray = cv.equalizeHist(img_gray) cas_object = cascade.detectMultiScale(img_gray) for (x, y, w, h) in cas_object: roi = cv.rectangle(img, (x,y), (x+w, y+h), (0,0,255), 2) cpoint_circle = cv.circle(img, (int(x+(w/2)), int(y+(h/2))), 3, (0,0,255), 3) if args.circle is True: roi = img[y:y+h, x:x+w] img = detect_circles(roi) cv.imshow('image', img) cv.waitKey(0) cv.destroyAllWindows() def dir_classifier(): imgs = [] for filename in os.listdir(args.dir): img = cv.imread(os.path.join(args.dir, filename)) if img is not None: img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) cas_object = cascade.detectMultiScale(img_gray) for (x, y, w, h) in cas_object: cv.rectangle(img, (x,y), (x+w, y+h), (0,0,255), 2) cpoint_circle = cv.circle(img, (int(x+(w/2)), int(y+(h/2))), 3, (0,0,255), 3) if args.circle is True: roi = img[y:y+h, x:x+w] img = detect_circles(roi) cv.imshow(str(filename), img) cv.waitKey(0) imgs.append(img) #print(imgs) #return imgs def vid_classifier(): vid = cv.VideoCapture(args.vid) if not vid.isOpened(): print("Could not open video") sys.exit() # Read the first frame _ , frame = vid.read() frame = scale(frame, args.scale) if not _: print("Cannot read video file") sys.exit() if args.save is not None and _ is True: out = save(frame=frame) if args.track is not None and _ is True: cas_roi = get_roi(frame) while not cas_roi: _, frame = vid.read() frame = scale(frame, args.scale) cas_roi = get_roi(frame) tracker = choose_tracker() tracker.init(frame, cas_roi) while(vid.isOpened()): _ , frame = vid.read() frame = scale(frame, args.scale) frame = get_cascade(frame) if args.track is not None: frame = tracking(vid=vid, tracker=tracker) if args.circle is True: roi = get_roi(frame) roi_circle = frame[int(roi[1]):int(roi[1] + roi[3]), int(roi[0]):int(roi[0] + roi[2])] frame = detect_circles(roi_circle) cv.imshow('video', frame) if args.save is not None: out.write(frame) if cv.waitKey(1) & 0xFF == ord('q'): break if args.save is not None: out.release() vid.release() cv.destroyAllWindows() def cam_classifier(): cam = cv.VideoCapture(0) if not cam.isOpened(): raise IOError("Cannot access camera") while(cam.isOpened()): _, frame = cap.read() frame = get_cascade(frame) cv2.imshow('camera', frame) if cv.waitKey(10) & 0xFF == ord('q'): break cam.release() cv.destroyAllWindows() if __name__ == "__main__": if args.img is not None: img_classifier() elif args.vid is not None: vid_classifier() elif args.dir is not None: dir_classifier() elif args.cam is not None: cam_clasifier() else: parser.print_help()