mask2former.py

from torch.utils.data import Dataset, DataLoader
from transformers import Mask2FormerForUniversalSegmentation, Mask2FormerImageProcessor, Mask2FormerConfig
import os
import math
import random
import numpy as np
import json
import cv2
from torch.nn.utils import clip_grad_norm_
import time
from tqdm import tqdm
import albumentations as A
import torch
import tree_commons as tc


class ImageAugmenter:

    def __init__(self):
        self.noise = np.random.randint(0, 256, (tc.CROPPED_IMAGE_HEIGHT, tc.CROPPED_IMAGE_WIDTH, 3), dtype=np.uint8)

    def _noise_indivi_trees_randomly(self, img_arr, mask, p_indivi_survival=1 / 3):
        p_indivi_kill = 1 - p_indivi_survival
        instance_ids = np.unique(mask)
        indivi_tree_ids = instance_ids[instance_ids % 2 == 1]
        indivi_kill_no = int(p_indivi_kill * len(indivi_tree_ids))
        if indivi_kill_no == 0:
            return img_arr, mask

        random.shuffle(indivi_tree_ids)
        not_survived_indivi_trees = indivi_tree_ids[:indivi_kill_no]

        not_survived_trees_mask_bool = np.isin(mask, not_survived_indivi_trees)
        not_survived_trees_mask = (not_survived_trees_mask_bool > 0).astype(np.uint8)

        survived_trees_mask = mask * (1 - not_survived_trees_mask)

        not_survived_trees_mask_3c = cv2.cvtColor(not_survived_trees_mask, cv2.COLOR_GRAY2RGB)

        blended = img_arr * (1 - not_survived_trees_mask_3c) + self.noise * not_survived_trees_mask_3c
        return blended, survived_trees_mask

    def _add_shadows(self, img_arr, mask):

        dx = np.random.choice([-40, -30, -25, 25, 30, 40]).item()
        dy = np.random.choice([-40, -30, -25, 25, 30, 40]).item()

        M = np.float32([[1, 0, dx], [0, 1, dy]]) # type: ignore

        shadow_mask = (mask > 0).astype(np.uint8)
        shadow_mask = cv2.warpAffine(shadow_mask, M, (mask.shape[1], mask.shape[0])) # type: ignore
        shadow_mask = cv2.GaussianBlur(shadow_mask, (13, 13), 70)
        mask_3c = cv2.cvtColor(shadow_mask, cv2.COLOR_GRAY2RGB)

        intensity = np.random.choice([0.45, 0.30, 0.35, 0.20, 0.25, 0.20, 0.15, 0.50])
        shadowed_arr = img_arr * (1 - intensity * mask_3c)

        return np.clip(shadowed_arr, 0, 255).astype(np.uint8)

    def augment_image(self, image_arr, mask):
        h, w = tc.CROPPED_IMAGE_HEIGHT, tc.CROPPED_IMAGE_WIDTH

        transform = A.Compose([
            A.RandomCrop(height=h, width=w, p=1.0, fill_mask=0),
            A.SquareSymmetry(p=1),
            A.RandomBrightnessContrast(
                brightness_limit=(-0.15, 0.15),
                contrast_limit=(-0.05, 0.05),
                p=0.40
            ),
            A.MaskDropout(max_objects=(2, 20), fill=0.0, fill_mask=0.0, p=0.20),
            A.Affine(scale=(0.90, 1.1), keep_ratio=True, p=0.20),
            A.CoarseDropout(num_holes_range=(1, 5), fill_mask=0, p=0.20),
            A.GaussianBlur(blur_limit=(3, 7), p=0.30),
            A.OneOf([
                A.ToGray(p=1.0),
                A.ChannelDropout(p=0.2)
            ], p=0.20)
        ])

        out = transform(image=image_arr, mask=mask)
        aug_img_arr, aug_mask = out["image"], out["mask"]

        if np.random.uniform(0, 1) <= 0.50:
            survival = np.random.choice([1/3, 1.5/3, 2/3])
            aug_img_arr, aug_mask = self._noise_indivi_trees_randomly(aug_img_arr, aug_mask, survival)

        if np.random.uniform(0, 1) <= 0.75:
            aug_img_arr = self._add_shadows(aug_img_arr, aug_mask)

        return aug_img_arr, aug_mask

    def random_crop_image(self, image_arr, mask):
        h, w = tc.CROPPED_IMAGE_HEIGHT, tc.CROPPED_IMAGE_WIDTH

        transform = A.Compose([
            A.RandomCrop(height=h, width=w, p=1.0, fill_mask=0),
            A.SquareSymmetry(p=0.25)
        ])

        out = transform(image=image_arr, mask=mask)
        cropped_img, cropped_mask = out["image"], out["mask"]

        return cropped_img, cropped_mask

