app.py

import cv2
import numpy as np
import math
import torch
import random

from torch.utils.data import DataLoader
from torchvision.transforms import Resize

torch.manual_seed(12345)
random.seed(12345)
np.random.seed(12345)

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

class WireframeExtractor:

    def __call__(self, image: np.ndarray):
        """
        Extract corners of wireframe from a barnacle image
        :param image: Numpy RGB image of shape (W, H, 3)
        :return [x1, y1, x2, y2]
        """
        h, w = image.shape[:2]
        imghsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
        hsvblur = cv2.GaussianBlur(imghsv, (9, 9), 0)

        lower = np.array([70, 20, 20])
        upper = np.array([130, 255, 255])

        color_mask = cv2.inRange(hsvblur, lower, upper)

        invert = cv2.bitwise_not(color_mask)

        contours, _ = cv2.findContours(invert, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

        max_contour = contours[0]
        largest_area = 0
        for index, contour in enumerate(contours):
            area = cv2.contourArea(contour)
            if area > largest_area:
                if cv2.pointPolygonTest(contour, (w / 2, h / 2), False) == 1:
                    largest_area = area
                    max_contour = contour

        x, y, w, h = cv2.boundingRect(max_contour)
        # return [x, y, x + w, y + h]
        return x,y,w,h

wireframe_extractor = WireframeExtractor()

def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)
    polygons = []
    color = []
    for ann in sorted_anns:
        m = ann['segmentation']
        img = np.ones((m.shape[0], m.shape[1], 3))
        color_mask = np.random.random((1, 3)).tolist()[0]
        for i in range(3):
            img[:,:,i] = color_mask[i]
        ax.imshow(np.dstack((img, m*0.35)))
    

# def find_contours(img, color):
#     low = color - 10
#     high = color + 10

#     mask = cv2.inRange(img, low, high)
#     contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

#     print(f"Total Contours: {len(contours)}")
#     nonempty_contours = list()
#     for i in range(len(contours)):
#         if hierarchy[0,i,3] == -1 and cv2.contourArea(contours[i]) > cv2.arcLength(contours[i], True):
#             nonempty_contours += [contours[i]]
#     print(f"Nonempty Contours: {len(nonempty_contours)}")
#     contour_plot = img.copy()
#     contour_plot = cv2.drawContours(contour_plot, nonempty_contours, -1, (0,255,0), -1)

#     sorted_contours = sorted(nonempty_contours, key=cv2.contourArea, reverse= True)

#     bounding_rects = [cv2.boundingRect(cnt) for cnt in contours]

#     for (i,c) in enumerate(sorted_contours):
#         M= cv2.moments(c)
#         cx= int(M['m10']/M['m00'])
#         cy= int(M['m01']/M['m00'])
#         cv2.putText(contour_plot, text= str(i), org=(cx,cy),
#                 fontFace= cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.25, color=(255,255,255),
#                 thickness=1, lineType=cv2.LINE_AA)

#     N = len(sorted_contours)
#     H, W, C = img.shape
#     boxes_array_xywh = [cv2.boundingRect(cnt) for cnt in sorted_contours]
#     boxes_array_corners = [[x, y, x+w, y+h] for x, y, w, h in boxes_array_xywh]
#     boxes = torch.tensor(boxes_array_corners)

#     labels = torch.ones(N)
#     masks = np.zeros([N, H, W])
#     for idx in range(len(sorted_contours)):
#         cnt = sorted_contours[idx]
#         cv2.drawContours(masks[idx,:,:], [cnt], 0, (255), -1)
#     masks = masks / 255.0
#     masks = torch.tensor(masks)

#     # for box in boxes:
#     #     cv2.rectangle(contour_plot, (box[0].item(), box[1].item()), (box[2].item(), box[3].item()), (255,0,0), 2)

#     return contour_plot, (boxes, masks)


# def get_dataset_x(blank_image, filter_size=50, filter_stride=2):
#     full_image_tensor = torch.tensor(blank_image).type(torch.FloatTensor).permute(2, 0, 1).unsqueeze(0)
#     num_windows_h = math.floor((full_image_tensor.shape[2] - filter_size) / filter_stride) + 1
#     num_windows_w = math.floor((full_image_tensor.shape[3] - filter_size) / filter_stride) + 1
#     windows = torch.nn.functional.unfold(full_image_tensor, (filter_size, filter_size), stride=filter_stride).reshape(
#         [1, 3, 50, 50, num_windows_h * num_windows_w]).permute([0, 4, 1, 2, 3]).squeeze()

