Merge pull request #881 from bmahabirbu/object-fix

fix: add new file format support (PNG)

Author: bmahabirbu

model_servers/object_detection_python/src/object_detection_server.py

@@ -8,6 +8,7 @@
 import os
 import io
 import shutil
+from typing import Optional
 
 
 app = FastAPI()
@@ -28,7 +29,7 @@
     model = AutoModelForObjectDetection.from_pretrained(model, revision=revision)
 
 class Item(BaseModel):
-    image: bytes 
+    image: bytes
 
 @app.get("/health")
 def tests_alive():
@@ -39,7 +40,7 @@ def detection(item: Item):
     b64_image = item.image
     b64_image = base64.b64decode(b64_image)
     bytes_io = io.BytesIO(b64_image)    
-    image = Image.open(bytes_io)
+    image = Image.open(bytes_io).convert("RGB")
     inputs = processor(images=image, return_tensors="pt")
     outputs = model(**inputs)
     target_sizes = torch.tensor([image.size[::-1]])
@@ -54,8 +55,9 @@ def detection(item: Item):
         label_confidence = f"Detected {model.config.id2label[label.item()]} with confidence {round(score.item(), 3)}"
         scores.append(label_confidence)
 
-    bytes_io = io.BytesIO() 
-    image.save(bytes_io, "JPEG")
+    bytes_io = io.BytesIO()
+    # Convert image format to PNG
+    image.save(bytes_io, format="PNG")
     img_bytes = bytes_io.getvalue()
     img_bytes = base64.b64encode(img_bytes).decode('utf-8')
     return {'image': img_bytes, "boxes": scores}

recipes/computer_vision/object_detection/app/object_detection_client.py

@@ -23,15 +23,22 @@
         scale_factor = 0.20
     img = img.resize((int(img.width * scale_factor) , 
                 int(img.height * scale_factor)))
-    window.image(img, use_column_width=True)  
-    # convert PIL image into bytes for post request 
+    window.image(img, use_column_width=True)
+
+    # Convert image to RGB
+    img = img.convert("RGB")
+    # Encode image
     bytes_io = io.BytesIO() 
-    if img.mode in ("RGBA", "P"): 
-        img = img.convert("RGB")
-    img.save(bytes_io, "JPEG")
+    # Convert image format to PNG
+    img.save(bytes_io, format="PNG")
     img_bytes = bytes_io.getvalue()
     b64_image = base64.b64encode(img_bytes).decode('utf-8')
-    data = {'image': b64_image}
+
+    # Prepare payload
+    data = {
+        'image': b64_image,
+    }
+
     response = requests.post(f'{endpoint}/detection', headers=headers,json=data, verify=False)
     # parse response and display outputs
     response_json = response.json()