Add dataset to training

Author: jrrodri

abraia/__init__.py

@@ -2,7 +2,7 @@
 from dotenv import load_dotenv
 load_dotenv()
 
-__version__ = '0.24.1'
+__version__ = '0.24.2'
 
 from . import config
 from .client import Abraia, APIError

abraia/training/__init__.py

@@ -1,27 +1,25 @@
 
 from ..client import Abraia
-from ..utils import get_color, url_path, make_dirs, load_image, load_url
+from ..utils import url_path, make_dirs
 from . import classify, detect
 
 import os
 import shutil
 import itertools
-from tqdm import tqdm
 from PIL import Image
-from transformers import pipeline
 from tqdm.contrib.concurrent import process_map
 from sklearn.model_selection import train_test_split
 
 
 abraia = Abraia()
 
 
-def load_projects():
+def list_datasets():
     folders = abraia.list_files()[1]
     return [folder['name'] for folder in folders if abraia.check_file(f"{folder['name']}/annotations.json")]
 
 
-def load_dataset(project):
+def list_images(project):
     files = abraia.list_files(f"{project}/")[0]
     dataset = [f for f in files if f['type'].startswith('image/')]
     for data in dataset:
@@ -165,42 +163,3 @@ def create_dataset(dataset, task, classes):
         for x in ['train', 'val', 'test']:
             save_data(data_annotations[x], f"{dataset}/{x}", classes, task)
         save_config(dataset, classes)
-
-
-# As the Grounding DINO model was trained with a "." after each text, we'll do the same here.
-def preprocess_caption(caption: str) -> str:
-    result = caption.lower().strip()
-    return result if result.endswith('.') else result + '.'
-
-
-def format_results(results, threshold=0.6):
-    r = []
-    for result in results:
-        score = result["score"]
-        if score > threshold:
-            label = result["label"].rpartition('.')[0]
-            xmin, ymin, xmax, ymax = result['box'].values()
-            r.append({"label": label, "score": score, "box": [xmin, ymin, xmax - xmin, ymax - ymin]})
-    return r
-
-
-def annotate_image(pipe, img, classes, threshold=0.3):
-    im = Image.fromarray(img)
-    labels = [preprocess_caption(txt) for txt in classes]
-    objects = format_results(pipe(im, candidate_labels=labels, threshold=threshold))
-    return objects    
-
-
-def annotate_images(project, classes):
-    dataset = load_dataset(project)
-    pipe = pipeline(task="zero-shot-object-detection", model="IDEA-Research/grounding-dino-tiny")
-    #pipe = pipeline(task="zero-shot-object-detection", model="google/owlv2-base-patch16-ensemble")
-    annotations = []
-    for row in tqdm(dataset):
-        url, filename = row['url'], row['name']
-        img = load_image(load_url(url))
-        objects = annotate_image(pipe, img, classes)
-        if objects:
-            annotation = {'url': url, 'filename': filename, 'objects': objects}
-            annotations.append(annotation)
-    save_annotations(project, annotations)

abraia/utils/dataset.py abraia/training/dataset.pyabraia/utils/dataset.py renamed to abraia/training/dataset.py

@@ -7,8 +7,11 @@
 import itertools
 import imagehash
 
+from tqdm import tqdm
 from PIL import Image
-from . import HEADERS
+from transformers import pipeline
+from ..utils import HEADERS, load_image, load_url
+from . import list_datasets, list_images, load_annotations, save_annotations
 
 GOOGLE_BASE_URL = 'https://www.google.com/search?q='
 GOOGLE_PICTURE_ID = '''&biw=1536&bih=674&tbm=isch&sxsrf=ACYBGNSXXpS6YmAKUiLKKBs6xWb4uUY5gA:1581168823770&source=lnms&sa=X&ved=0ahUKEwioj8jwiMLnAhW9AhAIHbXTBMMQ_AUI3QUoAQ'''
@@ -22,7 +25,7 @@ def download_page(url):
     return resp.text
 
 
-def save_image(link, file_path, timeout=10, max_size=1920):
+def save_image(link, output_dir, timeout=10, max_size=1920):
     resp = requests.get(link, headers=HEADERS, allow_redirects=True, timeout=timeout)
     kind = filetype.guess(resp.content)
     if kind and kind.mime.startswith('image'):
@@ -31,7 +34,7 @@ def save_image(link, file_path, timeout=10, max_size=1920):
         im = Image.open(d).convert('RGB')
         im.thumbnail([max_size, max_size])
         phash = str(imagehash.phash(im))
-        im.save(os.path.join(os.path.dirname(file_path), phash + '.jpg'))
+        im.save(os.path.join(output_dir, phash + '.jpg'))
     else:
         raise ValueError(f'Invalid image, not saving')
 
