prompt_editor / app_b.py
dsaigc's picture
Rename app.py to app_b.py
86e5926
import os
import gradio as gr
import cv2
import numpy as np
def read_files(source_file):
jpg_files = []
txt_files = []
for file in os.listdir(source_file):
if file.endswith(".jpg"):
jpg_files.append(os.path.join(source_file, file))
txt_file = os.path.join(source_file, file[:-4] + ".txt")
if os.path.exists(txt_file):
txt_files.append(txt_file)
else:
txt_files.append(None)
return jpg_files, txt_files
def load_jpg(jpg_file):
img = cv2.imread(jpg_file)
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
def load_txt(txt_file):
with open(txt_file) as f:
text = f.read().strip()
return text
def save_txt(txt_file, text):
with open(txt_file, "w") as f:
f.write(text)
def main(source_file, target_file):
jpg_files, txt_files = read_files(source_file)
preview_list = gr.inputs.CheckboxGroup(jpg_files, label="Preview List")
input_box = gr.inputs.Textbox(lines=5, placeholder="Enter text here", label="Text Input")
gradio_interface = gr.Interface(
fn=None,
inputs=[preview_list],
outputs=[gr.outputs.Image(type="numpy", label="Image Preview"), input_box],
layout="vertical",
theme="compact",
title="Image Text Editor",
description="Select an image from the preview list and edit the associated text."
)
@gradio_interface.func
def image_text_editor(preview_list):
selected_file = preview_list[0]
selected_index = jpg_files.index(selected_file)
jpg_file = jpg_files[selected_index]
txt_file = txt_files[selected_index]
img = load_jpg(jpg_file)
if txt_file is not None:
text = load_txt(txt_file)
input_box.placeholder = text
suggestions = [f"({i}) {word}" for i, word in enumerate(text.split(","))]
output_box = gr.outputs.CheckboxGroup(suggestions, label="Text Suggestions")
gradio_interface.set_output([img, input_box, output_box])
else:
gradio_interface.set_output([img, input_box])
@gradio_interface.func
def apply_suggestion(suggestion):
original_text = input_box.value
selected_index = int(suggestion[1:].split(")")[0])
words = original_text.split(",")
if selected_index < len(words):
words[selected_index] = ""
new_text = ",".join(words)
input_box.update(new_text)
@gradio_interface.func
def save_image_text():
selected_file = preview_list[0]
selected_index = jpg_files.index(selected_file)
jpg_file = jpg_files[selected_index]
txt_file = os.path.join(target_file, os.path.basename(jpg_file)[:-4] + ".txt")
save_txt(txt_file, input_box.value)
new_jpg_file = os.path.join(target_file, os.path.basename(jpg_file))
os.makedirs(os.path.dirname(new_jpg_file), exist_ok=True)
os.rename(jpg_file, new_jpg_file)
gradio_interface.set_output("Text and image saved successfully.")
gradio_interface.add_output("Text Suggestions", apply_suggestion)
gradio_interface.add_output("Save", save_image_text)
gradio_interface.launch()
if __name__ == "__main__":
source_file = gr.inputs.Textbox(lines=1,label="source_file")
target_file = gr.inputs.Textbox(lines=1,label="target_file")
interface = gr.Interface(
fn=main,
[source_file,target_file]
)
interface.launch()