first commit
This commit is contained in:
Semi212 committed
commit
54ccfb11fe
218 files changed
+19525
No files matched your search
Executable
+193
@@ -0,0 +1,193 @@
|
||||
'''
|
||||
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()
|
||||
Reference in new issue
Block a user