gened yolo dataset

.gitignore

@@ -171,4 +171,5 @@ lerped_keypoints/*
 avg_keypoints/*
 video_frames/*
 video/*
-new_yolo_keypoints/*
+new_yolo_keypoints/*
+new_yolo_keypoints_bak/*

openpose_gen.py

@@ -104,7 +104,7 @@ def upload_image_circular_queue(image_path, size, unqiue_id, server_address):
 
     return image_name
 
-def main():
+def visualize_for_fixed_dataset(json_file: str):
     directory = './fixed'
     json_files = [f for f in os.listdir(directory) if f.endswith('.json')]
     if not json_files:
@@ -134,5 +134,30 @@ def main():
 
         cv2.imwrite(image_path + '_' + str(i) + '.png', canvas)
 
+def visualize_for_new_yolo_dataset():
+    directory = './new_yolo_keypoints'
+    json_files = [f for f in os.listdir(directory) if f.endswith('.json')]
+    if not json_files:
+        print("No JSON files found in the directory.")
+        return
+    
+    json_file = os.path.join(directory, random.choice(json_files))
+    # create ./output directory if it does not exist
+    os.makedirs('output', exist_ok=True)
+    
+    frames = pjf.read_new_yolo_keypoints(json_file)
+    
+    for i, skeletons in enumerate(frames):
+        canvas = np.zeros((360, 640, 3), dtype=np.uint8)
+        for skeleton in skeletons:
+            keypoints = skeleton.keypoints
+            canvas = skel.draw_bodypose(canvas, keypoints, skel.yolo_coco_limbSeq, skel.yolo_coco_colors)
+        
+        print(os.path.join('./output', os.path.basename(json_file) + '_' + str(i) + '.png'))
+        cv2.imwrite(os.path.join('./output', 'a_test' + '_' + str(i) + '.png'), canvas)
+        
+def main():
+    visualize_for_new_yolo_dataset()
+
 if __name__ == '__main__':
     main()

process_json_file.py

@@ -31,6 +31,34 @@ def array_json_to_Skeleton_Seqences(json_file: str) -> List[skel.Skeleton_Seqenc
             skeleton_sequences[i].add_frame(skeleton)
     return skeleton_sequences
 
+def read_new_yolo_keypoints(json_file: str) -> List[List[skel.Skeleton]]:
+    with open(json_file, 'r') as file:
+        data = json.load(file)
+    
+    frames = []
+    for frame in data:
+        skeletons = []
+        for skeleton in frame:
+            keypoints = [skel.Keypoint(point['x'], point['y'], point['confidence']) for point in skeleton['keypoints']]
+            skeletons.append(skel.Skeleton(keypoints))
+        frames.append(skeletons)
+    return frames
+
+def add_neck_to_og_coco(folder: str):
+    for file in os.listdir(folder):
+        with open(os.path.join(folder, file), 'r') as f:
+            data = json.load(f)
+        print(f"Processing {file}")
+        for frame in data:
+            for skeleton in frame:
+                _skeleton = skeleton['keypoints']
+                # Add neck keypoint to the original COCO keypoints
+                if len(_skeleton) > 6:
+                    neck = {'x': (_skeleton[6]['x'] + _skeleton[7]['x']) / 2, 'y': (_skeleton[6]['y'] + _skeleton[7]['y']) / 2, 'confidence': (_skeleton[6]['confidence'] + _skeleton[7]['confidence']) / 2}
+                    _skeleton.insert(6, neck)
+        with open(os.path.join(folder, file), 'w') as f:
+            json.dump(data, f, indent=4)
+
 def Skeleton_Seqences_save_to_array_json(skeleton_sequences: List[skel.Skeleton_Seqence], json_file: str):
     # Ensure the directory exists
     os.makedirs(os.path.dirname(json_file), exist_ok=True)
@@ -287,11 +315,11 @@ def main():
     video_path = './video'
     video_frame_path = './video_frames'
     
-
+    add_neck_to_og_coco('./new_yolo_keypoints')
 
 
     
-    print("Done processing all rows.")
+    # print("Done processing all rows.")
     # canvas = np.zeros((360, 640, 3), dtype=np.uint8)
     # canvas = skel.draw_bodypose(canvas, frames[0][0].keypoints, skel.body_25_limbSeq, skel.body_25_colors)
     

skeleton_lib.py

@@ -19,6 +19,19 @@ coco_colors = [
     [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]
 ]
 
+yolo_coco_colors = [
+    [255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
+    [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
+    [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]
+]
+
+yolo_coco_limbSeq = [
+    [1, 2], [1, 5], [2, 3], [3, 4], 
+    [5, 6], [6, 7], [1, 8], [8, 9], 
+    [9, 10], [1, 11], [11, 12], [12, 13], 
+    [1, 0], [0, 14], [14, 16], [0, 15], 
+]
+
 body_25_limbSeq = [
     [1, 8], [1, 2], [1, 5], [2, 3], 
     [3, 4], [5, 6], [6, 7], [8, 9], 
@@ -176,28 +189,28 @@ def draw_bodypose(canvas: ndarray, keypoints: List[Keypoint], limbSeq, colors, x
     if keypoints is None or len(keypoints) == 0:
         return canvas
 
-    for (k1_index, k2_index), color in zip(limbSeq, colors):
+    # for (k1_index, k2_index), color in zip(limbSeq, colors):
-        keypoint1 = keypoints[k1_index]
+    #     keypoint1 = keypoints[k1_index]
-        keypoint2 = keypoints[k2_index]
+    #     keypoint2 = keypoints[k2_index]
 
-        if keypoint1 is None or keypoint2 is None or keypoint1.confidence == 0 or keypoint2.confidence == 0:
+    #     if keypoint1 is None or keypoint2 is None or keypoint1.confidence == 0 or keypoint2.confidence == 0 or keypoint1.x <= 0 or keypoint1.y <= 0 or keypoint2.x <= 0 or keypoint2.y <= 0:
-            # if keypoint1 is None or keypoint1.confidence == 0:
+    #         # if keypoint1 is None or keypoint1.confidence == 0:
-            #     print(f"keypoint failed: {k1_index}")
+    #         #     print(f"keypoint failed: {k1_index}")
-            # if keypoint2 is None or keypoint2.confidence == 0:
+    #         # if keypoint2 is None or keypoint2.confidence == 0:
-            #     print(f"keypoint failed: {k2_index}")
+    #         #     print(f"keypoint failed: {k2_index}")
-            continue
+    #         continue
 
-        Y = np.array([keypoint1.x, keypoint2.x]) * float(W)
+    #     Y = np.array([keypoint1.x, keypoint2.x]) * float(W)
-        X = np.array([keypoint1.y, keypoint2.y]) * float(H)
+    #     X = np.array([keypoint1.y, keypoint2.y]) * float(H)
-        mX = np.mean(X)
+    #     mX = np.mean(X)
-        mY = np.mean(Y)
+    #     mY = np.mean(Y)
-        length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
+    #     length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
-        angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
+    #     angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
-        polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth*stick_scale), int(angle), 0, 360, 1)
+    #     polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth*stick_scale), int(angle), 0, 360, 1)
-        cv2.fillConvexPoly(canvas, polygon, [int(float(c) * 0.6) for c in color])
+    #     cv2.fillConvexPoly(canvas, polygon, [int(float(c) * 0.6) for c in color])
 
     for keypoint, color in zip(keypoints, colors):
-        if keypoint is None or keypoint.confidence == 0:
+        if keypoint is None or keypoint.confidence == 0 or keypoint.x <= 0 or keypoint.y <= 0:
             continue
 
         x, y = keypoint.x, keypoint.y

yolo_main.py

@@ -38,7 +38,7 @@ def get_yoloed_frames(results, lerped_keypoints_path):
     return frames
 
 def process_clip(row, model):
-    clip_name = row['ClipName']
+    clip_name = row
     input_video_path = f"video_frames/{clip_name}"
     lerped_keypoints_path = f"./lerped_keypoints/{clip_name}.json"
     output_keypoints_path = f"./new_yolo_keypoints/{clip_name}.json"
@@ -47,9 +47,9 @@ def process_clip(row, model):
     os.makedirs(os.path.dirname(lerped_keypoints_path), exist_ok=True)
     os.makedirs(os.path.dirname(output_keypoints_path), exist_ok=True)
 
-    # # return if the file already exists
+    # return if the file already exists
-    # if os.path.exists(output_keypoints_path):
+    if os.path.exists(output_keypoints_path):
-    #     return
+        return
 
     results = model(input_video_path)
     frames = get_yoloed_frames(results, lerped_keypoints_path)
@@ -59,7 +59,7 @@ def process_clip(row, model):
         json.dump(frames, f, cls=skel.Encoder, indent=4)
 
 def process_rows_on_gpu(rows, model, device):
-    for _, row in rows.iterrows():
+    for row in rows:
         for _ in range(5): 
             try:
                 process_clip(row, model)
@@ -73,7 +73,11 @@ def process_rows_on_gpu(rows, model, device):
 def gen_yolo_skeletons(descriptor):
     num_gpus = torch.cuda.device_count()
     
-    rows_per_gpu = len(descriptor) // num_gpus
+    ClipName_list = descriptor['ClipName'].tolist()
+    finished_list = os.listdir('./new_yolo_keypoints')
+    ClipName_list = [clip for clip in ClipName_list if clip + '.json' not in finished_list]
+    
+    rows_per_gpu = len(ClipName_list) // num_gpus
     
     models = [YOLO("yolo11x-pose.pt").to(torch.device(f'cuda:{i}')) for i in range(num_gpus)]
 
@@ -81,8 +85,8 @@ def gen_yolo_skeletons(descriptor):
         futures = []
         for i in range(num_gpus):
             start_idx = i * rows_per_gpu
-            end_idx = (i + 1) * rows_per_gpu if i != num_gpus - 1 else len(descriptor)
+            end_idx = (i + 1) * rows_per_gpu if i != num_gpus - 1 else len(ClipName_list)
-            gpu_rows = descriptor.iloc[start_idx:end_idx]
+            gpu_rows = ClipName_list[start_idx:end_idx]
             futures.append(executor.submit(process_rows_on_gpu, gpu_rows, models[i], torch.device(f'cuda:{i}')))
                     
         for future in concurrent.futures.as_completed(futures):
@@ -91,12 +95,6 @@ def gen_yolo_skeletons(descriptor):
             except Exception as e:
                 print(f"Error processing rows on GPU: {e}")
 
-def gen_yolo_skeletons_single(descriptor):
-    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-    model = YOLO("yolo11x-pose.pt").to(device)
-
-    process_rows_on_gpu(descriptor, model, device)
-
 def main():
     model = YOLO("yolo11x-pose.pt")  # pretrained YOLO11n model
     descriptor = pd.read_csv('./ClipDescriptorKaggle_processed.csv')