|
import argparse
|
|
import os
|
|
import random
|
|
import cv2
|
|
import numpy as np
|
|
import pandas as pd
|
|
import tqdm
|
|
from PIL import Image
|
|
from retinaface.pre_trained_models import get_model
|
|
|
|
|
|
random.seed(1)
|
|
|
|
|
|
def gaussian(x_min, y_min, x_max, y_max):
|
|
x_min, x_max = sorted((x_min, x_max))
|
|
y_min, y_max = sorted((y_min, y_max))
|
|
x_mid, y_mid = (x_min + x_max) / 2, (y_min + y_max) / 2
|
|
x_sigma2, y_sigma2 = (
|
|
np.clip(np.square([x_max - x_min, y_max - y_min], dtype=float), 1, None) / 3
|
|
)
|
|
|
|
def _gaussian(_xs, _ys):
|
|
return np.exp(
|
|
-(np.square(_xs - x_mid) / x_sigma2 + np.square(_ys - y_mid) / y_sigma2)
|
|
)
|
|
|
|
return _gaussian
|
|
|
|
|
|
def plot_ori(label_path, data_dir):
|
|
df = pd.read_csv(
|
|
label_path,
|
|
names=[
|
|
|
|
"path",
|
|
"idx",
|
|
"body_bbox_x",
|
|
"body_bbox_y",
|
|
"body_bbox_w",
|
|
"body_bbox_h",
|
|
"eye_x",
|
|
"eye_y",
|
|
"gaze_x",
|
|
"gaze_y",
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"inout",
|
|
"meta0",
|
|
"meta1",
|
|
],
|
|
index_col=False,
|
|
encoding="utf-8-sig",
|
|
)
|
|
grouped = df.groupby("path")
|
|
|
|
output_dir = os.path.join(data_dir, "head_masks")
|
|
|
|
for image_name, group_df in tqdm.tqdm(grouped, desc="Generating masks with annotations: "):
|
|
if not os.path.exists(os.path.join(output_dir, image_name)):
|
|
w, h = Image.open(image_name).size
|
|
heatmap = np.zeros((h, w), dtype=np.float32)
|
|
indices = np.meshgrid(
|
|
np.linspace(0.0, float(w), num=w, endpoint=False),
|
|
np.linspace(0.0, float(h), num=h, endpoint=False),
|
|
)
|
|
for _, row in group_df.iterrows():
|
|
x_min, y_min, x_max, y_max = (
|
|
row["x_min"],
|
|
row["y_min"],
|
|
row["x_max"],
|
|
row["y_max"],
|
|
)
|
|
gauss = gaussian(x_min, y_min, x_max, y_max)
|
|
heatmap += gauss(*indices)
|
|
heatmap /= np.max(heatmap)
|
|
heatmap_image = Image.fromarray((heatmap * 255).astype(np.uint8), mode="L")
|
|
output_filename = os.path.join(output_dir, image_name)
|
|
os.makedirs(os.path.dirname(output_filename), exist_ok=True)
|
|
heatmap_image.save(output_filename)
|
|
|
|
|
|
def plot_gen(df, data_dir):
|
|
df = df[df["score"] > 0.8]
|
|
grouped = df.groupby("path")
|
|
|
|
output_dir = os.path.join(data_dir, "head_masks")
|
|
|
|
for image_name, group_df in tqdm.tqdm(grouped, desc="Generating masks with predictions: "):
|
|
w, h = Image.open(image_name).size
|
|
heatmap = np.zeros((h, w), dtype=np.float32)
|
|
indices = np.meshgrid(
|
|
np.linspace(0.0, float(w), num=w, endpoint=False),
|
|
np.linspace(0.0, float(h), num=h, endpoint=False),
|
|
)
|
|
for index, row in group_df.iterrows():
|
|
x_min, y_min, x_max, y_max = (
|
|
row["x_min"],
|
|
row["y_min"],
|
|
row["x_max"],
|
|
row["y_max"],
|
|
)
|
|
gauss = gaussian(x_min, y_min, x_max, y_max)
|
|
heatmap += gauss(*indices)
|
|
heatmap /= np.max(heatmap)
|
|
heatmap_image = Image.fromarray((heatmap * 255).astype(np.uint8), mode="L")
|
|
output_filename = os.path.join(output_dir, image_name)
|
|
os.makedirs(os.path.dirname(output_filename), exist_ok=True)
|
|
heatmap_image.save(output_filename)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--dataset_dir", help="Root directory of dataset")
|
|
parser.add_argument(
|
|
"--subset",
|
|
help="Subset of dataset to process",
|
|
choices=["train", "test"],
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
label_path = os.path.join(
|
|
args.dataset_dir, args.subset + "_annotations_release.txt"
|
|
)
|
|
|
|
column_names = [
|
|
"path",
|
|
"idx",
|
|
"body_bbox_x",
|
|
"body_bbox_y",
|
|
"body_bbox_w",
|
|
"body_bbox_h",
|
|
"eye_x",
|
|
"eye_y",
|
|
"gaze_x",
|
|
"gaze_y",
|
|
"bbox_x_min",
|
|
"bbox_y_min",
|
|
"bbox_x_max",
|
|
"bbox_y_max",
|
|
]
|
|
df = pd.read_csv(
|
|
label_path,
|
|
sep=",",
|
|
names=column_names,
|
|
usecols=column_names,
|
|
index_col=False,
|
|
)
|
|
df = df.groupby("path")
|
|
|
|
model = get_model("resnet50_2020-07-20", max_size=2048, device="cuda")
|
|
model.eval()
|
|
|
|
paths = list(df.groups.keys())
|
|
csv = []
|
|
for path in tqdm.tqdm(paths, desc="Predicting head bboxes: "):
|
|
img = cv2.imread(os.path.join(args.dataset_dir, path))
|
|
|
|
annotations = model.predict_jsons(img)
|
|
|
|
for annotation in annotations:
|
|
if len(annotation["bbox"]) == 0:
|
|
continue
|
|
|
|
csv.append(
|
|
[
|
|
path,
|
|
annotation["score"],
|
|
annotation["bbox"][0],
|
|
annotation["bbox"][1],
|
|
annotation["bbox"][2],
|
|
annotation["bbox"][3],
|
|
]
|
|
)
|
|
|
|
|
|
df = pd.DataFrame(
|
|
csv, columns=["path", "score", "x_min", "y_min", "x_max", "y_max"]
|
|
)
|
|
df.to_csv(os.path.join(args.dataset_dir, f"{args.subset}_head.csv"), index=False)
|
|
|
|
plot_gen(df, args.dataset_dir)
|
|
plot_ori(label_path, args.data_dir)
|
|
|
