import cv2 import pytesseract from PIL import Image, ImageDraw, ImageFont import numpy as np import argparse import io import base64 import time import logging from functools import lru_cache from concurrent.futures import ThreadPoolExecutor, TimeoutError # Setup logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class IPhoneEditor: def __init__(self, font_path="SF-Pro-Display-Regular.otf"): self.font_path = font_path self._cache = {} @lru_cache(maxsize=100) def _perform_ocr(self, image_bytes): """Perform OCR with caching and timeout""" def ocr_task(): # Convert bytes to PIL Image first image_stream = io.BytesIO(image_bytes) pil_image = Image.open(image_stream).convert('RGB') # Convert PIL Image to numpy array image_array = np.array(pil_image) # Convert RGB to BGR for OpenCV image_bgr = cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR) # Crop vertically by 50% height = image_bgr.shape[0] mid_point = height // 2 image_bgr = image_bgr[:mid_point, :] return pytesseract.image_to_data(image_bgr, output_type=pytesseract.Output.DICT) try: with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(ocr_task) result = future.result(timeout=30) # 30 second timeout return result except TimeoutError: logger.error("OCR operation timed out") return None except Exception as e: logger.error(f"OCR error: {str(e)}") return None @staticmethod def parse_anggota(anggota_str, ocr_count=None): """Parse anggota parameter and handle + prefix for addition""" if anggota_str.startswith('+'): # If starts with +, add to existing OCR count if ocr_count is None: logger.error("OCR count is None, cannot perform addition") return None try: addition = int(anggota_str[1:]) result = ocr_count + addition logger.info(f"Adding {addition} to OCR count {ocr_count} = {result}") return result except ValueError: logger.error(f"Invalid number format in anggota: {anggota_str}") return None else: # Direct number try: result = int(anggota_str) logger.info(f"Using direct anggota value: {result}") return result except ValueError: logger.error(f"Invalid number format in anggota: {anggota_str}") return None def process_image(self, image_path, anggota): start_time = time.time() image = cv2.imread(image_path) if image is None: logger.error("Failed to read image") return None # First, get the original OCR count for potential addition original_result = self._process_core(image, "0", show_preview=False, get_ocr_count_only=True) if original_result is None: logger.warning("Could not extract OCR count, using 0 as default") ocr_count = 0 else: ocr_count = original_result # Parse anggota parameter parsed_anggota = self.parse_anggota(anggota, ocr_count) if parsed_anggota is None: logger.error("Invalid anggota parameter") return None result = self._process_core(image, str(parsed_anggota), show_preview=True) end_time = time.time() logger.info(f"Total processing time: {end_time - start_time:.2f} seconds") return result def process_image_bytes(self, image_bytes, anggota): start_time = time.time() try: image_stream = io.BytesIO(image_bytes) pil_image = Image.open(image_stream).convert('RGB') image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR) # First, get the original OCR count for potential addition original_result = self._process_core(image, "0", show_preview=False, get_ocr_count_only=True) if original_result is None: logger.warning("Could not extract OCR count, using 0 as default") ocr_count = 0 else: ocr_count = original_result # Parse anggota parameter parsed_anggota = self.parse_anggota(anggota, ocr_count) if parsed_anggota is None: logger.error("Invalid anggota parameter") return None, None result, theme = self._process_core(image, str(parsed_anggota), show_preview=False, return_theme=True) if result is not None: pil_result = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) output_io = io.BytesIO() pil_result.save(output_io, format='PNG') img_b64 = base64.b64encode(output_io.getvalue()).decode('utf-8') end_time = time.time() logger.info(f"Total processing time: {end_time - start_time:.2f} seconds") return img_b64, theme logger.error("Result image is None") return None, None except Exception as e: logger.error(f"Error in process_image_bytes: {str(e)}") return None, None def _process_core(self, image, anggota, show_preview=False, return_theme=False, get_ocr_count_only=False): # Generate cache key image_hash = hash(image.tobytes()) cache_key = f"{image_hash}_{anggota}_{get_ocr_count_only}" # Check cache if cache_key in self._cache: logger.info("Using cached result") return self._cache[cache_key] try: # Convert image to bytes for OCR _, img_encoded = cv2.imencode('.png', image) img_bytes = img_encoded.tobytes() # Perform OCR with caching extracted_data = self._perform_ocr(img_bytes) if extracted_data is None: return None text_list = extracted_data['text'] # Cari semua pasangan angka + anggota/members/member candidates = [] for i, text in enumerate(text_list): if text.isdigit(): for offset in [1,2]: idx = i + offset if idx < len(text_list): next_text = text_list[idx].lower() if next_text in ["anggota", "members", "member"]: # Cari Grup/Group terdekat di atas group_idx = None for j in range(i-1, max(-1, i-10), -1): if text_list[j] in ["Grup", "Group"]: group_idx = j break # Cari split '-' di antara split_idx = None for j in range(group_idx+1 if group_idx is not None else i, i): if text_list[j] == "-": split_idx = j break # Simpan kandidat candidates.append({ 'group_idx': group_idx, 'split_idx': split_idx, 'number_idx': i, 'member_idx': idx }) # Pilih kandidat dengan group_idx valid dan posisi vertikal berdekatan best = None min_dist = 1e9 for c in candidates: if c['group_idx'] is not None: y_group = extracted_data['top'][c['group_idx']] y_member = extracted_data['top'][c['member_idx']] dist = abs(y_group - y_member) if dist < min_dist: min_dist = dist best = c if not best: logger.error("No valid text pattern found") return None group_idx = best['group_idx'] split_idx = best['split_idx'] number_idx = best['number_idx'] member_idx = best['member_idx'] lang = 'id' if text_list[group_idx] == "Grup" else 'en' # If only getting OCR count, return the original count if get_ocr_count_only: try: ocr_count = int(text_list[number_idx]) logger.info(f"Found OCR count: {ocr_count}") return ocr_count except ValueError: logger.warning("No valid OCR count found, returning 0") return 0 # Ambil posisi group_position = { "left": extracted_data['left'][group_idx], "top": extracted_data['top'][group_idx], "width": extracted_data['width'][group_idx], "height": extracted_data['height'][group_idx], } member_position = { "left": extracted_data['left'][member_idx], "top": extracted_data['top'][member_idx], "width": extracted_data['width'][member_idx], "height": extracted_data['height'][member_idx], } member_count_position = { "left": extracted_data['left'][number_idx], "top": extracted_data['top'][number_idx], "width": extracted_data['width'][number_idx], "height": extracted_data['height'][number_idx], } split_position = None if split_idx is not None: split_position = { "left": extracted_data['left'][split_idx], "top": extracted_data['top'][split_idx], "width": extracted_data['width'][split_idx], "height": extracted_data['height'][split_idx], } # Ambil warna background di sekitar member_position x = member_position['left'] + member_position['width'] + 10 y = member_position['top'] + member_position['height'] // 2 bg_color = image[y, x] rgb = (int(bg_color[0]), int(bg_color[1]), int(bg_color[2])) # Deteksi tema r, g, b = float(bg_color[0]), float(bg_color[1]), float(bg_color[2]) brightness = (r * 299 + g * 587 + b * 114) / 1000 is_dark = brightness < 128 theme = 'Dark Mode' if is_dark else 'Light Mode' font_color = (145, 144, 144, 255) if is_dark else (90, 94, 95, 255) margin = 10 # Masking area for pos in [group_position, split_position, member_position, member_count_position]: if pos: cv2.rectangle( image, (pos['left'] - margin, pos['top'] - margin), (pos['left'] + pos['width'] + margin, pos['top'] + pos['height'] + margin), rgb, -1, ) # Teks baru updated_member_count = { 'id': f"Grup · {anggota} anggota", 'en': f"Group · {anggota} members" }.get(lang, f"Group · {anggota} members") # Penyesuaian font size original_height = member_count_position['height'] original_width = member_count_position['width'] font_size = int(original_height * 1.9) image_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) draw = ImageDraw.Draw(image_pil) min_size = int(original_height * 1.4) max_size = int(original_height * 2.3) max_iterations = 5 # Reduced from 10 iteration = 0 while min_size <= max_size and iteration < max_iterations: font = ImageFont.truetype(self.font_path, font_size) text_bbox = draw.textbbox((0, 0), updated_member_count, font=font) text_height = text_bbox[3] - text_bbox[1] text_width = text_bbox[2] - text_bbox[0] if abs(text_height - original_height) <= 2 and text_width <= original_width * 2: break if text_height > original_height or text_width > original_width * 2: font_size = int(font_size * 0.963) else: font_size = int(font_size * 1.02) font_size = max(min_size, min(max_size, font_size)) iteration += 1 top_y = min(group_position['top'], member_position['top']) - margin bot_y = max(group_position['top'] + group_position['height'], member_position['top'] + member_position['height']) + margin left_x = min(group_position['left'], member_position['left']) - margin right_x = max(member_position['left'] + member_position['width'], group_position['left'] + group_position['width']) + margin center_x = (left_x + right_x) // 2 center_y = (top_y + bot_y) // 2 text_bbox = draw.textbbox((0, 0), updated_member_count, font=font) text_width = text_bbox[2] - text_bbox[0] text_height = text_bbox[3] - text_bbox[1] text_x = center_x - (text_width // 2) text_y = center_y - (text_height // 2) - 5 draw.text((text_x, text_y), updated_member_count, font=font, fill=font_color) image = cv2.cvtColor(np.array(image_pil), cv2.COLOR_BGR2RGB) cv2.imwrite('output.png', image) # Return result result = image if not return_theme else (image, theme) self._cache[cache_key] = result return result except Exception as e: logger.error(f"Error in _process_core: {str(e)}") return None if __name__ == '__main__': parser = argparse.ArgumentParser(description='Proses gambar grup iPhone') parser.add_argument('image_path', help='Path ke file gambar') parser.add_argument('anggota', help='Jumlah anggota (bisa menggunakan + untuk menambah ke jumlah yang ada, contoh: +5 untuk menambah 5)') args = parser.parse_args() editor = IPhoneEditor() editor.process_image(args.image_path, args.anggota) # Contoh penggunaan: # python iphone.py image.png 10 # Set jumlah anggota menjadi 10 # python iphone.py image.png +5 # Tambah 5 ke jumlah anggota yang ada di OCR # python iphone.py image.png +10 # Tambah 10 ke jumlah anggota yang ada di OCR