IMG_TRANSFORM = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=IMG_NORM_MEAN,
std=IMG_NORM_STDDEV,
inplace=True)])
val_ds = CocoDistillationDatasetAugmented(
COCO_DIR, "val2017", remove_images_without_annotations=False)
img_paths = [os.path.join(COCO_DIR, "images", "val2017",
"{:012d}.jpg".format(x)) for x in val_ds.ids]
# model
hhrnet = get_hrnet_w48_teacher(MODEL_PATH).to(DEVICE)
hhrnet.eval()
hm_parser = HeatmapParser(num_joints=NUM_HEATMAPS,
**HM_PARSER_PARAMS)
# main loop
all_preds = []
all_scores = []
for ii, imgpath in enumerate(img_paths):
img = Image.open(imgpath).convert("RGB")
print(ii, "processing", imgpath, img.size)
resized_img, center, scale = resize_align_multi_scale(np.array(img),
INPUT_SIZE, 1, 1)
t = IMG_TRANSFORM(resized_img).unsqueeze(0).to(DEVICE)
w, h = img.size
# run HHRNet inference
with torch.no_grad():
preds, refined = hhrnet(t)
hms = torch.nn.functional.interpolate(
HALF_PRECISION,
PARAM_INIT_FN,
TRAINABLE_STEM,
BATCHNORM_MOMENTUM)
# load pretrained attention part
# LOAD_TIMESTAMP, LOAD_EPOCH, LOAD_STEP = "21_May_2020_00:02:57.265", 3, 1349
# inpath = os.path.join(SNAPSHOT_DIR, "{}_epoch{}_step{}".format(
# LOAD_TIMESTAMP, LOAD_EPOCH, LOAD_STEP))
# student.load_state_dicts(inpath)
hm_parser = HeatmapParser(num_joints=NUM_HEATMAPS,
**HM_PARSER_PARAMS)
# LOG MODEL AND HYPERPARAMETERS
student_summary = ModuleSummary.get_model_summary(student, as_string=True)
txt_logger.info(student_summary)
tb_logger.add_text("Architecture summary", student_summary, 0)
# tb_logger.add_graph(student, DUMMY_INPUT)
HPARS_DICT = {"num heatmaps": NUM_HEATMAPS, # model arch
"AE dimensions": AE_DIMENSIONS,
"pretrained HRNet path": MODEL_PATH,
"half precision": HALF_PRECISION,
"param init fn": PARAM_INIT_FN,
**HM_PARSER_PARAMS,
# preprocessing/augmentation
"img norm mean": IMG_NORM_MEAN,
img_transform=IMG_NORMALIZE_TRANSFORM,
whitelist_ids=MINIVAL_IDS)
val_augm_dataset = CocoDistillationDatasetAugmented(
COCO_DIR,
"val2017",
os.path.join(COCO_DIR, "hrnet_predictions", "val2017"),
gt_stddevs_pix=[20.0, 9.0, 2.0],
img_transform=IMG_NORMALIZE_TRANSFORM,
overall_transform=OVERALL_HHRNET_TRANSFORM)
hm_parser = HeatmapParser(num_joints=17,
max_num_people=30,
detection_threshold=0.1,
tag_threshold=1.0,
use_detection_val=True,
ignore_too_much=False,
tag_per_joint=True,
nms_ksize=5,
nms_padding=2)
# print("Due to PLT bug type in this line and press c:\n",
# "import matplotlib; matplotlib.use('TkAgg')")
# breakpoint()
# "TRAINING" SET
for i in range(NUM_TRAIN_PLOTS):
print("TRAIN >>>", i)
img_id, img, mask, hms, teach_hms, teach_ae = val_augm_dataset[i]
# plot: img, mask, ground truths, teacher detection
matplotlib.use('TkAgg')