Update app.py

add old segmentation

app.py

@@ -11,6 +11,8 @@ from torchvision.ops import box_convert
 from torchvision.transforms.functional import to_tensor
 from torchvision.transforms import GaussianBlur
 
+from Ambrosia import pre_process_image
+
 
 # Define a custom transform for Gaussian blur
 def gaussian_blur(x, p=0.5, kernel_size_min=3, kernel_size_max=20, sigma_min=0.1, sigma_max=3):
@@ -55,62 +57,62 @@ def unkown_prob_calc(probs, wedge_threshold, wedge_magnitude=1, wedge='strict'):
     unknown_prob = 1-kown_prob
     return(unknown_prob)
 
-def load_image(image_source):
-    transform = T.Compose(
-        [
-            T.RandomResize([800], max_size=1333),
-            T.ToTensor(),
-            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-        ]
-    )
-    image_source = image_source.convert("RGB")
+# def load_image(image_source):
+#     transform = T.Compose(
+#         [
+#             T.RandomResize([800], max_size=1333),
+#             T.ToTensor(),
+#             T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+#         ]
+#     )
+#     image_source = image_source.convert("RGB")
     
-    image_transformed, _ = transform(image_source, None)
-    return image_transformed
-
-# load object detection model
-od_model = load_model(
-    model_checkpoint_path="groundingdino_swint_ogc.pth", 
-    model_config_path="GroundingDINO_SwinT_OGC.cfg.py", 
-    device="cpu")
-print("Object detection model loaded")
-
-def detect_objects(og_image, model=od_model, prompt="bug . insect", device="cpu"):
-    TEXT_PROMPT = prompt
-    BOX_TRESHOLD = 0.35
-    TEXT_TRESHOLD = 0.25
-    DEVICE = device  # cuda or cpu
+#     image_transformed, _ = transform(image_source, None)
+#     return image_transformed
+
+# # load object detection model
+# od_model = load_model(
+#     model_checkpoint_path="groundingdino_swint_ogc.pth", 
+#     model_config_path="GroundingDINO_SwinT_OGC.cfg.py", 
+#     device="cpu")
+# print("Object detection model loaded")
+
+# def detect_objects(og_image, model=od_model, prompt="bug . insect", device="cpu"):
+#     TEXT_PROMPT = prompt
+#     BOX_TRESHOLD = 0.35
+#     TEXT_TRESHOLD = 0.25
+#     DEVICE = device  # cuda or cpu
     
-    # Convert numpy array to PIL Image if needed
-    if isinstance(og_image, np.ndarray):
-        og_image_obj = Image.fromarray(og_image)
-    else:
-        og_image_obj = og_image  # Assuming og_image is already a PIL Image
-
-    # Transform the image
-    image_transformed = load_image(image_source = og_image_obj)
+#     # Convert numpy array to PIL Image if needed
+#     if isinstance(og_image, np.ndarray):
+#         og_image_obj = Image.fromarray(og_image)
+#     else:
+#         og_image_obj = og_image  # Assuming og_image is already a PIL Image
+
+#     # Transform the image
+#     image_transformed = load_image(image_source = og_image_obj)
     
-    # Your model prediction code here...
-    boxes, logits, phrases = grounding_dino_predict(
-        model=model,
-        image=image_transformed,
-        caption=TEXT_PROMPT,
-        box_threshold=BOX_TRESHOLD,
-        text_threshold=TEXT_TRESHOLD,
-        device=DEVICE)
-
-    # Use og_image_obj directly for further processing
-    height, width = og_image_obj.size
-    boxes_norm = boxes * torch.Tensor([height, width, height, width])
-    xyxy = box_convert(
-        boxes=boxes_norm,
-        in_fmt="cxcywh",
-        out_fmt="xyxy").numpy()
-    img_lst = []
-    for i in range(len(boxes_norm)):
-        crop_img = og_image_obj.crop((xyxy[i]))
-        img_lst.append(crop_img)
-    return (img_lst)
+#     # Your model prediction code here...
+#     boxes, logits, phrases = grounding_dino_predict(
+#         model=model,
+#         image=image_transformed,
+#         caption=TEXT_PROMPT,
+#         box_threshold=BOX_TRESHOLD,
+#         text_threshold=TEXT_TRESHOLD,
+#         device=DEVICE)
+
+#     # Use og_image_obj directly for further processing
+#     height, width = og_image_obj.size
+#     boxes_norm = boxes * torch.Tensor([height, width, height, width])
+#     xyxy = box_convert(
+#         boxes=boxes_norm,
+#         in_fmt="cxcywh",
+#         out_fmt="xyxy").numpy()
+#     img_lst = []
+#     for i in range(len(boxes_norm)):
+#         crop_img = og_image_obj.crop((xyxy[i]))
+#         img_lst.append(crop_img)
+#     return (img_lst)
 
 
 # load beetle classifier model
@@ -123,7 +125,11 @@ print("Classification model loaded")
 def predict_beetle(img):
     print("Detecting & classifying beetles...")
     # Split image into smaller images of detected objects
-    image_lst = detect_objects(og_image=img, model=od_model, prompt="bug . insect", device="cpu")
+    # image_lst = detect_objects(og_image=img, model=od_model, prompt="bug . insect", device="cpu")
+    pre_process = pre_process_image(manual_thresh_buffer=0.15, image = img) # use image_dir if directory of image used
+    pre_process.segment(cluster_num=2, 
+                        image_edge_buffer=50)
+    image_lst = pre_process.col_image_lst
     print("Objects detected")
     # get predictions for all segments
     conf_dict_lst = []