import mmcv
from mmdet.apis import inference_detector, init_detector
import cv2
import numpy as np
import pandas as pd
import os
import random

# Load the already trained model profile and weights
# config_file = 'C:/Users/18863662902/Desktop/mmdetection/work_dirs/yolact_r50_1xb8-55e_coco/yolact_r50_1xb8-55e_coco.py'
# checkpoint_file = 'C:/Users/18863662902/Desktop/mmdetection/work_dirs/yolact_r50_1xb8-55e_coco/epoch_400.pth'
# model = init_detector(config_file, checkpoint_file)
config_file = 'work_dirs/yolact_r50_1xb8-55e_coco/yolact_r50_1xb8-55e_coco.py'
checkpoint_file = 'work_dirs/yolact_r50_1xb8-55e_coco/epoch_400.pth'
model = init_detector(config_file, checkpoint_file, device='cpu')  # or device='cuda:0'
# Define the path and size of the image to be tested
image_folder = 'C:/Users/18863662902/Desktop/sum1203/'
# image_folder = 'C:/Users/18863662902/Desktop/val2017'
image_height = 720
image_width = 1280
real_height_cm = 15
real_width_cm = 27

# Calculate the true length corresponding to each pixel
pixel_size = real_width_cm / image_width

# Iterate over all images in the images folder
for filename in os.listdir(image_folder):
    if filename.endswith('.jpg') or filename.endswith('.png'):
        # Constructing the image path
        image_path = os.path.join(image_folder, filename)

        # Read image
        image = mmcv.imread(image_path)

        # Make inferences and get results
        result = inference_detector(model, image_path)
        bboxes = result.pred_instances.bboxes
        scores = result.pred_instances.scores.cpu().numpy()
        masks = result.pred_instances.masks
        labels = result.pred_instances.labels + 1

        # Filtering by scores
        index = scores > 0.3
        bboxes = bboxes[index].cpu().numpy()
        scores = scores[index]
        masks = masks[index].cpu().numpy()
        labels = labels[index]

        class_names = ['1', '2', '3', '4']  # Customize the category list according to your needs

        # Statistical Spot Information
        num_spots = 0
        spot_info = []
        for i, mask in enumerate(masks):
            if labels[i] == 1:  # Assume that the object corresponds to a category label of 1
                num_spots += 1
                # Calculate the pixel area and the true area of the spot
                pixel_area = np.sum(mask)
                real_area = pixel_area * (pixel_size ** 2)
                mask = mask.astype(np.uint8)

                # Get the bounding box coordinates of the spot
                contours, _ = cv2.findContours(mask, mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_SIMPLE)

                if len(contours) > 0:
                    for contour in contours:
                        # Convert contours to polygonal approximations
                        epsilon = 0.001 * cv2.arcLength(contour, True)
                        approx = cv2.approxPolyDP(contour, epsilon, True)  # Polygon Approximation with epsilon

                        # Generate random colors
                        color = tuple(random.randint(0, 255) for _ in range(3))

                        # Draw polygons and mark spots using random colors
                        cv2.fillPoly(image, [approx], color=color)
                        cv2.polylines(image, [approx], isClosed=True, color=color, thickness=1)

                    x, y, w, h = cv2.boundingRect(contours[0])

                    # Get the center point coordinates of the spot
                    center_x = x + w // 2
                    center_y = y + h // 2

                    # Label the pictures with serial numbers
                    cv2.putText(image, str(num_spots), (center_x - 3, center_y), cv2.FONT_HERSHEY_SIMPLEX, 0.6,
                                (0, 0, 0), 2)

                    # Add spot information to the list
                    spot_info.append([num_spots, pixel_area, real_area, center_x, center_y, x, y, w, h])
                else:
                    # Handling of contour not found
                    # Option to skip the current spot or take other actions
                    pass

        # Save statistics as CSV
        columns = ['Spot ID', 'Pixel Area', 'Real Area', 'Center X', 'Center Y', 'X', 'Y', 'Width', 'Height']
        spot_info_df = pd.DataFrame(spot_info, columns=columns)
        csv_filename = f'spot_info_{filename}.csv'
        csv_path = os.path.join('C:/Users/18863662902/Desktop/sum1203', csv_filename)
        spot_info_df.to_csv(csv_path, index=False)
        print(f"Spot information has been saved to{csv_filename}file.")

        # Adjusting the size of the displayed interface
        display_width = 1280
        display_height = 720
        resized_image = cv2.resize(image, (display_width, display_height))

        # Saving processed images
        image_output_path = os.path.join('C:/Users/18863662902/Desktop/sum1203', f'resized_{filename}')
        cv2.imwrite(image_output_path, resized_image)
        print("The processed images have been saved to", image_output_path)

        # Show picture
        cv2.imshow("Resized Image", resized_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()