Support for Hunyuan Video Avatar
This commit is contained in:
170
hyvideo/data_kits/audio_dataset.py
Normal file
170
hyvideo/data_kits/audio_dataset.py
Normal file
@@ -0,0 +1,170 @@
|
||||
import os
|
||||
import cv2
|
||||
import math
|
||||
import json
|
||||
import torch
|
||||
import random
|
||||
import librosa
|
||||
import traceback
|
||||
import torchvision
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from PIL import Image
|
||||
from einops import rearrange
|
||||
from torch.utils.data import Dataset
|
||||
from decord import VideoReader, cpu
|
||||
from transformers import CLIPImageProcessor
|
||||
import torchvision.transforms as transforms
|
||||
from torchvision.transforms import ToPILImage
|
||||
|
||||
|
||||
|
||||
def get_audio_feature(feature_extractor, audio_path):
|
||||
audio_input, sampling_rate = librosa.load(audio_path, sr=16000)
|
||||
assert sampling_rate == 16000
|
||||
|
||||
audio_features = []
|
||||
window = 750*640
|
||||
for i in range(0, len(audio_input), window):
|
||||
audio_feature = feature_extractor(audio_input[i:i+window],
|
||||
sampling_rate=sampling_rate,
|
||||
return_tensors="pt",
|
||||
).input_features
|
||||
audio_features.append(audio_feature)
|
||||
|
||||
audio_features = torch.cat(audio_features, dim=-1)
|
||||
return audio_features, len(audio_input) // 640
|
||||
|
||||
|
||||
class VideoAudioTextLoaderVal(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
image_size: int,
|
||||
meta_file: str,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.meta_file = meta_file
|
||||
self.image_size = image_size
|
||||
self.text_encoder = kwargs.get("text_encoder", None) # llava_text_encoder
|
||||
self.text_encoder_2 = kwargs.get("text_encoder_2", None) # clipL_text_encoder
|
||||
self.feature_extractor = kwargs.get("feature_extractor", None)
|
||||
self.meta_files = []
|
||||
|
||||
csv_data = pd.read_csv(meta_file)
|
||||
for idx in range(len(csv_data)):
|
||||
self.meta_files.append(
|
||||
{
|
||||
"videoid": str(csv_data["videoid"][idx]),
|
||||
"image_path": str(csv_data["image"][idx]),
|
||||
"audio_path": str(csv_data["audio"][idx]),
|
||||
"prompt": str(csv_data["prompt"][idx]),
|
||||
"fps": float(csv_data["fps"][idx])
|
||||
}
|
||||
)
|
||||
|
||||
self.llava_transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize((336, 336), interpolation=transforms.InterpolationMode.BILINEAR),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.48145466, 0.4578275, 0.4082107), (0.26862954, 0.26130258, 0.27577711)),
|
||||
]
|
||||
)
|
||||
self.clip_image_processor = CLIPImageProcessor()
|
||||
|
||||
self.device = torch.device("cuda")
|
||||
self.weight_dtype = torch.float16
|
||||
|
||||
|
||||
def __len__(self):
|
||||
return len(self.meta_files)
|
||||
|
||||
@staticmethod
|
||||
def get_text_tokens(text_encoder, description, dtype_encode="video"):
|
||||
text_inputs = text_encoder.text2tokens(description, data_type=dtype_encode)
|
||||
text_ids = text_inputs["input_ids"].squeeze(0)
|
||||
text_mask = text_inputs["attention_mask"].squeeze(0)
|
||||
return text_ids, text_mask
|
||||
|
||||
def get_batch_data(self, idx):
|
||||
meta_file = self.meta_files[idx]
|
||||
videoid = meta_file["videoid"]
|
||||
image_path = meta_file["image_path"]
|
||||
audio_path = meta_file["audio_path"]
|
||||
prompt = "Authentic, Realistic, Natural, High-quality, Lens-Fixed, " + meta_file["prompt"]
|
||||
fps = meta_file["fps"]
|
||||
|
||||
img_size = self.image_size
|
||||
ref_image = Image.open(image_path).convert('RGB')
|
||||
|
||||
# Resize reference image
|
||||
w, h = ref_image.size
|
||||
scale = img_size / min(w, h)
|
||||
new_w = round(w * scale / 64) * 64
|
||||
new_h = round(h * scale / 64) * 64
|
||||
|
||||
if img_size == 704:
|
||||
img_size_long = 1216
|
||||
if new_w * new_h > img_size * img_size_long:
|
||||
import math
|
||||
scale = math.sqrt(img_size * img_size_long / w / h)
|
||||
new_w = round(w * scale / 64) * 64
|
||||
new_h = round(h * scale / 64) * 64
|
||||
|
||||
ref_image = ref_image.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
ref_image = np.array(ref_image)
|
||||
ref_image = torch.from_numpy(ref_image)
|
||||
|
||||
audio_input, audio_len = get_audio_feature(self.feature_extractor, audio_path)
|
||||
audio_prompts = audio_input[0]
|
||||
|
||||
motion_bucket_id_heads = np.array([25] * 4)
|
||||
motion_bucket_id_exps = np.array([30] * 4)
|
||||
motion_bucket_id_heads = torch.from_numpy(motion_bucket_id_heads)
|
||||
motion_bucket_id_exps = torch.from_numpy(motion_bucket_id_exps)
|
||||
fps = torch.from_numpy(np.array(fps))
|
||||
|
||||
to_pil = ToPILImage()
|
||||
pixel_value_ref = rearrange(ref_image.clone().unsqueeze(0), "b h w c -> b c h w") # (b c h w)
|
||||
|
||||
pixel_value_ref_llava = [self.llava_transform(to_pil(image)) for image in pixel_value_ref]
|
||||
pixel_value_ref_llava = torch.stack(pixel_value_ref_llava, dim=0)
|
||||
pixel_value_ref_clip = self.clip_image_processor(
|
||||
images=Image.fromarray((pixel_value_ref[0].permute(1,2,0)).data.cpu().numpy().astype(np.uint8)),
|
||||
return_tensors="pt"
|
||||
).pixel_values[0]
|
||||
pixel_value_ref_clip = pixel_value_ref_clip.unsqueeze(0)
|
||||
|
||||
# Encode text prompts
|
||||
|
||||
text_ids, text_mask = self.get_text_tokens(self.text_encoder, prompt)
|
||||
text_ids_2, text_mask_2 = self.get_text_tokens(self.text_encoder_2, prompt)
|
||||
|
||||
# Output batch
|
||||
batch = {
|
||||
"text_prompt": prompt, #
|
||||
"videoid": videoid,
|
||||
"pixel_value_ref": pixel_value_ref.to(dtype=torch.float16), # 参考图,用于vae提特征 (1, 3, h, w), 取值范围(0, 255)
|
||||
"pixel_value_ref_llava": pixel_value_ref_llava.to(dtype=torch.float16), # 参考图,用于llava提特征 (1, 3, 336, 336), 取值范围 = CLIP取值范围
|
||||
"pixel_value_ref_clip": pixel_value_ref_clip.to(dtype=torch.float16), # 参考图,用于clip_image_encoder提特征 (1, 3, 244, 244), 取值范围 = CLIP取值范围
|
||||
"audio_prompts": audio_prompts.to(dtype=torch.float16),
|
||||
"motion_bucket_id_heads": motion_bucket_id_heads.to(dtype=text_ids.dtype),
|
||||
"motion_bucket_id_exps": motion_bucket_id_exps.to(dtype=text_ids.dtype),
|
||||
"fps": fps.to(dtype=torch.float16),
|
||||
"text_ids": text_ids.clone(), # 对应llava_text_encoder
|
||||
"text_mask": text_mask.clone(), # 对应llava_text_encoder
|
||||
"text_ids_2": text_ids_2.clone(), # 对应clip_text_encoder
|
||||
"text_mask_2": text_mask_2.clone(), # 对应clip_text_encoder
|
||||
"audio_len": audio_len,
|
||||
"image_path": image_path,
|
||||
"audio_path": audio_path,
|
||||
}
|
||||
return batch
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self.get_batch_data(idx)
|
||||
|
||||
|
||||
|
||||
|
||||
72
hyvideo/data_kits/audio_preprocessor.py
Normal file
72
hyvideo/data_kits/audio_preprocessor.py
Normal file
@@ -0,0 +1,72 @@
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import json
|
||||
import time
|
||||
import decord
|
||||
import einops
|
||||
import librosa
|
||||
import torch
|
||||
import random
|
||||
import argparse
|
||||
import traceback
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from PIL import Image
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
|
||||
def get_facemask(ref_image, align_instance, area=1.25):
|
||||
# ref_image: (b f c h w)
|
||||
bsz, f, c, h, w = ref_image.shape
|
||||
images = rearrange(ref_image, "b f c h w -> (b f) h w c").data.cpu().numpy().astype(np.uint8)
|
||||
face_masks = []
|
||||
for image in images:
|
||||
image_pil = Image.fromarray(image).convert("RGB")
|
||||
_, _, bboxes_list = align_instance(np.array(image_pil)[:,:,[2,1,0]], maxface=True)
|
||||
try:
|
||||
bboxSrc = bboxes_list[0]
|
||||
except:
|
||||
bboxSrc = [0, 0, w, h]
|
||||
x1, y1, ww, hh = bboxSrc
|
||||
x2, y2 = x1 + ww, y1 + hh
|
||||
ww, hh = (x2-x1) * area, (y2-y1) * area
|
||||
center = [(x2+x1)//2, (y2+y1)//2]
|
||||
x1 = max(center[0] - ww//2, 0)
|
||||
y1 = max(center[1] - hh//2, 0)
|
||||
x2 = min(center[0] + ww//2, w)
|
||||
y2 = min(center[1] + hh//2, h)
|
||||
|
||||
face_mask = np.zeros_like(np.array(image_pil))
|
||||
face_mask[int(y1):int(y2), int(x1):int(x2)] = 1.0
|
||||
face_masks.append(torch.from_numpy(face_mask[...,:1]))
|
||||
face_masks = torch.stack(face_masks, dim=0) # (b*f, h, w, c)
|
||||
face_masks = rearrange(face_masks, "(b f) h w c -> b c f h w", b=bsz, f=f)
|
||||
face_masks = face_masks.to(device=ref_image.device, dtype=ref_image.dtype)
|
||||
return face_masks
|
||||
|
||||
|
||||
def encode_audio(wav2vec, audio_feats, fps, num_frames=129):
|
||||
if fps == 25:
|
||||
start_ts = [0]
|
||||
step_ts = [1]
|
||||
elif fps == 12.5:
|
||||
start_ts = [0]
|
||||
step_ts = [2]
|
||||
num_frames = min(num_frames, 400)
|
||||
audio_feats = wav2vec.encoder(audio_feats.unsqueeze(0)[:, :, :3000], output_hidden_states=True).hidden_states
|
||||
audio_feats = torch.stack(audio_feats, dim=2)
|
||||
audio_feats = torch.cat([torch.zeros_like(audio_feats[:,:4]), audio_feats], 1)
|
||||
|
||||
audio_prompts = []
|
||||
for bb in range(1):
|
||||
audio_feats_list = []
|
||||
for f in range(num_frames):
|
||||
cur_t = (start_ts[bb] + f * step_ts[bb]) * 2
|
||||
audio_clip = audio_feats[bb:bb+1, cur_t: cur_t+10]
|
||||
audio_feats_list.append(audio_clip)
|
||||
audio_feats_list = torch.stack(audio_feats_list, 1)
|
||||
audio_prompts.append(audio_feats_list)
|
||||
audio_prompts = torch.cat(audio_prompts)
|
||||
return audio_prompts
|
||||
41
hyvideo/data_kits/data_tools.py
Normal file
41
hyvideo/data_kits/data_tools.py
Normal file
@@ -0,0 +1,41 @@
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
import imageio
|
||||
import torchvision
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8, quality=8):
|
||||
videos = rearrange(videos, "b c t h w -> t b c h w")
|
||||
outputs = []
|
||||
for x in videos:
|
||||
x = torchvision.utils.make_grid(x, nrow=n_rows)
|
||||
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
|
||||
if rescale:
|
||||
x = (x + 1.0) / 2.0 # -1,1 -> 0,1
|
||||
x = torch.clamp(x,0,1)
|
||||
x = (x * 255).numpy().astype(np.uint8)
|
||||
outputs.append(x)
|
||||
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
imageio.mimsave(path, outputs, fps=fps, quality=quality)
|
||||
|
||||
def pad_image(crop_img, size, color=(255, 255, 255), resize_ratio=1):
|
||||
crop_h, crop_w = crop_img.shape[:2]
|
||||
target_w, target_h = size
|
||||
scale_h, scale_w = target_h / crop_h, target_w / crop_w
|
||||
if scale_w > scale_h:
|
||||
resize_h = int(target_h*resize_ratio)
|
||||
resize_w = int(crop_w / crop_h * resize_h)
|
||||
else:
|
||||
resize_w = int(target_w*resize_ratio)
|
||||
resize_h = int(crop_h / crop_w * resize_w)
|
||||
crop_img = cv2.resize(crop_img, (resize_w, resize_h))
|
||||
pad_left = (target_w - resize_w) // 2
|
||||
pad_top = (target_h - resize_h) // 2
|
||||
pad_right = target_w - resize_w - pad_left
|
||||
pad_bottom = target_h - resize_h - pad_top
|
||||
crop_img = cv2.copyMakeBorder(crop_img, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=color)
|
||||
return crop_img
|
||||
1
hyvideo/data_kits/face_align/__init__.py
Normal file
1
hyvideo/data_kits/face_align/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
from .align import AlignImage
|
||||
34
hyvideo/data_kits/face_align/align.py
Normal file
34
hyvideo/data_kits/face_align/align.py
Normal file
@@ -0,0 +1,34 @@
|
||||
import os
|
||||
import sys
|
||||
import torch
|
||||
from .detface import DetFace
|
||||
|
||||
class AlignImage(object):
|
||||
def __init__(self, device='cuda', det_path=''):
|
||||
self.facedet = DetFace(pt_path=det_path, confThreshold=0.5, nmsThreshold=0.45, device=device)
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(self, im, maxface=False):
|
||||
bboxes, kpss, scores = self.facedet.detect(im)
|
||||
face_num = bboxes.shape[0]
|
||||
|
||||
five_pts_list = []
|
||||
scores_list = []
|
||||
bboxes_list = []
|
||||
for i in range(face_num):
|
||||
five_pts_list.append(kpss[i].reshape(5,2))
|
||||
scores_list.append(scores[i])
|
||||
bboxes_list.append(bboxes[i])
|
||||
|
||||
if maxface and face_num>1:
|
||||
max_idx = 0
|
||||
max_area = (bboxes[0, 2])*(bboxes[0, 3])
|
||||
for i in range(1, face_num):
|
||||
area = (bboxes[i,2])*(bboxes[i,3])
|
||||
if area>max_area:
|
||||
max_idx = i
|
||||
five_pts_list = [five_pts_list[max_idx]]
|
||||
scores_list = [scores_list[max_idx]]
|
||||
bboxes_list = [bboxes_list[max_idx]]
|
||||
|
||||
return five_pts_list, scores_list, bboxes_list
|
||||
283
hyvideo/data_kits/face_align/detface.py
Normal file
283
hyvideo/data_kits/face_align/detface.py
Normal file
@@ -0,0 +1,283 @@
|
||||
# -*- coding: UTF-8 -*-
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision
|
||||
|
||||
|
||||
def xyxy2xywh(x):
|
||||
# Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
|
||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||
y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center
|
||||
y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center
|
||||
y[:, 2] = x[:, 2] - x[:, 0] # width
|
||||
y[:, 3] = x[:, 3] - x[:, 1] # height
|
||||
return y
|
||||
|
||||
|
||||
def xywh2xyxy(x):
|
||||
# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
|
||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||
y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
|
||||
y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y
|
||||
y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x
|
||||
y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y
|
||||
return y
|
||||
|
||||
|
||||
def box_iou(box1, box2):
|
||||
# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
|
||||
"""
|
||||
Return intersection-over-union (Jaccard index) of boxes.
|
||||
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
|
||||
Arguments:
|
||||
box1 (Tensor[N, 4])
|
||||
box2 (Tensor[M, 4])
|
||||
Returns:
|
||||
iou (Tensor[N, M]): the NxM matrix containing the pairwise
|
||||
IoU values for every element in boxes1 and boxes2
|
||||
"""
|
||||
|
||||
def box_area(box):
|
||||
# box = 4xn
|
||||
return (box[2] - box[0]) * (box[3] - box[1])
|
||||
|
||||
area1 = box_area(box1.T)
|
||||
area2 = box_area(box2.T)
|
||||
|
||||
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
|
||||
inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) -
|
||||
torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
|
||||
# iou = inter / (area1 + area2 - inter)
|
||||
return inter / (area1[:, None] + area2 - inter)
|
||||
|
||||
|
||||
def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
|
||||
# Rescale coords (xyxy) from img1_shape to img0_shape
|
||||
if ratio_pad is None: # calculate from img0_shape
|
||||
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
|
||||
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
|
||||
else:
|
||||
gain = ratio_pad[0][0]
|
||||
pad = ratio_pad[1]
|
||||
|
||||
coords[:, [0, 2]] -= pad[0] # x padding
|
||||
coords[:, [1, 3]] -= pad[1] # y padding
|
||||
coords[:, :4] /= gain
|
||||
clip_coords(coords, img0_shape)
|
||||
return coords
|
||||
|
||||
|
||||
def clip_coords(boxes, img_shape):
|
||||
# Clip bounding xyxy bounding boxes to image shape (height, width)
|
||||
boxes[:, 0].clamp_(0, img_shape[1]) # x1
|
||||
boxes[:, 1].clamp_(0, img_shape[0]) # y1
|
||||
boxes[:, 2].clamp_(0, img_shape[1]) # x2
|
||||
boxes[:, 3].clamp_(0, img_shape[0]) # y2
|
||||
|
||||
|
||||
def scale_coords_landmarks(img1_shape, coords, img0_shape, ratio_pad=None):
|
||||
# Rescale coords (xyxy) from img1_shape to img0_shape
|
||||
if ratio_pad is None: # calculate from img0_shape
|
||||
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
|
||||
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
|
||||
else:
|
||||
gain = ratio_pad[0][0]
|
||||
pad = ratio_pad[1]
|
||||
|
||||
coords[:, [0, 2, 4, 6, 8]] -= pad[0] # x padding
|
||||
coords[:, [1, 3, 5, 7, 9]] -= pad[1] # y padding
|
||||
coords[:, :10] /= gain
|
||||
#clip_coords(coords, img0_shape)
|
||||
coords[:, 0].clamp_(0, img0_shape[1]) # x1
|
||||
coords[:, 1].clamp_(0, img0_shape[0]) # y1
|
||||
coords[:, 2].clamp_(0, img0_shape[1]) # x2
|
||||
coords[:, 3].clamp_(0, img0_shape[0]) # y2
|
||||
coords[:, 4].clamp_(0, img0_shape[1]) # x3
|
||||
coords[:, 5].clamp_(0, img0_shape[0]) # y3
|
||||
coords[:, 6].clamp_(0, img0_shape[1]) # x4
|
||||
coords[:, 7].clamp_(0, img0_shape[0]) # y4
|
||||
coords[:, 8].clamp_(0, img0_shape[1]) # x5
|
||||
coords[:, 9].clamp_(0, img0_shape[0]) # y5
|
||||
return coords
|
||||
|
||||
|
||||
def show_results(img, xywh, conf, landmarks, class_num):
|
||||
h,w,c = img.shape
|
||||
tl = 1 or round(0.002 * (h + w) / 2) + 1 # line/font thickness
|
||||
x1 = int(xywh[0] * w - 0.5 * xywh[2] * w)
|
||||
y1 = int(xywh[1] * h - 0.5 * xywh[3] * h)
|
||||
x2 = int(xywh[0] * w + 0.5 * xywh[2] * w)
|
||||
y2 = int(xywh[1] * h + 0.5 * xywh[3] * h)
|
||||
cv2.rectangle(img, (x1,y1), (x2, y2), (0,255,0), thickness=tl, lineType=cv2.LINE_AA)
|
||||
|
||||
clors = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(0,255,255)]
|
||||
|
||||
for i in range(5):
|
||||
point_x = int(landmarks[2 * i] * w)
|
||||
point_y = int(landmarks[2 * i + 1] * h)
|
||||
cv2.circle(img, (point_x, point_y), tl+1, clors[i], -1)
|
||||
|
||||
tf = max(tl - 1, 1) # font thickness
|
||||
label = str(conf)[:5]
|
||||
cv2.putText(img, label, (x1, y1 - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
|
||||
return img
|
||||
|
||||
|
||||
def make_divisible(x, divisor):
|
||||
# Returns x evenly divisible by divisor
|
||||
return (x // divisor) * divisor
|
||||
|
||||
|
||||
def non_max_suppression_face(prediction, conf_thres=0.5, iou_thres=0.45, classes=None, agnostic=False, labels=()):
|
||||
"""Performs Non-Maximum Suppression (NMS) on inference results
|
||||
Returns:
|
||||
detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
|
||||
"""
|
||||
|
||||
nc = prediction.shape[2] - 15 # number of classes
|
||||
xc = prediction[..., 4] > conf_thres # candidates
|
||||
|
||||
# Settings
|
||||
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
|
||||
# time_limit = 10.0 # seconds to quit after
|
||||
redundant = True # require redundant detections
|
||||
multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img)
|
||||
merge = False # use merge-NMS
|
||||
|
||||
# t = time.time()
|
||||
output = [torch.zeros((0, 16), device=prediction.device)] * prediction.shape[0]
|
||||
for xi, x in enumerate(prediction): # image index, image inference
|
||||
# Apply constraints
|
||||
# x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
|
||||
x = x[xc[xi]] # confidence
|
||||
|
||||
# Cat apriori labels if autolabelling
|
||||
if labels and len(labels[xi]):
|
||||
l = labels[xi]
|
||||
v = torch.zeros((len(l), nc + 15), device=x.device)
|
||||
v[:, :4] = l[:, 1:5] # box
|
||||
v[:, 4] = 1.0 # conf
|
||||
v[range(len(l)), l[:, 0].long() + 15] = 1.0 # cls
|
||||
x = torch.cat((x, v), 0)
|
||||
|
||||
# If none remain process next image
|
||||
if not x.shape[0]:
|
||||
continue
|
||||
|
||||
# Compute conf
|
||||
x[:, 15:] *= x[:, 4:5] # conf = obj_conf * cls_conf
|
||||
|
||||
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
|
||||
box = xywh2xyxy(x[:, :4])
|
||||
|
||||
# Detections matrix nx6 (xyxy, conf, landmarks, cls)
|
||||
if multi_label:
|
||||
i, j = (x[:, 15:] > conf_thres).nonzero(as_tuple=False).T
|
||||
x = torch.cat((box[i], x[i, j + 15, None], x[i, 5:15] ,j[:, None].float()), 1)
|
||||
else: # best class only
|
||||
conf, j = x[:, 15:].max(1, keepdim=True)
|
||||
x = torch.cat((box, conf, x[:, 5:15], j.float()), 1)[conf.view(-1) > conf_thres]
|
||||
|
||||
# Filter by class
|
||||
if classes is not None:
|
||||
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
|
||||
|
||||
# If none remain process next image
|
||||
n = x.shape[0] # number of boxes
|
||||
if not n:
|
||||
continue
|
||||
|
||||
# Batched NMS
|
||||
c = x[:, 15:16] * (0 if agnostic else max_wh) # classes
|
||||
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
|
||||
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
|
||||
#if i.shape[0] > max_det: # limit detections
|
||||
# i = i[:max_det]
|
||||
if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
|
||||
# update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
|
||||
iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
|
||||
weights = iou * scores[None] # box weights
|
||||
x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
|
||||
if redundant:
|
||||
i = i[iou.sum(1) > 1] # require redundancy
|
||||
|
||||
output[xi] = x[i]
|
||||
# if (time.time() - t) > time_limit:
|
||||
# break # time limit exceeded
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class DetFace():
|
||||
def __init__(self, pt_path, confThreshold=0.5, nmsThreshold=0.45, device='cuda'):
|
||||
assert os.path.exists(pt_path)
|
||||
|
||||
self.inpSize = 416
|
||||
self.conf_thres = confThreshold
|
||||
self.iou_thres = nmsThreshold
|
||||
self.test_device = torch.device(device if torch.cuda.is_available() else "cpu")
|
||||
self.model = torch.jit.load(pt_path).to(self.test_device)
|
||||
self.last_w = 416
|
||||
self.last_h = 416
|
||||
self.grids = None
|
||||
|
||||
@torch.no_grad()
|
||||
def detect(self, srcimg):
|
||||
# t0=time.time()
|
||||
|
||||
h0, w0 = srcimg.shape[:2] # orig hw
|
||||
r = self.inpSize / min(h0, w0) # resize image to img_size
|
||||
h1 = int(h0*r+31)//32*32
|
||||
w1 = int(w0*r+31)//32*32
|
||||
|
||||
img = cv2.resize(srcimg, (w1,h1), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# Convert
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # BGR to RGB
|
||||
|
||||
# Run inference
|
||||
img = torch.from_numpy(img).to(self.test_device).permute(2,0,1)
|
||||
img = img.float()/255 # uint8 to fp16/32 0-1
|
||||
if img.ndimension() == 3:
|
||||
img = img.unsqueeze(0)
|
||||
|
||||
# Inference
|
||||
if h1 != self.last_h or w1 != self.last_w or self.grids is None:
|
||||
grids = []
|
||||
for scale in [8,16,32]:
|
||||
ny = h1//scale
|
||||
nx = w1//scale
|
||||
yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
|
||||
grid = torch.stack((xv, yv), 2).view((1,1,ny, nx, 2)).float()
|
||||
grids.append(grid.to(self.test_device))
|
||||
self.grids = grids
|
||||
self.last_w = w1
|
||||
self.last_h = h1
|
||||
|
||||
pred = self.model(img, self.grids).cpu()
|
||||
|
||||
# Apply NMS
|
||||
det = non_max_suppression_face(pred, self.conf_thres, self.iou_thres)[0]
|
||||
# Process detections
|
||||
# det = pred[0]
|
||||
bboxes = np.zeros((det.shape[0], 4))
|
||||
kpss = np.zeros((det.shape[0], 5, 2))
|
||||
scores = np.zeros((det.shape[0]))
|
||||
# gn = torch.tensor([w0, h0, w0, h0]).to(pred) # normalization gain whwh
|
||||
# gn_lks = torch.tensor([w0, h0, w0, h0, w0, h0, w0, h0, w0, h0]).to(pred) # normalization gain landmarks
|
||||
det = det.cpu().numpy()
|
||||
|
||||
for j in range(det.shape[0]):
|
||||
# xywh = (xyxy2xywh(det[j, :4].view(1, 4)) / gn).view(4).cpu().numpy()
|
||||
bboxes[j, 0] = det[j, 0] * w0/w1
|
||||
bboxes[j, 1] = det[j, 1] * h0/h1
|
||||
bboxes[j, 2] = det[j, 2] * w0/w1 - bboxes[j, 0]
|
||||
bboxes[j, 3] = det[j, 3] * h0/h1 - bboxes[j, 1]
|
||||
scores[j] = det[j, 4]
|
||||
# landmarks = (det[j, 5:15].view(1, 10) / gn_lks).view(5,2).cpu().numpy()
|
||||
kpss[j, :, :] = det[j, 5:15].reshape(5, 2) * np.array([[w0/w1,h0/h1]])
|
||||
# class_num = det[j, 15].cpu().numpy()
|
||||
# orgimg = show_results(orgimg, xywh, conf, landmarks, class_num)
|
||||
return bboxes, kpss, scores
|
||||
Reference in New Issue
Block a user