# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: privateai
#     language: python
#     name: privateai
# ---

# + colab={"base_uri": "https://localhost:8080/", "height": 35} colab_type="code" executionInfo={"elapsed": 495, "status": "ok", "timestamp": 1559977269702, "user": {"displayName": "<NAME>\u00e1rez", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="1ghOuV1nJ8Uy" outputId="856d49bb-cdea-4049-85e2-31540f9c9bf4"
import os
os.getcwd()

# + [markdown] colab_type="text" id="S-uYW5ENLe0V"
# # Mount File System for Working on Colaboratory

# + [markdown] colab_type="text" id="pF272MmRLbvf"
# # Running or importing .py Files

# + colab={"base_uri": "https://localhost:8080/", "height": 129} colab_type="code" executionInfo={"elapsed": 38352, "status": "ok", "timestamp": 1559977307596, "user": {"displayName": "<NAME>\u00e1rez", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="C6pH56x2J8ud" outputId="21aedae5-0c63-4637-8331-a4a383d62b8f"
from google.colab import drive
drive.mount('/content/drive')

# + [markdown] colab_type="text" id="yF6_T35WDdRN"
# # Change working directory

# + colab={"base_uri": "https://localhost:8080/", "height": 35} colab_type="code" executionInfo={"elapsed": 39186, "status": "ok", "timestamp": 1559977308454, "user": {"displayName": "<NAME>\u00e1rez", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="4zAzX0u-KBAg" outputId="bd65b04e-740a-4282-dfcb-6fa81be062fe"
# !ls

# + colab={"base_uri": "https://localhost:8080/", "height": 256} colab_type="code" executionInfo={"elapsed": 40012, "status": "ok", "timestamp": 1559977309301, "user": {"displayName": "<NAME>\u00e1rez", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="tyrxIXPtDvhp" outputId="39862e61-f18e-4354-8d48-3538b6a765c4"
# !ls 'drive/My Drive/pytorch-EverybodyDanceNow/notebook'

# + colab={} colab_type="code" id="GSsevlhMDcJ0"
import os
os.chdir('drive/My Drive/pytorch-EverybodyDanceNow/notebook')

# + colab={} colab_type="code" id="uC5YXwhGJ8zw"
import os
import numpy as np
import torch
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
torch.multiprocessing.set_sharing_strategy('file_system')
torch.backends.cudnn.benchmark = True
torch.cuda.set_device(0)



import time
import pickle
import matplotlib.pyplot as plt
from collections import OrderedDict
from torch.autograd import Variable
from pathlib import Path

# %matplotlib inline

# + colab={"base_uri": "https://localhost:8080/", "height": 35} colab_type="code" executionInfo={"elapsed": 583, "status": "ok", "timestamp": 1559977328707, "user": {"displayName": "<NAME>\u00e1rez", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="QHonnhjbJ83k" outputId="5aca6dcf-e6d2-4138-bef7-7e68d0e59226"
pix2pixhd_dir = Path('../src/pix2pixHD/')

import sys
sys.path.append(str(pix2pixhd_dir))

print(str(pix2pixhd_dir))

# + colab={} colab_type="code" id="tWwknmiyJ8_k"
# %load_ext autoreload

# + colab={} colab_type="code" id="gT_dK2Z2J9DA"
# %autoreload 2

# + [markdown] colab_type="text" id="xpDpRZ3_95hm"
# # Download and extract video

# + colab={} colab_type="code" id="yUU8C0qz95hr"
import cv2
from pathlib import Path

# + colab={"base_uri": "https://localhost:8080/", "height": 54} colab_type="code" executionInfo={"elapsed": 444, "status": "ok", "timestamp": 1559977389983, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="2uDRO1tb95h3" outputId="50b81c64-2945-40e4-d279-61ec9883fc74"
save_dir = Path('../data/source/')
print("Source data dir: ",save_dir)

try:
    save_dir.mkdir(True)

except FileExistsError:
    print("Directory exists")
    

img_dir = save_dir.joinpath('images')

print(img_dir)

try:
    img_dir.mkdir(True)
except FileExistsError:
    print("Directory exists")





# + [markdown] colab_type="text" id="nWstnWlz95h-"
# # Get Frames from Video

