Spaces:
Build error
Build error
File size: 3,808 Bytes
0c2c19f bd1ea7f 0c2c19f bd1ea7f 0c2c19f bd1ea7f b2de2a8 bd1ea7f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
import os
import sys
from pathlib import Path
import cv2
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
import torch
from yolov5.utils.torch_utils import select_device, time_sync
from yolov5.utils.plots import Annotator, colors, save_one_box
from yolov5.utils.general import (check_img_size,
increment_path, non_max_suppression, scale_coords, xyxy2xywh)
from yolov5.utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, pil_to_cv
from yolov5.models.common import DetectMultiBackend
import torchvision
test_transforms = torchvision.transforms.Compose([
torchvision.transforms.ToPILImage(),
torchvision.transforms.transforms.ToTensor(),
torchvision.transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
torchvision.transforms.Resize((224, 224)),
])
def load_yolo_model(weights, device="cpu", imgsz=[1280, 1280]):
# Load model
device = select_device(device)
model = DetectMultiBackend(weights, device=device, dnn=False, data=ROOT / 'data/coco128.yaml')
stride, names, pt, jit, onnx, engine = model.stride, model.names, model.pt, model.jit, model.onnx, model.engine
imgsz = check_img_size(imgsz, s=stride) # check image size
half = False
# Half
half &= (pt or jit or onnx or engine) and device.type != 'cpu' # FP16 supported on limited backends with CUDA
if pt or jit:
model.model.half() if half else model.model.float()
model.warmup(imgsz=(1, 3, *imgsz), half=half)
return model, stride, names, pt, jit, onnx, engine
def predict(
age_model,
model, # model.pt path(s)
stride,
source=None, # PIL Image
imgsz=(640, 640), # inference size (height, width)
conf_thres=0.5, # confidence threshold
iou_thres=0.45, # NMS IOU threshold
max_det=1000, # maximum detections per image
device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu
classes=None, # filter by class: --class 0, or --class 0 2 3
agnostic_nms=False, # class-agnostic NMS
augment=False, # augmented inference
visualize=False, # visualize features
half=False, # use FP16 half-precision inference
):
im, im0 = pil_to_cv(source, img_size=imgsz[0], stride=stride)
im = torch.from_numpy(im).to(device)
im = im.half() if half else im.float() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
if len(im.shape) == 3:
im = im[None] # expand for batch dim
# Inference
visualize = False
pred = model(im, augment=augment, visualize=visualize)
# NMS
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
# Process predictions
preds = []
for i, det in enumerate(pred): # per image
# im0 = im0.copy()
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
for *xyxy, conf, cls in reversed(det):
face = im0[int(xyxy[1]):int(xyxy[3]),int(xyxy[0]):int(xyxy[2])]
face_img = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
im = test_transforms(face_img).unsqueeze_(0)
with torch.no_grad():
y = age_model(im)
age = y[0]
preds.append({"class": str(int(age)), "xmin": int(xyxy[0]), "ymin": int(xyxy[1]), "xmax": int(xyxy[2]),"ymax": int(xyxy[3]), "conf": float(conf)})
return preds |