194 lines
6.4 KiB
Python
Executable File
194 lines
6.4 KiB
Python
Executable File
'''
|
||
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()
|