pop_demo.py

import argparse
import ast
import asyncio
import base64
import json
import logging
import os
import sys
from argparse import Namespace
from io import BytesIO
from typing import Any

from dotenv import load_dotenv
from PIL import Image
from pybars import Compiler
from pydantic import TypeAdapter
from webui import webui

from eyepop import EyePopSdk, Job
from eyepop.data.data_types import TranscodeMode
from eyepop.data.types.asset import Area, RectangleArea
from eyepop.worker.worker_types import (
    ComponentParams,
    ContourFinderComponent,
    ContourType,
    CropForward,
    ForwardComponent,
    FullForward,
    InferenceComponent,
    MotionDetectConfig,
    MotionModel,
    Pop,
    TrackingComponent,
)

load_dotenv()

logging.basicConfig(level=logging.INFO)
log = logging.getLogger('eyepop.example')

script_dir = os.path.dirname(__file__)


pop_examples = {
    "vehicles": Pop(components=[
        InferenceComponent(
            ability='eyepop.vehicle:latest',
            categoryName="vehicles",
            confidenceThreshold=0.8,
            forward=CropForward(
                targets=[TrackingComponent(
                    maxAgeSeconds=5.0,
                    motionModel=MotionModel.CONSTANT_VELOCITY,
                    agnostic=True,
                    classHysteresis=True,
                ), InferenceComponent(
                    ability='eyepop.vehicle.licence-plate:latest',
                    topK=1,
                    forward=CropForward(
                        targets=[InferenceComponent(
                            ability='eyepop.text.recognize.landscape:latest',
                            categoryName="licence-plate"
                        )]
                    )
                )]
            )
        )
    ]),

    "person": Pop(components=[
        InferenceComponent(
            ability='eyepop.person:latest',
            categoryName="person"
        )
    ]),

    "2d-body-points": Pop(components=[
        InferenceComponent(
            ability='eyepop.person:latest',
            categoryName="person",
            forward=CropForward(
                maxItems=128,
                targets=[InferenceComponent(
                    ability='eyepop.person.2d-body-points:latest',
                    categoryName="2d-body-points",
                    confidenceThreshold=0.25
                )]
        ))
    ]),

    "faces": Pop(components=[
        InferenceComponent(
            ability='eyepop.person:latest',
            categoryName="person",
            forward=CropForward(
                maxItems=128,
                targets=[InferenceComponent(
                    ability='eyepop.person.face.short-range:latest',
                    categoryName="2d-face-points",
                    forward=CropForward(
                        boxPadding=1.5,
                        orientationTargetAngle=-90.0,
                        targets=[InferenceComponent(
                            ability='eyepop.person.face-mesh:latest',
                            categoryName="3d-face-mesh"
                        )]
                    )
                )]
            )
        )
    ]),

    "hands": Pop(components=[
        InferenceComponent(
            ability='eyepop.person:latest',
            categoryName="person",
            forward=CropForward(
                maxItems=128,
                boxPadding=0.25,
                targets=[InferenceComponent(
                    ability='eyepop.person.palm:latest',
                    forward=CropForward(
                        includeClasses=["hand circumference"],
                        orientationTargetAngle=-90.0,
                        targets=[InferenceComponent(
                            ability='eyepop.person.3d-hand-points:latest',
                            categoryName="3d-hand-points"
                        )]
                    )
                )]
            )
        )
    ]),

    "3d-body-points": Pop(components=[
        InferenceComponent(
            ability='eyepop.person:latest',
            categoryName="person",
            forward=CropForward(
                boxPadding=0.5,
                targets=[InferenceComponent(
                    ability='eyepop.person.pose:latest',
                    hidden=True,
                    forward=CropForward(
                        boxPadding=0.5,
                        orientationTargetAngle=-90.0,
                        targets=[InferenceComponent(
                            ability='eyepop.person.3d-body-points.heavy:latest',
                            categoryName="3d-body-points",
                            confidenceThreshold=0.25
                        )]
                    )
                )]
            )
        )
    ]),

    "text": Pop(components=[
        InferenceComponent(
            ability='eyepop.text:latest',
            categoryName="text",
            confidenceThreshold=0.7,
            forward=CropForward(
                maxItems=128,
                targets=[InferenceComponent(
                    ability='eyepop.text.recognize.landscape:latest',
                    confidenceThreshold=0.1
                )]
            )
        )
    ]),

    "sam1": Pop(components=[
        InferenceComponent(
            ability='eyepop.sam.small:latest',
            id=1,
            forward=FullForward(
                targets=[ContourFinderComponent(
                    contourType=ContourType.POLYGON,
                    areaThreshold=0.005
                )]
            )
        )
    ]),

    "sam2": Pop(components=[
        InferenceComponent(
            ability="eyepop.sam2.encoder.tiny:latest",
            id=1,
            hidden=True,
            forward=FullForward(
                targets=[InferenceComponent(
                    ability='eyepop.sam2.decoder:latest',
                    forward=FullForward(
                        targets=[ContourFinderComponent(
                            contourType=ContourType.POLYGON,
                            areaThreshold=0.005
                        )]
                    )
                )]
            )
        )
    ]),
    "image-contents": Pop(components=[
        InferenceComponent(
            id=1,
            ability='eyepop.image-contents:latest',
        )
    ]),
    "localize-objects": Pop(components=[
        InferenceComponent(
            id=1,
            ability='eyepop.localize-objects:latest',
        )
    ]),
    "localize-objects-plus": Pop(components=[
        InferenceComponent(
            id=1,
            ability='eyepop.localize-objects:latest',
            params={
                "prompts": [{"prompt": "person"}]
            },
            forward=CropForward(
                targets=[InferenceComponent(
                    ability='eyepop.image-contents:latest',
                    params={
                        "prompts": [{"prompt": "hair color blond"},{"prompt": "hair color brown"}]
                    }
                )],
            )
        )
    ]),
    "localize-objects-t4": Pop(components=[
        InferenceComponent(
            id=1,
            ability='eyepop.localize-objects:latest',
            params={
                "prompts": [{"prompt": "person"}]
            },
            forward=CropForward(
                targets=[InferenceComponent(
                    ability='eyepop.image-contents-t4:latest',
                    params={
                        "prompts": [{"prompt": "shirt color?"}]
                    }
                )],
            )
        )
    ]),
}

