diff --git a/backend/src/router.py b/backend/src/router.py index 0bc135aa..1376540f 100644 --- a/backend/src/router.py +++ b/backend/src/router.py @@ -7,6 +7,7 @@ from basegun_ml.classification import get_typology from basegun_ml.exceptions import MissingCard, MissingGun from basegun_ml.measure import get_lengths +from basegun_ml.ocr import is_alarm_weapon from fastapi import ( APIRouter, BackgroundTasks, @@ -238,3 +239,16 @@ async def expert_contact( """, attachements=files, ) + + +@router.post("/identification-alarm-gun") +async def image_alarm_gun( + image: UploadFile = File(...), +): + try: + img_bytes = image.file.read() + # Process image with ML models + return {"is_alarm_model": is_alarm_weapon(img_bytes) == "Alarm_model"} + + except Exception as e: + return {"is_alarm_model": False, "exception": e.__class__.__name__} diff --git a/backend/tests/images/alarm_gun.jpg b/backend/tests/images/alarm_gun.jpg new file mode 100644 index 00000000..936f38fc Binary files /dev/null and b/backend/tests/images/alarm_gun.jpg differ diff --git a/backend/tests/epaule_a_levier_sous_garde.jpg b/backend/tests/images/epaule_a_levier_sous_garde.jpg similarity index 100% rename from backend/tests/epaule_a_levier_sous_garde.jpg rename to backend/tests/images/epaule_a_levier_sous_garde.jpg diff --git a/backend/tests/images/low_quality.jpg b/backend/tests/images/low_quality.jpg new file mode 100644 index 00000000..3e5900cf Binary files /dev/null and b/backend/tests/images/low_quality.jpg differ diff --git a/backend/tests/images/no_text.jpg b/backend/tests/images/no_text.jpg new file mode 100644 index 00000000..b057df5e Binary files /dev/null and b/backend/tests/images/no_text.jpg differ diff --git a/backend/tests/revolver.jpg b/backend/tests/images/revolver.jpg similarity index 100% rename from backend/tests/revolver.jpg rename to backend/tests/images/revolver.jpg diff --git a/backend/tests/test_api.py b/backend/tests/test_api.py index 255386a5..84977e4d 100644 --- a/backend/tests/test_api.py +++ b/backend/tests/test_api.py @@ -44,7 +44,9 @@ def test_version(self, client): def test_upload(self, client): """Checks that the file upload works properly""" create_bucket() - path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "revolver.jpg") + path = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "images/revolver.jpg" + ) with open(path, "rb") as f: r = client.post( @@ -99,7 +101,7 @@ def test_headers(self, client): class TestUpload: def test_revolver_without_card(self, client): - with open("./tests/revolver.jpg", "rb") as f: + with open("./tests/images/revolver.jpg", "rb") as f: response = client.post( "/api/upload", files={"image": f}, @@ -115,7 +117,7 @@ def test_revolver_without_card(self, client): assert response.data["conf_card"] is None def test_semi_auto_without_card(self, client): - with open("./tests/epaule_a_levier_sous_garde.jpg", "rb") as f: + with open("./tests/images/epaule_a_levier_sous_garde.jpg", "rb") as f: response = client.post( "/api/upload", files={"image": f}, @@ -134,7 +136,7 @@ def test_semi_auto_without_card(self, client): class TestExpertContact: def test_success(self, faker, client): client.authenticate() - with open("./tests/revolver.jpg", "rb") as f: + with open("./tests/images/revolver.jpg", "rb") as f: response = client.post( "/api/expert-contact", files=[ @@ -178,3 +180,37 @@ def test_403(self, client): response = client.get("/api/contact-details") response.data = response.json() assert response.status_code == 403 + + +class TestAlarmGunUpload: + def test_alarm_gun(self, client): + with open("./tests/images/alarm_gun.jpg", "rb") as f: + response = client.post( + "/api/identification-alarm-gun", + files={"image": f}, + ) + response.data = response.json() + assert response.status_code == 200 + assert response.data["is_alarm_model"] is True + + def test_bad_quality(self, client): + with open("./tests/images/low_quality.jpg", "rb") as f: + response = client.post( + "/api/identification-alarm-gun", + files={"image": f}, + ) + response.data = response.json() + assert response.status_code == 200 + assert response.data["is_alarm_model"] is False + assert response.data["exception"] == "LowQuality" + + def test_missing_text(self, client): + with open("./tests/images/no_text.jpg", "rb") as f: + response = client.post( + "/api/identification-alarm-gun", + files={"image": f}, + ) + response.data = response.json() + assert response.status_code == 200 + assert response.data["is_alarm_model"] is False + assert response.data["exception"] == "MissingText"