Spaces:
Running
on
Zero
Running
on
Zero
import os | |
import re | |
import random | |
from dataclasses import dataclass | |
from functools import partial | |
import torch | |
import gradio as gr | |
import spaces | |
from datasets import load_dataset | |
from torch.utils.data import DataLoader | |
from transformers import AutoProcessor, Gemma3ForConditionalGeneration | |
from PIL import Image, ImageDraw | |
# --- Configuration --- | |
class Configuration: | |
dataset_id: str = "ariG23498/license-detection-paligemma" | |
model_id: str = "google/gemma-3-4b-pt" | |
checkpoint_id: str = "ariG23498/gemma-3-4b-pt-object-detection" | |
device: str = "cuda" if torch.cuda.is_available() else "cpu" | |
dtype: torch.dtype = torch.bfloat16 | |
batch_size: int = 4 | |
learning_rate: float = 2e-05 | |
epochs: int = 1 | |
# --- Utils --- | |
def parse_paligemma_label(label, width, height): | |
# Extract location codes | |
loc_pattern = r"<loc(\d{4})>" | |
locations = [int(loc) for loc in re.findall(loc_pattern, label)] | |
# Extract category (everything after the last location code) | |
category = label.split(">")[-1].strip() | |
# Order in PaliGemma format is: y1, x1, y2, x2 | |
y1_norm, x1_norm, y2_norm, x2_norm = locations | |
# Convert normalized coordinates to image coordinates | |
x1 = (x1_norm / 1024) * width | |
y1 = (y1_norm / 1024) * height | |
x2 = (x2_norm / 1024) * width | |
y2 = (y2_norm / 1024) * height | |
return category, [x1, y1, x2, y2] | |
def visualize_bounding_boxes(image, label, width, height): | |
# Copy image for drawing | |
draw_image = image.copy() | |
draw = ImageDraw.Draw(draw_image) | |
category, bbox = parse_paligemma_label(label, width, height) | |
draw.rectangle(bbox, outline="red", width=2) | |
draw.text((bbox[0], max(0, bbox[1] - 10)), category, fill="red") | |
return draw_image | |
def test_collate_function(batch_of_samples, processor, dtype): | |
images = [] | |
prompts = [] | |
for sample in batch_of_samples: | |
images.append([sample["image"]]) | |
prompts.append(f"{processor.tokenizer.boi_token} detect \n\n") | |
batch = processor(images=images, text=prompts, return_tensors="pt", padding=True) | |
batch["pixel_values"] = batch["pixel_values"].to(dtype) | |
return batch, images | |
# --- Initialize --- | |
cfg = Configuration() | |
processor = AutoProcessor.from_pretrained(cfg.checkpoint_id) | |
model = Gemma3ForConditionalGeneration.from_pretrained( | |
cfg.checkpoint_id, | |
torch_dtype=cfg.dtype, | |
device_map="cpu", | |
) | |
model.eval() | |
test_dataset = load_dataset(cfg.dataset_id, split="test") | |
def get_sample(): | |
sample = random.choice(test_dataset) | |
images = [[sample["image"]]] | |
prompts = [f"{processor.tokenizer.boi_token} detect \n\n"] | |
batch = processor(images=images, text=prompts, return_tensors="pt", padding=True) | |
batch["pixel_values"] = batch["pixel_values"].to(cfg.dtype) | |
return batch, sample["image"] | |
# --- Prediction Logic --- | |
def run_prediction(): | |
model.to(cfg.device) | |
batch, raw_image = get_sample() | |
batch = {k: v.to(cfg.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()} | |
with torch.no_grad(): | |
generation = model.generate(**batch, max_new_tokens=100) | |
decoded = processor.batch_decode(generation, skip_special_tokens=True)[0] | |
image = raw_image # β FIXED: raw_image is already a PIL.Image | |
width, height = image.size | |
result_image = visualize_bounding_boxes(image, decoded, width, height) | |
return result_image | |
# --- Gradio Interface --- | |
demo = gr.Interface( | |
fn=run_prediction, | |
inputs=[], | |
outputs=gr.Image(type="pil", label="Detected Bounding Box"), | |
title="Gemma3 Object Detector", | |
description="Click 'Generate' to visualize a prediction from a randomly sampled test image.", | |
) | |
if __name__ == "__main__": | |
demo.launch() | |