adding yolov8s model as well

.gitignore

@@ -1,2 +1,3 @@
 env
-__pycache__
+__pycache__
+gradio_cached_examples

README.md

@@ -23,6 +23,9 @@ The target species were:
 - **Culiseta** - Genus
 - **Aedes japonicus/Aedes koreicus** - Species complex (Differentiating between these two species is particularly challenging).
 
+> ***Note:** Only one mosquito will be annotated even if there are multiple mosquitoes in the image.*
+
+
 ## Experiment Details
 
 All the details regarding the experiments and source code for the models can be found in the [GitHub repository](https://github.com/HCA97/Mosquito-Classifiction/tree/main).

app.py

@@ -4,37 +4,44 @@ import gradio as gr
 import numpy as np
 import cv2
 
-from my_models import YOLOV5CLIPModel
+from my_models import YOLOV5CLIPModel, YOLOV8CLIPModel
 
 
 def annotated_image(
     image: np.ndarray, label: str, conf: float, bbox: list
 ) -> np.ndarray:
-    
     line_thickness = int(0.005 * max(image.shape[:2]))
     image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    image = cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (255, 0, 0), thickness=line_thickness)
-    image = cv2.putText(image, 
-                        f"{label} {conf:.2f}", 
-                        (bbox[0], max(bbox[1] - 2*line_thickness, 0)),
-                        cv2.FONT_HERSHEY_SIMPLEX,
-                        thickness=max(line_thickness//2, 1),
-                        lineType=cv2.LINE_AA,
-                        color=(0, 0, 0),
-                        fontScale=0.1*line_thickness) 
+    image = cv2.rectangle(
+        image,
+        (bbox[0], bbox[1]),
+        (bbox[2], bbox[3]),
+        (255, 0, 0),
+        thickness=line_thickness,
+    )
+    image = cv2.putText(
+        image,
+        f"{label} {conf:.2f}",
+        (bbox[0], max(bbox[1] - 2 * line_thickness, 0)),
+        cv2.FONT_HERSHEY_SIMPLEX,
+        thickness=max(line_thickness // 2, 1),
+        lineType=cv2.LINE_AA,
+        color=(0, 0, 0),
+        fontScale=0.1 * line_thickness,
+    )
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
     return image
 
 
 def detect_mosquito(image):
-    label, conf, bbox = YOLOV5CLIPModel().predict(image)
+    label, conf, bbox = YOLOV8CLIPModel().predict(image)
     return annotated_image(image, label, conf, bbox)
 
 
 description = """# [Mosquito Alert Competition 2023](https://www.aicrowd.com/challenges/mosquitoalert-challenge-2023) - 7th Place Solution
 
-Welcome to my Hugging Face Space showcasing the performance of our model.
+Welcome to my Hugging Face Space showcasing the performance of our model. 
 
 This competition focused on detecting and classifying various mosquito species. 
 
@@ -46,14 +53,20 @@ The target species were:
 - **Culiseta** - Genus
 - **Aedes japonicus/Aedes koreicus** - Species complex (Differentiating between these two species is particularly challenging).
 
+> ***Note:** Only one mosquito will be annotated even if there are multiple mosquitoes in the image.*
+
 ## Experiment Details
 
 All the details regarding the experiments and source code for the models can be found in the [GitHub repository](https://github.com/HCA97/Mosquito-Classifiction/tree/main).
 """
 
 iface = gr.Interface(
-    fn=detect_mosquito, description=description, inputs=gr.Image(), outputs=gr.Image(), allow_flagging="never", 
+    fn=detect_mosquito,
+    description=description,
+    inputs=gr.Image(),
+    outputs=gr.Image(),
+    allow_flagging="never",
     examples=[os.path.join("examples", f) for f in os.listdir("examples")],
-    cache_examples=True
+    cache_examples=True,
 )
 iface.launch()

my_models/__init__.py

@@ -1 +1,2 @@
 from .yolov5_clip_model import YOLOV5CLIPModel
+from .yolov8_clip_model import YOLOV8CLIPModel

my_models/yolo_weights/best-yolov8-s.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36080c806f3a8b501bc52c126d68427c61dd15b8dcff1423ae35163588a09583
+size 22484974

my_models/yolov8_clip_model.py

@@ -0,0 +1,135 @@
+from ultralytics import YOLO
+import numpy as np
+import time
+
+import torch
+
+torch.set_num_threads(2)
+
+from my_models.clip_model.data_loader import pre_process_foo
+from my_models.clip_model.classification import MosquitoClassifier
+
+IMG_SIZE = (224, 224)
+USE_CHANNEL_LAST = False
+DATASET = "laion"
+DEVICE = "cpu"
+PRESERVE_ASPECT_RATIO = False
+SHIFT = 0
+
+
+@torch.no_grad()
+def classify_image(det: YOLO, cls: MosquitoClassifier, image: np.ndarray):
+    s = time.time()
+    labels = [
+        "albopictus",
+        "culex",
+        "japonicus-koreicus",
+        "culiseta",
+        "anopheles",
+        "aegypti",
+    ]
+
+    results = det(image, verbose=True, device=DEVICE, max_det=1)
+    img_w, img_h, _ = image.shape
+    bbox = [0, 0, img_w, img_h]
+    label = "albopictus"
+    conf = 0.0
+
+    for result in results:
+        _bbox = [0, 0, img_w, img_h]
+        _label = "albopictus"
+        _conf = 0.0
+
+        bboxes_tmp = result.boxes.xyxy.tolist()
+        labels_tmp = result.boxes.cls.tolist()
+        confs_tmp = result.boxes.conf.tolist()
+
+        for bbox_tmp, label_tmp, conf_tmp in zip(bboxes_tmp, labels_tmp, confs_tmp):
+            if conf_tmp > _conf:
+                _bbox = bbox_tmp
+                _label = labels[int(label_tmp)]
+                _conf = conf_tmp
+
+        if _conf > conf:
+            bbox = _bbox
+            label = _label
+            conf = _conf
+
+    bbox = [int(float(mcb)) for mcb in bbox]
+
+    try:
+        if conf < 1e-4:
+            raise Exception
+        image_cropped = image[bbox[1] : bbox[3], bbox[0] : bbox[2], :]
+        bbox = [bbox[0] + SHIFT, bbox[1] + SHIFT, bbox[2] - SHIFT, bbox[3] - SHIFT]
+    except Exception as e:
+        print("Error", e)
+        image_cropped = image
+
+    if PRESERVE_ASPECT_RATIO:
+        w, h = image_cropped.shape[:2]
+        if w > h:
+            x = torch.unsqueeze(
+                pre_process_foo(
+                    (IMG_SIZE[0], max(int(IMG_SIZE[1] * h / w), 32)), DATASET
+                )(image_cropped),
+                0,
+            )
+        else:
+            x = torch.unsqueeze(
+                pre_process_foo(
+                    (max(int(IMG_SIZE[0] * w / h), 32), IMG_SIZE[1]), DATASET
+                )(image_cropped),
+                0,
+            )
+    else:
+        x = torch.unsqueeze(pre_process_foo(IMG_SIZE, DATASET)(image_cropped), 0)
+
+    x = x.to(device=DEVICE)
+
+    if USE_CHANNEL_LAST:
+        p = cls(x.to(memory_format=torch.channels_last))
+    else:
+        p = cls(x)
+    ind = torch.argmax(p).item()
+    label = labels[ind]
+
+    e = time.time()
+
+    print("Time ", 1000 * (e - s), "ms")
+    return {"name": label, "confidence": p.max().item(), "bbox": bbox}
+
+
+# getting mosquito_class name from predicted result
+def extract_predicted_mosquito_class_name(extractedInformation):
+    return extractedInformation.get("name", "albopictus")
+
+
+def extract_predicted_mosquito_bbox(extractedInformation):
+    return extractedInformation.get("bbox", [0, 0, 0, 0])
+
+
+class YOLOV8CLIPModel:
+    def __init__(self):
+        trained_model_path = "my_models/yolo_weights/best-yolov8-s.pt"
+        clip_model_path = f"my_models/clip_weights/best_clf.ckpt"
+        self.det = YOLO(trained_model_path, task="detect")
+
+        self.cls = MosquitoClassifier.load_from_checkpoint(
+            clip_model_path, head_version=7, map_location=torch.device(DEVICE)
+        ).eval()
+
+        if USE_CHANNEL_LAST:
+            self.cls.to(memory_format=torch.channels_last)
+
+    def predict(self, image):
+        predictedInformation = classify_image(self.det, self.cls, image)
+
+        mosquito_class_name_predicted = extract_predicted_mosquito_class_name(
+            predictedInformation
+        )
+        mosquito_class_bbox = extract_predicted_mosquito_bbox(predictedInformation)
+
+        bbox = bbox = [int(float(mcb)) for mcb in mosquito_class_bbox]
+
+        return mosquito_class_name_predicted, predictedInformation["confidence"], bbox