imageprocessing_denseposev2.py

# -*- coding: utf-8 -*-
"""imageprocessing_denseposeV2.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1Ru7Bd7HnAcabBx1PE6Rl-w0mRFU-OahT
"""

# If you're not on a GPU runtime: Runtime > Change runtime type > GPU (T4/A100 works). TPU won't work for Detectron2.
!nvidia-smi || true

# Install Detectron2 (latest from main) and DensePose project as a package
!pip install -q -U "git+https://github.com/facebookresearch/detectron2.git"
!pip install -q -U "git+https://github.com/facebookresearch/detectron2@main#subdirectory=projects/DensePose"

# Basics
!pip install -q opencv-python pillow numpy pyyaml

import os, urllib.request

os.makedirs("densepose_cfg", exist_ok=True)

# Grab canonical config files from Detectron2's DensePose project (raw GitHub)
cfg_base_url = "https://raw.githubusercontent.com/facebookresearch/detectron2/main/projects/DensePose/configs"
cfg_main = "densepose_rcnn_R_50_FPN_s1x.yaml"
cfg_base = "Base-DensePose-RCNN-FPN.yaml" # Add base config

def fetch(url, out):
    if not os.path.exists(out):
        print("Downloading", out)
        urllib.request.urlretrieve(url, out)

# Main model config
fetch(f"{cfg_base_url}/{cfg_main}", f"densepose_cfg/{cfg_main}")
# Base model config
fetch(f"{cfg_base_url}/{cfg_base}", f"densepose_cfg/{cfg_base}")


# Model weights (R50-FPN s1x model used in DensePose docs)
# If this ever 404s, open the URL in a browser and copy the redirected link you get.
weights_url = "https://dl.fbaipublicfiles.com/densepose/densepose_rcnn_R_50_FPN_s1x/165712039/model_final_162be9.pkl"
weights_path = "model_final_162be9.pkl"
fetch(weights_url, weights_path)

print("Config & weights ready:", os.path.exists(f"densepose_cfg/{cfg_main}"), os.path.exists(weights_path))

import torch
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from densepose import add_densepose_config  # provided by the DensePose package

cfg_path = "/content/densepose_cfg/densepose_rcnn_R_50_FPN_s1x.yaml"

cfg = get_cfg()
add_densepose_config(cfg)
cfg.merge_from_file(cfg_path)
cfg.MODEL.WEIGHTS = "model_final_162be9.pkl"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # tweak per your needs
cfg.INPUT.MIN_SIZE_TEST = 800

# GPU if available, otherwise CPU (slower)
predictor = DefaultPredictor(cfg)

print("Device:", "CUDA" if torch.cuda.is_available() else "CPU")

# Option A: Upload from your machine
from google.colab import files
uploaded = files.upload()  # pick one image
IMAGE_PATH = list(uploaded.keys())[0]

# --- Option B (instead): URL fetch ---
# import urllib.request
# IMAGE_PATH = "test.jpg"
# urllib.request.urlretrieve("https://path/to/your/person.jpg", IMAGE_PATH)

print("Using:", IMAGE_PATH)

import cv2, json, base64, io
import numpy as np
from PIL import Image
import torch
import os

# DensePose result extractor (handles Instances -> DensePoseResult + boxes)
from densepose.vis.extractor import DensePoseResultExtractor

# Coarse grouping of DensePose 24 part indices to limb labels (I-channel).
# Mapping is based on DensePose part index conventions discussed by the authors/users:
# 0=bg; 1,2=torso; 3/4=hands; 5/6=feet; 7-14=upper/lower legs (R/L); 15-22=upper/lower arms (R/L); 23=head; 24=neck.
# (DensePose site/papers explain the IUV format; original repo is archived; see docs.)
# We aggregate these into head/torso/left_arm/right_arm/left_leg/right_leg.
COARSE_GROUPS = {
    "head":   {23, 24},
    "torso":  {1, 2},
    "left_arm":  {16, 18, 20, 22, 4},   # upper L arm, lower L arm, left hand
    "right_arm": {15, 17, 19, 21, 3},   # upper R arm, lower R arm, right hand
    "left_leg":  {8, 10, 12, 14, 6},    # upper L leg, lower L leg, left foot
    "right_leg": {7,  9, 11, 13, 5},    # upper R leg, lower R leg, right foot
}


# Read image
im_bgr = cv2.imread(IMAGE_PATH)
if im_bgr is None:
    raise RuntimeError(f"Failed to read image: {IMAGE_PATH}")
H, W = im_bgr.shape[:2]

# Inference
with torch.no_grad():
    outputs = predictor(im_bgr)

instances = outputs["instances"].to("cpu")

