Browse Source

Merge pull request #235 from MadMaxInfinity/patch-2

Update cli.py to make POST method work
Daniel Gatis 3 years ago
parent
commit
25298b15ed
1 changed files with 30 additions and 4 deletions
  1. 30 4
      rembg/cli.py

+ 30 - 4
rembg/cli.py

@@ -9,7 +9,7 @@ import click
 import filetype
 import filetype
 import uvicorn
 import uvicorn
 from asyncer import asyncify
 from asyncer import asyncify
-from fastapi import Depends, FastAPI, File, Query
+from fastapi import Depends, FastAPI, File, Query, Form
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.cors import CORSMiddleware
 from starlette.responses import Response
 from starlette.responses import Response
 from tqdm import tqdm
 from tqdm import tqdm
@@ -294,10 +294,10 @@ def s(port: int, log_level: str) -> None:
             ),
             ),
             a: bool = Query(default=False, description="Enable Alpha Matting"),
             a: bool = Query(default=False, description="Enable Alpha Matting"),
             af: int = Query(
             af: int = Query(
-                default=240, ge=0, description="Alpha Matting (Foreground Threshold)"
+                default=240, ge=0, le=255, description="Alpha Matting (Foreground Threshold)"
             ),
             ),
             ab: int = Query(
             ab: int = Query(
-                default=10, ge=0, description="Alpha Matting (Background Threshold)"
+                default=10, ge=0, le=255, description="Alpha Matting (Background Threshold)"
             ),
             ),
             ae: int = Query(
             ae: int = Query(
                 default=10, ge=0, description="Alpha Matting (Erode Structure Size)"
                 default=10, ge=0, description="Alpha Matting (Erode Structure Size)"
@@ -311,6 +311,32 @@ def s(port: int, log_level: str) -> None:
             self.ae = ae
             self.ae = ae
             self.om = om
             self.om = om
 
 
+    class CommonQueryPostParams:
+        def __init__(
+            self,
+            model: ModelType = Form(
+                default=ModelType.u2net,
+                description="Model to use when processing image",
+            ),
+            a: bool = Form(default=False, description="Enable Alpha Matting"),
+            af: int = Form(
+                default=240, ge=0, le=255, description="Alpha Matting (Foreground Threshold)"
+            ),
+            ab: int = Form(
+                default=10, ge=0, le=255, description="Alpha Matting (Background Threshold)"
+            ),
+            ae: int = Form(
+                default=10, ge=0, description="Alpha Matting (Erode Structure Size)"
+            ),
+            om: bool = Form(default=False, description="Only Mask"),
+        ):
+            self.model = model
+            self.a = a
+            self.af = af
+            self.ab = ab
+            self.ae = ae
+            self.om = om
+            
     def im_without_bg(content: bytes, commons: CommonQueryParams) -> Response:
     def im_without_bg(content: bytes, commons: CommonQueryParams) -> Response:
         return Response(
         return Response(
             remove(
             remove(
@@ -355,7 +381,7 @@ def s(port: int, log_level: str) -> None:
             default=...,
             default=...,
             description="Image file (byte stream) that has to be processed.",
             description="Image file (byte stream) that has to be processed.",
         ),
         ),
-        commons: CommonQueryParams = Depends(),
+        commons: CommonQueryPostParams = Depends(),
     ):
     ):
         return await asyncify(im_without_bg)(file, commons)
         return await asyncify(im_without_bg)(file, commons)