jsonForSNTK_OCR / app.py
syedfaisalabrar's picture
Update app.py
cab605a verified
import os
import json
import cv2
import gradio as gr
import numpy as np
from PIL import Image
# Path to images folder
image_folder = "./images"
output_folder = "./jsons"
# Ensure output folder exists
os.makedirs(output_folder, exist_ok=True)
# Get all image files
image_files = sorted([f for f in os.listdir(image_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
image_index = 0 # Start at first image
# JSON template for annotation
default_json_template = {
"image_path": "",
"text": {
"TITLE": "",
"TEMPAT-TANGGAL": "",
"NOMOR REGISTRASI": "",
"No.": "",
"NAMA PEMILIK": "",
"ALAMAT": "",
"MERK": "",
"TYPE": "",
"JENIS": "",
"MODEL": "",
"TAHUN PEMBUATAN": "",
"ISI SILINDER/DAYA LISTRIK": "",
"NOMOR RANGKA": "",
"NOMOR MESIN": "",
"NIK/TDP/NIB/KITAS/KITAP": "",
"WARNA": "",
"BAHAN BAKAR": "",
"WARNA TNKB": "",
"TAHUN REGISTRASI": "",
"NOMOR BPKB": "",
"KODE LOKASI": "",
"NO URUT PENDAFTARAN": "",
"BERLAKU SAMPAI": ""
}
}
def load_image():
global image_index
if 0 <= image_index < len(image_files):
image_path = os.path.join(image_folder, image_files[image_index])
image = Image.open(image_path)
json_filename = f"license_{image_index+1:03d}.json"
json_path = os.path.join(output_folder, json_filename)
annotation_data = default_json_template.copy()
if os.path.exists(json_path):
with open(json_path, "r") as f:
annotation_data = json.load(f)
return image, annotation_data["text"], image_path
return None, {}, ""
def save_annotation(image_path, annotation_text):
global image_index
if 0 <= image_index < len(image_files):
json_filename = f"license_{image_index+1:03d}.json"
json_path = os.path.join(output_folder, json_filename)
annotation_data = {"image_path": image_path, "text": annotation_text}
with open(json_path, "w") as f:
json.dump(annotation_data, f, indent=4)
return f"βœ… Saved: {json_filename}"
return "❌ Error saving annotation"
def next_image():
global image_index
if image_index < len(image_files) - 1:
image_index += 1
return load_image()
def prev_image():
global image_index
if image_index > 0:
image_index -= 1
return load_image()
# Gradio Interface
with gr.Blocks() as demo:
with gr.Row():
image_display = gr.Image()
with gr.Column():
text_inputs = {key: gr.Textbox(label=key) for key in default_json_template["text"]}
save_button = gr.Button("Save")
status = gr.Textbox(label="Status", interactive=False)
with gr.Row():
prev_button = gr.Button("Previous")
next_button = gr.Button("Next")
# Load first image on startup
image, annotation, img_path = load_image()
image_display.value = image
for key, textbox in text_inputs.items():
textbox.value = annotation.get(key, "")
# Button Actions
def save_action():
annotation_text = {key: text_inputs[key].value for key in text_inputs}
return save_annotation(img_path, annotation_text)
save_button.click(save_action, outputs=status)
next_button.click(lambda: next_image(), outputs=[image_display, *text_inputs.values()])
prev_button.click(lambda: prev_image(), outputs=[image_display, *text_inputs.values()])
demo.launch()