@@ -104,8 +107,7 @@ def download(query, limit=100, output_dir='dataset', verbose=True):
                 seen.add(link)
                 if download_count < limit:
                     try:
-                        filename = f"{'-'.join(query.split(' '))}_{download_count}.jpg"
-                        save_image(link, os.path.join(output_dir, filename))
+                        save_image(link, output_dir)
                         download_count += 1
                         if verbose:
                             print(f"[%] Downloaded Image #{download_count} from {link}")
@@ -117,3 +119,48 @@ def download(query, limit=100, output_dir='dataset', verbose=True):
             ends[id] = True
             if set(ends) == {True}:
                 break
+
+def search_images(query, limit=100, output_dir='dataset', verbose=True):
+    """Search and download images from Google and Bing."""
+    download(query, limit=limit, output_dir=output_dir,  verbose=verbose)
+    files = [os.path.join(output_dir, f) for f in os.listdir(output_dir) if os.path.isfile(os.path.join(output_dir, f))]
+    return files
+
+
+# As the Grounding DINO model was trained with a "." after each text, we'll do the same here.
+def preprocess_caption(caption: str) -> str:
+    result = caption.lower().strip()
+    return result if result.endswith('.') else result + '.'
+
+
+def format_results(results, threshold=0.6):
+    r = []
+    for result in results:
+        score = result["score"]
+        if score > threshold:
+            label = result["label"].rpartition('.')[0]
+            xmin, ymin, xmax, ymax = result['box'].values()
+            r.append({"label": label, "score": score, "box": [xmin, ymin, xmax - xmin, ymax - ymin]})
+    return r
+
+
+def annotate_image(pipe, img, classes, threshold=0.3):
+    im = Image.fromarray(img)
+    labels = [preprocess_caption(txt) for txt in classes]
+    objects = format_results(pipe(im, candidate_labels=labels, threshold=threshold))
+    return objects
+
+
+def annotate_images(images, classes):
+    """Annotate a dataset using Grounding Dino."""
+    pipe = pipeline(task="zero-shot-object-detection", model="IDEA-Research/grounding-dino-tiny")
+    #pipe = pipeline(task="zero-shot-object-detection", model="google/owlv2-base-patch16-ensemble")
+    annotations = []
+    for row in tqdm(images):
+        url, filename = row['url'], row['name']
+        img = load_image(load_url(url))
+        objects = annotate_image(pipe, img, classes)
+        if objects:
+            annotation = {'url': url, 'filename': filename, 'objects': objects}
+            annotations.append(annotation)
+    return annotations

abraia/training/detect.py

@@ -1,6 +1,8 @@
 from ..client import Abraia
 
 import os
+import io
+import contextlib
 import numpy as np
 #os.environ['YOLO_VERBOSE'] = 'False'
 
@@ -29,18 +31,31 @@ def __init__(self, task, model_type='yolov8n'):
         model_name = build_model_name(model_type, task)
         self.model = YOLO(f"{model_name}.pt", verbose=False)
         self.model_name = model_name
+        self.metrics = {}
         self.task = task
 
     def train(self, dataset, epochs=100, batch=32, imgsz=640):
         data = f"{dataset}" if self.task == 'classify' else f"{dataset}/data.yaml"
         results = self.model.train(data=data, batch=batch, epochs=epochs, imgsz=imgsz)
+        # TODO: Merge with test and add parse metrics
         metrics = self.model.val(data=data)
+        # self.metrics = self.test(split='val')
         return metrics
 
+    def test(self, split='val'):
+        out = io.StringIO()
+        with contextlib.redirect_stderr(out):
+            metrics = self.model.val(split=split)
+        return {'mAP': metrics.box.map50, 'P': metrics.box.p, 'R': metrics.box.r, 
+                'confusionMatrix': metrics.confusion_matrix.matrix}
+
     def save(self, dataset, classes, imgsz=640, device="cpu"):
+        # TODO: Add versioning
         model_src = self.model.export(format="onnx", device=device)
         abraia.upload_file(model_src, f"{dataset}/{self.model_name}.onnx")
-        abraia.save_json(f"{dataset}/{self.model_name}.json", {'task': self.task, 'inputShape': [1, 3, imgsz, imgsz], 'classes': classes})
+        abraia.save_json(f"{dataset}/{self.model_name}.json", 
+                         {'task': self.task, 'inputShape': [1, 3, imgsz, imgsz], 
+                          'classes': classes, 'metrics': self.metrics})
 
     def run(self, img):
         objects = []

abraia/utils/draw.py

@@ -45,8 +45,20 @@ def draw_rectangle(img, rect, color, thickness = 2):
     return img
 
 