def list_of_points(arg: str) -> list[dict[str, Any]]:
    points = []
    points_as_tuples = ast.literal_eval(f'[{arg}]')
    for tuple in points_as_tuples:
        points.append({
            "x": tuple[0],
            "y": tuple[1]
        })
    return points


def list_of_boxes(arg: str) -> list[dict[str, Any]]:
    boxes = []
    boxes_as_tuples = ast.literal_eval(f'[{arg}]')
    for tuple in boxes_as_tuples:
        boxes.append({
            "topLeft": {
                "x": tuple[0],
                "y": tuple[1]
            },
            "bottomRight": {
                "x": tuple[2],
                "y": tuple[3]
            }
        })
    return boxes


def rectangle_roi(arg: str) -> Area:
    roi = ast.literal_eval(arg)

    return RectangleArea(
        x=roi[0],
        y=roi[1],
        width=roi[2],
        height=roi[3],
    )


def add_optional_tracking_to_component(component: ForwardComponent, tracking_args: Namespace):
    if tracking_args.tracking:
        tracking_component = TrackingComponent()
        if tracking_args.tracking_reid_model is not None:
            tracking_component.reidModel = tracking_args.tracking_reid_model
        if tracking_args.tracking_max_age is not None:
            tracking_component.maxAgeSeconds = tracking_args.tracking_max_age
        if tracking_args.tracking_iou_threshold is not None:
            tracking_component.iouThreshold = tracking_args.tracking_iou_threshold
        if tracking_args.tracking_sim_threshold is not None:
            tracking_component.simThreshold = tracking_args.tracking_sim_threshold
        if tracking_args.tracking_agnostic is not None:
            tracking_component.agnostic = tracking_args.tracking_agnostic
        if tracking_args.tracking_motion_model is not None:
            tracking_component.motionModel = MotionModel(tracking_args.tracking_motion_model)
        component.forward = CropForward(
            targets=[tracking_component]
        )


