No result

I convert pynb to python but I got no output:

from trt_pose.parse_objects import ParseObjects
from trt_pose.draw_objects import DrawObjects
import PIL.Image
import torchvision.transforms as transforms
import cv2
import time
from torch2trt import TRTModule
import torch2trt
import torch
import trt_pose.models
import json
import trt_pose.coco
import os
import sys


with open('human_pose.json', 'r') as f:
    human_pose = json.load(f)

topology = trt_pose.coco.coco_category_to_topology(human_pose)


num_parts = len(human_pose['keypoints'])
num_links = len(human_pose['skeleton'])

model = trt_pose.models.resnet18_baseline_att(
    num_parts, 2 * num_links).cuda().eval()


MODEL_WEIGHTS = 'resnet18_baseline_att_224x224_A_epoch_249.pth'

model.load_state_dict(torch.load(MODEL_WEIGHTS))


WIDTH = 224
HEIGHT = 224

data = torch.zeros((1, 3, HEIGHT, WIDTH)).cuda()


model_trt = torch2trt.torch2trt(
    model, [data], fp16_mode=True, max_workspace_size=1 << 25)


OPTIMIZED_MODEL = 'resnet18_baseline_att_224x224_A_epoch_249_trt.pth'

torch.save(model_trt.state_dict(), OPTIMIZED_MODEL)


model_trt = TRTModule()
model_trt.load_state_dict(torch.load(OPTIMIZED_MODEL))


t0 = time.time()
torch.cuda.current_stream().synchronize()
for i in range(50):
    y = model_trt(data)
torch.cuda.current_stream().synchronize()
t1 = time.time()

print(50.0 / (t1 - t0))


mean = torch.Tensor([0.485, 0.456, 0.406]).cuda()
std = torch.Tensor([0.229, 0.224, 0.225]).cuda()
device = torch.device('cuda')


def preprocess(image):
    global device
    device = torch.device('cuda')
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = PIL.Image.fromarray(image)
    image = transforms.functional.to_tensor(image).to(device)
    image.sub_(mean[:, None, None]).div_(std[:, None, None])
    return image[None, ...]


parse_objects = ParseObjects(topology)
draw_objects = DrawObjects(topology)


if __name__ == '__main__':
    v_f = sys.argv[1]
    cap = cv2.VideoCapture(v_f)

    while(True):
        ret, image = cap.read()
        if ret:
            data = preprocess(image)
            cmap, paf = model_trt(data)
            cmap, paf = cmap.detach().cpu(), paf.detach().cpu()
            # , cmap_threshold=0.15, link_threshold=0.15)
            counts, objects, peaks = parse_objects(cmap, paf)
            print(counts)
            print(objects)
            draw_objects(image, counts, objects, peaks)
            cv2.imshow('res', image)
            cv2.waitKey(1)

[gongysh2004]

haha, the secret is to resize the frame into 224,224 before calling model_trt

[tucachmo2202]

@lweingart ` import os import time import json import cv2 import torchvision.transforms as transforms import PIL.Image from trt_pose.draw_objects import DrawObjects from trt_pose.parse_objects import ParseObjects import trt_pose.coco import trt_pose.models import torch from torch2trt import TRTModule from PIL import Image

WIDTH = 256 HEIGHT = 256 mean = torch.Tensor([0.485, 0.456, 0.406]).cuda() std = torch.Tensor([0.229, 0.224, 0.225]).cuda() device = torch.device(‘cuda’)

with open(‘human_pose.json’, ‘r’) as f: human_pose = json.load(f)

topology = trt_pose.coco.coco_category_to_topology(human_pose) num_parts = len(human_pose[‘keypoints’]) num_links = len(human_pose[‘skeleton’])

OPTIMIZED_MODEL = ‘resnet50_baseline_att_256x256_A_epoch_249_trt.pth’ print(OPTIMIZED_MODEL)

model_trt = TRTModule() model_trt.load_state_dict(torch.load(OPTIMIZED_MODEL))

parse_objects = ParseObjects(topology) draw_objects = DrawObjects(topology)

def preprocess(image): global device device = torch.device(‘cuda’)

image = PIL.Image.fromarray(image)
image = transforms.functional.to_tensor(image).to(device)
image.sub_(mean[:, None, None]).div_(std[:, None, None])
return image[None, ...]

def process_img(image): ori_image = image.copy() image = cv2.resize(image, (WIDTH, HEIGHT)) data = preprocess(image) start = time.time() start_model = time.time() cmap, paf = model_trt(data) print("FPS model: ", 1.0/(time.time() - start_model)) cmap, paf = cmap.detach().cpu(), paf.detach().cpu() counts, objects, peaks = parse_objects(cmap, paf)#, cmap_threshold=0.15, link_threshold=0.15) print("FPS: ", 1.0/(time.time() - start)) draw_objects(ori_image, counts, objects, peaks) return ori_image

def predict_image(path = ‘1.jpg’): image = cv2.imread(path) img = process_img(image) cv2.imshow(“as”, img) cv2.waitKey() cv2.destroyAllWindows()

def predict_video(path_video): print(path_video) if os.path.exists(path_video): print(“exist path video”) vid = cv2.VideoCapture(path_video)

while(True):
    ret, frame = vid.read() 
    if not ret:
        # print("no frame")
        break
    
    frame = process_img(frame)
    frame = cv2.resize(frame, (frame.shape[1]//2, frame.shape[0]//2))
    cv2.imshow("as", frame)
    if cv2.waitKey(25) & 0xFF == ord('q'):
        break
    
vid.release()
cv2.destroyAllWindows() 

predict_video(‘Hog_2.mp4’) ` Hope it help you!