Spaces:
Runtime error
Runtime error
File size: 1,700 Bytes
f94ad82 84f5b3e f94ad82 05629a1 f94ad82 05629a1 f94ad82 84f5b3e |
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 |
from ultralytics import YOLO
import PIL
import gradio as gr
import numpy as np
import os
model = YOLO("best.pt")
def predict(input_img) -> tuple[np.ndarray | PIL.Image.Image | str, list[tuple[np.ndarray | tuple[int, int, int, int], str]]]:
res = model(input_img)
if len(res) == 0:
return input_img, "No watermark detected"
res = res[0]
# convert res.boxes.xyxy to a tuple of (x1, y1, x2, y2)
bbox = res.boxes.xyxy[0].tolist()
bbox = (int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]))
# convert res.boxes.cls to a string
label = res.boxes.cls[0]
str_label = "Watermark is a logo" if label == 0 else "Watermark is a text"
print(bbox, str_label)
return input_img, [(bbox, str_label)]
gradio_app = gr.Interface(
predict,
inputs=gr.Image(label="Upload your watermaked image", sources=['upload'], type="pil"),
# output displays the image with the bounding boxes
outputs=gr.AnnotatedImage(),
title="Detect Watermark in Images",
description="This demo use a YoloV8 Nano model from Ultralytics, fine-tuned on the PITA Dataset for watermarked images",
examples=[
os.path.join(os.path.dirname(__file__), "samples/example_text1.jpg"),
os.path.join(os.path.dirname(__file__), "samples/example_text2.jpg"),
os.path.join(os.path.dirname(__file__), "samples/example_text3.jpg"),
os.path.join(os.path.dirname(__file__), "samples/example_logo1.jpg"),
os.path.join(os.path.dirname(__file__), "samples/example_logo2.jpg"),
os.path.join(os.path.dirname(__file__), "samples/example_logo3.jpg"),
],
allow_flagging="never"
)
if __name__ == "__main__":
gradio_app.launch()
|