def get_contours(mask):
    all_contours = []
    instance_ids = np.unique(mask)
    instance_ids = instance_ids[1:]
    for instance_id in instance_ids:
        mask_bin = (mask == instance_id).astype(np.uint8)
        contours, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if contours:
            all_contours.append(contours[0])
    return all_contours


def get_model_input(img_arr, mask, image_id):
   
    img_arr = img_arr.astype(np.uint8)
   
    indiv_tree_mask = mask * (mask % 2 == 1).astype(np.uint8)
    indiv_tree_contours = get_contours(indiv_tree_mask)

    grp_tree_mask = mask * (mask % 2 == 0).astype(np.uint8)
    grp_tree_contours = get_contours(grp_tree_mask)

    indiv_tree_class_id = tc.LABEL_TO_CLASS_ID[tc.CLASS_INDIVIDUAL_TREE]
    grp_tree_class_id = tc.LABEL_TO_CLASS_ID[tc.CLASS_GROUP_TREES]

    mask = np.zeros((img_arr.shape[0], img_arr.shape[1]), dtype=np.int32)
    mask.fill(255)
    instance_id_to_class_id = {}
    instance_id = 1
    for contours, class_id in zip([indiv_tree_contours, grp_tree_contours], [indiv_tree_class_id, grp_tree_class_id]):
        for contour in contours:
            polygon = contour.squeeze(1)

            box = tc.get_bounding_box(polygon)
            if box[0] >= box[2] or box[1] >= box[3]:
                continue

            cv2.fillPoly(mask, [polygon], color=instance_id) # type: ignore
            instance_id_to_class_id[instance_id] = class_id
            instance_id += 1
            
    return {'image': img_arr, 'instance_id_to_class_id': instance_id_to_class_id, 'mask': np.astype(mask, np.int32), 'image_id' : image_id}



class ImageDataCache:

    def __init__(self):
        self._filename_to_img_data = self._get_image_data_by_filename()
        self._file_type_to_filenames = self._get_filenames_by_file_type(self._filename_to_img_data.keys())

    def _get_filenames_by_file_type(self, filenames):
        file_type_to_filenames = {}
        for prefix in tc.FILE_TYPES_PREFIX:
            file_type_to_filenames[prefix] = []
            
        for filename in filenames:
            for prefix in tc.FILE_TYPES_PREFIX:
                if filename.startswith(prefix):
                    file_type_to_filenames[prefix].append(filename)
                    break

        return file_type_to_filenames

    def sample_train_files(self, percents, no):
        sampled_files = []
        for file_type, filenames in self._file_type_to_filenames.items():
            sorted_filenames = sorted(filenames)
            sampled_files.extend(random.choices(sorted_filenames[:-2], k=int(math.ceil(percents[file_type] * no)) ))

        random.shuffle(sampled_files)
        return sampled_files

    def sample_validation_files(self, percents, no):
        sampled_files = []
        for file_type, filenames in self._file_type_to_filenames.items():
            sorted_filenames = sorted(filenames)
            sampled_files.extend(random.choices(sorted_filenames[-2:], k=int(math.ceil(percents[file_type] * no)) ))
            
        return sampled_files

    #individual tree ids are masked with odd numbers
    #group tree ids are masked with even numbers
    def _get_image_data_by_filename(self):
        H, W = tc.IMAGE_HEIGHT, tc.IMAGE_WIDTH

        with open(tc.TRAIN_ANNOTATIONS_PATH, 'r') as file:
            data = json.load(file)

        filename_to_img_data = {}
        images_data = data[tc.IMAGES_KEY]

        for image_data in images_data:
            filename = image_data[tc.FILENAME_KEY]
            img_arr = tc.get_train_image_arr(filename)

            indivi_tree_instance_id = 1
            grp_tree_instance_id = 2
            mask = np.zeros((H, W), np.int32)

            annotations = image_data[tc.ANNOTATIONS_KEY]
            
            for annotation in annotations:
                polygon = tc.get_polygon(annotation[tc.SEGMENTATION_KEY])

                if annotation[tc.CLASS_KEY] == tc.CLASS_INDIVIDUAL_TREE:
                    instance_id = indivi_tree_instance_id
                    indivi_tree_instance_id += 2
                else:
                    instance_id = grp_tree_instance_id
                    grp_tree_instance_id += 2
                
                cv2.fillPoly(mask, [polygon], instance_id) # type: ignore

            filename_to_img_data[filename] = (img_arr, mask)

        return filename_to_img_data

    def get_image_data(self, filename):
        return self._filename_to_img_data[filename]
    

