utils/image_batch_classes.py

from pathlib import Path
import json
import numpy as np
from PIL import Image as PILImage
import os
from rtree import index
from shapely.geometry import box

import matplotlib.pyplot as plt
import matplotlib.patches as patches


# Constants for category mappings
catmus_zones_mapping = {
    'DefaultLine': 'Main script black',
    'InterlinearLine': 'Gloss',
        'MainZone': 'Column',
    'DropCapitalZone': 'Plain initial- coloured',
     'StampZone': 'Illustrations',
    'GraphicZone': 'Illustrations',
    'MarginTextZone': 'Gloss',
    'MusicZone': 'Music',
    'NumberingZone': 'Page Number',
    'QuireMarksZone': 'Quire Mark',
    'RunningTitleZone': 'Running header',
    'TitlePageZone': 'Column'
}

coco_class_mapping = {
    'Border': 1,
    'Table': 2,
    'Diagram': 3,
    'Main script black': 4,
    'Main script coloured': 5,
    'Variant script black': 6,
    'Variant script coloured': 7,
    'Historiated': 8,
    'Inhabited': 9,
    'Zoo - Anthropomorphic': 10,
    'Embellished': 11,
    'Plain initial- coloured': 12,
    'Plain initial - Highlighted': 13,
    'Plain initial - Black': 14,
    'Page Number': 15,
    'Quire Mark': 16,
    'Running header': 17,
    'Catchword': 18,
    'Gloss': 19,
    'Illustrations': 20,
    'Column': 21,
    'GraphicZone': 22,
    'MusicLine': 23,
    'MusicZone': 24,
    'Music': 25
}


class Annotation:
    def __init__(self, annotation, image):
        self.name = annotation['name']
        self.cls = annotation['class']
        self.confidence = annotation['confidence']
        self.bbox = annotation['box']
        self.segments = annotation['segments'] if 'segments' in annotation else None
        #Annotation contains name, class, confidence, bbox and segments
        self.image = image
    
    def set_id(self, id):
        self.id = id

    def fix_empty_segments(self,x_coords,y_coords):
        self.segments = {'x': x_coords, 'y': y_coords}

    def segments_to_coco_format(self, segment_dict):
        coco_segment = []
        for x, y in zip(segment_dict['x'], segment_dict['y']):
            coco_segment.append(x)
            coco_segment.append(y)
        return [coco_segment]

    def bbox_to_coco_format(self, box):
        x = box['x1']
        y = box['y1']
        width = box['x2'] - box['x1']
        height = box['y2'] - box['y1']
        return [x, y, width, height]

    def polygon_area(self, segment_dict):
        #Showlace formula for area of polygon
        x = segment_dict['x']
        y = segment_dict['y']
        return 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))

    def unify_names(self):
        self.name = catmus_zones_mapping.get(self.name, self.name)

    def to_coco_format(self, current_annotation_id):
        cls_string = catmus_zones_mapping.get(self.name, self.name)
        cls_int = coco_class_mapping[cls_string]

        if self.segments:
            segmentation = self.segments_to_coco_format(self.segments)
            area = self.polygon_area(self.segments)

        else:
            segmentation = []
            width = self.bbox['x2'] - self.bbox['x1']
            height = self.bbox['y2'] - self.bbox['y1']
            area = width * height
        
        annotation_dict = {
            "id": current_annotation_id,
            "image_id": self.image.id,
            "category_id": cls_int,
            "segmentation": segmentation,
            "area": area,
            "bbox": self.bbox_to_coco_format(self.bbox),
            "iscrowd": 0,
            "attributes": {"occluded": False}
        }
        return annotation_dict
    

