infer.py
· 1008 B · Python
Raw
from pathlib import Path
import numpy as np
import openvino as ov
import openvino_genai as ov_genai
from PIL import Image
MODEL_DIR = Path("models/medgemma-1.5-4b-it-int4-ov")
IMAGE_PATH = Path("images/test.png")
DEVICE = "CPU" # 有 Intel 核显/独显也可以试 "GPU"
def main() -> None:
if not MODEL_DIR.exists():
raise FileNotFoundError(f"Model directory not found: {MODEL_DIR.resolve()}")
if not IMAGE_PATH.exists():
raise FileNotFoundError(f"Image not found: {IMAGE_PATH.resolve()}")
image = Image.open(IMAGE_PATH).convert("RGB")
image_tensor = ov.Tensor(np.array(image)[None])
pipe = ov_genai.VLMPipeline(str(MODEL_DIR), DEVICE)
prompt = "<start_of_image> Describe this medical image in a concise way."
result = pipe.generate(
prompt,
images=[image_tensor],
max_new_tokens=256,
)
if hasattr(result, "texts"):
print(result.texts[0])
else:
print(result)
if __name__ == "__main__":
main()
| 1 | from pathlib import Path |
| 2 | |
| 3 | import numpy as np |
| 4 | import openvino as ov |
| 5 | import openvino_genai as ov_genai |
| 6 | from PIL import Image |
| 7 | |
| 8 | |
| 9 | MODEL_DIR = Path("models/medgemma-1.5-4b-it-int4-ov") |
| 10 | IMAGE_PATH = Path("images/test.png") |
| 11 | DEVICE = "CPU" # 有 Intel 核显/独显也可以试 "GPU" |
| 12 | |
| 13 | |
| 14 | def main() -> None: |
| 15 | if not MODEL_DIR.exists(): |
| 16 | raise FileNotFoundError(f"Model directory not found: {MODEL_DIR.resolve()}") |
| 17 | |
| 18 | if not IMAGE_PATH.exists(): |
| 19 | raise FileNotFoundError(f"Image not found: {IMAGE_PATH.resolve()}") |
| 20 | |
| 21 | image = Image.open(IMAGE_PATH).convert("RGB") |
| 22 | image_tensor = ov.Tensor(np.array(image)[None]) |
| 23 | |
| 24 | pipe = ov_genai.VLMPipeline(str(MODEL_DIR), DEVICE) |
| 25 | |
| 26 | prompt = "<start_of_image> Describe this medical image in a concise way." |
| 27 | |
| 28 | result = pipe.generate( |
| 29 | prompt, |
| 30 | images=[image_tensor], |
| 31 | max_new_tokens=256, |
| 32 | ) |
| 33 | |
| 34 | if hasattr(result, "texts"): |
| 35 | print(result.texts[0]) |
| 36 | else: |
| 37 | print(result) |
| 38 | |
| 39 | |
| 40 | if __name__ == "__main__": |
| 41 | main() |