parser = argparse.ArgumentParser(
                    prog='Pop examples',
                    description='Demonstrates the composition of a Pop',
                    epilog='.')
parser.add_argument('-d', '--dump', required=False, help="dump composable pop definition to stdout", default=False, action="store_true")
parser.add_argument('-l', '--local-path', required=False, type=str, default=False, help="run the inference on a local file, or all files on a directory")
parser.add_argument('-a', '--asset-uuid', required=False, type=str, default=False, help="run the inference on an asset by its Uuid")
parser.add_argument('-u', '--url', required=False, type=str, default=False, help="run the inference on a remote Url")
parser.add_argument('-p', '--pop', required=False, type=str, help="run this pop", choices=list(pop_examples.keys()))
parser.add_argument('-s', '--session', required=False, type=str, help="Use an existing a session uuid and do not set a pop, use the sessions as preconfigured", default=None)
parser.add_argument('-m', '--model-uuid', required=False, type=str, action="append", help="run this model(s) by uuid")
parser.add_argument('-ma', '--model-alias', required=False, type=str, action="append", help="run this model(s) by its tagged alias")
parser.add_argument('-ms1', '--model-uuid-sam1', required=False, type=str, help="run this model by uuid and compose with SAM1 (EfficientSAM) and Contour Finder")
parser.add_argument('-ms2', '--model-uuid-sam2', required=False, type=str, help="run this model by uuid and compose with SAM2 and Contour Finder")
parser.add_argument('-po', '--points', required=False, type=list_of_points, help="List of POIs as coordinates like (x1, y1), (x2, y2) in the original image coordinate system")
parser.add_argument('-bo', '--boxes', required=False, type=list_of_boxes, help="List of POIs as boxes like (left1, top1, right1, bottom1), (left1, top1, right1, bottom1) in the original image coordinate system")
parser.add_argument('-sp', '--single-prompt', required=False, type=str, help="Single prompt to pass as parameter")
parser.add_argument('-sl', '--single-label', required=False, type=str, help="Single label to use as result for single prompt to pass as parameter")
parser.add_argument('-pr', '--prompt', required=False, type=str, help="Prompt to pass as parameter", action="append")
parser.add_argument('-v', '--visualize', required=False, help="show rendered output", default=False, action="store_true")
parser.add_argument('-o', '--output', required=False, help="print results to stdout", default=False, action="store_true")
parser.add_argument('-ds', '--dataset-uuid', required=False, type=str, help="Ingest all assets into a dataset uuid", default=None)
parser.add_argument('-tk', '--top-k', required=False, type=int, help="For --model-uuid and -model-alias apply this top-k filter", default=None)
parser.add_argument('-ct', '--confidence-threshold', required=False, type=float, help="For --model-uuid and -model-alias apply this confidence threshold filter", default=None)
parser.add_argument('--fps', required=False, type=str, help="For --model-uuid and -model-alias apply this target fps - e.g. 15/1", default=None)
parser.add_argument('-vl', '--video-chunk-length-seconds', required=False, type=float, help="For --model-uuid and -model-alias apply this video chunk length", default=None)
parser.add_argument('-vo', '--video-chunk-overlap', required=False, type=float, help="For --model-uuid and -model-alias apply this video chunk overlap", default=None)

# Optional tracking for simple pops my model uuid oder model alias
parser.add_argument('--tracking', required=False, help="Track objects in videos", default=False, action="store_true")
parser.add_argument('--tracking-reid-model', required=False, help="Use re-id model uuid for tracking", default=None, type=str)
parser.add_argument('--tracking-agnostic', required=False, help="Track objects class-agnostic", default=False, action="store_true")
parser.add_argument('--tracking-max-age', required=False, help="Max age in seconds for unmatched tracks", default=None, type=float)
parser.add_argument('--tracking-iou-threshold', required=False, help="IoU threshold to match tracks", default=None, type=float)
parser.add_argument('--tracking-sim-threshold', required=False, help="Similarity threshold to match tracks by re-id", default=None, type=float)
parser.add_argument('--tracking-motion-model', required=False, help="Pick a motion model for tracking, one of 'random_walk', 'constant_velocity' or 'constant_acceleration'", default=None, type=str)

