iphone.py

import cv2
import pytesseract
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import argparse
import io
import base64

class IPhoneEditor:
    def __init__(self, font_path="SFNSText-Regular.otf"):
        self.font_path = font_path

    def process_image(self, image_path, anggota):
        image = cv2.imread(image_path)
        if image is None:
            print("Error: Could not read image")
            return
        result = self._process_core(image, anggota, show_preview=True)
        return result

    def process_image_bytes(self, image_bytes, anggota):
        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)
            result = self._process_core(image, anggota, show_preview=False, return_theme=True)
            if result is None:
                print("Error: _process_core returned None")
                return None, None
            result_image, theme = result
            if result_image is not None:
                pil_result = Image.fromarray(cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB))
                output_io = io.BytesIO()
                pil_result.save(output_io, format='PNG')
                img_b64 = base64.b64encode(output_io.getvalue()).decode('utf-8')
                return img_b64, theme
            print("Error: result_image is None")
            return None, None
        except Exception as e:
            print(f"Error in process_image_bytes: {str(e)}")
            return None, None

    def _process_core(self, image, anggota, show_preview=False, return_theme=False):
        try:
            extracted_data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
            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:
                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'
            # 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 = 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.95)
                else:
                    font_size = int(font_size * 1.05)
                font_size = max(min_size, min(max_size, font_size))
                iteration += 1
            # Center posisi antara group dan member
            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)
            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)
            if show_preview:
                cv2.imshow('Preview (Tekan q untuk keluar)', image)
                while True:
                    key = cv2.waitKey(1) & 0xFF
                    if key == ord('q'):
                        break
                cv2.destroyAllWindows()
            if return_theme:
                return image, theme
            return image
        except Exception as e:
            print(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')
    args = parser.parse_args()
    editor = IPhoneEditor()
    editor.process_image(args.image_path, args.anggota)