# + colab={"base_uri": "https://localhost:8080/", "height": 72} colab_type="code" executionInfo={"elapsed": 134118, "status": "ok", "timestamp": 1559977577787, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="ynlqklIW95iB" outputId="df94a4c6-3413-45df-9aa4-8b5b3084d761"
print(str(save_dir.joinpath('file_name.mp4')))
cap = cv2.VideoCapture( str(save_dir.joinpath('file_name.mp4')))
print(cap.isOpened())

i = 0
while(cap.isOpened()):
    flag, frame = cap.read()
    
    if flag == False:
        break 
    
    #frame = cv2.resize(frame, dsize=(512, 512), interpolation=cv2.INTER_CUBIC)
    cv2.imwrite(str(img_dir.joinpath('{:05d}.png'.format(i))), frame)
    
    #if i >= 300:
    #    break
    i += 1
print("Number of frames:", i)

# + [markdown] colab_type="text" id="8IGkvNJk95iJ"
# # Pose estimation (OpenPose)

# + colab={} colab_type="code" id="8rky8wrY95iL"
import numpy as np
import matplotlib.pyplot as plt
import torch
from tqdm import tqdm

# %matplotlib inline

# + colab={} colab_type="code" id="F9q3gvet95iS"
openpose_dir = Path('../src/pytorch_Realtime_Multi-Person_Pose_Estimation/')

import sys
sys.path.append(str(openpose_dir))
sys.path.append('../src/utils')


# -

# # Install packages with pip 

import sys
# #!{sys.executable} -m pip install scikit-image

# + colab={} colab_type="code" id="VbMgYY1P95if"
# openpose
from network.rtpose_vgg import get_model
from evaluate.coco_eval import get_multiplier, get_outputs

# utils
from openpose_utils import remove_noise, get_pose

# + colab={"base_uri": "https://localhost:8080/", "height": 3801} colab_type="code" executionInfo={"elapsed": 11978, "status": "ok", "timestamp": 1559977881884, "user": {"displayName": "<NAME>\u00e1rez", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="2WtMvF4F95ip" outputId="a5523570-bd3b-419f-a51d-e3abe09ae2a8"
weight_name = openpose_dir.joinpath('network/weight/pose_model.pth')

model = get_model('vgg19')     
model.load_state_dict(torch.load(weight_name))
model = torch.nn.DataParallel(model).cuda()
model.float()
model.eval()


# + [markdown] colab_type="text" id="JGrHBf2W95i2"
# ## check

# + colab={"base_uri": "https://localhost:8080/", "height": 324} colab_type="code" executionInfo={"elapsed": 525, "status": "ok", "timestamp": 1559977915339, "user": {"displayName": "<NAME>\u00e1rez", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="4cFKPhOn95i5" outputId="ba441250-bef4-4b3b-995f-91dd1a285fb0"

# %matplotlib inline  

print(img_dir)

print(len(os.listdir(img_dir)))

img_path = sorted(img_dir.iterdir())[5]
img = cv2.imread(str(img_path))
shape_dst = np.min(img.shape[:2])
# offset
oh = (img.shape[0] - shape_dst) // 2
ow = (img.shape[1] - shape_dst) // 2

img = img[oh:oh+shape_dst, ow:ow+shape_dst]
img = cv2.resize(img, (512, 512))
          
plt.imshow(img[:,:,[2, 1, 0]]) # BGR -> RGB

# + colab={"base_uri": "https://localhost:8080/", "height": 287} colab_type="code" executionInfo={"elapsed": 3658, "status": "ok", "timestamp": 1559977921811, "user": {"displayName": "<NAME>\u00e1rez", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="XPl3570E95jD" outputId="8c98c1a3-a362-4206-edc7-48121d09b0a8"
multiplier = get_multiplier(img)
with torch.no_grad():
    paf, heatmap = get_outputs(multiplier, img, model, 'rtpose')
    
r_heatmap = np.array([remove_noise(ht)
                      for ht in heatmap.transpose(2, 0, 1)[:-1]])\
                     .transpose(1, 2, 0)
heatmap[:, :, :-1] = r_heatmap
param = {'thre1': 0.1, 'thre2': 0.05, 'thre3': 0.5}
label,cord = get_pose(param, heatmap, paf)