class Image:
    def __init__(self, image_path, image_id):
        self.path = image_path
        self.id = image_id
        self.filename = os.path.basename(image_path)
        self.width, self.height = self._get_image_dimensions()
        self.annotations = []
        self.spatial_index = index.Index()
        self.deleted_indices = []
        self.annotations_dict = {}

    def _get_image_dimensions(self):
        with PILImage.open(self.path) as img:
            return img.size
        

    def process_intersection(self, new_box, relevant_classes, overlap_threshold, percentage_dividend, index_to_remove=-1):
        """
        Processes intersection of a new bounding box with existing bounding boxes in the spatial index.

        :param new_box: The new bounding box to check for intersections.
        :param relevant_classes: List of relevant classes to consider for processing.
        :param overlap_threshold: Minimum overlap percentage threshold to consider an intersection.
        :param percentage_dividend: Criterion for calculating percentage overlap ('new_box', 'match_bbox', 'symmetric').
        :param index_to_remove: Index to remove from self.deleted_indices; if -1, remove the intersecting box.
        """
        # Find possible matches using spatial index
        possible_matches = self.spatial_index.intersection(new_box.bounds, objects=True)

        # Iterate over possible matches
        for match in possible_matches:
            # Filter matches based on relevant classes
            if match.object['class'] not in relevant_classes:
                continue

            # Create bounding box for the matched object
            match_bbox = box(*match.bbox)

            # Calculate the intersection area
            intersection_area = new_box.intersection(match_bbox).area

            # Calculate percentage intersection based on the specified dividend
            if percentage_dividend == 'new_box':
                percentage_intersection = intersection_area / new_box.area
            elif percentage_dividend == 'match_bbox':
                percentage_intersection = intersection_area / match_bbox.area
            elif percentage_dividend == 'symmetric':
                # Ensure that both percentages meet the threshold
                percentage_intersection = min(intersection_area / new_box.area, intersection_area / match_bbox.area)
            else:
                raise ValueError("Invalid percentage_dividend value. Must be 'new_box', 'match_bbox', or 'symmetric'.")

            # Append to deleted indices if conditions are met and avoid duplicates
            if percentage_intersection > overlap_threshold:
                to_remove = index_to_remove if index_to_remove != -1 else match.id
                if to_remove not in self.deleted_indices:
                    self.deleted_indices.append(to_remove)


    def process_defaultline(self,new_box,index):
        
        possible_matches = list(self.spatial_index.intersection(new_box.bounds, objects=True))
        #Remove default line if it intersects with any of the following
        variant_colored_matches = [match for match in possible_matches if match.object['class'] in ['Variant script coloured',
        'Variant script black','Main script coloured','NumberingZone','Diagram','MarginTextZone','RunningTitleZone','Table',
        'Quire Mark']]

        if variant_colored_matches:
            self.deleted_indices.append(index)
        else:
            for match in possible_matches:      
                #Remove Main Script Black if its area overlaps with the default line
                if match.object['class']=='Main script black':
                    match_bbox= box(*match.bbox)
                    intersection_area = new_box.intersection(match_bbox).area
                    percentage_intersection = (intersection_area / match_bbox.area)
                    if percentage_intersection > 0.6:
                        self.deleted_indices.append(match.id)
        
    
    def add_annotation(self, annotation):
        #Store indices to remove to remove them at the end
        pos = len(self.annotations)
        #Correct annotations with segments with empty coordinates
        minx,miny,maxx,maxy=annotation.bbox['x1'],annotation.bbox['y1'],annotation.bbox['x2'],annotation.bbox['y2']
        new_box = box(minx,miny,maxx,maxy)

        if annotation.segments: # Execute validations for segmentation models

            if not annotation.segments['x']:
                x_coords = [minx, minx, maxx, maxx, minx]  
                y_coords = [miny, maxy, maxy, miny, miny]
                annotation.fix_empty_segments(x_coords, y_coords)

            if annotation.name in ['Main script black','Main script coloured','Variant script black','Variant script coloured','Plain initial- coloured','Plain initial - Highlighted','Plain initial - Black']:
                self.process_intersection(new_box,['MarginTextZone','NumberingZone'],0.7,'new_box',pos)
            
            if annotation.name in ['Embellished','Plain initial- coloured','Plain initial - Highlighted','Plain initial - Black','Inhabited']:
                self.process_intersection(new_box,['DropCapitalZone','GraphicZone'],0.4,'symmetric')
                
            if annotation.name=='Page Number':
                self.process_intersection(new_box,['NumberingZone'],0.8,'new_box',pos)

            if annotation.name=='Music':
                self.process_intersection(new_box,['MusicZone','GraphicZone'],0.7,'new_box')

            if annotation.name=='Table':
                self.process_intersection(new_box,['MainZone','MarginTextZone'],0.4,'match_bbox')

            if annotation.name in ['Diagram','Illustrations']:
                self.process_intersection(new_box,['GraphicZone'],0.5,'new_box')

            if annotation.name=='DefaultLine':
                
                self.process_defaultline(new_box,pos)


        self.annotations.append(annotation)
        
        annotation.set_id(pos)
        self.spatial_index.insert(pos, new_box.bounds,obj={'class':annotation.name})

    def filter_annotations(self):
    # Convert delete_indices to a set for faster lookup
        delete_indices_set = set(self.deleted_indices)
        filtered_annotations = [item for index, item in enumerate(self.annotations) if index not in delete_indices_set]
        return filtered_annotations
    
    def unify_names(self):
        overlapping_classes = ['MainZone','MarginTextZone']
        for index, annotation in enumerate(self.annotations):
            if index not in self.deleted_indices and annotation.name in overlapping_classes:
                minx,miny,maxx,maxy=annotation.bbox['x1'],annotation.bbox['y1'],annotation.bbox['x2'],annotation.bbox['y2']
                new_box = box(minx,miny,maxx,maxy)

                possible_matches = self.spatial_index.intersection(new_box.bounds, objects=True)

                for match in possible_matches:

                    if match.id not in self.deleted_indices and match.object['class']==annotation.name and match.id!=index:
                        match_bbox= box(*match.bbox)


                        # Calculate the intersection area as a percentage of the smaller box area
                        if new_box.area > match_bbox.area:
                            intersection_area = new_box.intersection(match_bbox).area / match_bbox.area
                        else:
                            intersection_area = match_bbox.intersection(new_box).area / new_box.area

                        if intersection_area > 0.80:
                            delete_index = index if new_box.area < match_bbox.area else match.id
                            self.deleted_indices.append(delete_index)
                        
            annotation.unify_names()



    

    def to_coco_image_dict(self):
        return {
            "id": self.id,
            "width": self.width,
            "height": self.height,
            "file_name": self.filename,
            "license": 0,
            "flickr_url": "",
            "coco_url": "",
            "date_captured": 0
        }
    
    def plot_annotations(self):
    # Load the image
        with PILImage.open(self.path) as img:
            fig, ax = plt.subplots(1, figsize=(self.width / 100, self.height / 100), dpi=100)
            ax.imshow(img)
            
            for annotation in self.filter_annotations():
                if annotation.segments:
                    
                    # Plot polygon segments
                    x = annotation.segments['x']
                    y = annotation.segments['y']
                    # Close the polygon by appending the first point to the end
                    x.append(x[0])
                    y.append(y[0])

                    polygon = patches.Polygon(xy=list(zip(x, y)), closed=True, edgecolor='r', facecolor='none')
                    ax.add_patch(polygon)
                    # Annotate the polygon with the name
                    plt.text(x[0], y[0], annotation.name, color='red', fontsize=25, verticalalignment='top')
                else:
                    # Plot bounding box if no segments
                    bbox = annotation.bbox
                    x1, y1 = bbox['x1'], bbox['y1']
                    x2, y2 = bbox['x2'], bbox['y2']
                    rect = patches.Rectangle(
                        (x1, y1),
                        x2 - x1,
                        y2 - y1,
                        linewidth=1,
                        edgecolor='r',
                        facecolor='none'
                    )
                    ax.add_patch(rect)
                    # Annotate the bounding box with the name
                    plt.text(x1, y1, annotation.name, color='red', fontsize=25, verticalalignment='top')
            
            plt.title(f"Image ID: {self.id} - {self.filename}")
            plt.axis('off')  # Hide axes
            plt.show()



