Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,163 Bytes
b764ffe 73cd058 b764ffe db520f8 b764ffe 73cd058 0d5e9d4 b764ffe 0d5e9d4 db520f8 0d5e9d4 db520f8 0d5e9d4 492a9fd 0d5e9d4 492a9fd b764ffe 0d5e9d4 b764ffe |
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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import os
import requests
import torch
import datetime
import subprocess
CUSTOM_CSS = """
#output_box textarea {
font-family: IBM Plex Mono, ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace;
}
"""
# Ensure the model file is in the correct location
model_path = "yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
if not os.path.exists(model_path):
# Download the model file if it doesn't exist
model_url = "https://huggingface.co/DILHTWD/documentlayoutsegmentation_YOLOv8_ondoclaynet/resolve/main/yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
response = requests.get(model_url)
with open(model_path, "wb") as f:
f.write(response.content)
# Load the document segmentation model
docseg_model = YOLO(model_path)
def process_image(image):
# Convert image to the format YOLO model expects
image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
results = docseg_model(image)
# Extract annotated image from results
annotated_img = results[0].plot()
annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
# Prepare detected areas and labels as text output
detected_areas_labels = "\n".join(
[f"{box.label}: {box.conf:.2f}" for box in results[0].boxes]
)
return annotated_img, detected_areas_labels
zero = torch.Tensor([0]).cuda()
print(zero.device) # <-- 'cpu' 🤔
@spaces.GPU
def run_gpu() -> str:
print(zero.device) # <-- 'cuda:0' 🤗
output: str = ""
try:
output = subprocess.check_output(["nvidia-smi"], text=True)
except FileNotFoundError:
output = "nvidia-smi failed"
comment = (
datetime.datetime.now().replace(microsecond=0).isoformat().replace("T", " ")
)
return f"# {comment}\n\n{output}"
def run(check: bool) -> str:
if check:
return run_gpu()
else:
comment = (
datetime.datetime.now().replace(microsecond=0).isoformat().replace("T", " ")
)
return f"# {comment}\n\nThis is running on CPU\n\nClick on 'Run on GPU' below to move to GPU instantly and run nvidia-smi"
output = gr.Textbox(
label="Command Output", max_lines=32, elem_id="output_box", value=run(False)
)
with gr.Blocks(css=CUSTOM_CSS) as demo:
gr.Markdown("#### `zero-gpu`: how to run on serverless GPU for free on Spaces 🔥")
output.render()
check = gr.Checkbox(label="Run on GPU")
check.change(run, inputs=[check], outputs=output, every=1)
# Define the Gradio interface
with gr.Blocks() as interface:
gr.Markdown("### Document Segmentation using YOLOv8")
input_image = gr.Image(type="pil", label="Input Image")
output_image = gr.Image(type="pil", label="Annotated Image")
output_text = gr.Textbox(label="Detected Areas and Labels")
gr.Button("Run").click(
fn=process_image,
inputs=input_image,
outputs=[output_image, output_text]
)
demo.queue().launch(show_api=False)
interface.launch()
if __name__ == "__main__":
demo.launch()
interface.launch()
|