Views
No views yet
trust_remote_code=True is needed to run the model. It will only download the external custom codes from the original HuggingFaceM4/Florence-2-DocVQA.)1from transformers import AutoProcessor, AutoModelForCausalLM
2import matplotlib.pyplot as plt
3import matplotlib.patches as patches
4
5model_id = "ucsahin/Florence-2-large-TableDetection"
6model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, device_map="cuda") # load the model on GPU
7processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
8
9def run_example(task_prompt, image, max_new_tokens=128):
10 prompt = task_prompt
11 inputs = processor(text=prompt, images=image, return_tensors="pt")
12 generated_ids = model.generate(
13 input_ids=inputs["input_ids"].cuda(),
14 pixel_values=inputs["pixel_values"].cuda(),
15 max_new_tokens=max_new_tokens,
16 early_stopping=False,
17 do_sample=False,
18 num_beams=3,
19 )
20 generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
21 parsed_answer = processor.post_process_generation(
22 generated_text,
23 task=task_prompt,
24 image_size=(image.width, image.height)
25 )
26 return parsed_answer
27
28def plot_bbox(image, data):
29 # Create a figure and axes
30 fig, ax = plt.subplots()
31 # Display the image
32 ax.imshow(image)
33 # Plot each bounding box
34 for bbox, label in zip(data['bboxes'], data['labels']):
35 # Unpack the bounding box coordinates
36 x1, y1, x2, y2 = bbox
37 # Create a Rectangle patch
38 rect = patches.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=1, edgecolor='r', facecolor='none')
39 # Add the rectangle to the Axes
40 ax.add_patch(rect)
41 # Annotate the label
42 plt.text(x1, y1, label, color='white', fontsize=8, bbox=dict(facecolor='red', alpha=0.5))
43 # Remove the axis ticks and labels
44 ax.axis('off')
45 # Show the plot
46 plt.show()
47
48########### Inference
49from datasets import load_dataset
50
51dataset = load_dataset("ucsahin/pubtables-detection-1500-samples")
52
53example_id = 5
54image = dataset["train"][example_id]["image"]
55
56parsed_answer = run_example("<OD>", image=image)
57plot_bbox(image, parsed_answer["<OD>"])| Training Loss | Epoch | Step | Validation Loss |
|---|---|---|---|
| 1.3199 | 1.0 | 169 | 1.0372 |
| 0.7922 | 2.0 | 338 | 0.9169 |
| 0.6824 | 3.0 | 507 | 0.8411 |
| 0.6109 | 4.0 | 676 | 0.8168 |
| 0.5752 | 5.0 | 845 | 0.7915 |
| 0.5605 | 6.0 | 1014 | 0.7862 |
| 0.5291 | 7.0 | 1183 | 0.7740 |
| 0.517 | 8.0 | 1352 | 0.7683 |
| 0.5139 | 9.0 | 1521 | 0.7642 |
| 0.5005 | 10.0 | 1690 | 0.7601 |