File size: 2,561 Bytes
0125293 c751482 446a534 c850c28 c751482 2f801eb c751482 2f801eb c751482 0125293 c751482 336d245 c850c28 0125293 |
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 |
import gradio as gr
import os
import tempfile
import json
import zipfile
from tools.demo_api import build_predictor, run_detection
import shutil
from square_crop import run_square_crop
PREDICTOR = build_predictor(
exp_file = "exps/yolox_s.py",
model_name = "yolox_s",
ckpt_path = "best_ckpt.pth",
device="cpu",
fp16=False,
fuse=False,
trt=False,
conf=0.3,
nms=0.3,
tsize=640,
)
def process_yolox_api(files):
'''API endpoint for YOLOX processing.'''
if not files:
return [], "No files uploaded"
# Create temporary directories
temp_dir = tempfile.mkdtemp()
input_dir = os.path.join(temp_dir, "input")
output_dir = os.path.join(temp_dir, "output")
cropped_dir = os.path.join(temp_dir, "cropped")
os.makedirs(input_dir, exist_ok=True)
os.makedirs(output_dir, exist_ok=True)
os.makedirs(cropped_dir, exist_ok=True)
# Save uploaded files
for file in files:
if file is not None:
shutil.copy(file.name, input_dir)
try:
coco_json_path = run_detection(PREDICTOR, input_dir)
print("Beginning to crop")
cropped_paths = run_square_crop(input_dir, coco_json_path, cropped_dir)
# Create zip file
zip_path = os.path.join(temp_dir, "cropped_results.zip")
with zipfile.ZipFile(zip_path, 'w') as zipf:
for img_path in cropped_paths:
if os.path.exists(img_path):
zipf.write(img_path, os.path.basename(img_path))
return zip_path, cropped_paths, f"Processed {len(cropped_paths)} images"
except Exception as e:
return None, [], f"Error: {str(e)}"
# Create interface with API endpoint
with gr.Blocks() as yolox_demo:
gr.Markdown("# YOLOX Auto-Cropping Service")
with gr.Row():
with gr.Column():
files_input = gr.File(label="Upload Images", file_count="multiple", file_types=["image"])
process_btn = gr.Button("Process", variant="primary")
with gr.Column():
gallery_output = gr.Gallery(label="Cropped Images", columns=3, rows=2)
download_output = gr.File(label="Download Results")
status_output = gr.Textbox(label="Status", interactive=False)
process_btn.click(
fn=process_yolox_api,
inputs=[files_input],
outputs=[download_output, gallery_output, status_output],
api_name="yolox_process" # This creates the API endpoint
)
yolox_demo.launch() |