Refactor dataset adding Dataset class

Author: jrrodri

abraia/training/__init__.py

@@ -1,89 +1,24 @@
-
-from ..client import Abraia
-from ..utils import url_path, make_dirs
-from . import classify, detect
-
 import os
 import shutil
 import itertools
+
 from PIL import Image
-from tqdm.contrib.concurrent import process_map
-from sklearn.model_selection import train_test_split
 from typing import Dict, Any
+from tqdm.contrib.concurrent import process_map
 
-
-abraia = Abraia()
-
-
-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 list_images(project):
-    files = abraia.list_files(f"{project}/")[0]
-    dataset = [f for f in files if f['type'] in ['image/jpeg', 'image/png']]
-    for data in dataset:
-        data['url'] = url_path(f"{abraia.userid}/{data['path']}")
-    return dataset
-
-
-def load_annotations(project):
-    annotations = abraia.load_json(f"{project}/annotations.json")
-    for annotation in annotations:
-        annotation['path'] = f"{project}/{annotation['filename']}"
-        annotation['url'] = url_path(f"{abraia.userid}/{annotation['path']}")
-    return annotations
-
-
-def save_annotations(project, annotations):
-    abraia.save_json(f"{project}/annotations.json", annotations)
-
-
-def load_labels(annotations):
-    labels = []
-    for annotation in annotations:
-        for object in annotation.get('objects', []):
-            label = object.get('label')
-            if label and label not in labels:
-                labels.append(label)
-    return list(set(labels))
+from . import classify, detect
+from .dataset import list_datasets, load_dataset, search_images, list_models
+from ..utils import make_dirs, download_file
 
 
-def load_task(annotations):
-    classify, detect, segment = False, False, False
-    for annotation in annotations:
-        for object in annotation.get('objects', []):
-            if 'polygon' in object:
-                segment = True
-            elif 'box' in object:
-                detect = True
-            elif 'label' in object:
-                classify = True
-    if segment:
-        return 'segment'
-    if detect:
-        return 'detect'
-    if classify:
-        return 'classify'
-
-
-def load_tasks(annotations):
+def load_tasks(task):
     tasks = ['classify', 'detect', 'segment']
-    task = load_task(annotations)
     if task:
         idx = tasks.index(task)
         return tasks[:idx+1]
     return []
 
 
-def download_file(path, folder):
-    dest = os.path.join(folder, os.path.basename(path))
-    if not os.path.exists(dest):
-        abraia.download_file(path, dest)
-    return dest
-
-
 def save_annotation(annotation, folder, classes, task):
     if task == 'classify':
         for object in annotation.get('objects', []):
@@ -146,19 +81,10 @@ def save_config(dataset, classes):
         f.write(yaml_content)
 
 
-def split_dataset(annotations):
-    # TODO: Split dataset by classes to avoid class imbalance
-    backgrounds = [annotation for annotation in annotations if not annotation.get('objects')]
-    annotations = [annotation for annotation in annotations if annotation.get('objects')]
-    train, test = train_test_split(annotations, test_size=0.3)
-    val, test = train_test_split(test, test_size=0.5)
-    train.extend(backgrounds)
-    return train, val, test
-
-
 class ModelTrainer:
     """High-level trainer orchestrator using models and dataset utilities."""
-    def __init__(self, task: str, imgsz: int = 640):
+    def __init__(self, task: str, imgsz: int = None):
+        imgsz = imgsz or (224 if task == 'classify' else 640)
         self.task = task
         if task == 'classify':
             self.model = classify.Model()
@@ -167,15 +93,15 @@ def __init__(self, task: str, imgsz: int = 640):
 
     def prepare_dataset(self, project: str, classes, force: bool = False):
         if force or not os.path.exists(project):
-            annotations = load_annotations(project)
-            train, val, test = split_dataset(annotations)
+            dataset = load_dataset(project)
+            train, val, test = dataset.split()
             data_annotations = {'train': train, 'val': val, 'test': test}
-            #TODO: Download files in one single step
             for x in ['train', 'val', 'test']:
                 save_data(data_annotations[x], f"{project}/{x}", classes, self.task)
             save_config(project, classes)
 
-    def train(self, dataset: str, epochs: int, batch: int = 32, **kwargs) -> None:
+    def train(self, dataset: str, epochs: int = None, batch: int = 32) -> None:
+        epochs = epochs or (30 if self.task == 'classify' else 300)
         if self.task == 'classify':
             self.model.train(dataset, epochs=epochs)
         else:
@@ -193,4 +119,3 @@ def save(self, dataset: str, classes, device='cpu') -> None:
 
     def run(self, img):
         return self.model.run(img)
-    

abraia/training/dataset.py

@@ -9,8 +9,13 @@
 
 from tqdm import tqdm
 from PIL import Image
