Initial version supporting yolov5_s and yolov5_m

.gitignore

@@ -0,0 +1,2 @@
+flagged/
+__pycache__/

app.py

@@ -0,0 +1,35 @@
+import gradio as gr
+
+import yolov5
+
+models = {
+    "yolov5s": yolov5.load("akbojda/yolov5s-aquarium"),
+    "yolov5m": yolov5.load("akbojda/yolov5m-aquarium"),
+}
+
+def predict(img, model_type):
+    model = models[model_type]
+    results = model(img, size=640)
+    detection_img = results.render()[0]
+
+    return detection_img
+
+# Interface
+inputs = [
+    "image",
+    gr.Dropdown(["yolov5s", "yolov5m"], label="Model", value="yolov5s")
+]
+
+outputs = [
+    "image"
+]
+
+examples = [
+    ["examples/ex1.jpg", None],
+    ["examples/ex2.jpg", None],
+    ["examples/ex3.jpg", None],
+    ["examples/ex4.jpg", None],
+]
+
+iface = gr.Interface(fn=predict, inputs=inputs, outputs=outputs, examples=examples)
+iface.launch()

requirements.txt

@@ -0,0 +1 @@
+yolov5