# Optional motion detection parameters
parser.add_argument('--motion-detect', required=False, help="Skip video frames w/o detected motion", default=False, action="store_true")

# Optional global ROI parameters
parser.add_argument('--roi', required=False, type=rectangle_roi, help="Rectangular ROI as (x, y, width, height)")


main_args = parser.parse_args()

if not main_args.local_path and not main_args.url and not main_args.asset_uuid:
    print("Need something to run inference on; pass either --url or --local-path or --asset-uuid")
    parser.print_help()
    sys.exit(1)

if (not main_args.pop
        and not main_args.model_uuid
        and not main_args.model_alias
        and not main_args.model_uuid_sam1
        and not main_args.model_uuid_sam2
        and not main_args.session):
    print("Need something do do, pass either --pop or --model-uuid or --model-alias or --model-uuid-sam1 "
          "or --model-uuid-sam2 (or start with a preconfigured session with --session)")
    parser.print_help()
    sys.exit(1)

if main_args.pop:
    pop = pop_examples[main_args.pop]
elif main_args.model_uuid:
    pop = Pop(components=[
        InferenceComponent(
            id=i+1,
            abilityUuid=uuid,
        ) for i, uuid in enumerate(main_args.model_uuid)
    ])
    if main_args.top_k is not None:
        for c in pop.components:
            c.topK = main_args.top_k

    if main_args.confidence_threshold is not None:
        for c in pop.components:
            c.confidenceThreshold = main_args.confidence_threshold

    if main_args.video_chunk_length_seconds is not None:
        for c in pop.components:
            c.videoChunkLengthSeconds = main_args.video_chunk_length_seconds

    if main_args.video_chunk_overlap is not None:
        for c in pop.components:
            c.videoChunkOverlap = main_args.video_chunk_overlap

    add_optional_tracking_to_component(pop.components[0], main_args)
elif main_args.model_alias:
    pop = Pop(components=[
        InferenceComponent(
            id=i+1,
            ability=alias,
        ) for i, alias in enumerate(main_args.model_alias)
    ])

    if main_args.top_k is not None:
        for c in pop.components:
            c.topK = main_args.top_k

    if main_args.confidence_threshold is not None:
        for c in pop.components:
            c.confidenceThreshold = main_args.confidence_threshold

    if main_args.video_chunk_length_seconds is not None:
        for c in pop.components:
            c.videoChunkLengthSeconds = main_args.video_chunk_length_seconds

    if main_args.video_chunk_overlap is not None:
        for c in pop.components:
            c.videoChunkOverlap = main_args.video_chunk_overlap

    add_optional_tracking_to_component(pop.components[0], main_args)
elif main_args.model_uuid_sam1:
    pop = Pop(components=[
        InferenceComponent(
            abilityUuid=main_args.model_uuid_sam1,
            forward=CropForward(
                targets=[InferenceComponent(
                    ability='eyepop.sam.small:latest',
                    forward=FullForward(
                        targets=[ContourFinderComponent(
                            contourType=ContourType.POLYGON,
                            areaThreshold=0.005
                        )]
                    )
                )]
            )
        )
    ])
elif main_args.model_uuid_sam2:
    pop = Pop(components=[
        InferenceComponent(
            ability="eyepop.sam2.encoder.tiny:latest",
            hidden=True,
            forward=FullForward(
                targets=[InferenceComponent(
                    abilityUuid=main_args.model_uuid_sam2,
                    forward=CropForward(
                        targets=[InferenceComponent(
                            ability='eyepop.sam2.decoder:latest',
                            forward=FullForward(
                                targets=[ContourFinderComponent(
                                    contourType=ContourType.POLYGON,
                                    areaThreshold=0.005
                                )]
                            )
                        )]
                    )
                )]
            )
        )
    ])
elif main_args.session:
    pop = None
else:
    raise ValueError("pop or model required (or a preconfigured session)")