# Check if any instances (people) were detected and if DensePose results exist
if len(instances) == 0 or not instances.has("pred_densepose"):
    print("No people or DensePose results detected. Saving original image.")
    # Save the original image
    output_image_path = "output_image.jpg"
    Image.fromarray(cv2.cvtColor(im_bgr, cv2.COLOR_BGR2RGB)).save(output_image_path, "JPEG")
    print(f"Original image saved as {output_image_path}")

    # Create an empty JSON payload
    payload = {
        "source_image": IMAGE_PATH,
        "people": [],
        "densepose_model": {
            "config": os.path.basename(cfg_path),
            "weights": os.path.basename(weights_path),
            "framework": "detectron2 + densepose",
        },
        "message": "No people or DensePose results detected. Original image saved."
    }
    out_json = "densepose_singleframe.json"
    with open(out_json, "w") as f:
        json.dump(payload, f, indent=2)
    print(f"Wrote empty {out_json}")

else:
    scores = instances.scores.tolist() if instances.has("scores") else [1.0] * len(instances)

    # Extract DensePose results + boxes (xywh)
    extractor = DensePoseResultExtractor()
    dp_results, boxes_xywh = extractor(instances)


    os.makedirs("segmented_parts", exist_ok=True)

    people_json = []
    full_I_mask = np.zeros((H, W), dtype=np.uint8)  # optional composite mask

    for i, (dp_res, box_xywh, score) in enumerate(zip(dp_results, boxes_xywh, scores)):
        # dp_res.labels: [H_roi, W_roi], values in [0..24]; dp_res.uv: (2, H_roi, W_roi) but not needed for coarse parts.
        I_roi = dp_res.labels.cpu().numpy().astype(np.uint8)

        # Paste ROI labels into image coords using bbox (nearest-neighbor)
        x, y, w, h = [int(v) for v in box_xywh.tolist()]
        w = max(w, 1); h = max(h, 1)
        I_resized = cv2.resize(I_roi, (w, h), interpolation=cv2.INTER_NEAREST)

        # Compose per-person mask (same size as image; zero elsewhere)
        I_person = np.zeros((H, W), dtype=np.uint8)
        y2 = min(y+h, H); x2 = min(x+w, W)
        I_person[y:y2, x:x2] = I_resized[:(y2-y), :(x2-x)]

        # Update composite
        full_I_mask[y:y2, x:x2] = np.where(I_resized[:(y2-y), :(x2-x)]>0,
                                        I_resized[:(y2-y), :(x2-x)],
                                        full_I_mask[y:y2, x:x2])

        # Count pixels per coarse group and save masks
        parts_present = {}
        saved_parts = {}

        for name, ids in COARSE_GROUPS.items():
            # Extract only if present
            mask_bin = np.isin(I_person, list(ids)).astype(np.uint8)

            if mask_bin.sum() > 0:
                parts_present[name] = int(mask_bin.sum())

                # Apply mask on original image
                masked = im_bgr.copy()
                # Create a 3-channel mask from the single-channel binary mask
                mask_bin_3_channel = np.stack([mask_bin, mask_bin, mask_bin], axis=-1)
                masked = masked * mask_bin_3_channel # Apply the mask

                # Crop tight region of detected part
                ys, xs = np.where(mask_bin > 0)
                ymin, ymax = ys.min(), ys.max()
                xmin, xmax = xs.min(), xs.max()
                crop = masked[ymin:ymax, xmin:xmax]

                # Save as JPG
                out_path = f"segmented_parts/person{i}_{name}.jpg"
                Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)).save(out_path, "JPEG")
                saved_parts[name] = out_path


        person_entry = {
            "bbox_xywh": [x, y, w, h],
            "score": float(score),
            "parts_detected": parts_present,
            "part_files": saved_parts,  # file paths for cropped regions
        }
        people_json.append(person_entry)

    # Pack final JSON
    payload = {
        "source_image": IMAGE_PATH,
        "people": people_json,
        "densepose_model": {
            "config": os.path.basename(cfg_path),
            "weights": os.path.basename(weights_path),
            "framework": "detectron2 + densepose",
        },
    }

    out_json = "densepose_singleframe.json"
    with open(out_json, "w") as f:
        json.dump(payload, f, indent=2) # Use indent=2 for pretty printing the JSON

    print(f"Wrote {out_json} with {len(people_json)} person(s). Segmented parts in segmented_parts/")


    import matplotlib.pyplot as plt

    # Simple 25-color palette (0..24) for visualization
    palette = np.random.RandomState(123).randint(0, 255, (25, 3), dtype=np.uint8)
    palette[0] = 0  # background black

    # Build overlay for the composite mask
    color_mask = palette[full_I_mask]
    overlay = (0.6 * cv2.cvtColor(im_bgr, cv2.COLOR_BGR2RGB) + 0.4 * color_mask).astype(np.uint8)

    # Draw bboxes
    for p in payload["people"]:
        x, y, w, h = p["bbox_xywh"]
        cv2.rectangle(overlay, (x, y), (x+w, y+h), (255, 255, 255), 2)

    plt.figure(figsize=(10, 10))
    plt.title("DensePose parts overlay (with person boxes)")
    plt.axis("off")
    plt.imshow(overlay)
    plt.show()

from google.colab import files
files.download("densepose_singleframe.json")