Spaces:
Sleeping
Sleeping
File size: 3,102 Bytes
21cad26 5a163b8 c465528 21cad26 5a163b8 21cad26 5a163b8 21cad26 5a163b8 21cad26 5a163b8 c465528 a49a615 87d69ef c465528 a13abfa 5a163b8 c465528 a13abfa 5a163b8 c465528 5a163b8 |
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 |
import json
import skia
import gradio as gr
import os
import random
# Define the category dictionary
category_dict = {
0: "Algorithm",
1: "Caption",
2: "Equation",
3: "Figure",
4: "Footnote",
5: "List",
7: "Table",
8: "Text",
9: "Text-EQ",
10: "Title",
12: "PaperTitle",
13: "Code",
14: "Abstract"
}
# Function to draw bounding box and label on the image using Skia
def draw_annotation(canvas, bbox, category, source_code):
paint = skia.Paint(Color=skia.ColorRED, Style=skia.Paint.kStroke_Style, StrokeWidth=2)
text_paint = skia.Paint(Color=skia.ColorRED, TextSize=12)
# Unpack the bounding box coordinates
x_min, y_min, x_max, y_max = bbox
# Draw the bounding box
canvas.drawRect(skia.Rect.MakeLTRB(x_min, y_min, x_max, y_max), paint)
# Draw the category label
label = f"{category_dict[category]}: {source_code}"
canvas.drawSimpleText(label, x_min, y_min - 10, skia.Font(None, 12), text_paint)
# Load the annotations from a JSON file
def load_annotations(file_path):
with open(file_path, 'r') as file:
annotations = json.load(file)
return annotations
# Main function to render annotations onto images using Skia
def render_annotations(annotations, image_dir):
annotated_images = []
for annotation in annotations:
page_index = annotation['page_index']
image_path = f"{image_dir}/page_{page_index:04d}.jpg"
# Load the image using Skia
image = skia.Image.MakeFromEncoded(skia.Data.MakeFromFileName(image_path))
canvas = skia.Surface(image.width(), image.height()).getCanvas()
canvas.drawImage(image, 0, 0)
# Draw the annotation on the image
draw_annotation(canvas, annotation['bbox'], annotation['category'], annotation['source_code'])
# Save the annotated image to a buffer
surface = canvas.getSurface()
image_buffer = surface.makeImageSnapshot().encode(skia.kPNG_Codec)
annotated_images.append((image_buffer, f"Page {page_index}"))
return annotated_images
# Gradio interface function
def gradio_interface():
example_dir = "examples"
random_index = random.choice([i for i in range(len(os.listdir(example_dir)))])
selected_example_dir_path = os.path.join(example_dir, str(random_index))
# annotations_file = random.choice([f for f in os.listdir(example_dir) if f.endswith('.json')])
annotations_file = os.path.join(selected_example_dir_path, "reading_annotation.json")
annotations = load_annotations(os.path.join(example_dir, annotations_file))
annotated_images = render_annotations(annotations, example_dir)
return [gr.Image(value=image[0], label=image[1]) for image in annotated_images]
# Create Gradio interface
iface = gr.Interface(
fn=gradio_interface,
inputs=[],
outputs=gr.Gallery(label="Annotated Images"),
title="Paper Annotation Renderer",
description="Click to randomly choose an example and render annotations onto the images."
)
# Launch the Gradio interface
iface.launch() |