Spaces:
Running
Running
File size: 2,243 Bytes
d015578 |
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 |
# SPIGA library
import spiga.inference.config as model_cfg
from spiga.inference.framework import SPIGAFramework
# Demo modules
import spiga.demo.analyze.extract.processor as pr
class SPIGAProcessor(pr.Processor):
def __init__(self,
dataset='wflw',
features=('lnd', 'pose'),
gpus=[0]):
super().__init__()
# Configure and load processor
self.processor_cfg = model_cfg.ModelConfig(dataset)
self.processor = SPIGAFramework(self.processor_cfg, gpus=gpus)
# Define attributes
if 'lnd' in features:
self.attributes.append('landmarks')
self.attributes.append('landmarks_ids')
if 'pose' in features:
self.attributes.append('headpose')
def process_frame(self, frame, tracked_obj):
bboxes = []
for obj in tracked_obj:
x1, y1, x2, y2 = obj.bbox[:4]
bbox_wh = [x1, y1, x2-x1, y2-y1]
bboxes.append(bbox_wh)
features = self.processor.inference(frame, bboxes)
for obj_idx in range(len(tracked_obj)):
# Landmarks output
if 'landmarks' in self.attributes:
tracked_obj[obj_idx].landmarks = features['landmarks'][obj_idx]
tracked_obj[obj_idx].landmarks_ids = self.processor_cfg.dataset.ldm_ids
# Headpose output
if 'headpose' in self.attributes:
tracked_obj[obj_idx].headpose = features['headpose'][obj_idx]
return tracked_obj
def plot_features(self, image, features, plotter, show_attributes):
if 'landmarks' in self.attributes and 'landmarks' in show_attributes:
x1, y1, x2, y2 = features.bbox[:4]
thick = int(plotter.landmarks.thickness['lnd'] * (x2-x1)/200 + 0.5)
if thick == 0:
thick = 1
image = plotter.landmarks.draw_landmarks(image, features.landmarks, thick=thick)
if 'headpose' in self.attributes and 'headpose' in show_attributes:
image = plotter.hpose.draw_headpose(image, features.bbox[:5],
features.headpose[:3], features.headpose[3:], euler=True)
return image
|