1 files changed, 125 insertions, 0 deletions
diff --git a/data/label.py b/data/label.py
new file mode 100755
index 0000000..4ee6bd2
--- /dev/null
+++ b/data/label.py
@@ -0,0 +1,125 @@
+#!/usr/bin/python
+
+import cv
+import sys
+import math
+import pickle
+import signal
+
+sk_file = open(sys.argv[1])
+pic_dir = sys.argv[2].rstrip('/') + '/' + sys.argv[1].split('/')[-2] + '/'
+out_dir = sys.argv[1][0:sys.argv[1].rfind('/')+1]
+
+j1 = 'Head'
+j2 = 'ShoulderCenter'
+
+def display(lines):
+    global j1,j2
+    limbs = {}
+    pframe = 0
+    while cv.WaitKey(5) != -1: pass
+    for line in lines:
+        run, frame, nid, joint, state, x, y, z, X, Y = line.strip().split(',')
+        frame = int(frame)
+        if frame != pframe:
+            if j1 in limbs and j2 in limbs:
+                dx = limbs[j1][0] - limbs[j2][0]
+                dy = limbs[j1][1] - limbs[j2][1]
+
+                img = cv.LoadImage(pic_dir+str(pframe)+".jpg")
+                r = math.sqrt(dx*dx+dy*dy)
+                x1 = max(int(limbs[j1][0] - r),0)
+                y1 = max(int(limbs[j1][1] - r),0)
+                x2 = min(int(x1 + 2*r), img.width)
+                y2 = min(int(y1 + 3*r), img.height)
+
+                cv.Rectangle(img,(x1,y1),(x2,y2),(0,255,0),3,8,0)
+                #cv.SaveImage(out_dir+str(run)+"-"+str(frame)+'.jpg',img)
+                cv.ShowImage('video', img)
+                k = cv.WaitKey(5)
+                if k != -1:
+                    return chr(k&255)
+            pframe = frame 
+            limbs = {}
+        id = int(nid)
+
+        limbs[joint] = [int(X)*640.0/400.0,int(Y)*480.0/300.0]
+    return chr(cv.WaitKey()&255)
+
+def quit(save):
+    global people,labels
+    print save
+    if save:
+        pickle.dump(people,open(out_dir+'people.dump','w'))
+        pickle.dump(labels,open(out_dir+'labels.dump','w'))
+    sys.exit(0)
+
+prun = 0
+lines = []
+people = []
+#chars = '123456789qwerasdfzxcvQWERASDFZXCV'
+chars = 'abcdefghijklmnopqrstuvwxyz'
+try:
+    people = pickle.load(open(out_dir+'people.dump'))
+    people.sort()
+except:
+    pass
+labels = {}
+try:
+    labels = pickle.load(open(out_dir+'labels.dump'))
+except:
+    pass
+signal.signal(signal.SIGINT, lambda *args: quit(True))
+signal.signal(signal.SIGQUIT, lambda *args: quit(False))
+for line in sk_file:
+    run, frame, nid, joint, state, x, y, z, X, Y = line.strip().split(',')
+    if frame != 'Frame':
+        run = int(run)
+        if run != prun:
+            if prun != 0 and prun not in labels:
+                k = '\n'
+                while k == '\n':
+                    print "Run #",prun
+                    print "0: None"
+                    for i in range(len(people)):
+                        print chars[i]+": "+people[i]
+                    print "Enter: Replay"
+                    #print "N: New person"
+                    print "?: Unknown"
+                    k = display(lines)
+                    # Key pressed is SHIFT, grab next character and make it uppercase
+                    if k == '\xe1':
+                        k = chr(cv.WaitKey()&255).upper()
+
+                    # Key pressed is CTRL, grab next character and make quit if 'c','d', or '\' follows
+                    if k == '\xe3':
+                        k = chr(cv.WaitKey()&255)
+                        if k == 'c':
+                            quit(True)
+                        elif k == 'd' or k == '\\':
+                            quit(False)
+
+                    if k == '0':
+                        labels[prun] = 'None'
+                    elif k == '?':
+                        labels[prun] = '?'
+                    elif k in chars:
+                        labels[prun] = people[chars.find(k)]
+                    elif k == 'N':
+                        name = raw_input("Type name: ")
+                        if people != "":
+                            people += name.split(',')
+                            people.sort()
+                        #labels[prun] = name
+                        k = '\n'
+                    elif k != '~':
+                        k = '\n'
+                    if cv.WaitKey(1000) != -1:
+                        k = '\n'
+            prun = run
+            lines = []
+        lines += [line]
+quit(True)
+
+
+