# Generating JSON file compatible with ViP for HMDB51 dataset

import json
import os
import sys

splits      = ['train', 'test']
source_root = '$DATA_DIRECTORY/HMDB51'
target_root = '$JSON_DIRECTORY'

dataset_type = 'test'
base_path    = os.path.join('/z/dat/UCF101', dataset_type+'images')
actions      = os.listdir(base_path)

dest_path = 'test.json'
dest_data = [] 
count     = 0
samples   = []

label_dict = {'ApplyEyeMakeup':0
,'ApplyLipstick':1
,'Archery':2
,'BabyCrawling':3
,'BalanceBeam':4
,'BandMarching':5
,'BaseballPitch':6
,'Basketball':7
,'BasketballDunk':8
,'BenchPress':9
,'Biking':10
,'Billiards':11
,'BlowDryHair':12
,'BlowingCandles':13
,'BodyWeightSquats':14
,'Bowling':15
,'BoxingPunchingBag':16
,'BoxingSpeedBag':17
,'BreastStroke':18
,'BrushingTeeth':19
,'CleanAndJerk':20
,'CliffDiving':21
,'CricketBowling':22
,'CricketShot':23
,'CuttingInKitchen':24
,'Diving':25
,'Drumming':26
,'Fencing':27
,'FieldHockeyPenalty':28
,'FloorGymnastics':29
,'FrisbeeCatch':30
,'FrontCrawl':31
,'GolfSwing':32
,'Haircut':33
,'Hammering':34
,'HammerThrow':35
,'HandstandPushups':36
,'HandstandWalking':37
,'HeadMassage':38
,'HighJump':39
,'HorseRace':40
,'HorseRiding':41
,'HulaHoop':42
,'IceDancing':43
,'JavelinThrow':44
,'JugglingBalls':45
,'JumpingJack':46
,'JumpRope':47
,'Kayaking':48
,'Knitting':49
,'LongJump':50
,'Lunges':51
,'MilitaryParade':52
,'Mixing':53
,'MoppingFloor':54
,'Nunchucks':55
,'ParallelBars':56
,'PizzaTossing':57
,'PlayingCello':58
,'PlayingDaf':59
,'PlayingDhol':60
,'PlayingFlute':61
,'PlayingGuitar':62
,'PlayingPiano':63
,'PlayingSitar':64
,'PlayingTabla':65
,'PlayingViolin':66
,'PoleVault':67
,'PommelHorse':68
,'PullUps':69
,'Punch':70
,'PushUps':71
,'Rafting':72
,'RockClimbingIndoor':73
,'RopeClimbing':74
,'Rowing':75
,'SalsaSpin':76
,'ShavingBeard':77
,'Shotput':78
,'SkateBoarding':79
,'Skiing':80
,'Skijet':81
,'SkyDiving':82
,'SoccerJuggling':83
,'SoccerPenalty':84
,'StillRings':85
,'SumoWrestling':86
,'Surfing':87
,'Swing':88
,'TableTennisShot':89
,'TaiChi':90
,'TennisSwing':91
,'ThrowDiscus':92
,'TrampolineJumping':93
,'Typing':94
,'UnevenBars':95
,'VolleyballSpiking':96
,'WalkingWithDog':97
,'WallPushups':98
,'WritingOnBoard':99
,'YoYo':100}


for split in splits:
    base_path = os.path.join(source_root, split+'images')
    actions   = os.listdir(base_path) 

    dest_path = os.path.join(target_root, split+'.json')
    dest_data = [] 
    samples   = []
    count     = 0

    for action in actions:
        for video in os.listdir(os.path.join(base_path, action)):
            if not '.DS' in video:
                video_images = sorted(os.listdir(os.path.join(base_path, action, video.replace('.avi',''))))
                samples      = [os.path.join(base_path, action, video.replace('.avi',''), video_image) for video_image in video_images]
                    
                dest_data_lvl1           = {}
                dest_data_lvl1['frames'] = [] 
    
                for frame in samples:
                    dest_data_lvl1['frames'].append({'img_path': os.path.split(frame)[1],
                                                     'actions':[{'action_class': label_dict[action]}] })
    
                # END FOR
               
                dest_data_lvl1['base_path'] = os.path.split(frame)[0]
                dest_data.append(dest_data_lvl1)     
    
            # END IF
    
        # END FOR
    
    # END FOR
    
    with open(dest_path, 'w') as outfile:
        json.dump(dest_data, outfile, indent=4)