테스트 용 예시 코드
from pathlib import Path
import cv2
import torch
from models.experimental import attempt_load
from utils.datasets import LoadImages
from utils.general import non_max_suppression, scale_coords
device = 'cpu'
kpt_label = 5
def detect():
# Load model
model = attempt_load('yolov7-w6-face.pt', map_location='cpu')
stride = int(model.stride.max())
names = model.module.names if hasattr(model, 'module') else model.names # get class names
dataset = LoadImages('./crowd.jpg', img_size=640, stride=stride)
# Run inference
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3: img = img.unsqueeze(0)
# Inference
pred = model(img, augment=False)[0]
# Apply NMS
pred = non_max_suppression(pred, 0.25, 0.45, classes=0, agnostic=False, kpt_label=kpt_label)
# Process detections
for i, det in enumerate(pred): # detections per image
p, s, im0, frame = path, '', im0s.copy(), getattr(dataset, 'frame', 0)
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
if len(det):
# Rescale boxes from img_size to im0 size
scale_coords(img.shape[2:], det[:, :4], im0.shape, kpt_label=False)
scale_coords(img.shape[2:], det[:, 6:], im0.shape, kpt_label=kpt_label, step=3)
# Write results
for det_index, (*xyxy, conf, cls) in enumerate(reversed(det[:,:6])):
c = int(cls) # integer class
label = f'{names[c]} {conf:.2f}'
kpts = det[det_index, 6:]
# replace face region with blurred one
left_pos, top_pos, right_pos, bottom_pos = int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])
face_image = im0[top_pos:bottom_pos, left_pos:right_pos]
face_image = cv2.GaussianBlur(face_image, (99,99), 30)
im0[top_pos:bottom_pos, left_pos:right_pos] = face_image
cv2.imshow(str(p), im0)
if cv2.waitKey(10000) == ord('q'): break
if __name__ == '__main__':
with torch.no_grad():
detect()
