import cv2 import numpy as np import os from dotenv import load_dotenv from pathlib import Path import csv import re # .envファイルから環境変数を読み込む load_dotenv() # 画像ディレクトリのパスを取得 IMAGE_DIR = os.getenv("IMAGE_DIR") # グローバル変数 clicked_points = [] detected_centers = [] detected_areas = [] current_image = None current_image_name = "" all_results = [] # すべての画像の結果を保存するリスト def region_growing(image, seed_point, threshold=60, max_distance=50): height, width = image.shape[:2] segmented = np.zeros((height, width), dtype=np.uint8) # 8近傍 neighbors = [(0, 1), (1, 0), (0, -1), (-1, 0), (1, 1), (-1, -1), (1, -1), (-1, 1)] stack = [seed_point] segmented[seed_point[1], seed_point[0]] = 255 seed_color = image[seed_point[1], seed_point[0]] while stack: x, y = stack.pop() for dx, dy in neighbors: nx, ny = x + dx, y + dy if 0 <= nx < width and 0 <= ny < height: if segmented[ny, nx] == 0: # シード点からの距離を計算 distance = np.sqrt((nx - seed_point[0]) ** 2 + (ny - seed_point[1]) ** 2) if distance <= max_distance: pixel_color = image[ny, nx] if np.all(np.abs(seed_color.astype(int) - pixel_color.astype(int)) <= threshold): segmented[ny, nx] = 255 stack.append((nx, ny)) return segmented def detect_ping_pong_ball(image, click_point, max_distance): x, y = click_point hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) # 領域拡張法で似た色の領域を抽出（範囲制限付き） segmented = region_growing(hsv, (x, y), threshold=50, max_distance=max_distance) # ノイズ除去 kernel = np.ones((5, 5), np.uint8) segmented = cv2.morphologyEx(segmented, cv2.MORPH_OPEN, kernel) segmented = cv2.morphologyEx(segmented, cv2.MORPH_CLOSE, kernel) # 重心計算 M = cv2.moments(segmented) if M["m00"] > 0: cX = int(M["m10"] / M["m00"]) cY = int(M["m01"] / M["m00"]) area = cv2.countNonZero(segmented) print(f"Region detected: center=({cX}, {cY}), area={area}") return (cX, cY), area print(f"No region detected. Returning clicked point: {click_point}") return click_point, 0 def on_mouse_click(event, x, y, flags, param): global current_image, clicked_points, detected_centers, detected_areas if event == cv2.EVENT_LBUTTONDOWN: clicked_points.append((x, y)) print(f"Clicked point: ({x}, {y})") center, area = detect_ping_pong_ball(current_image, (x, y), max_distance=10) detected_centers.append(center) detected_areas.append(area) if center == (x, y): print("Warning: Detected center is the same as clicked point.") else: print(f"Detected center differs from clicked point. Offset: ({center[0]-x}, {center[1]-y})") # 検出結果を描画 draw_detection() elif event == cv2.EVENT_RBUTTONDOWN: if clicked_points: clicked_points.pop() detected_centers.pop() detected_areas.pop() print("Last click undone") draw_detection() def draw_detection(): global current_image, clicked_points, detected_centers, detected_areas image_copy = current_image.copy() for click, center, area in zip(clicked_points, detected_centers, detected_areas): cv2.circle(image_copy, center, int(np.sqrt(area / np.pi)), (0, 255, 0), 2) cv2.circle(image_copy, center, 2, (0, 0, 255), 3) cv2.circle(image_copy, click, 2, (255, 0, 0), -1) # クリックした点を青で表示 cv2.imshow("Image", image_copy) def start_detection(): global clicked_points, detected_centers, detected_areas clicked_points = [] detected_centers = [] detected_areas = [] cv2.setMouseCallback("Image", on_mouse_click) print("Detection started. Click on ping pong balls. Right-click to undo last click.") def save_results(): global clicked_points, detected_centers, current_image_name, detected_areas, all_results result = [current_image_name] for click, center in zip(clicked_points, detected_centers): result.extend([click[0], click[1], center[0], center[1]]) all_results.append(result) print(f"Results for {current_image_name} added to all_results") print("Detected centers:") for center in detected_centers: print(center) def save_all_results(): global all_results with open("ping-pong-detection.csv", "w", newline="") as csvfile: writer = csv.writer(csvfile) header = ["image_file_name"] for i in range(1, (len(max(all_results, key=len)) - 1) // 4 + 1): header.extend([f"Pos{i}_click_x", f"Pos{i}_click_y", f"Pos{i}_center_x", f"Pos{i}_center_y"]) writer.writerow(header) for result in all_results: writer.writerow(result) print("All results saved to ping-pong-detection.csv") def sort_key(filename): # ファイル名からangleの前の数値を抽出 match = re.search(r"_(-?\d+)_angle", filename) if match: return int(match.group(1)) return 0 # マッチしない場合は0を返す def main(): global current_image, current_image_name if not IMAGE_DIR: print("Error: IMAGE_DIR is not set in .env file") return image_dir = Path(IMAGE_DIR) if not image_dir.is_dir(): print(f"Error: {IMAGE_DIR} is not a valid directory") return # 画像ファイルを取得し、angleの前の数値でソート image_files = list(image_dir.glob("*.png")) + list(image_dir.glob("*.jpg")) image_files.sort(key=lambda x: sort_key(x.name)) cv2.namedWindow("Image") for image_file in image_files: current_image = cv2.imread(str(image_file)) current_image_name = image_file.stem if current_image is None: print(f"Error: Unable to read image {image_file}") continue cv2.imshow("Image", current_image) print(f"Current image: {image_file.name}") start_detection() while True: key = cv2.waitKey(1) & 0xFF if key == ord("f"): # 'f' key to finish and save results save_results() break elif key == 27: # ESC key to exit cv2.destroyAllWindows() save_all_results() return cv2.destroyAllWindows() save_all_results() if __name__ == "__main__": main()