class Solver: def __init__(self, dirname=DEFAULT_MODEL_DIR, gpu=-1, nms_thresh=DEFAULT_NMS_THRESH, score_thresh=DEFAULT_SCORE_THRESH): with open(os.path.join(dirname, "model.json"), 'r') as fp: metadata = json.load(fp) n_class = metadata['n_class'] n_channel = metadata['n_channel'] npz_file = metadata['file'] self.class_labels = metadata['class_labels'] self.model = SSD(n_class=n_class, n_channel=n_channel, nms_thresh=nms_thresh, score_thresh=score_thresh, grids=DEFAULT_GRIDS, aspect_ratios=DEFAULT_ASPECT_RATIOS, variance=DEFAULT_VARIANCE) chainer.serializers.load_npz(os.path.join(dirname, npz_file), self.model) if gpu >= 0: chainer.backends.cuda.get_device_from_id(gpu).use() self.model.to_gpu(gpu) @property def xp(self): return self.model.xp def solve(self, filename): xp = self.xp gif = cv2.VideoCapture(filename) _, color_image = gif.read(0) gray_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2GRAY) h, w = gray_image.shape[:2] img = xp.array(gray_image / 255.0, dtype=xp.float32).reshape(1, 1, h, w) output = self.model.predict(img) bbox, label, score = output[0][0], output[1][0], output[2][0] bbox = chainer.dataset.to_device(-1, bbox) label = chainer.dataset.to_device(-1, label) score = chainer.dataset.to_device(-1, score) if len(label) > NCHARS: indices = np.argsort(score)[-1:-NCHARS-1:-1] bbox = bbox[indices] label = label[indices] score = score[indices] bbox = np.vectorize(lambda v: int(v + 0.5), otypes=[int])(bbox) indices = np.argsort(bbox[:, 1]) text = ''.join([ self.class_labels[label[i]] for i in indices ]) return text, bbox[indices], score[indices]
# 予測結果のbbox print("予測結果") ################## # 画像を選択 index = image_path = os.path.join(image_dir, image_files[index]) img = cv2.imread(image_path) gt = annotation[index] answer = # 予測 confidence_threshold = out = ssd.predict(image_path) out = [pred for pred in out if pred["score"]>=confidence_threshold] out = sorted(out, key=lambda x:x["score"], reverse=True) # 評価 pred = evaluation = print('f1score:', evaluation) #表示 plot_bbox(img, gt, out) print() ################## # 画像を選択
print(os.listdir(models_dir)) model_path = os.path.join(models_dir, "ssd7.h5") ssd.load_weights(model_path) ################################################################## # 画像を1枚選択 index = 4 image_path = os.path.join(image_dir, image_files[index]) img = cv2.imread(image_path) gt = annotation[index] # 選択した画像に対してbboxの予測を行う output = ssd.predict(image_path) #print(output) # 確信度の閾値 confidence_threshold = 0.25 # スコアが閾値以上のものだけに絞る output_thresh = [prediction for prediction in output if prediction["score"]>= confidence_threshold] # スコア順にソート output_thresh = sorted(output_thresh, key=lambda x:x["score"], reverse=True) print(output_thresh) def plot_bbox(img, gt, out): # グラフサイズの指定 plt.figure(figsize=(11,11))
net = SSD(cuda=args.cuda, architecture='300_VGG16', num_classes=len(LogoDataset.CLASSES)) has_cuda = args.cuda and torch.cuda.is_available() if has_cuda: weights = torch.load(args.weights)['model'] else: weights = torch.load(args.weights, map_location='cpu')['model'] net = SSD.load(weights=weights) COLORMAP = [(255, 0, 0), (0, 255, 0), (0, 0, 255)] images = [] images = [cv2.imread(filename) for filename in glob.glob(args.test_filenames)] results = net.predict(images) for im, result_image in zip(images, results): for i, result in enumerate(result_image): print(LogoDataset.CLASSES[result['class']]) class_ = LogoDataset.CLASSES[result['class']] position = result['position'] confidence = int(100 * result['confidence']) cv2.rectangle(im, (int(position[0]), int(position[1])), (int(position[2]), int(position[3])), COLORMAP[i % len(COLORMAP)]) cv2.putText(im, '%s (%d%%)' % (class_, confidence), (int(position[0]), int(position[1])), cv2.FONT_HERSHEY_SIMPLEX, 2, COLORMAP[i % len(COLORMAP)], 2, cv2.LINE_AA)
#print(ssd.model.summary()) models_dir = os.path.join(".", "trained_models") #print(os.listdir(models_dir)) model_path = os.path.join(models_dir, "ssd7.h5") ssd.load_weights(model_path) # 画像を1枚選択 index = 4 image_path = os.path.join(image_dir, image_files[index]) img = cv2.imread(image_path) gt = annotation[index] # 選択した画像に対してbboxの予測を行う output = ssd.predict(image_path) #print(output) # 確信度の閾値 confidence_threshold = 0.25 # スコアが閾値以上のものだけに絞る output_thresh = [prediction for prediction in output if prediction["score"]>= confidence_threshold] # スコア順にソート output_thresh = sorted(output_thresh, key=lambda x:x["score"], reverse=True) print(output_thresh) def plot_bbox(img, gt, out): # グラフサイズの指定 plt.figure(figsize=(11,11))