Spaces:
Running
Running
import os | |
import sys | |
import time | |
from multiprocessing import Queue as pQueue | |
from threading import Thread | |
import cv2 | |
import numpy as np | |
import torch | |
import torch.multiprocessing as mp | |
import torch.utils.data as data | |
import torchvision.transforms as transforms | |
from PIL import Image | |
from torch.autograd import Variable | |
from SPPE.src.utils.eval import getPrediction, getMultiPeakPrediction | |
from SPPE.src.utils.img import load_image, cropBox, im_to_torch | |
from matching import candidate_reselect as matching | |
from opt import opt | |
from pPose_nms import pose_nms | |
from yolo.darknet import Darknet | |
from yolo.preprocess import prep_image, prep_frame | |
from yolo.util import dynamic_write_results | |
# import the Queue class from Python 3 | |
if sys.version_info >= (3, 0): | |
from queue import Queue, LifoQueue | |
# otherwise, import the Queue class for Python 2.7 | |
else: | |
from Queue import Queue, LifoQueue | |
if opt.vis_fast: | |
from fn import vis_frame_fast as vis_frame | |
else: | |
from fn import vis_frame | |
class Image_loader(data.Dataset): | |
def __init__(self, im_names, format='yolo'): | |
super(Image_loader, self).__init__() | |
self.img_dir = opt.inputpath | |
self.imglist = im_names | |
self.transform = transforms.Compose([ | |
transforms.ToTensor(), | |
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) | |
]) | |
self.format = format | |
def getitem_ssd(self, index): | |
im_name = self.imglist[index].rstrip('\n').rstrip('\r') | |
im_name = os.path.join(self.img_dir, im_name) | |
im = Image.open(im_name) | |
inp = load_image(im_name) | |
if im.mode == 'L': | |
im = im.convert('RGB') | |
ow = oh = 512 | |
im = im.resize((ow, oh)) | |
im = self.transform(im) | |
return im, inp, im_name | |
def getitem_yolo(self, index): | |
inp_dim = int(opt.inp_dim) | |
im_name = self.imglist[index].rstrip('\n').rstrip('\r') | |
im_name = os.path.join(self.img_dir, im_name) | |
im, orig_img, im_dim = prep_image(im_name, inp_dim) | |
# im_dim = torch.FloatTensor([im_dim]).repeat(1, 2) | |
inp = load_image(im_name) | |
return im, inp, orig_img, im_name, im_dim | |
def __getitem__(self, index): | |
if self.format == 'ssd': | |
return self.getitem_ssd(index) | |
elif self.format == 'yolo': | |
return self.getitem_yolo(index) | |
else: | |
raise NotImplementedError | |
def __len__(self): | |
return len(self.imglist) | |
class ImageLoader: | |
def __init__(self, im_names, batchSize=1, format='yolo', queueSize=50): | |
self.img_dir = opt.inputpath | |
self.imglist = im_names | |
self.transform = transforms.Compose([ | |
transforms.ToTensor(), | |
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) | |
]) | |
self.format = format | |
self.batchSize = batchSize | |
self.datalen = len(self.imglist) | |
leftover = 0 | |
if (self.datalen) % batchSize: | |
leftover = 1 | |
self.num_batches = self.datalen // batchSize + leftover | |
# initialize the queue used to store data | |
if opt.sp: | |
self.Q = Queue(maxsize=queueSize) | |
else: | |
self.Q = mp.Queue(maxsize=queueSize) | |
def start(self): | |
# start a thread to read frames from the file video stream | |
if self.format == 'ssd': | |
if opt.sp: | |
p = Thread(target=self.getitem_ssd, args=()) | |
else: | |
p = mp.Process(target=self.getitem_ssd, args=()) | |
elif self.format == 'yolo': | |
if opt.sp: | |
p = Thread(target=self.getitem_yolo, args=()) | |
else: | |
p = mp.Process(target=self.getitem_yolo, args=()) | |
else: | |
raise NotImplementedError | |
p.daemon = True | |
p.start() | |
return self | |
def getitem_ssd(self): | |
length = len(self.imglist) | |
for index in range(length): | |
im_name = self.imglist[index].rstrip('\n').rstrip('\r') | |
im_name = os.path.join(self.img_dir, im_name) | |
im = Image.open(im_name) | |
inp = load_image(im_name) | |
if im.mode == 'L': | |
im = im.convert('RGB') | |
ow = oh = 512 | |
im = im.resize((ow, oh)) | |
im = self.transform(im) | |
while self.Q.full(): | |
time.sleep(2) | |
self.Q.put((im, inp, im_name)) | |
def getitem_yolo(self): | |
for i in range(self.num_batches): | |
img = [] | |
orig_img = [] | |
im_name = [] | |
im_dim_list = [] | |
for k in range(i * self.batchSize, min((i + 1) * self.batchSize, self.datalen)): | |
inp_dim = int(opt.inp_dim) | |
im_name_k = self.imglist[k].rstrip('\n').rstrip('\r') | |
im_name_k = os.path.join(self.img_dir, im_name_k) | |
img_k, orig_img_k, im_dim_list_k = prep_image(im_name_k, inp_dim) | |
img.append(img_k) | |
orig_img.append(orig_img_k) | |
im_name.append(im_name_k) | |
im_dim_list.append(im_dim_list_k) | |
with torch.no_grad(): | |
# Human Detection | |
img = torch.cat(img) | |
im_dim_list = torch.FloatTensor(im_dim_list).repeat(1, 2) | |
im_dim_list_ = im_dim_list | |
while self.Q.full(): | |
time.sleep(2) | |
self.Q.put((img, orig_img, im_name, im_dim_list)) | |
def getitem(self): | |
return self.Q.get() | |
def length(self): | |
return len(self.imglist) | |
def len(self): | |
return self.Q.qsize() | |
class VideoLoader: | |
def __init__(self, path, batchSize=1, queueSize=50): | |
# initialize the file video stream along with the boolean | |
# used to indicate if the thread should be stopped or not | |
self.path = path | |
self.stream = cv2.VideoCapture(path) | |
assert self.stream.isOpened(), 'Cannot capture source' | |
self.stopped = False | |
self.batchSize = batchSize | |
self.datalen = int(self.stream.get(cv2.CAP_PROP_FRAME_COUNT)) | |
leftover = 0 | |
if (self.datalen) % batchSize: | |
leftover = 1 | |
self.num_batches = self.datalen // batchSize + leftover | |
# initialize the queue used to store frames read from | |
# the video file | |
if opt.sp: | |
self.Q = Queue(maxsize=queueSize) | |
else: | |
self.Q = mp.Queue(maxsize=queueSize) | |
def length(self): | |
return self.datalen | |
def start(self): | |
# start a thread to read frames from the file video stream | |
if opt.sp: | |
t = Thread(target=self.update, args=()) | |
t.daemon = True | |
t.start() | |
else: | |
p = mp.Process(target=self.update, args=()) | |
p.daemon = True | |
p.start() | |
return self | |
def update(self): | |
stream = cv2.VideoCapture(self.path) | |
assert stream.isOpened(), 'Cannot capture source' | |
for i in range(self.num_batches): | |
img = [] | |
orig_img = [] | |
im_name = [] | |
im_dim_list = [] | |
for k in range(i * self.batchSize, min((i + 1) * self.batchSize, self.datalen)): | |
inp_dim = int(opt.inp_dim) | |
(grabbed, frame) = stream.read() | |
# if the `grabbed` boolean is `False`, then we have | |
# reached the end of the video file | |
if not grabbed: | |
self.Q.put((None, None, None, None)) | |
print('===========================> This video get ' + str(k) + ' frames in total.') | |
sys.stdout.flush() | |
return | |
# process and add the frame to the queue | |
img_k, orig_img_k, im_dim_list_k = prep_frame(frame, inp_dim) | |
img.append(img_k) | |
orig_img.append(orig_img_k) | |
im_name.append(str(k) + '.jpg') | |
im_dim_list.append(im_dim_list_k) | |
with torch.no_grad(): | |
# Human Detection | |
img = torch.cat(img) | |
im_dim_list = torch.FloatTensor(im_dim_list).repeat(1, 2) | |
while self.Q.full(): | |
time.sleep(2) | |
self.Q.put((img, orig_img, im_name, im_dim_list)) | |
def videoinfo(self): | |
# indicate the video info | |
fourcc = int(self.stream.get(cv2.CAP_PROP_FOURCC)) | |
fps = self.stream.get(cv2.CAP_PROP_FPS) | |
frameSize = (int(self.stream.get(cv2.CAP_PROP_FRAME_WIDTH)), int(self.stream.get(cv2.CAP_PROP_FRAME_HEIGHT))) | |
return (fourcc, fps, frameSize) | |
def getitem(self): | |
# return next frame in the queue | |
return self.Q.get() | |
def len(self): | |
return self.Q.qsize() | |
class DetectionLoader: | |
def __init__(self, dataloder, batchSize=1, queueSize=1024): | |
# initialize the file video stream along with the boolean | |
# used to indicate if the thread should be stopped or not | |
self.det_model = Darknet("joints_detectors/Alphapose/yolo/cfg/yolov3-spp.cfg") | |
self.det_model.load_weights('joints_detectors/Alphapose/models/yolo/yolov3-spp.weights') | |
self.det_model.net_info['height'] = opt.inp_dim | |
self.det_inp_dim = int(self.det_model.net_info['height']) | |
assert self.det_inp_dim % 32 == 0 | |
assert self.det_inp_dim > 32 | |
self.det_model | |
self.det_model.eval() | |
self.stopped = False | |
self.dataloder = dataloder | |
self.batchSize = batchSize | |
self.datalen = self.dataloder.length() | |
leftover = 0 | |
if (self.datalen) % batchSize: | |
leftover = 1 | |
self.num_batches = self.datalen // batchSize + leftover | |
# initialize the queue used to store frames read from | |
# the video file | |
if opt.sp: | |
self.Q = Queue(maxsize=queueSize) | |
else: | |
self.Q = mp.Queue(maxsize=queueSize) | |
def start(self): | |
# start a thread to read frames from the file video stream | |
if opt.sp: | |
t = Thread(target=self.update, args=()) | |
t.daemon = True | |
t.start() | |
else: | |
p = mp.Process(target=self.update, args=(), daemon=True) | |
# p = mp.Process(target=self.update, args=()) | |
# p.daemon = True | |
p.start() | |
return self | |
def update(self): | |
# keep looping the whole dataset | |
for i in range(self.num_batches): | |
img, orig_img, im_name, im_dim_list = self.dataloder.getitem() | |
if img is None: | |
self.Q.put((None, None, None, None, None, None, None)) | |
return | |
with torch.no_grad(): | |
# Human Detection | |
img = img | |
prediction = self.det_model(img, CUDA=True) | |
# NMS process | |
dets = dynamic_write_results(prediction, opt.confidence, | |
opt.num_classes, nms=True, nms_conf=opt.nms_thesh) | |
if isinstance(dets, int) or dets.shape[0] == 0: | |
for k in range(len(orig_img)): | |
if self.Q.full(): | |
time.sleep(2) | |
self.Q.put((orig_img[k], im_name[k], None, None, None, None, None)) | |
continue | |
dets = dets.cpu() | |
im_dim_list = torch.index_select(im_dim_list, 0, dets[:, 0].long()) | |
scaling_factor = torch.min(self.det_inp_dim / im_dim_list, 1)[0].view(-1, 1) | |
# coordinate transfer | |
dets[:, [1, 3]] -= (self.det_inp_dim - scaling_factor * im_dim_list[:, 0].view(-1, 1)) / 2 | |
dets[:, [2, 4]] -= (self.det_inp_dim - scaling_factor * im_dim_list[:, 1].view(-1, 1)) / 2 | |
dets[:, 1:5] /= scaling_factor | |
for j in range(dets.shape[0]): | |
dets[j, [1, 3]] = torch.clamp(dets[j, [1, 3]], 0.0, im_dim_list[j, 0]) | |
dets[j, [2, 4]] = torch.clamp(dets[j, [2, 4]], 0.0, im_dim_list[j, 1]) | |
boxes = dets[:, 1:5] | |
scores = dets[:, 5:6] | |
for k in range(len(orig_img)): | |
boxes_k = boxes[dets[:, 0] == k] | |
if isinstance(boxes_k, int) or boxes_k.shape[0] == 0: | |
if self.Q.full(): | |
time.sleep(2) | |
self.Q.put((orig_img[k], im_name[k], None, None, None, None, None)) | |
continue | |
inps = torch.zeros(boxes_k.size(0), 3, opt.inputResH, opt.inputResW) | |
pt1 = torch.zeros(boxes_k.size(0), 2) | |
pt2 = torch.zeros(boxes_k.size(0), 2) | |
if self.Q.full(): | |
time.sleep(2) | |
self.Q.put((orig_img[k], im_name[k], boxes_k, scores[dets[:, 0] == k], inps, pt1, pt2)) | |
def read(self): | |
# return next frame in the queue | |
return self.Q.get() | |
def len(self): | |
# return queue len | |
return self.Q.qsize() | |
class DetectionProcessor: | |
def __init__(self, detectionLoader, queueSize=1024): | |
# initialize the file video stream along with the boolean | |
# used to indicate if the thread should be stopped or not | |
self.detectionLoader = detectionLoader | |
self.stopped = False | |
self.datalen = self.detectionLoader.datalen | |
# initialize the queue used to store data | |
if opt.sp: | |
self.Q = Queue(maxsize=queueSize) | |
else: | |
self.Q = pQueue(maxsize=queueSize) | |
def start(self): | |
# start a thread to read frames from the file video stream | |
if opt.sp: | |
# t = Thread(target=self.update, args=(), daemon=True) | |
t = Thread(target=self.update, args=()) | |
t.daemon = True | |
t.start() | |
else: | |
p = mp.Process(target=self.update, args=(), daemon=True) | |
# p = mp.Process(target=self.update, args=()) | |
# p.daemon = True | |
p.start() | |
return self | |
def update(self): | |
# keep looping the whole dataset | |
for i in range(self.datalen): | |
with torch.no_grad(): | |
(orig_img, im_name, boxes, scores, inps, pt1, pt2) = self.detectionLoader.read() | |
if orig_img is None: | |
self.Q.put((None, None, None, None, None, None, None)) | |
return | |
if boxes is None or boxes.nelement() == 0: | |
while self.Q.full(): | |
time.sleep(0.2) | |
self.Q.put((None, orig_img, im_name, boxes, scores, None, None)) | |
continue | |
inp = im_to_torch(cv2.cvtColor(orig_img, cv2.COLOR_BGR2RGB)) | |
inps, pt1, pt2 = crop_from_dets(inp, boxes, inps, pt1, pt2) | |
while self.Q.full(): | |
time.sleep(0.2) | |
self.Q.put((inps, orig_img, im_name, boxes, scores, pt1, pt2)) | |
def read(self): | |
# return next frame in the queue | |
return self.Q.get() | |
def len(self): | |
# return queue len | |
return self.Q.qsize() | |
class VideoDetectionLoader: | |
def __init__(self, path, batchSize=4, queueSize=256): | |
# initialize the file video stream along with the boolean | |
# used to indicate if the thread should be stopped or not | |
self.det_model = Darknet("yolo/cfg/yolov3-spp.cfg") | |
self.det_model.load_weights('models/yolo/yolov3-spp.weights') | |
self.det_model.net_info['height'] = opt.inp_dim | |
self.det_inp_dim = int(self.det_model.net_info['height']) | |
assert self.det_inp_dim % 32 == 0 | |
assert self.det_inp_dim > 32 | |
self.det_model | |
self.det_model.eval() | |
self.stream = cv2.VideoCapture(path) | |
assert self.stream.isOpened(), 'Cannot capture source' | |
self.stopped = False | |
self.batchSize = batchSize | |
self.datalen = int(self.stream.get(cv2.CAP_PROP_FRAME_COUNT)) | |
leftover = 0 | |
if (self.datalen) % batchSize: | |
leftover = 1 | |
self.num_batches = self.datalen // batchSize + leftover | |
# initialize the queue used to store frames read from | |
# the video file | |
self.Q = Queue(maxsize=queueSize) | |
def length(self): | |
return self.datalen | |
def len(self): | |
return self.Q.qsize() | |
def start(self): | |
# start a thread to read frames from the file video stream | |
t = Thread(target=self.update, args=()) | |
t.daemon = True | |
t.start() | |
return self | |
def update(self): | |
# keep looping the whole video | |
for i in range(self.num_batches): | |
img = [] | |
inp = [] | |
orig_img = [] | |
im_name = [] | |
im_dim_list = [] | |
for k in range(i * self.batchSize, min((i + 1) * self.batchSize, self.datalen)): | |
(grabbed, frame) = self.stream.read() | |
# if the `grabbed` boolean is `False`, then we have | |
# reached the end of the video file | |
if not grabbed: | |
self.stop() | |
return | |
# process and add the frame to the queue | |
inp_dim = int(opt.inp_dim) | |
img_k, orig_img_k, im_dim_list_k = prep_frame(frame, inp_dim) | |
inp_k = im_to_torch(orig_img_k) | |
img.append(img_k) | |
inp.append(inp_k) | |
orig_img.append(orig_img_k) | |
im_dim_list.append(im_dim_list_k) | |
with torch.no_grad(): | |
ht = inp[0].size(1) | |
wd = inp[0].size(2) | |
# Human Detection | |
img = Variable(torch.cat(img)) | |
im_dim_list = torch.FloatTensor(im_dim_list).repeat(1, 2) | |
im_dim_list = im_dim_list | |
prediction = self.det_model(img, CUDA=True) | |
# NMS process | |
dets = dynamic_write_results(prediction, opt.confidence, | |
opt.num_classes, nms=True, nms_conf=opt.nms_thesh) | |
if isinstance(dets, int) or dets.shape[0] == 0: | |
for k in range(len(inp)): | |
while self.Q.full(): | |
time.sleep(0.2) | |
self.Q.put((inp[k], orig_img[k], None, None)) | |
continue | |
im_dim_list = torch.index_select(im_dim_list, 0, dets[:, 0].long()) | |
scaling_factor = torch.min(self.det_inp_dim / im_dim_list, 1)[0].view(-1, 1) | |
# coordinate transfer | |
dets[:, [1, 3]] -= (self.det_inp_dim - scaling_factor * im_dim_list[:, 0].view(-1, 1)) / 2 | |
dets[:, [2, 4]] -= (self.det_inp_dim - scaling_factor * im_dim_list[:, 1].view(-1, 1)) / 2 | |
dets[:, 1:5] /= scaling_factor | |
for j in range(dets.shape[0]): | |
dets[j, [1, 3]] = torch.clamp(dets[j, [1, 3]], 0.0, im_dim_list[j, 0]) | |
dets[j, [2, 4]] = torch.clamp(dets[j, [2, 4]], 0.0, im_dim_list[j, 1]) | |
boxes = dets[:, 1:5].cpu() | |
scores = dets[:, 5:6].cpu() | |
for k in range(len(inp)): | |
while self.Q.full(): | |
time.sleep(0.2) | |
self.Q.put((inp[k], orig_img[k], boxes[dets[:, 0] == k], scores[dets[:, 0] == k])) | |
def videoinfo(self): | |
# indicate the video info | |
fourcc = int(self.stream.get(cv2.CAP_PROP_FOURCC)) | |
fps = self.stream.get(cv2.CAP_PROP_FPS) | |
frameSize = (int(self.stream.get(cv2.CAP_PROP_FRAME_WIDTH)), int(self.stream.get(cv2.CAP_PROP_FRAME_HEIGHT))) | |
return (fourcc, fps, frameSize) | |
def read(self): | |
# return next frame in the queue | |
return self.Q.get() | |
def more(self): | |
# return True if there are still frames in the queue | |
return self.Q.qsize() > 0 | |
def stop(self): | |
# indicate that the thread should be stopped | |
self.stopped = True | |
class WebcamLoader: | |
def __init__(self, webcam, queueSize=256): | |
# initialize the file video stream along with the boolean | |
# used to indicate if the thread should be stopped or not | |
self.stream = cv2.VideoCapture(int(webcam)) | |
assert self.stream.isOpened(), 'Cannot capture source' | |
self.stopped = False | |
# initialize the queue used to store frames read from | |
# the video file | |
self.Q = LifoQueue(maxsize=queueSize) | |
def start(self): | |
# start a thread to read frames from the file video stream | |
t = Thread(target=self.update, args=()) | |
t.daemon = True | |
t.start() | |
return self | |
def update(self): | |
# keep looping infinitely | |
while True: | |
# otherwise, ensure the queue has room in it | |
if not self.Q.full(): | |
# read the next frame from the file | |
(grabbed, frame) = self.stream.read() | |
# if the `grabbed` boolean is `False`, then we have | |
# reached the end of the video file | |
if not grabbed: | |
self.stop() | |
return | |
# process and add the frame to the queue | |
inp_dim = int(opt.inp_dim) | |
img, orig_img, dim = prep_frame(frame, inp_dim) | |
inp = im_to_torch(orig_img) | |
im_dim_list = torch.FloatTensor([dim]).repeat(1, 2) | |
self.Q.put((img, orig_img, inp, im_dim_list)) | |
else: | |
with self.Q.mutex: | |
self.Q.queue.clear() | |
def videoinfo(self): | |
# indicate the video info | |
fourcc = int(self.stream.get(cv2.CAP_PROP_FOURCC)) | |
fps = self.stream.get(cv2.CAP_PROP_FPS) | |
frameSize = (int(self.stream.get(cv2.CAP_PROP_FRAME_WIDTH)), int(self.stream.get(cv2.CAP_PROP_FRAME_HEIGHT))) | |
return (fourcc, fps, frameSize) | |
def read(self): | |
# return next frame in the queue | |
return self.Q.get() | |
def len(self): | |
# return queue size | |
return self.Q.qsize() | |
def stop(self): | |
# indicate that the thread should be stopped | |
self.stopped = True | |
class DataWriter: | |
def __init__(self, save_video=False, | |
savepath='examples/res/1.avi', fourcc=cv2.VideoWriter_fourcc(*'XVID'), fps=25, frameSize=(640, 480), | |
queueSize=1024): | |
if save_video: | |
# initialize the file video stream along with the boolean | |
# used to indicate if the thread should be stopped or not | |
self.stream = cv2.VideoWriter(savepath, fourcc, fps, frameSize) | |
assert self.stream.isOpened(), 'Cannot open video for writing' | |
self.save_video = save_video | |
self.stopped = False | |
self.final_result = [] | |
# initialize the queue used to store frames read from | |
# the video file | |
self.Q = Queue(maxsize=queueSize) | |
if opt.save_img: | |
if not os.path.exists(opt.outputpath + '/vis'): | |
os.mkdir(opt.outputpath + '/vis') | |
def start(self): | |
# start a thread to read frames from the file video stream | |
t = Thread(target=self.update, args=(), daemon=True) | |
# t = Thread(target=self.update, args=()) | |
# t.daemon = True | |
t.start() | |
return self | |
def update(self): | |
# keep looping infinitely | |
while True: | |
# if the thread indicator variable is set, stop the | |
# thread | |
if self.stopped: | |
if self.save_video: | |
self.stream.release() | |
return | |
# otherwise, ensure the queue is not empty | |
if not self.Q.empty(): | |
(boxes, scores, hm_data, pt1, pt2, orig_img, im_name) = self.Q.get() | |
orig_img = np.array(orig_img, dtype=np.uint8) | |
if boxes is None: | |
if opt.save_img or opt.save_video or opt.vis: | |
img = orig_img | |
if opt.vis: | |
cv2.imshow("AlphaPose Demo", img) | |
cv2.waitKey(30) | |
if opt.save_img: | |
cv2.imwrite(os.path.join(opt.outputpath, 'vis', im_name), img) | |
if opt.save_video: | |
self.stream.write(img) | |
else: | |
# location prediction (n, kp, 2) | score prediction (n, kp, 1) | |
if opt.matching: | |
preds = getMultiPeakPrediction( | |
hm_data, pt1.numpy(), pt2.numpy(), opt.inputResH, opt.inputResW, opt.outputResH, opt.outputResW) | |
result = matching(boxes, scores.numpy(), preds) | |
else: | |
preds_hm, preds_img, preds_scores = getPrediction( | |
hm_data, pt1, pt2, opt.inputResH, opt.inputResW, opt.outputResH, opt.outputResW) | |
result = pose_nms( | |
boxes, scores, preds_img, preds_scores) | |
result = { | |
'imgname': im_name, | |
'result': result | |
} | |
self.final_result.append(result) | |
if opt.save_img or opt.save_video or opt.vis: | |
img = vis_frame(orig_img, result) | |
if opt.vis: | |
cv2.imshow("AlphaPose Demo", img) | |
cv2.waitKey(30) | |
if opt.save_img: | |
cv2.imwrite(os.path.join(opt.outputpath, 'vis', im_name), img) | |
if opt.save_video: | |
self.stream.write(img) | |
else: | |
time.sleep(0.1) | |
def running(self): | |
# indicate that the thread is still running | |
time.sleep(0.2) | |
return not self.Q.empty() | |
def save(self, boxes, scores, hm_data, pt1, pt2, orig_img, im_name): | |
# save next frame in the queue | |
self.Q.put((boxes, scores, hm_data, pt1, pt2, orig_img, im_name)) | |
def stop(self): | |
# indicate that the thread should be stopped | |
self.stopped = True | |
time.sleep(0.2) | |
def results(self): | |
# return final result | |
return self.final_result | |
def len(self): | |
# return queue len | |
return self.Q.qsize() | |
class Mscoco(data.Dataset): | |
def __init__(self, train=True, sigma=1, | |
scale_factor=(0.2, 0.3), rot_factor=40, label_type='Gaussian'): | |
self.img_folder = '../data/coco/images' # root image folders | |
self.is_train = train # training set or test set | |
self.inputResH = opt.inputResH | |
self.inputResW = opt.inputResW | |
self.outputResH = opt.outputResH | |
self.outputResW = opt.outputResW | |
self.sigma = sigma | |
self.scale_factor = scale_factor | |
self.rot_factor = rot_factor | |
self.label_type = label_type | |
self.nJoints_coco = 17 | |
self.nJoints_mpii = 16 | |
self.nJoints = 33 | |
self.accIdxs = (1, 2, 3, 4, 5, 6, 7, 8, | |
9, 10, 11, 12, 13, 14, 15, 16, 17) | |
self.flipRef = ((2, 3), (4, 5), (6, 7), | |
(8, 9), (10, 11), (12, 13), | |
(14, 15), (16, 17)) | |
def __getitem__(self, index): | |
pass | |
def __len__(self): | |
pass | |
def crop_from_dets(img, boxes, inps, pt1, pt2): | |
''' | |
Crop human from origin image according to Dectecion Results | |
''' | |
imght = img.size(1) | |
imgwidth = img.size(2) | |
tmp_img = img | |
tmp_img[0].add_(-0.406) | |
tmp_img[1].add_(-0.457) | |
tmp_img[2].add_(-0.480) | |
for i, box in enumerate(boxes): | |
upLeft = torch.Tensor( | |
(float(box[0]), float(box[1]))) | |
bottomRight = torch.Tensor( | |
(float(box[2]), float(box[3]))) | |
ht = bottomRight[1] - upLeft[1] | |
width = bottomRight[0] - upLeft[0] | |
scaleRate = 0.3 | |
upLeft[0] = max(0, upLeft[0] - width * scaleRate / 2) | |
upLeft[1] = max(0, upLeft[1] - ht * scaleRate / 2) | |
bottomRight[0] = max( | |
min(imgwidth - 1, bottomRight[0] + width * scaleRate / 2), upLeft[0] + 5) | |
bottomRight[1] = max( | |
min(imght - 1, bottomRight[1] + ht * scaleRate / 2), upLeft[1] + 5) | |
try: | |
inps[i] = cropBox(tmp_img.clone(), upLeft, bottomRight, opt.inputResH, opt.inputResW) | |
except IndexError: | |
print(tmp_img.shape) | |
print(upLeft) | |
print(bottomRight) | |
print('===') | |
pt1[i] = upLeft | |
pt2[i] = bottomRight | |
return inps, pt1, pt2 | |