def __init__(self, cfg, instance_mode=ColorMode.IMAGE, parallel=False):
"""
Args:
cfg (CfgNode):
instance_mode (ColorMode):
parallel (bool): whether to run the model in different processes from visualization.
Useful since the visualization logic can be slow.
"""
#ipdb.set_trace()
self.metadata = MetadataCatalog.get(
cfg.DATASETS.TEST[0] if len(cfg.DATASETS.TEST) else "__unused")
#ipdb.set_trace()
self.cpu_device = torch.device("cpu")
self.instance_mode = instance_mode
self.parallel = parallel
#print(parallel)
if parallel:
num_gpu = torch.cuda.device_count()
self.predictor = AsyncPredictor(cfg, num_gpus=num_gpu)
else:
self.predictor = DefaultPredictor(cfg)
def __init__(self, args):
self.args = args
use_cuda = bool(strtobool(self.args.use_cuda))
if args.display:
cv2.namedWindow("test", cv2.WINDOW_NORMAL)
cv2.resizeWindow("test", args.display_width, args.display_height)
if not args.image_input:
self.vdo = cv2.VideoCapture()
cfg = get_cfg()
#cfg.merge_from_file("detectron2_repo/configs/COCO-Keypoints/keypoint_rcnn_X_101_32x8d_FPN_3x.yaml")
#cfg.MODEL.WEIGHTS = "detectron2://COCO-Keypoints/keypoint_rcnn_X_101_32x8d_FPN_3x/139686956/model_final_5ad38f.pkl"
cfg.merge_from_file("../detectron2_repo/configs/Misc/cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.yaml")
cfg.MODEL.WEIGHTS = args.detectron2_weights
#"detectron2://Misc/cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv/18131413/model_0039999_e76410.pkl"
cfg.MODEL.MASK_ON = False
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
#cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST = 0.5
self.predictor = DefaultPredictor(cfg)
self.deepsort = DeepSort(args.deepsort_checkpoint, use_cuda=use_cuda, extractor_type=args.extractor_type, game_id=args.game_id, team_0=args.team_0)
def execute(cls: type, args: argparse.Namespace):
logger.info(f"Loading config from {args.cfg}")
opts = []
cfg = cls.setup_config(args.cfg, args.model, args, opts)
logger.info(f"Loading model from {args.model}")
predictor = DefaultPredictor(cfg)
logger.info(f"Loading data from {args.input}")
file_list = cls._get_input_file_list(args.input)
if len(file_list) == 0:
logger.warning(f"No input images for {args.input}")
return
context = cls.create_context(args)
for file_name in file_list:
img = read_image(file_name,
format="BGR") # predictor expects BGR image.
with torch.no_grad():
outputs = predictor(img)["instances"]
cls.execute_on_outputs(context, {
"file_name": file_name,
"image": img
}, outputs)
cls.postexecute(context)
def __init__(self, cfg, instance_mode=ColorMode.IMAGE, parallel=False):
"""
Args:
cfg (CfgNode):
instance_mode (ColorMode):
parallel (bool): whether to run the model in different processes from visualization.
Useful since the visualization logic can be slow.
"""
DatasetCatalog.register("pfallcnt_pred", lambda d: [])
MetadataCatalog.get("pfallcnt_pred").set(thing_classes=["0", "1"],
thing_colors=[(0, 255, 0),
(255, 0, 0)])
self.metadata = MetadataCatalog.get("pfallcnt_pred")
self.cpu_device = torch.device("cpu")
self.instance_mode = instance_mode
self.parallel = parallel
if parallel:
num_gpu = torch.cuda.device_count()
self.predictor = AsyncPredictor(cfg, num_gpus=num_gpu)
else:
self.predictor = DefaultPredictor(cfg)
def __init__(self, config_file, model_file=None):
try:
from detectron2 import __version__ as version
if version != Detectron2Server._supported_version:
raise RuntimeError(
"Supported version is '{}', you have '{}'.".format(
self._supported_version, version))
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.engine.defaults import DefaultPredictor
except ImportError:
raise
self.currently_busy = Event()
self.cfg = get_cfg()
self.classes = _unknown_class()
try:
from fruit_detection.config import add_fruit_detection_config
from fruit_detection.datasets import register_data_set
add_fruit_detection_config(self.cfg)
self.cfg.merge_from_file(config_file)
register_data_set(self.cfg.DATASETS.TEST[0])
DatasetCatalog.get(self.cfg.DATASETS.TEST[0])
metadata = MetadataCatalog.get(self.cfg.DATASETS.TEST[0] if len(
self.cfg.DATASETS.TEST) else "__unused")
self.classes = metadata.get("thing_classes") or _unknown_class()
except ImportError:
self.cfg.merge_from_file(config_file)
if model_file is not None:
self.cfg.MODEL.WEIGHTS = model_file
self.cfg.freeze()
self.predictor = DefaultPredictor(self.cfg)
# Base class must be called at the end due to self.service_server.spin()
BaseDetectionServer.__init__(self)
def __init__(self, cfg, instance_mode=ColorMode.IMAGE, parallel=False):
"""
Args:
cfg (CfgNode):
instance_mode (ColorMode):
parallel (bool): whether to run the model in different processes from visualization.
Useful since the visualization logic can be slow.
"""
self.metadata = MetadataCatalog.get("__unused")
unified_label_file = json.load(
open(cfg.MULTI_DATASET.UNIFIED_LABEL_FILE))
self.metadata.thing_classes = [
'{}'.format([xx for xx in x['name'].split('_') if xx != ''][0]) \
for x in unified_label_file['categories']]
self.cpu_device = torch.device("cpu")
self.instance_mode = instance_mode
self.parallel = parallel
if parallel:
num_gpu = torch.cuda.device_count()
self.predictor = AsyncPredictor(cfg, num_gpus=num_gpu)
else:
self.predictor = DefaultPredictor(cfg)
def execute(args: dict):
print(f"Loading config from {args['cfg']}")
cfg = setup_config(args['cfg'], args['model'], args)
print(f"Loading model from {args['model']}")
predictor = DefaultPredictor(cfg)
print(f"Loading data from {args['input']}")
file_list = _get_input_file_list(args['input'])
if len(file_list) == 0:
print(f"No input images for {args['input']}")
return
context = create_context(args)
for file_name in file_list:
# predictor expects BGR image.
img = read_image(file_name, format="BGR")
with torch.no_grad():
outputs = predictor(img)["instances"]
execute_on_outputs(context, {
"file_name": file_name,
"image": img
}, outputs)
return context["results"]
def execute2(args: dict):
print(f"Loading config from {args['cfg']}")
cfg = setup_config(args['cfg'], args['model'], args)
print(f"Loading model from {args['model']}")
predictor = DefaultPredictor(cfg)
print(f"Loading data from {args['input']}")
pils = [args['input']]
if len(pils) == 0:
print(f"No input images for {args['input']}")
return
context = create_context(args)
for pil in pils:
# predictor expects BGR image.
img = convert_PIL_to_numpy(pil, format="BGR")
with torch.no_grad():
outputs = predictor(img)["instances"]
execute_on_outputs(context, {
"file_name": "pil",
"image": img
}, outputs)
return context["results"]
def main():
# Globals
global model, metadata
# Setup logger
mp.set_start_method("spawn", force=True)
setup_logger(name="fvcore")
logger = setup_logger()
logger.info("Arguments: " + str(args))
# Setup config and predictor
cfg = setup_cfg(args)
# model = VisualizationDemo(cfg)
model = DefaultPredictor(cfg)
metadata = MetadataCatalog.get(
cfg.DATASETS.TEST[0] if len(cfg.DATASETS.TEST) else "__unused"
)
# Run http app
if args.ngrok:
app.run()
else:
app.run(port=PORT, host=DOMAIN, debug=False)
def load_trained_nn(pickle_file):
"""
This function loads the trained nearual network from
the local .pkl file as a predictor
Input:
- pickle_file: .pkl file containing the configurations
of the neural network that can be loaded
into Detectron2 Default Predictor
Output:
- predictor: Detectron2 DefualtPredictor with the
weights ad parameters from the pickle
file
"""
# Load the configurations from the .pkl file
with open(pickle_file, 'rb') as file:
# Assert the input type
assert os.path.splitext(
pickle_file
)[1] == ".pkl", "Input has the wrong file type. A pickle file is required."
cfg = pickle.load(file)
# Use Default predictor to load the pretrained NN
predictor = DefaultPredictor(cfg)
return predictor
def do_panoptic_test(cfg, model):
categoryies=json.load(open("../data/panoptic_coco_categories.json",'r'))
os.makedirs("../data/detectron2_panoptic", exist_ok=True)
categoryies_dict={}
for category in categoryies:
categoryies_dict[category['id']]=category
id_generator=IdGenerator(categoryies_dict)
image_dict = {}
error_list=[]
for dataset_name in ['viroi_test']:#,'viroi_train']:#cfg.DATASETS.TRAIN:#cfg.DATASETS.TEST:
# data_loader = build_detection_test_loader(cfg, dataset_name)
thing_id_map = MetadataCatalog.get(dataset_name).get("thing_contiguous_id_to_class_id")
stuff_id_map = MetadataCatalog.get(dataset_name).get("stuff_contiguous_id_to_class_id")
test_images_dict=json.load(open(MetadataCatalog.get(dataset_name).get("instance_json_file"),'r'))
image_path = MetadataCatalog.get(dataset_name).get("image_path")
predictor=DefaultPredictor(cfg)
total = len(test_images_dict)
count=0
# for idx, inputs in enumerate(data_loader):
for image_id in test_images_dict:
image_info=test_images_dict[image_id]
img=read_image(image_path+"/"+image_info['image_name'],format="BGR")
count+=1
print(str(count)+"/"+str(total))
if True:
# try:
# print(inputs[0])
# predictions = model(inputs, "panoptic")[0] # 'sem_seg', 'instances', 'panoptic_seg'
predictions = predictor(img,0)
panoptic_seg, segments_info = predictions["panoptic_seg"] # seg, info
panoptic_seg=panoptic_seg.data.cpu().numpy()
panoptic_color_seg = np.zeros((panoptic_seg.shape[0], panoptic_seg.shape[1], 3)) #tensor
instance_dict={}
for info in segments_info:
if 'score' in info:
del info['score']
if 'area' in info:
del info['area']
bbox = info['box'] # x1,y1,x2,y2->y1,x1,y2,x2
info['bbox'] = [int(bbox[1]), int(bbox[0]), int(bbox[3]), int(bbox[2])]
del info['box']
class_id=info['class_id']
del info['category_id']
mask = info['mask'].data.cpu().numpy()
mask = np.asfortranarray(mask)
segmentation = maskUtils.encode(mask)
segmentation['counts'] = segmentation['counts'].decode('utf8')
info['segmentation'] = segmentation
instance_id, panoptic_color_seg[mask] = id_generator.get_id_and_color(categoryies[class_id - 1]['id'])
info['instance_id'] = instance_id
del info['mask']
instance_dict[str(instance_id)]=info
image_dict[image_id]={'instances':instance_dict,
"image_id":image_info['image_id'],
"height":image_info['height'],
"width":image_info['width'],
"image_name":image_info['image_name']
}
# print(image_dict)
Image.fromarray(panoptic_color_seg.astype(np.uint8)).save("../data/detectron2_panoptic/"+image_info["image_name"].replace("jpg","png"))
# except:
# print("ERROR - "+image_info['image_name'])
# error_list.append(image_info['image_name'])
json.dump(image_dict,open("../data/viroi_json/detectron2_"+dataset_name+"_images_dict.json",'w'))
json.dump(error_list,open("error_list.json",'w'))
def __init__(self, cfg):
self.predictor = DefaultPredictor(cfg)
super().__init__()
type=float,
default=0.5,
help="Minimum score for instance predictions to be shown",
)
parser.add_argument(
"--opts",
help=
"Modify config options using the command-line 'KEY VALUE' pairs",
default=[],
nargs=argparse.REMAINDER,
)
return parser
args = get_parser().parse_args()
cfg = setup_cfg(args)
model = DefaultPredictor(cfg)
# Load custom dataset
DatasetCatalog.register("my_dataset", my_dataset)
MetadataCatalog.get("my_dataset").thing_classes = [
'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train',
'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign',
'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep',
'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella',
'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork',
'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv',
'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
def init_model(init_args, context):
os.environ['CUDA_VISIBLE_DEVICES'] = str(context['gpu'])
context['predictor'] = DefaultPredictor(init_args['config'])
def predict(self, img):
predictor = DefaultPredictor(self.cfg)
with torch.no_grad():
outputs = predictor(
img)["instances"].pred_keypoints.cpu().detach()[0]
return KeyPointPredictor.format2viton(outputs)
def __init__(
self,
model=MODELS["chimps"],
model_cache_dir: Path = Path(".zamba_cache"),
download_region=RegionEnum("us"),
):
"""Create a DensePoseManager object.
Parameters
----------
model : dict, optional (default MODELS['chimps'])
A dictionary with the densepose model defintion like those defined in MODELS.
"""
if not DENSEPOSE_AVAILABLE:
raise ImportError(
"Densepose not installed. See: https://zamba.drivendata.org/docs/stable/models/densepose/#installation"
)
# setup configuration for densepose
self.cfg = get_cfg()
add_densepose_config(self.cfg)
self.cfg.merge_from_file(model["config"])
if not (model_cache_dir / model["weights"]).exists():
model_cache_dir.mkdir(parents=True, exist_ok=True)
self.cfg.MODEL.WEIGHTS = download_weights(model["weights"],
model_cache_dir,
download_region)
# automatically use CPU if no cuda available
if not torch.cuda.is_available():
self.cfg.MODEL.DEVICE = "cpu"
self.cfg.freeze()
logging.getLogger("fvcore").setLevel(
"CRITICAL") # silence noisy detectron2 logging
# set up predictor with the configuration
self.predictor = DefaultPredictor(self.cfg)
# we have a specific texture atlas for chimps with relevant regions
# labeled that we can use instead of the default segmentation
self.visualizer = model["viz_class"](
self.cfg,
device=self.cfg.MODEL.DEVICE,
**model.get("viz_class_kwargs", {}),
)
# set up utilities for use with visualizer
self.vis_extractor = create_extractor(self.visualizer)
self.vis_embedder = build_densepose_embedder(self.cfg)
self.vis_class_to_mesh_name = get_class_to_mesh_name_mapping(self.cfg)
self.vis_mesh_vertex_embeddings = {
mesh_name: self.vis_embedder(mesh_name).to(self.cfg.MODEL.DEVICE)
for mesh_name in self.vis_class_to_mesh_name.values()
if self.vis_embedder.has_embeddings(mesh_name)
}
if "anatomy_color_mapping" in model:
self.anatomy_color_mapping = pd.read_csv(
model["anatomy_color_mapping"], index_col=0)
else:
self.anatomy_color_mapping = None
def get_human_model(nms_thres=0.0,score_thres=0.995):
cfg_human.MODEL.ROI_HEADS.NMS_THRESH_TEST = nms_thres #IoU aka overlap suppression (suppress if overlap > threshold)
cfg_human.MODEL.ROI_HEADS.SCORE_THRESH_TEST = score_thres #this isnt a confidence score filter... but seems to correlate well anyways
human_model = DefaultPredictor(cfg_human)
return human_model
def get_clothes_model(nms_thres=0.2,score_thres=0.6):
cfg_clothes.MODEL.ROI_HEADS.NMS_THRESH_TEST = nms_thres
cfg_clothes.MODEL.ROI_HEADS.SCORE_THRESH_TEST = score_thres
clothes_model = DefaultPredictor(cfg_clothes)
return clothes_model
def main():
parser = argparse.ArgumentParser()
## Required parameters
parser.add_argument(
"--model",
default='helmet_tracking/model_final_13000.pth',
type=str, #required=True,
help="path to base D2 model or s3 location")
parser.add_argument(
"--reid_model",
default='helmet_tracking/ResNet_iter_25137.pth',
type=str, #required=True,
help="path to reid model or s3 location")
parser.add_argument(
"--output_dir",
default='/home/model_results',
type=str, #required=True,
help=
"The output directory where the model checkpoints and predictions will be written."
)
parser.add_argument("--bucket",
default='privisaa-bucket-virginia',
type=str)
parser.add_argument(
"--img_paths",
default='nfl-data/live_video',
type=str,
#required=True,
help="path to images or image location in s3")
parser.add_argument(
"--conf_thresh",
default=.5,
type=float, #required=True,
help="base D2 model")
parser.add_argument("--use_mask", default=0, type=int)
# parser.add_argument("--video", default=None, type=str, required=True, help="path to video for tracking job")
parser.add_argument(
"--d2_config",
default=
'/home/detectron2-ResNeSt/configs/quick_schedules/mask_rcnn_R_50_FPN_inference_acc_test.yaml',
type=str,
help="Detectron2 config file")
args = parser.parse_args()
#COCO-Detection/faster_cascade_rcnn_ResNeSt_101_FPN_syncbn_range-scale_1x.yaml'
confidence_threshold = .55
cfg = get_cfg()
config_file = args.d2_config
model_pth = args.model
reid_pth = args.reid_model
try:
os.mkdir('/home/video')
except:
pass
if os.path.exists(args.model) == False:
print('\n')
print('downloading d2 model from s3')
s3.download_file(args.bucket, args.model, '/home/d2_final.pth')
model_pth = '/home/d2_final.pth'
if os.path.exists(args.reid_model) == False:
print('\n')
print('downloading reid model from s3')
s3.download_file(args.bucket, args.reid_model, '/home/reid_final.pth')
reid_pth = '/home/reid_final.pth'
if os.path.exists(args.img_paths) == False:
objs = s3.list_objects(Bucket='privisaa-bucket-virginia',
Prefix=args.img_paths)['Contents']
keys = []
for key in objs:
if key['Size'] > 0:
keys.append(key['Key'])
folders = []
for key in keys:
folders.append(key.split('/')[-2])
folder = list(np.unique(folders))[-1]
print('\n')
print('Loading images from this video: ', folder)
for key in tqdm(keys):
if key.split('/')[-2] == folder:
s3.download_file(args.bucket, key,
f"/home/video/{key.split('/')[-1]}")
cfg.merge_from_file(config_file)
cfg.MODEL.RETINANET.SCORE_THRESH_TEST = args.conf_thresh
cfg.MODEL.WEIGHTS = model_pth #'/home/ubuntu/finetuned-fasterrcnn-cascade-d2-resnest-13000imgs-005lr-1class-tune2/model_final.pth'
#finetuned-fasterrcnn-cascade-d2-resnest-13000imgs-02lr/model_final.pth'
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = args.conf_thresh
cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = args.conf_thresh
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
model = DefaultPredictor(cfg)
#model.model.load_state_dict(torch.load('finetuned-detectron2-maskrcnn_13000imgs_03lr_3rdRun/model_final.pth')['model'])
eval_results = model.model.eval(
) # finetuned-detectron2-maskrcnn_fixed/model_final.pth')['model'])
# 'finetuned-detectron2-maskrcnn_fixed/model_final.pth'
# finetuned-detectron2-maskrcnn_5700imgs_025lr_furthertuned/model_final.pth
if args.use_mask == 0:
obj_detect = MRCNN_FPN(model, num_classes=1)
else:
obj_detect = FRCNN_FPN(model, num_classes=1)
#obj_detect.to(device)
# reid
reid_network = resnet50(pretrained=False,
output_dim=128) # need to insert dictionary here
reid_network.load_state_dict(
torch.load(
reid_pth, #tracktor['reid_weights'],'/home/ubuntu/code/tracking_wo_bnw/output/tracktor/reid/test/ResNet_iter_25137.pth'
map_location=device)) #lambda storage, loc: storage))
reid_network.eval()
#reid_network.cuda()
reid_network.to(device)
tracker = {
# FRCNN score threshold for detections
'detection_person_thresh': 0.5,
# FRCNN score threshold for keeping the track alive
'regression_person_thresh': 0.5,
# NMS threshold for detection
'detection_nms_thresh': 0.3,
# NMS theshold while tracking
'regression_nms_thresh': 0.6,
# motion model settings
'motion_model': {
'enabled': True,
# average velocity over last n_steps steps
'n_steps': 2,
# if true, only model the movement of the bounding box center. If false, width and height are also modeled.
'center_only': True
},
# DPM or DPM_RAW or 0, raw includes the unfiltered (no nms) versions of the provided detections,
# 0 tells the tracker to use private detections (Faster R-CNN)
'public_detections': False,
# How much last appearance features are to keep
'max_features_num': 20,
# Do camera motion compensation
'do_align': False,
# Which warp mode to use (cv2.MOTION_EUCLIDEAN, cv2.MOTION_AFFINE, ...)
'warp_mode': str(cv2.MOTION_EUCLIDEAN),
# maximal number of iterations (original 50)
'number_of_iterations': 500,
# Threshold increment between two iterations (original 0.001)
'termination_eps': 0.00001,
# Use siamese network to do reid
'do_reid': True,
# How much timesteps dead tracks are kept and cosidered for reid
'inactive_patience': 15,
# How similar do image and old track need to be to be considered the same person
'reid_sim_threshold': 20.0,
# How much IoU do track and image need to be considered for matching
'reid_iou_threshold': 0.05
}
tracker = IJP_tracker(obj_detect, reid_network, tracker)
# tracker = Tracker(obj_detect, reid_network, tracker)
def transform_img(i, pth='/home/ubuntu/videos/'):
if (i < 10):
ind = f'0000{i}'
elif (i < 100):
ind = f'000{i}'
else:
ind = f'00{i}'
frame = Image.open(f'{pth}split_frames_{ind}.jpg')
frame_ten = torch.tensor(np.reshape(np.array(frame),
(1, 3, 720, 1280)),
device=device,
dtype=torch.float32)
blob = {'img': frame_ten}
return blob
# cannot make frame size assumption...
# need to make this an argparse argument
pth = '/home/video/'
# vid_name = args.video #'57675_003286_Endzone'
img_paths = glob(f'{pth}/*') # {vid_name}
img_paths.sort()
# print(img_paths)
tracker.reset()
# results = []
# print('Starting tracking!')
# for i in tqdm(range(1,len(img_paths)+1)):
# #tracker.reset()
# blob = transform_img(i ,pth=pth)
# with torch.no_grad():
# tracker.step(blob)
track_dict = {}
print("\n")
print("###########################################")
print("############## BEGIN TRACKING #############")
print("###########################################")
print('\n')
for i in tqdm(range(1, len(img_paths) + 1)):
track_dict[i] = {}
blob = transform_img(i, pth=pth)
with torch.no_grad():
tracker.step(blob)
for tr in tracker.tracks:
track_dict[i][tr.id] = tr.pos[0].detach().cpu().numpy()
iou_dict = {}
iou_dict2 = {}
new_tracks = {}
missing_tracks = {}
for track in tqdm(track_dict):
if track == 1: # set dictionaries with 1st frame
new_tracks[track] = list(track_dict[track].keys())
missing_tracks[track] = set()
iou_dict[track] = {}
iou_dict[track] = {}
for tr in track_dict[track]:
iou_dict[track][tr] = {}
else:
new_tracks[track] = set(list(track_dict[track].keys())) - set(
new_tracks[1]) # - set(list(track_dict[track-1].keys()))
missing_tracks[track] = set(new_tracks[1]) - set(
list(track_dict[track - 1].keys()))
iou_dict[track] = {}
iou_dict2[track] = {}
for tr in track_dict[track]:
iou_dict[track][tr] = {}
iou_dict2[track][tr] = {}
for t in track_dict[track - 1]:
iou = bbox_iou(track_dict[track][tr],
track_dict[track - 1][t])
if iou > 0.001:
iou_dict[track][tr][t] = iou
if track > 2:
for t in track_dict[track - 2]:
iou = bbox_iou(track_dict[track][tr],
track_dict[track - 2][t])
if iou > 0.001:
iou_dict2[track][tr][t] = iou
tracks_to_delete = {}
tracks_to_change = {}
momentum_dict = {}
for track in track_dict:
if track == 1:
for tr in track_dict[track]:
momentum_dict[tr] = 0 # initialize momentum dict
elif len(track_dict[track]
) > 22: # if there are more than 22 annotations
tracks_to_delete[track] = []
for tr in track_dict[track]:
try:
momentum_dict[tr] += 1
except:
momentum_dict[tr] = 0
for ind in iou_dict[track]:
if (iou_dict[track][ind] == {}) & (
ind > 22
): # need to adjust this, right now just looking if there is ANY IoU
tracks_to_delete[track].append(ind)
else: # if less than 22 annotations
tracks_to_change[track] = {}
if new_tracks[track] != set(
): # if there are new tracks, check them
for tr in track_dict[track]: # update momentum dict
try:
momentum_dict[tr] += 1
except:
momentum_dict[tr] = 0
for newt in new_tracks[track]: # cycle through new tracks
# print('For track ',track,"and ID ",newt,iou_dict[track][newt])
if newt > 22: # if there is a new track and it's greater than 22, change it, figure out what to change to
if (missing_tracks[track] != set()) | (
missing_tracks[track - 1] != set()
): # if there are missing tracks, cycle through them and compare
mis_iou = {}
if missing_tracks[track] != set():
for mis in missing_tracks[track]:
try:
# iou = bbox_iou(track_dict[track-1][mis], track_dict[track][newt])
dist = distance.euclidean(
track_dict[track - 1][mis],
track_dict[track][newt])
except:
try:
# iou = bbox_iou(track_dict[track-2][mis], track_dict[track][newt])
dist = distance.euclidean(
track_dict[track - 2][mis],
track_dict[track][newt])
except:
try:
# iou = bbox_iou(track_dict[track-3][mis], track_dict[track][newt])
dist = distance.euclidean(
track_dict[track - 3][mis],
track_dict[track][newt])
except:
pass
mis_iou[mis] = dist
# tracks_to_change[track][newt] =
tracks_to_change[track][newt] = mis_iou
# if missing_tracks[track-1]!=set():
# for mis in missing_tracks[track-1]:
# iou = bbox_iou(track_dict[track-2][mis], track_dict[track][newt])
# mis_iou[mis] = iou
# tracks_to_change[track][newt] = mis_iou
# try:
# if max(tracks_to_change[track][newt].values())>.2:
# for t in tracks_to_change[trc][tr]:
to_ind = np.argmin(
list(tracks_to_change[track]
[newt].values()))
to_id = list(tracks_to_change[track]
[newt].keys())[to_ind]
to_pos = track_dict[track][newt]
del track_dict[track][newt]
track_dict[track][to_id] = to_pos
# need to send results to s3
result = tracker.get_results()
file = 'tracking_results.json'
tracking_dict = {}
for res in result:
tracking_dict[res] = {}
for r in result[res]:
tracking_dict[res][r] = list(result[res][r][0:4])
with open(file, 'w') as f:
json.dump(tracking_dict, f)
# now = str(datetime.datetime.now()).replace(' ','').replace(':','-')
k = f'nfl-data/tracking_results_{folder}.json'
s3.upload_file(Filename=file, Bucket=args.bucket, Key=k)
print(f'Tracking finished and results saved to: s3://{args.bucket}/{k}')
os.makedirs('/home/labeled_frames')
# create labeled images
print('\n')
print("###########################################")
print("############ Generating Video! ############")
print("###########################################")
print('\n')
print('...')
for j, pth in tqdm(enumerate(img_paths)):
fig, ax = plt.subplots(1, figsize=(24, 14))
img = Image.open(pth)
# Display the image
ax.imshow(np.array(img))
# Create a Rectangle patch
label_list = {}
for r in track_dict[j + 1]:
try:
res = track_dict[j + 1][r]
label_list[r] = res
except:
pass
for i, r in enumerate(label_list):
labs = label_list[r]
rect = patches.Rectangle((labs[0], labs[1]),
labs[2] - labs[0],
labs[3] - labs[1],
linewidth=1,
edgecolor='r',
facecolor='none') # 50,100),40,30
ax.add_patch(rect)
plt.text(labs[0] - 10, labs[1] - 10, f'H:{r}', fontdict=None)
plt.savefig(
f"/home/labeled_frames/{pth.split('/')[-1].replace('.jpg','.png')}"
)
plt.close()
# create video of labels
os.system(
'ffmpeg -r 15 -f image2 -s 1280x720 -i /home/labeled_frames/split_frames_%05d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p /home/labeled_frames.mp4'
)
k = f'nfl-data/tracking_results_{folder}.mp4'
s3.upload_file(Filename='/home/labeled_frames.mp4',
Bucket=args.bucket,
Key=k)
print('\n')
print(f'Video uploaded to: s3://{args.bucket}/{k}')
# for launching A2I job set a conditional here
s3_fname = f's3://{args.bucket}/{k}'
workteam = 'arn:aws:sagemaker:us-east-1:209419068016:workteam/private-crowd/ijp-private-workteam'
flowDefinitionName = 'ijp-video-flow-official'
humanTaskUiArn = 'arn:aws:sagemaker:us-east-1:209419068016:human-task-ui/ijp-video-task3'
#'s3://privisaa-bucket-virginia/nfl-data/nfl-frames/nfl-video-frame00001.png'
# create_workflow_definition_response = sagemaker_client.create_flow_definition(
# FlowDefinitionName= flowDefinitionName,
# RoleArn= role,
# HumanLoopConfig= {
# "WorkteamArn": workteam,
# "HumanTaskUiArn": humanTaskUiArn,
# "TaskCount": 1,
# "TaskDescription": "Identify if the labels in the video look correct.",
# "TaskTitle": "Video classification a2i demo"
# },
# OutputConfig={
# "S3OutputPath" : OUTPUT_PATH
# }
# )
# flowDefinitionArn = create_workflow_definition_response['FlowDefinitionArn']
inputContent = {
"initialValue": .2,
"taskObject":
s3_fname # the s3 object will be passed to the worker task UI to render
}
now = str(datetime.datetime.now()).replace(' ', '-').replace(':',
'-').replace(
'.', '-')
response = a2i.start_human_loop(
HumanLoopName=f'ijp-video-{now}',
FlowDefinitionArn=flowDefinitionArn,
HumanLoopInput={"InputContent": json.dumps(inputContent)},
DataAttributes={
'ContentClassifiers': ['FreeOfPersonallyIdentifiableInformation']
})
print(f'Launched A2I loop ijp-video-{now}')
sns.publish(
TopicArn='arn:aws:sns:us-east-1:209419068016:ijp-topic',
Message=
f'Your video inference is done! You can find the output here: s3://{args.bucket}/{k}',
Subject='Video labeling')
def do_relation_test_demo(cfg,image_path,visible=False,visible_num=5):
predictor = DefaultPredictor(cfg)
img = read_image(image_path, format="BGR")
pred_instances, results_dict = predictor(img, 0, mode="relation")
image_info={}
if len(pred_instances[0]) > 0:
pred_boxes = pred_instances[0].pred_boxes.tensor
height, width = pred_instances[0].image_size
ori_height, ori_width = img.shape[0], img.shape[1]
pred_classes = pred_instances[0].pred_classes
pred_boxes = torch.stack([pred_boxes[:, 1] * ori_height * 1.0 / height,
pred_boxes[:, 0] * ori_width * 1.0 / width,
pred_boxes[:, 3] * ori_height * 1.0 / height,
pred_boxes[:, 2] * ori_width * 1.0 / width], dim=1)
pred_classes = pred_classes.data.cpu().numpy()
pred_boxes = pred_boxes.data.cpu().numpy()
predicate_categories = results_dict['predicate_categories'][0].data.cpu().numpy().reshape(
len(pred_instances[0]), len(pred_instances[0]), 249)
pair_interest_pred = results_dict['pair_interest_pred'][0].data.cpu().numpy().reshape(len(pred_instances[0]),
len(pred_instances[0]))
pair_interest_pred_instance_pair = pair_interest_pred * (1 - np.eye(len(pred_instances[0])))
predicate_factor = pair_interest_pred_instance_pair.reshape(len(pred_instances[0]), len(pred_instances[0]), 1)
single_result = (predicate_factor * predicate_categories).reshape(-1)
single_result_indx = np.argsort(single_result)[::-1][:100]
single_index = []
for i in range(len(pred_instances[0])):
for j in range(len(pred_instances[0])):
for k in range(249):
single_index.append([i, j, k])
single_index = np.array(single_index)
locations = single_index[single_result_indx]
scores = single_result[single_result_indx]
image_info[image_path] = {
"relation_ids": (locations[:, 2] + 1).tolist(),
"subject_class_ids": pred_classes[locations[:, 0]].tolist(),
"subject_boxes": pred_boxes[locations[:, 0]].tolist(),
"object_class_ids": pred_classes[locations[:, 1]].tolist(),
"object_boxes": pred_boxes[locations[:, 1]].tolist(),
"scores": scores.tolist()
}
else:
image_info[image_path] = {
"relation_ids": [],
"subject_class_ids": [],
"subject_boxes": [],
"object_class_ids": [],
"object_boxes": [],
"scores": []
}
if visible:
subject_boxes = image_info[image_path]['subject_boxes']
object_boxes = image_info[image_path]['object_boxes']
subject_boxes_xyxy = []
object_boxes_xyxy = []
for sub_box, obj_box in zip(subject_boxes, object_boxes):
subject_boxes_xyxy.append([sub_box[1], sub_box[0], sub_box[3], sub_box[2]])
object_boxes_xyxy.append([obj_box[1], obj_box[0], obj_box[3], obj_box[2]])
subject_class_ids = image_info[image_path]['subject_class_ids']
object_class_ids = image_info[image_path]['object_class_ids']
scores = image_info[image_path]['scores']
relation_ids = image_info[image_path]['relation_ids']
subject_boxes = np.array(subject_boxes_xyxy)
object_boxes = np.array(object_boxes_xyxy)
subject_class_ids = np.array(subject_class_ids)
object_class_ids = np.array(object_class_ids)
relation_ids = np.array(relation_ids)
scores = np.array(scores)
sort_idx = np.argsort(-scores)[:visible_num]
triplets = Instances((img.shape[0], img.shape[1]))
triplets.subject_classes = torch.Tensor(subject_class_ids[sort_idx]).int() - 1
triplets.object_classes = torch.Tensor(object_class_ids[sort_idx]).int() - 1
triplets.subject_boxes = Boxes(torch.Tensor(subject_boxes[sort_idx, :]))
triplets.object_boxes = Boxes(torch.Tensor(object_boxes[sort_idx, :]))
triplets.relations = torch.Tensor(relation_ids[sort_idx]).int() - 1
triplets.scores = torch.Tensor(scores[sort_idx])
visualizer = Visualizer(img, MetadataCatalog.get(cfg.DATASETS.TEST[0]), instance_mode=ColorMode.IMAGE)
vis_output_instance = visualizer.draw_relation_predictions(triplets)
vis_output_instance.save(os.path.join(image_path.split("/")[-1].split(".")[0] + "_" + str(visible_num) + ".png"))
return image_info
def d2_train_model(regist_train_name, regist_val_name, train_json_path,
train_images_dir, val_json_path, val_images_dir,
ims_per_batch, model_lr, bach_size_per_img, max_train_iter,
num_workers, num_labels):
## 1. models:
model_name = "mask_rcnn_R_50_FPN_3x.yaml"
cfgFile = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
# work_root = os.getcwd()
work_root = "/home/user/qunosen/2_project/4_train/2_zhuchao/6_d2_final_train/2_train/1_train_2020_9_17/test6"
log_file = os.path.join(work_root, "log_dat.txt")
d2_start = time.time()
datetime_now = datetime.datetime.now()
log = ("###" * 100 + "\n") * 5 + " %s\n" % (
str(datetime_now)) + "model_name: %s ..." % model_name
print(log)
print(log, file=open(log_file, "a"))
log = "parameter setting:\n model_to_try:%s\n num_labels: %d\n ims_per_batch:%d\n num_workers:%d\n model_lr:%s\n max_train_iter:%d\n bach_size_per_img:%d\n" % \
(model_name, num_labels, ims_per_batch, num_workers, str(model_lr), max_train_iter, bach_size_per_img)
print(log)
print(log, file=open(log_file, "a"))
new_root = os.path.join(
work_root,
str(model_name) + "_%s_%s_%s_%s" %
(str(model_lr), str(bach_size_per_img), str(max_train_iter),
str(ims_per_batch)))
if not os.path.exists(new_root): os.makedirs(new_root)
os.chdir(new_root)
# register_coco_instances(regist_train_name, {}, train_json_path, train_images_dir)
# register_coco_instances(regist_val_name, {}, val_json_path, val_images_dir)
if regist_train_name in DatasetCatalog._REGISTERED:
log = 'regist_data exists before: %s , and try to del.... ' % regist_train_name
print(log)
print(log, file=open(log_file, "a"))
DatasetCatalog._REGISTERED.pop(regist_train_name)
DatasetCatalog._REGISTERED.pop(regist_val_name)
else:
log = 'regist_data : %s .... ' % regist_train_name
print(log)
print(log, file=open(log_file, "a"))
register_coco_instances(regist_train_name, {}, train_json_path,
train_images_dir)
register_coco_instances(regist_val_name, {}, val_json_path,
val_images_dir)
train_metadata = MetadataCatalog.get(regist_train_name)
val_metadata = MetadataCatalog.get(regist_val_name)
trainset_dicts = DatasetCatalog.get(regist_train_name)
################
# #### trainning:
cfg = get_cfg()
mode_config = cfgFile
log = "model_to_train: %s ..." % mode_config
print(log)
print(log, file=open(log_file, "a"))
cfg.merge_from_file(model_zoo.get_config_file(mode_config))
cfg.DATASETS.TRAIN = (regist_train_name, )
cfg.DATASETS.TEST = (regist_val_name,
) # no metrics implemented for this dataset
cfg.DATALOADER.NUM_WORKERS = num_workers
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
mode_config) ## out_model : ./output/model_final.pth
# cfg.MODEL.WEIGHTS = '/home/user/tmp/pycharm_project_310/1_detectron2/ImageDetectionAPI/d2_object_detection/pre_trained_model/model_final_a54504.pkl'
cfg.SOLVER.IMS_PER_BATCH = ims_per_batch
cfg.SOLVER.BASE_LR = model_lr
cfg.SOLVER.MAX_ITER = (
max_train_iter
) # 300 iterations seems good enough, but you can certainly train longer
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = (
bach_size_per_img) # faster, and good enough for this toy dataset
cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_labels # len(select_cats) # 5 classes ['chair', 'table', 'swivelchair', 'sofa', 'bed']
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
model_path = os.path.join(new_root, 'output/model_final.pth')
if os.path.exists(model_path):
log = "model_save: %s" % model_path
print(log)
print(log, file=open(log_file, "a"))
#### predict
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2 # set the testing threshold for this model
cfg.DATASETS.TEST = (regist_val_name, )
predictor = DefaultPredictor(cfg)
out_model_dir = os.path.join(new_root, "output")
out_dir = os.path.join(out_model_dir, 'result_' + str(model_name))
if not os.path.exists(out_dir): os.makedirs(out_dir)
test_dir = val_images_dir # os.path.join(work_dir,"./val_images")
# test_dir = "/home/user/qunosen/2_project/4_train/2_zhuchao/6_d2_final_train/1_data/1_data_2020_9_17/new_set_6/val"
imgs_list = [
os.path.join(test_dir, file_name)
for file_name in os.listdir(test_dir)
if file_name.endswith(".jpg") or file_name.endswith(".png")
or file_name.endswith(".bmp") or file_name.endswith(".jpeg")
]
for d in imgs_list:
im = cv2.imread(d)
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
metadata=train_metadata,
scale=0.9,
instance_mode=ColorMode.IMAGE_BW)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
predict_file = os.path.join(
out_dir,
os.path.splitext(os.path.basename(d))[0] + "_predict.png")
cv2.imwrite(predict_file, v.get_image()[:, :, ::-1])
if os.path.exists(predict_file):
print("Done: %s" % predict_file)
#### evaluate
evaluator = COCOEvaluator(regist_val_name,
cfg,
False,
output_dir="./output/")
val_loader = build_detection_test_loader(cfg, regist_val_name)
my_eval = inference_on_dataset(trainer.model, val_loader, evaluator)
print(my_eval)
log = ("%s evaluate: \n" % (model_name), my_eval)
print(log, file=open(log_file, "a"))
###############
DatasetCatalog._REGISTERED.pop(regist_train_name)
DatasetCatalog._REGISTERED.pop(regist_val_name)
log = "clean regist_data: %s and %s" % (regist_train_name,
regist_val_name)
print(log)
print(log, file=open(log_file, "a"))
d2_end = time.clock()
log = "model %s : it takes %s ." % (model_name, str(d2_end - d2_start))
print(log)
print(log, file=open(log_file, "a"))
os.chdir(work_root)
else:
print("NotFound: {}".format(model_path))
def predict(self, img):
cfg = self.setup_config()
predictor = DefaultPredictor(cfg)
with torch.no_grad():
outputs = predictor(img)["instances"]
return DensePosePredictor.execute_on_outputs(img, outputs)
nargs=argparse.REMAINDER,
)
return parser
if __name__ == "__main__":
mp.set_start_method("spawn", force=True)
args = get_parser().parse_args()
logger = setup_logger()
logger.info("Arguments: " + str(args))
cfg = setup_cfg(args)
_cpu_device = torch.device("cpu")
# demo = VisualizationDemo(cfg)
predictor = DefaultPredictor(cfg)
# Loading COCO validation images
args.data_type = 'val2017'
if 'test' not in args.data_type:
annotation_file = '{}/annotations/instances_{}.json'.format(
args.coco_path, args.data_type)
dataset_name = 'coco_2017_train'
else:
annotation_file = '{}/annotations/image_info_test-dev2017.json'.format(
args.coco_path)
dataset_name = 'coco_2017_val'
coco = COCO(annotation_file)
imgIds = coco.getImgIds()
seg_results = []
def __init__(self, maxImageDim=320):
super().__init__()
self.predictor = DefaultPredictor(cfg)
self.AREA_THRESHOLD = 40 * 40
self.MaxImageDim = maxImageDim
