add gpu support

app.py

@@ -10,22 +10,27 @@ 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()
+    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
 
@@ -33,25 +38,27 @@ 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)
+    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 = 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'])
+        cx = int(M['m10'] / M['m00'])
+        cy = int(M['m01'] / M['m00'])
         return cx, cy
 
     filter_width = 50
@@ -59,8 +66,8 @@ def count_barnacles(input_img, progress=gr.Progress()):
 
     def rev_window_transform(point):
         wx, wy = point
-        x = int(filter_width/2) + wx*filter_stride
-        y = int(filter_width/2) + wy*filter_stride
+        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)
@@ -69,8 +76,10 @@ def count_barnacles(input_img, progress=gr.Progress()):
 
     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
+        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"), gr.Image(type="numpy"))
-demo.queue(concurrency_count=10).launch()
+demo = gr.Interface(count_barnacles, gr.Image(shape=(500, 500), type="numpy"),
+                    outputs=[gr.Image(type="numpy"), "number"])
+demo.queue(concurrency_count=10).launch()

requirements.txt

@@ -1,5 +1,6 @@
 opencv-python
 numpy
+--extra-index-url https://download.pytorch.org/whl/cu113
 torch
 torchvision
 gradio