johnbradley's picture
Add initial changes to run drexel_metadata
a258e87
raw
history blame
1.65 kB
import os
import json
import numpy as np
import gradio as gr
import cv2
from drexel_metadata.gen_metadata import gen_metadata
from PIL import Image
def create_temp_file_path(prefix, suffix):
with tempfile.NamedTemporaryFile(prefix=prefix, suffix=suffix, delete=False) as tmpfile:
return tmpfile.name
def run_inference(input_img):
# input_mg: NumPy array with the shape (width, height, 3)
# Save input_mg as a temporary file
tmpfile = create_temp_file_path(prefix="input_", suffix, ".png")
im = Image.fromarray(input_img)
im.save(tmpfile)
# Create temp filenames for output images
visfname = create_temp_file_path(prefix="vis_", suffix=".png")
maskfname = create_temp_file_path(prefix="mask_", suffix=".png")
# Run inference
result = gen_metadata(tmpfile, device='cpu', maskfname=maskfname, visfname=visfname)
json_metadata = json.dumps(result)
# Cleanup
os.remove(tempfile)
return visfname, maskfname, json_metadata
def read_app_header_markdown():
with open('app_header.md') as infile:
return infile.read()
dm_app = gr.Interface(
fn=run_inference,
# Input shows markdown explaining and app and a single image upload panel
inputs=[
gr.Markdown(read_app_header_markdown()),
gr.Image()
],
# Output consists of a visualization image, a masked image, and JSON metadata
outputs=[
gr.Image(label='visualization'),
gr.Image(label='mask'),
gr.JSON(label="JSON metadata")
],
allow_flagging="never" # Do not save user's results or prompt for users to save the results
)
dm_app.launch()