class DynamicDataset(Dataset):

    def __init__(self, no, image_data_cache):
        self._no = no
        self._img_augmenter = ImageAugmenter()
        self._image_data_cache = image_data_cache

        self._percents = {tc.PREFIX_10CM_FILE_TYPE: 0.15, tc.PREFIX_20CM_FILE_TYPE: 0.15,
                          tc.PREFIX_40CM_FILE_TYPE: 0.20, tc.PREFIX_60CM_FILE_TYPE: 0.25, tc.PREFIX_80CM_FILE_TYPE: 0.25}

        self._train_datasets = []

    def __len__(self):
        return self._no

    def on_epoch_start(self):
        self._train_datasets = self._image_data_cache.sample_train_files(self._percents, self._no)

    def __getitem__(self, indx):
        filename = self._train_datasets[indx]
        img_arr, mask = self._image_data_cache.get_image_data(filename)
        aug_img_arr, aug_mask = self._img_augmenter.augment_image(img_arr, mask)

        model_inputs = get_model_input(aug_img_arr, aug_mask, indx)

        return model_inputs
    

class StaticDataset(Dataset):
    def __init__(self, no, image_data_cache):
        self._no = no

        self._image_augmenter = ImageAugmenter()

        self._percents = {
            tc.PREFIX_10CM_FILE_TYPE: 0.20, tc.PREFIX_20CM_FILE_TYPE: 0.20,
            tc.PREFIX_40CM_FILE_TYPE: 0.20, tc.PREFIX_60CM_FILE_TYPE: 0.20,
            tc.PREFIX_80CM_FILE_TYPE: 0.20
        }

        self._sampled_files = image_data_cache.sample_validation_files(self._percents, no)

        image_id = 1
        self._sampled_datasets = []
        for filename in self._sampled_files:
            img_arr, mask = image_data_cache.get_image_data(filename)
            cropped_image, cropped_mask = self._image_augmenter.random_crop_image(img_arr, mask)

            self._sampled_datasets.append((cropped_image, cropped_mask, image_id))

            image_id += 1

    def __len__(self):
        return self._no

    def __getitem__(self, indx):
        img_arr, mask, image_id = self._sampled_datasets[indx]
        return get_model_input(img_arr, mask, image_id)
    

def collate(datasets):
    
    img_arrs = []
    masks = []
    arr_instance_id_to_class_id = []
    image_ids = []
    for dataset in datasets:
        img_arrs.append(dataset['image'])
        masks.append(dataset['mask'])
        arr_instance_id_to_class_id.append(dataset['instance_id_to_class_id'])
        image_ids.append(dataset['image_id'])
        
    return img_arrs, masks, arr_instance_id_to_class_id, image_ids


