''' author ： Hatcher W Email： hcwang@semi.ac.cn use: 使用SSD算法对大数据进行筛选，先使用Hatcher is cool 将文件夹中所有的视频和图像的路径进行递归存放在videoroot.json和imageroot.json文件中 input:两个json文件 output：相同路径（改为data）有动物图片的结果存储 ''' import json import glob import os import time import torch from PIL import Image from vizer.draw import draw_boxes from ssd.config import cfg from ssd.data.datasets import COCODataset, VOCDataset import argparse import numpy as np from ssd.data.transforms import build_transforms from ssd.modeling.detector import build_detection_model from ssd.utils import mkdir from ssd.utils.checkpoint import CheckPointer from vediotoima import vedio_img os.environ['CUDA_VISIBLE_DEVICES'] = '5' FPS = 5 num_image = 380000 @torch.no_grad() def run_demo(cfg, ckpt, score_threshold, images_dir, output_dir, dataset_type,model): # if dataset_type == "voc": # class_names = VOCDataset.class_names # elif dataset_type == 'coco': # class_names = COCODataset.class_names # else: # raise NotImplementedError('Not implemented now.') # device = torch.device(cfg.MODEL.DEVICE) # # model = build_detection_model(cfg) # model = model.to(device) # checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR) # checkpointer.load(ckpt, use_latest=ckpt is None) # weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file() # print('Loaded weights from {}'.format(weight_file)) # # #image_paths = glob.glob(os.path.join(images_dir, '*.jpg')) global num_image device = torch.device(cfg.MODEL.DEVICE) image_paths = images_dir #mkdir(output_dir) cpu_device = torch.device("cpu") transforms = build_transforms(cfg, is_train=False) # model.eval() start = time.time() image_name = os.path.basename(image_paths) image = np.array(Image.open(image_paths).convert("RGB")) height, width = image.shape[:2] images = transforms(image)[0].unsqueeze(0) load_time = time.time() - start start = time.time() result = model(images.to(device))[0] inference_time = time.time() - start result = result.resize((width, height)).to(cpu_device).numpy() boxes, labels, scores = result['boxes'], result['labels'], result['scores'] indices = scores > score_threshold boxes = boxes[indices] #print(len(boxes)) #exit() labels = labels[indices] scores = scores[indices] meters = ' | '.join( [ 'objects {:02d}'.format(len(boxes)), 'load {:03d}ms'.format(round(load_time * 1000)), 'inference {:03d}ms'.format(round(inference_time * 1000)), 'FPS {}'.format(round(1.0 / inference_time)) ] ) print('({:04d}/{:04d}) {}: {}'.format(1, len(image_paths), image_name, meters)) #drawn_image = draw_boxes(image, boxes, labels, scores, class_names).astype(np.uint8) if(len(boxes)) : print(labels) Image.fromarray(image).save(os.path.join(output_dir,"{}.jpg".format(num_image))) num_image +=1 print(image_name+"is saved") def main(): parser = argparse.ArgumentParser(description="SSD Demo.") parser.add_argument( "--config-file", default="configs/vgg_ssd512_voc0712.yaml", metavar="FILE", help="path to config file", type=str, ) parser.add_argument("--ckpt", type=str, default="../data_MPKFNet/weigth/detection/detection.pth", help="Trained weights.") parser.add_argument("--score_threshold", type=float, default=0.3) parser.add_argument("--images_dir", default='../data_MPKFNet/result_raw/yang/', type=str, help='Specify a image dir to do prediction.') parser.add_argument("--output_dir", default='../data_MPKFNet/result/yangd/', type=str, help='Specify a image dir to save predicted images.') parser.add_argument("--dataset_type", default="voc", type=str, help='Specify dataset type. Currently support voc and coco.') parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() print(args) cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() print("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() print(config_str) print("Running with config:\n{}".format(cfg)) ######################################## with open ("./imageroot.json",'r+') as fvedio: vediolists = json.load(fvedio) fvedio.close() with open("./log_imglast.txt", 'r') as f: last_path = f.readline() print("lastpath is "+last_path) #exit() f.close() flag_detectlast = True ####### if args.dataset_type == "voc": class_names = VOCDataset.class_names elif args.dataset_type == 'coco': class_names = COCODataset.class_names else: raise NotImplementedError('Not implemented now.') device = torch.device(cfg.MODEL.DEVICE) model = build_detection_model(cfg) model = model.to(device) checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR) checkpointer.load(args.ckpt, use_latest=args.ckpt is None) weight_file = args.ckpt if args.ckpt else checkpointer.get_checkpoint_file() print('Loaded weights from {}'.format(weight_file)) model.eval() ##### for vediolist in vediolists: if flag_detectlast: if (vediolist == last_path) or (last_path == ""): flag_detectlast=False print("Detecting finish and begining") continue else: print(vediolist +"has been finished") continue vedio_path = vediolist args.images_dir = vedio_path args.output_dir = "/data1//Hatcher/model/data_MPKFNet/bigdata/ALL_image/" if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) run_demo(cfg=cfg, ckpt=args.ckpt, score_threshold=args.score_threshold, images_dir=args.images_dir, output_dir=args.output_dir, dataset_type=args.dataset_type, model = model ) with open("./log_imglast.txt",'w') as f: f.write(vediolist) f.close() if __name__ == '__main__': main()