Update app.py

app.py

@@ -224,10 +224,6 @@ def create_mask(image, x1, y1, x2, y2):
 
 
 def special_process_image_yolo(risk_level, image_path, point1, point2, thresholds=None):
-    import os
-    import torch
-    from ultralytics import YOLO
-    
     # デバイスの確認
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     print(f"Using device: {device}")
@@ -238,14 +234,10 @@ def special_process_image_yolo(risk_level, image_path, point1, point2, threshold
     # モデルファイルの存在確認
     if not os.path.isfile(model_path):
         raise FileNotFoundError(f"モデルファイル '{model_path}' が見つかりません。パスを確認してください。")
-    else:
-        print("モデルが正常にロードされ、デバイスに移動しました。")
-        # YOLOv8モデルをロードし、指定デバイスに移動
-        model = YOLO(model_path).to(device)
-        
-
-    model.to(device)  # モデルをGPUに移動
-
+    
+    # YOLOv8モデルをロードし、指定デバイスに移動
+    model = YOLO(model_path).to(device)
+    print("モデルが正常にロードされ、デバイスに移動しました。")
 
     # タイムスタンプを作成
     timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
@@ -255,56 +247,56 @@ def special_process_image_yolo(risk_level, image_path, point1, point2, threshold
         return 1 / (1 + np.exp(-k * (risk_level - r0)))
 
     decay_factor = logistic_decay(risk_level)
-    adjusted_thresholds = {key: max(value - decay_factor + 0.8, 0.01) / 2 for key, value in thresholds.items()}
-
-
+    adjusted_thresholds = {key: max(value - decay_factor + 0.8, 0.01) / 2 for key, value in (thresholds or {}).items()}
 
     # 画像の読み込みとRGB変換
     image = cv2.imread(image_path)
     image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    image_np = np.array(image_rgb, dtype=np.uint8)
 
     # 推論実行
     results = model(image_rgb)
 
     # 初期化したマスク画像
-    mask = np.zeros((image_np.shape[0], image_np.shape[1]), dtype=np.uint8)
+    mask = np.zeros(image.shape[:2], dtype=np.uint8)
 
     # 各検出結果に基づきマスク作成
     for box in results[0].boxes:
-        x1, y1, x2, y2 = map(int, box.xyxy[0])  # ボックスの座標
+        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)]
 
         # クラス名に基づいたしきい値
         threshold = adjusted_thresholds.get(object_type, 0.5)
-
         if confidence >= threshold:
             mask[y1:y2, x1:x2] = 255  # ボックス領域を白に設定
 
     # 絶対座標に変換した点の範囲を黒に設定
-    p1_x, p1_y = int(point1[0] * image_np.shape[1]), int(point1[1] * image_np.shape[0])
-    p2_x, p2_y = int(point2[0] * image_np.shape[1]), int(point2[1] * image_np.shape[0])
+    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_np.shape[1], max(p1_x, p2_x)), min(image_np.shape[0], max(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_np.copy()
+    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_pil.save(f"./debug_image_with_rectangle_{timestamp}.jpg")
+    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_pil.save(f"./final_mask_{timestamp}.jpg")
-
-    return f"./final_mask_{timestamp}.jpg"
-
+    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"):