maiz

app.py

@@ -1,135 +1,111 @@
-import gradio as gr
+import os
 import cv2
-import numpy as np
-import onnxruntime as ort
 import torch
-import subprocess
-import os
-from gradio_webrtc import WebRTC
-# --- 모델 로드 ---
-
-def convert_pt_to_onnx(pt_path="yolov5n.pt", onnx_path="yolov5n.onnx"):
-    if not os.path.exists(pt_path):
-        raise FileNotFoundError(f"{pt_path} 파일이 존재하지 않습니다.")
-    if os.path.exists(onnx_path):
-        print("ONNX 모델이 이미 존재합니다.")
-        return onnx_path
-
-    # PyTorch YOLOv5 export via subprocess (권장)
-    subprocess.run([
-        "python", "export.py",
-        "--weights", pt_path,
-        "--img", "640",
-        "--batch", "1",
-        "--device", "cpu",
-        "--include", "onnx",
-        "--simplify"
-    ], cwd="yolov5", check=True)
-
-    # 결과 파일 위치 확인
-    if not os.path.exists(onnx_path):
-        raise RuntimeError("ONNX 변환 실패. export.py 실행 결과를 확인하세요.")
-    return onnx_path
-session = None
-input_name = None
-output_names = []
-
-# session = ort.InferenceSession("yolov5n.onnx")
-# input_name = session.get_inputs()[0].name
-# output_names = [output.name for output in session.get_outputs()]
-COCO_CLASSES = ["person", "bicycle", "car", "motorbike", "bus", "truck", "traffic light", "stop sign"]  # 예시
-
-def on_model_upload(weight_file):
-    global session, input_name, output_names
-    session, input_name, output_names = load_model(weight_file)
-    return "모델이 성공적으로 로드되었습니다."
-
-def load_model(weight_file):
-    ext = os.path.splitext(weight_file.name)[1]
-    if ext == ".pt":
-        pt_path = weight_file.name
-        onnx_path = "yolov5n.onnx"
-        convert_pt_to_onnx(pt_path, onnx_path)
-    elif ext == ".onnx":
-        onnx_path = weight_file.name
+import onnx
+import numpy as np
+import gradio as gr
+from huggingface_hub import hf_hub_download
+
+# --- YOLOv5n ONNX 추론 클래스 (간단히 구현) ---
+import onnxruntime
+
+class YOLOv5nONNX:
+    def __init__(self, onnx_path):
+        self.session = onnxruntime.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
+        input_shape = self.session.get_inputs()[0].shape  # e.g. [1,3,640,640]
+        self.input_height, self.input_width = input_shape[2], input_shape[3]
+    
+    def preprocess(self, image):
+        # BGR to RGB, resize, normalize, transpose
+        img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        img = cv2.resize(img, (self.input_width, self.input_height))
+        img = img.astype(np.float32) / 255.0
+        img = np.transpose(img, (2,0,1))  # HWC to CHW
+        img = np.expand_dims(img, axis=0)
+        return img
+    
+    def postprocess(self, outputs, conf_threshold=0.3):
+        # 여기선 간단히 confidence 필터링 정도만, 보통은 NMS 필요
+        preds = outputs[0]
+        preds = preds[preds[:, 4] > conf_threshold]  # conf 필터링
+        return preds
+    
+    def detect_objects(self, image, conf_threshold=0.3):
+        input_tensor = self.preprocess(image)
+        outputs = self.session.run(None, {'images': input_tensor})
+        preds = self.postprocess(outputs, conf_threshold)
+
+        # 간단한 박스 그리기 예시 (여기서는 bbox 좌표 가정 필요, 실제 yolov5 onnx 출력 형태에 맞게 수정 필요)
+        for *box, conf, cls in preds:
+            x1, y1, x2, y2 = map(int, box)
+            cv2.rectangle(image, (x1,y1), (x2,y2), (0,255,0), 2)
+            cv2.putText(image, f'{int(cls)}:{conf:.2f}', (x1,y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
+        return image
+
+# --- PT → ONNX 변환 함수 ---
+def convert_pt_to_onnx(pt_path, onnx_path, input_size=(640, 640)):
+    model = torch.hub.load('ultralytics/yolov5', 'custom', path=pt_path, source='local', force_reload=True)
+    model.eval()
+    dummy_input = torch.randn(1, 3, *input_size)
+    torch.onnx.export(
+        model.model,
+        dummy_input,
+        onnx_path,
+        input_names=['images'],
+        output_names=['output'],
+        opset_version=12,
+        dynamic_axes={'images': {0: 'batch'}, 'output': {0: 'batch'}},
+    )
+    onnx_model = onnx.load(onnx_path)
+    onnx.checker.check_model(onnx_model)
+    print(f"Converted {pt_path} to {onnx_path}")
+
+# --- 모델 로딩 함수 ---
+model = None
+def load_model(weight_path):
+    global model
+    ext = os.path.splitext(weight_path)[1].lower()
+    if ext == '.pt':
+        onnx_path = weight_path.replace('.pt', '.onnx')
+        if not os.path.exists(onnx_path):
+            convert_pt_to_onnx(weight_path, onnx_path)
+        model = YOLOv5nONNX(onnx_path)
+    elif ext == '.onnx':
+        model = YOLOv5nONNX(weight_path)
     else:
-        raise ValueError("지원하지 않는 파일 형식입니다. .pt 또는 .onnx만 허용됩니다.")
+        raise ValueError("지원하지 않는 모델 파일 형식입니다. .pt 또는 .onnx만 가능")
+    print(f"Model loaded from {weight_path}")
 
-    session = ort.InferenceSession(onnx_path)
-    input_name = session.get_inputs()[0].name
-    output_names = [output.name for output in session.get_outputs()]
-    return session, input_name, output_names
+# --- Detection 함수 ---
+def detection(image, weight_file, conf_threshold=0.3):
+    global model
+    if weight_file is not None:
+        if model is None or weight_file.name != getattr(model, 'weight_path', None):
+            # 모델 새로 로드
+            load_model(weight_file.name)
+            model.weight_path = weight_file.name
 
-# --- 객체 탐지 함수 ---
-def detect_live(frame, conf_thresh=0.3):
-    if frame is None:
-        return None
+    if model is None:
+        return image
 
-    h0, w0 = frame.shape[:2]
-
-    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    img = cv2.resize(img, (640, 640))
-    img = img.astype(np.float32) / 255.0
-    img = img.transpose(2, 0, 1)
-    img = np.expand_dims(img, axis=0)
-
-    outputs = session.run(output_names, {input_name: img})
-    preds = outputs[0].squeeze()
-
-    boxes = []
-    for det in preds:
-        conf = det[4]
-        if conf < conf_thresh:
-            continue
-        scores = det[5:]
-        class_id = np.argmax(scores)
-        score = scores[class_id]
-        if score * conf < conf_thresh:
-            continue
-        cx, cy, w, h = det[:4]
-        x1 = int((cx - w/2) * w0 / 640)
-        y1 = int((cy - h/2) * h0 / 640)
-        x2 = int((cx + w/2) * w0 / 640)
-        y2 = int((cy + h/2) * h0 / 640)
-        boxes.append((x1, y1, x2, y2, conf*score, class_id))
-
-    for x1, y1, x2, y2, conf, class_id in boxes:
-        label = f"{COCO_CLASSES[class_id]} {conf:.2f}"
-        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
-        cv2.putText(frame, label, (x1, y1 - 10),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
-
-    return frame
-
-rtc_configuration = {
-    "iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]
-}
+    result_img = model.detect_objects(image, conf_threshold)
+    return result_img
 
 # --- Gradio UI ---
 with gr.Blocks() as demo:
-    gr.HTML("<h1>YOLOv5n 실시간 탐지 (ONNX or PT 업로드 가능)</h1>")
-
-    model_file = gr.File(label="YOLO 모델 업로드 (.pt 또는 .onnx)", file_types=[".pt", ".onnx"])
-    model_status = gr.Textbox(label="모델 상태", interactive=False)
-
-    with gr.Row():
-        webrtc_stream = WebRTC(
-        label="웹캠 스트림",
-        mode="sendrecv",
-        video_frame_callback=lambda frame: detect(frame, conf_slider.value),
-        rtc_configuration=rtc_configuration,
-        media_stream_constraints={"video": True, "audio": False},
-        key="webcam"
+    gr.Markdown("# YOLOv5n ONNX Detection with Weight Upload")
+
+    weight_file = gr.File(label="Upload YOLOv5n weights (.pt or .onnx)", file_types=['.pt', '.onnx'])
+    conf_threshold = gr.Slider(0.0, 1.0, value=0.3, step=0.05, label="Confidence Threshold")
+    input_image = gr.Image(source="webcam", streaming=True)
+    output_image = gr.Image()
+
+    input_image.stream(
+        fn=detection,
+        inputs=[input_image, weight_file, conf_threshold],
+        outputs=output_image,
+        every=0.1,
     )
-        conf_slider = gr.Slider(label="Confidence", minimum=0.1, maximum=1.0, value=0.3)
-
-    output_img = gr.Image(label="탐지 결과")
-
-    model_file.change(fn=on_model_upload, inputs=model_file, outputs=model_status)
-
-    gr.Live(fn=detect_live, inputs=[webcam, conf_slider], outputs=output_img)
-
 
-# --- 실행 ---
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch()