yolo_functions.py

import cv2
import numpy as np
import matplotlib.pyplot as plt
import random
import hashlib
import os
import numpy as np
import hashlib
import random
import matplotlib.pyplot as plt
import cv2
import easyocr 
import pytesseract



from ultralytics import YOLO 

# 1. Load a YOLOv8 segmentation model (pre-trained weights)
model = YOLO("best.pt") 

def get_label_color_id(label_id):
    """
    Generate a consistent BGR color for a numeric label_id by hashing the ID.
    This ensures that each numeric ID always maps to the same color.
    """
    label_str = str(int(label_id))
    # Use the MD5 hash of the label string as a seed
    seed_value = int(hashlib.md5(label_str.encode('utf-8')).hexdigest(), 16)
    random.seed(seed_value)
    # Return color in BGR format
    return (
        random.randint(50, 255),  # B
        random.randint(50, 255),  # G
        random.randint(50, 255)   # R
    )

def segment_large_image_with_tiles(
    model,
    large_image_path,
    tile_size=1080,
    overlap=60,  # Overlap in pixels
    alpha=0.4,
    display=True
):
    """
    1. Reads a large image from `large_image_path`.
    2. Tiles it into sub-images of size `tile_size` x `tile_size`,
       stepping by (tile_size - overlap) to have overlap regions.
    3. Runs `model.predict()` on each tile and accumulates all polygons (in global coords).
    4. For each class, merges overlapping polygons by:
       - filling them on a single-channel mask
       - finding final contours of the connected regions
    5. Draws merged polygons onto an overlay and alpha-blends with the original image.
    6. Returns the final annotated image (in RGB) and a dictionary of merged contours.
    """

    # Read the large image
    image_bgr = cv2.imread(large_image_path)
    if image_bgr is None:
        raise ValueError(f"Could not load image from {large_image_path}")

    # Convert to RGB (for plotting consistency)
    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    H, W, _ = image_rgb.shape

    # Dictionary to store raw polygon coords for each class
    # (before merging)
    class_mask_dict = {}

    # Step size with overlap
    step = tile_size - overlap if overlap < tile_size else tile_size

    # ------------------------
    # 1) Perform Tiled Inference
    # ------------------------
    for top in range(0, H, step):
        for left in range(0, W, step):
            bottom = min(top + tile_size, H)
            right = min(left + tile_size, W)

            tile_rgb = image_rgb[top:bottom, left:right]

            # Run YOLOv8 model prediction
            results = model.predict(tile_rgb)
            if len(results) == 0:
                continue

            # Typically, results[0] holds the main predictions
            pred = results[0]

            # Check if we have valid masks
            if (pred.masks is None) or (pred.masks.xy is None):
                continue

            tile_masks_xy = pred.masks.xy  # list of polygon coords
            tile_labels = pred.boxes.cls   # list of class IDs

            # Convert to numpy int if needed
            if hasattr(tile_labels, 'cpu'):
                tile_labels = tile_labels.cpu().numpy()
            tile_labels = tile_labels.astype(int).tolist()

            # Accumulate polygon coords in global space
            for label_id, polygon in zip(tile_labels, tile_masks_xy):
                # Convert polygon float coords to int points in shape (N,1,2)
                polygon_pts = polygon.reshape((-1, 1, 2)).astype(np.int32)

                # Offset the polygon to the large image coords
                polygon_pts[:, 0, 0] += left  # x-offset
                polygon_pts[:, 0, 1] += top   # y-offset

                if label_id not in class_mask_dict:
                    class_mask_dict[label_id] = []
                class_mask_dict[label_id].append(polygon_pts)

    # -----------------------------------------
    # 2) Merge Overlapping Polygons For Each Class
    #    by rasterizing them in a mask and then
    #    finding final contours
    # -----------------------------------------
    merged_class_mask_dict = {}
    for label_id, polygons_cv in class_mask_dict.items():
        # Create a blank mask (single channel) for the entire image
        mask = np.zeros((H, W), dtype=np.uint8)

        # Fill all polygons for this label on the mask
        for pts in polygons_cv:
            cv2.fillPoly(mask, [pts], 255)

        # Now findContours to get merged regions
        # Use RETR_EXTERNAL so we just get outer boundaries of each connected region
        contours, _ = cv2.findContours(
            mask,
            mode=cv2.RETR_EXTERNAL,
            method=cv2.CHAIN_APPROX_SIMPLE
        )

        # Store final merged contours
        merged_class_mask_dict[label_id] = contours

    # -----------------------
    # 3) Draw Merged Polygons
    # -----------------------
    overlay = image_rgb.copy()
    for label_id, contours in merged_class_mask_dict.items():
        color_bgr = get_label_color_id(label_id)
        for cnt in contours:
            # Fill each contour on the overlay
            cv2.fillPoly(overlay, [cnt], color_bgr)

    # 4) Alpha blend
    output = cv2.addWeighted(overlay, alpha, image_rgb, 1 - alpha, 0)

    # 5) Optional Display
    if display:
        plt.figure(figsize=(12, 12))
        plt.imshow(output)
        plt.axis('off')
        plt.title("Segmentation on Large Image (Overlapped Tiles + Merged Polygons)")
        plt.show()

    return output, merged_class_mask_dict

