File size: 2,503 Bytes
3820d5b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
"""Register Wizara Vision API endpoints on the Gradio Server app."""

from __future__ import annotations

import os
from typing import Any

from .endpoints import handle_detect, handle_ocr, handle_unsupported


def register_wizara_api(app, deps: dict[str, Any]) -> None:
    base_dir = os.path.dirname(os.path.abspath(deps["__file__"]))
    run_image_gpu = deps["run_image_gpu_api"]
    generate_prompt = deps["generate_raw_prompt"]
    parse_results = deps["parse_mixed_results"]

    shared = {
        "base_dir": base_dir,
        "run_image_gpu": run_image_gpu,
        "generate_prompt": generate_prompt,
        "parse_results": parse_results,
    }

    @app.api(name="detect")
    def detect_api(
        image: Any = None,
        categories: str = "objects",
        task_type: str = "Detection",
        model_mode: str = "hybrid",
        temp: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 20,
        short_size: int | None = None,
        advanced_settings: str = "",
    ) -> dict:
        """Detect objects in an uploaded image and return normalized JSON."""
        return handle_detect(
            image_file=image,
            categories=categories,
            task_type=task_type,
            model_mode=model_mode,
            temp=temp,
            top_p=top_p,
            top_k=top_k,
            short_size=short_size,
            advanced_settings=advanced_settings,
            **shared,
        )

    @app.api(name="ocr")
    def ocr_api(
        image: Any = None,
        model_mode: str = "hybrid",
        temp: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 20,
        short_size: int | None = None,
        advanced_settings: str = "",
    ) -> dict:
        """Run OCR localization using the existing LocateAnything OCR task."""
        return handle_ocr(
            image_file=image,
            model_mode=model_mode,
            temp=temp,
            top_p=top_p,
            top_k=top_k,
            short_size=short_size,
            advanced_settings=advanced_settings,
            **shared,
        )

    @app.api(name="caption")
    def caption_api() -> dict:
        return handle_unsupported("caption")

    @app.api(name="segment")
    def segment_api() -> dict:
        return handle_unsupported("segment")

    @app.api(name="count")
    def count_api() -> dict:
        return handle_unsupported("count")

    @app.api(name="classify")
    def classify_api() -> dict:
        return handle_unsupported("classify")