Spaces:
Runtime error
Runtime error
File size: 3,119 Bytes
5afcf8b fc1a0c8 5afcf8b fc1a0c8 5afcf8b fc1a0c8 5afcf8b fc1a0c8 5afcf8b fc1a0c8 5afcf8b fc1a0c8 5afcf8b fc1a0c8 5afcf8b fc1a0c8 5afcf8b fc1a0c8 5afcf8b fc1a0c8 5afcf8b fc1a0c8 5afcf8b fc1a0c8 5afcf8b fc1a0c8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
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)
model.to("cuda")
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=1024, shuffle=False)
model.eval()
predicted_labels_list = []
for data in progress.tqdm(test_dataloader):
with torch.no_grad():
data.to("cuda")
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, len(list(points))
demo = gr.Interface(count_barnacles, gr.Image(shape=(500, 500), type="numpy"),
outputs=[gr.Image(type="numpy"), "number"])
demo.queue(concurrency_count=10).launch()
|