def usable_data(img_results, image_1):
    """
    Extract bounding boxes, centers, and polygon areas from the segmentation
    results for a single image. Returns a dictionary keyed by label,
    with each value a list of object data: { 'bbox', 'center', 'area' }.
    """
    width, height = image_1.width, image_1.height
    image_data = {}
    for key in img_results.keys():
        image_data[key] = []
        for polygon in img_results[key]:
            polygon = np.array(polygon)

            # Handle varying polygon shapes
            # If shape is (N, 1, 2) e.g. from cv2 findContours
            if polygon.ndim == 3 and polygon.shape[1] == 1 and polygon.shape[2] == 2:
                polygon = polygon.reshape(-1, 2)
            elif polygon.ndim == 2 and polygon.shape[1] == 1:
                polygon = np.squeeze(polygon, axis=1)

            # Now we expect polygon to be (N, 2):
            xs = polygon[:, 0]
            ys = polygon[:, 1]

            # Bounding box
            xmin, xmax = xs.min(), xs.max()
            ymin, ymax = ys.min(), ys.max()
            bbox = (xmin, ymin, xmax, ymax)

            # Center
            centerX = (xmin + xmax) / 2.0
            centerY = (ymin + ymax) / 2.0
            x = width/2
            y = height/2
            # Direction
            dx = x - centerX
            dy = centerY - y  # Invert y-axis for proper orientation
            if dx > 0 and dy > 0:
                direction = "NE"
            elif dx > 0 and dy < 0:
                direction = "SE"
            elif dx < 0 and dy > 0:
                direction = "NW"
            elif dx < 0 and dy < 0:
                direction = "SW"
            elif dx == 0 and dy > 0:
                direction = "N"
            elif dx == 0 and dy < 0:
                direction = "S"
            elif dy == 0 and dx > 0:
                direction = "E"
            elif dy == 0 and dx < 0:
                direction = "W"
            else:
                direction = "Center"


            # Polygon area (Shoelace formula)
            # area = 0.5 * | x0*y1 + x1*y2 + ... + x_{n-1}*y0 - (y0*x1 + y1*x2 + ... + y_{n-1}*x0 ) |
            area = 0.5 * np.abs(
                np.dot(xs, np.roll(ys, 1)) - np.dot(ys, np.roll(xs, 1))
            )

            image_data[key].append({
                'bbox': bbox,
                'center': (centerX, centerY),
                'area': area,
                "direction": direction
            })
    return image_data

import cv2
import numpy as np
import matplotlib.pyplot as plt

def plot_differences_on_image1(
    image1_path,
    mask_dict1,  # e.g., label_name -> list of contours for image1
    image2_path,
    mask_dict2,  # e.g., label_name -> list of contours for image2
    display=True
):
    """
    Compare two images (and their object masks). Plot all differences on Image 1 only:
      - Red: Objects that are missing on Image 1 (present in Image 2 but not Image 1).
      - Green: Objects that are missing on Image 2 (present in Image 1 but not Image 2).

    :param image1_path: Path to the first image
    :param mask_dict1:  dict[label_name] = [contour1, contour2, ...] for the first image
    :param image2_path: Path to the second image
    :param mask_dict2:  dict[label_name] = [contour1, contour2, ...] for the second image
    :param display:     If True, shows the final overlay with matplotlib.
    :return: A tuple:
             - overlay1 (numpy array in RGB) with all differences highlighted
             - list_of_differences: Names of labels with differences
             - difference_masks: A dict with keys "missing_on_img1" and "missing_on_img2",
               where each key maps to a list of contours (original format) for the respective differences.
    """

    # Read both images
    img1_bgr = cv2.imread(image1_path)
    img2_bgr = cv2.imread(image2_path)
    if img1_bgr is None or img2_bgr is None:
        raise ValueError("Could not read one of the input images.")

    # Convert to RGB
    img1_rgb = cv2.cvtColor(img1_bgr, cv2.COLOR_BGR2RGB)
    img2_rgb = cv2.cvtColor(img2_bgr, cv2.COLOR_BGR2RGB)

    # Check matching dimensions
    H1, W1, _ = img1_rgb.shape
    H2, W2, _ = img2_rgb.shape
    if (H1 != H2) or (W1 != W2):
        raise ValueError("Images must be the same size to compare masks reliably.")

    # Prepare an overlay on top of Image 1
    overlay1 = img1_rgb.copy()

    # Take the union of all labels in both dictionaries
    all_labels = set(mask_dict1.keys()).union(set(mask_dict2.keys()))

    # Colors:
    RED = (255, 0, 0)    # (R, G, B)
    GREEN = (0, 255, 0)  # (R, G, B)

    # Track differences
    list_of_differences = []
    difference_masks = {
        "missing_on_img1": {},  # dict[label_name] = list of contours
        "missing_on_img2": {},  # dict[label_name] = list of contours
    }

    for label_id in all_labels:
        # Create binary masks for this label in each image
        mask1 = np.zeros((H1, W1), dtype=np.uint8)
        mask2 = np.zeros((H1, W1), dtype=np.uint8)

        # Fill polygons for label_id in Image 1
        if label_id in mask_dict1:
            for cnt in mask_dict1[label_id]:
                cv2.fillPoly(mask1, [cnt], 255)

        # Fill polygons for label_id in Image 2
        if label_id in mask_dict2:
            for cnt in mask_dict2[label_id]:
                cv2.fillPoly(mask2, [cnt], 255)

        # Missing on Image 1 (present in Image 2 but not in Image 1)
        # => mask2 AND (NOT mask1)
        missing_on_img1 = cv2.bitwise_and(mask2, cv2.bitwise_not(mask1))

        # Missing on Image 2 (present in Image 1 but not in Image 2)
        # => mask1 AND (NOT mask2)
        missing_on_img2 = cv2.bitwise_and(mask1, cv2.bitwise_not(mask2))

        # Extract contours of differences
        contours_missing_on_img1, _ = cv2.findContours(
            missing_on_img1, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        contours_missing_on_img2, _ = cv2.findContours(
            missing_on_img2, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )

        # Store contours in difference masks
        if contours_missing_on_img1:
            difference_masks["missing_on_img1"][label_id] = contours_missing_on_img1
        if contours_missing_on_img2:
            difference_masks["missing_on_img2"][label_id] = contours_missing_on_img2

        # If there are differences, track the label name
        if contours_missing_on_img1 or contours_missing_on_img2:
            list_of_differences.append(label_id)

        # Color them on the overlay of Image 1:
        for cnt in contours_missing_on_img1:
            cv2.drawContours(overlay1, [cnt], -1, RED, -1)  # highlight in red
        for cnt in contours_missing_on_img2:
            cv2.drawContours(overlay1, [cnt], -1, GREEN, -1)  # highlight in green

    # Display if required
    if display:
        plt.figure(figsize=(10, 8))
        plt.imshow(overlay1)
        plt.title("Differences on Image 1\n(Red: Missing on Image 1, Green: Missing on Image 2)")
        plt.axis("off")
        plt.show()

    return overlay1, list_of_differences, difference_masks

def preprocess_image(image_path):
    """
    1) Load and prepare the image for further analysis.
    2) Convert to grayscale, optionally binarize or threshold.
    3) Return the processed image.
    """
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Optional: adaptive thresholding for clearer linework
    # thresholded = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
    #                                     cv2.THRESH_BINARY, 11, 2)

    return gray

def detect_lines_and_grid(processed_image):
    """
    1) Detect major horizontal/vertical lines using Hough transform or morphological ops.
    2) Identify grid lines by analyzing line segments alignment.
    3) Returns lines or grid intersections.
    """
    edges = cv2.Canny(processed_image, 50, 150, apertureSize=3)

    # Hough line detection for demonstration
    lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=100,
                            minLineLength=100, maxLineGap=10)
    # Here you would parse out vertical/horizontal lines, cluster them, etc.

    return lines