#     dataset_images = [windows[idx] for idx in range(len(windows))]
#     dataset = list(dataset_images)
#     return dataset


# def get_dataset(labeled_image, blank_image, color, filter_size=50, filter_stride=2, label_size=5):
#     contour_plot, (blue_boxes, blue_masks) = find_contours(labeled_image, color)

#     mask = torch.sum(blue_masks, 0)

#     label_dim = int((labeled_image.shape[0] - filter_size) / filter_stride + 1)
#     labels = torch.zeros(label_dim, label_dim)
#     mask_labels = torch.zeros(label_dim, label_dim, filter_size, filter_size)

#     for lx in range(label_dim):
#         for ly in range(label_dim):
#             mask_labels[lx, ly, :, :] = mask[
#                                         lx * filter_stride: lx * filter_stride + filter_size,
#                                         ly * filter_stride: ly * filter_stride + filter_size
#                                         ]

#     print(labels.shape)
#     for box in blue_boxes:
#         x = int((box[0] + box[2]) / 2)
#         y = int((box[1] + box[3]) / 2)

#         window_x = int((x - int(filter_size / 2)) / filter_stride)
#         window_y = int((y - int(filter_size / 2)) / filter_stride)

#         clamp = lambda n, minn, maxn: max(min(maxn, n), minn)

#         labels[
#         clamp(window_y - label_size, 0, labels.shape[0] - 1):clamp(window_y + label_size, 0, labels.shape[0] - 1),
#         clamp(window_x - label_size, 0, labels.shape[0] - 1):clamp(window_x + label_size, 0, labels.shape[0] - 1),
#         ] = 1

#     positive_labels = labels.flatten() / labels.max()
#     negative_labels = 1 - positive_labels
#     pos_mask_labels = torch.flatten(mask_labels, end_dim=1)
#     neg_mask_labels = 1 - pos_mask_labels
#     mask_labels = torch.stack([pos_mask_labels, neg_mask_labels], dim=1)
#     dataset_labels = torch.tensor(list(zip(positive_labels, negative_labels)))
#     dataset = list(zip(
#         get_dataset_x(blank_image, filter_size=filter_size, filter_stride=filter_stride),
#         dataset_labels,
#         mask_labels
#     ))
#     return dataset, (labels, mask_labels)


# from torchvision.models.resnet import resnet50
# from torchvision.models.resnet import ResNet50_Weights

# print("Loading resnet...")
# model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
# hidden_state_size = model.fc.in_features
# model.fc = torch.nn.Linear(in_features=hidden_state_size, out_features=2, bias=True)
# model.to(device)
# model.load_state_dict(torch.load("model_best_epoch_4_59.62.pth", map_location=torch.device(device)))
# model.to(device)
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

model = sam_model_registry["default"](checkpoint="./sam_vit_h_4b8939.pth")
model.to(device)

predictor = SamPredictor(model)

mask_generator = SamAutomaticMaskGenerator(model)

import gradio as gr

import matplotlib.pyplot as plt
import io