print("Cord", cord, len(cord))
plt.imshow(label)

# + [markdown] colab_type="text" id="lqziN8I695jL"
# ## make label images for pix2pix

# + colab={"base_uri": "https://localhost:8080/", "height": 35} colab_type="code" executionInfo={"elapsed": 6248329, "status": "ok", "timestamp": 1559984191699, "user": {"displayName": "<NAME>\u00e1rez", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="CDGn0V3X95jM" outputId="27fb5e15-5f5d-4808-bfd3-0f54a26c5990"
print(save_dir)
# Source images
test_img_dir = save_dir.joinpath('test_img')
try:
    test_img_dir.mkdir(exist_ok=True)
except FileExistsError:
    print("Directory exists")
    

# Label images
test_label_dir = save_dir.joinpath('test_label_ori')
try:
    test_label_dir.mkdir(exist_ok=True)
except FileExistsError:
    print("Directory exists")

# Head images
test_head_dir  = save_dir.joinpath('test_head_ori')
try:
    test_head_dir.mkdir(exist_ok = True)
except FileExistsError:
    print("Directory exists")

NUM_FRAMES = len(os.listdir(str(img_dir)))
print("Number of frames", NUM_FRAMES)


pose_cords = []

for idx in tqdm(range(NUM_FRAMES)):
    img_path = img_dir.joinpath('{:05d}.png'.format(idx))
    img = cv2.imread(str(img_path))
    shape_dst = np.min(img.shape[:2])
    oh = (img.shape[0] - shape_dst) // 2
    ow = (img.shape[1] - shape_dst) // 2

    img = img[oh:oh+shape_dst, ow:ow+shape_dst]
    img = cv2.resize(img, (512, 512))
    multiplier = get_multiplier(img)
    with torch.no_grad(): 
        paf, heatmap = get_outputs(multiplier, img, model, 'rtpose')
    r_heatmap = np.array([remove_noise(ht)
                      for ht in heatmap.transpose(2, 0, 1)[:-1]])\
                     .transpose(1, 2, 0)
    heatmap[:, :, :-1] = r_heatmap
    param = {'thre1': 0.1, 'thre2': 0.05, 'thre3': 0.5}
    label, cord  = get_pose(param, heatmap, paf)
    index = 13
    crop_size =  25
    try: 
        head_cord = cord[index]
    except:
        head_cord = pose_cords[-1] # if there is no head point in this picture.
        #Use las frame head point
        
    pose_cords.extend([head_cord])
    head = img[ int(head_cord[1] - crop_size)  :int(head_cord[1] + crop_size),
                int(head_cord[0] - crop_size ) :int(head_cord[0] + crop_size),:
              ]
    
    #Show head
    plt.imshow(head)
    plt.savefig(str( test_head_dir.joinpath('pose_{}.jpg'.format(idx))))
    plt.clf()
    
    
    cv2.imwrite(str(test_img_dir.joinpath( '{:05}.png'.format(idx))), img)
    cv2.imwrite(str(test_label_dir.joinpath( '{:05}.png'.format(idx))), label)
    
    if idx % 100 == 0 and idx != 0:
        pose_cords_arr = np.array(pose_cords, dtype = np.int)
        # Save pose coordinates array.
        np.save(str(save_dir.joinpath('pose_source.npy')), pose_cords_arr)

pose_cords_arr = np.array( pose_cords, dtype = np.int )
np.save( str(save_dir.joinpath('pose_source.npy')), pose_cords_arr )
        
    
torch.cuda.empty_cache()

# + colab={"base_uri": "https://localhost:8080/", "height": 35} colab_type="code" executionInfo={"elapsed": 279, "status": "ok", "timestamp": 1559977930126, "user": {"displayName": "<NAME>\u00e1rez", "photoUrl": "", "userId": "17609825801325148392"}, "user_tz": 300} id="_G1ZoynL95jW" outputId="77e50ffe-1ac4-474f-eb27-c9c62b68fd1e"
len(os.listdir(str(img_dir)))

# + colab={} colab_type="code" id="elT-wai3J0fb"


StarcoderdataJupyterScriptsDedupFiltered