File size: 3,095 Bytes
e4b4ce0 948a4f1 e4b4ce0 948a4f1 917a01f 948a4f1 917a01f 948a4f1 e4b4ce0 948a4f1 e4b4ce0 240d96f e4b4ce0 948a4f1 74d81c7 2de3863 74d81c7 948a4f1 e4b4ce0 948a4f1 e4b4ce0 948a4f1 e4b4ce0 948a4f1 e4b4ce0 c8b1764 948a4f1 c8b1764 e4b4ce0 948a4f1 e4b4ce0 948a4f1 e4b4ce0 948a4f1 |
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 |
import os
import re
import subprocess
import numpy as np
from PIL import Image
import gradio as gr
import torch
from transformers import AutoProcessor, AutoModelForCausalLM
# Load model and processor, enabling trust_remote_code if needed
model_name = "PJMixers-Images/Florence-2-base-Castollux-v0.5"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).eval()
processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
# Set device (GPU if available)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
TITLE = f"# [{model_name}](https://huggingface.co/{model_name})"
def process_image(image):
"""
Process a single image to generate a caption.
Supports image input as file path, numpy array, or PIL Image.
"""
try:
# Convert input to PIL image if necessary
if isinstance(image, np.ndarray):
image = Image.fromarray(image)
elif isinstance(image, str):
image = Image.open(image)
if image.mode != "RGB":
image = image.convert("RGB")
# Prepare inputs for the model
inputs = processor(text="<CAPTION>", images=image, return_tensors="pt")
# Move tensors to the appropriate device
inputs = {k: v.to(device) for k, v in inputs.items()}
# Disable gradients during inference
with torch.no_grad():
generated_ids = model.generate(
input_ids=inputs["input_ids"],
pixel_values=inputs["pixel_values"],
max_new_tokens=1024,
num_beams=5,
do_sample=True,
)
# Decode and post-process the generated text
return processor.batch_decode(
generated_ids,
skip_special_tokens=False
)[0].replace('</s>', '').replace('<s>', '').replace('<pad>', '').strip()
except Exception as e:
return f"Error processing image: {e}"
# Custom CSS to style the output box
css = """
#output { height: 500px; overflow: auto; border: 1px solid #ccc; }
"""
with gr.Blocks(css=css) as demo:
gr.Markdown(TITLE)
with gr.Tab(label="Single Image Processing"):
with gr.Row():
with gr.Column():
input_img = gr.Image(label="Input Picture")
submit_btn = gr.Button(value="Submit")
with gr.Column():
output_text = gr.Textbox(label="Output Text")
gr.Examples(
[
["eval_img_1.jpg"],
["eval_img_2.jpg"],
["eval_img_3.jpg"],
["eval_img_4.jpg"],
["eval_img_5.jpg"],
["eval_img_6.jpg"],
["eval_img_7.png"],
["eval_img_8.jpg"],
],
inputs=[input_img],
outputs=[output_text],
fn=process_image,
label="Try captioning on below examples",
)
submit_btn.click(process_image, [input_img], [output_text])
if __name__ == "__main__":
demo.launch(debug=True)
|