def count_barnacles(image_raw, progress=gr.Progress()):
    progress(0, desc="Finding bounding wire")

    # crop image
    # h, w = raw_input_img.shape[:2]
    # imghsv = cv2.cvtColor(raw_input_img, cv2.COLOR_RGB2HSV)
    # hsvblur = cv2.GaussianBlur(imghsv, (9, 9), 0)

    # lower = np.array([70, 20, 20])
    # upper = np.array([130, 255, 255])

    # color_mask = cv2.inRange(hsvblur, lower, upper)

    # invert = cv2.bitwise_not(color_mask)

    # contours, _ = cv2.findContours(invert, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

    # max_contour = contours[0]
    # largest_area = 0
    # for index, contour in enumerate(contours):
    #     area = cv2.contourArea(contour)
    #     if area > largest_area:
    #         if cv2.pointPolygonTest(contour, (w / 2, h / 2), False) == 1:
    #             largest_area = area
    #             max_contour = contour

    # x, y, w, h = cv2.boundingRect(max_contour)


    image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB)
    corners = wireframe_extractor(image)
    cropped_image = image[corners[1]:corners[3], corners[0]:corners[2], :]
    cropped_image = cropped_image[100:400, 100:400]
    # print(cropped_image)


    # progress(0, desc="Generating Masks by point in window")

    # # get center point of windows
    # predictor.set_image(image)
    # mask_counter = 0
    # masks = []

    # for x in range(1,20, 2):
    #     for y in range(1,20, 2):
    #         point = np.array([[x*25, y*25]])
    #         input_label = np.array([1])
    #         mask, score, logit = predictor.predict(
    #             point_coords=point,
    #             point_labels=input_label,
    #             multimask_output=False,
    #         )
    #         if score[0] > 0.8:
    #             mask_counter += 1
    #             masks.append(mask)

    # return mask_counter

    mask_counter = 0
    good_masks = []
    coords = []
    progress(0, desc="Generating Masks")
    # masks = mask_generator.generate(cropped_image)
    masks = mask_generator.generate(cropped_image)
    for mask in masks:
        if mask['predicted_iou'] > 0.95:
            mask_counter += 1
            good_masks.append(mask)
            coords.append(mask['point_coords'])

    
    # Create a figure with a size of 10 inches by 10 inches
    fig = plt.figure(figsize=(10, 10))

    # Display the image using the imshow() function
    plt.imshow(cropped_image)

    # Call the custom function show_anns() to plot annotations on top of the image
    show_anns(good_masks)

    # Turn off the axis
    plt.axis('off')

    # Get the plot as a numpy array
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
    buf.seek(0)
    img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
    buf.close()

    # Decode the numpy array to an image
    annotated = cv2.imdecode(img_arr, 1)
    annotated = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)

    # Close the figure
    plt.close(fig)
    
    
    # cropped_copy = torch.transpose(cropped_image, 0, 2).to("cpu").detach().numpy().copy()
    return annotated, mask_counter


    # return len(masks)

    # progress(0, desc="Resizing Image")
    # cropped_img = raw_input_img[x:x+w, y:y+h]
    # cropped_image_tensor = torch.transpose(torch.tensor(cropped_img).to(device), 0, 2)
    # resize = Resize((1500, 1500))
    # input_img = cropped_image_tensor
    # blank_img_copy = torch.transpose(input_img, 0, 2).to("cpu").detach().numpy().copy()

    # progress(0, desc="Generating Windows")
    # test_dataset = get_dataset_x(input_img)
    # test_dataloader = DataLoader(test_dataset, batch_size=1024, shuffle=False)
    # model.eval()
    # predicted_labels_list = []
    # for data in progress.tqdm(test_dataloader):
    #     with torch.no_grad():
    #         data = data.to(device)
    #         predicted_labels_list += [model(data)]
    # predicted_labels = torch.cat(predicted_labels_list)
    # x = int(math.sqrt(predicted_labels.shape[0]))
    # predicted_labels = predicted_labels.reshape([x, x, 2]).detach()
    # label_img = predicted_labels[:, :, :1].cpu().numpy()
    # label_img -= label_img.min()
    # label_img /= label_img.max()
    # label_img = (label_img * 255).astype(np.uint8)
    # mask = np.array(label_img > 180, np.uint8)
    # contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)\

    # gt_contours = find_contours(labeled_input_img[x:x+w, y:y+h], cropped_img, np.array([59, 76, 160]))



    # def extract_contour_center(cnt):
    #     M = cv2.moments(cnt)
    #     cx = int(M['m10'] / M['m00'])
    #     cy = int(M['m01'] / M['m00'])
    #     return cx, cy

    # filter_width = 50
    # filter_stride = 2

    # def rev_window_transform(point):
    #     wx, wy = point
    #     x = int(filter_width / 2) + wx * filter_stride
    #     y = int(filter_width / 2) + wy * filter_stride
    #     return x, y

    # nonempty_contours = filter(lambda cnt: cv2.contourArea(cnt) != 0, contours)
    # windows = map(extract_contour_center, nonempty_contours)
    # points = list(map(rev_window_transform, windows))
    # for x, y in points:
    #     blank_img_copy = cv2.circle(blank_img_copy, (x, y), radius=4, color=(255, 0, 0), thickness=-1)
    # print(f"pointlist: {len(points)}")
    # return blank_img_copy, len(points)


demo = gr.Interface(count_barnacles,
                    inputs=[
                        gr.Image(shape=(500, 500), type="numpy", label="Input Image"),
                    ],
                    outputs=[
                        gr.Image(shape=(500, 500), type="numpy", label="Annotated Image"),
                        gr.Number(label="Predicted Number of Barnacles"),
                        # gr.Number(label="Actual Number of Barnacles"),
                        # gr.Number(label="Custom Metric")
                    ])
                    # examples="examples")
demo.queue(concurrency_count=10).launch()