Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from ultralytics import YOLO | |
| import spaces | |
| import torch | |
| import cv2 | |
| import numpy as np | |
| import os | |
| import requests | |
| # Define constants for the new model | |
| ENTITIES_COLORS = { | |
| "Caption": (191, 100, 21), | |
| "Footnote": (2, 62, 115), | |
| "Formula": (140, 80, 58), | |
| "List-item": (168, 181, 69), | |
| "Page-footer": (2, 69, 84), | |
| "Page-header": (83, 115, 106), | |
| "Picture": (255, 72, 88), | |
| "Section-header": (0, 204, 192), | |
| "Table": (116, 127, 127), | |
| "Text": (0, 153, 221), | |
| "Title": (196, 51, 2) | |
| } | |
| BOX_PADDING = 2 | |
| # Load pre-trained YOLOv8 models | |
| model_paths = { | |
| "YOLOv8x Model": "yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt", | |
| "YOLOv8m Model": "yolov8m-doclaynet.pt", | |
| "YOLOv8n Model": "yolov8n-doclaynet.pt", | |
| "DLA Model": "models/dla-model.pt" | |
| } | |
| # Ensure the model files are in the correct location | |
| for model_name, model_path in model_paths.items(): | |
| if not os.path.exists(model_path): | |
| # For demonstration, we only download the YOLOv8x model | |
| if model_name == "YOLOv8x Model": | |
| 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 models | |
| models = {name: YOLO(path) for name, path in model_paths.items()} | |
| # Get class names from the YOLOv8 models | |
| class_names = list(ENTITIES_COLORS.keys()) | |
| def process_image(image, model_choice): | |
| try: | |
| if "YOLOv8" in model_choice: | |
| # Use the selected YOLOv8 model | |
| model = models[model_choice] | |
| results = model(source=image, save=False, show_labels=True, show_conf=True, show_boxes=True) | |
| result = results[0] | |
| # Extract annotated image and labels with class names | |
| annotated_image = result.plot() | |
| detected_areas_labels = "\n".join([ | |
| f"{class_names[int(box.cls.item())].upper()}: {float(box.conf):.2f}" for box in result.boxes | |
| ]) | |
| return annotated_image, detected_areas_labels | |
| elif model_choice == "DLA Model": | |
| # Use the DLA model | |
| image_path = "input_image.jpg" # Temporary save the uploaded image | |
| cv2.imwrite(image_path, cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)) | |
| image = cv2.imread(image_path) | |
| results = models[model_choice].predict(source=image, conf=0.2, iou=0.8) | |
| boxes = results[0].boxes | |
| if len(boxes) == 0: | |
| return image | |
| for box in boxes: | |
| detection_class_conf = round(box.conf.item(), 2) | |
| cls = class_names[int(box.cls)] | |
| start_box = (int(box.xyxy[0][0]), int(box.xyxy[0][1])) | |
| end_box = (int(box.xyxy[0][2]), int(box.xyxy[0][3])) | |
| line_thickness = round(0.002 * (image.shape[0] + image.shape[1]) / 2) + 1 | |
| image = cv2.rectangle(img=image, | |
| pt1=start_box, | |
| pt2=end_box, | |
| color=ENTITIES_COLORS[cls], | |
| thickness=line_thickness) | |
| text = cls + " " + str(detection_class_conf) | |
| font_thickness = max(line_thickness - 1, 1) | |
| (text_w, text_h), _ = cv2.getTextSize(text=text, fontFace=2, fontScale=line_thickness/3, thickness=font_thickness) | |
| image = cv2.rectangle(img=image, | |
| pt1=(start_box[0], start_box[1] - text_h - BOX_PADDING*2), | |
| pt2=(start_box[0] + text_w + BOX_PADDING * 2, start_box[1]), | |
| color=ENTITIES_COLORS[cls], | |
| thickness=-1) | |
| start_text = (start_box[0] + BOX_PADDING, start_box[1] - BOX_PADDING) | |
| image = cv2.putText(img=image, text=text, org=start_text, fontFace=0, color=(255,255,255), fontScale=line_thickness/3, thickness=font_thickness) | |
| return cv2.cvtColor(image, cv2.COLOR_BGR2RGB), "Labels: " + ", ".join(class_names) | |
| else: | |
| return None, "Invalid model choice" | |
| except Exception as e: | |
| return None, f"Error processing image: {e}" | |
| # Create the Gradio Interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Document Segmentation Demo (ZeroGPU)") | |
| with gr.Row(): | |
| input_image = gr.Image(type="pil", label="Upload Image") | |
| output_image = gr.Image(type="pil", label="Annotated Image") | |
| model_choice = gr.Dropdown(list(model_paths.keys()), label="Select Model", value="YOLOv8x Model", scale=0.5) | |
| output_text = gr.Textbox(label="Detected Areas and Labels") | |
| btn = gr.Button("Run Document Segmentation") | |
| btn.click(fn=process_image, inputs=[input_image, model_choice], outputs=[output_image, output_text]) | |
| # Launch the demo with queuing | |
| demo.queue(max_size=1).launch() | |