Spaces:
Sleeping
Sleeping
File size: 4,400 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """Wizara Vision API endpoint handlers."""
from __future__ import annotations
import os
from typing import Any, Callable
from PIL import Image
from .responses import (
detections_to_objects,
error_response,
parse_advanced_settings,
success_response,
unsupported_response,
)
def resolve_image_path(image_file: Any, base_dir: str) -> str | None:
if isinstance(image_file, str):
candidate = image_file
if not os.path.isabs(candidate):
candidate = os.path.join(base_dir, candidate)
return candidate if os.path.exists(candidate) else None
if isinstance(image_file, dict):
path = image_file.get("path")
return path if path and os.path.exists(path) else None
path = getattr(image_file, "path", None)
return path if path and os.path.exists(path) else None
def merge_settings(
*,
model_mode: str,
temp: float,
top_p: float,
top_k: int,
short_size: int | None,
advanced_settings: str | None,
) -> dict[str, Any]:
settings = {
"model_mode": model_mode,
"temp": temp,
"top_p": top_p,
"top_k": top_k,
"short_size": short_size,
}
settings.update(parse_advanced_settings(advanced_settings))
return settings
def handle_detect(
*,
image_file: Any,
categories: str,
task_type: str,
model_mode: str,
temp: float,
top_p: float,
top_k: int,
short_size: int | None,
advanced_settings: str | None,
base_dir: str,
run_image_gpu: Callable[..., tuple],
generate_prompt: Callable[[str, str], str],
parse_results: Callable[[str, str], list[dict[str, Any]]],
) -> dict[str, Any]:
if not image_file:
return error_response("Image is required.", code="MISSING_IMAGE")
image_path = resolve_image_path(image_file, base_dir)
if not image_path:
return error_response("Invalid image file path.", code="INVALID_IMAGE")
settings = merge_settings(
model_mode=model_mode,
temp=temp,
top_p=top_p,
top_k=top_k,
short_size=short_size,
advanced_settings=advanced_settings,
)
category = categories.strip() or "objects"
task = task_type.strip() or "Detection"
question_override = settings.get("question_override")
if not question_override:
question_override = generate_prompt(task, category)
try:
image = Image.open(image_path).convert("RGB")
width, height = image.size
_, stats, raw_text, _, _ = run_image_gpu(
image_path,
category,
settings.get("model_mode", model_mode),
float(settings.get("temp", temp)),
float(settings.get("top_p", top_p)),
int(settings.get("top_k", top_k)),
settings.get("short_size", short_size),
question_override,
)
category_str = " ".join(part.strip() for part in category.split(",") if part.strip())
detections = parse_results(raw_text, category_str)
objects = detections_to_objects(detections)
return success_response(
image_width=width,
image_height=height,
objects=objects,
task=task.lower(),
extra={
"stats": stats,
"raw_text": raw_text,
"prompt": question_override,
},
)
except Exception as exc: # noqa: BLE001
return error_response(str(exc), code="INFERENCE_FAILED")
def handle_ocr(
*,
image_file: Any,
model_mode: str,
temp: float,
top_p: float,
top_k: int,
short_size: int | None,
advanced_settings: str | None,
base_dir: str,
run_image_gpu: Callable[..., tuple],
generate_prompt: Callable[[str, str], str],
parse_results: Callable[[str, str], list[dict[str, Any]]],
) -> dict[str, Any]:
return handle_detect(
image_file=image_file,
categories="text",
task_type="OCR",
model_mode=model_mode,
temp=temp,
top_p=top_p,
top_k=top_k,
short_size=short_size,
advanced_settings=advanced_settings,
base_dir=base_dir,
run_image_gpu=run_image_gpu,
generate_prompt=generate_prompt,
parse_results=parse_results,
)
def handle_unsupported(task: str) -> dict[str, Any]:
return unsupported_response(task)
|