Update seasonal_color_analysis/fe.py

seasonal_color_analysis/fe.py

@@ -1,6 +1,6 @@
 import os
 from io import BytesIO
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageOps
 import uuid
 import json
 import datetime
@@ -58,19 +58,18 @@ def get_season_description(season: str) -> tuple[str, str]:
 
 
 @st.cache_data
-def predict(img_bytes: bytes) -> tuple[np.ndarray | None, dict[str, float], np.ndarray, np.ndarray]:
-    with Image.open(BytesIO(img_bytes)) as img:
-        batch_boxes, proba_dicts, np_season_embeddings, np_facenet_embeddings = st.session_state["classifier"].predict([img.convert("RGB")])
-        return batch_boxes[0], proba_dicts[0], np_season_embeddings[0], np_facenet_embeddings[0]
+def predict(np_img: np.ndarray) -> tuple[np.ndarray | None, dict[str, float], np.ndarray, np.ndarray]:
+    img = Image.fromarray(np_img)
+    batch_boxes, proba_dicts, np_season_embeddings, np_facenet_embeddings = st.session_state["classifier"].predict([img])
+    return batch_boxes[0], proba_dicts[0], np_season_embeddings[0], np_facenet_embeddings[0]
 
 
 @st.cache_data
-def draw_bbox(img_bytes: bytes, bbox: np.ndarray) -> Image:
-    with Image.open(BytesIO(img_bytes)) as img:
-        _img = img.copy()
-        draw = ImageDraw.Draw(_img)
-        draw.rectangle(bbox.tolist(), outline="green", width=img.size[0] // 100)
-        return _img
+def draw_bbox(np_img: np.ndarray, bbox: np.ndarray) -> np.ndarray:
+    img = Image.fromarray(np_img)
+    draw = ImageDraw.Draw(img)
+    draw.rectangle(bbox.tolist(), outline="green", width=img.size[0] // 100)
+    return np.array(img)
 
 
 @st.cache_data
@@ -179,12 +178,17 @@ img_stream = st.file_uploader(
 )
 
 if img_stream is not None:
-    img_bytes = img_stream.getvalue()
-    bbox, proba_dict, np_season_embedding, np_facenet_embedding = predict(img_bytes)
+    with Image.open(img_stream) as img:
+        np_img = np.array(
+            ImageOps.exif_transpose(
+                img.convert("RGB")
+            )
+        )
+    bbox, proba_dict, np_season_embedding, np_facenet_embedding = predict(np_img)
     if bbox is None:
         col1, col2 = st.columns(2)
         with col1:
-            st.image(img_bytes, caption="Your image")
+            st.image(np_img, caption="Your image")
         with col2:
             st.write("⚠️\n\nIt was not possibile to detect any face in your image, try uploading another one\n\n⚠️")
     else:
@@ -196,10 +200,10 @@ if img_stream is not None:
         second_most_likely_prob = np.sort(probs)[-2]
         col1, col2 = st.columns(2)
         with col1:
-            st.image(img_bytes, caption="Your image")
+            st.image(np_img, caption="Your image")
         with col2:
-            img_w_bbox = draw_bbox(img_bytes, bbox)
-            st.image(np.array(img_w_bbox), caption="Detected face")
+            np_img_w_bbox = draw_bbox(np_img, bbox)
+            st.image(np_img_w_bbox, caption="Detected face")
         
         st.header("Your result")