image_processor.py

import cv2
import numpy as np
import os
from PIL import Image, ImageDraw, ImageOps

class ImageProcessor:
    def __init__(self):
        self.supported_formats = ['jpg', 'jpeg', 'png', 'bmp', 'gif', 'webp']
        
    def load_image(self, image_path, target_size=(2048, 2048)):
        try:
            if not os.path.exists(image_path):
                print(f"Image not found: {image_path}")
                return None
            
            image = cv2.imread(image_path)
            
            if image is None:
                print(f"cv2.imread failed, trying PIL: {image_path}")
                try:
                    from PIL import Image as PILImage
                    pil_img = PILImage.open(image_path)
                    pil_img = pil_img.convert('RGB')
                    image = np.array(pil_img)
                    print(f"Successfully loaded with PIL: {image_path}, size: {image.shape}")
                except Exception as pil_error:
                    print(f"PIL also failed: {pil_error}")
                    return None
            else:
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                print(f"Successfully loaded image: {image_path}, size: {image.shape}")
            
            h, w = image.shape[:2]
            if h > target_size[0] or w > target_size[1]:
                scale = min(target_size[0] / h, target_size[1] / w)
                new_h, new_w = int(h * scale), int(w * scale)
                image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
                print(f"Resized from {h}x{w} to {new_h}x{new_w}")
            
            return image
            
        except Exception as e:
            print(f"Error loading image {image_path}: {e}")
            return None
            
    def add_corners(self, image, radius, bg_color=(255, 255, 255)):
        if radius <= 0:
            return image
            
        try:
            from PIL import Image as PILImage, ImageDraw as PILImageDraw
            
            pil_image = PILImage.fromarray(image)
            width, height = pil_image.size
            
            mask = PILImage.new('L', (width, height), 0)
            draw = PILImageDraw.Draw(mask)
            
            try:
                draw.rounded_rectangle([(0, 0), (width, height)], radius=radius, fill=255)
            except AttributeError:
                print("rounded_rectangle not available, using ellipse method")
                for corner in [(0, 0), (width, 0), (0, height), (width, height)]:
                    if corner == (0, 0):
                        draw.pieslice([0, 0, radius*2, radius*2], 180, 270, fill=255)
                    elif corner == (width, 0):
                        draw.pieslice([width-radius*2, 0, width, radius*2], 270, 360, fill=255)
                    elif corner == (0, height):
                        draw.pieslice([0, height-radius*2, radius*2, height], 90, 180, fill=255)
                    elif corner == (width, height):
                        draw.pieslice([width-radius*2, height-radius*2, width, height], 0, 90, fill=255)
                draw.rectangle([radius, 0, width-radius, height], fill=255)
                draw.rectangle([0, radius, width, height-radius], fill=255)
            
            pil_image_rgba = pil_image.convert('RGBA')
            pil_image_rgba.putalpha(mask)

            background = PILImage.new('RGBA', (width, height), (bg_color[0], bg_color[1], bg_color[2], 255))
            composed = PILImage.alpha_composite(background, pil_image_rgba)
            result = np.array(composed.convert('RGB'))
            print(f"Added corners with radius {radius}, image size: {width}x{height}")
            return result
        except Exception as e:
            print(f"Error adding corners: {e}")
            import traceback
            traceback.print_exc()
            return image
            
    def create_collage(self, images, spacing=4, target_ratio=2.5, radius=4, is_video=None, folder_name="", max_width=2048, max_height=2048):
        print(f"create_collage called with {len(images)} images, ratio={target_ratio}, folder_name={folder_name}")
        
        if not images:
            print("No images provided, returning None")
            return None
            
        if is_video is None:
            is_video = [False] * len(images)
            
        try:
            import math
            
            collage_width = int(max(300, max_width))
            if collage_width <= spacing * 2:
                print("Invalid collage width, returning None")
                return None
                
            aspect_ratios = []
            for img in images:
                h, w = img.shape[:2]
                if h <= 0 or w <= 0:
                    aspect_ratios.append(1.0)
                else:
                    aspect_ratios.append(w / h)
            
            safe_ratio = max(0.5, min(3.0, float(target_ratio)))
            sum_ratios = sum(aspect_ratios) if aspect_ratios else 1.0
            if sum_ratios <= 0:
                sum_ratios = 1.0
                
            target_row_height = collage_width / math.sqrt(sum_ratios * safe_ratio)
            min_row_height = max(80, int(collage_width * 0.08))
            max_row_height = max(min_row_height + 1, int(collage_width * 0.35))
            target_row_height = int(round(max(min_row_height, min(max_row_height, target_row_height))))
            
            print(f"Target row height: {target_row_height}, collage_width: {collage_width}")
            
            header_height = 80 if folder_name else 0
            available_width = collage_width - spacing * 2
            
            def build_rows(row_height_target):
                rows = []
                current_row = []
                current_width = 0.0
                
                for idx, img in enumerate(images):
                    h, w = img.shape[:2]
                    if h <= 0 or w <= 0:
                        continue
                        
                    ratio = w / h
                    width_at_target = ratio * row_height_target
                    
                    if current_row and (current_width + width_at_target + spacing * len(current_row)) > available_width:
                        rows.append(current_row)
                        current_row = []
                        current_width = 0.0
                    
                    current_row.append((idx, img, ratio))
                    current_width += width_at_target
                
                if current_row:
                    rows.append(current_row)
                    
                if not rows:
                    return []
                    
                row_infos = []
                for row_index, row in enumerate(rows):
                    widths = [ratio * row_height_target for (_, _, ratio) in row]
                    sum_widths = sum(widths)
                    if sum_widths <= 0:
                        continue
                        
                    if row_index < len(rows) - 1:
                        scale = (available_width - spacing * (len(row) - 1)) / sum_widths
                    else:
                        scale = min(1.0, (available_width - spacing * (len(row) - 1)) / sum_widths)
                    
                    if scale <= 0:
                        scale = 1.0
                        
                    row_height = max(1, int(round(row_height_target * scale)))
                    row_widths = [max(1, int(round(w * scale))) for w in widths]
                    
                    if row_index < len(rows) - 1:
                        total_row_width = sum(row_widths) + spacing * (len(row) - 1)
                        diff = available_width - total_row_width
                        if diff != 0:
                            row_widths[-1] = max(1, row_widths[-1] + diff)
                    
                    row_infos.append((row, row_height, row_widths))
                
                return row_infos
            
            row_infos = build_rows(target_row_height)
            if not row_infos:
                print("No row info generated, returning None")
                return None
            
            if max_height is None:
                max_height = 0
            max_height = int(max_height)
            if max_height > 0:
                for _ in range(3):
                    total_rows_height = sum(info[1] for info in row_infos)
                    collage_height = header_height + spacing + total_rows_height + spacing * (len(row_infos) - 1) + spacing
                    if collage_height <= max_height:
                        break
                    available_height = max_height - header_height - spacing * 2 - spacing * (len(row_infos) - 1)
                    if available_height <= 0:
                        break
                    scale = available_height / max(1, total_rows_height)
                    new_target = max(10, int(target_row_height * scale))
                    if new_target >= target_row_height:
                        break
                    target_row_height = new_target
                    row_infos = build_rows(target_row_height)
                    if not row_infos:
                        break
                
            total_rows_height = sum(info[1] for info in row_infos)
            collage_height = header_height + spacing + total_rows_height + spacing * (len(row_infos) - 1) + spacing
            collage_height = int(collage_height)
            
            print(f"Collage dimensions: {collage_width}x{collage_height} (header: {header_height})")
            
            collage = np.ones((collage_height, collage_width, 3), dtype=np.uint8) * 255
            
            if folder_name:
                print(f"Adding folder name: {folder_name}")
                from PIL import Image as PILImage, ImageDraw as PILImageDraw, ImageFont as PILImageFont
                try:
                    pil_collage = PILImage.fromarray(collage)
                    draw = PILImageDraw.Draw(pil_collage)
                    
                    font_size = 40
                    font = None
                    font_paths = [
                        "C:\\Windows\\Fonts\\msyh.ttc",
                        "C:\\Windows\\Fonts\\simhei.ttf",
                        "C:\\Windows\\Fonts\\simsun.ttc",
                        "arial.ttf",
                        "C:\\Windows\\Fonts\\arial.ttf"
                    ]
                    
                    for font_path in font_paths:
                        try:
                            print(f"Trying font: {font_path}")
                            font = PILImageFont.truetype(font_path, font_size)
                            print(f"Successfully loaded font: {font_path}, size: {font_size}")
                            break
                        except Exception as font_error:
                            print(f"Failed to load font {font_path}: {font_error}")
                            continue
                    
                    if font is None:
                        print("Using default font")
                        font = PILImageFont.load_default()
                    
                    bbox = draw.textbbox((0, 0), folder_name, font=font)
                    text_width = bbox[2] - bbox[0]
                    text_height = bbox[3] - bbox[1]
                    text_x = (collage_width - text_width) // 2
                    text_y = (header_height - text_height) // 2
                    
                    print(f"Drawing text at ({text_x}, {text_y}), width: {text_width}, height: {text_height}")
                    draw.text((text_x, text_y), folder_name, fill=(0, 0, 0), font=font)
                    
                    collage = np.array(pil_collage)
                    print("Folder name added successfully")
                except Exception as e:
                    print(f"Error adding folder name: {e}")
                    import traceback
                    traceback.print_exc()
            
            print(f"Placing {len(images)} images in adaptive rows")
            
            y = header_height + spacing
            for row, row_height, row_widths in row_infos:
                x = spacing
                for (idx, img, _), new_w in zip(row, row_widths):
                    if new_w <= 0 or row_height <= 0:
                        continue
                    
                    resized = cv2.resize(img, (new_w, row_height), interpolation=cv2.INTER_AREA)
                    if radius > 0:
                        max_radius = max(1, min(new_w, row_height) // 2)
                        applied_radius = min(int(radius), max_radius)
                        resized = self.add_corners(resized, applied_radius)
                    collage[y:y+row_height, x:x+new_w] = resized
                    
                    if is_video[idx]:
                        self._add_video_label(collage, x, y, new_w)
                    
                    x += new_w + spacing
                    
                y += row_height + spacing
            
            print("Collage creation completed successfully")
            return collage
            
        except Exception as e:
            print(f"Error creating collage: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def _add_video_label(self, collage, x, y, width):
        try:
            from PIL import Image as PILImage, ImageDraw as PILImageDraw, ImageFont as PILImageFont
            
            label_width = 78
            label_height = 34
            padding = 10
            corner_radius = 6
            
            label_x = x + width - label_width - padding
            label_y = y + padding
            
            print(f"Adding video label at ({label_x}, {label_y}), image width: {width}")
            
            pil_collage = PILImage.fromarray(collage)
            draw = PILImageDraw.Draw(pil_collage)
            
            font = None
            font_paths = [
                "C:\\Windows\\Fonts\\arial.ttf",
                "arial.ttf",
                "C:\\Windows\\Fonts\\msyh.ttc",
            ]
            
            for font_path in font_paths:
                try:
                    font = PILImageFont.truetype(font_path, 22)
                    break
                except:
                    continue
            
            if font is None:
                font = PILImageFont.load_default()
            
            draw.rounded_rectangle(
                [label_x, label_y, label_x + label_width, label_y + label_height], 
                radius=corner_radius,
                fill=(33, 150, 243), 
                outline=(33, 150, 243)
            )
            
            text = "MP4"
            text_bbox = draw.textbbox((0, 0), text, font=font)
            text_w = text_bbox[2] - text_bbox[0]
            text_h = text_bbox[3] - text_bbox[1]
            
            text_x = label_x + (label_width - text_w) // 2
            text_y = label_y + (label_height - text_h) // 2 - 4
            
            print(f"Drawing MP4 text at ({text_x}, {text_y})")
            draw.text((text_x, text_y), text, fill=(255, 255, 255), font=font)
            
            collage[:] = np.array(pil_collage)
            print(f"Video label added successfully")
        except Exception as e:
            print(f"Error adding video label: {e}")
            import traceback
            traceback.print_exc()
            
    def calculate_optimal_layout(self, total_images):
        import math
        
        if total_images <= 0:
            return 1, 1
            
        cols = int(math.ceil(math.sqrt(total_images)))
        rows = int(math.ceil(total_images / cols))
        
        return rows, cols