def get_pretrained_model(device):
    ckpt_dir = 'facebook/mask2former-swin-tiny-coco-instance'
    cfg = Mask2FormerConfig.from_pretrained(ckpt_dir, num_labels=2, id2label=tc.CLASS_ID_TO_LABEL, label2id=tc.LABEL_TO_CLASS_ID)
    
    cfg.backbone_config.drop_path_rate = 0.02 # type: ignore
    cfg.num_queries = 350
    cfg.backbone_config.attention_probs_dropout_prob = 0.025 # type: ignore
    cfg.backbone_config.hidden_dropout_prob = 0.025 # type: ignore
    cfg.backbone_config.dropout = 0.025 # type: ignore
    cfg.backbone_config.id2label = tc.CLASS_ID_TO_LABEL # type: ignore
    cfg.backbone_config.label2id = tc.LABEL_TO_CLASS_ID # type: ignore
    cfg.image_size = tc.CROPPED_IMAGE_HEIGHT
    cfg.class_weight = 5.0
    return Mask2FormerForUniversalSegmentation.from_pretrained(ckpt_dir, config = cfg, ignore_mismatched_sizes=True).to(device) # type: ignore

def get_image_processor():
    image_processor = Mask2FormerImageProcessor.from_pretrained('facebook/mask2former-swin-tiny-coco-instance')
    image_processor.size = {'height' : tc.CROPPED_IMAGE_HEIGHT, 'width': tc.CROPPED_IMAGE_WIDTH}
    image_processor.num_labels = 3
    image_processor.do_resize = False
    image_processor.ignore_index = 255
    return image_processor

    

