From 8fb4f8753d5b993df4e6dc6a53b95419e16652f8 Mon Sep 17 00:00:00 2001 From: zaqxs123456 <18200842@life.hkbu.edu.hk> Date: Mon, 21 Oct 2024 12:38:29 +0800 Subject: [PATCH] gened yolo dataset --- .gitignore | 3 ++- openpose_gen.py | 27 +++++++++++++++++++++++- process_json_file.py | 32 +++++++++++++++++++++++++++-- skeleton_lib.py | 49 ++++++++++++++++++++++++++++---------------- yolo_main.py | 26 +++++++++++------------ 5 files changed, 101 insertions(+), 36 deletions(-) diff --git a/.gitignore b/.gitignore index 1a7f2c6..df9e8bf 100644 --- a/.gitignore +++ b/.gitignore @@ -171,4 +171,5 @@ lerped_keypoints/* avg_keypoints/* video_frames/* video/* -new_yolo_keypoints/* \ No newline at end of file +new_yolo_keypoints/* +new_yolo_keypoints_bak/* \ No newline at end of file diff --git a/openpose_gen.py b/openpose_gen.py index 41531ad..d6348ab 100644 --- a/openpose_gen.py +++ b/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() \ No newline at end of file diff --git a/process_json_file.py b/process_json_file.py index a5440fa..d541da2 100644 --- a/process_json_file.py +++ b/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) diff --git a/skeleton_lib.py b/skeleton_lib.py index 472b4a6..7c4525c 100644 --- a/skeleton_lib.py +++ b/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): - keypoint1 = keypoints[k1_index] - keypoint2 = keypoints[k2_index] + # for (k1_index, k2_index), color in zip(limbSeq, colors): + # keypoint1 = keypoints[k1_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 keypoint1.confidence == 0: - # print(f"keypoint failed: {k1_index}") - # if keypoint2 is None or keypoint2.confidence == 0: - # print(f"keypoint failed: {k2_index}") - continue + # 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: + # # print(f"keypoint failed: {k1_index}") + # # if keypoint2 is None or keypoint2.confidence == 0: + # # print(f"keypoint failed: {k2_index}") + # continue - Y = np.array([keypoint1.x, keypoint2.x]) * float(W) - X = np.array([keypoint1.y, keypoint2.y]) * float(H) - mX = np.mean(X) - mY = np.mean(Y) - 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])) - 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]) + # Y = np.array([keypoint1.x, keypoint2.x]) * float(W) + # X = np.array([keypoint1.y, keypoint2.y]) * float(H) + # mX = np.mean(X) + # mY = np.mean(Y) + # 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])) + # 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]) 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 diff --git a/yolo_main.py b/yolo_main.py index c2ec62d..9eb5468 100644 --- a/yolo_main.py +++ b/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 - # if os.path.exists(output_keypoints_path): - # return + # return if the file already exists + if os.path.exists(output_keypoints_path): + 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) - gpu_rows = descriptor.iloc[start_idx:end_idx] + end_idx = (i + 1) * rows_per_gpu if i != num_gpus - 1 else len(ClipName_list) + 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')