def run_ocr(processed_image, method='easyocr'):
    """
    1) Use an OCR engine to detect text (room labels, dimensions, etc.).
    2) 'method' can switch between Tesseract or EasyOCR.
    3) Return recognized text data (text content and bounding boxes).
    """
    text_data = []

    if method == 'easyocr':
        reader = easyocr.Reader(['en', 'ko'], gpu=False)
        result = reader.readtext(processed_image, detail=1, paragraph=False)
        # result structure: [ [bbox, text, confidence], ... ]
        for (bbox, text, conf) in result:
            text_data.append({'bbox': bbox, 'text': text, 'confidence': conf})
    else:
        # Tesseract approach
        config = r'--psm 6'
        tess_result = pytesseract.image_to_data(processed_image, config=config, output_type=pytesseract.Output.DICT)
        # parse data into a structured list
        for i in range(len(tess_result['text'])):
            txt = tess_result['text'][i].strip()
            if txt:
                x = tess_result['left'][i]
                y = tess_result['top'][i]
                w = tess_result['width'][i]
                h = tess_result['height'][i]
                conf = tess_result['conf'][i]
                text_data.append({
                    'bbox': (x, y, x+w, y+h),
                    'text': txt,
                    'confidence': conf
                })
    return text_data

def detect_symbols_and_rooms(processed_image):
    """
    1) Potentially run object detection (e.g., YOLO, Detectron2) to detect symbols:
       - Doors, balconies, fixtures, etc.
    2) Segment out rooms by combining wall detection + adjacency.
    3) Return data about room polygons, symbols, etc.
    """
    # Placeholder: real implementation would require a trained model or rule-based approach.
    # For demonstration, return empty data.
    rooms_data = []
    symbols_data = []
    return rooms_data, symbols_data



def blueprint_analyzer(image_path):
    """
    Orchestrate the entire pipeline on one image:
      1) Preprocess
      2) Detect structural lines
      3) OCR text detection
      4) Symbol/room detection
      5) Compute area differences or summarize
    """
    processed_img = preprocess_image(image_path)

    lines = detect_lines_and_grid(processed_img)
    text_data = run_ocr(processed_img, method='easyocr')


    return lines, text_data



system_prompt = """You are given two construction blueprint images along with their segmentation data.

Do not present any numeric bounding box or area values in your final answer.
Instead, produce a concise, high-level descriptive summary of the differences, using relative location references or known blueprint areas (e.g., “balcony,” “bathroom,” “central hallway,” etc.).
Treat two objects as identical (and thus ignore them) if:
They have the same label/class, and
Their center coordinates are very close.
If possible, provide an OCR-based overview of changed text or lines in those areas. For example, mention if the balcony area contains new textual annotations or if certain labels have been removed/added.
Output the result in brief, correct Markdown summarizing only the differences between the images (e.g., newly added structures, missing items, changed labeling or text).
Remember: No numeric bounding box or area data should be included in the final response. Use location references (“in the top-right corner,” “in the balcony,” etc.) and class names to describe changes.
"""

system_prompt_2 = """You are analyzing two construction blueprint images (Image 1 and Image 2). Each image has a set of detected objects, including “areas” like Balconies, Rooms, Hallways, etc., and smaller objects like Doors, Walls, or Stairs.

Key Points:

An object is considered to belong to an area if the object's center lies within or very close to that area’s bounding box.
Two objects in different images are considered the same object if:
They share the same label/class, and
Their centers are very close in coordinates. In such a case, ignore them (do not list them) because they have not changed significantly.
Focus only on describing the differences between Image 1 and Image 2, such as:
New objects or areas that appear in Image 2 but not in Image 1 (and vice versa).
Changes in labeling or text (e.g., from an OCR perspective).
Changes in object location or area assignment.
Do NOT output numeric bounding boxes, polygon areas, or center coordinates in your final explanation. Instead, provide a relative or area-based description (e.g., “The door is now located in the balcony,” “There are two new doors in the living room,” “A new label is added near the main hallway,” etc.).
Produce a concise and correct Markdown summary that highlights only significant differences.

"""

