Update app.py

app.py

@@ -58,62 +58,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
@@ -127,11 +127,13 @@ def predict_beetle(img):
     print("Detecting & classifying beetles...")
     start_time = time.perf_counter() # record how long it processes
     # Split image into smaller images of detected objects
-    # 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
+    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")
     end_time = time.perf_counter()
     processing_time = end_time - start_time