환경 설정 및 패키지 설치

conda activate dify
cp -r /tmp_package/* /root/bayes-tmp/
python3 /root/start_dify.py

cd /root/bayes-tmp/package/
tar -zxf labelme_dep.tar.gz
cd labelme_dep
dpkg -i *.deb
cd ..
tar -xf ./户型图片识别.tar.gz
tar -zxf py_deps.tar.gz
cd py_deps
dpkg -i ./*
cd ..
tar -zxf labelme.tar.gz
cd labelme
pip install *.whl

YOLO11 환경 구성

cd /root/bayes-tmp/package/
conda create -n yolo11 python=3.10 -y
conda activate yolo11
tar -zxf deps.tar.gz
cd deps
pip install *.whl
cd ..
tar -zxf reqs.tar.gz
cd reqs
pip install *.whl
cd ..
tar -zxf modelscope.tar.gz
cd modelscope
pip install *.whl
mkdir -p /root/bayes-tmp/Model_yo11
cp yolo11s.pt /root/bayes-tmp/Model_yo11
tar -zxf libgl1.tar.gz
cd libgl1
dpkg -i *.deb
python -c "import ultralytics; print(ultralytics.__version__)"

BGE 모델 구성 및 실행

cd /root/bayes-tmp/package/
conda create -n bge -y
conda activate bge
export HF_ENDPOINT=http://hf-mirror.com
export XINFERENCE_MODEL_SRC=modelscope
export XINFERENCE_HOME=/root/bayes-tmp
tar -zxf torchvision_torchaudio_torch.tar.gz
cd torchvision_torchaudio_torch
pip install *.whl
cd ..
tar -zxf Xinference.tar.gz
cd Xinference
xinference-local --host 0.0.0.0 --port 9997

데이터셋 준비 및 변환 스크립트

import os
import json
from random import shuffle
from PIL import Image

json_dir = "/root/bayes-tmp/mydataset/标注JSON"
image_dir = "/root/bayes-tmp/mydataset/원본이미지"
output_root = "/root/bayes-tmp/mydataset/dataset"

os.makedirs(output_root, exist_ok=True)

subdirs = ['images/train', 'images/val', 'labels/train', 'labels/val']
for subdir in subdirs:
    os.makedirs(os.path.join(output_root, subdir), exist_ok=True)

class_names = []

def convert_bbox(size, bbox):
    dw = 1.0 / size[0]
    dh = 1.0 / size[1]
    x = (bbox[0] + bbox[2]) / 2.0 * dw
    y = (bbox[1] + bbox[3]) / 2.0 * dh
    w = (bbox[2] - bbox[0]) * dw
    h = (bbox[3] - bbox[1]) * dh
    return [x, y, w, h]

files = [f for f in os.listdir(json_dir) if f.endswith(".json")]
shuffle(files)
split_idx = int(len(files) * 0.8)
train_files, val_files = files[:split_idx], files[split_idx:]

def process_set(file_list, subset):
    for file in file_list:
        with open(os.path.join(json_dir, file), 'r') as f:
            data = json.load(f)
        image_name = os.path.basename(data["imagePath"])
        image_path = os.path.join(image_dir, image_name)
        if not os.path.exists(image_path):
            print(f"❌ {image_path} 파일을 찾을 수 없습니다.")
            continue

        with Image.open(image_path) as img:
            width, height = img.size

        base_name = os.path.splitext(file)[0]
        txt_path = os.path.join(output_root, f"labels/{subset}/{base_name}.txt")
        with open(txt_path, 'w') as out_file:
            for shape in data['shapes']:
                label = shape['label'].strip()
                if label not in class_names:
                    class_names.append(label)
                class_id = class_names.index(label)

                points = shape['points']
                bbox = [min(p[0] for p in points), min(p[1] for p in points),
                       max(p[0] for p in points), max(p[1] for p in points)]
                yolo_box = convert_bbox((width, height), bbox)
                out_file.write(" ".join([str(class_id)] + [f"{x:.6f}" for x in yolo_box]) + "\n")

        dst_image_path = os.path.join(output_root, f"images/{subset}/{image_name}")
        os.symlink(os.path.abspath(image_path), dst_image_path)

process_set(train_files, "train")
process_set(val_files, "val")

with open(os.path.join(output_root, "classes.txt"), 'w') as f:
    f.write("\n".join(class_names))

yaml_content = [
    f"train: {os.path.join(output_root, 'images/train')}",
    f"val: {os.path.join(output_root, 'images/val')}",
    f"nc: {len(class_names)}",
    "names: " + str(class_names)
]
with open(os.path.join(output_root, "data.yaml"), 'w') as f:
    f.write("\n".join(yaml_content))

모델 훈련 스크립트

import torch
from ultralytics import YOLO

torch.cuda.set_per_process_memory_fraction(0.8, device=0)

model_path = "/root/bayes-tmp/Model_yo11/yolo11s.pt"

if not os.path.exists(model_path):
    raise FileNotFoundError(f"모델 파일이 존재하지 않습니다: {model_path}")

file_size = os.path.getsize(model_path) / (1024 * 1024)
if file_size < 10:
    raise ValueError(f"모델 파일 크기가 너무 작습니다: {file_size} MB")

try:
    model = YOLO(model_path, verbose=True)
    print("✅ 모델 로드 성공!")
except Exception as e:
    print(f"❌ 모델 로드 실패: {e}")
    raise

model.train(
    data="/root/bayes-tmp/mydataset/dataset/data.yaml",
    epochs=100,
    imgsz=640,
    batch=16,
    name='yolov5cls',
    workers=0,
    amp=False
)

FastAPI 기반 예측 서비스

import os
import uuid
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import FileResponse, JSONResponse
from ultralytics import YOLO

app = FastAPI()

BASE_DIR = "/root/bayes-tmp"
MODEL_PATH = os.path.join(BASE_DIR, "runs/detect/yolov5cls/weights/best.pt")
MODEL = YOLO(MODEL_PATH)

TEMP_DIR = os.path.join(BASE_DIR, "temp")
PREDICT_DIR = os.path.join(BASE_DIR, "predictions")

os.makedirs(TEMP_DIR, exist_ok=True)
os.makedirs(PREDICT_DIR, exist_ok=True)

@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
    original_name = os.path.basename(file.filename)
    name, ext = os.path.splitext(original_name)
    unique_name = f"{name}_{uuid.uuid4().hex}{ext}"
    original_path = os.path.join(TEMP_DIR, unique_name)

    with open(original_path, "wb") as f:
        f.write(await file.read())

    result = MODEL(original_path, save=True, save_dir=PREDICT_DIR, show_conf=False)
    detections = []
    for box in result[0].boxes:
        cls_id = int(box.cls.item())
        conf = float(box.conf.item())
        label = MODEL.names[cls_id]
        bbox = [round(x, 2) for x in box.xyxy[0].tolist()]
        detections.append({
            "label": label,
            "confidence": conf,
            "bbox": bbox
        })

    json_name = f"{os.path.splitext(unique_name)[0]}_result.json"
    json_path = os.path.join(PREDICT_DIR, json_name)
    with open(json_path, "w") as f:
        json.dump(detections, f, indent=4)

    return {
        "status": "success",
        "detections": detections,
        "json_url": f"http://localhost:8000/json/{json_name}"
    }

@app.get("/json/{filename}")
def get_json(filename: str):
    path = os.path.join(PREDICT_DIR, filename)
    if os.path.exists(path):
        with open(path, "r") as f:
            return JSONResponse(content=json.load(f))
    return JSONResponse(content={"error": "파일을 찾을 수 없습니다."}, status_code=404)

python yolo_api.py
export PATH="/opt/conda/envs/yolo11/bin/:$PATH"
cd package/
tar -zxf uvicorn.tar.gz
cd uvicorn
pip install *.whl
cd /root/bayes-tmp
python -m uvicorn yolo_api:app --host 0.0.0.0 --port 8000