File size: 4,220 Bytes
e4b4ce0 948a4f1 e4b4ce0 948a4f1 917a01f 948a4f1 917a01f 948a4f1 e4b4ce0 948a4f1 e4b4ce0 240d96f e4b4ce0 b329cda 948a4f1 b329cda 948a4f1 18390a1 948a4f1 b329cda 948a4f1 b329cda 948a4f1 74d81c7 2de3863 74d81c7 948a4f1 e4b4ce0 948a4f1 e4b4ce0 948a4f1 e4b4ce0 18390a1 e4b4ce0 948a4f1 b329cda 18390a1 b329cda 18390a1 b329cda 18390a1 b329cda e4b4ce0 c8b1764 b329cda c8b1764 b329cda e4b4ce0 948a4f1 e4b4ce0 b329cda 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 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 |
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, num_beams=5, min_p=0.0, top_p=1.0):
"""
Process a single image to generate a caption.
Supports image input as file path, numpy array, or PIL Image.
Generation settings (num_beams, min_p, top_p) can be customized.
"""
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=num_beams,
do_sample=True,
top_p=top_p,
min_p=min_p,
)
# 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")
with gr.Column():
output_text = gr.Textbox(label="Output Text")
submit_btn = gr.Button(value="Submit")
num_beams_slider = gr.Slider(
minimum=1,
maximum=5,
step=1,
value=5,
label="Number of Beams"
)
min_p_slider = gr.Slider(
minimum=0,
maximum=1,
step=0.01,
value=0.0,
label="Min-P"
)
top_p_slider = gr.Slider(
minimum=0,
maximum=1,
step=0.01,
value=1.0,
label="Top-P"
)
gr.Examples(
[
["eval_img_1.jpg", 5, 0.0, 1.0],
["eval_img_2.jpg", 5, 0.0, 1.0],
["eval_img_3.jpg", 5, 0.0, 1.0],
["eval_img_4.jpg", 5, 0.0, 1.0],
["eval_img_5.jpg", 5, 0.0, 1.0],
["eval_img_6.jpg", 5, 0.0, 1.0],
["eval_img_7.png", 5, 0.0, 1.0],
["eval_img_8.jpg", 5, 0.0, 1.0],
],
inputs=[input_img, num_beams_slider, min_p_slider, top_p_slider],
outputs=[output_text],
fn=process_image,
label="Try captioning on below examples",
)
submit_btn.click(
process_image,
[input_img, num_beams_slider, min_p_slider, top_p_slider],
[output_text]
)
if __name__ == "__main__":
demo.launch(debug=True)
|