# file name : reidai13_11.py
# reidai13_11 画像検出のAIモデルを使った画像の分析アプリの作成
# 左右にウィンドウが分かれないバージョン

import streamlit as st
import cv2
import numpy as np
import pandas as pd
from ultralytics import YOLO

st.set_page_config(page_title='物体検出アプリ',layout='wide')
st.title('物体検出アプリ')

# 1-1. YOLOモデルロード
@st.cache_resource                      # 1回読み込んだモデルを再利用
def load_model(): return YOLO('yolov8n.pt')
model = load_model()

# 1-2. 信頼度しきい値（0～1 を0.05 刻み）
conf_set=st.slider('信頼度しきい値',0.0,1.0,0.5,0.05)

# 1-3. 画像入力（インカメラor アップロード）
how=st.radio('画像の取得方法',
             ['インカメラで撮影','画像をアップロード'],
             horizontal=True)           # ラジオボタン、水平方向に選択肢を表示
if how == 'インカメラで撮影':
    uploaded_image = st.camera_input('インカメラ画像')
else:
    uploaded_image = st.file_uploader('画像ファイルを選択', 
                                      type=['png', 'jpg', 'jpeg'])
if uploaded_image is None:              # 画像が存在しない場合
    st.info(' 画像を撮影またはアップロードしてください。')
    st.stop()

# 2-1. バイト列→ OpenCV(BGR)
arr=np.frombuffer(uploaded_image.getvalue(),np.uint8)
cv2_img=cv2.imdecode(arr,cv2.IMREAD_COLOR)
if cv2_img is None:
    st.error('画像の読み込みに失敗しました。')
    st.stop()
# 大きすぎる画像は長辺1280px に縮小
h,w=cv2_img.shape[:2]
m=max(h,w)
if m > 1280:
    s=1280/m
    cv2_img=cv2.resize(cv2_img,(int(w*s),int(h*s)),
                       interpolation=cv2.INTER_AREA)
    
# 2-2. 物体検出
res=model.predict(cv2_img,conf=conf_set,verbose=False)[0]
bgr_img=res.plot()                      # 枠とラベルを書き込んだBGR 画像
rgb_img=cv2.cvtColor(bgr_img,cv2.COLOR_BGR2RGB) # BGR → RGB
st.subheader('物体検出（画像）')
st.image(rgb_img,caption=f'検出結果（しきい値{conf_set:.2f}）',
         use_container_width=True)
ok,buf=cv2.imencode('.png',bgr_img)     # BGR のままで PNG に変換
if ok: # 変換成功の場合ボタン表示
    st.download_button('結果画像をダウンロード',
        buf.tobytes(),file_name='detection_result.png',
        mime='image/png',use_container_width=True )
    
# 3. 物体検出データをデータフレームに変換し表示
detected_objects=[]
for i,box in enumerate(res.boxes.data.cpu().numpy(),start=1):
    x1,y1,x2,y2,conf,cls=box.astype(float)
    detected_objects.append([i,int(x1),int(y1),int(x2),int(y2),
                             round(conf,2),res.names[int(cls)]])
st.subheader('物体検出数値（データフレーム）')
if detected_objects:
    df=pd.DataFrame(detected_objects,
        columns=['ID','X1','Y1','X2','Y2','Confidence','Class'])
    # データフレームの表示
    st.dataframe(df,use_container_width=True)
    # 検出結果をCSVにダウンロード
    csv_bytes=df.to_csv(index=False).encode('utf-8')
    st.download_button(                 # ボタン表示
        '検出結果をCSVファイルにダウンロード',
        csv_bytes,file_name='detection_data.csv',
        mime='text/csv',use_container_width=True )
else:
    st.info('検出結果はありません。')