def checkpoint_state(epoch, model, optimizer):
    directory = tc.MASK2FORMER_CHECKPOINT_DIR
    checkpoint = {'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict()}

    new_file_name = f'checkpoint-{epoch}.pt'
    old_file_name = f'checkpoint-{epoch - 1}.pt'

    
    for item in os.listdir(directory):
        item_path = os.path.join(directory, item)
        if os.path.isfile(item_path) and item.startswith('checkpoint-') and item != old_file_name:
            os.remove(item_path)

    file_path = os.path.join(directory, new_file_name)

    torch.save(checkpoint, file_path)

def setup_state_from_checkpoint(model, optimizer):
    directory = tc.MASK2FORMER_CHECKPOINT_DIR
    latest_epoch = -1
    for item in os.listdir(directory):
        item_path = os.path.join(directory, item)
        if os.path.isfile(item_path) and item.startswith('checkpoint-'):
            epoch = int(item.removeprefix('checkpoint-').removesuffix('.pt'))
            latest_epoch = max(epoch, latest_epoch)

    if latest_epoch == -1:
        print('no existing checkpoint found')
        return 0

    print(f'resuming states with {latest_epoch} checkpoint')
    checkpoint_pt = torch.load(os.path.join(directory, f'checkpoint-{latest_epoch}.pt'), map_location='cuda')
    model.load_state_dict(checkpoint_pt['model_state'])
    if 'optimizer_state' in checkpoint_pt:
        optimizer.load_state_dict(checkpoint_pt['optimizer_state'])
    

    best_val = math.inf
    if os.path.isfile(tc.MASK2FORMER_TRAIN_BEST_WEIGHT_LOC):
        ckpt = torch.load(tc.MASK2FORMER_TRAIN_BEST_WEIGHT_LOC, map_location='cpu')
        best_val = ckpt['best_val']

    return latest_epoch + 1, best_val

def load_latest_model_weights(model):
    directory = tc.MASK2FORMER_CHECKPOINT_DIR
    latest_epoch = -1
    for item in os.listdir(directory):
        item_path = os.path.join(directory, item)
        if os.path.isfile(item_path) and item.startswith('checkpoint-'):
            latest_epoch = int(item.removeprefix('checkpoint-').removesuffix('.pt'))

    if latest_epoch == -1:
        return

    checkpoint_pt = torch.load(os.path.join(directory, f'checkpoint-{latest_epoch}.pt'), map_location='cuda')
    model.load_state_dict(checkpoint_pt["model_state"])


class ValidationDatasetMetric:

    def __init__(self, model, validation_dataset, image_processor):
        self.model = model
        self.validation_dataset = validation_dataset
        batch_size = 1
        self.validation_dataset_dataloader = DataLoader(validation_dataset, batch_size=batch_size, in_order=False, collate_fn=collate, pin_memory=True)
        self.image_processor = image_processor


    def get_prediction_loss(self):
        self.model.eval()
        self.model.to('cuda')

        step = 0
        loss_sum = 0
        with torch.no_grad():
            for (img_arrs, masks, arr_instance_id_to_class_id, _) in self.validation_dataset_dataloader:

                processed = self.image_processor(images=img_arrs, segmentation_maps=masks,
                                instance_id_to_semantic_id=arr_instance_id_to_class_id, return_tensors="pt").data

                model_inputs = {
                                  'pixel_values': processed['pixel_values'].to('cuda'),
                                  'pixel_mask': processed['pixel_mask'].to('cuda'),
                                  'class_labels': [label.to('cuda') for label in processed["class_labels"]],  
                                  'mask_labels':  [label.to('cuda') for label in processed["mask_labels"]],
                                }

                outputs = self.model(**model_inputs)
                loss = outputs.loss
                loss_sum+= loss.item()
                step+= 1

                del outputs, processed

        avg_loss = (loss_sum / step if step > 0 else float("nan"))
        return avg_loss


    def evaluate(self):
        validation_loss = self.get_prediction_loss()
        return validation_loss


from typing import cast
def train(model, train_dataset, validation_dataset, image_processor):

    batch_size = 1
    dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate, in_order=False, 
                        pin_memory=True)
    
    validation_dataset_metric = ValidationDatasetMetric(model, validation_dataset, image_processor)

    names = []
    for indx, (name, _) in enumerate(model.named_parameters()):
        names.append(f'{name} {indx}')


    param_indices = list(range(0, len(names)))
    trainable_params = []
    for indx, param in enumerate(model.parameters()):
        param.requires_grad = False
        if indx in param_indices:
            param.requires_grad = True
            trainable_params.append(param)

    optimizer = torch.optim.AdamW(trainable_params, lr=2.5e-5, weight_decay=1e-4)
    best_val = math.inf

    resume = False

    start_epoch = 1
    if resume:
        start_epoch, best_val = setup_state_from_checkpoint(model, optimizer)  # type: ignore

    optimizer.param_groups[0]['lr'] = 1e-5

    end_epoch = start_epoch + 50
    accumulation_steps = 1
    
    total_instances = len(dataloader)
    for epoch in range(start_epoch, end_epoch):

        model.train()
        optimizer.zero_grad()
        loop = tqdm(dataloader, total=len(dataloader), desc=f"Epoch {epoch}")
        start = time.perf_counter()

        step = 0
        loss_sum = 0
        
        cast(DynamicDataset, dataloader.dataset).on_epoch_start()
        for (img_arrs, masks, arr_instance_id_to_class_id, _) in loop:

            processed = image_processor(images = img_arrs, segmentation_maps = masks,
                                    instance_id_to_semantic_id = arr_instance_id_to_class_id, return_tensors = "pt").data

            model_inputs = {
                            'pixel_values': processed['pixel_values'].to('cuda'),
                            'pixel_mask': processed['pixel_mask'].to('cuda'),
                            'class_labels': [label.to('cuda') for label in processed["class_labels"]],  
                            'mask_labels':  [label.to('cuda') for label in processed["mask_labels"]],
                            }
            
            output = model(**model_inputs)

            loss_sum+= output.loss.item()
            loss = output.loss / accumulation_steps
            loss.backward()

            if (step + 1) % accumulation_steps == 0 or (step + 1) == int(math.ceil(total_instances/batch_size)):
                gradient_before_clipping = clip_grad_norm_(trainable_params, max_norm=1.5, foreach=True)
                optimizer.step()
                optimizer.zero_grad()
                loop.set_postfix({
                        'gradient': gradient_before_clipping.item(),
                        'step': step,
                        'loss' : loss.item(),
                        'learning_rates': '[' + ','.join(
                            [str(param_group['lr']) for param_group in optimizer.state_dict()['param_groups']]) + ']'
                    }
                )

            step += 1

            del model_inputs, output, loss
        
        avg_loss = loss_sum / step
        validation_loss = validation_dataset_metric.evaluate()
        if validation_loss < best_val:
            torch.save({'model_state': model.state_dict(), 'best_val': validation_loss}, tc.MASK2FORMER_TRAIN_BEST_WEIGHT_LOC)

        
        end = time.perf_counter()
        print(f'{{ epoch {epoch} avg_loss: {avg_loss} took {end - start:.3f} seconds }}')
        checkpoint_state(epoch, model, optimizer)

