commit 6604923b2674c3b0d423283aef59769cdbfb9286 Author: freund Date: Tue Aug 23 19:13:13 2022 +0200 init diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..81ef964 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +/venv/ +/data/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..c8e5900 --- /dev/null +++ b/README.md @@ -0,0 +1,199 @@ +# Drohnen KI + +Dieses Projekt widmet sich der Erkennung von Flächenbelägen auf Drohnenbildern. + +## Installation + +Die Installation einer virtuellen Umgebung lässt sich durch folgende Zeilen Konsolencode bewerkstelligen. + +```shell +python3 -m venv venv +source venv/bin/activate +pip install 'git+https://github.com/facebookresearch/detectron2.git' +pip install -r requirements.txt +``` + +Da sich Detectron2 nicht in den Standard pip-Bibliotheken verfügbar ist, muss dieses Modul separat installiert werden. + +Nach Abschluss dieser Befehle ist die virtuelle Umgebung einsatzbereit. +Durch Ansprechen des Befehls `deactivate` gelangt man wieder aus der Umgebung heraus. +Um eine erneute Nutzung dieser virtuellen Umgebung durchzuführen, reicht es, den Befehl `source venv/bin/activate` +auszuführen. + +## Nutzung + +Dieser Abschnitt geht davon aus, dass die Umgebung korrekt installiert ist. + +### Erkennung + +Zur Erkennung einzelner Bilder kann die Datei `predict.py` verwendet werden. + +```shell +python predict.py \ + --source data/images/westend/cropped \ + --output_dir data/images/Markierungen +``` + +Die Parameter können beliebig angepasst und in allen Kombinationen übergeben werden. + +| Parameter | Standardwert | Beschreibung | +|---------------|-----------------------|-------------------------------------------------------------------------| +| category_json | | Falls andere Klassen verwendet werden, können sie hier übergeben werden | +| source | data/images/original | Ein Ordner, in dem sich alle zu bearbeitenden Bilder befinden | +| output_dir | data/images/predicted | Ein Ordner in dem alle erkannten Bilder abgelegt werden | + +### Training + +Zur Nutzung des Trainingsskripts wird eine GPU empfohlen. + +Das Training lässt sich über die Datei `train.py` starten. +Dazu dient folgender Befehl: + +```shell +python train.py --data_json data/json/train_data.json +``` + +Der Pfad `data/json/train_data.json` kann hierbei durch eine andere Trainingsdatei ersetzt werden. +Zur Anpassung weiterer Einstellungen (wie die Anzahl der zu trainierenden Epochen) stehen Konfigurationsoptionen direkt +in der `train.py`-Datei bereit. + +#### Trainingsdaten + +Die Trainingsdaten werden im folgenden JSON-Format erwartet: + +```JSON +{ + "train_images": [], + "test_images": [], + "categories": [] +} +``` + +Der Bereich `categories` ist ein Array von Strings. +Dies sind die Namen der Klassen, auf die in den image-Objekten referenziert werden. +Beispiel: + +```JSON +[ + "Baumbestand", + "Festweg", + "Pflaster", + "Wiese", + "Wasser", + "Gullydeckel" +] +``` + +Sowohl `train_images`, als auch `test_images` halten Bilddaten in dem folgenden Format vor: + +```JSON +{ + "image_id": 0, + "width": 1000, + "height": 1000, + "file_name": "data/images/westend/cropped/WestendDOP2_0_9000.tif", + "annotations": [] +} +``` + +Der Parameter `annotations` ist eine Liste aus Objekten: + +```JSON +{ + "category_name": "Gullydeckel", + "ignore": 0, + "iscrowd": 0, + "area": 1256.6370614359173, + "bbox": [ + 878.0, + 526.0, + 918.0, + 566.0 + ], + "bbox_mode": 0, + "segmentation": [ + [ + 918.0, + 546.0, + 917.61571, + 549.90181, + 916.47759, + 553.65367, + 914.62939, + 557.1114, + 912.14214, + 560.14214, + 909.1114, + 562.62939, + 905.65367, + 564.47759, + 901.90181, + 565.61571, + 898.0, + 566.0, + 894.09819, + 565.61571, + 890.34633, + 564.47759, + 886.8886, + 562.62939, + 883.85786, + 560.14214, + 881.37061, + 557.1114, + 879.52241, + 553.65367, + 878.38429, + 549.90181, + 878.0, + 546.0, + 878.38429, + 542.09819, + 879.52241, + 538.34633, + 881.37061, + 534.8886, + 883.85786, + 531.85786, + 886.8886, + 529.37061, + 890.34633, + 527.52241, + 894.09819, + 526.38429, + 898.0, + 526.0, + 901.90181, + 526.38429, + 905.65367, + 527.52241, + 909.1114, + 529.37061, + 912.14214, + 531.85786, + 914.62939, + 534.8886, + 916.47759, + 538.34633, + 917.61571, + 542.09819 + ] + ], + "category_id": 5 +} +``` + +`segmentation` folgt dem xyxy-Muster. Das bedeutet, dass ein Objekt, welches von (0, 10) über (0, 50) zu (20, 30) geht, +diese Liste generieren würde: + +```JSON +[ + 0, + 10, + 0, + 50, + 20, + 30 +] +``` + diff --git a/model_v0/model_final.pth b/model_v0/model_final.pth new file mode 100644 index 0000000..b1a816d Binary files /dev/null and b/model_v0/model_final.pth differ diff --git a/model_v1/model_final.pth b/model_v1/model_final.pth new file mode 100644 index 0000000..4e53ac4 Binary files /dev/null and b/model_v1/model_final.pth differ diff --git a/source/.dockerignore b/source/.dockerignore new file mode 100644 index 0000000..a67e991 --- /dev/null +++ b/source/.dockerignore @@ -0,0 +1,4 @@ +/data/ +/venv/ +/depr/ +/output/ \ No newline at end of file diff --git a/source/Dockerfile b/source/Dockerfile new file mode 100644 index 0000000..99916d0 --- /dev/null +++ b/source/Dockerfile @@ -0,0 +1,12 @@ +FROM python:3.10 + +RUN pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu +RUN pip install --user 'git+https://github.com/facebookresearch/detectron2.git' +RUN apt-get update +RUN apt-get install ffmpeg libsm6 libxext6 -y +RUN pip install opencv-python==4.6.0.66 pymongo==4.2.0 + +COPY ../model_v1 output/ +COPY predict.py . + +ENTRYPOINT python predict.py diff --git a/source/api.yaml b/source/api.yaml new file mode 100644 index 0000000..5f0bd9f --- /dev/null +++ b/source/api.yaml @@ -0,0 +1,56 @@ +openapi: 3.0.3 + +info: + title: Stadt MG - Drohne + description: Stadt MG - Drohne + version: 1.0.0 + +tags: + - name: Drohne + +servers: + - url: 'https://drohne.masasana.ai' + +paths: + /detect: + post: + tags: + - Drohne + requestBody: + content: + image/*: + schema: + type: string + format: binary + responses: + 200: + content: + application/json: + schema: + type: string + format: uuid + description: ok + + /image/{image_id}: + get: + tags: + - Drohne + parameters: + - name: image_id + required: true + in: path + schema: + type: string + format: uuid + responses: + 200: + content: + image/*: + schema: + type: string + format: binary + description: ok + 204: + description: The image is still in processing and no content can be provided just yet. + 404: + description: This id doesn't exist diff --git a/source/exploration.py b/source/exploration.py new file mode 100644 index 0000000..2f9e173 --- /dev/null +++ b/source/exploration.py @@ -0,0 +1,7 @@ +from prepare_training_data import crop_image + +if __name__ == '__main__': + crop_image( + "../data/images/tmp/westend/westend.jpg", + "data/images/tmp/westend/cropped" + ) diff --git a/source/predict.py b/source/predict.py new file mode 100644 index 0000000..f7f2149 --- /dev/null +++ b/source/predict.py @@ -0,0 +1,157 @@ +import argparse +import json +import logging +import os + +import cv2 +from detectron2 import model_zoo +from detectron2.config import get_cfg +from detectron2.data import MetadataCatalog +from detectron2.engine import DefaultPredictor +from detectron2.utils.visualizer import Visualizer +import numpy as np +from pymongo import MongoClient + +# Setup logging +logging.basicConfig(format='[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s') +logger = logging.getLogger("Prediction") +logger.setLevel(logging.DEBUG) + + +def predict( + category_json="", + source="data/images/original", + output_dir="data/images/predicted", + image_id="", +): + if not category_json: + categories = [ + "Baumbestand", + "Festweg", + "Pflaster", + "Wiese", + "Wasser", + "Gullydeckel" + ] + else: + with open(category_json, "r") as file: + data = json.load(file) + categories = data["categories"] + logger.debug(f"Found categories: {categories}") + + metadata = MetadataCatalog.get("predict") + metadata.set(thing_classes=categories) + + cfg = get_cfg() + cfg.merge_from_file( + model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) + + cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") + cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 + cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(categories) + cfg.MODEL.DEVICE = 'cpu' + predictor = DefaultPredictor(cfg) + + if source.startswith("mongo://"): + logger.info("MONGOOOOO") + client = MongoClient(source[8:]) + collection = client.get_database("stadtmg").get_collection("predictions") + + db_object = collection.find_one({"id": image_id}) + logger.debug(db_object) + image_binary = db_object.get("input") + nparr = np.fromstring(image_binary, np.uint8) + image = cv2.imdecode(nparr, cv2.IMREAD_COLOR) + result_image = analyse_image(image, predictor, metadata) + + is_success, im_buf_arr = cv2.imencode(".png", result_image) + byte_im = im_buf_arr.tobytes() + + collection.update_one({"id": image_id}, {"$set": { + "output": byte_im + }}) + logger.info(collection.find_one({"id": image_id}).keys()) + + else: + for path, dirs, files in os.walk(source): + for file_name in files: + image_path = os.path.join(path, file_name) + logger.debug("Opening image") + + logger.info(f"Analysing image '{image_path}'") + im = cv2.imread(image_path) + result_image = analyse_image( + im, + predictor, + metadata, + ) + out_path = os.path.join(output_dir, path[len(source):]) + os.makedirs(out_path, exist_ok=True) + cv2.imwrite(os.path.join(out_path, file_name), result_image[:, :, ::-1]) + + logger.info("Done.") + + +def analyse_image( + image, + predictor, + metadata, +): + if image is None: + return + + logger.debug("Predicting") + outputs = predictor(image) + instances = outputs["instances"] + logger.debug(instances.pred_classes) + logger.debug(instances.pred_boxes) + + logger.debug("Drawing masks") + v = Visualizer(image[:, :, ::-1], metadata=metadata, scale=1.2) + out = v.draw_instance_predictions(instances.to("cpu")) + logger.debug("Saving image") + result_image = out.get_image() + return result_image + + +def parse_opt(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--category_json", + type=str, + default="", + help="Path to a .json file containing the definition of categories", + ) + parser.add_argument( + "--source", + type=str, + default="data/images/original", + help="Path to images that should be detected" + ) + parser.add_argument( + "--output_dir", + type=str, + default="data/images/predicted", + help="Path to folder where the detected images will be placed" + ) + parser.add_argument( + "--image_id", + type=str, + default="", + help="ID of the object in MongoDB that contains the image" + ) + opt = parser.parse_args() + return opt + + +if __name__ == '__main__': + logger.info("Hello") + opt = parse_opt() + logger.info(vars(opt)) + predict(**vars(opt)) + +# if __name__ == '__main__': +# predict( +# image_dir="../data/images/tmp/westend/cropped", +# output_dir="../data/images/predicted/westend", +# ) diff --git a/source/prepare_training_data.py b/source/prepare_training_data.py new file mode 100644 index 0000000..96c3bc9 --- /dev/null +++ b/source/prepare_training_data.py @@ -0,0 +1,393 @@ +import json +import logging +import os +from itertools import product +from math import pi, cos, sin +from random import random +from uuid import uuid4 + +import geopandas as gpd +import pandas as pd +from PIL import Image, ImageDraw +from detectron2.structures import BoxMode +from shapely.geometry import box, Polygon, MultiPolygon + +# To enable loading big images +Image.MAX_IMAGE_PIXELS = None + +# Setup logging +logging.basicConfig(format='[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s') +logger = logging.getLogger("Data Preparation") +logger.setLevel(logging.DEBUG) + + +def generate_circle( + center_x=0, + center_y=0, + r=20, + n=32, +): + """ + This function creates an array of `n` xy points that are evenly spaced on a circle. + The circle is defined by its center point at ( `center_x` | `center_y` ) and the radius `r`. + + The result is given in xyxy format. + + :param center_x: Center's x coordinate + :param center_y: Center's y coordinate + :param r: Radius + :param n: Number of points on the circle + :return: List of xy-coordinates on the circle + """ + points = [ + ( + round(cos(2 * pi / n * x) * r + center_x, 5), + round(sin(2 * pi / n * x) * r + center_y, 5), + ) + for x in range(n) + ] + points = [p for x in points for p in x] + return points + + +def crop_image( + filename, + clean_image_dir, + d=1000, + save_empty=True, +): + name, ext = os.path.splitext(filename.split(os.path.sep)[-1]) + logger.debug("Opening image") + img = Image.open(filename) + logger.debug("Image opened") + w, h = img.size + + os.makedirs(clean_image_dir, exist_ok=True) + grid = product(range(0, h - h % d, d), range(0, w - w % d, d)) + for i, j in grid: + view_box = (j, i, j + d, i + d) + part = img.crop(view_box) + colors = part.getcolors() + if not save_empty and colors is not None and len(colors) <= 1: + continue + + image_out_path = os.path.join(clean_image_dir, f'{name}_{i}_{j}{ext}') + part.save(image_out_path) + logger.debug("done") + + +def westend( + filename, + clean_image_dir, + d=1000, + r=20, + save_clean_images=False, + marked_image_dir=None, + training_object=None, +): + category_translations = { + 'Asphalt': "Festweg", + 'Bepflanzte_Flaechen': "Baumbestand", + 'Beton': "Festweg", + 'Noppenpflaster': "Pflaster", + 'Pflaster': "Pflaster", + 'Platten': "Pflaster", + 'Rasen': "Wiese", + 'Rasengittersteine': "Pflaster", + 'Rippenpflaster': "Pflaster", + 'Sand': "Festweg", + 'unbefestigt': "Festweg", + } + + reference_points = pd.read_csv("../data/Referenzpunkte.csv") + name, ext = os.path.splitext(filename.split(os.path.sep)[-1]) + img = Image.open(filename) + w, h = img.size + corners = reference_points.head(4) + left = min(corners["X"]) + right = max(corners["X"]) + bottom = min(corners["Y"]) + top = max(corners["Y"]) + x_scale = (right - left) / w + y_scale = (bottom - top) / h + draw = ImageDraw.Draw(img) + + # + gullys = gpd.read_file("../data/GeoJson/Kanaldeckel_im_Bild.json") + gullys["X_"] = gullys["X"].apply(lambda x: int((x - left) / x_scale)) + gullys["Y_"] = gullys["Y"].apply(lambda y: int((y - top) / y_scale)) + # + + flaechenbelaege = gpd.read_file("../data/GeoJson/Flaechenbelaege.json") + + if training_object is None: + training_object = dict( + train_images=[], + test_images=[], + categories=[], + ) + + train_images = training_object["train_images"] + test_images = training_object["test_images"] + categories = training_object["categories"] + + grid = product(range(0, h - h % d, d), range(0, w - w % d, d)) + for i, j in grid: + if random() < 0.8: + images = train_images + else: + images = test_images + + view_box = (j, i, j + d, i + d) + colors = img.crop(view_box).getcolors() + if colors is not None and len(colors) <= 1: + # Skip empty and single-colour images + continue + + image_out_path = os.path.join(clean_image_dir, f'{name}_{i}_{j}{ext}') + if save_clean_images: + img.crop(view_box).save(image_out_path) + + annotations = [] + + # + for _, row in flaechenbelaege.iterrows(): + kind = row["Art"] + category_name = category_translations.get(kind, "") + first_obj: Polygon = row["geometry"] + + scaled = list(map( + lambda xy: ((xy[0] - left) / x_scale, (xy[1] - top) / y_scale), + first_obj.exterior.coords + )) + + belag = Polygon(scaled) + intersection = belag.intersection(box(*view_box)) + if not intersection.bounds: + continue + + if isinstance(intersection, MultiPolygon): + intersects = intersection.geoms + else: + intersects = [intersection] + for intersect in intersects: + assert isinstance(intersect, Polygon) + + global_segmentation = list(intersect.exterior.coords) + global_boundaries = intersect.bounds + if marked_image_dir is not None: + category_id = categories.index(category_name) + draw.polygon( + global_segmentation, + fill=( + 255 * (category_id & 1 << 2 > 0), + 255 * (category_id & 1 << 1 > 0), + 255 * (category_id & 1 << 0 > 0), + ), + ) + + local_segmentation = [x for xs in global_segmentation for x in xs] + local_segmentation = [[ + x - (i * (c % 2)) - (j * (1 - c % 2)) + for c, x in enumerate(local_segmentation) + ]] + local_boundaries = [ + x - (i * (c % 2)) - (j * (1 - c % 2)) + for c, x in enumerate(global_boundaries) + ] + + annotations.append(dict( + category_id=categories.index(category_name), + category_name=category_name, + ignore=0, + iscrowd=0, + bbox=local_boundaries, + bbox_mode=BoxMode.XYXY_ABS, + segmentation=local_segmentation, + )) + # + + # + gullys_in_box = gullys[ + (j <= gullys["X_"]) + & (gullys["X_"] <= j + d) + & (i <= gullys["Y_"]) + & (gullys["Y_"] <= i + d) + ][["X_", "Y_"]] + gullys_in_box["x"] = (gullys_in_box["X_"] - j) / 1 + gullys_in_box["y"] = (gullys_in_box["Y_"] - i) / 1 + if len(gullys_in_box): + for _, row in gullys_in_box.iterrows(): + circle = generate_circle( + center_x=row["x"], + center_y=row["y"], + r=r, + ) + circle = [min(max(0., p), d) for p in circle] + annotations.append(dict( + category_id=categories.index("Gullydeckel"), + category_name="Gullydeckel", + ignore=0, + iscrowd=0, + bbox=[ + max(row["x"] - r, 0), + max(row["y"] - r, 0), + min(row["x"] + r, d), + min(row["y"] + r, d), + ], + bbox_mode=BoxMode.XYXY_ABS, + segmentation=[circle], + )) + if marked_image_dir is not None: + draw.ellipse( + (row["X_"] - r, row["Y_"] - r, row["X_"] + r, row["Y_"] + r), + fill=( + 255 * (5 & 1 << 2 > 0), + 255 * (5 & 1 << 1 > 0), + 255 * (5 & 1 << 0 > 0), + ), + ) + # + + if annotations: + images.append( + dict( + image_id=str(uuid4()), + width=d, + height=d, + file_name=image_out_path, + annotations=annotations, + ) + ) + + return training_object + + +def label_studio( + json_file, + image_dir, + training_object=None, + train_split=0.8, + marked_image_dir=None, +): + colors = { + "Festweg": "grey", + "Gullydeckel": "red", + "Wiese": "green", + "Baumbestand": "brown", + "Pflaster": "orange", + "Wasser": "blue", + } + + if training_object is None: + training_object = dict( + train_images=[], + test_images=[], + categories=[], + ) + + with open(json_file, "r") as file: + data = json.load(file) + + train_images = training_object["train_images"] + test_images = training_object["test_images"] + categories = training_object["categories"] + + for obj in data: + if random() < train_split: + images = train_images + else: + images = test_images + + annotations = [] + filename = obj["data"]["image"].split("=")[1] # Might fail if filename has = in it + ext = os.path.splitext(filename)[-1] + if not ext.lower() in [".jpg", ".png", ".tif"]: + logger.debug(f"Found extension '{ext}', skipping...") + continue + image_path = image_dir + filename + image = Image.open(image_path) + width, height = image.size + draw = ImageDraw.Draw(image) + + for task in obj["annotations"]: + for annotation in task["result"]: + points = annotation["value"]["points"] + for point in points: + point[0] *= annotation["original_width"] / 100 + point[1] *= annotation["original_height"] / 100 + category_name = annotation["value"]["polygonlabels"][0] + category_id = categories.index(category_name) + + polygon = Polygon(points) + seg = [round(x, 4) for xs in polygon.exterior.coords for x in xs] + bounds = [round(x, 4) for x in polygon.bounds] + annotations.append(dict( + category_name=category_name, + ignore=0, + iscrowd=0, + bbox=bounds, + bbox_mode=BoxMode.XYXY_ABS, + segmentation=[seg], + category_id=category_id, + )) + + if marked_image_dir is not None: + draw.polygon( + polygon.exterior.coords, + fill=colors.get(category_name, "#000") + ) + if marked_image_dir is not None: + marked_path = os.path.join(marked_image_dir, filename) + os.makedirs(os.path.dirname(marked_path), exist_ok=True) + image.save(marked_path) + + if annotations: + images.append( + dict( + image_id=str(uuid4()), + width=width, + height=height, + file_name=image_path, + annotations=annotations, + ) + ) + + return training_object + + +if __name__ == '__main__': + marked_images_dir = "../data/images/marked" + + categories = [ + "Baumbestand", + "Festweg", + "Pflaster", + "Wiese", + "Wasser", + "Gullydeckel", + ] + training_object = dict( + train_images=[], + test_images=[], + categories=categories, + ) + + logger.info("Converting data from WestendDOP2.tif") + training_object = westend( + filename="../data/images/westend/WestendDOP2.tif", + clean_image_dir="../data/images/westend/cropped/", + training_object=training_object, + ) + + logger.info("Converting data from LabelStudio") + training_object = label_studio( + json_file="../data/json/project-4-at-2022-07-26-15-52-273e43ac.json", + image_dir="../data/images/label_studio/", + training_object=training_object, + ) + + logger.info("Saving results to disk") + with open(os.path.join("../data", "json", "_train_data.json"), "w") as json_file: + json.dump(training_object, json_file, indent=2) + logger.info("Done.") diff --git a/source/requirements.txt b/source/requirements.txt new file mode 100644 index 0000000..405b41f --- /dev/null +++ b/source/requirements.txt @@ -0,0 +1,15 @@ +cython==0.29.30 +geojson==2.5.0 +geopandas==0.11.0 +matplotlib==3.5.3 +Pillow~=9.2.0 +torch>=1.12.0 + +torchvision~=0.13.0 +detectron2~=0.6 +pandas~=1.4.3 + +opencv-python==4.6.0.66 +shapely==1.8.2 +pymongo==4.2.0 +kubernetes==24.2.0 \ No newline at end of file diff --git a/source/train.py b/source/train.py new file mode 100644 index 0000000..0935bcc --- /dev/null +++ b/source/train.py @@ -0,0 +1,54 @@ +import argparse +import json +import os + +from detectron2.config import get_cfg +from detectron2.data import DatasetCatalog, MetadataCatalog +from detectron2.engine import DefaultTrainer +from detectron2.model_zoo import model_zoo + + +def train( + data_json="data/json/train_data.json", +): + with open(data_json, "r") as file: + data = json.load(file) + + DatasetCatalog.register("train", lambda: data["train_images"]) + DatasetCatalog.register("test", lambda: data["test_images"]) + MetadataCatalog.get("train").set(thing_classes=data["categories"]) + + cfg = get_cfg() + cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) + cfg.DATASETS.TRAIN = ("train",) + cfg.DATASETS.TEST = ("test",) + cfg.DATALOADER.NUM_WORKERS = 4 + cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") + cfg.SOLVER.IMS_PER_BATCH = 2 + cfg.SOLVER.BASE_LR = 0.00025 + cfg.SOLVER.MAX_ITER = 10000 + cfg.SOLVER.STEPS = [] + cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 + cfg.MODEL.ROI_HEADS.NUM_CLASSES = 6 + + os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) + trainer = DefaultTrainer(cfg) + trainer.resume_or_load(resume=True) + trainer.train() + + +def parse_opt(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--data_json", + type=str, + default="data/json/train_data.json", + help="Path to a .json file containing the definition of train data", + ) + opt = parser.parse_args() + return opt + + +if __name__ == '__main__': + opt = parse_opt() + train(**vars(opt)) diff --git a/source/utils.py b/source/utils.py new file mode 100644 index 0000000..2ad0397 --- /dev/null +++ b/source/utils.py @@ -0,0 +1,2 @@ +def setup_dataset(): + pass \ No newline at end of file