add demo

app.py

@@ -0,0 +1,76 @@
+import cv2
+import numpy as np
+import math
+import torch
+import random
+
+from torch.utils.data import DataLoader
+
+torch.manual_seed(12345)
+random.seed(12345)
+np.random.seed(12345)
+
+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
+
+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)
+
+import gradio as gr
+
+
+def count_barnacles(input_img, progress=gr.Progress()):
+    progress(0, desc="Loading Image")
+    test_dataset = get_dataset_x(input_img)
+    test_dataloader = DataLoader(test_dataset, batch_size=256, shuffle=False)
+    model.eval()
+    predicted_labels_list = []
+    for data in progress.tqdm(test_dataloader):
+        with torch.no_grad():
+            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)
+    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 = map(rev_window_transform, windows)
+
+    blank_img_copy = input_img.copy()
+    for x, y in points:
+        blank_img_copy = cv2.circle(blank_img_copy, (x,y), radius=4, color=(255, 0, 0), thickness=-1)
+    return blank_img_copy
+
+demo = gr.Interface(count_barnacles, gr.Image(shape=(500, 500), type="numpy"), gr.Image(type="numpy"))
+demo.queue(concurrency_count=10).launch()

requirements.txt

@@ -0,0 +1,5 @@
+opencv-python
+numpy
+torch
+torchvision
+gradio