Update app.py

app.py

@@ -1,32 +1,16 @@
-import os
-import dill
-import timm
 import random
 import numpy as np
 import gradio as gr
-from fastai.tabular.all import *
-from fastai.vision.all import *
-from fastai.vision.utils import get_image_files
-from fastai.learner import load_learner
-from Ambrosia import pre_process_image
-from huggingface_hub import from_pretrained_fastai, push_to_hub_fastai, hf_hub_download
+from huggingface_hub import from_pretrained_fastai
+from PIL import Image
+from groundingdino.util.inference import load_model
+from groundingdino.util.inference import predict as grounding_dino_predict
+import groundingdino.datasets.transforms as T
+import torch
+from torchvision.ops import box_convert
+from torchvision.transforms.functional import to_tensor
 from torchvision.transforms import GaussianBlur
 
-# Set the token
-huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-if huggingface_token is None:
-    raise ValueError("Hugging Face token not found. Please set the HUGGINGFACE_TOKEN environment variable.")
-
-
-# # Define a custom transform for Gaussian blur
-# def gaussian_blur(x, p=0.5, kernel_size_min=3, kernel_size_max=9, sigma_min=0.1, sigma_max=2):
-#     if x.ndim == 4:
-#         for i in range(x.shape[0]):
-#             if random.random() < p:
-#                 kernel_size = random.randrange(kernel_size_min, kernel_size_max + 1, 2)
-#                 sigma = random.uniform(sigma_min, sigma_max)
-#                 x[i] = GaussianBlur(kernel_size=kernel_size, sigma=sigma)(x[i])
-#     return x
 
 # 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):
@@ -38,6 +22,12 @@ def gaussian_blur(x, p=0.5, kernel_size_min=3, kernel_size_max=20, sigma_min=0.1
                 x[i] = GaussianBlur(kernel_size=kernel_size, sigma=sigma)(x[i])
     return x
 
+# Custom Label Function
+def custom_label_func(fpath):
+    # this directs the labels to be 2 levels up from the image folder
+    label = fpath.parents[2].name
+    return label
+
 # this function only describes how much a singular value in al ist stands out.
 # if all values in the lsit are high or low this is 1
 # the smaller the proportiopn of number of disimilar vlaues are to other more similar values the lower this number
@@ -65,26 +55,78 @@ def unkown_prob_calc(probs, wedge_threshold, wedge_magnitude=1, wedge='strict'):
     unknown_prob = 1-kown_prob
     return(unknown_prob)
 
-# load model
-learn = from_pretrained_fastai("ChristopherMarais/beetle-model")
-# learn = load_learner(
-#     hf_hub_download('ChristopherMarais/Andrew_Alpha_model', filename="model.pkl")
-#     )
+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")
+
+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)
     
+    # 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
+repo_id="ChristopherMarais/beetle-model"
+bc_model = from_pretrained_fastai(repo_id)
 # get class names
-labels = np.append(np.array(learn.dls.vocab), "Unknown")
-
-def predict(img):
-    # Segment image into smaller images
-    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)
-    # get predictions for all segments
+labels = np.append(np.array(bc_model.dls.vocab), "Unknown")
+
+def predict_beetle(img):
+    # Split image into smaller images of detected objects
+    image_lst = detect_objects(og_image=img, model=od_model, prompt="bug . insect", device="cpu")
+        # get predictions for all segments
     conf_dict_lst = []
     output_lst = []
-    img_cnt = len(pre_process.col_image_lst)
+    img_cnt = len(image_lst)
     for i in range(0,img_cnt):
-        prob_ar = np.array(learn.predict(pre_process.col_image_lst[i])[2])
+        prob_ar = np.array(bc_model.predict(image_lst[i])[2])
         unkown_prob = unkown_prob_calc(probs=prob_ar, wedge_threshold=0.85, wedge_magnitude=5, wedge='dynamic')
         prob_ar = np.append(prob_ar, unkown_prob)
         prob_ar = np.around(prob_ar*100, decimals=1)
@@ -92,19 +134,20 @@ def predict(img):
         conf_dict = {labels[i]: float(prob_ar[i]) for i in range(len(prob_ar))}
         conf_dict = dict(sorted(conf_dict.items(), key=lambda item: item[1], reverse=True))
         conf_dict_lst.append(str(conf_dict))
-        result = list(zip(pre_process.col_image_lst, conf_dict_lst))
-                
+        result = list(zip(image_lst, conf_dict_lst))
     return(result)
 
+
+# gradio app
 with gr.Blocks() as demo:
     with gr.Column(variant="panel"):
         with gr.Row(variant="compact"):
             inputs = gr.Image()
-            btn = gr.Button("Classify").style(full_width=False)
+            btn = gr.Button("Classify")
 
         gallery = gr.Gallery(
             label="Show images", show_label=True, elem_id="gallery"
-        ).style(grid=[8], height="auto")
+        )
 
-    btn.click(predict, inputs, gallery)
-    demo.launch()
+    btn.click(predict_beetle, inputs, gallery)
+demo.launch()