diff --git a/src/mot/serving/inference.py b/src/mot/serving/inference.py index 2c3dfdb..08b53c4 100644 --- a/src/mot/serving/inference.py +++ b/src/mot/serving/inference.py @@ -78,8 +78,16 @@ def process_video(): return images_folder def process_zip(): - images_folder = "{}_split".format(filename) - zipfile.ZipFile(full_filepath).extractall(images_folder) + images_folder = os.path.join(upload_folder, "{}_split".format(filename)) + with zipfile.ZipFile(full_filepath, 'r') as zip_obj: + file_names = zip_obj.namelist() + for file_name in file_names: + zip_obj.extract(file_name, images_folder) + if os.path.basename(file_name): + shutil.move( + os.path.join(images_folder, file_name), + os.path.join(images_folder, os.path.basename(file_name)) + ) return images_folder if file.mimetype == "": @@ -103,13 +111,15 @@ def process_zip(): logger.info("{} images to analyze on {} CPUs.".format(len(image_paths), CPU_COUNT)) try: inference_outputs = [] - with multiprocessing.Pool(CPU_COUNT) as p: - inference_outputs = list( - tqdm( - p.imap(_process_image, image_paths), - total=len(image_paths), - ) - ) + for image_path in image_paths: + inference_outputs.append(_process_image(image_path)) + # with multiprocessing.Pool(CPU_COUNT) as p: + # inference_outputs = list( + # tqdm( + # p.imap(_process_image, image_paths), + # total=len(image_paths), + # ) + # ) except ValueError as e: return {"error": str(e)} logger.info("Object detection on video {} finished.".format(full_filepath)) @@ -223,4 +233,5 @@ def predict_and_format_image( "score": score, } detected_trash.append(trash_json) + return detected_trash diff --git a/tests/tests_mot/serving/test_app.py b/tests/tests_mot/serving/test_app.py index 595b673..a703b41 100644 --- a/tests/tests_mot/serving/test_app.py +++ b/tests/tests_mot/serving/test_app.py @@ -51,7 +51,7 @@ def test_app_post_tracking(mock_server_result, tmpdir): zip_path += ".zip" app_folder = os.path.join(tmpdir, "app_folder") - os.makedirs(os.path.join(tmpdir, "app_folder")) + os.makedirs(app_folder) app.config["UPLOAD_FOLDER"] = app_folder with app.test_client() as c: response = c.post("/tracking", data={"file": (open(zip_path, "rb"), "toto.zip")}) @@ -111,19 +111,11 @@ def test_app_post_image(mock_server_result, tmpdir): output = response.get_json() assert response.status_code == 200 expected_output = { - "detected_trash": - [ - { - "box": [0.0, 0.0, 0.1, 0.05], - "label": "bottles", - "score": 0.7 - }, { - "box": [0.0, 0.0, 0.1, 0.1], - "label": "fragments", - "score": 0.6 - } - ], - "image_path": os.path.join(tmpdir, TMP_IMAGE_NAME) + "detected_trash": [{ + "box": [0.0, 0.0, 0.1, 0.05], + "label": "bottles", + "score": 0.7 + }] } assert output == expected_output diff --git a/tests/tests_mot/serving/test_inference.py b/tests/tests_mot/serving/test_inference.py index 74c0629..c3deef1 100644 --- a/tests/tests_mot/serving/test_inference.py +++ b/tests/tests_mot/serving/test_inference.py @@ -3,11 +3,13 @@ from unittest import mock import cv2 +import ffmpeg import numpy as np import pytest -from werkzeug import FileStorage +from werkzeug.datastructures import FileStorage -from mot.serving.inference import handle_post_request, predict_and_format_image, process_image +from mot.serving.inference import (_process_image, detect_and_track_images, + predict_and_format_image) HOME = os.path.expanduser("~") PATH_TO_TEST_VIDEO = os.path.join(HOME, ".mot/tests/test_video.mp4") @@ -18,6 +20,7 @@ def mock_post_tensorpack_localizer(*args, **kwargs): boxes = [[0, 0, 120, 40], [0, 0, 120, 80]] scores = [[0.7, 0.01, 0.01], [0.1, 0.1, 0.6]] classes = [1, 3] + class Response(mock.Mock): json_text = { 'outputs': @@ -40,104 +43,14 @@ def json(self): @mock.patch('requests.post', side_effect=mock_post_tensorpack_localizer) -def test_handle_post_request_image(mock_server_result): - image = np.ones((300, 200, 3)) - data = "{}".format(json.dumps({"image": image.tolist()})) - data = data.encode('utf-8') - m = mock.MagicMock() # here we mock flask.request - m.data = data - with mock.patch("mot.serving.inference.request", m): - output = handle_post_request() - expected_output = { - "detected_trash": - [ - { - "box": [0.0, 0.0, 0.1, 0.05], - "label": "bottles", - "score": 0.71 - }, { - "box": [0.0, 0.0, 0.1, 0.1], - "label": "fragments", - "score": 0.71 - } - ] - } - assert output == expected_output - - -def test_handle_post_request_wrong_image(): - data = "{}".format(json.dumps({"image:0": [[[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]]})) - data = data.encode('utf-8') - m = mock.MagicMock() # here we mock flask.request - m.data = data - with pytest.raises(ValueError): - with mock.patch("mot.serving.inference.request", m): - output = handle_post_request() - - -def test_handle_post_request_video(): - # TODO test good behavior when implemented - data = "{}".format(json.dumps({"video": [[[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]]})) - data = data.encode('utf-8') - m = mock.MagicMock() # here we mock flask.request - m.data = data - with pytest.raises(NotImplementedError): - with mock.patch("mot.serving.inference.request", m): - output = handle_post_request() - - -@mock.patch('requests.post', side_effect=mock_post_tensorpack_localizer) -def test_handle_post_request_file_image(mock_server_result, tmpdir): - data = np.ones((300, 200, 3)) - filename = "test.jpg" - filepath = os.path.join(tmpdir, filename) - cv2.imwrite(filepath, data) - m = mock.MagicMock() - files = {"file": FileStorage(open(filepath, "rb"), content_type='image/jpg')} - m.files = files - - upload_folder = os.path.join(tmpdir, "upload") - os.mkdir(upload_folder) - with open(os.path.join(upload_folder, filename), "w+") as f: - f.write( - "this file should be deleted by the handle post request to be replaced" - " by the image we are uploading." - ) - - with mock.patch("mot.serving.inference.request", m): - output = handle_post_request(upload_folder=upload_folder) - - expected_output = { - "image": - output["image"], - "detected_trash": - [ - { - "box": [0.0, 0.0, 0.1, 0.05], - "label": "bottles", - "score": 0.71 - }, { - "box": [0.0, 0.0, 0.1, 0.1], - "label": "fragments", - "score": 0.71 - } - ] - } - assert output == expected_output - - -@mock.patch('requests.post', side_effect=mock_post_tensorpack_localizer) -def test_handle_post_request_file_video(mock_server_result, tmpdir): - m = mock.MagicMock() - files = {"file": FileStorage(open(PATH_TO_TEST_VIDEO, "rb"), content_type='video/mkv')} - m.files = files - m.form = {"fps": 2, "foo": "bar"} +def test_detect_and_track_images_video(mock_server_result, tmpdir): split_frames_folder = os.path.join( tmpdir, "{}_split".format(os.path.basename(PATH_TO_TEST_VIDEO)) ) os.mkdir(split_frames_folder) # this folder should be deleted by handle post request - with mock.patch("mot.serving.inference.request", m): - output = handle_post_request(upload_folder=str(tmpdir)) + output = detect_and_track_images( + FileStorage(open(PATH_TO_TEST_VIDEO, "rb")), upload_folder=str(tmpdir), fps=2 + ) assert len(output["detected_trash"]) == 2 assert "id" in output["detected_trash"][0] @@ -151,13 +64,10 @@ def test_handle_post_request_file_video(mock_server_result, tmpdir): assert "video_id" in output @mock.patch('requests.post', side_effect=mock_post_tensorpack_localizer) -def test_handle_post_request_file_zip(mock_server_result, tmpdir): - m = mock.MagicMock() - files = {"file": FileStorage(open(PATH_TO_TEST_ZIP, "rb"), content_type='application/zip')} - m.files = files - m.form = {"fps": 2, "foo": "bar"} - with mock.patch("mot.serving.inference.request", m): - output = handle_post_request(upload_folder=str(tmpdir)) +def test_detect_and_track_images_zip(mock_server_result, tmpdir): + output = detect_and_track_images( + FileStorage(open(PATH_TO_TEST_ZIP, "rb")), upload_folder=str(tmpdir), fps=2 + ) assert len(output["detected_trash"]) == 2 assert "id" in output["detected_trash"][0] @@ -176,13 +86,9 @@ def test_handle_post_request_file_other(tmpdir): filepath = os.path.join(tmpdir, filename) with open(filepath, "w") as f: f.write("mock data") - m = mock.MagicMock() - files = {"file": FileStorage(open(filepath, "rb"), content_type='poulet/pdf')} - m.files = files upload_folder = os.path.join(tmpdir, "upload_folder") - with pytest.raises(NotImplementedError): - with mock.patch("mot.serving.inference.request", m): - output = handle_post_request(upload_folder=upload_folder) + with pytest.raises(ffmpeg._run.Error): + detect_and_track_images(FileStorage(open(filepath, "rb")), upload_folder=upload_folder) assert os.path.isdir( upload_folder ) # the upload_folder should be created by handle post request @@ -193,59 +99,54 @@ def test_process_image(mock_server_result, tmpdir): image = np.ones((300, 200, 3)) image_path = os.path.join(tmpdir, "image.jpg") cv2.imwrite(image_path, image) - predictions = process_image(image_path) - assert predictions == { + output = _process_image(image_path) + expected_output = { 'output/boxes:0': [[0, 0, 0.1, 0.05], [0, 0, 0.1, 0.1]], - 'output/scores:0': [[0.71, 0.1, 0.1], [0.2, 0.05, 0.71]], + 'output/scores:0': [[0.7, 0.01, 0.01], [0.1, 0.1, 0.6]], 'output/labels:0': [1, 3] } + assert output == expected_output @mock.patch('requests.post', side_effect=mock_post_tensorpack_localizer) def test_predict_and_format_image(mock_server_result, tmpdir): image = np.ones((300, 200, 3)) predictions = predict_and_format_image(image) - assert predictions == [ - { - "box": [0, 0, 0.1, 0.05], - "label": "bottles", - "score": 0.71 - }, { - "box": [0, 0, 0.1, 0.1], - "label": "fragments", - "score": 0.71 - } - ] + assert predictions == [{"box": [0, 0, 0.1, 0.05], "label": "bottles", "score": 0.7}] # tesing with different class names class_names = ["others", "fragments", "chicken", "bottles"] - predictions = predict_and_format_image(image, class_names) - assert predictions == [ + output = predict_and_format_image(image, class_names) + expected_output = [ { "box": [0, 0, 0.1, 0.05], "label": "others", - "score": 0.71 + "score": 0.7 }, { "box": [0, 0, 0.1, 0.1], "label": "chicken", - "score": 0.71 + "score": 0.6 } ] + assert output == expected_output # testing with different thresholds class_to_threshold = {"bottles": 0.8, "others": 0.3, "fragments": 0.3} - predictions = predict_and_format_image(image, class_to_threshold=class_to_threshold) - assert predictions == [{"box": [0, 0, 0.1, 0.1], "label": "fragments", "score": 0.71}] + output = predict_and_format_image(image, class_to_threshold=class_to_threshold) + expected_output = [{"box": [0, 0, 0.1, 0.1], "label": "fragments", "score": 0.6}] + assert output == expected_output # testing with different thresholds class_to_threshold = {"bottles": 0.8, "others": 0.3, "fragments": 0.8} - predictions = predict_and_format_image(image, class_to_threshold=class_to_threshold) - assert predictions == [] + output = predict_and_format_image(image, class_to_threshold=class_to_threshold) + expected_output = [] + assert output == expected_output def mock_post_tensorpack_localizer_error(*args, **kwargs): + class Response(mock.Mock): - json_text = {'error': "¯\(°_o)/¯"} + json_text = {'error': "¯\(°_o)/¯"} @property def text(self): @@ -261,33 +162,28 @@ def json(self): @mock.patch('requests.post', side_effect=mock_post_tensorpack_localizer_error) def test_handle_post_request_file_error(mock_server_result, tmpdir): # videos - m = mock.MagicMock() - files = {"file": FileStorage(open(PATH_TO_TEST_VIDEO, "rb"), content_type='video/mkv')} - m.files = files - m.form = {"fps": 2, "foo": "bar"} split_frames_folder = os.path.join( tmpdir, "{}_split".format(os.path.basename(PATH_TO_TEST_VIDEO)) ) os.mkdir(split_frames_folder) # this folder should be deleted by handle post request - with mock.patch("mot.serving.inference.request", m): - outputs = handle_post_request(upload_folder=str(tmpdir)) - assert "error" in outputs - assert outputs["error"].endswith("¯\(°_o)/¯") + outputs = detect_and_track_images( + FileStorage(open(PATH_TO_TEST_VIDEO, "rb")), upload_folder=str(tmpdir) + ) + assert "error" in outputs + assert outputs["error"].endswith("¯\(°_o)/¯") # images data = np.ones((300, 200, 3)) filename = "test.jpg" filepath = os.path.join(tmpdir, filename) cv2.imwrite(filepath, data) - m = mock.MagicMock() - files = {"file": FileStorage(open(filepath, "rb"), content_type='image/jpg')} - m.files = files - with mock.patch("mot.serving.inference.request", m): - outputs = handle_post_request(upload_folder=str(tmpdir)) - assert "error" in outputs - assert outputs["error"].endswith("¯\(°_o)/¯") + outputs = detect_and_track_images(FileStorage(open(filepath, "rb")), upload_folder=str(tmpdir)) + assert "error" in outputs + assert outputs["error"].endswith("¯\(°_o)/¯") + def mock_post_tensorpack_localizer_unknown(*args, **kwargs): + class Response(mock.Mock): json_text = {'unknown': "¯\_(ツ)_/¯"} @@ -305,28 +201,21 @@ def json(self): @mock.patch('requests.post', side_effect=mock_post_tensorpack_localizer_unknown) def test_handle_post_request_file_unknwon(mock_server_result, tmpdir): # videos - m = mock.MagicMock() - files = {"file": FileStorage(open(PATH_TO_TEST_VIDEO, "rb"), content_type='video/mkv')} - m.files = files - m.form = {"fps": 2, "foo": "bar"} split_frames_folder = os.path.join( tmpdir, "{}_split".format(os.path.basename(PATH_TO_TEST_VIDEO)) ) os.mkdir(split_frames_folder) # this folder should be deleted by handle post request - with mock.patch("mot.serving.inference.request", m): - outputs = handle_post_request(upload_folder=str(tmpdir)) - assert "error" in outputs - assert outputs["error"].endswith("{'unknown': '¯\\\\_(ツ)_/¯'}") + outputs = detect_and_track_images( + FileStorage(open(PATH_TO_TEST_VIDEO, "rb")), upload_folder=str(tmpdir) + ) + assert "error" in outputs + assert outputs["error"].endswith("{'unknown': '¯\\\\_(ツ)_/¯'}") # images data = np.ones((300, 200, 3)) filename = "test.jpg" filepath = os.path.join(tmpdir, filename) cv2.imwrite(filepath, data) - m = mock.MagicMock() - files = {"file": FileStorage(open(filepath, "rb"), content_type='image/jpg')} - m.files = files - with mock.patch("mot.serving.inference.request", m): - outputs = handle_post_request(upload_folder=str(tmpdir)) - assert "error" in outputs - assert outputs["error"].endswith("{'unknown': '¯\\\\_(ツ)_/¯'}") + outputs = detect_and_track_images(FileStorage(open(filepath, "rb")), upload_folder=str(tmpdir)) + assert "error" in outputs + assert outputs["error"].endswith("{'unknown': '¯\\\\_(ツ)_/¯'}")