Update app.py

app.py

@@ -194,38 +194,119 @@ def analyze_mask_combined(original_image_path, mask_image_path, small_threshold,
 
 
 
+
+import os
+import numpy as np
+import cv2
+import torch
+from datetime import datetime
+from PIL import Image
+from ultralytics import YOLO
+
+# オブジェクト選択を動的に行う関数
+def decide_to_object(risk_level):
+    '''
+    tex = [
+        'text','Name tag', 'License plate', 'Mail', 'Documents', 'QR codes',
+        'barcodes', 'Map', 'Digital screens', 'information board',
+        'signboard', 'poster', 'sign', 'utility pole'
+    ]
+    '''
+    tex = [
+        'text', 'License plate', 'Digital screens',
+        'signboard', 'poster', 'sign', 'logo', 'card', 'window', 'mirror',
+        'Famous landmark', 'cardboard', 'manhole', 'utility pole'
+    ]
+    
+    # この配列の要素の順番を変えると消える順番が変わる。
+    risk_level = int(risk_level / 20) * (len(tex) / 10)  # 個数決定(1/2)
+    return tex[:int(risk_level) + 1]
+
+# マスクを生成する関数
 def create_mask(image, x1, y1, x2, y2):
-    # Create a black image with the same size as the input image
     mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
-
-    # Draw a white rectangle on the mask where the object is located
     cv2.rectangle(mask, (int(x1), int(y1)), (int(x2), int(y2)), 255, -1)
-
     return mask
 
+# 特殊な処理を行う関数
+def special_process_image_yolo(risk_level, image_path, point1, point2, thresholds=None):
+    # デバイスの確認
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
 
-#この下で消去対象を決定
-def decide_to_object(risk_level):
+    # モデルファイルのパス
+    model_path = './1113.pt'
 
-    tex = [
-    'text', 'License plate', 'digital screens',
-    'signboard', 'poster', 'sign', 'logo', 'cardboard', 'manhole', 'electricity pole'
-]
+    # モデルファイルの存在確認
+    if not os.path.isfile(model_path):
+        raise FileNotFoundError(f"モデルファイル '{model_path}' が見つかりません。パスを確認してください。")
 
+    # YOLOv8モデルをロードし、指定デバイスに移動
+    model = YOLO(model_path).to(device)
+    print("モデルが正常にロードされ、デバイスに移動しました。")
 
-    #この配列の要素の順番を変えると消える順番が変わる。
-    risk_level = int(risk_level / 10)*(len(tex)/10)#個数決定(1/2)
-    return tex[:int(risk_level)+1]
+    # タイムスタンプを作成
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
 
-def create_mask(image, x1, y1, x2, y2):
-    # Create a black image with the same size as the input image
-    mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
+    # リスクレベルに基づいた減衰率の計算
+    def logistic_decay(risk_level, k=0.1, r0=50):
+        return 1 / (1 + np.exp(-k * (risk_level - r0)))
 
-    # Draw a white rectangle on the mask where the object is located
-    cv2.rectangle(mask, (int(x1), int(y1)), (int(x2), int(y2)), 255, -1)
+    decay_factor = logistic_decay(risk_level)
+    adjusted_thresholds = {key: max(value - decay_factor + 0.8, 0.01) / 2 for key, value in (thresholds or {}).items()}
 
-    return mask
+    # 画像の読み込みとRGB変換
+    image = cv2.imread(image_path)
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
+    # 推論実行
+    results = model(image_rgb)
+
+    # 初期化したマスク画像
+    mask = np.zeros(image.shape[:2], dtype=np.uint8)
+
+    # リスクレベルに基づいて消去対象のオブジェクトを決定
+    target_objects = decide_to_object(risk_level)
+
+    # 各検出結果に基づきマスク作成
+    for box in results[0].boxes:
+        x1, y1, x2, y2 = map(int, box.xyxy[0])
+        confidence = box.conf[0]
+        class_id = box.cls[0]
+        object_type = model.names[int(class_id)]
+
+        # クラス名が消去対象に含まれている場合のみ処理
+        if object_type in target_objects:
+            threshold = adjusted_thresholds.get(object_type, 0.5)
+            if confidence >= threshold:
+                mask = create_mask(image, x1, y1, x2, y2)
+
+    # 絶対座標に変換した点の範囲を黒に設定
+    p1_x, p1_y = int(point1[0] * image.shape[1]), int(point1[1] * image.shape[0])
+    p2_x, p2_y = int(point2[0] * image.shape[1]), int(point2[1] * image.shape[0])
+    x_min, y_min = max(0, min(p1_x, p2_x)), max(0, min(p1_y, p2_y))
+    x_max, y_max = min(image.shape[1], max(p1_x, p2_x)), min(image.shape[0], max(p1_y, p2_y))
+    mask[y_min:y_max, x_min:x_max] = 0  # 範囲を黒に設定
+
+    # デバッグ用に白い長方形を描画
+    debug_image = image_rgb.copy()
+    cv2.rectangle(debug_image, (x_min, y_min), (x_max, y_max), (255, 255, 255), 2)
+
+    # デバッグ画像とマスク画像を保存
+    save_dir = "./saved_images"
+    os.makedirs(save_dir, exist_ok=True)
+    debug_image_pil = Image.fromarray(debug_image)
+    debug_image_path = os.path.join(save_dir, f"debug_image_with_rectangle_{timestamp}.jpg")
+    debug_image_pil.save(debug_image_path)
+
+    mask_image_pil = Image.fromarray(mask)
+    mask_image_path = os.path.join(save_dir, f"final_mask_{timestamp}.jpg")
+    mask_image_pil.save(mask_image_path)
+
+    print(f"デバッグ画像が {debug_image_path} に保存されました。")
+    print(f"マスク画像が {mask_image_path} に保存されました。")
+
+    return mask_image_path
 
 def convert_image_format(input_path, output_format="png"):
     """