run.py 5.52 KB
Newer Older
nikhil_rayaprolu's avatar
nikhil_rayaprolu committed
1
2
3
4
5
6
7
8
9
import random
import json
import numpy as np
import argparse
import base64
import glob
import os
import traceback
from PIL import Image
10
11
12
use_aicrowd_helpers = True
if use_aicrowd_helpers:
    import aicrowd_helpers
Ubuntu's avatar
Ubuntu committed
13
14
# from eval import evaluate
from api_ml_bc.aicrowd_bridge import evaluate
nikhil_rayaprolu's avatar
nikhil_rayaprolu committed
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

"""
Expected ENVIRONMENT Variables
* AICROWD_TEST_IMAGES_PATH : abs path to  folder containing all the test images
* AICROWD_PREDICTIONS_OUTPUT_PATH : path where you are supposed to write the output predictions.json
"""

# Configuration Variables
padding = 50
SEGMENTATION_LENGTH = 10
MAX_NUMBER_OF_ANNOTATIONS = 10
# List of valid categories to choose from
VALID_CATEGORIES =  [1565, 1010, 1085, 2053, 1310, 2578, 1154, 2620, 1566, 1151, 2939, \
    1040, 1070, 2580, 2512, 1056, 1069, 2131, 2521, 2022, 1026, 1068, 1022, 2750, 1468, \
    1013, 1078, 2738, 1061, 2618, 1311, 1163, 2504, 2498, 1788, 2099, 1032, 1505, 1058, \
    1554]


# Helper functions
def bounding_box_from_points(points):
    """
        This function only supports the `poly` format.
    """
    points = np.array(points).flatten()
    even_locations = np.arange(points.shape[0]/2) * 2
    odd_locations = even_locations + 1
    X = np.take(points, even_locations.tolist())
    Y = np.take(points, odd_locations.tolist())
    bbox = [X.min(), Y.min(), X.max()-X.min(), Y.max()-Y.min()]
    bbox = [int(b) for b in bbox]
    return bbox

def single_segmentation(image_width, image_height, number_of_points=10):
    points = []
    for k in range(number_of_points):
        # Choose a random x-coordinate
        random_x = int(random.randint(0, image_width))
        # Choose a random y-coordinate
        random_y = int(random.randint(0, image_height))
        #Flatly append them to the list of points
        points.append(random_x)
        points.append(random_y)
    return [points]

def single_annotation(image_id, number_of_points=10):
    width, height = get_image_width_height(image_id)
    _result = {}
    _result["image_id"] = image_id
    """
    Valid Categories are embedded in the annotations.json of the training set
    """
Ubuntu's avatar
Ubuntu committed
66
67
    _result["category_id"] = random.choice(VALID_CATEGORIES)

nikhil_rayaprolu's avatar
nikhil_rayaprolu committed
68
69
70
71
72
73
74
75
76
77
    _result["score"] = np.random.rand() # a random score between 0 and 1

    _result["segmentation"] = single_segmentation(width, height, number_of_points=number_of_points)
    _result["bbox"] = bounding_box_from_points(_result["segmentation"])
    return _result


def get_image_id_from_image_path(image_path):
    """
    Returns the image_id from the image_path of a file in the test set
Ubuntu's avatar
Ubuntu committed
78
    image_path : /home/testsdir/098767.jpg
nikhil_rayaprolu's avatar
nikhil_rayaprolu committed
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
    image_id : 98767
    """
    filename = os.path.basename(image_path)
    filename = filename.replace(".jpg", "")
    image_id = int(filename)
    return image_id

def gather_images(test_images_path):
    images = glob.glob(os.path.join(
        test_images_path, "*.jpg"
    ))
    return images

def gather_image_ids(test_images_path):
    images = gather_images(test_images_path)
    filenames = [os.path.basename(image_path).replace(".jpg","") for image_path in images]
    image_ids = [int(x) for x in filenames]
    return image_ids

def get_image_path(image_id):
    test_images_path = os.getenv("AICROWD_TEST_IMAGES_PATH", False)
    return "{}.jpg".format(os.path.join(test_images_path, str(image_id).zfill(6)))

def get_image_width_height(image_id):
    image_path = get_image_path(image_id)
    im = Image.open(image_path)
    width, height = im.size
    im.close()
    return width, height

def gather_input_output_path():
    test_images_path = os.getenv("AICROWD_TEST_IMAGES_PATH", False)
    assert test_images_path != False, "Please provide the path to the test images using the environment variable : AICROWD_TEST_IMAGES_PATH"

    predictions_output_path = os.getenv("AICROWD_PREDICTIONS_OUTPUT_PATH", False)
    assert predictions_output_path != False, "Please provide the output path (for writing the predictions.json) using the environment variable : AICROWD_PREDICTIONS_OUTPUT_PATH"

    return test_images_path, predictions_output_path

def run():
    ########################################################################
    # Register Prediction Start
    ########################################################################
122
123
    if use_aicrowd_helpers:
        aicrowd_helpers.execution_start()
nikhil_rayaprolu's avatar
nikhil_rayaprolu committed
124
125
126
127
128
129
130
131
132
133
134
135
136
137

    ########################################################################
    # Gather Input and Output paths from environment variables
    ########################################################################
    test_images_path, predictions_output_path = gather_input_output_path()

    ########################################################################
    # Gather Image IDS
    ########################################################################
    image_ids = gather_image_ids(test_images_path)

    ########################################################################
    # Generate Predictions
    ########################################################################
nikhil_rayaprolu's avatar
nikhil_rayaprolu committed
138
    evaluate(test_images_path, predictions_output_path)
nikhil_rayaprolu's avatar
nikhil_rayaprolu committed
139
140
141
    ########################################################################
    # Register Prediction Complete
    ########################################################################
142
143
144
145
    if use_aicrowd_helpers:
        aicrowd_helpers.execution_success({
            "predictions_output_path" : predictions_output_path
        })
nikhil_rayaprolu's avatar
nikhil_rayaprolu committed
146
147
148
149
150
151
152

if __name__ == "__main__":
    try:
        run()
    except Exception as e:
        error = traceback.format_exc()
        print(error)
153
154
        if use_aicrowd_helpers:
            aicrowd_helpers.execution_error(error)