from sahi import DetectionModel
from typing import Any
from sahi.prediction import ObjectPrediction
from sahi.utils.compatibility import fix_full_shape_list, fix_shift_amount_list
from sahi.predict import get_sliced_prediction
class Mask2FormerSahi(DetectionModel):
    
    def __init__(
        self,
        model_path: str | None = None,
        model: Any | None = None,
        processor: Any | None = None,
        config_path: str | None = None,
        device: str | None = None,
        mask_threshold: float = 0.5,
        confidence_threshold: float = 0.3,
        category_mapping: dict | None = None,
        category_remapping: dict | None = None,
        load_at_init: bool = True,
        image_size: int | None = None,
        token: str | None = None,
    ):
        self._processor = processor
        self._image_shapes: list = []
        self._token = token
        
        super().__init__(
            model_path,
            model,
            config_path,
            device,
            mask_threshold,
            confidence_threshold,
            category_mapping,
            category_remapping,
            load_at_init,
            image_size,
        )

    @property
    def processor(self):
        return self._processor

    @property
    def image_shapes(self):
        return self._image_shapes

    @property
    def num_categories(self) -> int:
        return self.model.config.num_labels

    def load_model(self):
        self.set_model(self.model, self._processor)

    def set_model(self, model: Any, processor: Any = None, **kwargs):
        self.model = model
        self.model.to(self.device)
        self.category_mapping = self.model.config.id2label

    def perform_inference(self, image: list | np.ndarray):
        import torch

        if self.model is None or self.processor is None:
            raise RuntimeError(f'{self.model is None} {self.processor is None}')

        with torch.no_grad():
            inputs = self.processor(images=image, return_tensors="pt")
            inputs["pixel_values"] = inputs.pixel_values.to(self.device)
            if hasattr(inputs, "pixel_mask"):
                inputs["pixel_mask"] = inputs.pixel_mask.to(self.device)
            
            outputs = self.model(**inputs)

        if isinstance(image, list):
            self._image_shapes = [(img.shape[0], img.shape[1]) for img in image]
        else:
            self._image_shapes = [(image.shape[0], image.shape[1])]

        n_image = len(self._image_shapes)
        target_sizes = self._image_shapes
        
        post_processed_outputs = self.processor.post_process_instance_segmentation(
                                outputs,threshold=self.confidence_threshold, mask_threshold=self.mask_threshold, 
                                return_binary_maps=True, target_sizes=target_sizes)
        
        self._original_predictions = post_processed_outputs
    
    def perform_batch_inference(self, images: list[np.ndarray]) -> None:
        return self.perform_inference(images)

    def get_polygonal_predictions(self, post_processed_output) -> tuple:
       
        scores=[]
        cat_ids=[]
        #returns polygons as list of lists where inner list is flattened as [x1, y1, ...,xn,yn]
        polygonal_segments=[]
      
        segments = post_processed_output['segmentation']
        segments_info = post_processed_output['segments_info']
    
        for segment, segment_info in zip(segments, segments_info):
            mask = segment.cpu().numpy().astype(np.uint8)
            score = segment_info['score']
            lbl_indx = segment_info['label_id']
            class_id = tc.INDEX_TO_CLASS_ID[lbl_indx]
            
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            if contours:
                polygonal_segments.append([contours[0].squeeze(1).flatten().tolist()])
                    
            scores.append(score)
            cat_ids.append(class_id)
                
        return scores, cat_ids, polygonal_segments

    #peroforms batch inference on the list of slices
    #shit amount [[shift_x, shift_y],...] gives how much each slice must be shifted
    #full_shape [[width, height]] gives each slices dimension
    def _create_object_prediction_list_from_original_predictions(
        self,
        shift_amount_list: list[list[int]] | None = [[0, 0]],
        full_shape_list: list[list[int]] | None = None,
    ):
        
        original_predictions = self._original_predictions

        shift_amount_list = fix_shift_amount_list(shift_amount_list)
        full_shape_list = fix_full_shape_list(full_shape_list)

        n_image = len(original_predictions)
        object_prediction_list_per_image = []
        for image_ind in range(n_image):
            image_height, image_width  = self.image_shapes[image_ind]
            scores, cat_ids, segments = self.get_polygonal_predictions(original_predictions[image_ind])

            object_prediction_list = []

            shift_amount = shift_amount_list[image_ind]
            full_shape = None if full_shape_list is None else full_shape_list[image_ind]
            
            #iterate each polygonal segment
            for ind in range(len(segments)):
                category_id = cat_ids[ind]
                segment = segments[ind]
                score = scores[ind]
                
                #8 represents the numer of x,y coords of a polygon
                #we need atleast8 to have a polygon
                if len(segment[0]) >= 8:
                    object_prediction = ObjectPrediction(
                        bbox=None,
                        segmentation=segment,
                        category_id=category_id,
                        category_name=self.category_mapping[category_id],
                        shift_amount=shift_amount,
                        score=score,
                        full_shape=full_shape,
                    )
                    object_prediction_list.append(object_prediction)
            object_prediction_list_per_image.append(object_prediction_list)

        self._object_prediction_list_per_image = object_prediction_list_per_image


