Untitled

from mmpose.apis import init_pose_model, inference_top_down_pose_model, vis_pose_result
import os
import cv2
import matplotlib.pyplot as plt

def inference_and_render_skeleton(image_path, output_path):

    # Paths
    config_file = 'configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/ochuman/ViTPose_base_ochuman_256x192.py'
    checkpoint_file = '/content/drive/MyDrive/ViTPose/models/vitpose_base_coco_aic_mpii.pth'
    image_folder = image_path
    output_folder = output_path
    os.makedirs(output_folder, exist_ok=True)

    # Initialize the model
    pose_model = init_pose_model(config_file, checkpoint_file, device='cpu') # cuda:0

    line_thickness = 3
    keypoint_radius = 4

    output_data = []

    for image_name in os.listdir(image_folder):

        print(image_name)

        image_path = os.path.join(image_folder, image_name)
        if os.path.isdir(image_path):
            continue

        # Read image
        image = cv2.imread(image_path)
        height, width, _ = image.shape

        print(width, height)

        # Dummy bounding box for the whole image (if no person detection)
        dummy_bboxes = [{'bbox': [0, 0, image.shape[1], image.shape[0], 1]}]

        # Perform inference
        pose_results, _ = inference_top_down_pose_model(
            pose_model,
            image_path,
            person_results=dummy_bboxes,
            format='xywh',
            bbox_thr=0.5,
            dataset='TopDownOCHumanDataset',  # Set for OCHuman
            return_heatmap=False
        )

        vis_result = vis_pose_result(
            pose_model,
            image_path,
            pose_results,
            dataset='TopDownOCHumanDataset',
            show=False,
            kpt_score_thr=0.5,  # Confidence threshold for keypoints
            radius=keypoint_radius,
            thickness=line_thickness
        )

        vis_result_rgb = cv2.cvtColor(vis_result, cv2.COLOR_BGR2RGB)

        # Show the image in the notebook
        plt.imshow(vis_result_rgb)
        plt.axis('off')
        plt.show()

        cv2.imwrite(os.path.join(output_folder, image_name), vis_result)

        print(f"Processed {image_name}")



        for result in pose_results:  # Loop over all detected results

            result['keypoints'] = result['keypoints'].astype(np.float64)
            keypoints = []
            bbox = [float('inf'), float('inf'), -float('inf'), -float('inf')]

            for kp in result['keypoints']:  # Keypoints is a list of (x, y, visibility)
                x, y, visibility = kp
                keypoints.extend([x, y, visibility])
                if visibility >= 0.5:
                      bbox[0] = min(bbox[0], x)
                      bbox[1] = min(bbox[1], y)
                      bbox[2] = max(bbox[2], x)
                      bbox[3] = max(bbox[3], y)

            center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
            scale = [bbox[2] - bbox[0], bbox[3] - bbox[1]]


            # center, scale, score doesn't affect mAP calculation. Use dummy values
            output_entry = {
                "category_id": 1,
                "center": [0, 0],
                "image_id": name2id.get(image_name, 1),
                "keypoints": keypoints,
                "scale": [0, 0],
                "score": 1.0
            }

            output_data.append(output_entry)




    with open(os.path.join(output_folder, 'annotations.json'), 'w') as f:
        json.dump(output_data, f, indent=4)

        print(f"Annotations saved")