Update app.py

app.py

@@ -13,6 +13,30 @@ def load_model_and_mtcnn(model_path):
     mtcnn = MTCNN(keep_all=True, device=device)
     return model, device, mtcnn
 
+def detect_and_process_skin(image_bytes):
+    """Detects faces in an image, crops the skin region, and returns it as an image object."""
+    # Load image from bytes
+    img = Image.open(io.BytesIO(image_bytes))
+    img_np = np.array(img)
+    img_rgb = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
+
+    # Detect faces in the image
+    detections = mtcnn.detect_faces(img_rgb)
+
+    # Check if any faces were detected
+    if detections:
+        x, y, width, height = detections[0]['box']
+
+        # Crop the face region
+        face_img_np = img_np[y:y+height, x:x+width]
+
+        # Convert to PIL Image for return
+        pil_img = Image.fromarray(face_img_np)
+        return pil_img
+    else:
+        # Return original image if no face was detected
+        return img
+
 # Function to preprocess the image and return both the tensor and the final PIL image for display
 def preprocess_image(image, mtcnn, device):
     processed_image = image  # Initialize with the original image
@@ -56,9 +80,14 @@ uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png
 if uploaded_file is not None:
     image = Image.open(uploaded_file).convert("RGB")
     st.image(image, caption='Uploaded Image', use_column_width=True)
+    image1 = image.getvalue()
+    image_ten = detect_and_process_skin(image1)
     image_tensor, final_image = preprocess_image(image, mtcnn, device)
-    predicted_class, probabilities = predict(image_tensor, model, device)
+    predicted_class, probabilities = predict(image_ten, model, device)
     
     st.write(f"Predicted class: {predicted_class.item()}")
     # Display the final processed image
-    st.image(final_image, caption='Processed Image', use_column_width=True)
+   # st.image(final_image, caption='Processed Image', use_column_width=True)
+    img_bytes = io.BytesIO()
+        detect_and_process_skin(image1.getvalue()).save(img_bytes, format='JPEG')
+        st.image(img_bytes.getvalue(), width=250, caption="Processed Image")