def Init_Net(self, net_type): print("==========Init Net==========") if net_type=="mb2-ssd": model_path = "./models/mobilenet2-ssd.pth" label_path = "./models/voc-model-labels_mb2.txt" self.class_names = [name.strip() for name in open(label_path).readlines()] net = create_mobilenetv2_ssd_lite(len(self.class_names), is_test=True) net.load(model_path) self.predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200) elif net_type=="mb3-large-ssd": model_path = "./models/mobilenet3-large-ssd.pth" label_path = "./models/voc-model-labels_mb3.txt" self.class_names = [name.strip() for name in open(label_path).readlines()] net = create_mobilenetv3_ssd_lite("Large", len(self.class_names), is_test=True) net.load(model_path) self.predictor = create_mobilenetv3_ssd_lite_predictor(net, candidate_size=200)
if not picked_box_probs: return torch.tensor([]), torch.tensor([]), torch.tensor([]) picked_box_probs = torch.cat(picked_box_probs) picked_box_probs[:, 0] *= width picked_box_probs[:, 1] *= height picked_box_probs[:, 2] *= width picked_box_probs[:, 3] *= height return picked_box_probs[:, :4], torch.tensor( picked_labels), picked_box_probs[:, 4] class_names = [name.strip() for name in open(label_path).readlines()] if ~PRED: net = create_mobilenetv3_ssd_lite(len(class_names), is_test=True) net.load(model_path) dummy_input = torch.randn(1, 3, config.image_size, config.image_size) input_names = ["input0"] output_names = ["output0"] # print("strt input") torch_out = torch.onnx._export( net, # model being run dummy_input, # model input (or a tuple for multiple inputs) "mobilenet_v3_ssd_lite.onnx", # where to save the model (can be a file or file-like object) verbose=False, export_params=True, input_names=input_names, output_names=output_names
config = vgg_ssd_config elif args.net == 'mb1-ssd': create_net = create_mobilenetv1_ssd config = mobilenetv1_ssd_config elif args.net == 'mb1-ssd-lite': create_net = create_mobilenetv1_ssd_lite config = mobilenetv1_ssd_config elif args.net == 'sq-ssd-lite': create_net = create_squeezenet_ssd_lite config = squeezenet_ssd_config elif args.net == 'mb2-ssd-lite': create_net = lambda num: create_mobilenetv2_ssd_lite( num, width_mult=args.mb2_width_mult) config = mobilenetv1_ssd_config elif args.net == 'mb3-ssd-lite': create_net = lambda num: create_mobilenetv3_ssd_lite(num) config = mobilenetv1_ssd_config else: logging.fatal("The net type is wrong.") parser.print_help(sys.stderr) sys.exit(1) train_transform = TrainAugmentation(config.image_size, config.image_mean, config.image_std) target_transform = MatchPrior(config.priors, config.center_variance, config.size_variance, 0.5) test_transform = TestTransform(config.image_size, config.image_mean, config.image_std) logging.info("Prepare training datasets.") datasets = []
logger.info(f"Use GPU {DEVICE}({args.gpu_ids}).") if args.net == 'mb1-ssd': create_net = create_mobilenetv1_ssd config = mobilenetv1_ssd_config elif args.net == 'mb1-ssd-lite': create_net = create_mobilenetv1_ssd_lite config = mobilenetv1_ssd_config elif args.net == 'mb2-ssd-lite': create_net = lambda num: create_mobilenetv2_ssd_lite( num, width_mult=args.mb2_width_mult) config = mobilenetv1_ssd_config elif args.net == 'mb3s-ssd-lite': create_net = lambda num: create_mobilenetv3_ssd_lite( num, model_mode="SMALL", width_mult=args.mb3_width_mult, device=DEVICE) create_net_test = lambda num: create_mobilenetv3_ssd_lite( num, model_mode="SMALL", width_mult=args.mb3_width_mult, is_test=True, device=DEVICE) config = mobilenetv1_ssd_config elif args.net == 'mb3l-ssd-lite': create_net = lambda num: create_mobilenetv3_ssd_lite( num, model_mode="LARGE", width_mult=args.mb3_width_mult, device=DEVICE)
def main(args): net_type = args.net_type img_folder=args.img_folder model_path = args.weights_path label_path = args.label_path class_names = [name.strip() for name in open(label_path).readlines()] out_path=args.out_path if not os.path.exists(out_path): os.mkdir(out_path) num_gpus=torch.cuda.device_count() device='cuda' if num_gpus else 'cpu' if net_type == 'vgg16-ssd': net = create_vgg_ssd(len(class_names), is_test=True) elif net_type == 'mb1-ssd': net = create_mobilenetv1_ssd(len(class_names), is_test=True) elif net_type == 'mb1-ssd-lite': net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True) elif net_type == 'mb2-ssd-lite': net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True) elif net_type == 'sq-ssd-lite': net = create_squeezenet_ssd_lite(len(class_names), is_test=True) elif net_type == 'mb3-ssd-lite': net = create_mobilenetv3_ssd_lite(len(class_names), is_test=True) else: print("The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite.") sys.exit(1) net.load(model_path) if net_type == 'vgg16-ssd': predictor = create_vgg_ssd_predictor(net, candidate_size=20,device=device) elif net_type == 'mb1-ssd': predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=20,device=device) elif net_type == 'mb1-ssd-lite': predictor = create_mobilenetv1_ssd_lite_predictor(net, candidate_size=20,device=device) elif net_type == 'mb2-ssd-lite': predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=20,device=device) elif net_type == 'sq-ssd-lite': predictor = create_squeezenet_ssd_lite_predictor(net, candidate_size=20,device=device) elif net_type == 'mb3-ssd-lite': predictor = create_mobilenetv3_ssd_lite_predictor(net, candidate_size=10,device=device) else: print("The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite.") sys.exit(1) timer = Timer() img_names=glob.glob(img_folder+os.sep+'*.jpg') if len(img_names)==0: print('No imagesfound in {}'.format(img_folder)) exit(-1) for img_name in img_names: image=cv2.imread(img_name) timer.start() boxes,labels,probs=predictor.predict(image,10,0.3) interval = timer.end() print('Time: {:.2f}s, Detect Objects: {:d}.'.format(interval, labels.size(0))) label_text=[] for i in range(boxes.size(0)): box=boxes[i,:] label = f"{class_names[labels[i]]}: {probs[i]:.2f}" label_text.append(label) cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (255, 255, 0), 4) cv2.putText(image, label,(box[0]+20, box[1]+40),cv2.FONT_HERSHEY_SIMPLEX,1, (255, 0, 255), 2) if args.store_result: new_name='{}/{}'.format(out_path,img_name.split('/')[-1]) cv2.imwrite(new_name,image) if not label_text: result_label='empty' else: result_label=label_text[0] with open(os.path.join(out_path,'rest_result.csv'),'a+') as result_writer: result_writer.write(img_name.split('/')[-1]+','+result_label+'\n') cv2.imshow('result', image) if cv2.waitKey(1) & 0xFF == ord('q'): break cv2.destroyAllWindows()
def main(args): net_type = args.net_type model_path = args.weights_path label_path = args.label_path class_names = [name.strip() for name in open(label_path).readlines()] num_classes = len(class_names) if args.live: cap = cv2.VideoCapture(0) cap.set(3, 640) cap.set(4, 480) else: cap = cv2.VideoCapture(args.video_path) out_video = args.out_video Fourcc = cv2.VideoWriter_fourcc('M', 'P', '4', 'V') writer = cv2.VideoWriter(out_video, fourcc=Fourcc, fps=15, frameSize=(640, 480)) num_gpus = torch.cuda.device_count() device = 'cuda' if num_gpus else 'cpu' if net_type == 'vgg16-ssd': net = create_vgg_ssd(len(class_names), is_test=True) elif net_type == 'mb1-ssd': net = create_mobilenetv1_ssd(len(class_names), is_test=True) elif net_type == 'mb1-ssd-lite': net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True) elif net_type == 'mb2-ssd-lite': net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True) elif net_type == 'sq-ssd-lite': net = create_squeezenet_ssd_lite(len(class_names), is_test=True) elif net_type == 'mb3-ssd-lite': net = create_mobilenetv3_ssd_lite(len(class_names), is_test=True) else: print( "The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite." ) sys.exit(1) net.load(model_path) if net_type == 'vgg16-ssd': predictor = create_vgg_ssd_predictor(net, candidate_size=200) elif net_type == 'mb1-ssd': predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=200) elif net_type == 'mb1-ssd-lite': predictor = create_mobilenetv1_ssd_lite_predictor(net, candidate_size=200) elif net_type == 'mb2-ssd-lite': predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200) elif net_type == 'sq-ssd-lite': predictor = create_squeezenet_ssd_lite_predictor(net, candidate_size=200) elif net_type == 'mb3-ssd-lite': predictor = create_mobilenetv3_ssd_lite_predictor(net, candidate_size=10, device=device) else: print( "The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite." ) sys.exit(1) timer = Timer() while True: ret, orig_image = cap.read() if orig_image is None: print('END') break #continue image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB) timer.start() boxes, labels, probs = predictor.predict(image, 10, 0.3) interval = timer.end() print('Time: {:.2f}s, Detect Objects: {:d}.'.format( interval, labels.size(0))) for i in range(boxes.size(0)): box = boxes[i, :] label = f"{class_names[labels[i]]}: {probs[i]:.2f}" cv2.rectangle(orig_image, (box[0], box[1]), (box[2], box[3]), (255, 255, 0), 4) cv2.putText( orig_image, label, (box[0] + 20, box[1] + 40), cv2.FONT_HERSHEY_SIMPLEX, 1, # font scale (255, 0, 255), 2) # line type writer.write(orig_image) cv2.imshow('annotated', orig_image) if cv2.waitKey(1) & 0xFF == ord('q'): break cap.release() writer.release() cv2.destroyAllWindows()