-def draw_filled_rectangle(img, rect, color):
-    return draw_rectangle(img, rect, color, -1)
+# def draw_filled_rectangle(img, rect, color):
+#     return draw_rectangle(img, rect, color, -1)
+
+
+def draw_filled_rectangle(img, rect, color, opacity = 1):
+    x, y, w, h = np.round(rect).astype(np.int32)
+    pt1, pt2 = (x, y), (x + w, y + h)
+    if opacity == 1:
+        cv2.rectangle(img, pt1, pt2, color, -1)
+    else:
+        img_copy = img.copy()
+        cv2.rectangle(img_copy, pt1, pt2, color, -1)
+        cv2.addWeighted(img_copy, opacity, img, 1 - opacity, 0, img)
+    return img
 
 
 def draw_polygon(img, polygon, color, thickness = 2):
@@ -55,12 +67,39 @@ def draw_polygon(img, polygon, color, thickness = 2):
     return img
 
 
+# def draw_filled_polygon(img, polygon, color, opacity = 1):
+#     points = np.round(polygon).astype(np.int32)
+#     cv2.fillPoly(img, [points], color)
+#     return img
+
+
 def draw_filled_polygon(img, polygon, color, opacity = 1):
     points = np.round(polygon).astype(np.int32)
-    cv2.fillPoly(img, [points], color)
+    if opacity == 1:
+        cv2.fillPoly(img, [points], color)
+    else:
+        img_copy = img.copy()
+        cv2.fillPoly(img_copy, [points], color)
+        cv2.addWeighted(img_copy, opacity, img, 1 - opacity, 0, img)
     return img
 
 
+# def draw_blurred_mask(img, mask):
+#     w_k = int(0.1 * max(img.shape[:2]))
+#     w_k = w_k + 1 if w_k % 2 == 0 else w_k
+#     blurred_img = cv2.GaussianBlur(img, (w_k, w_k), 0)
+#     mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
+#     img = np.where(mask==0, img, blurred_img)
+#     return img
+
+
+# def draw_blurred_polygon(img, polygon):
+#     points = np.round(polygon).astype(np.int32)
+#     mask = np.zeros(img.shape, dtype=np.uint8)
+#     mask = cv2.fillPoly(mask, [points], 255)
+#     return draw_blurred_mask(img, mask)
+
+
 def draw_blurred_mask(img, mask):
     w_k = int(0.1 * max(img.shape[:2]))
     w_k = w_k + 1 if w_k % 2 == 0 else w_k
@@ -82,16 +121,24 @@ def draw_overlay_mask(img, mask, color = (255, 0, 0), opacity = 1):
     return img_copy
 
 
+def calculate_contrast_text_color(background_color):
+    r, g, b = background_color
+    brightness = (r * 299 + g * 587 + b * 114) / 1000
+    return (0, 0, 0) if brightness > 150 else (255, 255, 255)
+
+
 def draw_text(img, text, point, background_color = None, text_color = (255, 255, 255), 
               text_scale = 0.8, padding = 6):
-    text_font, text_thickness = cv2.FONT_HERSHEY_DUPLEX, 1
-    w, h = cv2.getTextSize(text, text_font, text_scale, text_thickness)[0]
-    width, height = w + 2 * padding, h + 2 * padding
-    x, y = point[0], max(point[1] - height, 0)
-    org = (x + padding, y + padding + h)
-    if background_color is not None:
-        img = draw_filled_rectangle(img, [x, y, width, height], background_color)
-    cv2.putText(img, text, org, text_font, text_scale, text_color, text_thickness, cv2.LINE_AA)
+    if text:
+        text_font, text_thickness = cv2.FONT_HERSHEY_DUPLEX, 1
+        w, h = cv2.getTextSize(text, text_font, text_scale, text_thickness)[0]
+        width, height = w + 2 * padding, h + 2 * padding
+        x, y = point[0], max(point[1] - height, 0)
+        org = (x + padding, y + padding + h)
+        if background_color is not None:
+            img = draw_filled_rectangle(img, [x, y, width, height], background_color)
+            text_color = calculate_contrast_text_color(background_color)
+        cv2.putText(img, text, org, text_font, text_scale, text_color, text_thickness, cv2.LINE_AA)
     return img
 
 
@@ -147,10 +194,11 @@ def render_counter(img, line, text='', color=(0, 0, 255)):
     return img
 
 
-def render_region(img, region, text='', color=(0, 0, 0)):
+def render_region(img, region, text='', color=(255, 0, 0)):
     point = np.min(region, axis=0).astype(np.int32)
     thickness = calculate_optimal_thickness(img.shape[:2])
     text_scale = calculate_optimal_text_scale(img.shape[:2])
     draw_polygon(img, region, color=color, thickness=thickness)
+    draw_filled_polygon(img, region, color=color, opacity=0.2)
     draw_text(img, text, point, background_color=color, text_scale=text_scale)
     return img