change app.py

app.py

@@ -66,122 +66,6 @@ def show_anns(anns):
             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")
@@ -215,72 +99,14 @@ def check_circularity(segmentation):
 def count_barnacles(image_raw, split_num, 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)
-    # image = Image.fromarray(image_raw)
-    # image = image[:,:,::-1]
-    # image = image_raw
-    # print(image.shape)
-    # print(type(image))
-    # print(image.dtype)
-    # print(image)
     corners = wireframe_extractor(image_raw)
     print(corners) # (0, 0, 1254, 1152)
 
     cropped_image = image_raw[corners[1]:corners[3]+corners[1], corners[0]:corners[2]+corners[0], :]
 
     print(cropped_image.shape)
-    # 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
-    split_num = 2
+    split_num = 5
 
     x_inc = int(cropped_image.shape[0]/split_num)
     y_inc = int(cropped_image.shape[1]/split_num)
@@ -301,23 +127,17 @@ def count_barnacles(image_raw, split_num, progress=gr.Progress()):
             # plt.figure()
             # plt.imshow(small_image)
             # plt.axis('on')
-
+            progress(0, desc=f"Encoding crop {r*split_num + c}/{split_num ** 2}")
+            mask_generator.predictor.set_image(small_image)
+            progress(0, desc=f"Generating masks for crop {r*split_num + c}/{split_num ** 2}")
             masks = mask_generator.generate(small_image)
-
+            num_masks = len(masks)
             
-            for mask in masks:
+            for idx, mask in enumerate(masks):
+                progress(float(idx)/float(num_masks), desc=f"Processing masks for crop {r*split_num + c}/{split_num ** 2}")
                 circular = check_circularity(mask['segmentation'])
                 if circular and mask['area']>500 and mask['area'] < 10000:
                     mask_counter += 1
-                    # if cropped_image.shape != image_raw.shape:
-                    #     add_to_row = [False] * corners[0] 
-                    #     temp = [False]*(corners[2]+corners[0])
-                    #     temp = [temp]*corners[1]
-                    #     new_seg = np.array(temp)
-                    #     for row in mask['segmentation']:
-                    #         row = np.append(add_to_row, row)
-                    #         new_seg = np.vstack([new_seg, row])
-                    #     mask['segmentation'] = new_seg
                     good_masks.append(mask)
                     box = mask['bbox']
                     centers.append((box[0] + box[2]/2 + corners[0] + startx, box[1] + box[3]/2 + corners[1] + starty))
@@ -358,63 +178,6 @@ def count_barnacles(image_raw, split_num, progress=gr.Progress()):
     # return annotated, mask_counter, centers
     return fig, mask_counter, centers
 
-
-    # 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(type="numpy", label="Input Image"),