feat(front): add app.py

app.py

@@ -0,0 +1,42 @@
+from ultralytics import YOLO
+import PIL
+import gradio as gr
+import numpy as np
+import os
+
+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="Update your watermaked image", sources=['upload'], type="pil"),
+    # output displays the image with the bounding boxes
+    outputs=gr.AnnotatedImage(),
+    title="Detect Watermark in 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__":
+    model = YOLO("best.pt")
+    gradio_app.launch(debug=True)

requirements.txt

@@ -0,0 +1,2 @@
+gradio
+ultralytics