system_prompt_3 = """You are analyzing two construction blueprint images (Image 1 and Image 2). For each image, you have:

A set of objects (walls, doors, stairs, etc.) along with information on their labels and centers.
A set of “areas” (e.g., “Balcony,” “Living Room,” “Hallway,” “Bathroom,” etc.) with bounding boxes to identify where each area is located.
Task Requirements:
Identify differences between Image 1 and Image 2:
Newly added objects in Image 2 that were not in Image 1.
Missing objects in Image 2 that were in Image 1.
Objects that have changed location or have changed labels.
Text or label changes, if available.
For missing or newly added objects, describe their location in terms of relative position or known areas (not raw coordinates):
For example, say “the missing doors were originally near the top-left corner, adjacent to the main hallway,” or “new walls have been added in the southeast corner, near the living room.”
Avoid including numeric bounding boxes, polygon areas, or centers in the final explanation.
If two objects (one in Image 1 and one in Image 2) have the same label and nearly identical centers, consider them the same object and do not report them as a difference.
Whenever possible, use known area labels to describe positions (e.g., “within the dining area,” “just north of the bathroom,” “adjacent to the balcony,” etc.).
Return a concise and correct Markdown summary with these differences, focusing on where changes occur.
"""

system_prompt_4 = """You are given two sets of data from two blueprint images (Image 1 and Image 2). Along with each image’s extracted objects, you have:
A set of objects (walls, doors, stairs, etc.) along with information on their labels and centers.
A set of “areas” (e.g., “Balcony,” “Living Room,” “Hallway,” “Bathroom,” etc.) with bounding boxes to identify where each area is located.

A “nearest reference area” for each object, including a small textual description of distance and direction (e.g., “Door #2 is near the Balcony to the east”).
Identifications of which objects match across the two images (same label and close centers).
Your Task
Ignore any objects that match between the two images (same label, nearly identical location).
Summarize the differences: newly added or missing objects, label changes, and any changes in object location.
Use the relative position data (distance/direction text) to describe where each new or missing object is/was in terms of known areas (e.g., “the missing wall in the northern side of the corridor,” “the new door near the balcony,” etc.).
Do not output raw numeric distances, bounding boxes, or polygon areas in your final summary. Instead, give a natural-language location description (e.g., “near the east side of the main hallway,” “slightly south of the balcony,” etc.).
Provide your answer in a concise Markdown format, focusing only on significant differences."""

# user_prompt = f"""I have two construction blueprint images, Image 1 and Image 2, and here are their segmentation results (with bounding boxes, centers, and areas). Please compare them and provide a short Markdown summary of the differences, ignoring any objects that match in both images:

# Image 1:
# image: {image_1}

# json
# Copy
# {image_1_data}
# Image 2:
# image: {image_2}
# json
# Copy
# {image_2_data}

# Please:

# Compare the two images in terms of architectural/structural changes.
# Ignore objects that appear in both images (same label & near-identical centers).
# Refer to changes in relative location or in known blueprint areas (e.g. “balcony,” “living room,” “main hallway”), not numeric bounding boxes or polygon areas.
# Include mentions of new text or lines if any appear based on an OCR-like analysis.
# Output only the differences in a concise Markdown summary."""

# user_prompt_2 = f"""I have two construction blueprint images, Image 1 and Image 2, and here are their segmentation results (with bounding boxes, centers, and areas). Please compare them and provide a short Markdown summary of the differences, ignoring any objects that match in both images:

#     Image 1:
#     image: {image_1}
    
#     json
#     Copy
#     {image_1_data}
#     Image 2:
#     image: {image_2}
#     json
#     Copy
#     {image_2_data}
    
#     Please:
    
#     Ignore objects that appear in both images with matching labels and nearly identical centers.
#     Use the bounding boxes of recognized “areas” (like “Balcony,” “Living Room,” “Bathroom,” etc.) to determine which area new or changed objects belong to. For instance, if a door’s center is inside or very close to the balcony’s bounding box, treat that door as being “in the balcony.”
#     Do not display any raw bounding box coordinates, center points, or numeric area values in your final response.
#     Summarize only the differences (e.g., newly added objects, missing objects, changed textual labels) in a brief Markdown format.
#     Mention if there are text/label changes (e.g., from an OCR perspective) in any particular area or region"""