def predict(model, image_processor, run_on_a_subset_of_eval_images):
    state = torch.load(tc.MASK2FORMER_INFERENCE_BEST_WEIGHT_LOC, map_location='cuda')
    model.load_state_dict(state['model_state'])
    sahi_model = Mask2FormerSahi(model=model, processor=image_processor, 
                            mask_threshold=0.75, confidence_threshold=0.80, 
                            device='cuda', image_size=tc.CROPPED_IMAGE_HEIGHT)
    files = []
    for filename in os.listdir(tc.EVALUATION_IMAGES_PATH):
        if filename.startswith('60cm') or not run_on_a_subset_of_eval_images:
            files.append(filename)

    files = files[:10]

    results = []
    for indx, filename in enumerate(files):
        print(f'remaining is {len(files) - indx}')
        filepath = os.path.join(tc.EVALUATION_IMAGES_PATH, filename)

        result = get_sliced_prediction(
            image=filepath,
            detection_model=sahi_model,
            slice_height=tc.CROPPED_IMAGE_HEIGHT,        
            slice_width=tc.CROPPED_IMAGE_WIDTH,         
            overlap_height_ratio=0.2,
            overlap_width_ratio=0.2, 
            perform_standard_pred = False
        )
        results.append(result)
    
    for filename, result in zip(files, results):
        annotations = []
        for ann in result.to_coco_predictions():
            annotation = {tc.CLASS_KEY:ann['category_name'], 'confidence_score': ann['score'], tc.SEGMENTATION_KEY:ann['segmentation'][0]}
            annotations.append(annotation)
        
        image_data = {tc.ANNOTATIONS_KEY :  annotations}
        tc.show_img_with_annotations(tc.get_evaluation_image_arr(filename), image_data)
    


if __name__ == '__main__':
    model:Mask2FormerForUniversalSegmentation = get_pretrained_model('cuda')
    image_data_cache = ImageDataCache()
    train_dataset = DynamicDataset(1900, image_data_cache)
    validation_dataset = StaticDataset(30, image_data_cache)

    model.state_dict()['criterion.empty_weight'][1] = 5

    image_processor = get_image_processor()

    train(model, train_dataset,validation_dataset, image_processor)

    # import matplotlib.pyplot as plt
    # plt.rcParams['figure.max_open_warning'] = 200

    # predict(model, image_processor, run_on_a_subset_of_eval_images= True)