initial commit

app.py

@@ -0,0 +1,84 @@
+import io
+import math
+from decimal import ROUND_HALF_UP, Decimal
+
+import av
+import requests
+import streamlit as st
+import torch
+from matplotlib import pyplot as plt
+from PIL import Image, ImageDraw, ImageFont
+from streamlit_webrtc import webrtc_streamer
+
+# モデルの読み込み
+yolov5n = torch.hub.load("ultralytics/yolov5", "yolov5n", pretrained=True)
+yolov5s = torch.hub.load("ultralytics/yolov5", "yolov5s", pretrained=True)
+truetype_url = "https://github.com/JotJunior/PHP-Boleto-ZF2/blob/master/public/assets/fonts/arial.ttf?raw=true"
+r = requests.get(truetype_url, allow_redirects=True)
+classes = list(yolov5n.model.names.values())
+
+
+def object_detection(image: Image) -> Image:
+    # 物体検出の実行
+    pred = model(image)
+    # 色の一覧を作成
+    cmap = plt.get_cmap("hsv", len(model.model.names))
+    # フォントサイズ設定
+    sqrt = math.sqrt(image.size[0] * image.size[1] / 10000)
+    size = int(sqrt * 5)
+    font = ImageFont.truetype(io.BytesIO(r.content), size=size)
+    # BBoxの線の太さ設定
+    rec_width = int(sqrt / 2)
+
+    # 検出結果の描画
+    for detections in pred.xyxy:
+        for detection in detections:
+            class_id = int(detection[5])
+            class_name = str(model.model.names[class_id])
+            bbox = [int(x) for x in detection[:4].tolist()]
+            conf = float(detection[4])
+            # 閾値以上のconfidenceの場合のみ描画
+            if conf >= threshold:
+                color = cmap(class_id, bytes=True)
+                draw = ImageDraw.Draw(image)
+                draw.rectangle(bbox, outline=color, width=rec_width)
+                conf_str = Decimal(str(conf * 100)).quantize(
+                    Decimal("0.01"), rounding=ROUND_HALF_UP
+                )
+                draw.text(
+                    [bbox[0] + 5, bbox[1] + 10],
+                    f"{class_name} {conf_str}%",
+                    fill=color,
+                    font=font,
+                )
+
+    return image
+
+
+def callback(frame: av.VideoFrame) -> av.VideoFrame:
+    img = frame.to_image()
+
+    img = object_detection(image=img)
+
+    return av.VideoFrame.from_image(img)
+
+
+# Streamlitの画面設定
+st.set_page_config(page_title="Real-time object detection", page_icon=":shark:")
+
+# サイドバー表示
+classes_str = "\n".join(f"- {item}" for item in classes)
+st.sidebar.markdown(f"データセットに含まれるクラス一覧:\n{classes_str}")
+
+# メイン画面表示
+st.title("Real-time object detection")
+model_name = st.selectbox("Model", ["yolov5n", "yolov5s"])
+model = yolov5n if model_name == "yolov5n" else yolov5s
+threshold = st.slider("Confidence threshold", 0.0, 1.0, 0.25, 0.01)
+
+webrtc_streamer(
+    key="object_detection",
+    video_frame_callback=callback,
+    rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]},
+    media_stream_constraints={"video": True, "audio": False},
+)

packages.txt

@@ -0,0 +1 @@
+libgl1

yolov5n.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f180cf23ba0717ada0badd6c685026d73d48f184d00fc159c2641284b2ac0a3
+size 4062133

yolov5s.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b3b748c1e592ddd8868022e8732fde20025197328490623cc16c6f24d0782ee
+size 14808437