import os import time from openai import OpenAI from volcenginesdkarkruntime import Ark import base64 import io from PIL import Image, ImageDraw import cv2 import numpy as np def encode_image(image): if type(image) == str: try: with open(image, "rb") as image_file: encoding = base64.b64encode(image_file.read()).decode('utf-8') except Exception as e: print(e) with open(image, "r", encoding="utf-8") as image_file: encoding = base64.b64encode(image_file.read()).decode('utf-8') return encoding else: buffered = io.BytesIO() image.save(buffered, format="PNG") return base64.b64encode(buffered.getvalue()).decode('utf-8') def image_mask(image_path: str, bbox_normalized: tuple[int, int, int, int]) -> Image.Image: """Creates a mask on the image in the specified normalized bounding box.""" image = Image.open(image_path) masked_image = image.copy() w, h = image.size # Convert normalized coordinates to pixel coordinates for drawing bbox_pixels = ( int(bbox_normalized[0] * w / 1000), int(bbox_normalized[1] * h / 1000), int(bbox_normalized[2] * w / 1000), int(bbox_normalized[3] * h / 1000) ) draw = ImageDraw.Draw(masked_image) draw.rectangle(bbox_pixels, fill=(255, 255, 255)) # Pure white return masked_image def projection_analysis(image_path: str, bbox_normalized: tuple[int, int, int, int]) -> dict: """ Performs projection analysis on a specified normalized bounding box area. All returned coordinates are also normalized. """ image = cv2.imread(image_path) if image is None: print(f"Error: Failed to read image {image_path}") return {} h, w = image.shape[:2] # Convert normalized bbox to pixel coordinates for cropping bbox_pixels = ( int(bbox_normalized[0] * w / 1000), int(bbox_normalized[1] * h / 1000), int(bbox_normalized[2] * w / 1000), int(bbox_normalized[3] * h / 1000) ) x1, y1, x2, y2 = bbox_pixels roi = image[y1:y2, x1:x2] if roi.size == 0: print(f"Error: Invalid bbox region {bbox_pixels}") return {} gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY) _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) # Perform projection analysis (this part operates on pixels within the ROI) horizontal_projection = np.sum(binary, axis=1) vertical_projection = np.sum(binary, axis=0) # Find groups and convert their coordinates back to normalized space horizontal_groups = _find_groups_and_normalize(horizontal_projection, 'horizontal', bbox_normalized, w, h) vertical_groups = _find_groups_and_normalize(vertical_projection, 'vertical', bbox_normalized, w, h) return { 'horizontal_groups': horizontal_groups, 'vertical_groups': vertical_groups, 'bbox_normalized': bbox_normalized, } def _find_groups_and_normalize(projection: np.ndarray, direction: str, bbox_normalized: tuple[int, int, int, int], image_width: int, image_height: int, min_group_size_px: int = 5, threshold_ratio: float = 0.1) -> list: """ Finds contiguous groups from projection data and returns them in normalized coordinates. """ threshold = np.max(projection) * threshold_ratio non_zero_indices = np.where(projection > threshold)[0] if len(non_zero_indices) == 0: return [] groups_px = [] start_px = non_zero_indices[0] for i in range(1, len(non_zero_indices)): if non_zero_indices[i] > non_zero_indices[i-1] + 1: if non_zero_indices[i-1] - start_px >= min_group_size_px: groups_px.append((start_px, non_zero_indices[i-1])) start_px = non_zero_indices[i] if non_zero_indices[-1] - start_px >= min_group_size_px: groups_px.append((start_px, non_zero_indices[-1])) # Convert pixel groups (relative to ROI) to normalized coordinates (relative to full image) norm_groups = [] roi_x1_norm, roi_y1_norm, roi_x2_norm, roi_y2_norm = bbox_normalized roi_w_norm = roi_x2_norm - roi_x1_norm roi_h_norm = roi_y2_norm - roi_y1_norm roi_w_px = int(roi_w_norm * image_width / 1000) roi_h_px = int(roi_h_norm * image_height / 1000) for start_px, end_px in groups_px: if direction == 'horizontal': start_norm = roi_y1_norm + int(start_px * roi_h_norm / roi_h_px) end_norm = roi_y1_norm + int(end_px * roi_h_norm / roi_h_px) norm_groups.append((roi_x1_norm, roi_x2_norm, start_norm, end_norm)) else: # vertical start_norm = roi_x1_norm + int(start_px * roi_w_norm / roi_w_px) end_norm = roi_x1_norm + int(end_px * roi_w_norm / roi_w_px) norm_groups.append((start_norm, end_norm, roi_y1_norm, roi_y2_norm)) return norm_groups def visualize_projection_analysis(image_path: str, analysis_result: dict, save_path: str = None) -> str: """ Visualizes the results of a completed projection analysis. This function takes the analysis result dictionary and draws it on the image. """ if not analysis_result: print("Error: Analysis result is empty.") return "" image = cv2.imread(image_path) if image is None: print(f"Error: Failed to read image for visualization: {image_path}") return "" h, w = image.shape[:2] vis_image = image.copy() bbox_normalized = analysis_result.get('bbox_normalized') if not bbox_normalized: print("Error: 'bbox_normalized' not found in analysis result.") return "" # Convert normalized bbox to pixel coordinates for drawing the main ROI x1, y1, x2, y2 = ( int(bbox_normalized[0] * w / 1000), int(bbox_normalized[1] * h / 1000), int(bbox_normalized[2] * w / 1000), int(bbox_normalized[3] * h / 1000) ) cv2.rectangle(vis_image, (x1, y1), (x2, y2), (0, 255, 0), 2) # Green for main ROI # Draw horizontal groups (Blue) for i, group_norm in enumerate(analysis_result.get('horizontal_groups', [])): g_x1, g_y1, g_x2, g_y2 = ( int(group_norm[0] * w / 1000), int(group_norm[1] * h / 1000), int(group_norm[2] * w / 1000), int(group_norm[3] * h / 1000) ) cv2.rectangle(vis_image, (g_x1, g_y1), (g_x2, g_y2), (255, 0, 0), 1) cv2.putText(vis_image, f'H{i}', (g_x1, g_y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) # Draw vertical groups (Red) for i, group_norm in enumerate(analysis_result.get('vertical_groups', [])): g_x1, g_y1, g_x2, g_y2 = ( int(group_norm[0] * w / 1000), int(group_norm[1] * h / 1000), int(group_norm[2] * w / 1000), int(group_norm[3] * h / 1000) ) cv2.rectangle(vis_image, (g_x1, g_y1), (g_x2, g_y2), (0, 0, 255), 1) cv2.putText(vis_image, f'V{i}', (g_x1, g_y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1) if save_path is None: base_name = os.path.splitext(os.path.basename(image_path))[0] save_path = f"data/{base_name}_projection_analysis.png" os.makedirs(os.path.dirname(save_path), exist_ok=True) if cv2.imwrite(save_path, vis_image): print(f"Projection analysis visualization saved to: {save_path}") return save_path else: print("Error: Failed to save visualization") return "" class Bot: def __init__(self, key_path, patience=3) -> None: if os.path.exists(key_path): with open(key_path, "r") as f: self.key = f.read().replace("\n", "") else: self.key = key_path self.patience = patience def ask(self): raise NotImplementedError def try_ask(self, question, image_encoding=None, verbose=False): for i in range(self.patience): try: return self.ask(question, image_encoding, verbose) except Exception as e: print(e, "waiting for 5 seconds") time.sleep(5) return None class Doubao(Bot): def __init__(self, key_path, patience=3, model="doubao-1.5-thinking-vision-pro-250428") -> None: super().__init__(key_path, patience) self.client = Ark(api_key=self.key) self.model = model def ask(self, question, image_encoding=None, verbose=False): if image_encoding: content = { "role": "user", "content": [ {"type": "text", "text": question}, { "type": "image_url", "image_url": { "url": f"data:image/png;base64,{image_encoding}", }, }, ], } else: content = {"role": "user", "content": question} response = self.client.chat.completions.create( model=self.model, messages=[content], max_tokens=4096, temperature=0, ) response = response.choices[0].message.content if verbose: print("####################################") print("question:\n", question) print("####################################") print("response:\n", response) # print("seed used: 42") # img = base64.b64decode(image_encoding) # img = Image.open(io.BytesIO(img)) # img.show() return response class Qwen_2_5_VL(Bot): def __init__(self, key_path, patience=3, model="qwen2.5-vl-32b-instruct") -> None: super().__init__(key_path, patience) self.client = OpenAI(api_key=self.key, base_url="https://dashscope.aliyuncs.com/compatible-mode/v1") self.name = model def ask(self, question, image_encoding=None, verbose=False): if image_encoding: content = { "role": "user", "content": [ {"type": "text", "text": question}, { "type": "image_url", "image_url": { "url": f"data:image/png;base64,{image_encoding}" } } ] } else: content = {"role": "user", "content": question} response = self.client.chat.completions.create( model=self.name, messages=[content], max_tokens=4096, temperature=0, seed=42, ) response = response.choices[0].message.content if verbose: print("####################################") print("question:\n", question) print("####################################") print("response:\n", response) print("seed used: 42") return response