80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
import json
|
|
import math
|
|
from pathlib import Path
|
|
from ultralytics import YOLO
|
|
import requests
|
|
import argparse
|
|
|
|
from clearml import Task # 1. ClearML 임포트
|
|
|
|
# 방법 A) 가장 최신 트렌드인 NMS-Free 임베디드 특화 모델 로드
|
|
model = YOLO("./weights/yolo26m.pt") # Medium 크기 가중치 자동 다운로드
|
|
|
|
|
|
def init_acai_task(output):
|
|
task = Task.init(
|
|
project_name="Normal_Object_Detection",
|
|
task_name="model-yolo26-human"
|
|
)
|
|
|
|
result_data = {"output": output, "status": "PASS"}
|
|
task.upload_artifact(name="final_result", artifact_object=result_data)
|
|
|
|
|
|
def xyxy_to_xywh(box: dict) -> dict:
|
|
x1 = math.floor(box["x1"])
|
|
y1 = math.floor(box["y1"])
|
|
x2 = math.ceil(box["x2"])
|
|
y2 = math.ceil(box["y2"])
|
|
return {"x": x1, "y": y1, "w": x2 - x1, "h": y2 - y1}
|
|
|
|
|
|
# 방법 B) 복잡한 구도나 정밀 식별에 강한 어텐션 기반 모델 로드
|
|
# model = YOLO("yolo12m.pt")
|
|
|
|
def main(image_path: str):
|
|
# uv run main.py --image_url "https://acai.ketidev.kr:20443/detect/image/202606/20260619_145116_image.jpg"
|
|
# 이미지 원격 추론 테스트 (COCO 80종 기본 탐지 가능)
|
|
# results = model("https://acai.ketidev.kr:20443/detect/image/202606/20260619_145116_image.jpg")
|
|
results = model(image_path)
|
|
|
|
result = results[0]
|
|
|
|
image_path = Path(result.path)
|
|
marked_path = image_path.with_name(f"{image_path.stem}_marked{image_path.suffix}")
|
|
result.save(filename=str(marked_path))
|
|
|
|
detections = []
|
|
for det in result.summary():
|
|
detections.append({**det, "box": xyxy_to_xywh(det["box"])})
|
|
|
|
output = {
|
|
"path": result.path,
|
|
"marked_path": str(marked_path),
|
|
"shape": {"height": result.orig_shape[0], "width": result.orig_shape[1]},
|
|
"speed_ms": result.speed,
|
|
"detections": detections,
|
|
}
|
|
print(json.dumps(output, indent=2, ensure_ascii=False))
|
|
|
|
init_acai_task(output)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--image_url", type=str)
|
|
args = parser.parse_args()
|
|
|
|
image_url = args.image_url
|
|
|
|
if image_url is None:
|
|
print("Image path is required")
|
|
exit(1)
|
|
|
|
response = requests.get(image_url)
|
|
response.raise_for_status()
|
|
image_data = response.content
|
|
with open("image.jpg", "wb") as f:
|
|
f.write(image_data)
|
|
image_path = "image.jpg"
|
|
|
|
main(image_path) |