params = None
if main_args.points:
    params = [
        ComponentParams(componentId=1, values={
          "roi": {
              "points": main_args.points
          }
        })
    ]
elif main_args.boxes:
    params = [
        ComponentParams(componentId=1, values={
            "roi": {
                "boxes": main_args.boxes
            }
        })
    ]
elif main_args.prompt is not None and len(main_args.prompt) > 0:
    params = [
        ComponentParams(componentId=1, values={
            "prompts": [{"prompt": p} for p in main_args.prompt]
        })
    ]
elif main_args.single_prompt is not None:
    params = [
        ComponentParams(componentId=1, values={
            "prompt": main_args.single_prompt,
            "label": main_args.single_label if main_args.single_label else main_args.single_prompt
        })
    ]

async def main(args) -> tuple[dict[str, Any] | None, str | None]:
    visualize_prediction = None
    visualize_path = None
    example_image_src = None
    motion_detect = MotionDetectConfig(motionGap=1) if args.motion_detect else None

    if args.dump and pop:
        print("Pop:", pop.model_dump_json(indent=2))
        if params:
            print("Params:", TypeAdapter(list[ComponentParams]).dump_json(params, indent=2).decode("utf-8"))

    async with EyePopSdk.async_worker(dataset_uuid=args.dataset_uuid, session_uuid=args.session) as endpoint:
        if pop:
            await endpoint.set_pop(pop)

        if args.local_path:
            if not os.path.exists(args.local_path):
                log.warning(f"local path {args.local_path} does not exist")
                sys.exit(1)
            if os.path.isfile(args.local_path):
                local_files = [args.local_path]
            else:
                local_files = []

                for f in os.listdir(args.local_path):
                    local_file = os.path.join(args.local_path, f)
                    if os.path.isfile(local_file):
                        local_files.append(local_file)
            jobs = []
            async def on_ready(job: Job, path: str):
                nonlocal visualize_prediction
                nonlocal visualize_path
                while result := await job.predict():
                    visualize_prediction = result
                    visualize_path = path
                    if args.output:
                        print(path, json.dumps(result, indent=2))
            for local_file in local_files:
                job = await endpoint.upload(local_file, params=params, motion_detect=motion_detect, roi=args.roi, fps=args.fps)
                jobs.append(on_ready(job, local_file))
            await asyncio.gather(*jobs)
            if args.visualize and visualize_prediction is not None:
                image = Image.open(visualize_path)
                buffer = BytesIO()
                image.save(buffer, format="PNG")
                example_image_src = f"data:image/png;base64, {base64.b64encode(buffer.getvalue()).decode()}"
        elif args.url:
            job = await endpoint.load_from(args.url, params=params, motion_detect=motion_detect, roi=args.roi, fps=args.fps)
            while result := await job.predict():
                visualize_prediction = result
                if args.output:
                    print(json.dumps(result, indent=2))
            if args.visualize:
                example_image_src = args.url
        elif args.asset_uuid:
            job = await endpoint.load_asset(args.asset_uuid, params=params, motion_detect=motion_detect, roi=args.roi, fps=args.fps)
            while result := await job.predict():
                visualize_prediction = result
                if args.output:
                    print(json.dumps(result, indent=2))
            if args.visualize:
                async with EyePopSdk.dataEndpoint(is_async=True) as dataEndpoint:
                    buffer = await dataEndpoint.download_asset(
                        args.asset_uuid,
                        transcode_mode=TranscodeMode.image_original_size
                    ).read()
                    example_image_src = f"data:image/jpeg;base64, {base64.b64encode(buffer).decode()}"
    return visualize_prediction, example_image_src

visualize_result, example_image_src = asyncio.run(main(main_args))
if main_args.visualize:
    with open(os.path.join(script_dir, 'viewer.html')) as file:
        compiler = Compiler()
        html_template = compiler.compile(file.read())

    preview = html_template({
        'image_src': example_image_src,
        'result_json': json.dumps(visualize_result)
    })
    window = webui.window()
    window.set_root_folder('.')
    window.show(preview)
    webui.wait()