-from ..utils import HEADERS, load_image, load_url, list_dir
-from . import list_datasets, list_images, load_annotations, save_annotations
+from sklearn.model_selection import train_test_split
+
+from ..client import Abraia
+from ..utils import HEADERS, load_image, load_url, list_dir, url_path
+from ..inference.detect import segment_objects
+
+abraia = Abraia()
 
 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'''
@@ -24,7 +29,7 @@ def download_page(url):
     return resp.text
 
 
-def save_image(link, output_dir, timeout=10, max_size=1920):
+def save_image_file(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'):
@@ -39,27 +44,27 @@ def save_image(link, output_dir, timeout=10, max_size=1920):
 
 
 def get_filter(shorthand):
-        if shorthand == "line" or shorthand == "linedrawing":
-            return "+filterui:photo-linedrawing"
-        elif shorthand == "photo":
-            return "+filterui:photo-photo"
-        elif shorthand == "clipart":
-            return "+filterui:photo-clipart"
-        elif shorthand == "gif" or shorthand == "animatedgif":
-            return "+filterui:photo-animatedgif"
-        elif shorthand == "transparent":
-            return "+filterui:photo-transparent"
-        else:
-            return ""
-        
+    if shorthand == "line" or shorthand == "linedrawing":
+        return "+filterui:photo-linedrawing"
+    elif shorthand == "photo":
+        return "+filterui:photo-photo"
+    elif shorthand == "clipart":
+        return "+filterui:photo-clipart"
+    elif shorthand == "gif" or shorthand == "animatedgif":
+        return "+filterui:photo-animatedgif"
+    elif shorthand == "transparent":
+        return "+filterui:photo-transparent"
+    else:
+        return ""
+
 
 def scan_bing_page(html):
     links = re.findall('murl":"(.*?)"', html)
     for link in links:
         link = link.replace(" ", "%20")
         yield link
 
-        
+
 def search_bing(query, limit=50, adult='off', filter=''):
     for page_counter in range(100):
         # Parse the page source and download pics
@@ -106,7 +111,7 @@ def download(query, limit=100, output_dir='dataset', verbose=True):
                 seen.add(link)
                 if download_count < limit:
                     try:
-                        save_image(link, output_dir)
+                        save_image_file(link, output_dir)
                         download_count += 1
                         if verbose:
                             print(f"[%] Downloaded Image #{download_count} from {link}")
@@ -119,12 +124,20 @@ def download(query, limit=100, output_dir='dataset', verbose=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)
     return list_dir(output_dir)
 
 
+def download_file(path, folder):
+    dest = os.path.join(folder, os.path.basename(path))
+    if not os.path.exists(dest):
+        abraia.download_file(path, dest)
+    return dest
+
+
 # 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()
@@ -149,17 +162,116 @@ def annotate_image(pipe, img, classes, threshold=0.3):
     return objects
 
 
-def annotate_images(images, classes):
+def annotate_images(images, classes, segment=False):
     """Annotate a dataset using Grounding Dino."""
     from transformers import pipeline
     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:
+            if segment:
+                try:
+                    objects = segment_objects(img, objects)
+                except:
+                    continue
             annotation = {'url': url, 'filename': filename, 'objects': objects}
             annotations.append(annotation)
     return annotations
+
+
+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 list_models(project):
+    files = abraia.list_files(f"{project}/")[0]
+    return [f['name'] for f in files if f['name'].endswith('.onnx')]
+
+
+def load_annotations(project):
+    annotations = abraia.load_json(f"{project}/annotations.json")
+    for annotation in annotations:
+        annotation['path'] = f"{project}/{annotation['filename']}"
+        annotation['url'] = url_path(f"{abraia.userid}/{annotation['path']}")
+    return annotations
+
+
+def load_labels(annotations):
+    labels = []
+    for annotation in annotations:
+        for object in annotation.get('objects', []):
+            label = object.get('label')
+            if label and label not in labels:
+                labels.append(label)
+    return list(set(labels))
+
+
+def load_task(annotations):
+    classify, detect, segment = False, False, False
+    for annotation in annotations:
+        for object in annotation.get('objects', []):
+            if 'polygon' in object:
+                segment = True
+            elif 'box' in object:
+                detect = True
+            elif 'label' in object:
+                classify = True
+    return 'segment' if segment else 'detect' if detect else 'classify' if classify else ''
+
+
+def list_images(project):
+    files = abraia.list_files(f"{project}/")[0]
+    files = [f for f in files if f['type'] in ['image/jpeg', 'image/png']]
+    for data in files:
+        data['url'] = url_path(f"{abraia.userid}/{data['path']}")
+    return files
+
+
+def save_annotations(project, annotations):
+    abraia.save_json(f"{project}/annotations.json", annotations)
+
+
+class Dataset:
+    def __init__(self, project):
+        self.project = project
+        self.annotations = []
+        self.classes = []
+        self.task = ''
+        self.images = []
+
+    def load(self):
+        if self.project in list_datasets():
+            self.annotations = load_annotations(self.project)
+            self.classes = load_labels(self.annotations)
+            self.task = load_task(self.annotations)
+            self.images = list_images(self.project)
+        return self
+    
+    def annotate(self, label, segment=False):
+        annotated_filenames = {a['filename'] for a in self.annotations}
+        images = [img for img in self.images if img['name'] not in annotated_filenames]
+        new_annotations = annotate_images(images, [label], segment=segment)
+        self.annotations.extend(new_annotations)
+        return self.annotations
+
+    def save(self, annotations=None):
+        if annotations is not None:
+            self.annotations = annotations
+        save_annotations(self.project, self.annotations)
+
+    def split(self):
+        # TODO: Split dataset by classes to avoid class imbalance
+        backgrounds = [annotation for annotation in self.annotations if not annotation.get('objects')]
+        annotations = [annotation for annotation in self.annotations if annotation.get('objects')]
+        train, test = train_test_split(annotations, test_size=0.3)
+        val, test = train_test_split(test, test_size=0.5)
+        train.extend(backgrounds)
+        return train, val, test
+
+        
+def load_dataset(project):
+    return Dataset(project).load()