def save_violation(cfg, violation_frame, img_stack, img_result, frame_num): if violation_frame != -1: global im_name global plate_word img_stack.append(img_result) f = frame_num index = [0, f // 2, f, int(f * 1.5), f * 2] # 且湊滿frame_num*2張frame,更改旗標,清空 if len(img_stack) == f * 2 + 1: tmp_path = './violation/' + str(violation_frame) + '/' for i in index: save_index = violation_frame - f + i img_path = tmp_path + str(save_index) + '.png' im_name = str(violation_frame) cv2.imwrite(img_path, img_stack[i]) #車牌辨識##############charnet########## charnet = CharNet() charnet.load_state_dict(torch.load(cfg.WEIGHT)) charnet.eval() charnet.cuda() #裁切圖片再丟進去辨識 im_result = img_stack[i][360:1080, 640:1920] im, scale_w, scale_h, original_w, original_h = resize( im_result, size=cfg.INPUT_SIZE) start_time = time.time() with torch.no_grad(): char_bboxes, char_scores, word_instances = charnet( im, scale_w, scale_h, original_w, original_h) plate_word = save_word_recognition(word_instances, im_name, "./charnet_result", cfg.RESULTS_SEPARATOR) #draw the box and result #draw_box_word(im_name,im_result,"./charnet_result") print("--- %s seconds ---" % (time.time() - start_time)) ######################### print( '==============>>>>>> Sucess save Violation Image folder : ', violation_frame, "in", im_name) f = frame_num * -1 img_stack = img_stack[f:] violation_frame = -1 print("license number ----------------------", plate_word) #########change dir name for folder_name in os.listdir('./violation'): if folder_name == im_name and len(plate_word) > 0: os.rename('./violation/{}'.format(im_name), './violation/{}'.format(plate_word)) # 如果沒有違規 else: if len(img_stack) < frame_num: img_stack.append(img_result) else: f = (frame_num * -1) + 1 img_stack = img_stack[f:] img_stack.append(img_result) #只保留最新frame_num筆img_result return violation_frame, img_stack
if __name__ == '__main__': parser = argparse.ArgumentParser(description="Test") parser.add_argument("config_file", help="path to config file", type=str) parser.add_argument("image_dir", type=str) parser.add_argument("results_dir", type=str) args = parser.parse_args() cfg.merge_from_file(args.config_file) cfg.freeze() print(cfg) charnet = CharNet() charnet.load_state_dict(torch.load(cfg.WEIGHT)) charnet.eval() charnet.cuda() for im_name in tqdm(sorted(os.listdir(args.image_dir))): #print("Processing {}...".format(im_name)) im_file = os.path.join(args.image_dir, im_name) im_original = cv2.imread(im_file) im, scale_w, scale_h, original_w, original_h = resize( im_original, size=cfg.INPUT_SIZE) with torch.no_grad(): char_bboxes, char_scores, word_instances = charnet( im, scale_w, scale_h, original_w, original_h) save_word_recognition(word_instances, os.path.splitext(im_name)[0], args.results_dir, cfg.RESULTS_SEPARATOR)
class CharNetRunner: """Warper for charnet. Handles extracting words on windows from the panorama Public Functions: get_absolute_window_words - extract words from a given window of the panorama clean_duplicate_words - remove duplicated from words extracted in the same window new_words_only - filter out words that have been extracted in previous windows Variables: config_file_path - path for the charnet config file charnet - charnet instance """ def __init__(self, config_file): self.config_file_path = config_file # ----- CharNet configuration and initialization ----- cfg.merge_from_file(config_file) cfg.freeze() print(cfg) self.charnet = CharNet() self.charnet.load_state_dict(torch.load(cfg.WEIGHT)) self.charnet.eval() if torch.cuda.is_available(): print("[*] Using cuda!") self.charnet.cuda() def get_absolute_window_words(self, pano_windows, window): """Get the words of the given window. Every word position is relative to the whole panorama, i.e. absolute position :param pano_windows: ImageWindows instance of the panorama :param window: Window instance of the current window, must be a member of the pano_windows :return: WordsInstance list of recognized words in the current window """ words = [] im, scale_w, scale_h, window_w, window_h = self.__resize(window.im) with torch.no_grad(): # char_bboxes, char_scores, word_instances = ... _, _, word_instances = self.charnet(im, scale_w, scale_h, window_w, window_h) for word in word_instances: # To combat google's watermark of street-view messing with the words if word.text == 'GOOGLE': continue old_word_bbox = word.word_bbox.copy() # update absolute position word.word_bbox[::2] = [ x_coord + window.pos_x for x_coord in word.word_bbox[::2] ] word.word_bbox[1::2] = [ y_coord + window.pos_y for y_coord in word.word_bbox[1::2] ] word_abs = word # open a new window for near-border words if self.__word_is_near_border(old_word_bbox, 50, window_w, window_h): zoom_w = pano_windows.get_window_at_pos( word.word_bbox[0], word.word_bbox[1], 50) z_im, z_scale_w, z_scale_h, z_window_w, z_window_h = self.__resize( zoom_w.im) with torch.no_grad(): _, _, z_word_instances = self.charnet( z_im, z_scale_w, z_scale_h, z_window_w, z_window_h) for z_word in z_word_instances: # Swap only the word that intersects z_word.word_bbox[::2] = [ x_coord + zoom_w.pos_x for x_coord in z_word.word_bbox[::2] ] z_word.word_bbox[1::2] = [ y_coord + zoom_w.pos_y for y_coord in z_word.word_bbox[1::2] ] if self._do_words_intersect(word, z_word): word_abs = z_word # save only the new word from the window break words.append(word_abs) return words @staticmethod def __resize(im, size=cfg.INPUT_SIZE): h, w, _ = im.shape scale = max(h, w) / float(size) image_resize_height = int( round(h / scale / cfg.SIZE_DIVISIBILITY) * cfg.SIZE_DIVISIBILITY) image_resize_width = int( round(w / scale / cfg.SIZE_DIVISIBILITY) * cfg.SIZE_DIVISIBILITY) scale_h = float(h) / image_resize_height scale_w = float(w) / image_resize_width im = cv2.resize(im, (image_resize_width, image_resize_height), interpolation=cv2.INTER_LINEAR) return im, scale_w, scale_h, w, h @staticmethod def _do_words_intersect(s_word, o_word): """Checks if the bounding boxes of 2 words intersects (goes ontop of each-other) :param s_word: first WordInstance :param o_word: second WordInstance :return: Boolean of weather an intersection occurs """ s_poly = Polygon([(x_pos, y_pos) for x_pos, y_pos in zip( s_word.word_bbox[::2], s_word.word_bbox[1::2])]) o_poly = Polygon([(x_pos, y_pos) for x_pos, y_pos in zip( o_word.word_bbox[::2], o_word.word_bbox[1::2])]) return s_poly.intersects(o_poly) @staticmethod def __word_is_near_border(bbox, margin, window_w, window_h): """Check if a given bounding box (of a word) is near the border of the window :param bbox: bounding box (of a word recognized by the network) :param margin: margin amount (in pixels) from the side of the window :param window_w: window's width :param window_h: window's height :return: Boolean of weather the word is near the border """ # [0][1] [2][3] # # [4][5] [6][7] if bbox[2] > window_w - margin or bbox[6] > window_w - margin: return True if bbox[5] > window_h - margin or bbox[7] > window_h - margin: return True return False @staticmethod def clean_duplicate_words(words): """Remove words that appeared in multiple windows, but are in face the same instance This might happen when a word is near 2 margins, and so 2 windows will open for it seperatly, and so it will be recognized twice. The comparison is based on intersection, i.e. the absolute position of the word in the panorama :param words: WordsInstance list :return: filtered WordsInstance list """ clean_words = [] state = np.ones(len(words)) while state.any(): base_idx = state.nonzero()[0][0] for idx in range(len(words)): base_word = words[base_idx] new_word = words[idx] # if 2 different words (both still viable) intersect if idx != base_idx and state[idx] and \ CharNetRunner._do_words_intersect(base_word, new_word): # decide if to keep left/right word if base_word.text_score > new_word.text_score: state[idx] = False else: state[base_idx] = False break # if all other intersections were less accurate if state[base_idx]: clean_words.append(words[base_idx]) state[base_idx] = False return clean_words @staticmethod def new_words_only(base_words, window_words): """Filter previously recognized words from this window's recognized words this are the same words position-wise. :param base_words: previously recognized words in a list of WordsInstance :param window_words: this window's list of WordsInstance recognized words :return: filtered WordsInstance list """ new_words = [] if base_words: for wword in window_words: dont_include = False for bword in base_words: if CharNetRunner._do_words_intersect(wword, bword): dont_include = True break if not dont_include: new_words.append(wword) else: new_words = window_words return new_words