class ImageBatch:
    def __init__(self, image_folder, catmus_labels_folder, emanuskript_labels_folder,zone_labels_folder):
        self.image_folder = image_folder
        self.catmus_labels_folder = catmus_labels_folder
        self.emanuskript_labels_folder = emanuskript_labels_folder
        self.zone_labels_folder = zone_labels_folder
        self.images = []



    def load_images(self):
        image_paths = [
            str(path).replace('\\', '/') 
            for path in Path(self.image_folder).glob('*') 
            if path.is_file()  # Ensure only files are processed
        ]
        image_paths = sorted(image_paths)
        
        for image_id, image_path in enumerate(image_paths, start=1):
            print(f"Processing image: {image_path}")  # Print the image path
            self.images.append(Image(image_path, image_id))


    def load_annotations(self):
        for image in self.images:
            image_basename = os.path.splitext(image.filename)[0]

            catmus_json_path = f'{self.catmus_labels_folder}/{image_basename}.json'
            emanuskript_json_path = f'{self.emanuskript_labels_folder}/{image_basename}.json'
            zone_json_path = f'{self.zone_labels_folder}/{image_basename}.json'

            with open(catmus_json_path) as f:
                catmus_predictions = json.load(f)

            with open(emanuskript_json_path) as f:
                emanuskripts_predictions = json.load(f)

            with open(zone_json_path) as f:
                zone_predictions = json.load(f)

            for annotation_data in zone_predictions + emanuskripts_predictions + catmus_predictions :
                
                if annotation_data['name'] =='Variant script black' and len(annotation_data['segments']['x'])<3:
                    pass
                else:
                    annotation = Annotation(annotation_data, image)
                    image.add_annotation(annotation)

    def unify_names(self):
        for image in self.images:
            image.unify_names()

    def create_coco_dict(self):
        coco_dict = {
            "licenses": [{"name": "", "id": 0, "url": ""}],
            "info": {
                "contributor": "",
                "date_created": "",
                "description": "",
                "url": "",
                "version": "",
                "year": ""
            },
            "categories": [
                {"id": coco_id, "name": cls_name, "supercategory": ""}
                for cls_name, coco_id in coco_class_mapping.items()
            ],
            "annotations": [annotation.to_coco_format(annotation_id) for image in self.images for annotation_id, annotation in enumerate(image.filter_annotations(), start=1)],
            "images": [image.to_coco_image_dict() for image in self.images]
        }
        return coco_dict

    def save_coco_file(self, output_file):
        coco_dict = self.create_coco_dict()
        with open(output_file, 'w') as f:
            json.dump(coco_dict, f, indent=4)

    def return_coco_file(self):
        coco_dict = self.create_coco_dict()
        return coco_dict