stzhao's picture
Update app.py
a13abfa verified
raw
history blame
2.83 kB
import json
import skia
import gradio as gr
# 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(annotations_file, image_dir):
annotations = load_annotations(annotations_file.name)
annotated_images = render_annotations(annotations, image_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=[
gr.File(label="Upload Annotations JSON File"),
gr.Textbox(label="Image Directory")
],
outputs=gr.Gallery(label="Annotated Images"),
title="Paper Annotation Renderer",
description="Upload an annotations JSON file and specify the image directory to render annotations onto the images."
)
# Launch the Gradio interface
iface.launch()