Spaces:
Runtime error
Runtime error
| 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() | |